#!/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.
import colmap_processing.camera_models as camera_models
from postflight_scripts.ins import ANINS
from postflight_scripts.utilities import draw_keypoints
from colmap_processing.image_renderer import save_gif


bridge = CvBridge()

# Bag information
nav_bag_fname = '/home/user/adapt_ws/data/2022_jan_flight/nav_data.bag'
out_dir = os.path.splitext(nav_bag_fname)[0]

camera_model_fname = '/home/user/adapt_ws/data/2022_jan_flight/camera_model.yaml'

num_features = 10000
# ----------------------------------------------------------------------------

ins = ANINS.from_ros_bag(nav_bag_fname)
cm = camera_models.load_from_file(camera_model_fname, ins)


def load_image(fname):
    img = cv2.imread(fname)[:, :, ::-1]
    fname = os.path.split(fname)[1]
    t = float(os.path.splitext(fname)[0])/1000000
    return img, t


fname1 = '/host_filesystem/mnt/data2tb/libraries/adapt/adapt_ros_ws/data/2022_jan_flight/images/2022-01-24-11-40-16/1643057096439870.jpg'
fname2 = '/host_filesystem/mnt/data2tb/libraries/adapt/adapt_ros_ws/data/2022_jan_flight/images/2022-01-24-11-40-16/1643057096584634.jpg'
fname1 = '/host_filesystem/mnt/data2tb/libraries/adapt/adapt_ros_ws/data/2022_jan_flight/images/2022-01-24-11-40-16/1643057239756352.jpg'
fname2 = '/host_filesystem/mnt/data2tb/libraries/adapt/adapt_ros_ws/data/2022_jan_flight/images/2022-01-24-11-40-16/1643057239901116.jpg'

img1, t1 = load_image(fname1)
img2, t2 = load_image(fname2)

gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

# Find the keypoints and descriptors and match them.
orb = cv2.ORB_create(nfeatures=num_features, edgeThreshold=21,
                     patchSize=31, nlevels=16,
                     scoreType=cv2.ORB_FAST_SCORE, fastThreshold=10)

kp1, des1 = orb.detectAndCompute(gray1, None)
kp2, des2 = orb.detectAndCompute(gray2, None)
pts1 = np.float32([kp_.pt for kp_ in kp1]).T
pts2 = np.float32([kp_.pt for kp_ in kp2]).T

P = cm.get_camera_pose(t1)
R = P[:, :3]
tvec = P[:, 3]
cam_pos = np.dot(-R.T, tvec)

# ----------------------------------------------------------------------------
# Visualize that we can dead reckon.
ray_pos, ray_dir = cm.unproject(pts1, t1)
pts2_ = cm.project(ray_pos + ray_dir*30, t2)

plt.figure(num=None, figsize=(15.3, 10.7), dpi=80);
plt.imshow(img1)
plt.plot(pts1[0], pts1[1], 'ro')
plt.figure(num=None, figsize=(15.3, 10.7), dpi=80);
plt.imshow(img2)
plt.plot(pts2_[0], pts2_[1], 'ro')
plt.xlim([0, img2.shape[1]])
plt.ylim([img2.shape[0], 0])


out1 = draw_keypoints(gray1, pts1.T, radius=10, copy=True)
out2 = draw_keypoints(gray2, pts2_.T, radius=10, copy=True)


fname = '/host_filesystem/mnt/data2tb/libraries/adapt/adapt_ros_ws/data/2022_jan_flight/multiframe.gif'
save_gif([out1, out2], fname)
# ----------------------------------------------------------------------------


kp1 = cv2.goodFeaturesToTrack(gray1, 1000, qualityLevel=0.1, minDistance=10)
kp1 = np.squeeze(kp1).T

fast = cv2.FastFeatureDetector_create(threshold=15, nonmaxSuppression=True)
kp1 = fast.detect(gray1, None)
pts1 = np.float32([kp_.pt for kp_ in kp1]).T
plt.figure(num=None, figsize=(15.3, 10.7), dpi=80);
plt.imshow(draw_keypoints(gray1, pts1.T, radius=10, copy=True))

bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
matches = bf.match(des1, des2)
