# Copyright 2023 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.
# ==============================================================================
"""Convex collisions."""

import functools
from typing import Tuple

import jax
from jax import numpy as jp
from mujoco.mjx._src import math
from mujoco.mjx._src import mesh
# pylint: disable=g-importing-member
from mujoco.mjx._src.collision_types import Collision
from mujoco.mjx._src.collision_types import ConvexInfo
from mujoco.mjx._src.collision_types import FunctionKey
from mujoco.mjx._src.collision_types import GeomInfo
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import GeomType
from mujoco.mjx._src.types import Model
# pylint: enable=g-importing-member


def collider(ncon: int):
  """Wraps collision functions for use by collision_driver."""

  def wrapper(func):
    def collide(
        m: Model, d: Data, key: FunctionKey, geom: jax.Array
    ) -> Collision:
      g1, g2 = geom.T
      infos = [
          GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1]),
          GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2]),
      ]
      in_axes = [0, 0]
      for i in [0, 1]:
        if key.types[i] == GeomType.BOX:
          infos[i] = mesh.box(infos[i])
          in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
              pos=0, mat=0, face=0, vert=0
          )
        elif key.types[i] == GeomType.MESH:
          infos[i] = mesh.convex(m, key.data_ids[i], infos[i])
          in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
              pos=0, mat=0
          )
      dist, pos, frame = jax.vmap(func, in_axes=in_axes)(*infos)
      if ncon > 1:
        return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
      return dist, pos, frame

    collide.ncon = ncon
    return collide

  return wrapper


def _closest_segment_point_plane(
    a: jax.Array, b: jax.Array, p0: jax.Array, plane_normal: jax.Array
) -> jax.Array:
  """Gets the closest point between a line segment and a plane.

  Args:
    a: first line segment point
    b: second line segment point
    p0: point on plane
    plane_normal: plane normal

  Returns:
    closest point between the line segment and the plane
  """
  # Parametrize a line segment as S(t) = a + t * (b - a), plug it into the plane
  # equation dot(n, S(t)) - d = 0, then solve for t to get the line-plane
  # intersection. We then clip t to be in [0, 1] to be on the line segment.
  n = plane_normal
  d = jp.sum(p0 * n)  # shortest distance from origin to plane
  denom = jp.sum(n * (b - a))
  t = (d - jp.sum(n * a)) / (denom + 1e-6 * (denom == 0.0))
  t = jp.clip(t, 0, 1)
  segment_point = a + t * (b - a)

  return segment_point


def _closest_triangle_point(
    p0: jax.Array, p1: jax.Array, p2: jax.Array, pt: jax.Array
) -> jax.Array:
  """Gets the closest point between a triangle and a point in space.

  Args:
    p0: triangle point
    p1: triangle point
    p2: triangle point
    pt: point to test

  Returns:
    closest point on the triangle w.r.t point pt
  """
  # Parametrize the triangle s.t. a point inside the triangle is
  # Q = p0 + u * e0 + v * e1, when 0 <= u <= 1, 0 <= v <= 1, and
  # 0 <= u + v <= 1. Let e0 = (p1 - p0) and e1 = (p2 - p0).
  # We analytically minimize the distance between the point pt and Q.
  e0 = p1 - p0
  e1 = p2 - p0
  a = e0.dot(e0)
  b = e0.dot(e1)
  c = e1.dot(e1)
  d = pt - p0
  # The determinant is 0 only if the angle between e1 and e0 is 0
  # (i.e. the triangle has overlapping lines).
  det = a * c - b * b
  u = (c * e0.dot(d) - b * e1.dot(d)) / det
  v = (-b * e0.dot(d) + a * e1.dot(d)) / det
  inside = (0 <= u) & (u <= 1) & (0 <= v) & (v <= 1) & (u + v <= 1)
  closest_p = p0 + u * e0 + v * e1
  d0 = (closest_p - pt).dot(closest_p - pt)

  # If the closest point is outside the triangle, it must be on an edge, so we
  # check each triangle edge for a closest point to the point pt.
  closest_p1, d1 = math.closest_segment_point_and_dist(p0, p1, pt)
  closest_p = jp.where((d0 < d1) & inside, closest_p, closest_p1)
  min_d = jp.where((d0 < d1) & inside, d0, d1)

  closest_p2, d2 = math.closest_segment_point_and_dist(p1, p2, pt)
  closest_p = jp.where(d2 < min_d, closest_p2, closest_p)
  min_d = jp.minimum(min_d, d2)

  closest_p3, d3 = math.closest_segment_point_and_dist(p2, p0, pt)
  closest_p = jp.where(d3 < min_d, closest_p3, closest_p)

  return closest_p


def _closest_segment_triangle_points(
    a: jax.Array,
    b: jax.Array,
    p0: jax.Array,
    p1: jax.Array,
    p2: jax.Array,
    triangle_normal: jax.Array,
) -> Tuple[jax.Array, jax.Array]:
  """Gets the closest points between a line segment and triangle.

  Args:
    a: first line segment point
    b: second line segment point
    p0: triangle point
    p1: triangle point
    p2: triangle point
    triangle_normal: normal of triangle

  Returns:
    closest point on the triangle w.r.t the line segment
  """
  # The closest triangle point is either on the edge or within the triangle.
  # First check triangle edges for the closest point.
  # TODO(robotics-simulation): consider vmapping over closest point functions
  seg_pt1, tri_pt1 = math.closest_segment_to_segment_points(a, b, p0, p1)
  d1 = (seg_pt1 - tri_pt1).dot(seg_pt1 - tri_pt1)
  seg_pt2, tri_pt2 = math.closest_segment_to_segment_points(a, b, p1, p2)
  d2 = (seg_pt2 - tri_pt2).dot(seg_pt2 - tri_pt2)
  seg_pt3, tri_pt3 = math.closest_segment_to_segment_points(a, b, p0, p2)
  d3 = (seg_pt3 - tri_pt3).dot(seg_pt3 - tri_pt3)

  # Next, handle the case where the closest triangle point is inside the
  # triangle. Either the line segment intersects the triangle or a segment
  # endpoint is closest to a point inside the triangle.
  seg_pt4 = _closest_segment_point_plane(a, b, p0, triangle_normal)
  tri_pt4 = _closest_triangle_point(p0, p1, p2, seg_pt4)
  d4 = (seg_pt4 - tri_pt4).dot(seg_pt4 - tri_pt4)

  # Get the point with minimum distance from the line segment point to the
  # triangle point.
  distance = jp.array([[d1, d2, d3, d4]])
  min_dist = jp.amin(distance)
  mask = (distance == min_dist).T
  seg_pt = jp.array([seg_pt1, seg_pt2, seg_pt3, seg_pt4]) * mask
  tri_pt = jp.array([tri_pt1, tri_pt2, tri_pt3, tri_pt4]) * mask
  seg_pt = jp.sum(seg_pt, axis=0) / jp.sum(mask)
  tri_pt = jp.sum(tri_pt, axis=0) / jp.sum(mask)

  return seg_pt, tri_pt


def _manifold_points(
    poly: jax.Array, poly_mask: jax.Array, poly_norm: jax.Array
) -> jax.Array:
  """Chooses four points on the polygon with approximately maximal area."""
  dist_mask = jp.where(poly_mask, 0.0, -1e6)
  a_idx = jp.argmax(dist_mask)
  a = poly[a_idx]
  # choose point b furthest from a
  b_idx = (((a - poly) ** 2).sum(axis=1) + dist_mask).argmax()
  b = poly[b_idx]
  # choose point c furthest along the axis orthogonal to (a-b)
  ab = jp.cross(poly_norm, a - b)
  ap = a - poly
  c_idx = (jp.abs(ap.dot(ab)) + dist_mask).argmax()
  c = poly[c_idx]
  # choose point d furthest from the other two triangle edges
  ac = jp.cross(poly_norm, a - c)
  bc = jp.cross(poly_norm, b - c)
  bp = b - poly
  dist_bp = jp.abs(bp.dot(bc)) + dist_mask
  dist_ap = jp.abs(ap.dot(ac)) + dist_mask
  d_idx = (dist_bp + dist_ap).argmax() % poly.shape[0]
  return jp.array([a_idx, b_idx, c_idx, d_idx])


@collider(ncon=4)
def plane_convex(plane: GeomInfo, convex: ConvexInfo) -> Collision:
  """Calculates contacts between a plane and a convex object."""
  vert = convex.vert

  # get points in the convex frame
  plane_pos = convex.mat.T @ (plane.pos - convex.pos)
  n = convex.mat.T @ plane.mat[:, 2]
  support = (plane_pos - vert) @ n
  idx = _manifold_points(vert, support > jp.maximum(0, support.max() - 1e-4), n)
  pos = vert[idx]

  # convert to world frame
  pos = convex.pos + pos @ convex.mat.T
  n = plane.mat[:, 2]

  frame = jp.stack([math.make_frame(n)] * 4, axis=0)
  unique = jp.tril(idx == idx[:, None]).sum(axis=1) == 1
  dist = jp.where(unique, -support[idx], 1)
  pos = pos - 0.5 * dist[:, None] * n
  return dist, pos, frame


@collider(ncon=1)
def sphere_convex(sphere: GeomInfo, convex: ConvexInfo) -> Collision:
  """Calculates contact between a sphere and a convex object."""
  faces = convex.face
  normals = convex.face_normal

  # Put sphere in convex frame.
  sphere_pos = convex.mat.T @ (sphere.pos - convex.pos)

  # Get support from face normals.
  @jax.vmap
  def get_support(faces, normal):
    pos = sphere_pos - normal * sphere.size[0]
    return jp.dot(pos - faces[0], normal)

  support = get_support(faces, normals)
  has_separating_axis = jp.any(support >= 0)

  # Pick the face with the best separating axis.
  best_idx = support.argmax()
  face = faces[best_idx]
  face_normal = normals[best_idx]

  # Get closest point between the polygon face and the sphere center point.
  # Project the sphere center point onto poly plane. If it's inside polygon
  # side planes, then we're done.
  pt = _project_pt_onto_plane(sphere_pos, face[0], face_normal)
  edge_p0 = jp.roll(face, 1, axis=0)
  edge_p1 = face
  side_normals = jax.vmap(jp.cross, in_axes=[0, None])(
      edge_p1 - edge_p0,
      face_normal,
  )
  edge_dist = jax.vmap(
      lambda plane_pt, plane_norm: (pt - plane_pt).dot(plane_norm)
  )(edge_p0, side_normals)
  pt_on_face = jp.all(edge_dist <= 0)  # lte to handle degenerate edges

  # If the point is outside side planes, project onto the closest side plane
  # that the point is in front of.
  degenerate_edge = jp.all(side_normals == 0, axis=1)
  behind = edge_dist < 0.0
  edge_dist = jp.where(degenerate_edge | behind, 1e12, edge_dist)
  idx = edge_dist.argmin()
  edge_pt = math.closest_segment_point(edge_p0[idx], edge_p1[idx], pt)
  pt = jp.where(pt_on_face, pt, edge_pt)

  # Get the normal, dist, and contact position.
  pt_normal, d = math.normalize_with_norm(pt - sphere_pos)
  # Ensure normal points towards convex centroid.
  inside = jp.dot(pt, pt_normal) > 0
  sign = jp.where(inside, -1, 1)
  n = jp.where(pt_on_face | (d < 1e-6), -face_normal, sign * pt_normal)
  d *= sign

  spt = sphere_pos + n * sphere.size[0]
  dist = jp.where(has_separating_axis, 1.0, d - sphere.size[0])
  pos = (pt + spt) * 0.5

  # Go back to world frame.
  n = convex.mat @ n
  pos = convex.mat @ pos + convex.pos

  return dist, pos, math.make_frame(n)


@collider(ncon=2)
def capsule_convex(cap: GeomInfo, convex: ConvexInfo) -> Collision:
  """Calculates contacts between a capsule and a convex object."""
  # Get convex transformed normals, faces, and vertices.
  faces = convex.face
  normals = convex.face_normal

  # Put capsule in convex frame.
  cap_pos = convex.mat.T @ (cap.pos - convex.pos)
  axis, length = cap.mat[:, 2], cap.size[1]
  axis = convex.mat.T @ axis
  seg = axis * length
  cap_pts = jp.array([
      cap_pos - seg,
      cap_pos + seg,
  ])

  # Get support from face normals.
  @jax.vmap
  def get_support(face, normal):
    pts = cap_pts - normal * cap.size[0]
    sup = jax.vmap(lambda x: jp.dot(x - face[0], normal))(pts)
    return sup.min()

  support = get_support(faces, normals)
  has_support = jp.all(support < 0)

  # Pick the face with minimal penetration.
  best_idx = support.argmax()
  face = faces[best_idx]
  normal = normals[best_idx]

  # Clip the segment against side planes and create two contact points against
  # the face.
  edge_p0 = jp.roll(face, 1, axis=0)
  edge_p1 = face
  side_planes = jax.vmap(jp.cross, in_axes=[0, None])(
      edge_p1 - edge_p0,
      normal,
  )
  cap_pts_clipped, mask = _clip_edge_to_planes(
      cap_pts[0], cap_pts[1], edge_p0, side_planes
  )
  cap_pts_clipped = cap_pts_clipped - normal * cap.size[0]
  face_pts = jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])(
      cap_pts_clipped, face[0], normal
  )
  # Create variables for the face contact.
  pos = (cap_pts_clipped + face_pts) * 0.5
  contact_normal = -jp.stack([normal] * 2, 0)
  face_penetration = jp.where(
      mask & has_support, jp.dot(face_pts - cap_pts_clipped, normal), -1
  )

  # Pick a potential shallow edge contact.
  def get_edge_axis(edge):
    edge_closest_pt, cap_closest_pt = math.closest_segment_to_segment_points(
        edge[0], edge[1], cap_pts[0], cap_pts[1]
    )
    edge_dir = edge_closest_pt - cap_closest_pt
    degenerate_edge_dir = jp.sum(jp.square(edge_dir)) < 1e-6
    edge_axis, edge_dist = math.normalize_with_norm(edge_dir)
    return (
        edge_dist,
        edge_axis,
        degenerate_edge_dir,
        edge_closest_pt,
        cap_closest_pt,
    )

  edge = jp.take(convex.vert, convex.edge, axis=0)
  edge_face_normal = convex.edge_face_normal  # pytype: disable=attribute-error

  res = jax.vmap(get_edge_axis)(edge.reshape(-1, 2, 3))
  e_idx = jp.abs(res[0]).argmin()
  (
      edge_dist,
      edge_axis,
      degenerate_edge_dir,
      edge_closest_pt,
      cap_closest_pt,
  ) = jax.tree_util.tree_map(lambda x, i=e_idx: jp.take(x, i, axis=0), res)

  edge_face_normals = edge_face_normal[e_idx]
  edge_voronoi_front = ((edge_face_normals @ edge_axis) < 0).all()
  shallow = ~degenerate_edge_dir & edge_voronoi_front
  edge_penetration = jp.where(shallow, cap.size[0] - edge_dist, -1)

  # Determine edge contact position.
  edge_pos = (
      edge_closest_pt + (cap_closest_pt + edge_axis * cap.size[0])
  ) * 0.5
  edge_dir_parallel_to_face = (
      jp.abs(edge_axis.dot(normal)) > 0.99
  ) & ~degenerate_edge_dir
  min_face_penetration = face_penetration.min()
  has_edge_contact = (
      (edge_penetration > 0)
      # prefer edge contact if the edge is smaller than face penetration
      & jp.where(
          min_face_penetration > 0,
          edge_penetration < min_face_penetration,
          True,
      )
      # prefer face contact if the edge axis is parallel to the face normal
      & ~edge_dir_parallel_to_face
      # make sure we have a shallow contact
      & edge_voronoi_front
  )

  # Get the contact info.
  pos = jp.where(has_edge_contact, pos.at[0].set(edge_pos), pos)
  n = jp.where(
      has_edge_contact, contact_normal.at[0].set(edge_axis), contact_normal
  )

  # Go back to world frame.
  pos = convex.pos + pos @ convex.mat.T
  n = n @ convex.mat.T

  dist = -jp.where(
      has_edge_contact, jp.array([edge_penetration, -1]), face_penetration
  )
  frame = jax.vmap(math.make_frame)(n)
  return dist, pos, frame


def _project_pt_onto_plane(
    pt: jax.Array, plane_pt: jax.Array, plane_normal: jax.Array
) -> jax.Array:
  """Projects a point onto a plane along the plane normal."""
  dist = (pt - plane_pt).dot(plane_normal)
  return pt - dist * plane_normal


def _project_poly_onto_plane(
    poly: jax.Array, plane_pt: jax.Array, plane_normal: jax.Array
) -> jax.Array:
  """Projects a polygon onto a plane using the plane normal."""
  return jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])(
      poly, plane_pt, math.normalize(plane_normal)
  )


def _project_poly_onto_poly_plane(
    poly1: jax.Array, norm1: jax.Array, poly2: jax.Array, norm2: jax.Array
) -> jax.Array:
  """Projects poly1 onto the poly2 plane along poly1's normal."""
  d = poly2[0].dot(norm2)
  denom = norm1.dot(norm2)
  t = (d - poly1.dot(norm2)) / (denom + 1e-6 * (denom == 0.0))
  new_poly = poly1 + t.reshape(-1, 1) * norm1
  return new_poly


def _point_in_front_of_plane(
    plane_pt: jax.Array, plane_normal: jax.Array, pt: jax.Array
) -> jax.Array:
  """Checks if a point is strictly in front of a plane."""
  return (pt - plane_pt).dot(plane_normal) > 1e-6


def _clip_edge_to_planes(
    edge_p0: jax.Array,
    edge_p1: jax.Array,
    plane_pts: jax.Array,
    plane_normals: jax.Array,
) -> Tuple[jax.Array, jax.Array]:
  """Clips an edge against side planes.

  We return two clipped points, and a mask to include the new edge or not.

  Args:
    edge_p0: the first point on the edge
    edge_p1: the second point on the edge
    plane_pts: side plane points
    plane_normals: side plane normals

  Returns:
    new_ps: new edge points that are clipped against side planes
    mask: a boolean mask, True if an edge point is a valid clipped point and
    False otherwise
  """
  p0, p1 = edge_p0, edge_p1
  p0_in_front = jax.vmap(jp.dot)(p0 - plane_pts, plane_normals) > 1e-6
  p1_in_front = jax.vmap(jp.dot)(p1 - plane_pts, plane_normals) > 1e-6

  # Get candidate clipped points along line segment (p0, p1) by clipping against
  # all clipping planes.
  candidate_clipped_ps = jax.vmap(
      _closest_segment_point_plane, in_axes=[None, None, 0, 0]
  )(p0, p1, plane_pts, plane_normals)

  def clip_edge_point(p0, p1, p0_in_front, clipped_ps):
    @jax.vmap
    def choose_edge_point(in_front, clipped_p):
      return jp.where(in_front, clipped_p, p0)

    # Pick the clipped point if p0 is in front of the clipping plane. Otherwise
    # keep p0 as the edge point.
    new_edge_ps = choose_edge_point(p0_in_front, clipped_ps)

    # Pick the clipped point that is most along the edge direction.
    # This degenerates to picking the original point p0 if p0 is *not* in front
    # of any clipping planes.
    dists = jp.dot(new_edge_ps - p0, p1 - p0)
    new_edge_p = new_edge_ps[jp.argmax(dists)]
    return new_edge_p

  # Clip each edge point.
  new_p0 = clip_edge_point(p0, p1, p0_in_front, candidate_clipped_ps)
  new_p1 = clip_edge_point(p1, p0, p1_in_front, candidate_clipped_ps)
  clipped_pts = jp.array([new_p0, new_p1])

  # Keep the original points if both points are in front of any of the clipping
  # planes, rather than creating a new clipped edge. If the entire subject edge
  # is in front of any clipping plane, we need to grab an edge from the clipping
  # polygon instead.
  both_in_front = p0_in_front & p1_in_front
  mask = ~jp.any(both_in_front)
  new_ps = jp.where(mask, clipped_pts, jp.array([p0, p1]))
  # Mask out crossing clipped edge points.
  mask = jp.where((p0 - p1).dot(new_ps[0] - new_ps[1]) < 0, False, mask)
  return new_ps, jp.array([mask, mask])


def _clip(
    clipping_poly: jax.Array,
    subject_poly: jax.Array,
    clipping_normal: jax.Array,
    subject_normal: jax.Array,
) -> Tuple[jax.Array, jax.Array]:
  """Clips a subject polygon against a clipping polygon.

  A parallelized clipping algorithm for convex polygons. The result is a set of
  vertices on the clipped subject polygon in the subject polygon plane.

  Args:
    clipping_poly: the polygon that we use to clip the subject polygon against
    subject_poly: the polygon that gets clipped
    clipping_normal: normal of the clipping polygon
    subject_normal: normal of the subject polygon

  Returns:
    clipped_pts: points on the clipped polygon
    mask: True if a point is in the clipping polygon, False otherwise
  """
  # Get clipping edge points, edge planes, and edge normals.
  clipping_p0 = jp.roll(clipping_poly, 1, axis=0)
  clipping_plane_pts = clipping_p0
  clipping_p1 = clipping_poly
  clipping_plane_normals = jax.vmap(jp.cross, in_axes=[0, None])(
      clipping_p1 - clipping_p0,
      clipping_normal,
  )

  # Get subject edge points, edge planes, and edge normals.
  subject_edge_p0 = jp.roll(subject_poly, 1, axis=0)
  subject_plane_pts = subject_edge_p0
  subject_edge_p1 = subject_poly
  subject_plane_normals = jax.vmap(jp.cross, in_axes=[0, None])(
      subject_edge_p1 - subject_edge_p0,
      subject_normal,
  )

  # Clip all edges of the subject poly against clipping side planes.
  clipped_edges0, masks0 = jax.vmap(
      _clip_edge_to_planes, in_axes=[0, 0, None, None]
  )(
      subject_edge_p0,
      subject_edge_p1,
      clipping_plane_pts,
      clipping_plane_normals,
  )

  # Project the clipping poly onto the subject plane.
  clipping_p0_s = _project_poly_onto_poly_plane(
      clipping_p0, clipping_normal, subject_poly, subject_normal
  )
  clipping_p1_s = _project_poly_onto_poly_plane(
      clipping_p1, clipping_normal, subject_poly, subject_normal
  )

  # Clip all edges of the clipping poly against subject planes.
  clipped_edges1, masks1 = jax.vmap(
      _clip_edge_to_planes, in_axes=[0, 0, None, None]
  )(clipping_p0_s, clipping_p1_s, subject_plane_pts, subject_plane_normals)

  # Merge the points and reshape.
  clipped_edges = jp.concatenate([clipped_edges0, clipped_edges1])
  masks = jp.concatenate([masks0, masks1])
  clipped_points = clipped_edges.reshape((-1, 3))
  mask = masks.reshape(-1)

  return clipped_points, mask


def _create_contact_manifold(
    clipping_poly: jax.Array,
    subject_poly: jax.Array,
    clipping_norm: jax.Array,
    subject_norm: jax.Array,
    sep_axis: jax.Array,
) -> Tuple[jax.Array, jax.Array, jax.Array]:
  """Creates a contact manifold between two convex polygons.

  The polygon faces are expected to have a counter clockwise winding order so
  that clipping plane normals point away from the polygon center.

  Args:
    clipping_poly: the reference polygon to clip the contact against.
    subject_poly: the subject polygon to clip contacts onto.
    clipping_norm: the clipping polygon normal.
    subject_norm: the subject polygon normal.
    sep_axis: the separating axis

  Returns:
    tuple of dist, pos, and normal
  """
  # Clip the subject (incident) face onto the clipping (reference) face.
  # The incident points are clipped points on the subject polygon.
  poly_incident, mask = _clip(
      clipping_poly, subject_poly, clipping_norm, subject_norm
  )
  # The reference points are clipped points on the clipping polygon.
  poly_ref = _project_poly_onto_plane(
      poly_incident, clipping_poly[0], clipping_norm
  )
  behind_clipping_plane = _point_in_front_of_plane(
      clipping_poly[0], -clipping_norm, poly_incident
  )
  mask = mask & behind_clipping_plane

  # Choose four contact points.
  best = _manifold_points(poly_ref, mask, clipping_norm)
  contact_pts = jp.take(poly_ref, best, axis=0)
  mask_pts = jp.take(mask, best, axis=0)
  penetration_dir = jp.take(poly_incident, best, axis=0) - contact_pts
  penetration = penetration_dir.dot(-clipping_norm)

  dist = jp.where(mask_pts, -penetration, jp.ones_like(penetration))
  pos = contact_pts
  normal = -jp.stack([sep_axis] * 4, 0)
  return dist, pos, normal


def _sat_bruteforce(
    faces_a: jax.Array,
    faces_b: jax.Array,
    vertices_a: jax.Array,
    vertices_b: jax.Array,
    normals_a: jax.Array,
    normals_b: jax.Array,
    unique_edges_a: jax.Array,
    unique_edges_b: jax.Array,
) -> Tuple[jax.Array, jax.Array, jax.Array]:
  """Runs the Separating Axis Test for a pair of hulls.

  Given two convex hulls, the Separating Axis Test finds a separating axis
  between all edge pairs and face pairs. Edge pairs create a single contact
  point and face pairs create a contact manifold (up to four contact points).
  We return both the edge and face contacts. Valid contacts can be checked with
  dist < 0. Resulting edge contacts should be preferred over face contacts.

  This method checks all separating axes via a brute force support function, and
  is thus costly to run over large meshes, but is more performant for smaller
  meshes (boxes, tetrahedra, etc.).

  Args:
    faces_a: Faces for hull A.
    faces_b: Faces for hull B.
    vertices_a: Vertices for hull A.
    vertices_b: Vertices for hull B.
    normals_a: Normal vectors for hull A faces.
    normals_b: Normal vectors for hull B faces.
    unique_edges_a: Unique edges for hull A.
    unique_edges_b: Unique edges for hull B.

  Returns:
    tuple of dist, pos, and normal
  """
  # get the separating axes
  v_norm = jax.vmap(math.normalize)
  edge_dir_a = v_norm(unique_edges_a[:, 0] - unique_edges_a[:, 1])
  edge_dir_b = v_norm(unique_edges_b[:, 0] - unique_edges_b[:, 1])
  edge_dir_a_r = jp.tile(edge_dir_a, reps=(unique_edges_b.shape[0], 1))
  edge_dir_b_r = jp.repeat(edge_dir_b, repeats=unique_edges_a.shape[0], axis=0)
  edge_axes = jax.vmap(jp.cross)(edge_dir_a_r, edge_dir_b_r)
  degenerate_edge_axes = (edge_axes**2).sum(axis=1) < 1e-6
  edge_axes = jax.vmap(lambda x: math.normalize(x, axis=0))(edge_axes)
  n_norm = normals_a.shape[0] + normals_b.shape[0]
  degenerate_axes = jp.concatenate(
      [jp.array([False] * n_norm), degenerate_edge_axes]
  )

  axes = jp.concatenate([normals_a, normals_b, edge_axes])

  # for each separating axis, get the support
  @jax.vmap
  def get_support(axis, is_degenerate):
    # the matmul here is more performant with vmap(dot)
    dot = functools.partial(jp.dot, precision=jax.lax.Precision.HIGH)
    support_a = jax.vmap(dot, in_axes=[None, 0])(axis, vertices_a)
    support_b = jax.vmap(dot, in_axes=[None, 0])(axis, vertices_b)
    dist1 = support_a.max() - support_b.min()
    dist2 = support_b.max() - support_a.min()
    sign = jp.where(dist1 > dist2, -1, 1)
    dist = jp.minimum(dist1, dist2)
    dist = jp.where(~is_degenerate, dist, 1e6)  # degenerate axis
    return dist, sign

  support, sign = get_support(axes, degenerate_axes)

  # choose the best separating axis
  best_idx = jp.argmin(support)
  best_sign = sign[best_idx]
  best_axis = axes[best_idx]
  is_edge_contact = best_idx >= (normals_a.shape[0] + normals_b.shape[0])

  # get the (reference) face most aligned with the separating axis
  dist_a = normals_a @ best_axis
  dist_b = normals_b @ best_axis
  a_max = dist_a.argmax()
  b_max = dist_b.argmax()
  a_min = dist_a.argmin()
  b_min = dist_b.argmin()

  ref_face = jp.where(best_sign > 0, faces_a[a_max], faces_b[b_max])
  ref_face_norm = jp.where(best_sign > 0, normals_a[a_max], normals_b[b_max])
  incident_face = jp.where(best_sign > 0, faces_b[b_min], faces_a[a_min])
  incident_face_norm = jp.where(
      best_sign > 0, normals_b[b_min], normals_a[a_min]
  )

  dist, pos, normal = _create_contact_manifold(
      ref_face,
      incident_face,
      ref_face_norm,
      incident_face_norm,
      -best_sign * best_axis,
  )

  # For edge contacts, we use the clipped face point, mainly for performance
  # reasons. For small penetration, the clipped face point is roughly the edge
  # contact point.
  idx = dist.argmin()
  dist = jp.where(
      is_edge_contact,
      jp.array([dist[idx], 1, 1, 1]),
      dist,
  )
  pos = jp.where(is_edge_contact, jp.tile(pos[idx], (4, 1)), pos)

  return dist, pos, normal


def _arcs_intersect(
    a: jax.Array, b: jax.Array, c: jax.Array, d: jax.Array
) -> jax.Array:
  """Tests if arcs AB and CD on the unit sphere intersect."""
  ba, dc = jp.cross(b, a), jp.cross(d, c)
  cba, dba = jp.dot(c, ba), jp.dot(d, ba)
  adc, bdc = jp.dot(a, dc), jp.dot(b, dc)
  return (cba * dba < 0) & (adc * bdc < 0) & (cba * bdc > 0)


def _sat_gaussmap(
    centroid_a: jax.Array,
    faces_a: jax.Array,
    faces_b: jax.Array,
    vertices_a: jax.Array,
    vertices_b: jax.Array,
    normals_a: jax.Array,
    normals_b: jax.Array,
    edges_a: jax.Array,
    edges_b: jax.Array,
    edge_face_normals_a: jax.Array,
    edge_face_normals_b: jax.Array,
) -> Tuple[jax.Array, jax.Array, jax.Array]:
  """Runs the Separating Axis Test for a pair of hulls.

  Runs the separating axis test for all faces. Tests edge separating axes via
  edge intersections on gauss maps for all edge pairs. h/t to Dirk Gregorius
  for the implementation details and gauss map trick.

  Args:
    centroid_a: Centroid of hull A.
    faces_a: Faces for hull A.
    faces_b: Faces for hull B.
    vertices_a: Vertices for hull A.
    vertices_b: Vertices for hull B.
    normals_a: Normal vectors for hull A faces.
    normals_b: Normal vectors for hull B faces.
    edges_a: Edges for hull A.
    edges_b: Edges for hull B.
    edge_face_normals_a: Face normals for edges in hull A.
    edge_face_normals_b: Face normals for edges in hull B.

  Returns:
    tuple of dist, pos, and normal
  """
  # Handle face separating axes.
  axes = jp.concatenate([normals_a, -normals_b])

  @jax.vmap
  def get_support(axis):
    # the matmul here is more performant with vmap(dot)
    dot = functools.partial(jp.dot, precision=jax.lax.Precision.HIGH)
    support_a = jax.vmap(dot, in_axes=[None, 0])(axis, vertices_a)
    support_b = jax.vmap(dot, in_axes=[None, 0])(axis, vertices_b)
    dist = support_a.max() - support_b.min()
    separating = dist < 0
    dist = jp.where(dist < 0, 1e6, dist)
    return dist, separating

  support, separating = get_support(axes)
  is_face_separating = separating.any()

  # choose the best separating axis
  best_idx = jp.argmin(support)
  best_axis = axes[best_idx]

  # get the (reference) face most aligned with the separating axis
  dist_a = normals_a @ best_axis
  dist_b = normals_b @ -best_axis
  face_a_idx = dist_a.argmax()
  face_b_idx = dist_b.argmax()

  cond = best_idx < normals_a.shape[0]
  ref_face = jp.where(cond, faces_a[face_a_idx], faces_b[face_b_idx])
  incident_face = jp.where(cond, faces_b[face_b_idx], faces_a[face_a_idx])
  ref_face_norm = jp.where(cond, normals_a[face_a_idx], normals_b[face_b_idx])
  incident_face_norm = jp.where(
      cond, normals_b[face_b_idx], normals_a[face_a_idx]
  )

  dist, pos, normal = _create_contact_manifold(
      ref_face,
      incident_face,
      ref_face_norm,
      incident_face_norm,
      -best_axis,
  )

  # Handle edge separating axes by checking all edge pairs.
  a_idx = jp.tile(jp.arange(edges_a.shape[0]), reps=edges_b.shape[0])
  b_idx = jp.repeat(
      jp.arange(edges_b.shape[0]), repeats=edges_a.shape[0], axis=0
  )
  normal_a_1 = edge_face_normals_a[a_idx, 0]
  normal_a_2 = edge_face_normals_a[a_idx, 1]
  normal_b_1 = edge_face_normals_b[b_idx, 0]
  normal_b_2 = edge_face_normals_b[b_idx, 1]
  is_minkowski_face = jax.vmap(_arcs_intersect)(
      normal_a_1, normal_a_2, -normal_b_1, -normal_b_2
  )

  # get distances
  edge_a_dir = jax.vmap(math.normalize)(edges_a[:, 0] - edges_a[:, 1])[a_idx]
  edge_b_dir = jax.vmap(math.normalize)(edges_b[:, 0] - edges_b[:, 1])[b_idx]
  edges_a, edges_b = edges_a[a_idx], edges_b[b_idx]
  edge_a_pt, edge_a_pt_2 = edges_a[:, 0], edges_a[:, 1]
  edge_b_pt, edge_b_pt_2 = edges_b[:, 0], edges_b[:, 1]

  def get_normals(a_dir, a_pt, b_dir):
    edge_axis = jp.cross(a_dir, b_dir)
    degenerate_edge_axis = jp.sum(edge_axis**2) < 1e-6
    edge_axis = math.normalize(edge_axis)
    # correct normal to point from a to b, object b is at the origin
    sign = jp.where(jp.dot(edge_axis, a_pt - centroid_a) > 0.0, 1.0, -1.0)
    return edge_axis * sign, degenerate_edge_axis

  edge_axes, degenerate_edge_axes = jax.vmap(get_normals)(
      edge_a_dir, edge_a_pt, edge_b_dir)
  edge_dist = jax.vmap(jp.dot)(edge_axes, edge_b_pt - edge_a_pt)
  # handle degenerate axis
  edge_dist = jp.where(degenerate_edge_axes, -jp.inf, edge_dist)
  # ensure edges create minkowski face
  edge_dist = jp.where(is_minkowski_face, edge_dist, -jp.inf)

  best_edge_idx = edge_dist.argmax()
  best_edge_dist = edge_dist[best_edge_idx]
  is_edge_contact = jp.where(
      dist.max() < 0, best_edge_dist > dist.max() - 1e-6,
      (best_edge_dist < 0) & ~jp.isinf(best_edge_dist)
  )
  is_edge_contact = is_edge_contact & ~is_face_separating
  normal = jp.where(is_edge_contact, edge_axes[best_edge_idx], normal)
  dist = jp.where(
      is_edge_contact,
      jp.array([best_edge_dist, 1, 1, 1]),
      dist,
  )
  a_closest, b_closest = math.closest_segment_to_segment_points(
      edge_a_pt[best_edge_idx], edge_a_pt_2[best_edge_idx],
      edge_b_pt[best_edge_idx], edge_b_pt_2[best_edge_idx])
  pos = jp.where(
      is_edge_contact,
      jp.tile(0.5 * (a_closest + b_closest), (4, 1)), pos)

  return dist, pos, normal


@collider(ncon=4)
def convex_convex(c1: ConvexInfo, c2: ConvexInfo) -> Collision:
  """Calculates contacts between two convex objects."""
  # pad face vertices so that we can broadcast between geom1 and geom2
  # face has shape (n_face, n_vert, 3)
  nvert1, nvert2 = c1.face.shape[1], c2.face.shape[1]
  if nvert1 < nvert2:
    face = jp.pad(c1.face, ((0, 0), (0, nvert2 - nvert1), (0, 0)), 'edge')
    c1 = c1.replace(face=face)
  elif nvert2 < nvert1:
    face = jp.pad(c2.face, ((0, 0), (0, nvert1 - nvert2), (0, 0)), 'edge')
    c2 = c2.replace(face=face)

  # ensure that the first object has fewer verts
  swapped = c1.vert.shape[0] > c2.vert.shape[0]
  if swapped:
    c1, c2 = c2, c1

  faces1 = c1.face
  faces2 = c2.face

  to_local_pos = c2.mat.T @ (c1.pos - c2.pos)
  to_local_mat = c2.mat.T @ c1.mat

  faces1 = to_local_pos + faces1 @ to_local_mat.T
  normals1 = c1.face_normal @ to_local_mat.T
  normals2 = c2.face_normal

  vertices1 = to_local_pos + c1.vert @ to_local_mat.T
  vertices2 = c2.vert

  unique_edges1 = jp.take(vertices1, c1.edge_dir, axis=0)
  unique_edges2 = jp.take(vertices2, c2.edge_dir, axis=0)

  edges1 = jp.take(vertices1, c1.edge, axis=0)
  edges2 = jp.take(vertices2, c2.edge, axis=0)

  edge_face_normals1 = c1.edge_face_normal @ to_local_mat.T
  edge_face_normals2 = c2.edge_face_normal

  enable_bruteforce = (
      unique_edges1.shape[0] * unique_edges2.shape[0]
      < edges1[0].shape[0] * edges2[0].shape[0]
  )
  if enable_bruteforce:
    dist, pos, normal = _sat_bruteforce(
        faces1,
        faces2,
        vertices1,
        vertices2,
        normals1,
        normals2,
        unique_edges1,
        unique_edges2,
    )
  else:
    dist, pos, normal = _sat_gaussmap(
        to_local_pos,
        faces1,
        faces2,
        vertices1,
        vertices2,
        normals1,
        normals2,
        edges1,
        edges2,
        edge_face_normals1,
        edge_face_normals2,
    )

  # Go back to world frame.
  pos = c2.pos + pos @ c2.mat.T
  normal = normal @ c2.mat.T
  normal = -normal if swapped else normal
  frame = jax.vmap(math.make_frame)(normal)

  return dist, pos, frame