import numpy as np
import chumpy as ch

try:
    from opendr.geometry import Rodrigues  # This lets us import blmath.geometry.normals.NormalizedNx3 etc. whether or not opendr is available. pylint: disable=unused-import
except ImportError:
    class Rodrigues(ch.Ch):
        dterms = 'rt'

        def compute_r(self):
            import cv2
            return cv2.Rodrigues(self.rt.r)[0]

        def compute_dr_wrt(self, wrt):
            import cv2
            if wrt is self.rt:
                return cv2.Rodrigues(self.rt.r)[1].T


def rodrigues(r, calculate_jacobian=False):
    r = np.array(r, dtype=np.double)
    eps = np.finfo(np.double).eps

    if np.all(r.shape == (3, 1)) or np.all(r.shape == (1, 3)) or np.all(r.shape == (3,)):
        r = r.flatten()
        theta = np.linalg.norm(r)
        if theta < eps:
            r_out = np.eye(3)
            if calculate_jacobian:
                jac = np.zeros((3, 9))
                jac[0, 5] = jac[1, 6] = jac[2, 1] = -1
                jac[0, 7] = jac[1, 2] = jac[2, 3] = 1

        else:
            c = np.cos(theta)
            s = np.sin(theta)
            c1 = 1. - c
            itheta = 1.0 if theta == 0.0 else 1.0 / theta
            r *= itheta
            I = np.eye(3)
            rrt = np.array([r * r[0], r * r[1], r * r[2]])
            _r_x_ = np.array([
                [0, -r[2], r[1]],
                [r[2], 0, -r[0]],
                [-r[1], r[0], 0],
            ])
            r_out = c*I + c1*rrt + s*_r_x_
            if calculate_jacobian:
                drrt = np.array([
                    [r[0]+r[0], r[1], r[2], r[1], 0, 0, r[2], 0, 0],
                    [0, r[0], 0, r[0], r[1]+r[1], r[2], 0, r[2], 0],
                    [0, 0, r[0], 0, 0, r[1], r[0], r[1], r[2]+r[2]],
                ])
                d_r_x_ = np.array([
                    [0, 0, 0, 0, 0, -1, 0, 1, 0],
                    [0, 0, 1, 0, 0, 0, -1, 0, 0],
                    [0, -1, 0, 1, 0, 0, 0, 0, 0],
                ])
                I = np.array([I.flatten(), I.flatten(), I.flatten()])
                ri = np.array([[r[0]], [r[1]], [r[2]]])
                a0 = -s*ri
                a1 = (s - 2*c1*itheta)*ri
                a2 = np.ones((3, 1))*c1*itheta
                a3 = (c - s*itheta)*ri
                a4 = np.ones((3, 1))*s*itheta
                jac = a0*I + a1*rrt.flatten() + a2*drrt + a3*_r_x_.flatten() + a4*d_r_x_
    elif np.all(r.shape == (3, 3)):
        u, _, v = np.linalg.svd(r)
        r = np.dot(u, v)
        rx = r[2, 1] - r[1, 2]
        ry = r[0, 2] - r[2, 0]
        rz = r[1, 0] - r[0, 1]
        s = np.linalg.norm(np.array([rx, ry, rz])) * np.sqrt(0.25)
        c = np.clip((np.sum(np.diag(r)) - 1) * 0.5, -1, 1)
        theta = np.arccos(c)
        if s < 1e-5:
            if c > 0:
                r_out = np.zeros((3, 1))
            else:
                rx, ry, rz = np.clip(np.sqrt((np.diag(r) + 1) * 0.5), 0, np.inf)
                if r[0, 1] < 0:
                    ry = -ry
                if r[0, 2] < 0:
                    rz = -rz
                if np.abs(rx) < np.abs(ry) and np.abs(rx) < np.abs(rz) and ((r[1, 2] > 0) != (ry*rz > 0)):
                    rz = -rz

                r_out = np.array([[rx, ry, rz]]).T
                theta /= np.linalg.norm(r_out)
                r_out *= theta
            if calculate_jacobian:
                jac = np.zeros((9, 3))
                if c > 0:
                    jac[1, 2] = jac[5, 0] = jac[6, 1] = -0.5
                    jac[2, 1] = jac[3, 2] = jac[7, 0] = 0.5
        else:
            vth = 1.0 / (2.0 * s)
            if calculate_jacobian:
                dtheta_dtr = -1./s
                dvth_dtheta = -vth*c/s
                d1 = 0.5*dvth_dtheta*dtheta_dtr
                d2 = 0.5*dtheta_dtr
                dvardR = np.array([
                    [0, 0, 0, 0, 0, 1, 0, -1, 0],
                    [0, 0, -1, 0, 0, 0, 1, 0, 0],
                    [0, 1, 0, -1, 0, 0, 0, 0, 0],
                    [d1, 0, 0, 0, d1, 0, 0, 0, d1],
                    [d2, 0, 0, 0, d2, 0, 0, 0, d2],
                ])
                dvar2dvar = np.array([
                    [vth, 0, 0, rx, 0],
                    [0, vth, 0, ry, 0],
                    [0, 0, vth, rz, 0],
                    [0, 0, 0, 0, 1],
                ])
                domegadvar2 = np.array([
                    [theta, 0, 0, rx*vth],
                    [0, theta, 0, ry*vth],
                    [0, 0, theta, rz*vth],
                ])
                jac = np.dot(np.dot(domegadvar2, dvar2dvar), dvardR)
                for ii in range(3):
                    jac[ii] = jac[ii].reshape((3, 3)).T.flatten()
                jac = jac.T
            vth *= theta
            r_out = np.array([[rx, ry, rz]]).T * vth
    else:
        raise Exception("rodrigues: input matrix must be 1x3, 3x1 or 3x3.")
    if calculate_jacobian:
        return r_out, jac
    else:
        return r_out


def as_rotation_matrix(r):
    if np.all(r.shape == (3, 1)) or np.all(r.shape == (1, 3)) or np.all(r.shape == (3,)):
        return rodrigues(r)
    elif np.all(r.shape == (3, 3)):
        return rodrigues(rodrigues(r))
    else:
        raise Exception("rodrigues: input matrix must be 1x3, 3x1 or 3x3.")

numpy.array numpy.linalg.norm numpy.sin numpy.dot numpy.arccos numpy.zeros numpy.ones numpy.eye numpy.finfo numpy.cos numpy.linalg.svd numpy.all numpy.sqrt numpy.abs numpy.diag