mujoco

// Copyright 2021 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "engine/engine_collision_convex.h"

#include <float.h>
#include <stddef.h>

#include <ccd/ccd.h>
#include <ccd/vec3.h>

#include <mujoco/mjdata.h>
#include <mujoco/mjmacro.h>
#include <mujoco/mjmodel.h>
#include "engine/engine_collision_gjk.h"
#include "engine/engine_macro.h"
#include "engine/engine_memory.h"
#include "engine/engine_inline.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
#include "engine/engine_util_spatial.h"

#define mjMINVAL2 (mjMINVAL * mjMINVAL)

// CCD internal buffer used for batched processing; if NULL, stack memory allocated on each
// mjc_penetration call
static mjTHREADLOCAL void* ccd_buffer = NULL;

// set CCD internal buffer
void mjc_setCCDBuffer(void* buffer) {
  ccd_buffer = buffer;
}

// ccd prism first dir
static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) {
  ccdVec3Set(vec, 0, 0, 1);
}

// wrapper around libccd; returns number of collisions found
static inline int _libccd_wrapper(const mjModel* m, mjCCDObj* obj1, mjCCDObj* obj2,
                                  mjPreContact* con, mjtNum margin) {
  ccd_t ccd;
  CCD_INIT(&ccd);
  ccd.mpr_tolerance = m->opt.ccd_tolerance;
  ccd.epa_tolerance = m->opt.ccd_tolerance;  // use MPR tolerance for EPA
  ccd.max_iterations = m->opt.ccd_iterations;
  ccd.support1 = mjccd_support;
  ccd.support2 = mjccd_support;
  ccd.center1 = mjccd_center;
  ccd.center2 = mjccd_center;
  if (obj1->geom_type == mjGEOM_HFIELD || obj2->geom_type == mjGEOM_HFIELD) {
    ccd.first_dir = prism_firstdir;
  } else {
    ccd.first_dir = ccdFirstDirDefault;
  }

  ccd_real_t ccd_depth;
  ccd_vec3_t ccd_dir, ccd_pos;
  int ret = ccdMPRPenetration(obj1, obj2, &ccd, &ccd_depth, &ccd_dir, &ccd_pos);
  if (ret == 0) {
    if (ccdVec3Eq(&ccd_dir, ccd_vec3_origin)) {
      return 0;
    }
    con[0].dist = margin - ccd_depth;
    mji_copy3(con[0].normal, ccd_dir.v);
    mji_copy3(con[0].pos, ccd_pos.v);
    mji_zero3(con[0].tangent);
    return 1;
  }
  return 0;
}


// find penetration info between two geoms; returns number of collisions found
static int mjc_penetration(const mjModel* m, mjData* d, mjCCDObj* obj1, mjCCDObj* obj2,
                           mjPreContact* con, int ncon, mjtNum margin) {
  if (mjDISABLED(mjDSBL_NATIVECCD)) {
    return _libccd_wrapper(m, obj1, obj2, con, margin);
  }

  // nativeccd
  mjCCDConfig config;
  mjCCDStatus status;
  mjtNum dist;
  int nwitness = 0;
  void* buffer = ccd_buffer;

  // set config
  config.max_iterations = m->opt.ccd_iterations;
  config.tolerance = m->opt.ccd_tolerance;
  config.max_contacts = ncon;
  config.dist_cutoff = 0;  // no geom distances needed
  config.npolygonmax = m->npolygonmax;
  config.nmeshdegmax = m->nmeshdegmax;
  if (buffer) {
    config.buffer = buffer;
  } else {
    mj_markStack(d);
    int npolygonmax = mjDISABLED(mjDSBL_MULTICCD) ? 0 : m->npolygonmax;
    int nmeshdegmax = mjDISABLED(mjDSBL_MULTICCD) ? 0 : m->nmeshdegmax;
    config.buffer = mj_stackAllocByte(d, mjc_ccdSize(npolygonmax, nmeshdegmax,
                                                     config.max_iterations), sizeof(mjtNum));
  }

  if ((dist = mjc_ccd(&config, &status, obj1, obj2)) < 0) {
    nwitness = status.nx;
    for (int i = 0; i < nwitness; i++) {
      con[i].dist = margin + dist;
      con[i].pos[0] = 0.5*(status.x1[3*i + 0] + status.x2[3*i + 0]);
      con[i].pos[1] = 0.5*(status.x1[3*i + 1] + status.x2[3*i + 1]);
      con[i].pos[2] = 0.5*(status.x1[3*i + 2] + status.x2[3*i + 2]);
      mji_sub3(con[i].normal, status.x1 + 3*i, status.x2 + 3*i);
      mju_normalize3(con[i].normal);
      mji_zero3(con[i].tangent);
    }
  }
  if (!buffer) {
    mj_freeStack(d);
  }
  return nwitness;
}


// ccd center function
void mjccd_center(const void *obj, ccd_vec3_t *center) {
  mjc_center(center->v, (const mjCCDObj*) obj);
}

// center function for convex collision algorithms
void mjc_center(mjtNum res[3], const mjCCDObj *obj) {
  int g = obj->geom;
  int f = obj->flex;
  int e = obj->elem;
  int v = obj->vert;

  if (obj->geom_type == mjGEOM_HFIELD) {
    mju_zero3(res);
    for (int i=0; i < 6; i++) {
      mji_addTo3(res, obj->data.hfield.prism[i]);
     }
    mju_scl3(res, res, 1.0/6.0);
    return;
  }

  // return geom position
  if (g >= 0) {
    mji_copy3(res, obj->pos);
    return;
  }

  // return flex element position
  if (e >= 0) {
    mji_copy3(res, obj->data.flex.aabb + 6*(obj->data.flex.elemadr[f]+e));
    return;
  }

  // return flex vertex position
  if (f >= 0) {
    mji_copy3(res, obj->data.flex.vert_xpos + 3*(obj->data.flex.vertadr[f]+v));
    return;
  }
}

// ------------------------------------ Support functions -----------------------------------------

// transform a vector from global to local frame
static inline void mulMatTVec3(mjtNum res[3], const mjtNum mat[9], const mjtNum dir[3]) {
  // perform matT * dir
  res[0] = mat[0]*dir[0] + mat[3]*dir[1] + mat[6]*dir[2];
  res[1] = mat[1]*dir[0] + mat[4]*dir[1] + mat[7]*dir[2];
  res[2] = mat[2]*dir[0] + mat[5]*dir[1] + mat[8]*dir[2];
}


// transform a vector from local to global frame
static inline void localToGlobal(mjtNum res[3], const mjtNum mat[9], const mjtNum dir[3],
                                 const mjtNum pos[3]) {
  // perform mat * dir + pos
  res[0] = mat[0]*dir[0] + mat[1]*dir[1] + mat[2]*dir[2];
  res[1] = mat[3]*dir[0] + mat[4]*dir[1] + mat[5]*dir[2];
  res[2] = mat[6]*dir[0] + mat[7]*dir[1] + mat[8]*dir[2];
  res[0] += pos[0];
  res[1] += pos[1];
  res[2] += pos[2];
}


// point support function
void mjc_pointSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  mji_copy3(res, obj->pos);
}


// sphere support function
static void mjc_sphereSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  // sphere data
  const mjtNum* pos = obj->pos;
  mjtNum radius = obj->size[0];

  res[0] = radius*dir[0] + pos[0];
  res[1] = radius*dir[1] + pos[1];
  res[2] = radius*dir[2] + pos[2];
}


// line support function (capsule)
void mjc_lineSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  // capsule data
  const mjtNum* mat = obj->mat;
  const mjtNum* pos = obj->pos;
  mjtNum length = obj->size[1];

  mjtNum dot = mat[2]*dir[0] + mat[5]*dir[1] + mat[8]*dir[2];
  mjtNum scl = dot >= 0 ? length : -length;

  // transform result to global frame
  res[0] = mat[2]*scl + pos[0];
  res[1] = mat[5]*scl + pos[1];
  res[2] = mat[8]*scl + pos[2];
}


// capsule support function
static void mjc_capsuleSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  // capsule data
  const mjtNum* mat = obj->mat;
  const mjtNum* pos = obj->pos;
  mjtNum radius = obj->size[0];
  mjtNum length = obj->size[1];

  // rotate dir to geom local frame
  mjtNum local_dir[3], local_supp[3];
  mulMatTVec3(local_dir, mat, dir);

  // start with sphere
  local_supp[0] = local_dir[0] * radius;
  local_supp[1] = local_dir[1] * radius;
  local_supp[2] = local_dir[2] * radius;

  // add cylinder contribution
  local_supp[2] += (local_dir[2] >= 0 ? length : -length);

  // transform result to global frame
  localToGlobal(res, mat, local_supp, pos);
}


// ellipsoid support function
static void mjc_ellipsoidSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  // ellipsoid data
  const mjtNum* mat = obj->mat;
  const mjtNum* pos = obj->pos;
  const mjtNum* size = obj->size;

  // rotate dir to geom local frame
  mjtNum local_dir[3], local_supp[3];
  mulMatTVec3(local_dir, mat, dir);

  // find support point on unit sphere: scale dir by ellipsoid sizes
  local_supp[0] = local_dir[0] * size[0];
  local_supp[1] = local_dir[1] * size[1];
  local_supp[2] = local_dir[2] * size[2];

  mjtNum norm2 = local_supp[0]*local_supp[0] + local_supp[1]*local_supp[1] + local_supp[2]*local_supp[2];

  // too small to normalize
  if (norm2 < mjMINVAL2) {
    res[0] = mat[0]*size[0] + pos[0];
    res[1] = mat[3]*size[0] + pos[1];
    res[2] = mat[6]*size[0] + pos[2];
    return;
  }

  // normalize and transform to ellipsoid
  mjtNum norm_inv = 1/mju_sqrt(norm2);
  local_supp[0] *= norm_inv * size[0];
  local_supp[1] *= norm_inv * size[1];
  local_supp[2] *= norm_inv * size[2];

  // transform result to global frame
  localToGlobal(res, mat, local_supp, pos);
}


// cylinder support function
static void mjc_cylinderSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  // cylinder data
  const mjtNum* mat = obj->mat;
  const mjtNum* pos = obj->pos;
  const mjtNum* size = obj->size;

  // rotate dir to geom local frame
  mjtNum local_dir[3], local_supp[3];
  mulMatTVec3(local_dir, mat, dir);

  mjtNum n2 = local_dir[0]*local_dir[0] + local_dir[1]*local_dir[1];
  mjtNum scl = n2 >= mjMINVAL2 ? size[0] / mju_sqrt(n2) : 0;
  local_supp[0] = scl * local_dir[0];
  local_supp[1] = scl * local_dir[1];

  // set result in Z direction
  local_supp[2] = local_dir[2] >= 0 ? size[1] : -size[1];

  // transform result to global frame
  localToGlobal(res, mat, local_supp, pos);
}


// box support function
static void mjc_boxSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  // box data
  const mjtNum* mat = obj->mat;
  const mjtNum* pos = obj->pos;
  const mjtNum* size = obj->size;

  // rotate dir to geom local frame
  mjtNum local_dir[3], local_supp[3];
  mulMatTVec3(local_dir, mat, dir);

  // find support point in local frame
  local_supp[0] = local_dir[0] >= 0 ? size[0] : -size[0];
  local_supp[1] = local_dir[1] >= 0 ? size[1] : -size[1];
  local_supp[2] = local_dir[2] >= 0 ? size[2] : -size[2];

  // mark the index of the corner of the box for fast lookup
  obj->vertindex  = (local_supp[0] > 0) ? 1 : 0;
  obj->vertindex |= (local_supp[1] > 0) ? 2 : 0;
  obj->vertindex |= (local_supp[2] > 0) ? 4 : 0;

  // transform support point to global frame
  localToGlobal(res, mat, local_supp, pos);
}


// dot product between mjtNum and float
static inline mjtNum dot3f(const mjtNum a[3], const float b[3]) {
  return a[0]*(mjtNum)b[0] + a[1]*(mjtNum)b[1] + a[2]*(mjtNum)b[2];
}


// mesh support function via exhaustive search
static void mjc_meshSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  const mjtNum* mat = obj->mat;
  const mjtNum* pos = obj->pos;
  const float* verts = obj->data.mesh.vert;
  int nverts = obj->data.mesh.nvert;

  mjtNum local_dir[3];
  mulMatTVec3(local_dir, mat, dir);

  mjtNum max = -FLT_MAX;
  int imax = 0;

  // used cached results from previous search
  if (obj->vertindex >= 0) {
    imax = obj->vertindex;
    max = dot3f(local_dir, verts + 3*imax);
  }

  // search all vertices, find maximum dot product
  for (int i=0; i < nverts; i++) {
    mjtNum vdot = dot3f(local_dir, verts + 3*i);

    // update max
    if (vdot > max) {
      max = vdot;
      imax = i;
    }
  }

  // record vertex index of maximum
  obj->vertindex = imax;

  local_dir[0] = (mjtNum)verts[3*imax + 0];
  local_dir[1] = (mjtNum)verts[3*imax + 1];
  local_dir[2] = (mjtNum)verts[3*imax + 2];

  // transform result to global frame
  localToGlobal(res, mat, local_dir, pos);
}


// mesh support function via hill climbing
static void mjc_hillclimbSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  int numvert = obj->data.mesh.graph[0];
  const int* vert_edgeadr = obj->data.mesh.graph + 2;
  const int* vert_globalid = obj->data.mesh.graph  + 2 + numvert;
  const int* edge_localid = obj->data.mesh.graph + 2 + 2*numvert;
  const float* verts = obj->data.mesh.vert;
  const mjtNum* pos = obj->pos;
  const mjtNum* mat = obj->mat;

  // rotate dir to geom local frame
  mjtNum local_dir[3];
  mulMatTVec3(local_dir, mat, dir);

  int prev = -1;
  int imax = obj->meshindex >= 0 ? obj->meshindex : 0;
  mjtNum max = dot3f(local_dir, verts + 3*vert_globalid[imax]);

  // hillclimb until no change
  while (imax != prev) {
    prev = imax;
    int subidx;
    for (int i = vert_edgeadr[imax]; (subidx = edge_localid[i]) >= 0; i++) {
      mjtNum vdot = dot3f(local_dir, verts + 3*vert_globalid[subidx]);
      if (vdot > max) {
        max = vdot;
        imax = subidx;  // update maximum vertex index
      }
    }
  }

  // record vertex index of maximum (local id)
  obj->meshindex = imax;

  // get resulting support vertex
  obj->vertindex = imax = vert_globalid[imax];
  local_dir[0] = (mjtNum)verts[3*imax + 0];
  local_dir[1] = (mjtNum)verts[3*imax + 1];
  local_dir[2] = (mjtNum)verts[3*imax + 2];

  // transform result to global frame
  localToGlobal(res, mat, local_dir, pos);
}


// prism support function
static void mjc_prism_support(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  int istart, ibest;
  mjtNum best, tmp;
  mjtNum (*prism)[3] = obj->data.hfield.prism;

  // find best vertex in halfspace determined by dir.z
  istart = dir[2] < 0 ? 0 : 3;
  ibest = istart;
  best = mju_dot3(prism[istart], dir);
  for (int i=1; i < 3; i++) {
    if ((tmp = mju_dot3(prism[istart + i], dir)) > best) {
      ibest = istart + i;
      best = tmp;
    }
  }

  // copy best point
  mji_copy3(res, prism[ibest]);
}


// flex support function
static void mjc_flexSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
  int f = obj->flex;
  int dim = obj->data.flex.dim[f];

  // flex element
  if (obj->elem >= 0) {
    int e = obj->elem;
    const int* edata = obj->data.flex.elem + obj->data.flex.elemdataadr[f] + e*(dim+1);
    const mjtNum* vert = obj->data.flex.vert_xpos + 3*obj->data.flex.vertadr[f];

    // find element vertex with largest projection along dir
    mji_copy3(res, vert+3*edata[0]);
    mjtNum best = mju_dot3(res, dir);
    for (int i=1; i <= dim; i++) {
      mjtNum dot = mju_dot3(vert+3*edata[i], dir);

      // better vertex found: assign
      if (dot > best) {
        best = dot;
        mji_copy3(res, vert+3*edata[i]);
      }
    }

    // add radius and margin/2
    mji_addToScl3(res, dir, obj->data.flex.xradius[f] + 0.5*obj->margin);
    return;
  }

  // flex vertex
  else {
    const mjtNum* vert = obj->data.flex.vert_xpos + 3*(obj->data.flex.vertadr[f] + obj->vert);
    mji_addScl3(res, vert, dir, obj->data.flex.xradius[f] + 0.5*obj->margin);
    return;
  }
}


// libccd support function
void mjccd_support(const void *_obj, const ccd_vec3_t *_dir, ccd_vec3_t *vec) {
  mjCCDObj *obj = (mjCCDObj *)_obj;
  mjtNum *res = vec->v;
  const mjtNum *dir = _dir->v;
  int g = obj->geom;

  if (g < 0) {
    int f = obj->flex;
    int dim = obj->data.flex.dim[f];

    // flex element
    if (obj->elem >= 0) {
      int e = obj->elem;
      const int* edata = obj->data.flex.elem + obj->data.flex.elemdataadr[f] + e*(dim+1);
      const mjtNum* vert = obj->data.flex.vert_xpos + 3*obj->data.flex.vertadr[f];

      // find element vertex with largest projection along dir
      mji_copy3(res, vert+3*edata[0]);
      mjtNum best = mju_dot3(res, dir);
      for (int i=1; i <= dim; i++) {
        mjtNum dot = mju_dot3(vert+3*edata[i], dir);

        // better vertex found: assign
        if (dot > best) {
          best = dot;
          mji_copy3(res, vert+3*edata[i]);
        }
      }

      // add radius and margin/2
      mji_addToScl3(res, dir, obj->data.flex.xradius[f] + 0.5*obj->margin);
      return;
    }

    // flex vertex
    else {
      const mjtNum* vert = obj->data.flex.vert_xpos + 3*(obj->data.flex.vertadr[f] + obj->vert);
      mji_addScl3(res, vert, dir, obj->data.flex.xradius[f] + 0.5*obj->margin);
      return;
    }
  }

  const float* vertdata;
  int ibest, numvert, change, locid;
  mjtNum tmp, vdot;

  const mjtNum* size = obj->size;          // geom sizes
  mjtNum local_dir[3];                    // direction in geom local frame

  // rotate dir to geom local frame
  mju_mulMatTVec3(local_dir, obj->mat, dir);

  // compute result according to geom type
  switch ((mjtGeom) obj->geom_type) {
  case mjGEOM_SPHERE:
    mji_scl3(res, local_dir, size[0]);
    break;

  case mjGEOM_CAPSULE:
    // start with sphere
    mji_scl3(res, local_dir, size[0]);

    // add cylinder contribution
    res[2] += mju_sign(local_dir[2]) * size[1];
    break;

  case mjGEOM_ELLIPSOID:
    // find support point on unit sphere: scale dir by ellipsoid sizes and renormalize
    for (int i=0; i < 3; i++) {
      res[i] = local_dir[i] * size[i];
    }
    mju_normalize3(res);

    // transform to ellipsoid
    for (int i=0; i < 3; i++) {
      res[i] *= size[i];
    }
    break;

  case mjGEOM_CYLINDER:
    // set result in XY plane: support on circle
    tmp = mju_sqrt(local_dir[0]*local_dir[0] + local_dir[1]*local_dir[1]);
    if (tmp > mjMINVAL) {
      res[0] = local_dir[0]/tmp*size[0];
      res[1] = local_dir[1]/tmp*size[0];
    } else {
      res[0] = res[1] = 0;
    }

    // set result in Z direction
    res[2] = mju_sign(local_dir[2]) * size[1];
    break;

  case mjGEOM_BOX:
    for (int i=0; i < 3; i++) {
      res[i] = mju_sign(local_dir[i]) * size[i];
    }
    break;

  case mjGEOM_MESH:
  case mjGEOM_SDF:
    // init search
    vertdata = obj->data.mesh.vert;
    tmp = -1E+10;
    ibest = -1;

    // no graph data: exhaustive search
    if (obj->data.mesh.graph == NULL) {
      // search all vertices, find best
      for (int i=0; i < obj->data.mesh.nvert; i++) {
        // vdot = dot(vertex, dir)
        vdot = local_dir[0] * (mjtNum)vertdata[3*i] +
               local_dir[1] * (mjtNum)vertdata[3*i+1] +
               local_dir[2] * (mjtNum)vertdata[3*i+2];

        // update best
        if (vdot > tmp) {
          tmp = vdot;
          ibest = i;
        }
      }

      // record best vertex index, in globalid format
      obj->meshindex = ibest;
    }

    // hill-climb using graph data
    else {
      // get info
      numvert = obj->data.mesh.graph[0];
      const int* vert_edgeadr = obj->data.mesh.graph + 2;
      const int* vert_globalid = obj->data.mesh.graph + 2 + numvert;
      const int* edge_localid = obj->data.mesh.graph + 2 + 2*numvert;

      // init with first vertex in convex hull or warmstart
      ibest = obj->meshindex < 0 ? 0 : obj->meshindex;
      tmp = local_dir[0] * (mjtNum)vertdata[3*vert_globalid[ibest]+0] +
            local_dir[1] * (mjtNum)vertdata[3*vert_globalid[ibest]+1] +
            local_dir[2] * (mjtNum)vertdata[3*vert_globalid[ibest]+2];

      // hill-climb until no change
      change = 1;
      while (change) {
        // look for improvement in ibest neighborhood
        change = 0;
        int i = vert_edgeadr[ibest];
        while ((locid=edge_localid[i]) >= 0) {
          // vdot = dot(vertex, local_dir)
          vdot = local_dir[0] * (mjtNum)vertdata[3*vert_globalid[locid]] +
                 local_dir[1] * (mjtNum)vertdata[3*vert_globalid[locid]+1] +
                 local_dir[2] * (mjtNum)vertdata[3*vert_globalid[locid]+2];

          // update best
          if (vdot > tmp) {
            tmp = vdot;
            ibest = locid;
            change = 1;
          }

          // advance to next edge
          i++;
        }
      }

      // record best vertex index, in locid format
      obj->meshindex = ibest;

      // map best index to globalid
      ibest = vert_globalid[ibest];
    }

    // sanity check, SHOULD NOT OCCUR
    if (ibest < 0) {
      mju_warning("mesh_support could not find support vertex");
      mju_zero3(res);
    }

    // copy best vertex
    else {
      for (int i=0; i < 3; i++) {
        res[i] = (mjtNum)vertdata[3*ibest + i];
      }
    }
    break;

  case mjGEOM_HFIELD:
    mjc_prism_support(res, obj, dir);
    return;

  default:
    mjERROR("ccd support function is undefined for geom type %d", obj->geom_type);
  }

  // add local_dir*margin/2 to result
  for (int i=0; i < 3; i++) {
    res[i] += local_dir[i] * obj->margin/2;
  }

  // rotate result to global frame
  mju_mulMatVec3(res, obj->mat, res);

  // add geom position
  mji_addTo3(res, obj->pos);
}

// ------------------------------------------------------------------------------------------------

// initialize a CCD object
void mjc_initCCDObj(mjCCDObj* obj, const mjModel* m, const mjData* d, int g, mjtNum margin) {
  int graphadr, vertadr, polyadr;
  obj->geom = g;
  obj->margin = margin;
  obj->center = mjc_center;
  obj->vertindex = -1;
  obj->meshindex = -1;
  obj->flex = -1;
  obj->elem = -1;
  obj->vert = -1;
  mju_zero4(obj->rotate);
  obj->rotate[0] = 1;
  if (g >= 0) {
    mju_copy(obj->size, m->geom_size+3*g, 3);
    mju_copy(obj->pos, d->geom_xpos+3*g, 3);
    mju_copy(obj->mat, d->geom_xmat+9*g, 9);
    obj->geom_type = m->geom_type[g];
    switch ((mjtGeom) obj->geom_type) {
    case mjGEOM_ELLIPSOID:
      obj->support = mjc_ellipsoidSupport;
      break;
    case mjGEOM_MESH:
    case mjGEOM_SDF:
      graphadr = m->mesh_graphadr[m->geom_dataid[g]];
      vertadr = m->mesh_vertadr[m->geom_dataid[g]];
      polyadr = m->mesh_polyadr[m->geom_dataid[g]];
      if (graphadr < 0 || m->mesh_vertnum[m->geom_dataid[g]] < mjMESH_HILLCLIMB_MIN) {
        obj->data.mesh.graph = NULL;
        obj->support = mjc_meshSupport;
      } else {
        obj->data.mesh.graph = m->mesh_graph + graphadr;
        obj->support = mjc_hillclimbSupport;
      }
      obj->data.mesh.vert = m->mesh_vert + 3*vertadr;
      obj->data.mesh.nvert = m->mesh_vertnum[m->geom_dataid[g]];
      obj->data.mesh.mpolymapadr = m->mesh_polymapadr + vertadr;
      obj->data.mesh.mpolymapnum = m->mesh_polymapnum + vertadr;
      obj->data.mesh.polymap = m->mesh_polymap;
      obj->data.mesh.polynormal = m->mesh_polynormal + 3*polyadr;
      obj->data.mesh.polyvertadr = m->mesh_polyvertadr + polyadr;
      obj->data.mesh.polyvertnum = m->mesh_polyvertnum + polyadr;
      obj->data.mesh.polyvert = m->mesh_polyvert;
      obj->data.mesh.mesh_polynum = m->mesh_polynum[m->geom_dataid[g]];
      break;
    case mjGEOM_SPHERE:
      obj->support = mjc_sphereSupport;
      break;
    case mjGEOM_CAPSULE:
      obj->support = mjc_capsuleSupport;
      break;
    case mjGEOM_CYLINDER:
      obj->support = mjc_cylinderSupport;
      break;
    case mjGEOM_BOX:
      obj->support = mjc_boxSupport;
      break;
    case mjGEOM_HFIELD:
      obj->center = mjc_center;
      obj->support = mjc_prism_support;

      int hid = m->geom_dataid[g];
      obj->data.hfield.hfield_nrow = m->hfield_nrow[hid];
      obj->data.hfield.hfield_ncol = m->hfield_ncol[hid];
      mju_copy(obj->size, m->hfield_size + 4*hid, 4);
      obj->data.hfield.hfield_data = m->hfield_data + m->hfield_adr[hid];
      break;
    default:
      obj->support = NULL;
      break;
    }
  } else {
    obj->geom_type = mjGEOM_FLEX;
    obj->data.flex.dim = m->flex_dim;
    obj->support = mjc_flexSupport;
    obj->data.flex.aabb = d->flexelem_aabb;
    obj->data.flex.elemadr = m->flex_elemadr;
    obj->data.flex.vert_xpos = d->flexvert_xpos;
    obj->data.flex.vertadr = m->flex_vertadr;
    obj->data.flex.xradius = m->flex_radius;
    obj->data.flex.elemdataadr = m->flex_elemdataadr;
    obj->data.flex.elem = m->flex_elem;
  }
}


// set flex data for CCD object
static void mjc_setCCDObjFlex(mjCCDObj* obj, int flex, int elem, int vert) {
  obj->flex = flex;
  obj->elem = elem;
  obj->vert = vert;
}


// compare new contact to previous contacts, return 1 if it is far from all of them
static int mjc_isDistinctContact(const mjPreContact* con, int ncon, mjtNum tolerance) {
  const mjtNum* last_pos = con[ncon - 1].pos;
  for (int i=0; i < ncon-1; i++) {
    if (mju_dist3(con[i].pos, last_pos) <= tolerance) {
      return 0;
    }
  }
  return 1;
}


// in-place rotation of spatial frame around given point of origin
static void mju_rotateFrame(const mjtNum origin[3], const mjtNum rot[9],
                            mjtNum xmat[9], mjtNum xpos[3]) {
  mjtNum mat[9], vec[3], rel[3];

  // rotate frame: xmat = rot*xmat
  mju_mulMatMat3(mat, rot, xmat);
  mju_copy(xmat, mat, 9);

  // vector to rotation origin: rel = origin - xpos
  mji_sub3(rel, origin, xpos);

  // displacement of origin due to rotation: vec = rot*rel - rel
  mju_mulMatVec3(vec, rot, rel);
  mju_subFrom3(vec, rel);

  // correct xpos by subtracting displacement: xpos = xpos - vec
  mji_subFrom3(xpos, vec);
}


// return number of contacts supported by a single pass of narrowphase
static int maxContacts(const mjModel* m, const mjCCDObj* obj1, const mjCCDObj* obj2) {
  // single pass not supported for margins
  if (obj1->margin > 0 || obj2->margin > 0) {
    return 1;
  }

  // can return 8 contacts for box-box collision in one pass
  int type1 = obj1->geom_type;
  int type2 = obj2->geom_type;
  if (type1 == mjGEOM_BOX && type2 == mjGEOM_BOX) {
    return 8;
  }

  // reduce mesh collisions to 4 contacts max
  if (type1 == mjGEOM_BOX || type1 == mjGEOM_MESH) {
    if (type2 == mjGEOM_BOX || type2 == mjGEOM_MESH) {
      return mjDISABLED(mjDSBL_MULTICCD) ? 1 : 4;
    }
  }

  // not supported for other geom types
  return 1;
}


// multi-point convex-convex collision, using libccd
int mjc_Convex(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
  // init ccd objects
  mjCCDObj obj1, obj2;
  mjc_initCCDObj(&obj1, m, d, g1, margin);
  mjc_initCCDObj(&obj2, m, d, g2, margin);
  int max_contacts = maxContacts(m, &obj1, &obj2);

  // find initial contact
  int ncon = mjc_penetration(m, d, &obj1, &obj2, con, max_contacts, margin);
  if (mjDISABLED(mjDSBL_NATIVECCD) && ncon && g1 >= 0 && g2 >= 0) {
    mjc_fixNormal(m, d, con, g1, g2);
  }

  // no additional contacts needed
  if (!mjDISABLED(mjDSBL_NATIVECCD) && max_contacts > 1) {
    return ncon;
  }

  // look for additional contacts
  if (ncon == 1 && !mjDISABLED(mjDSBL_MULTICCD)
      && m->geom_type[g1] != mjGEOM_ELLIPSOID && m->geom_type[g1] != mjGEOM_SPHERE
      && m->geom_type[g2] != mjGEOM_ELLIPSOID && m->geom_type[g2] != mjGEOM_SPHERE) {
    // multiCCD parameters
    const mjtNum relative_tolerance = 1e-3;
    const mjtNum perturbation_angle = 1e-3;

    // complete frame of initial contact
    mjtNum frame[9];
    mji_copy3(frame, con[0].normal);
    mju_zero(frame+3, 6);
    mju_makeFrame(frame);

    // tolerance for determining if newly found contacts are distinct
    const mjtNum tolerance = relative_tolerance * mju_min(m->geom_rbound[g1], m->geom_rbound[g2]);

    // axes and rotation angles for perturbation test
    mjtNum* axes[2] = {frame+3, frame+6};
    mjtNum angles[2] = {-perturbation_angle, perturbation_angle};

    // rotate both geoms, search for new contacts
    for (int axis_id = 0; axis_id < 2; ++axis_id) {
      for (int angle_id = 0; angle_id < 2; ++angle_id) {
        mjtNum* axis = axes[axis_id];
        mjtNum angle = angles[angle_id];

        // make rotation matrix rot
        mjtNum quat[4], rot[9];
        mji_axisAngle2Quat(quat, axis, angle);
        mju_quat2Mat(rot, quat);

        // rotate g1 around initial contact point
        mju_rotateFrame(con[0].pos, rot, obj1.mat, obj1.pos);

        // inversely rotate g2 around initial contact point
        mjtNum invrot[9];
        mju_transpose(invrot, rot, 3, 3);
        mju_rotateFrame(con[0].pos, invrot, obj2.mat, obj2.pos);

        // search for new contact
        int n = mjc_penetration(m, d, &obj1, &obj2, con + ncon, 1, margin);
        if (mjDISABLED(mjDSBL_NATIVECCD) && n && g1 >= 0 && g2 >= 0) {
          mjc_fixNormal(m, d, con + ncon, g1, g2);
        }

        // check new contact
        if (n && mjc_isDistinctContact(con, ncon + 1, tolerance)) {
          // set penetration of new point to equal that of initial point
          con[ncon].dist = con[0].dist;
          // add new point
          ncon += 1;
        }

        // reset positions and orientations of g1 and g2
        mji_copy3(obj1.pos, d->geom_xpos+3*g1);
        mji_copy9(obj1.mat, d->geom_xmat+9*g1);
        mji_copy3(obj2.pos, d->geom_xpos+3*g2);
        mji_copy9(obj2.mat, d->geom_xmat+9*g2);
      }
    }
  }
  return ncon;
}


// parameters for plane-mesh extra contacts
const int maxplanemesh = 3;
const mjtNum tolplanemesh = 0.3;

// add one plane-mesh contact
static int addplanemesh(mjPreContact* con, const float vertex[3],
                        const mjtNum pos1[3], const mjtNum normal1[3],
                        const mjtNum pos2[3], const mjtNum mat2[9],
                        const mjtNum first[3], mjtNum rbound) {
  // compute point in global coordinates
  mjtNum pnt[3], v[3] = {vertex[0], vertex[1], vertex[2]};
  mju_mulMatVec3(pnt, mat2, v);
  mju_addTo3(pnt, pos2);

  // skip if too close to first contact
  if (mju_dist3(pnt, first) < tolplanemesh*rbound) {
    return 0;
  }

  // pnt-pos difference vector
  mjtNum dif[3];
  mji_sub3(dif, pnt, pos1);

  // set distance
  con[0].dist = mju_dot3(normal1, dif);

  // set position
  mji_copy3(con[0].pos, pnt);
  mji_addToScl3(con[0].pos, normal1, -0.5*con[0].dist);

  // set frame
  mji_copy3(con[0].normal, normal1);
  mji_zero3(con[0].tangent);

  return 1;
}


// plane-convex collision, using libccd
int mjc_PlaneConvex(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
  const mjtNum* pos1  = d->geom_xpos + 3*g1;
  const mjtNum* mat1  = d->geom_xmat + 9*g1;
  const mjtNum* pos2  = d->geom_xpos + 3*g2;
  const mjtNum* mat2  = d->geom_xmat + 9*g2;

  mjtNum dif[3], normal[3] = {mat1[2], mat1[5], mat1[8]};
  ccd_vec3_t ccd_dir, ccd_vec;
  mjCCDObj obj;
  mjc_initCCDObj(&obj, m, d, g2, 0);
  // get support point in -normal direction
  ccdVec3Set(&ccd_dir, -mat1[2], -mat1[5], -mat1[8]);
  mjccd_support(&obj, &ccd_dir, &ccd_vec);

  // compute normal distance, return if too far
  mji_sub3(dif, ccd_vec.v, pos1);
  con[0].dist = mju_dot3(normal, dif);
  if (con[0].dist > margin) {
    return 0;
  }

  // fill in contact data
  mji_copy3(con[0].pos, ccd_vec.v);
  mji_addToScl3(con[0].pos, normal, -0.5*con[0].dist);
  mji_copy3(con[0].normal, normal);
  mji_zero3(con[0].tangent);

  //--------------- add all/connected vertices below margin
  float* vertdata;
  int graphadr, numvert, locid;
  int *vert_edgeadr, *vert_globalid, *edge_localid;
  mjtNum vdot;
  int count = 1, g = g2;

  // g is an ellipsoid: no need for further mesh-specific processing
  if (m->geom_dataid[g] == -1) {
    return count;
  }

  // init
  vertdata = m->mesh_vert + 3*m->mesh_vertadr[m->geom_dataid[g]];

  // express dir in geom local frame
  mjtNum locdir[3];
  mju_mulMatTVec3(locdir, d->geom_xmat+9*g, ccd_dir.v);

  // inclusion threshold along locdir, relative to geom2 center
  mji_sub3(dif, pos2, pos1);
  mjtNum threshold = mju_dot3(normal, dif) - margin;

  // no graph data: exhaustive search
  if (m->mesh_graphadr[m->geom_dataid[g]] < 0) {
    // search all vertices, find best
    for (int i=0; i < m->mesh_vertnum[m->geom_dataid[g]] && count < maxplanemesh; i++) {
      // vdot = dot(vertex, dir)
      vdot = locdir[0] * (mjtNum)vertdata[3*i] +
             locdir[1] * (mjtNum)vertdata[3*i+1] +
             locdir[2] * (mjtNum)vertdata[3*i+2];

      // detect contact, skip best
      if (vdot > threshold && i != obj.meshindex) {
        count += addplanemesh(con+count, vertdata+3*i,
                              pos1, normal, pos2, mat2,
                              con->pos, m->geom_rbound[g2]);
      }
    }
  }

  // use graph data
  else if (obj.meshindex >= 0) {
    // get info
    graphadr = m->mesh_graphadr[m->geom_dataid[g]];
    numvert = m->mesh_graph[graphadr];
    vert_edgeadr = m->mesh_graph + graphadr + 2;
    vert_globalid = m->mesh_graph + graphadr + 2 + numvert;
    edge_localid = m->mesh_graph + graphadr + 2 + 2*numvert;

    // look for contacts in ibest neighborhood
    int i = vert_edgeadr[obj.meshindex];
    while ((locid=edge_localid[i]) >= 0 && count < maxplanemesh) {
      // vdot = dot(vertex, dir)
      vdot = locdir[0] * (mjtNum)vertdata[3*vert_globalid[locid]] +
             locdir[1] * (mjtNum)vertdata[3*vert_globalid[locid]+1] +
             locdir[2] * (mjtNum)vertdata[3*vert_globalid[locid]+2];

      // detect contact
      if (vdot > threshold) {
        count += addplanemesh(con+count, vertdata+3*vert_globalid[locid],
                              pos1, normal, pos2, mat2,
                              con->pos, m->geom_rbound[g2]);
      }

      // advance to next edge
      i++;
    }
  }

  return count;
}


//----------------------------  heightfield collisions ---------------------------------------------


// add vertex to prism
static inline void addVert(mjCCDObj* obj, mjtNum x, mjtNum y, mjtNum z) {
  // move old data
  mji_copy3(obj->data.hfield.prism[0], obj->data.hfield.prism[1]);
  mji_copy3(obj->data.hfield.prism[1], obj->data.hfield.prism[2]);
  mji_copy3(obj->data.hfield.prism[3], obj->data.hfield.prism[4]);
  mji_copy3(obj->data.hfield.prism[4], obj->data.hfield.prism[5]);

  // add new vertex at last position
  obj->data.hfield.prism[2][0] = obj->data.hfield.prism[5][0] = x;
  obj->data.hfield.prism[2][1] = obj->data.hfield.prism[5][1] = y;
  obj->data.hfield.prism[5][2] = z;
}


// add vertex to prism
static inline void addPrismVert(mjCCDObj* obj, int r, int c, int i, mjtNum dx, mjtNum dy, mjtNum margin) {
  // move old data
  mji_copy3(obj->data.hfield.prism[0], obj->data.hfield.prism[1]);
  mji_copy3(obj->data.hfield.prism[1], obj->data.hfield.prism[2]);
  mji_copy3(obj->data.hfield.prism[3], obj->data.hfield.prism[4]);
  mji_copy3(obj->data.hfield.prism[4], obj->data.hfield.prism[5]);

  int dr = 1 - i;

  // add new vertex at last position
  obj->data.hfield.prism[2][0] = obj->data.hfield.prism[5][0] = dx*c - obj->size[0];
  obj->data.hfield.prism[2][1] = obj->data.hfield.prism[5][1] = dy*(r + dr) - obj->size[1];
  obj->data.hfield.prism[5][2] = obj->data.hfield.hfield_data[(r + dr)*obj->data.hfield.hfield_ncol + c]*obj->size[2];

  // factor in margin
  obj->data.hfield.prism[5][2] += margin;
}


// entry point for heightfield collisions
int mjc_ConvexHField(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                     mjtNum margin) {
  // hfield frame
  const mjtNum* pos1 = d->geom_xpos + 3*g1;
  const mjtNum* mat1 = d->geom_xmat + 9*g1;

  // geom2 frame
  mjtNum* pos2 = d->geom_xpos + 3*g2;
  mjtNum* mat2 = d->geom_xmat + 9*g2;

  // hfield data
  int hid = m->geom_dataid[g1];
  int nrow = m->hfield_nrow[hid];
  int ncol = m->hfield_ncol[hid];
  mjtNum size0 = m->hfield_size[4*hid + 0], size1 = m->hfield_size[4*hid + 1];
  mjtNum size2 = m->hfield_size[4*hid + 2], size3 = m->hfield_size[4*hid + 3];

  // try early return using box-sphere test

  // express geom2 pos in hfield frame
  mjtNum local_pos[3] = {pos2[0] - pos1[0], pos2[1] - pos1[1], pos2[2] - pos1[2]};
  mju_mulMatTVec3(local_pos, mat1, local_pos);

  // sphere radius is geom2 rbound + margin
  mjtNum radius = m->geom_rbound[g2] + margin;

  // box-sphere test
  if ((size0 < local_pos[0] - radius) || (-size0 > local_pos[0] + radius) ||
      (size1 < local_pos[1] - radius) || (-size1 > local_pos[1] + radius) ||
      (size2 < local_pos[2] - radius) || (-size3 > local_pos[2] + radius)) {
    return 0;
  }

  // ccd set up
  mjCCDObj obj1, obj2;
  mjc_initCCDObj(&obj1, m, d, g1, 0);
  mjc_initCCDObj(&obj2, m, d, g2, 0);


  // try early return using AABB box-box test

  // express geom2 mat in hfield frame
  mjtNum mat[9];
  mji_mulMatTMat3(mat, mat1, mat2);

  mji_copy9(obj2.mat, mat);
  mji_copy3(obj2.pos, local_pos);

  mjtNum local_dir[3] = {0, 0, 0}, res[3];

  // get support point in +X
  local_dir[0] = 1;
  obj2.support(res, &obj2, local_dir);
  mjtNum xmax = res[0];

  // get support point in -X
  local_dir[0] = -1;
  obj2.support(res, &obj2, local_dir);
  mjtNum xmin = res[0];

  local_dir[0] = 0;

  // get support point in +Y
  local_dir[1] = 1;
  obj2.support(res, &obj2, local_dir);
  mjtNum ymax = res[1];

  // get support point in -Y
  local_dir[1] = -1;
  obj2.support(res, &obj2, local_dir);
  mjtNum ymin = res[1];

  local_dir[1] = 0;

  // get support point in +Z
  local_dir[2] = 1;
  obj2.support(res, &obj2, local_dir);
  mjtNum zmax = res[2];

  // get support point in -Z
  local_dir[2] = -1;
  obj2.support(res, &obj2, local_dir);
  mjtNum zmin = res[2];

  // AABB box-box test
  if ((xmin - margin > size0) || (xmax + margin < -size0) ||
      (ymin - margin > size1) || (ymax + margin < -size1) ||
      (zmin - margin > size2) || (zmax + margin < -size3)) {
    return 0;
  }

  // compute sub-grid bounds
  int cmin = (int) mju_floor((xmin + size0) / (2.0*size0) * (ncol-1));
  int cmax = (int) mju_ceil ((xmax + size0) / (2.0*size0) * (ncol-1));
  int rmin = (int) mju_floor((ymin + size1) / (2.0*size1) * (nrow-1));
  int rmax = (int) mju_ceil ((ymax + size1) / (2.0*size1) * (nrow-1));
  cmin = mjMAX(0, cmin);
  cmax = mjMIN(ncol-1, cmax);
  rmin = mjMAX(0, rmin);
  rmax = mjMIN(nrow-1, rmax);

  // geom margin needed for actual collision test
  obj2.margin = margin;

  // compute real-valued grid step
  mjtNum dx = (2.0*size0) / (ncol-1);
  mjtNum dy = (2.0*size1) / (nrow-1);

  // set zbottom value using base size
  mjtNum (*prism)[3] = obj1.data.hfield.prism;
  prism[0][2] = prism[1][2] = prism[2][2] = -size3;

  // process all prisms in subgrid
  int ncon = 0;
  for (int r=rmin; r < rmax; r++) {
    addPrismVert(&obj1, r, cmin, 0, dx, dy, margin);
    addPrismVert(&obj1, r, cmin, 1, dx, dy, margin);
    for (int c=cmin + 1; c <= cmax; c++) {
      for (int i=0; i < 2; i++) {
        // send vertex to prism constructor
        addPrismVert(&obj1, r, c, i, dx, dy, margin);

        // prism height test
        if (prism[3][2] < zmin && prism[4][2] < zmin && prism[5][2] < zmin) {
          continue;
        }

        // run penetration function, save contact
        if (mjc_penetration(m, d, &obj1, &obj2, con + ncon, 1, 0.0)) {
          // transform to global coordinates
          mji_copy3(local_dir, con[ncon].normal);
          mji_copy3(local_pos, con[ncon].pos);
          mji_mulMatVec3(con[ncon].normal, mat1, local_dir);
          mji_mulMatVec3(con[ncon].pos, mat1, local_pos);
          mji_addTo3(con[ncon].pos, pos1);

          // force out of all loops if max contacts reached
          if (++ncon >= mjMAXCONPAIR) {
            r = rmax+1;
            c = cmax+1;
            i = 3;
            break;
          }
        }
      }
    }
  }

  if (mjDISABLED(mjDSBL_NATIVECCD)) {
    // fix contact normals
    for (int i=0; i < ncon; i++) {
      mjc_fixNormal(m, d, con + i, g1, g2);
    }
  }

  return ncon;
}


//--------------------------- fix contact frame normal ---------------------------------------------

// compute normal for point outside ellipsoid, using ray-projection SQP
static int mjc_ellipsoidInside(mjtNum nrm[3], const mjtNum pos[3], const mjtNum size[3]) {
  // algorithm constants
  const int maxiter = 30;
  const mjtNum tolerance = 1e-6;

  // precompute quantities
  mjtNum S2inv[3] = {1/(size[0]*size[0]), 1/(size[1]*size[1]), 1/(size[2]*size[2])};
  mjtNum C = pos[0]*pos[0]*S2inv[0] + pos[1]*pos[1]*S2inv[1] + pos[2]*pos[2]*S2inv[2] - 1;
  if (C > 0) {
    return 0;
  }

  // normalize initial normal (just in case)
  mju_normalize3(nrm);

  // main iteration
  int iter;
  for (iter=0; iter < maxiter; iter++) {
    // coefficients and determinant of quadratic
    mjtNum A = nrm[0]*nrm[0]*S2inv[0] + nrm[1]*nrm[1]*S2inv[1] + nrm[2]*nrm[2]*S2inv[2];
    mjtNum B = pos[0]*nrm[0]*S2inv[0] + pos[1]*nrm[1]*S2inv[1] + pos[2]*nrm[2]*S2inv[2];
    mjtNum det = B*B - A*C;
    if (det < mjMINVAL || A < mjMINVAL) {
      return (iter > 0);
    }

    // ray intersection with ellipse: pos + x*nrm, x>=0
    mjtNum x = (-B + mju_sqrt(det))/A;
    if (x < 0) {
      return (iter > 0);
    }

    // new point on ellipsoid
    mjtNum pnt[3];
    mji_addScl3(pnt, pos, nrm, x);

    // normal at new point
    mjtNum newnrm[3] = {pnt[0]*S2inv[0], pnt[1]*S2inv[1], pnt[2]*S2inv[2]};
    mju_normalize3(newnrm);

    // save change and assign
    mjtNum change = mju_dist3(nrm, newnrm);
    mji_copy3(nrm, newnrm);

    // terminate if converged
    if (change < tolerance) {
      break;
    }
  }

  return 1;
}


// compute normal for point inside ellipsoid, using diagonal QCQP
static int mjc_ellipsoidOutside(mjtNum nrm[3], const mjtNum pos[3], const mjtNum size[3]) {
  // algorithm constants
  const int maxiter = 30;
  const mjtNum tolerance = 1e-6;

  // precompute quantities
  mjtNum S2[3] = {size[0]*size[0], size[1]*size[1], size[2]*size[2]};
  mjtNum PS2[3] = {pos[0]*pos[0]*S2[0], pos[1]*pos[1]*S2[1], pos[2]*pos[2]*S2[2]};

  // main iteration
  mjtNum la = 0;
  int iter;
  for (iter=0; iter < maxiter; iter++) {
    // precompute 1/(s^2+la)
    mjtNum R[3] = {1/(S2[0]+la), 1/(S2[1]+la), 1/(S2[2]+la)};

    // value
    mjtNum val = PS2[0]*R[0]*R[0] + PS2[1]*R[1]*R[1] + PS2[2]*R[2]*R[2] - 1;
    if (val < tolerance) {
      break;
    }

    // derivative
    mjtNum deriv = -2*(PS2[0]*R[0]*R[0]*R[0] + PS2[1]*R[1]*R[1]*R[1] + PS2[2]*R[2]*R[2]*R[2]);
    if (deriv > -mjMINVAL) {
      break;
    }

    // delta
    mjtNum delta = -val/deriv;
    if (delta < tolerance) {
      break;
    }

    // update
    la += delta;
  }

  // compute normal given lambda
  nrm[0] = pos[0]/(S2[0]+la);
  nrm[1] = pos[1]/(S2[1]+la);
  nrm[2] = pos[2]/(S2[2]+la);
  mju_normalize3(nrm);

  return 1;
}


// fix normals if required
void mjc_fixNormal(const mjModel* m, const mjData* d, mjPreContact* con, int g1, int g2) {
  mjtNum dst1, dst2;

  // get geom ids and types
  int gid[2] = {g1, g2};
  mjtGeom type[2];
  for (int i=0; i < 2; i++) {
    if (gid[i] < 0) {
      type[i] = mjGEOM_NONE;
    } else {
      type[i] = m->geom_type[gid[i]];
    }

    // set to mjGEOM_NONE if type cannot be processed
    if (type[i] != mjGEOM_SPHERE && type[i] != mjGEOM_CAPSULE &&
        type[i] != mjGEOM_ELLIPSOID && type[i] != mjGEOM_CYLINDER) {
      type[i] = mjGEOM_NONE;
    }
  }

  // neither type can be processed: nothing to do
  if (type[0] == mjGEOM_NONE && type[1] == mjGEOM_NONE) {
    return;
  }

  // init normals
  mjtNum normal[2][3] = {
    {con->normal[0], con->normal[1], con->normal[2]},
    {-con->normal[0], -con->normal[1], -con->normal[2]}
  };

  // process geoms in type range
  int processed[2] = {0, 0};
  for (int i=0; i < 2; i++) {
    if (type[i] != mjGEOM_NONE) {
      // get geom mat and size
      mjtNum* mat = d->geom_xmat + 9*gid[i];
      mjtNum* size = m->geom_size + 3*gid[i];

      // map contact point and normal to local frame
      mjtNum dif[3], pos1[3], nrm[3];
      mju_sub3(dif, con->pos, d->geom_xpos+3*gid[i]);
      mju_mulMatTVec3(pos1, mat, dif);
      mju_mulMatTVec3(nrm, mat, normal[i]);

      // process according to type
      switch (type[i]) {
      case mjGEOM_SPHERE:
        mji_copy3(nrm, pos1);
        processed[i] = 1;
        break;

      case mjGEOM_CAPSULE:
        // Z: bottom cap
        if (pos1[2] < -size[1]) {
          nrm[2] = pos1[2]+size[1];
        }

        // Z: top cap
        else if (pos1[2] > size[1]) {
          nrm[2] = pos1[2]-size[1];
        }

        // Z: cylinder
        else {
          nrm[2] = 0;
        }

        // copy XY
        nrm[0] = pos1[0];
        nrm[1] = pos1[1];
        processed[i] = 1;
        break;

      case mjGEOM_ELLIPSOID:
        // guard against invalid ellipsoid size (just in case)
        if (size[0] < mjMINVAL || size[1] < mjMINVAL || size[2] < mjMINVAL) {
          break;
        }

        // compute elliptic distance^2
        dst1 = pos1[0]*pos1[0]/(size[0]*size[0]) +
               pos1[1]*pos1[1]/(size[1]*size[1]) +
               pos1[2]*pos1[2]/(size[2]*size[2]);

        // dispatch to inside or outside solver
        if (dst1 <= 1) {
          processed[i] = mjc_ellipsoidInside(nrm, pos1, size);
        } else {
          processed[i] = mjc_ellipsoidOutside(nrm, pos1, size);
        }
        break;

      case mjGEOM_CYLINDER:
        // skip if within 5% length of flat wall
        if (mju_abs(pos1[2]) > 0.95*size[1]) {
          break;
        }

        // compute distances to flat and round wall
        dst1 = mju_abs(size[1]-mju_abs(pos1[2]));
        dst2 = mju_abs(size[0]-mju_norm(pos1, 2));

        // require 4x closer to round than flat wall
        if (dst1 < 0.25*dst2) {
          break;
        }

        // set normal for round wall
        nrm[0] = pos1[0];
        nrm[1] = pos1[1];
        nrm[2] = 0;
        processed[i] = 1;
        break;

      default:
        // do nothing: only sphere, capsule, ellipsoid and cylinder are processed
        break;
      }

      // normalize and map normal to global frame
      if (processed[i]) {
        mju_normalize3(nrm);
        mji_mulMatVec3(normal[i], mat, nrm);
      }
    }
  }

  // both processed: average
  if (processed[0] && processed[1]) {
    mji_sub3(con->normal, normal[0], normal[1]);
    mju_normalize3(con->normal);
  }

  // first processed: copy
  else if (processed[0]) {
    mji_copy3(con->normal, normal[0]);
  }

  // second processed: copy reverse
  else if (processed[1]) {
    mji_scl3(con->normal, normal[1], -1);
  }
}


//----------------------------  flex collisions ---------------------------------------------

// geom-elem or elem-elem or vert-elem convex collision using ccd
int mjc_ConvexElem(const mjModel* m, mjData* d, mjPreContact* con, int g1, int f1, int e1, int v1,
                   int f2, int e2, mjtNum margin) {
  mjCCDObj obj1, obj2;
  mjc_initCCDObj(&obj1, m, d, g1, margin);
  mjc_initCCDObj(&obj2, m, d, -1, margin);
  mjc_setCCDObjFlex(&obj1, f1, e1, v1);
  mjc_setCCDObjFlex(&obj2, f2, e2, -1);

  // find contacts
  int ncon = mjc_penetration(m, d, &obj1, &obj2, con, 1, margin);

  // fix normals for 2D flex
  if (ncon && !mjDISABLED(mjDSBL_NATIVECCD)) {
    // check if either object is a 2D flex
    int isflex2d = 0;
    if (f1 >= 0 && m->flex_dim[f1] == 2) isflex2d = 1;
    if (f2 >= 0 && m->flex_dim[f2] == 2) isflex2d = 1;

    if (isflex2d) {
      for (int i = 0; i < ncon; i++) {
        mjc_fixNormal(m, d, con + i, g1, -1);
      }
    }
  }
  return ncon;
}


// test a height field and a flex element for collision
int mjc_HFieldElem(const mjModel* m, mjData* d, mjPreContact* con, int g, int f, int e,
                   mjtNum margin) {
  mjtNum vec[3], dx, dy;
  mjtNum xmin, xmax, ymin, ymax, zmin, zmax;
  int dr[2], cnt, rmin, rmax, cmin, cmax;
  mjCCDObj obj1;
  obj1.center = mjc_center;
  obj1.support = mjc_prism_support;

  // get hfield info
  int hid = m->geom_dataid[g];
  int nrow = m->hfield_nrow[hid];
  int ncol = m->hfield_ncol[hid];
  const mjtNum* hpos = d->geom_xpos + 3*g;
  const mjtNum* hmat = d->geom_xmat + 9*g;
  const mjtNum* hsize = m->hfield_size + 4*hid;
  const float* hdata = m->hfield_data + m->hfield_adr[hid];

  // get elem indo
  int dim = m->flex_dim[f];
  const int* edata = m->flex_elem + m->flex_elemdataadr[f] + e*(dim+1);
  mjtNum* evert[4] = {NULL, NULL, NULL, NULL};
  for (int i=0; i <= dim; i++) {
    evert[i] = d->flexvert_xpos + 3*(m->flex_vertadr[f] + edata[i]);
  }
  mjtNum* ecenter = d->flexelem_aabb + 6*(m->flex_elemadr[f]+e);

  // ccd-related
  mjCCDObj obj2;
  mjc_initCCDObj(&obj2, m, d, -1, margin);
  mjc_setCCDObjFlex(&obj2, f, e, -1);
  //------------------------------------- AABB computation, box-box test

  // save elem vertices, transform to hfield frame
  mjtNum savevert[4][3];
  for (int i=0; i <= dim; i++) {
    mji_copy3(savevert[i], evert[i]);
    mji_sub3(vec, evert[i], hpos);
    mji_mulMatTVec3(evert[i], hmat, vec);
  }

  // save elem center, transform to hfield frame
  mjtNum savecenter[3];
  mji_copy3(savecenter, ecenter);
  mji_sub3(vec, ecenter, hpos);
  mji_mulMatTVec3(ecenter, hmat, vec);

  // compute elem bounding box (in hfield frame)
  xmin = xmax = evert[0][0];
  ymin = ymax = evert[0][1];
  zmin = zmax = evert[0][2];
  for (int i=1; i <= dim; i++) {
    xmin = mju_min(xmin, evert[i][0]);
    xmax = mju_max(xmax, evert[i][0]);
    ymin = mju_min(ymin, evert[i][1]);
    ymax = mju_max(ymax, evert[i][1]);
    zmin = mju_min(zmin, evert[i][2]);
    zmax = mju_max(zmax, evert[i][2]);
  }

  // box-box test
  if ((xmin-margin > hsize[0]) || (xmax+margin < -hsize[0]) ||
      (ymin-margin > hsize[1]) || (ymax+margin < -hsize[1]) ||
      (zmin-margin > hsize[2]) || (zmax+margin < -hsize[3])) {
    // restore vertices and center
    for (int i=0; i <= dim; i++) {
      mji_copy3(evert[i], savevert[i]);
    }
    mji_copy3(ecenter, savecenter);

    return 0;
  }

  // compute sub-grid bounds
  cmin = (int) mju_floor((xmin + hsize[0]) / (2*hsize[0]) * (ncol-1));
  cmax = (int) mju_ceil ((xmax + hsize[0]) / (2*hsize[0]) * (ncol-1));
  rmin = (int) mju_floor((ymin + hsize[1]) / (2*hsize[1]) * (nrow-1));
  rmax = (int) mju_ceil ((ymax + hsize[1]) / (2*hsize[1]) * (nrow-1));
  cmin = mjMAX(0, cmin);
  cmax = mjMIN(ncol-1, cmax);
  rmin = mjMAX(0, rmin);
  rmax = mjMIN(nrow-1, rmax);

  //------------------------------------- collision testing

  // compute real-valued grid step, and triangulation direction
  dx = (2.0*hsize[0]) / (ncol-1);
  dy = (2.0*hsize[1]) / (nrow-1);
  dr[0] = 1;
  dr[1] = 0;

  // set zbottom value using base size
  mjtNum (*prism)[3] = obj1.data.hfield.prism;
  prism[0][2] = prism[1][2] = prism[2][2] = -hsize[3];

  // process all prisms in sub-grid
  cnt = 0;
  for (int r=rmin; r < rmax; r++) {
    int nvert = 0;
    for (int c=cmin; c <= cmax; c++) {
      for (int k=0; k < 2; k++) {
        // send vertex to prism constructor
        addVert(&obj1, dx*c-hsize[0], dy*(r+dr[k])-hsize[1],
                hdata[(r+dr[k])*ncol+c]*hsize[2]+margin);

        // check for enough vertices
        if (++nvert > 2) {
          // prism height test
          if (prism[3][2] < zmin && prism[4][2] < zmin && prism[5][2] < zmin) {
            continue;
          }

          // run ccd, save contact
          if (mjc_penetration(m, d, &obj1, &obj2, con + cnt, 1, 0.0)) {
            // transform to global coordinates
            mji_zero3(con[cnt].tangent);
            mju_mulMatVec3(con[cnt].normal, hmat, con[cnt].normal);
            mju_mulMatVec3(con[cnt].pos, hmat, con[cnt].pos);
            mji_addTo3(con[cnt].pos, hpos);

            // count, stop if max number reached
            cnt++;
            if (cnt >= mjMAXCONPAIR) {
              r = rmax+1;
              c = cmax+1;
              k = 3;
              break;
            }
          }
        }
      }
    }
  }

  // restore elem vertices and center
  for (int i=0; i <= dim; i++) {
    mji_copy3(evert[i], savevert[i]);
  }
  mji_copy3(ecenter, savecenter);

  return cnt;
}