from __future__ import division, print_function
import os
import glob
import numpy as np
import cv2
import matplotlib.pyplot as plt
import pycolmap

# colmap_processing imports.
import colmap_processing.camera_models as camera_models
from colmap_processing.platform_pose import PlatformPoseInterp
from colmap_processing.image_renderer import save_gif
from colmap_processing.slam import draw_keypoints


# Bag information
odometry_txt = '/mnt/homenas2/noaa_adapt/adapt_flights/2021-11-04/odometry.txt'
image_glob = '/mnt/homenas2/noaa_adapt/adapt_flights/2021-11-04/images/2021-11-04-13-25-46/*.jpg'
camera_model_fname = '/mnt/homenas2/noaa_adapt/adapt_flights/2022_jan_flight/camera_model.yaml'

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

ins = PlatformPoseInterp.from_odometry_llh_txt(odometry_txt)
cm = camera_models.load_from_file(camera_model_fname, ins)
image_fnames = sorted(glob.glob(image_glob))


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


img1, t1 = load_image(image_fnames[100])
img2, t2 = load_image(image_fnames[101])
img3, t3 = load_image(image_fnames[102])

plt.figure();   plt.imshow(img1);   plt.figure();   plt.imshow(img2)

gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
gray3 = cv2.cvtColor(img3, 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)
kp3, des3 = orb.detectAndCompute(gray3, None)

bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
matches12 = bf.match(des1, des2)
matches23 = bf.match(des2, des3)


[kp1[m.queryIdx].pt for m in matches]

[kp1[m.queryIdx].pt for m in matches]

pts1 = np.float32([kp1[m.queryIdx].pt for m in matches12])
pts2 = np.float32([kp2[m.trainIdx].pt for m in matches12])

answer = pycolmap.homography_matrix_estimation(pts1, pts2)
answer['inliers']



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)
