#!/usr/bin/env python3

r"""Tests special-case projection functions

Simple things like project_lonlat(), project_stereographic(), etc.

I do 3 things:

Here I make sure the projection functions return the correct values. This is a
regression test, so the "right" values were recorded at some point, and any
deviation is flagged.

I make sure that project(unproject(x)) == x

I run a gradient check. I do these for the simple project_...()
function AND the generic project() function.

"""

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
import testutils
from test_calibration_helpers import grad

if len(sys.argv) != 2:
    raise Exception(
        "Need one argument on the commandline: the projection type. Currently I support 'pinhole','latlon','lonlat','stereographic'"
    )
if sys.argv[1] == "pinhole" or sys.argv[1] == "latlon" or sys.argv[1] == "lonlat":
    # pixels/rad
    fx, fy = 3000.0, 2000.0

    # pixel where latlon = (0,0) projects to. May be negative
    cx, cy = (-10000.0, 4000.0)

    # a few points, some wide, some not. Some behind the camera
    p = np.array(((1.0, 2.0, 10.0), (-1.1, 0.3, -1.0), (-0.9, -1.5, -1.0)))

    if sys.argv[1] == "pinhole":
        unproject_is_normalized = False
    else:
        unproject_is_normalized = True

    if sys.argv[1] == "pinhole":
        # pinhole projects ahead only
        p[:, 2] = abs(p[:, 2])

    if sys.argv[1] == "pinhole":
        lensmodel = "LENSMODEL_PINHOLE"
        func_project = mrcal.project_pinhole
        func_unproject = mrcal.unproject_pinhole
        name = "pinhole"

        q_projected_ref = np.array(
            [[-9700.0, 4400.0], [-13300.0, 4600.0], [-12700.0, 1000.0]]
        )

    elif sys.argv[1] == "lonlat":
        lensmodel = "LENSMODEL_LONLAT"
        func_project = mrcal.project_lonlat
        func_unproject = mrcal.unproject_lonlat
        name = "lonlat"

        q_projected_ref = np.array(
            [
                [-9700.99404253, 4392.88198287],
                [-16925.83416075, 4398.25498944],
                [-17226.33265541, 2320.61601685],
            ]
        )

    elif sys.argv[1] == "latlon":
        lensmodel = "LENSMODEL_LATLON"
        func_project = mrcal.project_latlon
        func_unproject = mrcal.unproject_latlon
        name = "latlon"

        q_projected_ref = np.array(
            [
                [-9706.7632608, 4394.7911197],
                [-12434.4909092, 9700.27171822],
                [-11389.09468198, -317.59786068],
            ]
        )

elif sys.argv[1] == "stereographic":
    lensmodel = "LENSMODEL_STEREOGRAPHIC"
    func_project = mrcal.project_stereographic
    func_unproject = mrcal.unproject_stereographic
    name = "stereographic"

    fx, fy, cx, cy = 1512.0, 1112, 500.0, 333.0

    # a few points, some wide, some not. Some behind the camera
    p = np.array(((1.0, 2.0, 10.0), (-1.1, 0.3, -1.0), (-0.9, -1.5, -1.0)))

    q_projected_ref = np.array(
        [
            [649.35582325, 552.6874014],
            [-5939.33490417, 1624.58376866],
            [-2181.52681292, -2953.8803086],
        ]
    )

    unproject_is_normalized = False

else:
    raise Exception(
        "Unknown projection type. Currently I support 'lonlat','stereographic'"
    )


intrinsics = (lensmodel, np.array((fx, fy, cx, cy)))

q_projected = func_project(p, intrinsics[1])
testutils.confirm_equal(
    q_projected, q_projected_ref, msg=f"project_{name}()", worstcase=True, relative=True
)

testutils.confirm_equal(
    mrcal.project(p, *intrinsics),
    q_projected,
    msg=f"project({name}) returns the same as project_{name}()",
    worstcase=True,
    relative=True,
)

v_unprojected = func_unproject(q_projected, intrinsics[1])
if unproject_is_normalized:
    testutils.confirm_equal(
        nps.mag(v_unprojected),
        1.0,
        msg=f"unproject_{name}() returns normalized vectors",
        worstcase=True,
        relative=True,
    )
    testutils.confirm_equal(
        v_unprojected,
        p / nps.dummy(nps.mag(p), axis=-1),
        msg=f"unproject_{name}()",
        worstcase=True,
        relative=True,
    )
else:
    cos = nps.inner(v_unprojected, p) / (nps.mag(p) * nps.mag(v_unprojected))
    cos = np.clip(cos, -1, 1)
    testutils.confirm_equal(
        np.arccos(cos),
        np.zeros((p.shape[0],), dtype=float),
        msg=f"unproject_{name}()",
        worstcase=True,
    )

    # Not normalized by default. Make sure that if I ask for it to be
    # normalized, that it is
    testutils.confirm_equal(
        nps.mag(mrcal.unproject(q_projected, *intrinsics, normalize=True)),
        1.0,
        msg=f"unproject({name},normalize = True) returns normalized vectors",
        worstcase=True,
        relative=True,
    )
    testutils.confirm_equal(
        nps.mag(
            mrcal.unproject(
                q_projected, *intrinsics, normalize=True, get_gradients=True
            )[0]
        ),
        1.0,
        msg=f"unproject({name},normalize = True, get_gradients=True) returns normalized vectors",
        worstcase=True,
        relative=True,
    )


testutils.confirm_equal(
    mrcal.unproject(q_projected, *intrinsics),
    v_unprojected,
    msg=f"unproject({name}) returns the same as unproject_{name}()",
    worstcase=True,
    relative=True,
)

testutils.confirm_equal(
    mrcal.project(mrcal.unproject(q_projected, *intrinsics), *intrinsics),
    q_projected,
    msg="project(unproject()) is an identity",
    worstcase=True,
    relative=True,
)

testutils.confirm_equal(
    func_project(func_unproject(q_projected, intrinsics[1]), intrinsics[1]),
    q_projected,
    msg=f"project_{name}(unproject_{name}()) is an identity",
    worstcase=True,
    relative=True,
)


# Now gradients for project()
ipt = 1
_, dq_dp_reported = func_project(p[ipt], intrinsics[1], get_gradients=True)
dq_dp_observed = grad(lambda p: func_project(p, intrinsics[1]), p[ipt])
testutils.confirm_equal(
    dq_dp_reported,
    dq_dp_observed,
    msg=f"project_{name}() dq/dp",
    worstcase=True,
    relative=True,
    eps=1e-5,
)
_, dq_dp_reported, dq_di_reported = mrcal.project(
    p[ipt], *intrinsics, get_gradients=True
)
dq_dp_observed = grad(lambda p: mrcal.project(p, *intrinsics), p[ipt])
dq_di_observed = grad(
    lambda intrinsics_data: mrcal.project(p[ipt], intrinsics[0], intrinsics_data),
    intrinsics[1],
)
testutils.confirm_equal(
    dq_dp_reported,
    dq_dp_observed,
    msg=f"project({name}) dq/dp",
    worstcase=True,
    relative=True,
    eps=1e-5,
)
testutils.confirm_equal(
    dq_di_reported,
    dq_di_observed,
    msg=f"project({name}) dq/di",
    worstcase=True,
    relative=True,
    eps=1e-5,
)

# Now gradients for unproject()
ipt = 1
_, dv_dq_reported = func_unproject(q_projected[ipt], intrinsics[1], get_gradients=True)
dv_dq_observed = grad(lambda q: func_unproject(q, intrinsics[1]), q_projected[ipt])
testutils.confirm_equal(
    dv_dq_reported,
    dv_dq_observed,
    msg=f"unproject_{name}() dv/dq",
    worstcase=True,
    relative=True,
    eps=2e-6,
)


for normalize in (False, True):
    v_unprojected, dv_dq_reported, dv_di_reported = mrcal.unproject(
        q_projected[ipt], *intrinsics, get_gradients=True, normalize=normalize
    )
    dv_dq_observed = grad(
        lambda q: mrcal.unproject(q, *intrinsics, normalize=normalize), q_projected[ipt]
    )
    dv_di_observed = grad(
        lambda intrinsics_data: mrcal.unproject(
            q_projected[ipt], intrinsics[0], intrinsics_data, normalize=normalize
        ),
        intrinsics[1],
    )
    testutils.confirm_equal(
        dv_dq_reported,
        dv_dq_observed,
        msg=f"unproject({name}, normalize={normalize}) dv/dq",
        worstcase=True,
        relative=True,
        eps=1e-5,
    )
    testutils.confirm_equal(
        dv_di_reported,
        dv_di_observed,
        msg=f"unproject({name}, normalize={normalize}) dv/di",
        worstcase=True,
        relative=True,
        eps=1e-5,
    )

    v_unprojected_inplace = v_unprojected.copy() * 0
    dv_dq_reported_inplace = dv_dq_reported.copy() * 0
    dv_di_reported_inplace = dv_di_reported.copy() * 0

    mrcal.unproject(
        q_projected[ipt],
        *intrinsics,
        get_gradients=True,
        normalize=normalize,
        out=[v_unprojected_inplace, dv_dq_reported_inplace, dv_di_reported_inplace],
    )
    testutils.confirm_equal(
        v_unprojected_inplace,
        v_unprojected,
        msg=f"unproject({name}, normalize={normalize}) works in-place: v_unprojected",
        worstcase=True,
        relative=True,
    )
    testutils.confirm_equal(
        dv_dq_reported_inplace,
        dv_dq_reported,
        msg=f"unproject({name}, normalize={normalize}) works in-place: dv_dq",
        worstcase=True,
        relative=True,
    )
    testutils.confirm_equal(
        dv_di_reported_inplace,
        dv_di_reported,
        msg=f"unproject({name}, normalize={normalize}) works in-place: dv_di",
        worstcase=True,
        relative=True,
    )

testutils.finish()
