import numpy as np
import gtsam

from read_write_model import read_model

cameras, images, points3d = read_model("lost_underconstrained_data")

noisemodel = gtsam.noiseModel.Isotropic.Sigma(2, 1e-4)
error = []
for jj in points3d.keys():
    if len(points3d[jj].image_ids) < 5: 
        continue
    kpt_vec, cam_vec = [], []
    for idx, kk in enumerate(points3d[jj].image_ids):
        # Add keypoint measurements.
        pp = points3d[jj].point2D_idxs[idx];  # keypoint index
        kpt_vec.append(gtsam.Point2(images[kk].xys[pp].T.astype("double")))

        # Add cameras.
        R_BC = images[kk].qvec2rotmat().T
        r_CB_B = -R_BC @ images[kk].tvec[..., None]
        fx, fy, cx, cy = cameras[images[kk].camera_id].params
        cam_vec.append(
            gtsam.PinholeCameraCal3_S2(
                gtsam.Pose3(gtsam.Rot3(R_BC), gtsam.Point3(r_CB_B.flatten())),
                  gtsam.Cal3_S2(fx, fy, 0.0, cx, cy),
            )
        )

    # Triangulate.        
    kpt_vec = gtsam.Point2Vector(kpt_vec)
    cam_vec = gtsam.CameraSetCal3_S2(cam_vec)
    xyz_est = gtsam.triangulatePoint3(cam_vec, kpt_vec, rank_tol=1e-9, optimize=False, model=noisemodel, useLOST=True)
    # xyz_est = gtsam.triangulatePoint3(cam_vec, kpt_vec, rank_tol=1e-9, optimize=True, model=noisemodel)
    error.append(np.linalg.norm(points3d[jj].xyz - xyz_est))
    print(xyz_est)

print("Mean error:", np.mean(error))
print(error)