import numpy as np
import scipy.spatial
import scipy.linalg

def nullspace(A, atol=1e-13, rtol=0):
  u, s, vh = np.linalg.svd(A)
  tol = max(atol, rtol * s[0])
  nnz = (s >= tol).sum()
  ns = vh[nnz:].conj().T
  return ns

def nearest_orthogonal_matrix(R):
  U,S,Vt = np.linalg.svd(R)
  return U @ np.eye(3,dtype=R.dtype) @ Vt

def power_iters(A, n_iters=10):
  b = np.random.uniform(-1,1, size=(A.shape[0], A.shape[1], 1))
  for iter in range(n_iters):
    b = A @ b
    b = b / np.linalg.norm(b, axis=1, keepdims=True)
  return b

def rayleigh_quotient(A, b):
  return (b.transpose(0,2,1) @ A @ b) / (b.transpose(0,2,1) @ b)


def cross_prod_mat(x):
  x = x.reshape(-1,3)
  X = np.empty((x.shape[0],3,3), dtype=x.dtype)
  X[:,0,0] = 0
  X[:,0,1] = -x[:,2]
  X[:,0,2] = x[:,1]
  X[:,1,0] = x[:,2]
  X[:,1,1] = 0
  X[:,1,2] = -x[:,0]
  X[:,2,0] = -x[:,1]
  X[:,2,1] = x[:,0]
  X[:,2,2] = 0
  return X.squeeze()

def hat_operator(x):
  return cross_prod_mat(x)

def vee_operator(X):
  X = X.reshape(-1,3,3)
  x = np.empty((X.shape[0], 3), dtype=X.dtype)
  x[:,0] = X[:,2,1]
  x[:,1] = X[:,0,2]
  x[:,2] = X[:,1,0]
  return x.squeeze()


def rot_x(x, dtype=np.float32):
  x = np.array(x, copy=False)
  x = x.reshape(-1,1)
  R = np.zeros((x.shape[0],3,3), dtype=dtype)
  R[:,0,0] = 1
  R[:,1,1] = np.cos(x).ravel()
  R[:,1,2] = -np.sin(x).ravel()
  R[:,2,1] = np.sin(x).ravel()
  R[:,2,2] = np.cos(x).ravel()
  return R.squeeze()

def rot_y(y, dtype=np.float32):
  y = np.array(y, copy=False)
  y = y.reshape(-1,1)
  R = np.zeros((y.shape[0],3,3), dtype=dtype)
  R[:,0,0] = np.cos(y).ravel()
  R[:,0,2] = np.sin(y).ravel()
  R[:,1,1] = 1
  R[:,2,0] = -np.sin(y).ravel()
  R[:,2,2] = np.cos(y).ravel()
  return R.squeeze()

def rot_z(z, dtype=np.float32):
  z = np.array(z, copy=False)
  z = z.reshape(-1,1)
  R = np.zeros((z.shape[0],3,3), dtype=dtype)
  R[:,0,0] = np.cos(z).ravel()
  R[:,0,1] = -np.sin(z).ravel()
  R[:,1,0] = np.sin(z).ravel()
  R[:,1,1] = np.cos(z).ravel()
  R[:,2,2] = 1
  return R.squeeze()

def xyz_from_rotm(R):
  R = R.reshape(-1,3,3)
  xyz = np.empty((R.shape[0],3), dtype=R.dtype)
  for bidx in range(R.shape[0]):
    if R[bidx,0,2] < 1:
      if R[bidx,0,2] > -1:
        xyz[bidx,1] = np.arcsin(R[bidx,0,2])
        xyz[bidx,0] = np.arctan2(-R[bidx,1,2], R[bidx,2,2])
        xyz[bidx,2] = np.arctan2(-R[bidx,0,1], R[bidx,0,0])
      else:
        xyz[bidx,1] = -np.pi/2
        xyz[bidx,0] = -np.arctan2(R[bidx,1,0],R[bidx,1,1])
        xyz[bidx,2] = 0
    else:
      xyz[bidx,1] = np.pi/2
      xyz[bidx,0] = np.arctan2(R[bidx,1,0], R[bidx,1,1])
      xyz[bidx,2] = 0
  return xyz.squeeze()

def zyx_from_rotm(R):
  R = R.reshape(-1,3,3)
  zyx = np.empty((R.shape[0],3), dtype=R.dtype)
  for bidx in range(R.shape[0]):
    if R[bidx,2,0] < 1:
      if R[bidx,2,0] > -1:
        zyx[bidx,1] = np.arcsin(-R[bidx,2,0])
        zyx[bidx,0] = np.arctan2(R[bidx,1,0], R[bidx,0,0])
        zyx[bidx,2] = np.arctan2(R[bidx,2,1], R[bidx,2,2])
      else:
        zyx[bidx,1] = np.pi / 2
        zyx[bidx,0] = -np.arctan2(-R[bidx,1,2], R[bidx,1,1])
        zyx[bidx,2] = 0
    else:
      zyx[bidx,1] = -np.pi / 2
      zyx[bidx,0] = np.arctan2(-R[bidx,1,2], R[bidx,1,1])
      zyx[bidx,2] = 0
  return zyx.squeeze()

def rotm_from_xyz(xyz):
  xyz = np.array(xyz, copy=False).reshape(-1,3)
  return (rot_x(xyz[:,0]) @ rot_y(xyz[:,1]) @ rot_z(xyz[:,2])).squeeze()

def rotm_from_zyx(zyx):
  zyx = np.array(zyx, copy=False).reshape(-1,3)
  return (rot_z(zyx[:,0]) @ rot_y(zyx[:,1]) @ rot_x(zyx[:,2])).squeeze()

def rotm_from_quat(q):
  q = q.reshape(-1,4)
  w, x, y, z = q[:,0], q[:,1], q[:,2], q[:,3]
  R = np.array([
    [1 - 2*y*y - 2*z*z, 2*x*y - 2*z*w, 2*x*z + 2*y*w],
    [2*x*y + 2*z*w, 1 - 2*x*x - 2*z*z, 2*y*z - 2*x*w],
    [2*x*z - 2*y*w, 2*y*z + 2*x*w, 1 - 2*x*x - 2*y*y]
  ], dtype=q.dtype)
  R = R.transpose((2,0,1))
  return R.squeeze()

def rotm_from_axisangle(a):
  # exponential
  a = a.reshape(-1,3)
  phi = np.linalg.norm(a, axis=1).reshape(-1,1,1)
  iphi = np.zeros_like(phi)
  np.divide(1, phi, out=iphi, where=phi != 0)
  A = cross_prod_mat(a) * iphi
  R = np.eye(3, dtype=a.dtype) + np.sin(phi) * A + (1 - np.cos(phi)) * A @ A
  return R.squeeze()

def rotm_from_lookat(dir, up=None):
  dir = dir.reshape(-1,3)
  if up is None:
    up = np.zeros_like(dir)
    up[:,1] = 1
  dir /= np.linalg.norm(dir, axis=1, keepdims=True)
  up /= np.linalg.norm(up, axis=1, keepdims=True)
  x = dir[:,None,:] @ cross_prod_mat(up).transpose(0,2,1)
  y = x @ cross_prod_mat(dir).transpose(0,2,1)
  x = x.squeeze()
  y = y.squeeze()
  x /= np.linalg.norm(x, axis=1, keepdims=True)
  y /= np.linalg.norm(y, axis=1, keepdims=True)
  R = np.empty((dir.shape[0],3,3), dtype=dir.dtype)
  R[:,0,0] = x[:,0]
  R[:,0,1] = y[:,0]
  R[:,0,2] = dir[:,0]
  R[:,1,0] = x[:,1]
  R[:,1,1] = y[:,1]
  R[:,1,2] = dir[:,1]
  R[:,2,0] = x[:,2]
  R[:,2,1] = y[:,2]
  R[:,2,2] = dir[:,2]
  return R.transpose(0,2,1).squeeze()

def rotm_distance_identity(R0, R1):
  # https://link.springer.com/article/10.1007%2Fs10851-009-0161-2
  # in [0, 2*sqrt(2)]
  R0 = R0.reshape(-1,3,3)
  R1 = R1.reshape(-1,3,3)
  dists = np.linalg.norm(np.eye(3,dtype=R0.dtype) - R0 @ R1.transpose(0,2,1), axis=(1,2))
  return dists.squeeze()

def rotm_distance_geodesic(R0, R1):
  # https://link.springer.com/article/10.1007%2Fs10851-009-0161-2
  # in [0, pi)
  R0 = R0.reshape(-1,3,3)
  R1 = R1.reshape(-1,3,3)
  RtR = R0 @ R1.transpose(0,2,1)
  aa = axisangle_from_rotm(RtR)
  S = cross_prod_mat(aa).reshape(-1,3,3)
  dists = np.linalg.norm(S, axis=(1,2))
  return dists.squeeze()



def axisangle_from_rotm(R):
  # logarithm of rotation matrix
  # R = R.reshape(-1,3,3)
  # tr = np.trace(R, axis1=1, axis2=2)
  # phi = np.arccos(np.clip((tr - 1) / 2, -1, 1))
  # scale = np.zeros_like(phi)
  # div = 2 * np.sin(phi)
  # np.divide(phi, div, out=scale, where=np.abs(div) > 1e-6)
  # A = (R - R.transpose(0,2,1)) * scale.reshape(-1,1,1)
  # aa = np.stack((A[:,2,1], A[:,0,2], A[:,1,0]), axis=1)
  # return aa.squeeze()
  R = R.reshape(-1,3,3)
  omega = np.empty((R.shape[0], 3), dtype=R.dtype)
  omega[:,0] = R[:,2,1] - R[:,1,2]
  omega[:,1] = R[:,0,2] - R[:,2,0]
  omega[:,2] = R[:,1,0] - R[:,0,1]
  r = np.linalg.norm(omega, axis=1).reshape(-1,1)
  t = np.trace(R, axis1=1, axis2=2).reshape(-1,1)
  omega = np.arctan2(r, t-1) * omega
  aa = np.zeros_like(omega)
  np.divide(omega, r, out=aa, where=r != 0)
  return aa.squeeze()

def axisangle_from_quat(q):
  q = q.reshape(-1,4)
  phi = 2 * np.arccos(q[:,0])
  denom = np.zeros_like(q[:,0])
  np.divide(1, np.sqrt(1 - q[:,0]**2), out=denom, where=q[:,0] != 1)
  axis = q[:,1:] * denom.reshape(-1,1)
  denom = np.linalg.norm(axis, axis=1).reshape(-1,1)
  a = np.zeros_like(axis)
  np.divide(phi.reshape(-1,1) * axis, denom, out=a, where=denom != 0)
  aa = a.astype(q.dtype)
  return aa.squeeze()

def axisangle_apply(aa, x):
  # working only with single aa and single x at the moment
  xshape = x.shape
  aa = aa.reshape(3,)
  x = x.reshape(3,)
  phi = np.linalg.norm(aa)
  e = np.zeros_like(aa)
  np.divide(aa, phi, out=e, where=phi != 0)
  xr = np.cos(phi) * x + np.sin(phi) * np.cross(e, x) + (1 - np.cos(phi)) * (e.T @ x) * e
  return xr.reshape(xshape)


def exp_so3(R):
  w = axisangle_from_rotm(R)
  return w

def log_so3(w):
  R = rotm_from_axisangle(w)
  return R

def exp_se3(R, t):
  R = R.reshape(-1,3,3)
  t = t.reshape(-1,3)

  w = exp_so3(R).reshape(-1,3)

  phi = np.linalg.norm(w, axis=1).reshape(-1,1,1)
  A = cross_prod_mat(w)
  Vi = np.eye(3, dtype=R.dtype) - A/2 + (1 - (phi * np.sin(phi) / (2 * (1 - np.cos(phi))))) / phi**2 * A @ A
  u = t.reshape(-1,1,3) @ Vi.transpose(0,2,1)

  # v = (u, w)
  v = np.empty((R.shape[0],6), dtype=R.dtype)
  v[:,:3] = u.squeeze()
  v[:,3:] = w

  return v.squeeze()

def log_se3(v):
  # v = (u, w)
  v = v.reshape(-1,6)
  u = v[:,:3]
  w = v[:,3:]

  R = log_so3(w)

  phi = np.linalg.norm(w, axis=1).reshape(-1,1,1)
  A = cross_prod_mat(w)
  V = np.eye(3, dtype=v.dtype) + (1 - np.cos(phi)) / phi**2 * A + (phi - np.sin(phi)) / phi**3 * A @ A
  t = u.reshape(-1,1,3) @ V.transpose(0,2,1)

  return R.squeeze(), t.squeeze()


def quat_from_rotm(R):
  R = R.reshape(-1,3,3)
  q = np.empty((R.shape[0], 4,), dtype=R.dtype)
  q[:,0] = np.sqrt( np.maximum(0, 1 + R[:,0,0] + R[:,1,1] + R[:,2,2]) )
  q[:,1] = np.sqrt( np.maximum(0, 1 + R[:,0,0] - R[:,1,1] - R[:,2,2]) )
  q[:,2] = np.sqrt( np.maximum(0, 1 - R[:,0,0] + R[:,1,1] - R[:,2,2]) )
  q[:,3] = np.sqrt( np.maximum(0, 1 - R[:,0,0] - R[:,1,1] + R[:,2,2]) )
  q[:,1] *= np.sign(q[:,1] * (R[:,2,1] - R[:,1,2]))
  q[:,2] *= np.sign(q[:,2] * (R[:,0,2] - R[:,2,0]))
  q[:,3] *= np.sign(q[:,3] * (R[:,1,0] - R[:,0,1]))
  q /= np.linalg.norm(q,axis=1,keepdims=True)
  return q.squeeze()

def quat_from_axisangle(a):
  a = a.reshape(-1, 3)
  phi = np.linalg.norm(a, axis=1)
  iphi = np.zeros_like(phi)
  np.divide(1, phi, out=iphi, where=phi != 0)
  a = a * iphi.reshape(-1,1)
  theta = phi / 2.0
  r = np.cos(theta)
  stheta = np.sin(theta)
  q = np.stack((r, stheta*a[:,0], stheta*a[:,1], stheta*a[:,2]), axis=1)
  q /= np.linalg.norm(q, axis=1).reshape(-1,1)
  return q.squeeze()

def quat_identity(n=1, dtype=np.float32):
  q = np.zeros((n,4), dtype=dtype)
  q[:,0] = 1
  return q.squeeze()

def quat_conjugate(q):
  shape = q.shape
  q = q.reshape(-1,4).copy()
  q[:,1:] *= -1
  return q.reshape(shape)

def quat_product(q1, q2):
  # q1 . q2 is equivalent to R(q1) @ R(q2)
  shape = q1.shape
  q1, q2 = q1.reshape(-1,4), q2.reshape(-1, 4)
  q = np.empty((max(q1.shape[0], q2.shape[0]), 4), dtype=q1.dtype)
  a1,b1,c1,d1 = q1[:,0], q1[:,1], q1[:,2], q1[:,3]
  a2,b2,c2,d2 = q2[:,0], q2[:,1], q2[:,2], q2[:,3]
  q[:,0] = a1 * a2 - b1 * b2 - c1 * c2 - d1 * d2
  q[:,1] = a1 * b2 + b1 * a2 + c1 * d2 - d1 * c2
  q[:,2] = a1 * c2 - b1 * d2 + c1 * a2 + d1 * b2
  q[:,3] = a1 * d2 + b1 * c2 - c1 * b2 + d1 * a2
  return q.squeeze()

def quat_apply(q, x):
  xshape = x.shape
  x = x.reshape(-1, 3)
  qshape = q.shape
  q = q.reshape(-1, 4)

  p = np.empty((x.shape[0], 4), dtype=x.dtype)
  p[:,0] = 0
  p[:,1:] = x

  r = quat_product(quat_product(q, p), quat_conjugate(q))
  if r.ndim == 1:
    return r[1:].reshape(xshape)
  else:
    return r[:,1:].reshape(xshape)


def quat_random(rng=None, n=1):
  # http://planning.cs.uiuc.edu/node198.html
  if rng is not None:
    u = rng.uniform(0, 1, size=(3,n))
  else:
    u = np.random.uniform(0, 1, size=(3,n))
  q = np.array((
    np.sqrt(1 - u[0]) * np.sin(2 * np.pi * u[1]),
    np.sqrt(1 - u[0]) * np.cos(2 * np.pi * u[1]),
    np.sqrt(u[0]) * np.sin(2 * np.pi * u[2]),
    np.sqrt(u[0]) * np.cos(2 * np.pi * u[2])
  )).T
  q /= np.linalg.norm(q,axis=1,keepdims=True)
  return q.squeeze()

def quat_distance_angle(q0, q1):
  # https://math.stackexchange.com/questions/90081/quaternion-distance
  # https://link.springer.com/article/10.1007%2Fs10851-009-0161-2
  q0 = q0.reshape(-1,4)
  q1 = q1.reshape(-1,4)
  dists = np.arccos(np.clip(2 * np.sum(q0 * q1, axis=1)**2 - 1, -1, 1))
  return dists

def quat_distance_normdiff(q0, q1):
  # https://link.springer.com/article/10.1007%2Fs10851-009-0161-2
  # \phi_4
  # [0, 1]
  q0 = q0.reshape(-1,4)
  q1 = q1.reshape(-1,4)
  return 1 - np.sum(q0 * q1, axis=1)**2

def quat_distance_mineucl(q0, q1):
  # https://link.springer.com/article/10.1007%2Fs10851-009-0161-2
  # http://users.cecs.anu.edu.au/~trumpf/pubs/Hartley_Trumpf_Dai_Li.pdf
  q0 = q0.reshape(-1,4)
  q1 = q1.reshape(-1,4)
  diff0 = ((q0 - q1)**2).sum(axis=1)
  diff1 = ((q0 + q1)**2).sum(axis=1)
  return np.minimum(diff0, diff1)

def quat_slerp_space(q0, q1, num=100, endpoint=True):
  q0 = q0.ravel()
  q1 = q1.ravel()
  dot = q0.dot(q1)
  if dot < 0:
    q1 *= -1
    dot *= -1
  t = np.linspace(0, 1, num=num, endpoint=endpoint, dtype=q0.dtype)
  t = t.reshape((-1,1))
  if dot > 0.9995:
    ret = q0 + t * (q1 - q0)
    return ret
  dot = np.clip(dot, -1, 1)
  theta0 = np.arccos(dot)
  theta = theta0 * t
  s0 = np.cos(theta) - dot * np.sin(theta) / np.sin(theta0)
  s1 = np.sin(theta) / np.sin(theta0)
  return (s0 * q0) + (s1 * q1)

def cart_to_spherical(x):
  shape = x.shape
  x = x.reshape(-1,3)
  y = np.empty_like(x)
  y[:,0] = np.linalg.norm(x, axis=1)  # r
  y[:,1] = np.arccos(x[:,2] / y[:,0]) # theta
  y[:,2] = np.arctan2(x[:,1], x[:,0]) # phi
  return y.reshape(shape)

def spherical_to_cart(x):
  shape = x.shape
  x = x.reshape(-1,3)
  y = np.empty_like(x)
  y[:,0] = x[:,0] * np.sin(x[:,1]) * np.cos(x[:,2])
  y[:,1] = x[:,0] * np.sin(x[:,1]) * np.sin(x[:,2])
  y[:,2] = x[:,0] * np.cos(x[:,1])
  return y.reshape(shape)

def spherical_random(r=1, n=1):
  # http://mathworld.wolfram.com/SpherePointPicking.html
  # https://math.stackexchange.com/questions/1585975/how-to-generate-random-points-on-a-sphere
  x = np.empty((n,3))
  x[:,0] = r
  x[:,1] = 2 * np.pi * np.random.uniform(0,1, size=(n,))
  x[:,2] = np.arccos(2 * np.random.uniform(0,1, size=(n,)) - 1)
  return x.squeeze()

def color_pcl(pcl, K, im, color_axis=0, as_int=True, invalid_color=[0,0,0]):
  uvd = K @ pcl.T
  uvd /= uvd[2]
  uvd = np.round(uvd).astype(np.int32)
  mask = np.logical_and(uvd[0] >= 0, uvd[1] >= 0)
  color = np.empty((pcl.shape[0], 3), dtype=im.dtype)
  if color_axis == 0:
    mask = np.logical_and(mask, uvd[0] < im.shape[2])
    mask = np.logical_and(mask, uvd[1] < im.shape[1])
    uvd = uvd[:,mask]
    color[mask,:] = im[:,uvd[1],uvd[0]].T
  elif color_axis == 2:
    mask = np.logical_and(mask, uvd[0] < im.shape[1])
    mask = np.logical_and(mask, uvd[1] < im.shape[0])
    uvd = uvd[:,mask]
    color[mask,:] = im[uvd[1],uvd[0], :]
  else:
    raise Exception('invalid color_axis')
  color[np.logical_not(mask),:3] = invalid_color
  if as_int:
    color = (255.0 * color).astype(np.int32)
  return color

def center_pcl(pcl, robust=False, copy=False, axis=1):
  if copy:
    pcl = pcl.copy()
  if robust:
    mu = np.median(pcl, axis=axis, keepdims=True)
  else:
    mu = np.mean(pcl, axis=axis, keepdims=True)
  return pcl - mu

def to_homogeneous(x):
  # return np.hstack((x, np.ones((x.shape[0],1),dtype=x.dtype)))
  return np.concatenate((x, np.ones((*x.shape[:-1],1),dtype=x.dtype)), axis=-1)

def from_homogeneous(x):
  return x[:,:-1] / x[:,-1]

def project_uvn(uv, Ki=None):
  if uv.shape[1] == 2:
    uvn = to_homogeneous(uv)
  else:
    uvn = uv
  if uvn.shape[1] != 3:
    raise Exception('uv should have shape Nx2 or Nx3')
  if Ki is None:
    return uvn
  else:
    return uvn @ Ki.T

def project_uvd(uv, depth, K=np.eye(3), R=np.eye(3), t=np.zeros((3,1)), ignore_negative_depth=True, return_uvn=False):
  Ki = np.linalg.inv(K)

  if ignore_negative_depth:
    mask = depth >= 0
    uv = uv[mask,:]
    d = depth[mask]
  else:
    d = depth.ravel()

  uv1 = to_homogeneous(uv)

  uvn1 = uv1 @ Ki.T
  xyz = d.reshape(-1,1) * uvn1
  xyz = (xyz - t.reshape((1,3))) @ R

  if return_uvn:
    return xyz, uvn1
  else:
    return xyz

def project_depth(depth, K, R=np.eye(3,3), t=np.zeros((3,1)), ignore_negative_depth=True, return_uvn=False):
  u, v = np.meshgrid(range(depth.shape[1]), range(depth.shape[0]))
  uv = np.hstack((u.reshape(-1,1), v.reshape(-1,1)))
  return project_uvd(uv, depth.ravel(), K, R, t, ignore_negative_depth, return_uvn)


def project_xyz(xyz, K=np.eye(3), R=np.eye(3,3), t=np.zeros((3,1))):
  uvd = K @ (R @ xyz.T + t.reshape((3,1)))
  uvd[:2] /= uvd[2]
  return uvd[:2].T, uvd[2]


def relative_motion(R0, t0, R1, t1, Rt_from_global=True):
  t0 = t0.reshape((3,1))
  t1 = t1.reshape((3,1))
  if Rt_from_global:
    Rr = R1 @ R0.T
    tr = t1 - Rr @ t0
  else:
    Rr = R1.T @ R0
    tr = R1.T @ (t0 - t1)
  return Rr, tr.ravel()


def translation_to_cameracenter(R, t):
  t = t.reshape(-1,3,1)
  R = R.reshape(-1,3,3)
  C = -R.transpose(0,2,1) @ t
  return C.squeeze()

def cameracenter_to_translation(R, C):
  C = C.reshape(-1,3,1)
  R = R.reshape(-1,3,3)
  t = -R @ C
  return t.squeeze()

def decompose_projection_matrix(P, return_t=True):
  if P.shape[0] != 3 or P.shape[1] != 4:
    raise Exception('P has to be 3x4')
  M = P[:, :3]
  C = -np.linalg.inv(M) @ P[:, 3:]

  R,K = np.linalg.qr(np.flipud(M).T)
  K = np.flipud(K.T)
  K = np.fliplr(K)
  R = np.flipud(R.T)

  T = np.diag(np.sign(np.diag(K)))
  K = K @ T
  R = T @ R

  if np.linalg.det(R) < 0:
    R *= -1

  K /= K[2,2]
  if return_t:
    return K, R, cameracenter_to_translation(R, C)
  else:
    return K, R, C


def compose_projection_matrix(K=np.eye(3), R=np.eye(3,3), t=np.zeros((3,1))):
  return K @ np.hstack((R, t.reshape((3,1))))



def point_plane_distance(pts, plane):
  pts = pts.reshape(-1,3)
  return np.abs(np.sum(plane[:3] * pts, axis=1) + plane[3]) / np.linalg.norm(plane[:3])

def fit_plane(pts):
  pts = pts.reshape(-1,3)
  center = np.mean(pts, axis=0)
  A = pts - center
  u, s, vh = np.linalg.svd(A, full_matrices=False)
  # if pts.shape[0] > 100:
  #   import ipdb; ipdb.set_trace()
  plane = np.array([*vh[2], -vh[2].dot(center)])
  return plane

def tetrahedron(dtype=np.float32):
  verts = np.array([
    (np.sqrt(8/9), 0, -1/3), (-np.sqrt(2/9), np.sqrt(2/3), -1/3),
    (-np.sqrt(2/9), -np.sqrt(2/3), -1/3), (0, 0, 1)], dtype=dtype)
  faces = np.array([(0,1,2), (0,2,3), (0,1,3), (1,2,3)], dtype=np.int32)
  normals = -np.mean(verts, axis=0) + verts
  normals /= np.linalg.norm(normals, axis=1).reshape(-1,1)
  return verts, faces, normals

def cube(dtype=np.float32):
  verts = np.array([
    [-0.5,-0.5,-0.5], [-0.5,0.5,-0.5], [0.5,0.5,-0.5], [0.5,-0.5,-0.5],
    [-0.5,-0.5,0.5], [-0.5,0.5,0.5], [0.5,0.5,0.5], [0.5,-0.5,0.5]], dtype=dtype)
  faces = np.array([
    (0,1,2), (0,2,3), (4,5,6), (4,6,7),
    (0,4,7), (0,7,3), (1,5,6), (1,6,2),
    (3,2,6), (3,6,7), (0,1,5), (0,5,4)], dtype=np.int32)
  normals = -np.mean(verts, axis=0) + verts
  normals /= np.linalg.norm(normals, axis=1).reshape(-1,1)
  return verts, faces, normals

def octahedron(dtype=np.float32):
  verts = np.array([
    (+1,0,0), (0,+1,0), (0,0,+1),
    (-1,0,0), (0,-1,0), (0,0,-1)], dtype=dtype)
  faces = np.array([
    (0,1,2), (1,2,3), (3,2,4), (4,2,0),
    (0,1,5), (1,5,3), (3,5,4), (4,5,0)], dtype=np.int32)
  normals = -np.mean(verts, axis=0) + verts
  normals /= np.linalg.norm(normals, axis=1).reshape(-1,1)
  return verts, faces, normals

def icosahedron(dtype=np.float32):
  p = (1 + np.sqrt(5)) / 2
  verts = np.array([
    (-1,0,p), (1,0,p), (1,0,-p), (-1,0,-p),
    (0,-p,1), (0,p,1), (0,p,-1), (0,-p,-1),
    (-p,-1,0), (p,-1,0), (p,1,0), (-p,1,0)
    ], dtype=dtype)
  faces = np.array([
    (0,1,4), (0,1,5), (1,4,9), (1,9,10), (1,10,5), (0,4,8), (0,8,11), (0,11,5),
    (5,6,11), (5,6,10), (4,7,8), (4,7,9),
    (3,2,6), (3,2,7), (2,6,10), (2,10,9), (2,9,7), (3,6,11), (3,11,8), (3,8,7),
    ], dtype=np.int32)
  normals = -np.mean(verts, axis=0) + verts
  normals /= np.linalg.norm(normals, axis=1).reshape(-1,1)
  return verts, faces, normals

def xyplane(dtype=np.float32, z=0, interleaved=False):
  if interleaved:
    eps = 1e-6
    verts = np.array([
      (-1,-1,z), (-1,1,z), (1,1,z),
      (1-eps,1,z), (1-eps,-1,z), (-1-eps,-1,z)], dtype=dtype)
    faces = np.array([(0,1,2), (3,4,5)], dtype=np.int32)
  else:
    verts = np.array([(-1,-1,z), (-1,1,z), (1,1,z), (1,-1,z)], dtype=dtype)
    faces = np.array([(0,1,2), (0,2,3)], dtype=np.int32)
  normals = np.zeros_like(verts)
  normals[:,2] = -1
  return verts, faces, normals

def mesh_independent_verts(verts, faces, normals=None):
  new_verts = []
  new_normals = []
  for f in faces:
    new_verts.append(verts[f[0]])
    new_verts.append(verts[f[1]])
    new_verts.append(verts[f[2]])
    if normals is not None:
      new_normals.append(normals[f[0]])
      new_normals.append(normals[f[1]])
      new_normals.append(normals[f[2]])
  new_verts = np.array(new_verts)
  new_faces = np.arange(0, faces.size, dtype=faces.dtype).reshape(-1,3)
  if normals is None:
    return new_verts, new_faces
  else:
    new_normals = np.array(new_normals)
    return new_verts, new_faces, new_normals


def stack_mesh(verts, faces):
  n_verts = 0
  mfaces = []
  for idx, f in enumerate(faces):
    mfaces.append(f + n_verts)
    n_verts += verts[idx].shape[0]
  verts = np.vstack(verts)
  faces = np.vstack(mfaces)
  return verts, faces

def normalize_mesh(verts):
  # all the verts have unit distance to the center (0,0,0)
  return verts / np.linalg.norm(verts, axis=1, keepdims=True)


def mesh_triangle_areas(verts, faces):
  a = verts[faces[:,0]]
  b = verts[faces[:,1]]
  c = verts[faces[:,2]]
  x = np.empty_like(a)
  x = a - b
  y = a - c
  t = np.empty_like(a)
  t[:,0] = (x[:,1] * y[:,2] - x[:,2] * y[:,1]);
  t[:,1] = (x[:,2] * y[:,0] - x[:,0] * y[:,2]);
  t[:,2] = (x[:,0] * y[:,1] - x[:,1] * y[:,0]);
  return np.linalg.norm(t, axis=1) / 2

def subdivde_mesh(verts_in, faces_in, n=1):
  for iter in range(n):
    verts = []
    for v in verts_in:
      verts.append(v)
    faces = []
    verts_dict = {}
    for f in faces_in:
      f = np.sort(f)
      i0,i1,i2 = f
      v0,v1,v2 = verts_in[f]

      k = i0*len(verts_in)+i1
      if k in verts_dict:
        i01 = verts_dict[k]
      else:
        i01 = len(verts)
        verts_dict[k] = i01
        v01 = (v0 + v1) / 2
        verts.append(v01)

      k = i0*len(verts_in)+i2
      if k in verts_dict:
        i02 = verts_dict[k]
      else:
        i02 = len(verts)
        verts_dict[k] = i02
        v02 = (v0 + v2) / 2
        verts.append(v02)

      k = i1*len(verts_in)+i2
      if k in verts_dict:
        i12 = verts_dict[k]
      else:
        i12 = len(verts)
        verts_dict[k] = i12
        v12 = (v1 + v2) / 2
        verts.append(v12)

      faces.append((i0,i01,i02))
      faces.append((i01,i1,i12))
      faces.append((i12,i2,i02))
      faces.append((i01,i12,i02))

    verts_in = np.array(verts, dtype=verts_in.dtype)
    faces_in = np.array(faces, dtype=np.int32)
  return verts_in, faces_in


def mesh_adjust_winding_order(verts, faces, normals):
  n0 = normals[faces[:,0]]
  n1 = normals[faces[:,1]]
  n2 = normals[faces[:,2]]
  fnormals = (n0 + n1 + n2) / 3

  v0 = verts[faces[:,0]]
  v1 = verts[faces[:,1]]
  v2 = verts[faces[:,2]]

  e0 = v1 - v0
  e1 = v2 - v0
  fn = np.cross(e0, e1)

  dot = np.sum(fnormals * fn, axis=1)
  ma = dot < 0

  nfaces = faces.copy()
  nfaces[ma,1], nfaces[ma,2] = nfaces[ma,2], nfaces[ma,1]

  return nfaces


def pcl_to_shapecl(verts, colors=None, shape='cube', width=1.0):
  if shape == 'tetrahedron':
    cverts, cfaces, _ = tetrahedron()
  elif shape == 'cube':
    cverts, cfaces, _ = cube()
  elif shape == 'octahedron':
    cverts, cfaces, _ = octahedron()
  elif shape == 'icosahedron':
    cverts, cfaces, _ = icosahedron()
  else:
    raise Exception('invalid shape')

  sverts = np.tile(cverts, (verts.shape[0], 1))
  sverts *= width
  sverts += np.repeat(verts, cverts.shape[0], axis=0)

  sfaces = np.tile(cfaces, (verts.shape[0], 1))
  sfoffset = cverts.shape[0] * np.arange(0, verts.shape[0])
  sfaces += np.repeat(sfoffset, cfaces.shape[0]).reshape(-1,1)

  if colors is not None:
    scolors = np.repeat(colors, cverts.shape[0], axis=0)
  else:
    scolors = None

  return sverts, sfaces, scolors