import numpy as np
import cv2
from BRIEF import briefLite, briefMatch
from scipy.spatial import distance

def computeH(p1, p2):
    '''
    INPUTS
        p1 and p2 - Each are size (2 x N) matrices of corresponding (x, y)'  
                coordinates between two images
    OUTPUTS
        H2to1 - a 3 x 3 matrix encoding the homography that best matches the linear 
                equation
    '''
    
    assert(p1.shape[1] == p2.shape[1])
    assert(p1.shape[0] == 2)
    #############################
    N = p2.shape[1]
    p2 = np.concatenate((p2, np.ones((1, N))), axis=0)

    # A size = 2N * 9
    A = np.zeros((N * 2, 9))
    for i in range(N):
        A[i * 2, 0: 3] = 0
        A[i * 2, 3: 6] = -1 * p2[:, i]
        A[i * 2, 6: 9] = p1[1, i] * p2[:, i]
        A[i * 2 + 1, 0: 3] = 1 * p2[:, i]
        A[i * 2 + 1, 3: 6] = 0
        A[i * 2 + 1, 6: 9] = -p1[0, i] * p2[:, i]

    u, s, v = np.linalg.svd(A, full_matrices=True)

    # Reshape h (9x1) back to H (3x3)
    H2to1 = np.reshape(v[-1, :], (3, 3))

    return H2to1


def ransacH(matches, locs1, locs2, num_iter=5000, tol=2):
    '''
    Returns the best homography by computing the best set of matches using RANSAC
    
    INPUTS
        locs1 and locs2 - matrices specifying point locations in each of the images
        matches         - matrix specifying matches between these two sets of point locations
        nIter           - number of iterations to run RANSAC
        tol             - tolerance value for considering a point to be an inlier

    OUTPUTS
        bestH - homography matrix with the most inliers found during RANSAC
    '''

    ###########################
    max_match_num = 0
    for g in range(num_iter):
        matches = matches.copy()

        # Randomly grab 4 matches
        # Same as computing homography, 8 unknowns require 8 points which is 4 match pairs
        # So this is essentially "randomly" grabbing a set of homography out of all the match pairs
        initial_index = np.random.randint(0, matches.shape[0], 4, dtype='int')

        # Separate the points into 2 matrices
        p1 = np.zeros((2, 4))
        p2 = np.zeros((2, 4))
        for i, pt in enumerate(initial_index):
            p1[0, i] = locs1[matches[pt, 0], 0]
            p1[1, i] = locs1[matches[pt, 0], 1]
            p2[0, i] = locs2[matches[pt, 1], 0]
            p2[1, i] = locs2[matches[pt, 1], 1]

        # Calculate homography using these 4 pairs
        H_curr = computeH(p1, p2)
        accepted_pairs = 0


        p1_acceptable = []
        p2_acceptable = []

        for i in range(matches.shape[0]):
            # homography is the relation of two views
            # The calculated homography "supposedly" will tell us the coordinates of a point in view1 if given the
            # coordinates of the point in view2, with the relationship:
            # λx = Hu, λ is the scale
            # λ * point_view1 = H * point_view2

            # Using the current set of homography calculated, we want to see if all match point pairs has a similar
            # enough relationship as this current homography

            # point1_BRIEFmatch and point2 are matches from the inaccurate "BRIEF match"
            # Theoretically, H * point2 / λ should be mapped to point1
            # we have to turn point1_BRIEFmatch into homogeneous coordinates
            point1_BRIEFmatch = np.array([locs1[matches[i, 0],0], locs1[matches[i, 0],1], 1])
            # point2 = [p2x, p2y, 1], because how homography is set up,
            # we have to turn point2 into homogeneous coordinates
            point2 = np.array([locs2[matches[i, 1],0], locs2[matches[i, 1],1], 1])

            # point_view1 = H * point_view2 / λ
            # the x and y calculated are the point coordinates in view1 from this homography
            # (homography that is generated from 4 random pairs in all the matches)
            x = np.dot(H_curr[0],point2) / np.dot(H_curr[2], point2) # np.dot(H_curr[2],point2) is the scale
            y = np.dot(H_curr[1],point2) / np.dot(H_curr[2], point2)
            point1_homography = np.array([x, y, 1])

            # calculate the distance between the coordinates from our current homography,
            # and the match point from briefMatch/briefLite
            dist = distance.euclidean(point1_BRIEFmatch, point1_homography)

            # if the distance between the briefMatch point and the homography output point is
            # within our acceptable tolerance, keep this match point pair
            if dist <= tol:
                accepted_pairs = accepted_pairs + 1
                p1_acceptable.append(locs1[matches[i, 0]][0:2])
                p2_acceptable.append(locs2[matches[i, 1]][0:2])

        p1_acceptable = np.transpose(p1_acceptable)
        p2_acceptable = np.transpose(p2_acceptable)

        # After trying with num_iter homographies generated by different match pairs,
        # save the homography that has the most acceptable point pairs
        if accepted_pairs > max_match_num:
            max_match_num = accepted_pairs
            bestH = H_curr
    
    return bestH


if __name__ == '__main__':
    im1 = cv2.imread('../data/model_chickenbroth.jpg')
    im2 = cv2.imread('../data/chickenbroth_01.jpg')
    locs1, desc1 = briefLite(im1)
    locs2, desc2 = briefLite(im2)
    matches = briefMatch(desc1, desc2)
    ransacH(matches, locs1, locs2, num_iter=5000, tol=2)