#!/usr/bin/env python3

r"""Basic camera-calibration test

I observe, with noise, a number of chessboards from various angles with several
cameras. And I make sure that I can more or less compute the camera intrinsics
and extrinsics

"""

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 sample_dqref
import copy

# I want the RNG to be deterministic
np.random.seed(0)

############# Set up my world, and compute all the perfect positions, pixel
############# observations of everything
models_ref = (
    mrcal.cameramodel(f"{testdir}/data/cam0.opencv8.cameramodel"),
    mrcal.cameramodel(f"{testdir}/data/cam0.opencv8.cameramodel"),
    mrcal.cameramodel(f"{testdir}/data/cam1.opencv8.cameramodel"),
    mrcal.cameramodel(f"{testdir}/data/cam1.opencv8.cameramodel"),
)

imagersizes = nps.cat(*[m.imagersize() for m in models_ref])
lensmodel = models_ref[0].intrinsics()[0]
# I have opencv8 models_ref, but let me truncate to opencv4 models_ref to keep this
# simple and fast
lensmodel = "LENSMODEL_OPENCV4"
for m in models_ref:
    m.intrinsics(intrinsics=(lensmodel, m.intrinsics()[1][:8]))
Nintrinsics = mrcal.lensmodel_num_params(lensmodel)

Ncameras = len(models_ref)
Nframes = 50

models_ref[0].extrinsics_rt_fromref(np.zeros((6,), dtype=float))
models_ref[1].extrinsics_rt_fromref(np.array((0.08, 0.2, 0.02, 1.0, 0.9, 0.1)))
models_ref[2].extrinsics_rt_fromref(np.array((0.01, 0.07, 0.2, 2.1, 0.4, 0.2)))
models_ref[3].extrinsics_rt_fromref(np.array((-0.1, 0.08, 0.08, 4.4, 0.2, 0.1)))


pixel_uncertainty_stdev = 1.5
object_spacing = 0.1
object_width_n = 10
object_height_n = 9
calobject_warp_ref = np.array((0.002, -0.005))

# shapes (Nframes, Ncameras, Nh, Nw, 2),
#        (Nframes, 4,3)
q_ref, Rt_ref_board_ref = mrcal.synthesize_board_observations(
    models_ref,
    object_width_n=object_width_n,
    object_height_n=object_height_n,
    object_spacing=object_spacing,
    calobject_warp=calobject_warp_ref,
    rt_ref_boardcenter=np.array((0.0, 0.0, 0.0, -2, 0, 4.0)),
    rt_ref_boardcenter__noiseradius=np.array(
        (
            np.pi / 180.0 * 30.0,
            np.pi / 180.0 * 30.0,
            np.pi / 180.0 * 20.0,
            2.5,
            2.5,
            2.0,
        )
    ),
    Nframes=Nframes,
)
frames_ref = mrcal.rt_from_Rt(Rt_ref_board_ref)

############# I have perfect observations in q_ref. I corrupt them by noise
# weight has shape (Nframes, Ncameras, Nh, Nw),
weight01 = (np.random.rand(*q_ref.shape[:-1]) + 1.0) / 2.0  # in [0,1]
weight0 = 0.2
weight1 = 1.0
weight = weight0 + (weight1 - weight0) * weight01

# I want observations of shape (Nframes*Ncameras, Nh, Nw, 3) where each row is
# (x,y,weight)
observations_ref = nps.clump(nps.glue(q_ref, nps.dummy(weight, -1), axis=-1), n=2)

q_noise, observations = sample_dqref(
    observations_ref, pixel_uncertainty_stdev, make_outliers=True
)

# Now I make some of the observations bogus, and mark them as input outliers.
# The solve should be robust to that, but any code that uses the bogus data
# DESPITE it being marked as bogus will generate a test failure
#
# Let's pretend the center of the chessboard has an apriltag, so all those
# observations are bogus. I block out a 5x5 chunk in the center
i0 = object_height_n // 2
j0 = object_width_n // 2
observations[..., i0 - 2 : i0 + 3, j0 - 2 : j0 + 3, 2] = -1.0  # weight<=0: outlier
observations[
    ..., i0 - 2 : i0 + 3, j0 - 2 : j0 + 3, :2
] = -100.0  # all the values are bogus

############# Now I pretend that the noisy observations are all I got, and I run
############# a calibration from those

# Dense observations. All the cameras see all the boards
indices_frame_camera = np.zeros((Nframes * Ncameras, 2), dtype=np.int32)
indices_frame = indices_frame_camera[:, 0].reshape(Nframes, Ncameras)
indices_frame.setfield(
    nps.outer(np.arange(Nframes, dtype=np.int32), np.ones((Ncameras,), dtype=np.int32)),
    dtype=np.int32,
)
indices_camera = indices_frame_camera[:, 1].reshape(Nframes, Ncameras)
indices_camera.setfield(
    nps.outer(np.ones((Nframes,), dtype=np.int32), np.arange(Ncameras, dtype=np.int32)),
    dtype=np.int32,
)

indices_frame_camintrinsics_camextrinsics = nps.glue(
    indices_frame_camera, indices_frame_camera[:, (1,)] - 1, axis=-1
)


intrinsics_data, extrinsics_rt_fromref, frames_rt_toref = mrcal.seed_stereographic(
    imagersizes=imagersizes,
    focal_estimate=1500,
    indices_frame_camera=indices_frame_camera,
    observations=observations,
    object_spacing=object_spacing,
)

# I have a stereographic intrinsics estimate. Mount it into a full distortiony
# model, seeded with random numbers
intrinsics = np.zeros((Ncameras, Nintrinsics), dtype=float)
intrinsics[:, :4] = intrinsics_data
intrinsics[:, 4:] = np.random.random((Ncameras, intrinsics.shape[1] - 4)) * 1e-6

optimization_inputs = dict(
    intrinsics=intrinsics,
    extrinsics_rt_fromref=extrinsics_rt_fromref,
    frames_rt_toref=frames_rt_toref,
    points=None,
    observations_board=observations,
    indices_frame_camintrinsics_camextrinsics=indices_frame_camintrinsics_camextrinsics,
    observations_point=None,
    indices_point_camintrinsics_camextrinsics=None,
    lensmodel=lensmodel,
    calobject_warp=None,
    imagersizes=imagersizes,
    calibration_object_spacing=object_spacing,
    verbose=False,
    do_apply_regularization=True,
)

# Solve this thing incrementally
optimization_inputs["do_optimize_intrinsics_core"] = False
optimization_inputs["do_optimize_intrinsics_distortions"] = False
optimization_inputs["do_optimize_extrinsics"] = True
optimization_inputs["do_optimize_frames"] = True
optimization_inputs["do_optimize_calobject_warp"] = False
mrcal.optimize(**optimization_inputs, do_apply_outlier_rejection=True)

optimization_inputs["do_optimize_intrinsics_core"] = True
optimization_inputs["do_optimize_intrinsics_distortions"] = False
optimization_inputs["do_optimize_extrinsics"] = True
optimization_inputs["do_optimize_frames"] = True
optimization_inputs["do_optimize_calobject_warp"] = False
mrcal.optimize(**optimization_inputs, do_apply_outlier_rejection=True)

testutils.confirm_equal(
    mrcal.num_states(**optimization_inputs),
    4 * Ncameras + 6 * (Ncameras - 1) + 6 * Nframes,
    msg="num_states()",
)
testutils.confirm_equal(
    mrcal.num_states_intrinsics(**optimization_inputs),
    4 * Ncameras,
    msg="num_states_intrinsics()",
)
testutils.confirm_equal(
    mrcal.num_intrinsics_optimization_params(**optimization_inputs),
    4,
    msg="num_intrinsics_optimization_params()",
)
testutils.confirm_equal(
    mrcal.num_states_extrinsics(**optimization_inputs),
    6 * (Ncameras - 1),
    msg="num_states_extrinsics()",
)
testutils.confirm_equal(
    mrcal.num_states_frames(**optimization_inputs),
    6 * Nframes,
    msg="num_states_frames()",
)
testutils.confirm_equal(
    mrcal.num_states_points(**optimization_inputs), 0, msg="num_states_points()"
)
testutils.confirm_equal(
    mrcal.num_states_calobject_warp(**optimization_inputs),
    0,
    msg="num_states_calobject_warp()",
)

testutils.confirm_equal(
    mrcal.num_measurements_boards(**optimization_inputs),
    object_width_n * object_height_n * 2 * Nframes * Ncameras,
    msg="num_measurements_boards()",
)
testutils.confirm_equal(
    mrcal.num_measurements_points(**optimization_inputs),
    0,
    msg="num_measurements_points()",
)
testutils.confirm_equal(
    mrcal.num_measurements_regularization(**optimization_inputs),
    Ncameras * 2,
    msg="num_measurements_regularization()",
)


optimization_inputs["do_optimize_intrinsics_core"] = True
optimization_inputs["do_optimize_intrinsics_distortions"] = True
optimization_inputs["do_optimize_extrinsics"] = True
optimization_inputs["do_optimize_frames"] = True
optimization_inputs["do_optimize_calobject_warp"] = True

optimization_inputs["calobject_warp"] = np.array((0.001, 0.001))
stats = mrcal.optimize(**optimization_inputs, do_apply_outlier_rejection=True)

rmserr = stats["rms_reproj_error__pixels"]


testutils.confirm_equal(
    mrcal.state_index_intrinsics(2, **optimization_inputs),
    8 * 2,
    msg="state_index_intrinsics()",
)
testutils.confirm_equal(
    mrcal.state_index_extrinsics(2, **optimization_inputs),
    8 * Ncameras + 6 * 2,
    msg="state_index_extrinsics()",
)
testutils.confirm_equal(
    mrcal.state_index_frames(2, **optimization_inputs),
    8 * Ncameras + 6 * (Ncameras - 1) + 6 * 2,
    msg="state_index_frames()",
)
testutils.confirm_equal(
    mrcal.state_index_calobject_warp(**optimization_inputs),
    8 * Ncameras + 6 * (Ncameras - 1) + 6 * Nframes,
    msg="state_index_calobject_warp()",
)

testutils.confirm_equal(
    mrcal.measurement_index_boards(2, **optimization_inputs),
    object_width_n * object_height_n * 2 * 2,
    msg="measurement_index_boards()",
)
testutils.confirm_equal(
    mrcal.measurement_index_regularization(**optimization_inputs),
    object_width_n * object_height_n * 2 * Nframes * Ncameras,
    msg="measurement_index_regularization()",
)


############# Calibration computed. Now I see how well I did
models_solved = [
    mrcal.cameramodel(optimization_inputs=optimization_inputs, icam_intrinsics=i)
    for i in range(Ncameras)
]

# if 0:
#     for i in range(1,Ncameras):
#         models_solved[i].write(f'/tmp/tst{i}.cameramodel')


testutils.confirm_equal(rmserr, 0, eps=2.5, msg="Converged to a low RMS error")

testutils.confirm_equal(
    optimization_inputs["calobject_warp"],
    calobject_warp_ref,
    eps=2e-3,
    msg="Recovered the calibration object shape",
)

testutils.confirm_equal(
    np.std(mrcal.residuals_board(optimization_inputs, x=stats["x"])),
    pixel_uncertainty_stdev,
    eps=pixel_uncertainty_stdev * 0.1,
    msg="Residual have the expected distribution",
)

# Checking the extrinsics. These aren't defined absolutely: each solve is free
# to put the observed frames anywhere it likes. The projection-diff code
# computes a transformation to address this. Here I simply look at the relative
# transformations between cameras, which would cancel out any extra
# transformations, AND since camera0 is fixed at the identity transformation, I
# can simply look at each extrinsics transformation.
for icam in range(1, len(models_ref)):
    Rt_extrinsics_err = mrcal.compose_Rt(
        models_solved[icam].extrinsics_Rt_fromref(),
        models_ref[icam].extrinsics_Rt_toref(),
    )

    testutils.confirm_equal(
        nps.mag(Rt_extrinsics_err[3, :]),
        0.0,
        eps=0.05,
        msg=f"Recovered extrinsic translation for camera {icam}",
    )

    testutils.confirm_equal(
        (np.trace(Rt_extrinsics_err[:3, :]) - 1) / 2.0,
        1.0,
        eps=np.cos(1.0 * np.pi / 180.0),  # 1 deg
        msg=f"Recovered extrinsic rotation for camera {icam}",
    )

Rt_frame_err = mrcal.compose_Rt(
    mrcal.Rt_from_rt(optimization_inputs["frames_rt_toref"]),
    mrcal.invert_Rt(Rt_ref_board_ref),
)

testutils.confirm_equal(
    np.max(nps.mag(Rt_frame_err[..., 3, :])),
    0.0,
    eps=0.08,
    msg="Recovered frame translation",
)
testutils.confirm_equal(
    np.min((nps.trace(Rt_frame_err[..., :3, :]) - 1) / 2.0),
    1.0,
    eps=np.cos(1.0 * np.pi / 180.0),  # 1 deg
    msg="Recovered frame rotation",
)


# Checking the intrinsics. Each intrinsics vector encodes an implicit
# transformation. I compute and apply this transformation when making my
# intrinsics comparisons. I make sure that within some distance of the pixel
# center, the projections match up to within some number of pixels
Nw = 60


def projection_diff(models_ref, max_dist_from_center):
    lensmodels = [model.intrinsics()[0] for model in models_ref]
    intrinsics_data = [model.intrinsics()[1] for model in models_ref]

    # v  shape (...,Ncameras,Nheight,Nwidth,...)
    # q0 shape (...,         Nheight,Nwidth,...)
    v, q0 = mrcal.sample_imager_unproject(
        Nw, None, *imagersizes[0], lensmodels, intrinsics_data, normalize=True
    )

    W, H = imagersizes[0]
    focus_center = None
    focus_radius = -1
    if focus_center is None:
        focus_center = ((W - 1.0) / 2.0, (H - 1.0) / 2.0)
    if focus_radius < 0:
        focus_radius = min(W, H) / 6.0

    implied_Rt10 = mrcal.implied_Rt10__from_unprojections(
        q0, v[0, ...], v[1, ...], focus_center=focus_center, focus_radius=focus_radius
    )

    q1 = mrcal.project(
        mrcal.transform_point_Rt(implied_Rt10, v[0, ...]),
        lensmodels[1],
        intrinsics_data[1],
    )
    diff = nps.mag(q1 - q0)

    # zero-out everything too far from the center
    center = (imagersizes[0] - 1.0) / 2.0
    diff[nps.norm2(q0 - center) > max_dist_from_center * max_dist_from_center] = 0
    # gp.plot(diff,
    #         ascii = True,
    #         using = mrcal.imagergrid_using(imagersizes[0], Nw),
    #         square=1, _with='image', tuplesize=3, hardcopy='/tmp/yes.gp', cbmax=3)

    return diff


for icam in range(len(models_ref)):
    diff = projection_diff((models_ref[icam], models_solved[icam]), 800)

    testutils.confirm_equal(
        diff, 0, worstcase=True, eps=6.0, msg=f"Recovered intrinsics for camera {icam}"
    )


# It would be nice to check the outlier detections, but this is iffy. Here I'm
# generating 1% outliers (hard-coded in sample_dqref()), but the outlier
# rejection is overly aggressive. I'm currently seeing 4.4%:
#
#   np.count_nonzero(observations[...,2]<=0) / observations[...,0].ravel().size
#
# The outlier rejection scheme just cuts out 3sigma residuals and above, so it's
# not great. I'm not entirely sure why it's over-reporting the outliers here,
# but I should investigate that at the same time as I overhaul the outlier
# rejection scheme (presumably to use one of my flavors of Cook's D factor)

# I test make_perfect_observations(). Doing it here is easy; doing it elsewhere
# it much more work
if True:
    optimization_inputs_perfect = copy.deepcopy(optimization_inputs)

    mrcal.make_perfect_observations(
        optimization_inputs_perfect, observed_pixel_uncertainty=0
    )
    x = mrcal.optimizer_callback(
        **optimization_inputs_perfect, no_jacobian=True, no_factorization=True
    )[1]

    Nmeas = mrcal.num_measurements_boards(**optimization_inputs_perfect)
    if Nmeas > 0:
        i_meas0 = mrcal.measurement_index_boards(0, **optimization_inputs_perfect)
        testutils.confirm_equal(
            x[i_meas0 : i_meas0 + Nmeas],
            0,
            worstcase=True,
            eps=1e-8,
            msg="make_perfect_observations() works for boards",
        )
    else:
        testutils.confirm(False, msg="Nmeasurements_boards <= 0")
testutils.finish()
