import numpy as np
import sophuspy as sp


def ExpSO3(w):
    d = np.linalg.norm(w)
    d2 = d*d
    W = np.array([[0.0, -w[2], w[1]],
                  [w[2], 0.0, -w[0]],
                  [-w[1], w[0], 0.0]])
    if d < 1e-5:
        return np.eye(3) + W + 0.5*W@W
    else:
        return np.eye(3) + W*np.sin(d)/d + W@W*(1.-np.cos(d))/d2


def LogSO3(R):
    tr = R.trace()
    w = np.array([(R[2, 1]-R[1, 2])/2, (R[0, 2]-R[2, 0])/2, (R[1, 0]-R[0, 1])/2])
    costheta = (tr-1.0)*0.5
    if costheta > 1 or costheta < -1:
        return w
    theta = np.arccos(costheta)
    s = np.sin(theta)
    if abs(s) < 1e-5:
        return w
    else:
        return theta*w/s


def rightJ(vec):
    d = np.linalg.norm(vec)
    d2 = d*d
    if d < 1e-5:
        return np.eye(3)

    W = sp.SO3.hat(vec)
    return np.eye(3) - W*(1.0-np.cos(d))/d2 + W*W*(d-np.sin(d))/(d2*d)


def invRightJ(vec):
    d = np.linalg.norm(vec)
    d2 = d*d
    if d < 1e-5:
        return np.eye(3)

    W = sp.SO3.hat(vec)
    return np.eye(3) + W/2 + W*W*(1.0/d2 - (1.0+np.cos(d))/(2.0*d*np.sin(d)))
