import numpy as np
import cv2 as cv
from scipy import ndimage, signal
from scipy.io import loadmat
import matplotlib.pyplot as plt
import random
import math

from ransac import *  # the code in ransac.py was not written by me. Source info in the file


def diff_x(img):
    # From a2
    return signal.convolve2d(img, np.array([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]]) * 1 / 3, mode='same')


def diff_y(img):
    # From a2
    return signal.convolve2d(img, np.array([[1, 1, 1], [0, 0, 0], [-1, -1, -1]]) * 1 / 3, mode='same')


def gauss_pyramid(img, num_levels=7, orig_size=1024):
    # From a2
    """ Generate the images that create the Gaussian pyramid"""
    ret = []
    for level in range(0, num_levels):
        sigma = 2 ** level
        size = orig_size // sigma
        ret.append(cv.resize(ndimage.gaussian_filter(img, sigma=sigma), (size, size)))
    return ret


def dog_pyramid(img, num_levels=6, orig_size=1024):
    # From a2
    """ Generate the difference of Gaussian pyramid """
    ret = []
    gauss = gauss_pyramid(img, num_levels=num_levels + 1, orig_size=orig_size)
    for level in range(0, num_levels):
        size = orig_size // 2 ** level
        ret.append(gauss[level].astype(float) - cv.resize(gauss[level + 1], (size, size)).astype(float))
    return ret


def detect_keypoint(img, num_levels=5, orig_size=1024, thresh=-29):
    # From a2
    """ Find keypoint locations on an image using DoG (pass in greyscale image)
        Returns a dictionary of (i,j): level """
    tracker = {}
    dog = dog_pyramid(img, num_levels=num_levels, orig_size=orig_size)
    dog_local_min = np.zeros((orig_size, orig_size), dtype=int)

    for scale in range(0, num_levels):  # Find keypoints at different scales
        size = orig_size // 2 ** scale
        ratio = orig_size / size
        for level in range(0, scale + 1):
            curr = cv.resize(dog[level], (size, size))  # resize the image to current size
            for i in range(2, size - 3):
                for j in range(2, size - 3):
                    val = curr[i, j]
                    if val < thresh \
                            and val == np.amin(curr[i - 1:i + 2, j - 1:j + 2]) \
                            and val < dog_local_min[i, j]:  # Find local minima, compare to local minima at each level
                        a = np.hstack((curr[i - 2:i + 3, j - 2], curr[i - 2:i + 3, j + 2], curr[i - 2, j - 1:j + 2],
                                       curr[i + 2, j - 1:j + 2]))
                        if np.sum(a) > (-1 * thresh):
                            dog_local_min[i, j] = val
                            remap_i, remap_j = (int(i * ratio), int(j * ratio))
                            if np.count_nonzero(np.hstack((img[remap_i - 2:remap_i + 3, remap_j - 2],
                                                           img[remap_i - 2:remap_i + 3, remap_j + 2],
                                                           img[remap_i - 2, remap_j - 1:remap_j + 2],
                                                           img[remap_i + 2, remap_j - 1:remap_j + 2]))) > 12:
                                # do not include points that are surrounded by pure black pixels
                                # this is likely to be an artificial border (e.g. generated by rotation)
                                tracker[(remap_i, remap_j)] = level
    return tracker


def gauss_weights(sigma, size=16):
    gauss_1d = signal.gaussian(size, std=sigma)
    window = np.array([gauss_1d * i for i in gauss_1d])
    return window / np.sum(window)  # re-weight to sum to 1


def keypoint_local(img, thresh=None):
    """ Calculate information listed in Q1
        weights is the Gaussian weights to be used
        Returns dictionary of (i, j, sigma): [magnitude, orientation, weighted magnitude] where each item in the list
        is a 16x16 numpy array """
    # Find magnitude, orientation for each pixel in each level of Gaussian pyramid
    points = detect_keypoint(img) if thresh is None else detect_keypoint(img, thresh=thresh)
    local_vals = {}
    weights = [gauss_weights(5) for ro in range(1, 6)]
    gauss = gauss_pyramid(img, num_levels=5)
    diff_levels = [(diff_x(gauss_img), diff_y(gauss_img)) for gauss_img in gauss]
    magnitude = [np.sqrt(i_x ** 2 + i_y ** 2) for i_x, i_y in diff_levels]
    orientation = [np.arctan2(i_x, i_y) for i_x, i_y in diff_levels]
    # Find 1, 2, 3 for each 16x16 area around each pixel
    for i, j in points:
        level = points[(i, j)]
        size = magnitude[level].shape[0]
        sigma = 2 ** level
        vals = [np.zeros((16, 16), dtype=float), np.zeros((16, 16), dtype=float)]
        start_row, end_row = max(0, i - 8), min(size, i + 8)
        start_col, end_col = max(0, j - 8), min(size, j + 8)
        for row in range(start_row, end_row):
            for col in range(start_col, end_col):
                offset_row = 8 - (i - row)
                offset_col = 8 - (j - col)
                vals[0][offset_row, offset_col] = magnitude[level][row, col]
                vals[1][offset_row, offset_col] = orientation[level][row, col]
        vals.append(vals[0] * weights[level])
        local_vals[(i, j, sigma)] = vals
    return local_vals


def sift(img, thresh=None):
    """ Returns the SIFT feature vector for Q2 """
    local_vals = keypoint_local(img, thresh=thresh)
    sift_list = []

    for point in local_vals:
        # Convert rad to degree in orientation
        mag, orient, weighted_mag = local_vals[point]
        orient = (orient * 180 / np.pi) + 180
        # Iterate over orientation and weighted magnitude and sort magnitude into bins
        result = np.zeros(36, dtype=float)
        for i in range(16):
            for j in range(16):
                hist_bin = int(orient[i, j] // 10)
                index = hist_bin if hist_bin != 36 else 0
                result[index] += weighted_mag[i, j]
        if np.sum(result != 0):
            peak = np.argmax(result)
            result = np.roll(result, -peak)  # Cyclic shift counterclockwise so that peak is first
            temp = list(point)
            temp.extend(result)
            sift_list.append(temp)

    return sift_list


def sift_normal(img):
    """ Take a list of sift feature vectors and normalize each vector
        Returns a tuple of (labels, normalized vectors)
        labels are a list of tuples (i, j, sigma) """
    sift_result = sift(img)
    labels = [vector[0:3] for vector in sift_result]
    normals = [vector[3:39] / np.sum(vector[3:39]) for vector in sift_result]
    return labels, normals


def sift_match(img1, img2, ctr=None, theta=None, s=None, thresh=0.95, search_sz=100, debug=False):
    """ Perform matching between img and a transformed img using Bhattacharyya coefficient
        To limit matching to a certain area, specify ctr, theta, and s """
    k1, v1 = sift_normal(img1)
    k2, v2 = sift_normal(img2)
    matches = find_matches(img1, img2, k1, k2, v1, v2, ctr=ctr, theta=theta, s=s, search_sz=search_sz, thresh=thresh,
                           debug=debug)

    return matches


def find_matches(img1, img2, k1, k2, v1, v2, ctr=None, theta=None, s=None, search_sz=100, thresh=0.95, debug=False):
    if debug:
        img1 = cv.cvtColor(img1, cv.COLOR_GRAY2RGB)
        img2 = cv.cvtColor(img2, cv.COLOR_GRAY2RGB)

    matches = {}
    for v in range(len(k1)):
        new_pnt = transform_point((k1[v][1], k1[v][0]), ctr, theta, s)  # points in k1 are listed as (i, j, sigma)
        if new_pnt[0] > img2.shape[1] or new_pnt[1] > img2.shape[0]:
            continue
        if ctr is not None:
            valid_k2, valid_v2 = [], []
            for w in range(len(k2)):  # find sift vectors within the search box
                vector = k2[w]
                if new_pnt[0] - search_sz < vector[1] < new_pnt[0] + search_sz \
                        and new_pnt[1] - search_sz < vector[0] < new_pnt[1] + search_sz:
                    valid_k2.append(vector)
                    valid_v2.append(v2[w])
        else:
            valid_k2, valid_v2 = k2, v2
        bc = np.sum(np.sqrt(v1[v] * valid_v2), axis=1) if valid_v2 else []
        max_index = np.argmax(bc) if len(bc) > 0 else -1
        if max_index != -1 and bc[max_index] > thresh:
            key = [k1[v][1], k1[v][0], k1[v][2]]
            max_val = valid_k2[max_index]
            value = [max_val[1], max_val[0], max_val[2]]
            matches[tuple(key)] = tuple(value)

            if debug:
                fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(20, 11))
                ax[0].set_title("Image 1")
                ax[0].imshow(img1)
                ax[1].set_title("Image 2")
                ax[1].imshow(img2)
                ax[0].scatter(k1[v][1], k1[v][0], c='r')
                ax[1].scatter(new_pnt[0] - search_sz, new_pnt[1] - search_sz, c='b')
                ax[1].scatter(new_pnt[0] + search_sz, new_pnt[1] - search_sz, c='b')
                ax[1].scatter(new_pnt[0] - search_sz, new_pnt[1] + search_sz, c='b')
                ax[1].scatter(new_pnt[0] + search_sz, new_pnt[1] + search_sz, c='b')
                for p in valid_k2:
                    ax[1].scatter(p[1], p[0], c='y')
                ax[1].scatter(new_pnt[0], new_pnt[1], c='r')
                ax[1].scatter(value[0], value[1], c='g')

                plt.show()
    return matches


def orb_match(img1, img2):
    """ A SIFT matching algorithm I found online for use in Part II
        Source: https://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_feature2d/py_matcher/py_matcher.html"""

    # Initiate SIFT detector
    orb = cv.ORB_create()

    # find the keypoints and descriptors with SIFT
    kp1, des1 = orb.detectAndCompute(img1, None)  # queryImage
    kp2, des2 = orb.detectAndCompute(img2, None)  # trainImage

    # create BFMatcher object
    bf = cv.BFMatcher(cv.NORM_HAMMING, crossCheck=True)

    # Match descriptors.
    matches = bf.match(des1, des2)

    # Sort them in the order of their distance.
    return kp1, kp2, sorted(matches, key=lambda x: x.distance)


def format_orb_match(kp1, kp2, matches):
    """ Take matches generated by ORB and format to fit data structure used by HomographyModel """
    match_list = np.array([[kp1[match.queryIdx].pt, kp2[match.trainIdx].pt] for match in matches])
    return match_list


class HomographyModel:
    """ For use in the ransac() function. Fits the points to a homography matrix.
        Written after looking at LinearLeastSquaresModel in ransac.py """

    def __init__(self, input_columns, output_columns, debug=False):
        self.input_columns = input_columns
        self.output_columns = output_columns
        self.debug = debug

    def fit(self, data):
        """ Return a list of the values in the homography matrix computed using data """
        m_list1 = np.vstack([data[:, i] for i in self.input_columns])
        m_list2 = np.vstack([data[:, i] for i in self.output_columns])
        matrix = []
        for i in range(len(m_list1)):
            x, y = m_list1[i]
            x_p, y_p = m_list2[i]
            matrix.append([x, y, 1, 0, 0, 0, -x_p * x, -x_p * y, -x_p])
            matrix.append([0, 0, 0, x, y, 1, -y_p * x, -y_p * y, -y_p])
        a = np.array(matrix)
        eigenValues, eigenVectors = np.linalg.eig(np.matmul(a.T, a))
        h = eigenVectors[:, np.argmin(eigenValues)]  # eigenvector with smallest eigenvalue
        return h

    def get_error(self, data, h):
        """ Given a homography matrix model, compute the SSD for the transformed point and the actual point """
        m_list1 = np.vstack([data[:, i] for i in self.input_columns])
        m_list2 = np.vstack([data[:, i] for i in self.output_columns])
        error = []
        for i in range(len(m_list1)):
            x_orig, y_orig = m_list1[i]
            x_actual, y_actual = m_list2[i]
            original_vector = np.array([x_orig, y_orig, 1])
            x_predicted = np.sum(h[0:3] * original_vector) / np.sum(h[6:9] * original_vector)
            y_predicted = np.sum(h[3:6] * original_vector) / np.sum(h[6:9] * original_vector)
            error.append((x_actual - x_predicted) ** 2 + (y_actual - y_predicted) ** 2)
        return np.array(error)


def transform_img(img, x0, y0, theta, s):
    """ As described in Q3, take an image an rotate by theta, scale by s
        theta is in degrees, positive for rotate clockwise and negative for counterclockwise """
    rot_matrix = cv.getRotationMatrix2D((x0, y0), -theta, s)
    shape = (img.shape[1], img.shape[0])
    return cv.warpAffine(img, rot_matrix, shape)


def transform_point(pnt, ctr, theta, s):
    """ Transform pnt by rotating by theta around ctr and scaling by s
        pnt and ctr are given as (x, y)
        Same as transform_img, positive theta rotates theta degrees clockwise
        Transformed result returned as (x, y) """
    rad = np.deg2rad(theta)
    rot_matrix = [[np.cos(rad), -np.sin(rad)], [np.sin(rad), np.cos(rad)]]
    vector = np.array(pnt)  # convert to numpy array of [x, y]
    centre = np.array(ctr)
    origin_vector = vector - centre  # set centre as the origin
    new_vector = np.matmul(rot_matrix, (origin_vector * s))  # scale and rotate
    retrans_vector = new_vector + centre  # unset centre as origin
    return tuple(retrans_vector)


def homography_img(img, h, output_shape=None, align_left=True):
    """ Transform the img to a new img using the homography matrix """
    h_matrix = np.resize(h / np.sum(h), (3, 3))
    top_left = homography_pnt((0, 0, 1), h)
    top_right = homography_pnt((img.shape[1], 0, 1), h)
    bottom_left = homography_pnt((0, img.shape[0], 1), h)
    bottom_right = homography_pnt((img.shape[1], img.shape[0], 1), h)

    x_shift = -min(0, top_left[0], bottom_left[0])  # shift right if x min is negative
    y_shift = -min(0, top_left[1], top_right[1])  # shift up if y min is negative

    x_max = math.ceil(max(top_right[0], bottom_right[0]) + x_shift)
    y_max = math.ceil(max(bottom_left[1], bottom_right[1]) + y_shift)
    shape = (x_max, y_max) if output_shape is None else output_shape
    if not align_left:
        width = max(top_right[0] - top_left[0], bottom_right[0] - bottom_left[0])
        if width < output_shape[0]:
            x_shift += shape[0] - width

    translation_matrix = np.array([[1, 0, x_shift],
                                   [0, 1, y_shift],
                                   [0, 0, 1]])

    return cv.warpPerspective(img, np.matmul(translation_matrix, h_matrix), shape), x_shift, y_shift


def homography_pnt(vector, h):
    """ Return the new point (x', y') given (x, y, 1) and h """
    return np.sum(h[0:3] * vector) / np.sum(h[6:9] * vector), np.sum(h[3:6] * vector) / np.sum(h[6:9] * vector)


def merge_left(img1, img2, h, matches_lst, x_shift_cum=0, y_shift_cum=0):
    """ Merge images according to their given homography, by transforming the left image """
    img1, img1_x_shift, img1_y_shift = homography_img(img1, h, (2000, 1200))
    p1, p2 = matches_lst[0]
    h_x, h_y = homography_pnt((p1[0] + x_shift_cum, p1[1] + y_shift_cum, 1), h)
    h_x += img1_x_shift
    h_y += img1_y_shift
    img2_x_shift = int(h_x - p2[0])
    img2_y_shift = int(h_y - p2[1])

    img1_slice = img1[img2_y_shift:img2_y_shift + img2.shape[0], img2_x_shift:img2_x_shift + img2.shape[1], :]
    img1_slice[img1_slice < 3] = \
        img2[0:img1_slice.shape[0], 0:img1_slice.shape[1], 0:img1_slice.shape[2]][img1_slice < 3]

    # Show img1 for debugging purposes
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(20, 11))
    ax.imshow(img1)
    plt.show()

    return img1, img2_x_shift, img2_y_shift


def merge_right(img1, img2, h, matches_lst, x_shift_cum=0, y_shift_cum=0):
    """ Merge images according to their given homography, by transforming the right image """
    h_inv = np.reshape(np.linalg.inv(np.reshape(h, (3, 3))), (9,))
    h_inv = h_inv / np.sum(h_inv)
    img2, img2_x_shift, img2_y_shift = homography_img(img2, h_inv, output_shape=(2000, 800), align_left=False)
    p1, p2 = matches_lst[0]
    h_x, h_y = homography_pnt((p2[0] + x_shift_cum, p2[1] + y_shift_cum, 1), h_inv)
    h_x += img2_x_shift
    h_y += img2_y_shift
    img1_x_shift = int(h_x - p1[0])
    img1_y_shift = int(h_y - p1[1])

    img2[img1_y_shift:img1_y_shift + img1.shape[0], img1_x_shift:img1_x_shift + img1.shape[1], :] = img1
    return img2, img1_x_shift, img1_y_shift


def panorama(img_lst_rgb, img_lst_bw):
    """ Create a panorama """
    # combine left and right match

    # Find keypoints using orb

    img_lst_bw_cpy = img_lst_bw[0:6] + [np.flip(img_lst_bw[i], 1) for i in range(len(img_lst_bw) - 1, 4, -1)]
    img_lst_rgb_cpy = img_lst_rgb[0:5] + [np.flip(img_lst_rgb[i], 1) for i in range(len(img_lst_rgb) - 1, 4, -1)]

    # Do not match image 6 and image 10
    match_lst = [orb_match(img_lst_bw_cpy[i], img_lst_bw_cpy[i + 1]) for i in range(len(img_lst_bw_cpy) - 1) if i != 5]
    results_lst = [format_orb_match(kp_1, kp_2, matches_lst) for kp_1, kp_2, matches_lst in match_lst]

    # Find inliers using RANSAC
    debug = False
    n_inputs = 1
    n_outputs = 1
    input_columns = range(n_inputs)  # the first columns of the array
    output_columns = [n_inputs + i for i in range(n_outputs)]  # the last columns of the array
    model = HomographyModel(input_columns, output_columns, debug=debug)
    ransac_lst = [ransac(r, model, 4, 2000, 4000, 50, debug=debug, return_all=True) for r in results_lst]
    h_lst = [ransac_val[0] for ransac_val in ransac_lst]
    ransac_match_lst = []
    for i in range(len(ransac_lst)):
        idx_lst = ransac_lst[i][1]
        ransac_match_lst.append([results_lst[i][idx] for idx in idx_lst])

    left_match, x_shift_cum_left, y_shift_cum_left = merge_left(img_lst_rgb_cpy[0], img_lst_rgb_cpy[1],
                                                                h_lst[0], ransac_match_lst[0])
    for i in range(1, 4):
        curr_result = merge_left(left_match, img_lst_rgb_cpy[i + 1], h_lst[i], ransac_match_lst[i],
                                 x_shift_cum=x_shift_cum_left, y_shift_cum=y_shift_cum_left)
        left_match = curr_result[0]
        x_shift_cum_left = curr_result[1]
        y_shift_cum_left = curr_result[2]

    right_match, x_shift_cum_right, y_shift_cum_right = merge_left(img_lst_rgb_cpy[5], img_lst_rgb_cpy[6],
                                                                   h_lst[5], ransac_match_lst[5])
    for i in range(6, 9):
        curr_result = merge_left(right_match, img_lst_rgb_cpy[i + 1], h_lst[i], ransac_match_lst[i],
                                 x_shift_cum=x_shift_cum_right, y_shift_cum=y_shift_cum_right)
        right_match = curr_result[0]
        x_shift_cum_right = curr_result[1]
        y_shift_cum_right = curr_result[2]

    last_nonzero_idx = len(np.nonzero(np.sum(np.sum(right_match, axis=2), axis=0))[0])
    right_match = np.flip(right_match[:, :last_nonzero_idx, :], 1)
    y_shift = np.nonzero(right_match[:, 0, 0])[0][0]

    final_match = merge_left(left_match, right_match, h_lst[4], ransac_match_lst[4], x_shift_cum=x_shift_cum_left,
                             y_shift_cum=y_shift)[0]

    return final_match


def convert_1024(img):
    """ Convert an img to 1024x1024 while keeping aspect ratio """
    ratio = 1024 / max(img.shape[0:2])
    height = img.shape[0]
    width = img.shape[1]
    img_new = cv.resize(img, (int(width * ratio), int(height * ratio)))

    # insert horizontal or vertical padding
    if img.shape[0] < img.shape[1]:  # more columns than rows (wider)
        buf_len_top = math.ceil((1024 - img_new.shape[0]) / 2)
        buf_len_bot = math.floor((1024 - img_new.shape[0]) / 2)
        buffer_top = np.zeros((buf_len_top, 1024, 3), dtype=np.uint8) if len(img.shape) == 3 \
            else np.zeros((buf_len_top, 1024), dtype=np.uint8)
        buffer_bot = np.zeros((buf_len_bot, 1024, 3), dtype=np.uint8) if len(img.shape) == 3 \
            else np.zeros((buf_len_bot, 1024), dtype=np.uint8)
        img_new = np.concatenate((buffer_top, img_new, buffer_bot), axis=0)
    elif img.shape[0] > img.shape[1]:  # more rows than columns (taller)
        buf_len_left = math.ceil((1024 - img_new.shape[1]) / 2)
        buf_len_right = math.floor((1024 - img_new.shape[1]) / 2)
        buffer_left = np.zeros((1024, buf_len_left, 3), dtype=np.uint8) if len(img.shape) == 3 \
            else np.zeros((1024, buf_len_left), dtype=np.uint8)
        buffer_right = np.zeros((1024, buf_len_right, 3), dtype=np.uint8) if len(img.shape) == 3 \
            else np.zeros((1024, buf_len_right), dtype=np.uint8)
        img_new = np.concatenate((buffer_left, img_new, buffer_right), axis=1)

    return img_new


def plot_points(name, img, points):
    """ Plot points on img, where points is an iterable (i, j): level """
    # blue (level 1), green (level 2), yellow (level 3), magenta (level 4), red (level 5)
    colours = ['b', 'g', 'y', 'm', 'r']
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 10))
    ax.set_title("{}".format(name))
    ax.imshow(img)
    for point in points:
        ax.scatter([point[1]], [point[0]], c=colours[points[point]])
    plt.show()


def plot_patch(p, values):
    """ Plot the results of Q1 for a specified point
        Point is a tuple (i, j, sigma), values is the three 16x16 grids """
    # Portion of the code used to plot numbers is from:
    # https://matplotlib.org/3.1.1/gallery/images_contours_and_fields/image_annotated_heatmap.html
    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(30, 11))
    fig.suptitle("16x16 grid for ({}, {}), sigma={}".format(p[1], p[0], p[2]))
    ax[0].set_title("Magnitude")
    ax[1].set_title("Orientation")
    ax[2].set_title("Weighted Magnitude")

    for grid in range(3):
        ax[grid].imshow(values[grid])
        for i in range(16):
            for j in range(16):
                ax[grid].text(j, i, round(values[grid][i, j], 2), ha="center", va="center", color="w")
    plt.show()


def plot_histogram(vector):
    """ Given a feature vector as described in Q2, plot the vector as a histogram """
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 10))
    ax.set_title("Degree Histogram for ({}, {}), sigma={}".format(vector[1], vector[0], vector[2]))
    x = np.arange(36)
    vals = vector[3:39]
    plt.bar(x, vals)
    plt.show()


def plot_transform(img, arg1, arg2):
    """ Plot the original image and 2 transformed (rotated and scaled) images
        arg1, arg2 = [x0, y0, theta, s] """
    t1 = transform_img(img, arg1[0], arg1[1], arg1[2], arg1[3])
    t2 = transform_img(img, arg2[0], arg2[1], arg2[2], arg2[3])
    c1 = (arg1[0], arg1[1])
    c2 = (arg2[0], arg2[1])
    p0 = (350, 350)
    p1 = transform_point(p0, (arg1[0], arg1[1]), arg1[2], arg1[3])
    p2 = transform_point(p0, (arg2[0], arg2[1]), arg2[2], arg2[3])

    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(30, 11))
    fig.suptitle("Scaled and rotated images")
    ax[0].set_title("Original")
    ax[0].imshow(img)

    ax[1].set_title("(x,y)=({},{}), Theta={}, Scale={}".format(arg1[0], arg1[1], arg1[2], arg1[3]))
    ax[1].imshow(t1)
    ax[1].scatter(c1[0], c1[1], c='r', s=250)

    ax[2].set_title("(x,y)=({},{}), Theta={}, Scale={}".format(arg2[0], arg2[1], arg2[2], arg2[3]))
    ax[2].imshow(t2)
    ax[2].scatter(c2[0], c2[1], c='r', s=250)

    plt.show()


def plot_match(img1, img2, matches):
    """ Plot the matching points between two images """

    size = img1.shape[0]
    buf_len = 20
    buffer = np.full((size, buf_len, 3), 255)
    concat_img = np.concatenate((cv.cvtColor(img1, cv.COLOR_GRAY2RGB), buffer, cv.cvtColor(img2, cv.COLOR_GRAY2RGB)),
                                axis=1)

    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(20, 11))
    fig.suptitle("SIFT feature matching ({} matches)".format(len(matches)))
    ax.imshow(concat_img)

    colour_list = ['r', 'g', 'b', 'c', 'm', 'y']
    show = len(matches)
    curr = 0
    for p0 in matches:  # plot the points in matching colours on the two images
        if curr >= show:
            break
        colour = random.randint(0, len(colour_list) - 1)
        p1 = matches[p0]
        plt.plot((p0[0], p1[0] + size + buf_len), (p0[1], p1[1]), marker='o', ls='-', c=colour_list[colour])
        curr += 1

    plt.show()


def plot_sift_comp(img, k1, f1, k1_mat, f1_mat):
    """ Plot a comparison of the features detected with my sift and provided sift
        @imgs1, imgs2: [original, transformed] """
    i_01 = cv.cvtColor(img, cv.COLOR_GRAY2RGB)
    i_00 = cv.resize(i_01, (768, 1024))

    fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(20, 30))
    fig.suptitle("SIFT feature matching comparison")
    ax[0, 0].imshow(i_00)
    ax[0, 1].imshow(i_01)
    ax[0, 0].set_title("Original (my SIFT)")
    ax[0, 1].set_title("Original (provided SIFT)")
    colours = ['b', 'g', 'y', 'm', 'r']

    for point in k1:
        ax[0, 0].scatter([point[1] - 128], [point[0]], c='b')
    ax[0, 1].scatter(k1_mat[0], k1_mat[1], c='b')

    vector = f1[215]
    x = np.arange(36)
    values = vector[3:39]
    ax[1, 0].bar(x, values)
    ax[1, 0].set_title("Histogram for ({}, {}) (my SIFT)".format(vector[1] - 128, vector[0]))
    ax[0, 0].scatter(vector[1] - 128, vector[0], c='r')

    v_index = 20
    vector_mat = f1_mat[:, v_index]
    vector_mat = np.sum(np.reshape(vector_mat, (16, 8)), axis=0)
    x_mat = np.arange(len(vector_mat))
    ax[1, 1].bar(x_mat, vector_mat)
    ax[1, 1].set_title("Histogram for ({}, {}) (provided SIFT)".format(k1_mat[0][v_index], k1_mat[1][v_index]))
    ax[0, 1].scatter(k1_mat[0][v_index], k1_mat[1][v_index], c='r')

    plt.show()


def plot_orb_match(img1, img2, kp1, kp2, matches, name, num_drawn=None):
    """ Plot matches found using orb
        num_drawn in the number drawn, by default all matches are shown """
    draw = matches[:] if num_drawn is None else matches[:num_drawn]
    img3 = cv.drawMatches(img1, kp1, img2, kp2, draw, None, flags=2)
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(20, 10))
    ax.imshow(img3)
    ax.set_title(name)
    plt.show()


def plot_ransac_match(img1, img2, kp1, kp2, matches, ransac_idx, name, num_drawn=None):
    """ Plot inlier matches found using ransac
        num_drawn is the number drawn, by default all ransac matches are shown """
    draw_orb = matches[:] if num_drawn is None else matches[:num_drawn]
    img3 = cv.drawMatches(img1, kp1, img2, kp2, draw_orb, None, flags=2)

    draw_ransac = get_ransac_matches(matches, ransac_idx)
    if num_drawn is not None:
        draw_ransac = draw_ransac[:num_drawn]
    img4 = cv.drawMatches(img1, kp1, img2, kp2, draw_ransac, None, flags=2)

    fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(20, 30))
    fig.suptitle(name)
    ax[0].imshow(img3)
    ax[0].set_title("Without RANSAC ({} matches)".format(len(matches)))
    ax[1].imshow(img4)
    ax[1].set_title("With RANSAC ({} matches)".format(len(ransac_idx)))
    plt.show()


def get_ransac_matches(matches, ransac_idx):
    ransac_matches = [matches[index] for index in ransac_idx]
    return ransac_matches


def plot_homography_warp(img1, h, name, output_shape=None, align_left=True):
    """ Plot the original image and the warped image """
    img_rgb = cv.cvtColor(img1, cv.COLOR_GRAY2RGB) if len(img1.shape) == 2 else img1
    img_warp = homography_img(img1, h, output_shape=output_shape, align_left=align_left)[0]
    fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(20, 16))
    fig.suptitle(name)
    ax[0].imshow(img_rgb)
    ax[0].set_title("Original")
    ax[1].imshow(img_warp)
    ax[1].set_title("After transformation")
    plt.show()


if __name__ == '__main__':
    # Part 1
    test_q1 = True
    test_q2 = True
    test_q3 = True
    test_q4 = True

    # Part 2
    test_q7 = True
    test_q8 = True
    test_q9 = True
    test_q10 = True

    school = cv.cvtColor(cv.imread('images/UofT.jpg'), cv.COLOR_BGR2RGB)
    school_bw = cv.cvtColor(school, cv.COLOR_RGB2GRAY)
    school_bw2rgb = cv.cvtColor(school_bw, cv.COLOR_GRAY2RGB)

    apartment = [cv.cvtColor(cv.imread('images/my_apartment/image_{}.png'.format(i)), cv.COLOR_BGR2RGB)
                 for i in range(1, 11)]
    apartment_bw = [cv.cvtColor(img, cv.COLOR_RGB2GRAY) for img in apartment]

    if test_q1:
        keypoints = detect_keypoint(school_bw)
        local_info = keypoint_local(school_bw)
        point_info = list(local_info)[0]
        plot_points("School: {} keypoints".format(len(keypoints)), school_bw2rgb, keypoints)
        plot_patch(point_info, local_info[point_info])

    if test_q2:
        sift_vectors = sift(school_bw)
        plot_histogram(sift_vectors[0])

    if test_q3:
        plot_transform(school_bw2rgb, [450, 500, 30, 1.5], [512, 512, 270, 0.5])

        keypoints1 = detect_keypoint(school_bw)
        t1 = transform_img(school_bw, 512, 512, 270, 0.5)
        keypoints2 = detect_keypoint(t1)
        plot_points("School: {} keypoints".format(len(keypoints1)), school_bw2rgb, keypoints1)
        plot_points("School transformed: {} keypoints".format(len(keypoints2)), cv.cvtColor(t1, cv.COLOR_GRAY2RGB),
                    keypoints2)

    if test_q4:
        t1 = transform_img(school_bw, 450, 500, 30, 1.5)
        t2 = transform_img(school_bw, 512, 512, 270, 0.5)
        matches1 = sift_match(school_bw, t1, ctr=(450, 500), theta=30, s=1.5, search_sz=200)
        matches2 = sift_match(school_bw, t2, ctr=(512, 512), theta=270, s=0.5, search_sz=100)

        plot_match(school_bw, t1, matches1)
        plot_match(school_bw, t2, matches2)

    if test_q7:
        sift_features = loadmat('sift_features.mat')
        theta = int(sift_features["theta"][0, 0])
        apartment_1_1024 = convert_1024(apartment_bw[0])
        t1 = transform_img(apartment_bw[0], 300, 400, theta, 1)
        t1_1024 = convert_1024(t1)
        keypoints = detect_keypoint(apartment_1_1024, thresh=-21)
        sift_vectors = sift(apartment_1_1024, thresh=-21)
        plot_sift_comp(apartment_bw[0], keypoints, sift_vectors,
                       sift_features["keypoints_1"], sift_features["features_2"])

    if test_q8:
        orb1 = orb_match(apartment_bw[0], apartment_bw[1])
        orb2 = orb_match(apartment_bw[8], apartment_bw[9])
        plot_orb_match(apartment_bw[0], apartment_bw[1], orb1[0], orb1[1], orb1[2],
                       "Images 1 and 2 using ORB (all matches)")
        plot_orb_match(apartment_bw[8], apartment_bw[9], orb2[0], orb2[1], orb2[2],
                       "Images 9 and 10 using ORB (top 10 matches)", num_drawn=10)

    if test_q9:
        kp1, kp2, matches = orb_match(apartment_bw[7], apartment_bw[8])
        results = format_orb_match(kp1, kp2, matches)

        debug = False
        n_inputs = 1
        n_outputs = 1
        input_columns = range(n_inputs)  # the first columns of the array
        output_columns = [n_inputs + i for i in range(n_outputs)]  # the last columns of the array
        model = HomographyModel(input_columns, output_columns, debug=debug)
        ransac_fit, ransac_idx = ransac(results, model,
                                        4, 2000, 4000, 50,  # misc. parameters
                                        debug=debug, return_all=True)

        plot_ransac_match(apartment_bw[7], apartment_bw[8], kp1, kp2, matches, ransac_idx,
                          "Images 1 and 2 using ORB (all matches)")
        plot_homography_warp(apartment[7], ransac_fit, "Image 1 Homography")

    if test_q10:
        img_matched = panorama(apartment, apartment_bw)
        fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(22, 11))
        ax.imshow(img_matched)
        ax.set_title("Apartment Panorama")
        plt.show()
