from pathlib import Path
import argparse
import cv2
import matplotlib.cm as cm
import torch
import numpy as np
from utils.nnmatching import NNMatching
from utils.misc import (AverageTimer, VideoStreamer, make_matching_plot_fast, frame2tensor)

torch.set_grad_enabled(False)


def compute_essential(matched_kp1, matched_kp2, K):
    pts1 = cv2.undistortPoints(matched_kp1,cameraMatrix=K, distCoeffs = (-0.117918271740560,0.075246403574314,0,0))
    pts2 = cv2.undistortPoints(matched_kp2,cameraMatrix=K, distCoeffs = (-0.117918271740560,0.075246403574314,0,0))
    K_1 = np.eye(3)
    # Estimate the homography between the matches using RANSAC
    ransac_model, ransac_inliers = cv2.findEssentialMat(pts1, pts2, K_1, method=cv2.RANSAC, prob=0.999, threshold=0.001, maxIters=10000)
    if ransac_inliers is None or ransac_model.shape != (3,3):
        ransac_inliers = np.array([])
        ransac_model = None
    return ransac_model, ransac_inliers, pts1, pts2


sizer = (960, 640)
focallength_x = 4.504986436499113e+03/(6744/sizer[0])
focallength_y = 4.513311442889859e+03/(4502/sizer[1])
K = np.eye(3)
K[0,0] = focallength_x
K[1,1] = focallength_y
K[0,2] = 3.363322177533149e+03/(6744/sizer[0])# * 0.5
K[1,2] = 2.291824660547715e+03/(4502/sizer[1])# * 0.5


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='DarkFeat demo',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        '--input', type=str,
        help='path to an image directory')
    parser.add_argument(
        '--output_dir', type=str, default=None,
        help='Directory where to write output frames (If None, no output)')

    parser.add_argument(
        '--image_glob', type=str, nargs='+', default=['*.ARW'],
        help='Glob if a directory of images is specified')
    parser.add_argument(
        '--resize', type=int, nargs='+', default=[640, 480],
        help='Resize the input image before running inference. If two numbers, '
             'resize to the exact dimensions, if one number, resize the max '
             'dimension, if -1, do not resize')
    parser.add_argument(
        '--force_cpu', action='store_true',
        help='Force pytorch to run in CPU mode.')
    parser.add_argument('--model_path', type=str,
                        help='Path to the pretrained model')

    opt = parser.parse_args()
    print(opt)

    assert len(opt.resize) == 2
    print('Will resize to {}x{} (WxH)'.format(opt.resize[0], opt.resize[1]))

    device = 'cuda' if torch.cuda.is_available() and not opt.force_cpu else 'cpu'
    print('Running inference on device \"{}\"'.format(device))
    matching = NNMatching(opt.model_path).eval().to(device)
    keys = ['keypoints', 'scores', 'descriptors']

    vs = VideoStreamer(opt.input, opt.resize, opt.image_glob)
    frame, ret = vs.next_frame()
    assert ret, 'Error when reading the first frame (try different --input?)'

    frame_tensor = frame2tensor(frame, device)
    last_data = matching.darkfeat({'image': frame_tensor})
    last_data = {k+'0': [last_data[k]] for k in keys}
    last_data['image0'] = frame_tensor
    last_frame = frame
    last_image_id = 0

    if opt.output_dir is not None:
        print('==> Will write outputs to {}'.format(opt.output_dir))
        Path(opt.output_dir).mkdir(exist_ok=True)

    timer = AverageTimer()

    while True:
        frame, ret = vs.next_frame()
        if not ret:
            print('Finished demo_darkfeat.py')
            break
        timer.update('data')
        stem0, stem1 = last_image_id, vs.i - 1

        frame_tensor = frame2tensor(frame, device)
        pred = matching({**last_data, 'image1': frame_tensor})
        kpts0 = last_data['keypoints0'][0].cpu().numpy()
        kpts1 = pred['keypoints1'][0].cpu().numpy()
        matches = pred['matches0'][0].cpu().numpy()
        confidence = pred['matching_scores0'][0].cpu().numpy()
        timer.update('forward')

        valid = matches > -1
        mkpts0 = kpts0[valid]
        mkpts1 = kpts1[matches[valid]]

        E, inliers, pts1, pts2 = compute_essential(mkpts0, mkpts1, K)
        color = cm.jet(np.clip(confidence[valid][inliers[:, 0].astype('bool')] * 2 - 1, -1, 1))

        text = [
            'DarkFeat',
            'Matches: {}'.format(inliers.sum())
        ]

        out = make_matching_plot_fast(
            last_frame, frame, mkpts0[inliers[:, 0].astype('bool')], mkpts1[inliers[:, 0].astype('bool')], color, text,
            path=None, small_text=' ')

        if opt.output_dir is not None:
            stem = 'matches_{:06}_{:06}'.format(stem0, stem1)
            out_file = str(Path(opt.output_dir, stem + '.png'))
            print('Writing image to {}'.format(out_file))
            cv2.imwrite(out_file, out)