#!/usr/bin/env python
"""
Library handling projection operations of a standard camera model.
"""
from __future__ import division, print_function
import os
import numpy as np
import cv2
import StringIO
from PIL import Image as PILImage
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
import glob
from scipy import interpolate
from scipy.optimize import minimize_scalar
import copy

# ROS imports
from rosbag import Bag
from cv_bridge import CvBridge
from tf.transformations import quaternion_multiply, quaternion_matrix, \
    quaternion_from_matrix, quaternion_from_euler, quaternion_inverse, \
    euler_matrix, quaternion_slerp

# ADAPT imports.
from sensor_models.nav_conversions import enu_to_llh, llh_to_enu
from sensor_models.nav_state import NavStateINS, trajectory_to_kmz
from sensor_models.utilities import rotation_between_quats, \
    draw_moving_coordinate_system
from colmap_processing.calibration import horn, get_rvec_btw_times
from colmap_processing.colmap_interface import read_images_binary, \
    read_points3D_binary, read_cameras_binary, qvec2rotmat, \
    standard_cameras_from_colmap
from colmap_processing.colmap_interface import Image as ColmapImage
from postflight_scripts.ins import ANINS


bridge = CvBridge()

# Bag information
nav_bag_fname = '/home/user/adapt_ws/data/2022_jan_flight/nav_data.bag'
nav_bag_fname = '/host_filesystem/mnt/homenas2/noaa_adapt/adapt_flights/2022-04-21/nav_data.bag'
out_dir = os.path.splitext(nav_bag_fname)[0]

# ----------------------------------------------------------------------------

# You should have a colmap directory where all of the Colmap-generated files
# reside.

colmap_dir = '/home/user/adapt_ws/data/2022_jan_flight'
colmap_dir = '/host_filesystem/mnt/data2tb/adapt/2022-04-21'

# Directory with all of the raw images.
colmap_images_subdir = 'images0'

# Sub-directory containing the images.bin and cameras.bin. Set to '' if in the
# top-level Colmap directory.
#sparse_recon_subdir = 'sparse/1'
sparse_recon_subdir = 'sparse/0'
sparse_recon_subdir = 'snapshots/0462855338'
aligned_sparse_recon_subdir = 'aligned'
# ----------------------------------------------------------------------------


ins = ANINS.from_ros_bag(nav_bag_fname)


if False:
    # Write navigation trajectory to kmz file to view in Google Earth.
    filename = nav_bag_fname.replace('.bag', '.kmz')
    trajectory_to_kmz(filename, pos_times, lats, lons, alts,
                      pos_covs=pos_covs, min_dt=0.25, name='flight_path')

# We take the INS-reported position (converted from latitude, longitude, and
# altitude into easting/northing/up coordinates) and assign it to each image.
ins.print_origin()
# ----------------------------------------------------------------------------


# --------------------- Read Existing Colmap Reconstruction ------------------
# Read in the Colmap details of all images.
images_bin_fname = '%s/%s/images.bin' % (colmap_dir, sparse_recon_subdir)
colmap_images = read_images_binary(images_bin_fname)
camera_bin_fname = '%s/%s/cameras.bin' % (colmap_dir, sparse_recon_subdir)
colmap_cameras = read_cameras_binary(camera_bin_fname)

# If Colmap didn't use a single camera for all images, let's force all to use
# one of the redundant versions.
use_camera_id = 1
image_times = {}
for ind in colmap_images:
    colmap_images[ind] = ColmapImage(colmap_images[ind].id,
                                     colmap_images[ind].qvec,
                                     colmap_images[ind].tvec,
                                     use_camera_id,
                                     colmap_images[ind].name,
                                     colmap_images[ind].xys,
                                     colmap_images[ind].point3D_ids)
    img_fname = os.path.split(colmap_images[ind].name)[1]
    image_times[ind] = float(os.path.splitext(img_fname)[0])/1000000

sfm_cm = standard_cameras_from_colmap(colmap_cameras, colmap_images,
                                      image_times)[use_camera_id]


image_times = sfm_cm.platform_pose_provider.times
nav_times = ins.times
image_times = image_times[image_times >= min(nav_times)]
image_times = image_times[image_times <= max(nav_times)]

ins_cm = copy.deepcopy(sfm_cm)
ins_cm._platform_pose_provider = ins
# ----------------------------------------------------------------------------

# ----------------------------------------------------------------------------
# Compare linear acceleration between SfM and INS.

# Estimate acceleration from SfM.
sfm_accel = []
sfm_times = []
for i in range(len(image_times) - 2):
    t1 = image_times[i]
    t2 = image_times[i + 1]
    t3 = image_times[i + 2]
    if abs(t1 - t2) > 0.5 or abs(t1 - t3) > 0.5:
        continue

    pos1 = sfm_cm.platform_pose_provider.pose(t=t1)[0]
    pos2 = sfm_cm.platform_pose_provider.pose(t=t2)[0]
    pos3 = sfm_cm.platform_pose_provider.pose(t=t3)[0]
    v1 = (pos2 - pos1)/(t2 - t1)
    v2 = (pos3 - pos2)/(t3 - t2)
    a = (v2 - v1)/((t2 + t1)/2 - (t3 + t2)/2)
    sfm_times.append(t2)
    sfm_accel.append(a)

sfm_times = np.array(sfm_times)
sfm_accel = np.array(sfm_accel).T

data = ins.get_raw_imu()
imu_times = data[0]
imu_accel = data[:4]

mag_sfm = np.sqrt(np.sum(sfm_accel[1:]**2, axis=0))
mag_ins = np.sqrt(np.sum(imu_accel[1:]**2, axis=0))

plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
plt.rc('font', **{'size': 20})
plt.rc('axes', linewidth=4)
plt.plot(sfm_times, mag_sfm, label='SfM')
plt.plot(imu_times, mag_ins, label='INS')
plt.xlabel('Time (s)', fontsize=40)
plt.ylabel('Acceleration (m/s^2)', fontsize=40)
plt.legend()
plt.tight_layout()


# ----------------------------------------------------------------------------
def err(dt):
    err_ = 0
    step = 1
    for i in range(len(image_times) - step):
        tvec1 = get_rvec_btw_times(sfm_cm, image_times[i],
                                   image_times[i + step])
        tvec2 = get_rvec_btw_times(ins_cm, image_times[i],
                                   image_times[i + step] + dt)
        err_ += abs(np.linalg.norm(tvec1) - np.linalg.norm(tvec2))

    print(dt, err_)

    return err_


dts = np.linspace(-0.1, 0.1, 50)
errs = [err(dt) for dt in dts]
plt.plot(dts, errs)
dt = dts[np.argmin(errs)]
res = minimize_scalar(err, [-0.1, 0, 0.1])
dt = 0


sfm_tvecs = []
ins_tvecs = []
imu_tvecs = []
step = 1
measured_times = []
for i in range(len(image_times) - step):
    if abs(image_times[i + step] - image_times[i]) > 2:
        continue

    rotv = []
    ts = np.linspace(image_times[i], image_times[i + step] + dt, 100)
    for t in ts:
        rotv_ = ins.get_raw_imu_at_time(t)
        if rotv_ is not None:
            rotv.append(rotv_[1])

    if len(rotv) == 0:
        continue
    rotv = np.mean(rotv, axis=0)*(ts[-1] - ts[0])
    while True:
        mag = np.linalg.norm(rotv)
        if mag > np.pi:
            rotv = rotv - np.pi*rotv/mag
        else:
            break

    imu_tvecs.append(rotv)
    sfm_tvecs.append(get_rvec_btw_times(sfm_cm, image_times[i],
                                        image_times[i + step]))
    ins_tvecs.append(get_rvec_btw_times(ins_cm, image_times[i],
                                        image_times[i + step] + dt))
    measured_times.append(image_times[i])

    if np.linalg.norm(sfm_tvecs[-1]) > 2:
        raise Exception()


sfm_tvecs = np.array(sfm_tvecs).T
ins_tvecs = np.array(ins_tvecs).T
imu_tvecs = np.array(imu_tvecs).T
measured_times = np.array(measured_times)

mag_sfm = np.sqrt(np.sum(sfm_tvecs**2, axis=0))
mag_ins = np.sqrt(np.sum(ins_tvecs**2, axis=0))
mag_imu = np.sqrt(np.sum(imu_tvecs**2, axis=0))

if False:
    plt.plot(mag_sfm)
    plt.plot(mag_ins)
    plt.plot(mag_imu)

ind = np.where(mag_ins > 1e-1)[0][0]
sfm_tvecs = sfm_tvecs[:, ind:]
ins_tvecs = ins_tvecs[:, ind:]
imu_tvecs = imu_tvecs[:, ind:]
measured_times = measured_times[ind:]
mag_sfm = mag_sfm[ind:]
mag_ins = mag_ins[ind:]
mag_imu = mag_imu[ind:]
measured_times = measured_times - measured_times[0]

if False:
    sfm_tvecs = sfm_tvecs/mag_sfm
    ins_tvecs = ins_tvecs/mag_ins

ind = np.logical_and(mag_sfm > np.percentile(mag_sfm, 80),
                     mag_sfm < np.percentile(mag_sfm, 100))
R = horn(imu_tvecs[:, ind], sfm_tvecs[:, ind], fit_scale=False,
         fit_translation=False)[1]

#R = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
#R = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])

plt.figure(num=None, figsize=(15.3, 10.7), dpi=80);
for i in range(3):
    plt.subplot(3, 1, i+1)
    plt.plot(measured_times, sfm_tvecs[i], label='SfM')
    plt.plot(measured_times, sfm_tvecs[i], 'k.')
    #plt.plot(measured_times, np.dot(R, ins_tvecs)[i], label='INS')
    plt.plot(measured_times, np.dot(R, imu_tvecs)[i], label='IMU')
    plt.plot(measured_times, np.dot(R, imu_tvecs)[i], 'k.')
    plt.ylabel('%s-Axis Rotation' % ['X', 'Y', 'Z'][i])
    plt.xlim([2530, 2560])
    plt.legend()

R_ = np.identity(4);    R_[:3, :3] = R
cam_quat = quaternion_from_matrix(R_)
ins_cm.cam_quat = cam_quat

ins_cm.save_to_file('%s/camera_model.yaml' % colmap_dir)
# ----------------------------------------------------------------------------



sfm_poses = []
img_times = []

for image_num in colmap_images:
    image = colmap_images[image_num]

    img_fname = os.path.split(image.name)[1]
    t = float(os.path.splitext(img_fname)[0])/1000000
    img_times.append(t)

    # Query Colmaps pose for the camera.
    R = qvec2rotmat(image.qvec)
    pos = -np.dot(R.T, image.tvec)

    # The qvec used by Colmap is a (w, x, y, z) quaternion
    # representing the rotation of a vector defined in the world
    # coordinate system into the camera coordinate system. However,
    # the 'camera_models' module assumes (x, y, z, w) quaternions
    # representing a coordinate system rotation. Also, the quaternion
    # used by 'camera_models' represents a coordinate system rotation
    # versus the coordinate system transform of Colmap's convention,
    # so we need an inverse.

    #quat = transformations.quaternion_inverse(image.qvec)
    quat = image.qvec / np.linalg.norm(image.qvec)
    quat[0] = -quat[0]

    quat = [quat[1], quat[2], quat[3], quat[0]]

    sfm_pose = np.hstack([pos, quat])

    sfm_poses.append(sfm_pose)

img_times = np.array(img_times)
ind = np.argsort(img_times)
img_times = img_times[ind]
sfm_poses = np.array(sfm_poses)[ind]
ind = np.argsort(sfm_poses[:, 0])
sfm_poses = sfm_poses[ind]

fig = plt.figure()
ax = plt.axes(projection='3d')
plt.plot(sfm_poses[:, 0], sfm_poses[:, 1], sfm_poses[:, 2])

fig = plt.figure()
ax = plt.axes(projection='3d')
plt.plot(sfm_poses[:, 0], sfm_poses[:, 1], sfm_poses[:, 2], 'ro')

fig = plt.figure()
plt.plot(sfm_poses[:, 0], sfm_poses[:, 1], 'k.')

fig = plt.figure()
plt.plot(sfm_poses[:, 0], sfm_poses[:, 3], 'k.')


# Visualize coordinate system.
plt.close('all')
quats = sfm_poses[::20, 4:]
rotation_matrices = [quaternion_matrix(quaternion_inverse(quat))[:3,:3]
                     for quat in quats]
draw_moving_coordinate_system(sfm_poses[::20, 1:4], rotation_matrices,
                              arrow_scale=5)


# ----------------------------------------------------------------------------
# Look at velocity magnitudes

# ----------------------------------------------------------------------------

# ----------------------------------------------------------------------------
if False:
    img_times = sfm_cm.platform_pose_provider.times
    nav_times = ins.times
    img_times = img_times[img_times >= min(nav_times)]
    img_times = img_times[img_times <= max(nav_times)]

    xyz1 = []
    xyz2 = []
    for i in range(len(img_times)):
        t = img_times[i]
        xyz1.append(sfm_cm.platform_pose_provider.pose(t)[0])
        xyz2.append(ins.pose(t)[0])

    xyz1 = np.array(xyz1)
    xyz2 = np.array(xyz2)

    plt.figure(num=None, figsize=(15.3, 10.7), dpi=80);
    plt.subplot(211)
    plt.plot(img_times, xyz1[:, 1], 'ro', label='sfm')
    plt.legend(fontsize=18)
    plt.subplot(212)
    plt.plot(img_times, xyz2[:, 1], 'bo', label='ins')
    plt.legend(fontsize=18)

    # Solve for alignment transform and time offset.
    toffset = np.linspace(-5, 5, 100)


    # Look at velocity magnitudes
    # ------------------------------------------------------------------------
    def get_tform(dt):
        xyz1 = []
        xyz2 = []
        for i in range(len(img_times)):
            t = img_times[i]
            tmp = ins.pose(t + dt, time_thresh=0.1)
            if tmp is None:
                continue

            xyz2.append(tmp[0])
            xyz1.append(sfm_cm.platform_pose_provider.pose(t)[0])

        xyz1 = np.array(xyz1)
        xyz2 = np.array(xyz2)
        s1 = np.sum(np.diff(xyz1, axis=0)**2, axis=1)
        s2 = np.sum(np.diff(xyz2, axis=0)**2, axis=1)
        s1= s1*np.mean(s2)/np.mean(s1)
        err = np.sqrt(np.mean((s2 - s1)**2))

        print(dt, err)
        #plt.plot(s1)
        #plt.plot(s2)
        return xyz1, xyz2, s, R, t, err

    errs = []
    for dt in toffset:
        xyz1, xyz2, s, R, t, err = get_tform(dt)
        errs.append(err)

    ind = np.argmin(errs)
    dt = toffset[ind]
    print('Offset:', dt)

    plt.plot(toffset, errs)

    # ------------------------------------------------------------------------


    def get_tform(dt):
        xyz1 = []
        xyz2 = []
        for i in range(len(img_times)):
            t = img_times[i]
            xyz1.append(sfm_cm.platform_pose_provider.pose(t)[0])
            xyz2.append(ins.pose(t + dt)[0])

        xyz1 = np.array(xyz1)
        xyz2 = np.array(xyz2)
        s, R, t = horn(xyz1.T, xyz2.T)
        xyz1 = np.dot(R, s*xyz1.T).T + t
        err = np.mean((xyz2 - xyz1)**2)
        print(dt, err)
        return xyz1, xyz2, s, R, t, err

    errs = []
    for dt in toffset:
        xyz1, xyz2, s, R, t, err = get_tform(dt)
        errs.append(err)

    plt.plot(toffset, errs)

    ind = np.argmin(errs)
    dt = toffset[ind]
    xyz1, xyz2, s, R, t, err = get_tform(dt)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    plt.plot(xyz1[:, 0], xyz1[:, 1], xyz1[:, 2], 'ro', label='sfm')
    plt.plot(xyz2[:, 0], xyz2[:, 1], xyz2[:, 2], 'bo', label='ins')
    plt.legend(fontsize=18)

    plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
    ind = np.round(np.linspace(0, len(xyz1) - 1, 200)).astype(int)
    plt.plot(xyz1[ind, 0], xyz1[ind, 1], 'ro')
    plt.plot(xyz2[ind, 0], xyz2[ind, 1], 'bo')
    for i in ind:
        plt.plot([xyz1[i, 0], xyz2[i, 0]],
                 [xyz1[i, 1], xyz2[i, 1]], 'g-')

    fig = plt.figure()
    plt.plot(xyz1[:, 2], 'ro', label='sfm')
    plt.plot(xyz2[:, 2], 'bo', label='ins')
    plt.legend(fontsize=18)

    err = np.sqrt(np.sum((xyz1 - xyz2)**2, axis=1))
    fig = plt.figure()
    plt.plot(np.sort(err))


    # Colmap then uses this pairing to solve for a similarity transform to
    # best-match the SfM poses it recovered into these positions. All Colmap
    # coordinates in this aligned version of its reconstruction will then be in
    # easting/northing/up meters coordinates.
    align_fname = '%s/image_locations.txt' % colmap_dir
    with open(align_fname, 'w') as fo:
        for i in colmap_images:
            image = colmap_images[i]
            img_fname = os.path.split(image.name)[1]
            t = float(os.path.splitext(img_fname)[0])/1000000
            d = np.abs(img_times - t)
            ind = np.argmin(d)
            if d[ind] > 1e-5:
                continue

            name = image.name
            pos = xyz1[ind]
            fo.write('%s %0.8f %0.8f %0.8f\n' % (name, pos[0], pos[1], pos[2]))

    print('source model_aligner.sh %s %s %s %s' % (colmap_dir, 'sparse/0',
                                            'image_locations.txt',
                                            'aligned'))
# ----------------------------------------------------------------------------

# ----------------------------------------------------------------------------


# ------------------------------ Display Position ----------------------------



# ----------------------------------------------------------------------------



# ----------------------------------------------------------------------------
# Compare angulare rates.

def compare_angular_rates(sfm_poses, nav_state_provider, dt,
                          plot_results=False):
    sfm_rot_vecs = []
    ins_rot_vecs = []
    imu_rot_vecs = []

    for i in range(len(sfm_poses) - 1):
        # SFM
        t1, t2 = sfm_poses[i:i + 2, 0]
        quat1, quat2 = sfm_poses[i : i + 2, 4:]
        sfm_rot_vecs.append(np.product(rotation_between_quats(quat1, quat2)))

        # Correct for image -> INS time offset.
        t1 = t1 + dt
        t2 = t2 + dt

        # INS
        quat1 = nav_state_provider.quat(t1)
        quat2 = nav_state_provider.quat(t2)
        ins_rot_vecs.append(np.product(rotation_between_quats(quat1, quat2)))

        ts = np.linspace(t1, t2, 100)

        try:
            imu_del = np.sum(angular_vel_v_time(ts), 1)*(ts[1] - ts[0])
        except:
            imu_del = [0, 0, 0]

        imu_rot_vecs.append(imu_del)

    sfm_rot_vecs = np.array(sfm_rot_vecs)
    ins_rot_vecs = np.array(ins_rot_vecs)
    imu_rot_vecs = np.array(imu_rot_vecs)

    mag_sfm = np.sqrt(np.sum(sfm_rot_vecs**2, axis=1))*180/np.pi
    mag_ins = np.sqrt(np.sum(ins_rot_vecs**2, axis=1))*180/np.pi
    mag_imu = np.sqrt(np.sum(imu_rot_vecs**2, axis=1))*180/np.pi

    R = horn(sfm_rot_vecs.T, ins_rot_vecs.T, fit_scale=False,
             fit_translation=False)[1]
    sfm_rot_vecs1 = np.dot(R, sfm_rot_vecs.T).T

    if True:
        # Error from matching vectors
        err = np.mean((ins_rot_vecs - sfm_rot_vecs1)**2)
    else:
        # Error from matching vector magnitudes
        err = np.mean((mag_sfm - mag_ins)**2)

    if plot_results:
        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
        plt.semilogy(mag_sfm, '.',
                 label='SFM Rotation Degrees')
        plt.semilogy(mag_ins,  '.',
                 label='INS Rotation Degrees')
        plt.semilogy(mag_imu,  '.',
                 label='IMU Rotation Degrees')
        plt.xlabel('SfM Index')
        plt.legend(fontsize=16)

    return sfm_rot_vecs, ins_rot_vecs, R, err


# We view the full set of sfm_poses to check that it and the INS results
# appear valid over the whole timespan. It is possible that the INS
# doesn't initialize fully during the initial period.
ret = compare_angular_rates(sfm_poses, nav_state_provider, dt=0,
                            plot_results=True)

# Crop down to the valid region
i1 = 0
i2 = -1
ret = compare_angular_rates(sfm_poses[i1:i2], nav_state_provider, dt=0,
                            plot_results=True)


dts = np.linspace(-1, 1, 100)
errs = []
for dt in dts:
    errs.append(compare_angular_rates(sfm_poses[i1:i2], nav_state_provider, dt)[3])

plt.plot(dts, errs)




R = horn(ins_rot_vecs[:].T, sfm_rot_vecs[:].T, fit_translation=False)


width = 5
err = []
for i in range(len(ins_rot_vecs) - width):
    R = horn(ins_rot_vecs[i:i+width].T, sfm_rot_vecs[i:i+width].T,
             fit_translation=False)
    err.append(np.mean((np.dot(R, ins_rot_vecs[i:i+width].T) - sfm_rot_vecs[i:i+width].T)**2))

plt.close('all')
plt.plot(err)
R = horn(ins_rot_vecs[:].T, sfm_rot_vecs[:].T, fit_translation=False)


sfm_rot_vecs = []
ins_rot_vecs = []
for i in range(len(sfm_poses) - 1):
    # SFM
    t1, t2 = sfm_poses[i:i + 2, 0]
    quat1, quat2 = sfm_poses[i : i + 2, 4:]
    sfm_rot_vecs.append(np.product(rotation_between_quats(quat1, quat2)))

    # INS
    quat1 = nav_state_provider.quat(t1)
    quat2 = nav_state_provider.quat(t2)
    ins_rot_vecs.append(np.product(rotation_between_quats(quat1, quat2)))

# ----------------------------------------------------------------------------


# ----------------------------------------------------------------------------
points_bin_fname = '%s/%s/points3D.bin' % (colmap_dir, sparse_recon_subdir)
points3d = read_points3D_binary(points_bin_fname)

# Load in all of the Colmap results into more-convenient structures.
points_per_image = {}
camera_from_camera_str = {}
for image_num in colmap_images:
    image = colmap_images[image_num]
    camera_str = image.name.split('/')[0]
    camera_from_camera_str[camera_str] = colmap_cameras[image.camera_id]

    xys = image.xys
    pt_ids = image.point3D_ids
    ind = pt_ids != -1
    pt_ids = pt_ids[ind]
    xys = xys[ind]
    xyzs = np.array([points3d[pt_id].xyz for pt_id in pt_ids])
    try:
        img_fname = os.path.split(image.name)[1]
        t = float(os.path.splitext(img_fname)[0])/1000000
        points_per_image[image.name] = (xys, xyzs, t)
    except KeyError:
        pass

# Loop over all images and apply the camera model to project 3-D points
# into the image and compare to the measured versions to calculate
# reprojection error.
err = []
for i in range(len(img_fnames)):
    print('%i/%i' % (i + 1, len(img_fnames)))
    fname = img_fnames[i]
    sfm_pose = sfm_poses[i]
    camera_str = os.path.split(fname)[0]

    colmap_camera = camera_from_camera_str[camera_str]

    if colmap_camera.model == 'OPENCV':
            fx, fy, cx, cy, d1, d2, d3, d4 = colmap_camera.params
    elif colmap_camera.model == 'PINHOLE':
        fx, fy, cx, cy = colmap_camera.params
        d1 = d2 = d3 = d4 = 0

    K = K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
    dist = np.array([d1, d2, d3, d4])

    cm = StandardCamera(colmap_camera.width, colmap_camera.height, K, dist,
                        [0, 0, 0], [0, 0, 0, 1], None, frame_id=None,
                        nav_state_provider=NavStateFixed(*sfm_pose))
    xy, xyz, t = points_per_image[fname]
    err_ = np.sqrt(np.sum((xy - cm.project(xyz.T, t).T)**2, axis=1))
    err = err + err_.tolist()

plt.hist(err, 1000)
# ----------------------------------------------------------------------------