#!/usr/bin/env python3

import sys
import numpy as np
import numpysane as nps
import os

testdir = os.path.dirname(os.path.realpath(__file__))

# I import the LOCAL mrcal since that's what I'm testing
sys.path[:0] = (f"{testdir}/..",)
import mrcal

from testutils import *
from test_calibration_helpers import grad, grad__r_from_R

from test_poseutils_helpers import R_from_r, compose_r


def wrap_r_unconditional(r, dr_dX=None):
    """Unwrap a Rodrigues vector r

    returns r'     if dr_dX is None
    returns dr'_dX if dr_dX is not None

    I define rotations with a Rodrigues vector r = th v. Where v is the unit
    vector representing the rotation axis; and th is a scalar representing the
    magnitude or this rotation, in radians. Thus rotations th v and (th + 2pi*n)
    v are identical for all integer n. Furthermore the rotation (pi+th) v is
    equivalent to a rotation (pi-th) (-v). So usually we have th = mag(r) in
    [0,pi].

    For various analyses I want to convert a rotation r to its equivalent
    rotation 2*pi rad away:

      r' = v (th - 2pi)
         = r/magr (magr - 2pi)
         = r - r/magr 2pi
         = r (1 - 2pi/magr)

    Let

      k = 1 - 2pi/magr

    So

      r' = r k

    As noted above, usually magr <= pi, so k < 0. Thus

      magr' = -magr k

    Let's apply this wrapping a second time:

      r'' = r' (1 - 2pi/magr')
          = r k (1 + 2pi/magr/k)
          = r k (1 + (1-k) /k)
          = r (k + 1 - k)
          = r

    So this wrapping operation switches back/forth between two modes

    """

    if dr_dX is None:
        if nps.mag(r) == 0:
            return r
        return r - r / nps.mag(r) * 2.0 * np.pi

    if nps.mag(r) == 0:
        return dr_dX

    # drw_dX = d( r - r/nps.mag(r) * 2.*np.pi )/dX
    #        = dr_dX - 2pi (dr_dX/nps.mag(r) + r d/dx (1/magr))
    #        = dr_dX - 2pi (dr_dX/nps.mag(r) - r /magr^2 /2magr 2rt dr_dX )
    #        = dr_dX - 2pi (dr_dX - r rt /norm2(r) dr_dX ) /nps.mag(r)
    #        = dr_dX + 2pi /nps.mag(r) (r rt /norm2(r) - I) dr_dX

    # dr_dX has shape (3,...) so r rt dr_dX also has shape (3,...) I clump
    # the trailing shapes so that dr_dX.shape is (3,N), and then reshape it
    # at the end
    s = dr_dX.shape[1:]
    # shape (3,N)
    dr_dX = nps.clump(dr_dX, n=-(dr_dX.ndim - 1))

    drw_dX = dr_dX + 2.0 * np.pi / nps.mag(r) * nps.matmult(
        nps.outer(r, r) / nps.norm2(r) - np.eye(r.size), dr_dX
    )
    return drw_dX.reshape((3,) + s)


def wrap_r(r, *, r_match_direction=None, dr_dX=None):
    """Unwrap a Rodrigues vector r

    returns r'     if dr_dX is None
    returns dr'_dX if dr_dX is not None

    This function only wraps the argument if it needs to: if mag(r) > np.pi or
    if we're pointing opposite r_match_direction

    """

    if r_match_direction is not None:
        if nps.inner(r, r_match_direction) > 0:
            return r if dr_dX is None else dr_dX
        return wrap_r_unconditional(r, dr_dX=dr_dX)
    if nps.mag(r) <= np.pi:
        return r if dr_dX is None else dr_dX
    return wrap_r_unconditional(r, dr_dX=dr_dX)


################### Check the behavior around the th=0, th=180, th=360
################### singularities. Gradients and values should be correct
axes = np.array(
    (
        (1.0, 2.0, 0.1),
        (1.0, 0, 0),
        (0, 0, -1.0),
    )
)
axes /= nps.dummy(nps.mag(axes), -1)  # normalize
for axis in axes:
    for th0 in (-np.pi, 0, np.pi):
        for dth in (-1e-4, -1e-10, 0, 1e-10, 1e-4):
            r = (th0 + dth) * axis

            ######### R_from_r, r_from_R
            if True:
                R, dR_dr = mrcal.R_from_r(r, get_gradients=True)

                R_ref = R_from_r(r)
                dR_dr__ref = grad(R_from_r, r)

                confirm_equal(
                    R,
                    R_ref,
                    msg=f"R_from_r result near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                )
                confirm_equal(
                    dR_dr,
                    dR_dr__ref,
                    msg=f"R_from_r J near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                )

                # I check R_roundtrip. The dr/dR computation assumes this
                r_roundtrip, dr_dR = mrcal.r_from_R(R, get_gradients=True)
                R_roundtrip, dR_dr = mrcal.R_from_r(r_roundtrip, get_gradients=True)

                confirm_equal(
                    mrcal.compose_r(r_roundtrip, -r),
                    0,
                    eps=1e-8,
                    msg=f"roundtrip r result near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                )
                confirm_equal(
                    nps.matmult(R_roundtrip, mrcal.invert_R(R)) - np.eye(3),
                    0,
                    eps=1e-8,
                    msg=f"roundtrip R result near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                )

                dr_dR__ref = grad__r_from_R(R_ref)

                confirm_equal(
                    dr_dR,
                    dr_dR__ref,
                    relative=True,
                    worstcase=True,
                    eps=1e-3,
                    reldiff_eps=1e-5,
                    msg=f"r_from_R J near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                )

            ######### compose_r
            if True:
                r1_simple = np.array(
                    (-0.02, -1.2, 0.4),
                )

                inv_r1_simple_r = mrcal.compose_r(-r1_simple, r)

                # it isn't possible to correctly decide if we should or should
                # not wrap the result for ALL cases (some cases have mag(r)=pi
                # exactly, up to machine precision). So I try both, and pick the
                # closer one
                r_roundtrip = mrcal.compose_r(r1_simple, inv_r1_simple_r)
                r_wrapped = wrap_r_unconditional(r)
                confirm_equal(
                    r_roundtrip,
                    r
                    if nps.norm2(r_roundtrip - r) < nps.norm2(r_roundtrip - r_wrapped)
                    else r_wrapped,
                    worstcase=True,
                    eps=1e-6,
                    msg="compose()",
                )

                for r0, r1 in (
                    (r, r1_simple),
                    (r, -r),
                    (r, r),
                    (r1_simple, inv_r1_simple_r),
                ):
                    ###### r01
                    r01, dr01_dr0, dr01_dr1 = mrcal.compose_r(
                        r0, r1, get_gradients=True
                    )
                    r01_ref = compose_r(r0, r1)
                    confirm_equal(
                        r01,
                        r01_ref,
                        worstcase=True,
                        relative=True,
                        eps=1e-3,
                        msg=f"compose_r(r0,r1) near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                    )
                    dr01_dr0__ref = grad(
                        lambda r0: compose_r(r0, r1),
                        r0,
                        forward_differences=True,
                        switch=wrap_r_unconditional,
                        step=1e-7,
                    )
                    dr01_dr1__ref = grad(
                        lambda r1: compose_r(r0, r1),
                        r1,
                        forward_differences=True,
                        switch=wrap_r_unconditional,
                        step=1e-7,
                    )
                    confirm_equal(
                        dr01_dr0,
                        dr01_dr0__ref,
                        worstcase=True,
                        relative=True,
                        eps=1e-2,
                        reldiff_eps=1e-3,
                        msg=f"compose_r(r0,r1) dr01_dr0 near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                    )
                    confirm_equal(
                        dr01_dr1,
                        dr01_dr1__ref,
                        worstcase=True,
                        relative=True,
                        eps=1e-2,
                        reldiff_eps=1e-3,
                        msg=f"compose_r(r0,r1) dr01_dr1 near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                    )

                    if r0 is r1:
                        continue

                    ###### r10
                    r10, dr10_dr1, dr10_dr0 = mrcal.compose_r(
                        r1, r0, get_gradients=True
                    )
                    r10_ref = compose_r(r1, r0)
                    confirm_equal(
                        r10,
                        r10_ref,
                        worstcase=True,
                        relative=True,
                        eps=1e-3,
                        msg=f"compose_r(r1,r0) near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                    )
                    dr10_dr0__ref = grad(
                        lambda r0: compose_r(r1, r0),
                        r0,
                        forward_differences=True,
                        switch=wrap_r_unconditional,
                        step=1e-7,
                    )
                    dr10_dr1__ref = grad(
                        lambda r1: compose_r(r1, r0),
                        r1,
                        forward_differences=True,
                        switch=wrap_r_unconditional,
                        step=1e-7,
                    )
                    confirm_equal(
                        dr10_dr0,
                        dr10_dr0__ref,
                        worstcase=True,
                        relative=True,
                        eps=1e-2,
                        reldiff_eps=1e-3,
                        msg=f"compose_r(r1,r0 dr10_dr0 near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                    )
                    confirm_equal(
                        dr10_dr1,
                        dr10_dr1__ref,
                        worstcase=True,
                        relative=True,
                        eps=1e-2,
                        reldiff_eps=1e-3,
                        msg=f"compose_r(r1,r0 dr10_dr1 near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                    )

            ######### rotate_point_r
            if True:
                # Simple reference implementation. Should move this to
                # test_poseutils_helpers.py. And test_poseutils_lib.py should
                # use it
                def rotate_point_r(r, p):
                    return nps.inner(p, R_from_r(r))

                p = np.array(
                    (3.0, -0.2, -0.9),
                )

                pt, dpt_dr, dpt_dp = mrcal.rotate_point_r(r, p, get_gradients=True)
                pt_ref = rotate_point_r(r, p)

                dpt_dr__ref = grad(lambda r: rotate_point_r(r, p), r)
                dpt_dp__ref = grad(lambda p: rotate_point_r(r, p), p)
                confirm_equal(
                    pt,
                    pt_ref,
                    msg=f"rotate_point_r(r,p) r near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                )
                confirm_equal(
                    dpt_dr,
                    dpt_dr__ref,
                    msg=f"rotate_point_r(r,p) dpt_dr near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                )
                confirm_equal(
                    dpt_dp,
                    dpt_dp__ref,
                    msg=f"rotate_point_r(r,p) dpt_dp near a singularity. axis={axis}, th0={th0:.2f}, dth={dth}",
                )

finish()
