import os
import sys

import numpy as np

"""
Adapted from the homework assignment from Computer Vision I by Stefan Roth
Teaching assistants Shweta Mahajan and Krishnakant Singh
From the solution of group 22 (Merve Bektas, Danail Iordanov and Nick Steyer)
"""

def ransac_iters(p, k, z):
    """ Computes the required number of iterations for RANSAC.

    Args:
        p: probability that any given correspondence is valid
        k: number of pairs
        z: total probability of success after all iterations

    Returns:
        minimum number of required iterations
    """

    def log(base, argument):
        """ Computes the logarithm to a given base for a given argument

        :param base: The base of the logarithm
        :param argument: The argument of the logarithm
        :return:    The logarithm
        """

        return np.log(argument) / np.log(base)

    """
    Reference: Lecture 8 Slide 23
    Probability that we have no outlier in an iteration: p ** k
    Probability that we have an outlier in an iteration: 1 - p ** k
    Probability that we have an outlier in each iteration: (1 - p ** k) ** number_of_iterations
    (1 - p ** k) ** number_of_iterations should be smaller than or equal to 1 - z
    So number_of_iterations must be at least the logarithm of (1 - z) to the base (1 - p ** k)
    We always round up to the nearest integer using numpy.ceil because we can only execute full loops.
    """
    return int(np.ceil(log(1 - p ** k, 1 - z)))


def ransac(pairs, n_iters, k, threshold):
    """ RANSAC algorithm.

    Args:
        pairs: (l, 4) numpy array containing matched keypoint pairs
        n_iters: number of ransac iterations
        threshold: inlier detection threshold

    Returns:
        H: (3, 3) numpy array, best homography observed during RANSAC
        max_inliers: number of inliers N
        inliers: (N, 4) numpy array containing the coordinates of the inliers
        :param pairs:
        :param n_iters:
        :param threshold:
        :param k:
    """

    h = {}
    max_inliers = 0
    inliers = {}

    """
    We need to split pairs in two arrays with two columns so that we can pass it to pick_samples.
    """
    p1, p2 = np.hsplit(pairs, 2)

    """
    We execute the ransac loop n_iters times, so that we have a good chance to have a valid homography.
    """
    for iteration in range(n_iters):
        """
        First, we pick a sample of k corresponding point pairs
        """
        sample1, sample2 = pick_samples(p1, p2, k)

        """
        The points in the samples then have to be conditioned, so their values are between -1 and 1.
        We also retrieve the condition matrices from this function.
        """
        sample1_conditioned, t_1 = condition_points(sample1)
        sample2_conditioned, t_2 = condition_points(sample2)

        """
        Using the conditioned coordinates, we calculate the homographies both with respect to the conditioned points
        and the unconditioned points.
        """
        h_current, hc = compute_homography(sample1_conditioned, sample2_conditioned, t_1, t_2)

        """
        We then use the unconditioned points (as clarified in the office hours) to compute the homography distances.
        """
        homography_distances = compute_homography_distance(h_current, p1, p2)

        """
        Using the find_inliers function, we retrieve the inliers for our given homography as well as the number 
        of inliers we currently have.
        """
        number_inliers_current, inliers_current = find_inliers(pairs, homography_distances, threshold)

        """
        If the number of inliers is higher than for a previously computed homography (i.e. our current homography is
        better), we save our current homography as well as the number of inliers and the inliers themselves.
        """
        if number_inliers_current > max_inliers:
            h = h_current
            max_inliers = number_inliers_current
            inliers = inliers_current

        report_progress(round((iteration + 1) / n_iters * 100))

    """
    We then return the best homography we found as well as the number of inliers found with it and the inliers 
    themselves.
    """
    return h, max_inliers, inliers


def recompute_homography(inliers):
    """ Recomputes the homography matrix based on all inliers.

    Args:
        inliers: (N, 4) numpy array containing coordinate pairs of the inlier points

    Returns:
        H: (3, 3) numpy array, recomputed homography matrix
    """

    """
    We split the inliers matrix into two arrays with 2 columns each (one array for the points in the left image 
    and one for the points in the right image).
    """
    p1, p2 = np.hsplit(inliers, 2)

    """
    We then condition these points.
    """
    p1_conditioned, t_1 = condition_points(p1)
    p2_conditioned, t_2 = condition_points(p2)

    """
    Using the conditioned points, we recompute the homography. The new homography should give better results than
    the previously computed one, since we only use inliers (correct correspondences) here instead of all 
    correspondences including the wrong ones.
    """
    h, hc = compute_homography(p1_conditioned, p2_conditioned, t_1, t_2)

    return h


def pick_samples(p1, p2, k):
    """ Randomly select k corresponding point pairs.

    Args:
        p1: (n, 2) numpy array, given points in first image
        p2: (m, 2) numpy array, given points in second image
        k:  number of pairs to select

    Returns:
        sample1: (k, 2) numpy array, selected k pairs in left image
        sample2: (k, 2) numpy array, selected k pairs in right image
    """

    """
    We assume that m = n (as confirmed in the office hours)
    """

    """
    Reference: https://numpy.org/doc/stable/reference/random/index.html
    We calculate the number of points n (= m) and generate k unique random integers in [0; n) (random_numbers).
    We then use the random numbers as indices to select points from p1 and the corresponding points from p2.
    """
    n = p1.shape[0]
    generator = np.random.default_rng()
    random_numbers = generator.choice(n, size=k, replace=False)
    # random_numbers = np.array([96, 93, 63, 118])
    sample1 = p1[random_numbers]
    sample2 = p2[random_numbers]

    return sample1, sample2


def condition_points(points):
    """ Conditioning: Normalization of coordinates for numeric stability
    by subtracting the mean and dividing by half of the component-wise
    maximum absolute value.
    Further, turns coordinates into homogeneous coordinates.
    Args:
        points: (l, 2) numpy array containing not normalized cartesian coordinates.

    Returns:
        ps: (l, 3) numpy array containing normalized points in homogeneous coordinates.
        T: (3, 3) numpy array, transformation matrix for conditioning
    """

    """
    First, we calculate half of the component-wise maximum absolute value for all points, as discussed in the 
    office hours.
    """
    s = np.max(np.abs(points), axis=0) / 2
    sx = s[0]
    sy = s[1]

    """
    We then compute the component-wise mean of all points.
    """
    t = np.mean(points, axis=0)
    tx = t[0]
    ty = t[1]

    """
    We then calculate the transformation matrix T like on slide 18 in Lecture 8.
    We use the x components of s and t in the first row and the y components of s and t in the second row.
    """
    t = np.array(([1 / sx, 0, - tx / sx],
                  [0, 1 / sy, - ty / sy],
                  [0, 0, 1]))

    """
    We now iterate over each point and transform it into homogeneous coordinates.
    Then we transform the point using T and append the result to the list ps.
    Finally, the list ps is converted to a numpy array and returned together with the transformation matrix T.
    """
    ps = []
    for point in points:
        homogeneous_point = np.append(point, 1)
        homogeneous_transformed_point = np.matmul(t, homogeneous_point)
        ps.append(homogeneous_transformed_point)

    return np.array(ps), t


def compute_homography(p1, p2, t1, t2):
    """ Estimate homography matrix from point correspondences of conditioned coordinates.
    Both returned matrices should be normalized so that the bottom right value equals 1.
    You may use np.linalg.svd for this function.

    Args:
        p1: (l, 3) numpy array, the conditioned homogeneous coordinates of interest points in img1
        p2: (l, 3) numpy array, the conditioned homogeneous coordinates of interest points in img2
        t1: (3,3) numpy array, conditioning matrix for p1
        t2: (3,3) numpy array, conditioning matrix for p2

    Returns:
        H: (3, 3) numpy array, homography matrix with respect to unconditioned coordinates
        HC: (3, 3) numpy array, homography matrix with respect to the conditioned coordinates
    """

    """
    Reference: Lecture 8 Slide 9, 13 ff.
    We did something similar in Assignment 1
    """

    """
    We calculate the number of rows in p1 (which must be the same as number of rows in p2).
    """
    n = p1.shape[0]
    # Create a (2*N, 9) matrix filled with zeros
    a = np.zeros((2 * n, 9))

    # We iterate from 0 to N (N is the number of correspondences)
    for i in range(0, n):
        # The point from the first image (divided by the last component so that x and y have the non-homogeneous
        # value. We could cut off the last component (which is now 1), but that step is not necessary here.)
        point1 = p1[i] / p1[i][2]

        # The corresponding point from the second image (also divided by the last component to get the
        # non-homogeneous values for x and y)
        point2 = p2[i] / p2[i][2]

        # We always calculate two rows of the matrix A in each loop iteration
        # Here, we calculate the first row according to the lecture slides Lecture 8 slide 14
        first_row = np.concatenate((np.zeros(3), point1, - point2[1] * point1))

        # Calculate the second row as in Lecture 8 slide 14
        second_row = np.concatenate((- point1, np.zeros(3), point2[0] * point1))

        # Calculate the row indices at matrix A, where to store the two calculated rows
        # When we iterate from 0 to N, we get indices 0, 2, 4, 6... etc. for the first row index
        first_row_index = i * 2

        # The second row index is simply the first row index plus one, so the rows are below each other
        second_row_index = first_row_index + 1

        """
        Finally, we save the values in the corresponding rows in matrix A
        """
        a[first_row_index] = first_row
        a[second_row_index] = second_row

    """
    Now we perform a single value decomposition on A and retrieve the last right singular vector
    """
    u, s, vt = np.linalg.svd(a)
    h = vt[-1]

    """
    Since we used conditioned points, we retrieve the homography matrix HC regarding the conditioned coordinates 
    by reshaping the last right singular vector .
    """
    hc = np.reshape(h, (3, 3))

    """
    The matrix is then normalized so that the bottom right value equals 1.
    """
    hc = hc / hc[2, 2]

    """
    We calculate the homography matrix with respect to the unconditioned coordinates by using the formula from 
    lecture 8 on slide 18.
    """
    h = np.matmul(np.linalg.inv(t2), np.matmul(hc, t1))

    """
    The homography matrix is also normalized.
    """
    h = h / h[2, 2]

    return h, hc


def compute_homography_distance(h, p1, p2):
    """ Computes the pairwise symmetric homography distance.

    Args:
        h: (3, 3) numpy array, homography matrix
        p1: (l, 2) numpy array, interest points in img1
        p2: (l, 2) numpy array, interest points in img2

    Returns:
        dist: (l, ) numpy array containing the distances
    """

    """
    Determine the number of points in p1 (should be the same as in p2).
    """
    number_of_points = p1.shape[0]

    """
    Calculate the inverse of the homography H.
    """
    h_inverse = np.linalg.inv(h)

    """
    Create a zero-filled array which we later fill with the distances.
    """
    distances = np.zeros(number_of_points)

    """
    We calculate the transformed points for the given homography.
    """
    p1_transformed = transform_pts(np.array(p1), h)
    p2_transformed = transform_pts(np.array(p2), h_inverse)

    """
    We calculate the symmetric squared distances using the formula in the assignment sheet for every point pair.
    """
    for index in range(0, number_of_points):
        x1 = p1[index]
        x2 = p2[index]
        x1_transformed = p1_transformed[index]
        x2_transformed = p2_transformed[index]
        distances[index] = np.linalg.norm(x1_transformed - x2) ** 2 + np.linalg.norm(x1 - x2_transformed) ** 2

    return distances


def find_inliers(pairs, dist, threshold):
    """ Return and count inliers based on the homography distance.

    Args:
        pairs: (l, 4) numpy array containing keypoint pairs
        dist: (l, ) numpy array, homography distances for k points
        threshold: inlier detection threshold

    Returns:
        N: number of inliers
        inliers: (n, 4)
    """

    inliers_list = []

    """
    We iterate over each pair and use its index in the array and its value.
    We then use the index to retrieve the corresponding distance for the pair in dist. If the distance is smaller
    than or equal to the given threshold, we add the pair to the inliers_list.  
    """
    for pair_index, pair_value in enumerate(pairs):
        if dist[pair_index] <= threshold:
            inliers_list.append(pair_value)

    """
    We then get the number of inliers as the length of the list. 
    Furthermore, we convert the inliers_list to a numpy array.
    Both values are then returned.
    """
    n = len(inliers_list)
    inliers = np.array(inliers_list)

    return n, inliers


def transform_pts(p, h):
    """ Transform p through the homography matrix H.

    Args:
        p: (l, 2) numpy array, interest points
        h: (3, 3) numpy array, homography matrix

    Returns:
        points: (l, 2) numpy array, transformed points
    """

    """
    Reference: Lecture 8 Slide 8
    """

    points_list = []

    """
    Iterate over every point in p.
    For each point, we first append a 1 to transform the point into homogeneous coordinates.
    Then, we perform a matrix multiplication of the point with the homography H, as shown on slide 8 in lecture 8.
    The transformed point is then converted back into non-homogeneous coordinates by dividing it by the last 
    value and then cutting off the last value (which is then 1 again).
    The transformed points are returned as a numpy array.
    """
    for point_index, point_value in enumerate(p):
        homogeneous_p = np.append(point_value, 1)
        homogeneous_transformed_p = np.matmul(h, homogeneous_p)
        normalized_homogeneous_transformed_p = homogeneous_transformed_p / homogeneous_transformed_p[2]
        transformed_p = normalized_homogeneous_transformed_p[0:2]
        points_list.append(transformed_p)

    return np.array(points_list)


def report_progress(percentage):
    sys.stdout.write('\r')
    sys.stdout.write("[%-100s] %d%%" % ('=' * percentage, percentage))
    sys.stdout.flush()


def main():
    threshold = 0.02  # inlier threshold
    p = 0.35  # probability that any given correspondence is valid
    z = 0.99  # total probability of success after all iterations

    if len(sys.argv) == 3:
        threshold = float(sys.argv[2])
    elif len(sys.argv) == 4:
        threshold = float(sys.argv[2])
        p = float(sys.argv[3])
    elif len(sys.argv) == 5:
        threshold = float(sys.argv[2])
        p = float(sys.argv[3])
        z = float(sys.argv[4])
    elif len(sys.argv) != 2:
        print('Usage: hgcalc <pairs_file_path> [threshold [valid_correspondence_probability [total_probability]]]')
        return

    calculate_and_save(sys.argv[1], threshold, p, z)


def calculate_and_save(pairs_file_path, threshold, p, z):
    homography_file_path = '.\\homography.csv'
    pairs = np.loadtxt(pairs_file_path)
    k = 4  # number of samples drawn at each iteration
    n_iters = ransac_iters(p, k, z)
    print('Input file:', pairs_file_path)
    output_path = os.path.abspath(homography_file_path)
    print('Output file:', output_path)
    print('Threshold: %s, Inlier probability: %s, Success probability: %s' % (threshold, p, z))
    print('Calculating %s iterations' % n_iters)
    h, num_inliers, inliers = ransac(pairs, n_iters, k, threshold)
    print()
    print('Percentage of inliers:', "{0:.0%}".format(num_inliers / len(pairs)))

    # recompute homography matrix based on inliers
    h = recompute_homography(inliers)
    np.savetxt(homography_file_path, h)
    print('Homography saved at', output_path)


if __name__ == "__main__":
    main()