#!/usr/bin/python3

import sys
import numpy as np
import numpysane as nps
import gnuplotlib as gp

import pyopengv
import mrcal
from mrcal.utils import _sorted_eig

# 6ft
baseline = 6*12*2.54/100

# shape (Nimages=2, Npoints, Nxy=2)
q0_q1estimate_q1 = np.array(( \
                              [1667.33555875, 873.75422761,  1635.1480094,  764.73188304,  1634.42603067, 762.9145074 ],
                              [1651.22932438, 597.29294354,  1694.29003556, 483.66051127,  1695.24788451, 480.53042433],
                              [479.12304278,  459.47592743,  503.81013295,  224.94857087,  501.43836221,  227.02909723],
                              [1713.47755101, 1274.14990286, 1651.7598256,  1175.40154221, 1648.26427625, 1172.46634256],
                              [2117.72865246, 1425.35833012, 2043.66738197, 1357.46883217, 2039.40979881, 1355.77964729],
                              [2303.76180594, 1350.68778224, 2233.99016423, 1296.08388873, 2233.53172132, 1296.18265221],
                              [2169.27443859, 1596.40530304, 2078.33654808, 1531.37538255, 2083.82913855, 1532.14687599],
                              [1883.59263645, 1874.83122376, 1779.11311358, 1792.90977969, 1780.69886878, 1791.84423636],
                              [1723.31155023, 1933.00732172, 1624.76836389, 1840.4004719,  1623.75087896, 1840.96649026],
                              [1218.72294546, 2098.03747716, 1208.03753971, 1991.18341968, 1208.1542519,  1990.26499417],
                              [1167.50308963, 2611.65584817, 3301.35517957, 2736.08371383, 3302.36350656, 2738.36081631],
                              [1139.48389554, 2378.62442443, 1090.78514192, 2277.16868773, 1098.62296769, 2276.34354627],
                              [2699.19675445, 2415.14848965, 2619.38490824, 2394.85734231, 2621.32696316, 2394.84054909],
                              [2722.3940691,  2535.83231239, 2665.67250793, 2519.36476237, 2666.60388253, 2518.40371615],
                              [2894.38848038, 2662.08352918, 2857.79392479, 2652.93489028, 2859.50967363, 2653.78796184],
                              [4186.1762928,  2155.24893422, 4171.53847057, 2228.43804541, 4173.68227017, 2226.23375795],
                              [4312.4275096,  2144.27056755, 4290.47077624, 2226.60831763, 4289.68956009, 2223.7795986 ],
                              [5810.97456112, 1366.6362612,  5703.02062212, 1582.5441392,  5700.98270992, 1581.54872233],
                              [5638.98014983, 766.48554948,  5642.63960539, 1061.071722,   5640.96848272, 1062.74035658],
                              [5884.91434412, 1939.14858979, 5757.82976588, 2126.80113723, 5757.52979879, 2127.81086967],
                              [5572.7947573,  2220.67098501, -1.0,          -1.0,          5619.35819897, 2284.87209398],
                              [3884.47331949, 2722.23151915, 4915.86009802, 2510.66026109, 4917.44044553, 2511.07917974],
                              [ 196.21859106, 2534.6153689,  1548.48293462, 2925.71142905, 1550.76587178, 2925.52198467],
                              [2828.68230912, 1457.09683885, 2747.56449323, 1442.90122106, 2747.63192025, 1442.72479765],
                              [2817.86660034, 1572.0137447,  2728.63700285, 1556.46616332, 2728.68750646, 1556.33523169],
                              [2942.92323318, 1367.86724136, 2872.62112607, 1363.13536877, 2872.46997806, 1363.45798433],
                              [3086.90735639, 1373.27509575, 3015.25328568, 1378.00696835, 3015.72254709, 1378.06237561],
                              [3506.38294069, 1619.17153289, 3425.64410723, 1656.17683156, 3425.87250498, 1656.83487376],
                              [3434.98023817, 1673.83172314, 3348.85881581, 1703.97422097, 3347.5876093,  1704.71786723],
                              [4009.41012529, 1694.50086451, 3918.08560669, 1766.63478993, 3918.42876635, 1766.8819185 ],
                              [4549.84202509, 1264.89536278, 4477.30464905, 1386.19736137, 4475.12668105, 1385.98601876],
                             ))
q0q1 = q0_q1estimate_q1[:,(0,1,4,5)]
q = np.ascontiguousarray(nps.xchg(q0q1.reshape((len(q0q1),2,2),),
                                  0,1))

model0 = mrcal.cameramodel('dance3-centralized.cameramodel')
model1 = mrcal.cameramodel('dance4-centralized.cameramodel')


################# seed
mask_far = (820 < q[0,:,1]) * (q[0,:,1] < 1750)
q0,q1 = q
v0 = mrcal.unproject(q0, *model0.intrinsics(),
                     normalize = True)
v1 = mrcal.unproject(q1, *model1.intrinsics(),
                     normalize = True)

R01 = mrcal.align_procrustes_vectors_R01(v0[mask_far],
                                         v1[mask_far])
# Have an estimate for R01. I refine it

R01 = \
    pyopengv.relative_pose_eigensolver(v0, v1,
                                       # seed
                                       R01)

n = np.cross(v0, mrcal.rotate_point_R(R01, v1))
l,v = _sorted_eig(np.sum(nps.outer(n,n), axis=0))
# t is the eigenvector corresponding to the smallest eigenvalue
t = v[:,0]
# t.x should be <0 in THIS dataset only
if t[0] > 0:
    t *= -1.
t *= baseline
Rt01 = nps.glue(R01, t, axis=-2)
rt_10 = mrcal.rt_from_Rt(mrcal.invert_Rt(Rt01))
######### testing this solution
p = mrcal.triangulate_geometric(v0, v1,
                                v_are_local = True,
                                Rt01        = Rt01 )
mask_divergent = (nps.norm2(p) == 0)

th = \
    np.arccos(nps.inner(v0,mrcal.rotate_point_R(Rt01[:3,:],v1)))
mask_too_divergent = (th>0.1) * mask_divergent

if any(mask_divergent):
    raise Exception("Some points are divergent")


rt_0r = model0.extrinsics_rt_fromref()
rt_1r = mrcal.compose_rt(rt_10, rt_0r)

model1.extrinsics_rt_fromref(rt_1r)
model1.write("0.cameramodel")