from __future__ import print_function
from __future__ import division

import numpy as np
import cv2 as cv
import csv

import time

import argparse

alpha_slider_max = 100
title_window = 'Image registration'


def sift_algorithm(img_object_path, img_scene_path, dir_, method_name, nfeatures, nOctaveLayers, contrastThreshold, edgeThreshold, sigma):
    start_time = time.time()

    MIN_MATCH_COUNT = 10 # minimal number of matches

    img_object = cv.imread(img_object_path,0) # queryImage
    img_scene = cv.imread(img_scene_path,0) # trainImage

    if img_object is None or img_scene is None:
        print('Could not open or find the images!')
        exit(0)

    # Initiate SIFT detector
    sift = cv.xfeatures2d_SIFT.create(nfeatures, nOctaveLayers, contrastThreshold, edgeThreshold, sigma)

    # find the keypoints and descriptors with SIFT
    kp1, des1 = sift.detectAndCompute(img_object,None)
    kp2, des2 = sift.detectAndCompute(img_scene,None)

    # BFMatcher with default params
    bf = cv.BFMatcher()
    matches = bf.knnMatch(des1, des2, k=2)

    good = []
    threshold = 0.5

    while len(good) < 200:
        for m, n in matches:
            if m.distance < threshold * n.distance:
                good.append(m)
        threshold = threshold + 0.025



    if len(good) > MIN_MATCH_COUNT:

        src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)

        distance_a = abs(src_pts[0, 0, 0] - dst_pts[0, 0, 0])
        distance_b = abs(src_pts[0, 0, 1] - dst_pts[0, 0, 1])

        M, mask = cv.findHomography(src_pts, dst_pts, cv.RANSAC,5.0)
        matchesMask = mask.ravel().tolist()

        # draw borders of the base image
        h,w = img_object.shape
        pts = np.float32([[0, 0], [0, h-1], [w-1, h-1], [w-1, 0]]).reshape(-1, 1, 2)
        dst = cv.perspectiveTransform(pts, M)

        img_scene = cv.polylines(img_scene, [np.int32(dst)], True, 255, 3, cv.LINE_AA)

        # REGISTRATION
        width = img_object.shape[1]
        height = img_object.shape[0]

        transformed_img1 = cv.warpPerspective(img_object, M, (width, height))

        reg_img1 = np.zeros(transformed_img1.shape)

        for i in range(transformed_img1.shape[0]):
            for j in range(transformed_img1.shape[1]):

                if transformed_img1[i, j] == 0:
                    reg_img1[i, j] = transformed_img1[i, j] + img_scene[i, j]

                elif transformed_img1[i, j] != 0:
                    reg_img1[i, j] = transformed_img1[i, j]

        result_reg_img1 = img_object_path[:-5] + "reg_" + method_name + ".tif"
        cv.imwrite(result_reg_img1, reg_img1)




    else:
        print("Not enough matches are found - %d/%d" % (len(good), MIN_MATCH_COUNT))
        matchesMask = None


        src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)

        distance_a = abs(src_pts[0, 0, 0] - dst_pts[0, 0, 0])
        distance_b = abs(src_pts[0, 0, 1] - dst_pts[0, 0, 1])

        if len(good) == 0:
            src_pts = np.float32([0, 0]).reshape(-1, 1, 2)
            dst_pts = np.float32([0, 0]).reshape(-1, 1, 2)
            distance_a = abs(src_pts[0, 0, 0] - dst_pts[0, 0, 0])
            distance_b = abs(src_pts[0, 0, 1] - dst_pts[0, 0, 1])

        dst = np.float32([0,0]).reshape(-1, 1, 2)

        # REGISTRATION
        width = img_object.shape[1]
        height = img_object.shape[0]

        transformed_img2 = cv.warpPerspective(img_object, M, (width, height))

        reg_img2 = np.zeros(transformed_img2.shape)

        for i in range(transformed_img2.shape[0]):
            for j in range(transformed_img2.shape[1]):

                if transformed_img2[i, j] == 0:
                    reg_img2[i, j] = transformed_img2[i, j] + img_scene[i, j]

                elif transformed_img2[i, j] != 0:
                    reg_img2[i, j] = transformed_img2[i, j]

        result_reg_img2 = img_object_path[:-5] + "reg_" + method_name + ".tif"
        cv.imwrite(result_reg_img2, reg_img2)





    draw_params = dict(matchColor=(255, 255, 0),  # draw matches in green color
                       singlePointColor=None,
                       matchesMask=matchesMask,  # draw only inliers
                       flags=2)

    img_matches = cv.drawMatches(img_object, kp1, img_scene, kp2, good, None, **draw_params)

    result_img_path = img_object_path[:-5] + "result" + method_name + ".tif"

    cv.imwrite(result_img_path, img_matches)


    time_ex = time.time() - start_time


    with open(str(dir_) + '_results' + method_name + '.csv', 'a', newline='') as f:
        thewriter = csv.writer(f)
        thewriter.writerow([img_object_path[-8:-4], len(matches), len(good), threshold, time_ex, distance_a, distance_b,
                            nfeatures, nOctaveLayers, contrastThreshold, edgeThreshold, sigma])



    # TRACKBAR
    def on_trackbar(val):
        alpha = val / alpha_slider_max
        beta = (1.0 - alpha)
        dst = cv.addWeighted(src1, alpha, src2, beta, 0.0)
        cv.imshow(title_window, dst)

    parser = argparse.ArgumentParser(description='Code for Adding a Trackbar to our applications tutorial.')
    args = parser.parse_args()

    src1 = img_scene
    src2 = transformed_img1

    if src1 is None:
        print('Could not open or find the image: ', args.input1)
        exit(0)
    if src2 is None:
        print('Could not open or find the image: ', args.input2)
        exit(0)


    toolbar_decision = input("Do you want to show image registration by SIFT algorithm? <y/n>: ")
    if toolbar_decision == 'y':

        cv.namedWindow(title_window)
        trackbar_name = 'Alpha x %d' % alpha_slider_max
        cv.createTrackbar(trackbar_name, title_window, 0, alpha_slider_max, on_trackbar)

        # Show some stuff
        on_trackbar(0)
        # Wait until user press some key
        print('Press some key to close the window.')
        cv.waitKey()
        cv.destroyWindow(title_window)

    else:
        # cv.DestroyAllWindows()
        print(' ')


def surf_algorithm(img1_path, img2_path, dir_, method_name, hessianThreshold, nOctaves, nOctaveLayers, extended, upright):
    start_time = time.time()

    MIN_MATCH_COUNT = 10  # minimal number of matches

    img_object = cv.imread(img1_path, 0)
    img_scene = cv.imread(img2_path, 0)


    # Detect the keypoints using SURF Detector
    surf = cv.xfeatures2d_SURF.create(hessianThreshold, nOctaves, nOctaveLayers, extended, upright)
    kp1, des1 = surf.detectAndCompute(img_object, None)
    kp2, des2 = surf.detectAndCompute(img_scene, None)


    # Matching descriptor vectors with a FLANN based matcher
    # Since SURF is a floating-point descriptor NORM_L2 is used
    matcher = cv.DescriptorMatcher_create(cv.DescriptorMatcher_FLANNBASED)
    knn_matches = matcher.knnMatch(des1, des2, 2)

    # Filter matches using the Lowe's ratio test
    good_matches = []
    threshold = 0.5

    while len(good_matches) < 100:
        for m, n in knn_matches:
            if m.distance < threshold * n.distance:
                good_matches.append(m)
        threshold = threshold + 0.025

    if len(good_matches)>MIN_MATCH_COUNT:
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1,1,2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1,1,2)

        distance_a = abs(src_pts[0, 0, 0] - dst_pts[0, 0, 0])
        distance_b = abs(src_pts[0, 0, 1] - dst_pts[0, 0, 1])

        M, mask = cv.findHomography(src_pts, dst_pts, cv.RANSAC,5.0)
        matchesMask = mask.ravel().tolist()


        # Draw borders of the base image
        h,w = img_object.shape
        pts = np.float32([[0, 0], [0, h-1], [w-1, h-1], [w-1, 0]]).reshape(-1, 1, 2)
        dst = cv.perspectiveTransform(pts, M)

        img_scene = cv.polylines(img_scene, [np.int32(dst)], True, 255, 3, cv.LINE_AA)

        # REGISTRATION
        width = img_object.shape[1]
        height = img_object.shape[0]

        transformed_img1 = cv.warpPerspective(img_object, M, (width, height))

        reg_img1 = np.zeros(transformed_img1.shape)

        for i in range(transformed_img1.shape[0]):
            for j in range(transformed_img1.shape[1]):

                if transformed_img1[i, j] == 0:
                    reg_img1[i, j] = transformed_img1[i, j] + img_scene[i, j]

                elif transformed_img1[i, j] != 0:
                    reg_img1[i, j] = transformed_img1[i, j]

        result_reg_img1 = img1_path[:-5] + "reg_" + method_name + ".tif"
        cv.imwrite(result_reg_img1, reg_img1)




    else:
        print("Not enough matches are found - %d/%d" % (len(good_matches), MIN_MATCH_COUNT))
        matchesMask = None

        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)

        distance_a = abs(src_pts[0, 0, 0] - dst_pts[0, 0, 0])
        distance_b = abs(src_pts[0, 0, 1] - dst_pts[0, 0, 1])


        if len(good_matches) == 0:
            src_pts = np.float32([0, 0]).reshape(-1, 1, 2)
            dst_pts = np.float32([0, 0]).reshape(-1, 1, 2)

            distance_a = abs(src_pts[0, 0, 0] - dst_pts[0, 0, 0])
            distance_b = abs(src_pts[0, 0, 1] - dst_pts[0, 0, 1])

        dst = np.float32([0, 0]).reshape(-1, 1, 2)

        # REGISTRATION
        width = img_object.shape[1]
        height = img_object.shape[0]

        transformed_img2 = cv.warpPerspective(img_object, M, (width, height))

        reg_img2 = np.zeros(transformed_img2.shape)

        for i in range(transformed_img2.shape[0]):
            for j in range(transformed_img2.shape[1]):

                if transformed_img2[i, j] == 0:
                    reg_img2[i, j] = transformed_img2[i, j] + img_scene[i, j]

                elif transformed_img2[i, j] != 0:
                    reg_img2[i, j] = transformed_img2[i, j]

        result_reg_img2 = img1_path[:-5] + "reg_" + method_name + ".tif"
        cv.imwrite(result_reg_img2, reg_img2)



    draw_params = dict(matchColor=(0, 255, 0),  # draw matches in green color
                       singlePointColor=None,
                       matchesMask=matchesMask,  # draw only inliers
                       flags=2)

    # Draw matches
    img_matches = np.empty((max(img_object.shape[0], img_scene.shape[0]), img_object.shape[1]+img_scene.shape[1], 3), dtype=np.uint8)
    # cv.drawMatches(img_object, keypoints1, img_scene, keypoints2, good_matches, img_matches, flags=cv.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
    cv.drawMatches(img_object, kp1, img_scene, kp2, good_matches, img_matches, **draw_params)

    # Show detected matches
    result_img_path = img1_path[:-5] + "result" + method_name + ".tif"
    cv.imwrite(result_img_path, img_matches)

    time_ex = time.time() - start_time

    # Make a csv file with results

    with open(str(dir_) + '_results' + method_name + '.csv', 'a', newline='') as f:
        thewriter = csv.writer(f)
        thewriter.writerow([img1_path[-8:-4], len(knn_matches), len(good_matches), threshold, time_ex, distance_a, distance_b,
                            hessianThreshold, nOctaves, nOctaveLayers, extended, upright])

    # TRACKBAR

    def on_trackbar(val):
        alpha = val / alpha_slider_max
        beta = (1.0 - alpha)
        dst = cv.addWeighted(src1, alpha, src2, beta, 0.0)
        cv.imshow(title_window, dst)

    parser = argparse.ArgumentParser(description='Code for Adding a Trackbar to our applications tutorial.')
    args = parser.parse_args()

    src1 = img_scene
    src2 = transformed_img1

    if src1 is None:
        print('Could not open or find the image: ', args.input1)
        exit(0)
    if src2 is None:
        print('Could not open or find the image: ', args.input2)
        exit(0)

    toolbar_decision = input("Do you want to show image registration by SURF algorithm? <y/n>: ")
    if toolbar_decision == 'y':

        cv.namedWindow(title_window)
        trackbar_name = 'Alpha x %d' % alpha_slider_max
        cv.createTrackbar(trackbar_name, title_window, 0, alpha_slider_max, on_trackbar)

        # Show some stuff
        on_trackbar(0)
        # Wait until user press some key
        print('Press some key to close the window.')
        cv.waitKey()
        cv.destroyWindow(title_window)

    else:
        # cv.DestroyAllWindows()
        print(' ')




def akaze_algorithm(img1_path,img2_path, dir_, method_name, descriptor_type, descriptor_size, descriptor_channels, threshold_akaze,
                    nOctaves, nOctaveLayers, diffusivity):

    start_time = time.time()

    MIN_MATCH_COUNT = 10  # minimal number of matches

    img_object = cv.imread(img1_path, 0) # queryImage
    img_scene = cv.imread(img2_path, 0)  # trainImage


    if img_object is None or img_scene is None:
        print('Could not open or find the images!')
        exit(0)

    # Detect the keypoints using AKAZE Detector, compute the descriptors
    akaze = cv.AKAZE_create(descriptor_type, descriptor_size, descriptor_channels, threshold_akaze, nOctaves, nOctaveLayers, diffusivity)
    kp1, des1 = akaze.detectAndCompute(img_object, None)
    kp2, des2 = akaze.detectAndCompute(img_scene, None)


    # Matching descriptor vectors with a BF based matcher

    bf = cv.DescriptorMatcher_create(cv.DescriptorMatcher_BRUTEFORCE_HAMMING)
    # using Hamm.dist, AKAZE is using binary descriptors by default.


    matches = bf.knnMatch(des1, des2, k=2)

    good_matches = []
    threshold = 0.5

    while len(good_matches) < 200:
        for m, n in matches:
            if m.distance < threshold * n.distance:
                good_matches.append(m)
        threshold = threshold + 0.025


    if len(good_matches)>MIN_MATCH_COUNT:
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1,1,2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1,1,2)

        distance_a = abs(src_pts[0, 0, 0] - dst_pts[0, 0, 0])
        distance_b = abs(src_pts[0, 0, 1] - dst_pts[0, 0, 1])

        M, mask = cv.findHomography(src_pts, dst_pts, cv.RANSAC,5.0)
        matchesMask = mask.ravel().tolist()


        # draw borders of the base image
        h,w = img_object.shape
        pts = np.float32([[0, 0], [0, h-1], [w-1, h-1], [w-1, 0]]).reshape(-1, 1, 2)
        dst = cv.perspectiveTransform(pts, M)

        img_scene = cv.polylines(img_scene, [np.int32(dst)], True, 255, 3, cv.LINE_AA)

        # REGISTRATION
        width = img_object.shape[1]
        height = img_object.shape[0]

        transformed_img1 = cv.warpPerspective(img_object, M, (width, height))

        reg_img1 = np.zeros(transformed_img1.shape)

        for i in range(transformed_img1.shape[0]):
            for j in range(transformed_img1.shape[1]):

                if transformed_img1[i, j] == 0:
                    reg_img1[i, j] = transformed_img1[i, j] + img_scene[i, j]

                elif transformed_img1[i, j] != 0:
                    reg_img1[i, j] = transformed_img1[i, j]

        result_reg_img1 = img1_path[:-5] + "reg_" + method_name + ".tif"
        cv.imwrite(result_reg_img1, reg_img1)


    else:
        print("Not enough matches are found - %d/%d" % (len(good_matches), MIN_MATCH_COUNT))
        matchesMask = None


        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)

        distance_a = abs(src_pts[0, 0, 0] - dst_pts[0, 0, 0])
        distance_b = abs(src_pts[0, 0, 1] - dst_pts[0, 0, 1])

        if len(good_matches) == 0:
            src_pts = np.float32([0, 0]).reshape(-1, 1, 2)
            dst_pts = np.float32([0, 0]).reshape(-1, 1, 2)

            distance_a = abs(src_pts[0, 0, 0] - dst_pts[0, 0, 0])
            distance_b = abs(src_pts[0, 0, 1] - dst_pts[0, 0, 1])

        dst = np.float32([0, 0]).reshape(-1, 1, 2)

        # REGISTRATION
        width = img_object.shape[1]
        height = img_object.shape[0]

        transformed_img2 = cv.warpPerspective(img_object, M, (width, height))

        reg_img2 = np.zeros(transformed_img2.shape)

        for i in range(transformed_img2.shape[0]):
            for j in range(transformed_img2.shape[1]):

                if transformed_img2[i, j] == 0:
                    reg_img2[i, j] = transformed_img2[i, j] + img_scene[i, j]

                elif transformed_img2[i, j] != 0:
                    reg_img2[i, j] = transformed_img2[i, j]

        result_reg_img2 = img1_path[:-5] + "reg_" + method_name + ".tif"
        cv.imwrite(result_reg_img2, reg_img2)

    draw_params = dict(matchColor=(0, 0, 255),  # draw matches in green color
                       singlePointColor=None,
                       matchesMask=matchesMask,  # draw only inliers
                       flags=2)

    # Draw matches
    img_matches = np.empty((max(img_object.shape[0], img_scene.shape[0]), img_object.shape[1]+img_scene.shape[1], 3), dtype=np.uint8)
    cv.drawMatches(img_object, kp1, img_scene, kp2, good_matches, img_matches, **draw_params)

    result_img_path = img1_path[:-5] + "result" + method_name + ".tif"
    cv.imwrite(result_img_path, img_matches)

    time_ex = time.time() - start_time

    # Make a csv file with results
    with open(str(dir_) + '_results' + method_name + '.csv', 'a', newline='') as f:
        thewriter = csv.writer(f)
        thewriter.writerow([img1_path[-8:-4], len(matches), len(good_matches), threshold, time_ex, distance_a, distance_b,
                            descriptor_type, descriptor_size, descriptor_channels, threshold_akaze, nOctaves, nOctaveLayers, diffusivity])


    # TRACKBAR
    def on_trackbar(val):
        alpha = val / alpha_slider_max
        beta = (1.0 - alpha)
        dst = cv.addWeighted(src1, alpha, src2, beta, 0.0)
        cv.imshow(title_window, dst)

    parser = argparse.ArgumentParser(description='Code for Adding a Trackbar to our applications tutorial.')
    args = parser.parse_args()

    src1 = img_scene
    src2 = transformed_img1

    if src1 is None:
        print('Could not open or find the image: ', args.input1)
        exit(0)
    if src2 is None:
        print('Could not open or find the image: ', args.input2)
        exit(0)


    toolbar_decision = input("Do you want to show image registration by AKAZE algorithm? <y/n>: ")
    if toolbar_decision == 'y':

        cv.namedWindow(title_window)
        trackbar_name = 'Alpha x %d' % alpha_slider_max
        cv.createTrackbar(trackbar_name, title_window, 0, alpha_slider_max, on_trackbar)

        # Show some stuff
        on_trackbar(0)
        # Wait until user press some key
        print('Press some key to close the window.')
        cv.waitKey()
        cv.destroyWindow(title_window)

    else:
        # cv.DestroyAllWindows()
        print(' ')