# -*- coding: utf-8 -*-
import sys
import numpy as np
from scipy.ndimage import filters
from PIL import Image
import matplotlib.pylab as pyl


def harris_matrix_algorithm(img, sigma=2):
    """ step 1.1 - algorithm to compute harris corner detector response matrix """

    ix = np.zeros(img.shape)
    filters.gaussian_filter(img, 1, (0, 1), output=ix)  # (1,1) -> (sigma,sigma)
    iy = np.zeros(img.shape)
    filters.gaussian_filter(img, 1, (1, 0), output=iy)  # (1,1) -> (sigma,sigma)

    # compute components of the Harris matrix
    A = filters.gaussian_filter(ix * ix, sigma)
    B = filters.gaussian_filter(ix * iy, sigma)
    C = filters.gaussian_filter(iy * iy, sigma)

    # calculate trace and determinant
    det_M = (A * C) - (B ** 2)
    tr_M = (A + C)

    # return response matrix R
    return det_M / tr_M


def get_harris_interest_points(harris_im, min_distance=10, threshold=0.1):
    """	step 1.2 - select Harris corner points above a threshold and perform nonmax
        suppression in a region +/- min_distance.
        min_distance is the minimum number of pixels separating corner and image boundary (spec = 11)"""

    # find top corner candidates above a threshold
    corner_threshold = harris_im.max() * threshold
    harris_im_th = (harris_im > corner_threshold)  # expect True/False

    # find the co-ordinates of these candidates
    coords = np.array(harris_im_th.nonzero()).T
    # and their values
    candidate_values = np.array([harris_im[c[0], c[1]] for c in coords])
    indices = np.argsort(candidate_values)

    # store allowed point locations in a boolean image
    allowed_locations = np.zeros(harris_im.shape, dtype='bool')
    allowed_locations[min_distance:-min_distance, min_distance:-min_distance] = True

    # select best points, using nonmax suppression based on allowed_locations array (taking min_distance into account)
    filtered_coords = []
    for i in indices[::-1]:  # back to front
        r, c = coords[i]  # note reversed indices
        if allowed_locations[r, c]:
            filtered_coords.append((r, c))
            allowed_locations[r - min_distance:r + min_distance + 1, c - min_distance:c + min_distance + 1] = False

    return filtered_coords


def plot_harris_interest_points(img, hips):
    """ plot interest points """
    pyl.figure('Harris points/corners')
    pyl.imshow(img, cmap='gray')
    pyl.plot([p[1] for p in hips], [p[0] for p in hips], 'ro')
    pyl.axis('off')
    pyl.show()


def descriptors(img, hips, wid=5):
    """ step 2 - return pixel values around harris interest points using width 2*wid+1 (=11, as per spec) """

    descriptor = []
    for coords in hips:
        patch = img[coords[0] - wid:coords[0] + wid + 1, coords[1] - wid:coords[1] + wid + 1].flatten()
        # mean must be subtracted first then divided by norm
        # doing all in one line caused errors
        patch -= np.mean(patch)
        patch /= np.linalg.norm(patch)
        descriptor.append(patch)

    return descriptor


def match(descriptor1, descriptor2, threshold=0.95):
    # convert descriptors to arrays
    m1 = np.asarray(descriptor1, dtype=np.float32)
    m2 = np.asarray(descriptor2, dtype=np.float32).T

    # calculate the response matrix by finding the inner product
    response_matrix = np.dot(m1, m2)
    print "Max is: " + str(m1.max())
    print "Max is: " + str(m2.max())
    print "Max is: " + str(response_matrix.max())
    img = Image.fromarray(response_matrix * 255)
    pairs = []
    for r in range(response_matrix.shape[0]):
        row_max = response_matrix[r][0]

        for c in range(response_matrix.shape[1]):
            # pick only one
            if response_matrix[r][c] > threshold and response_matrix[r][c] > row_max:
                pairs.append((r, c))
            else:
                response_matrix[r][c] = 0

    # plot
    img2 = Image.fromarray(response_matrix * 255)
    pyl.subplot(121)
    pyl.imshow(img)
    pyl.subplot(122)
    pyl.imshow(img2)
    pyl.show()
    return pairs


def plot_matches(image1, image2, hips_image1, hips_image2, pairs):
    """ step 3.3 -  output a concatenation of images (appended side by side) showing the 'true' matches between points. """
    rows1 = image1.shape[0]
    rows2 = image2.shape[0]

    # select image with least rows and pad accordingly (wrt to larger image)
    if rows1 < rows2:
        image1 = np.concatenate((image1, np.zeros((rows2 - rows1, image1.shape[1]))), axis=0)
    elif rows2 < rows1:
        image2 = np.concatenate((image2, np.zeros((rows1 - rows2, image2.shape[1]))), axis=0)

    # create new image with two input images appended side-by-side, then plot matches
    image3 = np.concatenate((image1, image2), axis=1)

    pyl.figure('Both images side by side with matches plotted')  # note incorrect matches! - use RANSAC
    pyl.imshow(image3, cmap="gray")
    column1 = image1.shape[1]

    for i in range(len(pairs)):
        index1, index2 = pairs[i]
        pyl.plot([hips_image1[index1][1], hips_image2[index2][1] + column1],
                 [hips_image1[index1][0], hips_image2[index2][0]], 'c')
    pyl.axis('off')
    pyl.show()


def basic_RANSAC(matches, coord_list1, coord_list2, match_dist=1.6):
    """Carry out a simple version of the RANSAC procedure for calculating
       the appropriate translation needed to map image 2 onto image 1.
       N.B., this isn't true RANSAC because it doesn't carry out *random*
       consensus, instead it works exhaustively: for each offset in the
       list of possible offsets it checks to see how many other offsets
       agree with it (that point's consensus set), then returns the average
       offset of all the points in the largest consensus set.  Hence it is
       not performing *random* consensus, but rather *exhaustive*
       consensus.

       matches - List of matches, i.e., index pairs, (i1,i2). i1 is the
       index of the matching coordinate from coord_list1 and i2 the index
       of the coordinate in coord_list2 that it matches.  This list of
       matches is generated by a previous call to "match".

       coord_list1, coord_list2 -- lists of interest point coordinates in
       the first and second images.

       match_dist -- maximum distance between offsets for them to be
       considered part of a consensus set.

       Returns -- (row_offset,column_offset,match_count) -- a tuple,
       the first two elements are the best_offset relating the second image to
       the first, the third is the number of point-pairs supporting this
       offset according to this "Exhaustive RANSAC" matcher, i.e., the
       "strength" of this offset.
    """
    d2 = match_dist ** 2  ##  Square the match distance.
    ## Build a list of offsets from the lists of matching points for the
    ## first and second images.
    offsets = np.zeros((len(matches), 2))
    for i in range(len(matches)):
        index1, index2 = matches[i]
        offsets[i, 0] = coord_list1[index1][0] - coord_list2[index2][0]
        offsets[i, 1] = coord_list1[index1][1] - coord_list2[index2][1]

    ## Run the comparison.  best_match_count keeps track of the size of the
    ## largest consensus set, and (best_row_offset,best_col_offset) the
    ## current offset associated with the largest consensus set found so far.
    best_match_count = -1
    best_row_offset, best_col_offset = 1e6, 1e6
    for i in range(len(offsets)):
        match_count = 1.0
        offi0 = offsets[i, 0]
        offi1 = offsets[i, 1]
        ## Only continue into j loop looking for consensus if this point hasn't
        ## been found and folded into a consensus set earlier.  Just improves
        ## efficiency.
        if (offi0 - best_row_offset) ** 2 + (offi1 - best_col_offset) ** 2 >= d2:
            sum_row_offsets, sum_col_offsets = offi0, offi1
            for j in range(len(matches)):
                if j != i:
                    offj0 = offsets[j, 0]
                    offj1 = offsets[j, 1]
                    if (offi0 - offj0) ** 2 + (offi1 - offj1) ** 2 < d2:
                        sum_row_offsets += offj0
                        sum_col_offsets += offj1
                        match_count += 1.0
            if match_count >= best_match_count:
                best_row_offset = sum_row_offsets / match_count
                best_col_offset = sum_col_offsets / match_count
                best_match_count = match_count
    return best_row_offset, best_col_offset, best_match_count


def appendimage(image1, image2, row_offset, col_offset):
    row_offset = int(row_offset)
    col_offset = int(col_offset)
    canvas = Image.new(image1.mode, (image1.width + abs(col_offset), image1.width + abs(row_offset)))
    canvas.paste(image1, (0, canvas.height - image1.height))
    canvas.paste(image2, (col_offset, canvas.height - image1.height + row_offset))

    pyl.figure('Final Composite Image')
    pyl.imshow(canvas)
    pyl.axis('off')
    pyl.show()


# load images (accept cmdline arguments) 	# pass the two images with .py at cmd
print "Loading images..."
image1 = np.array(Image.open(sys.argv[1]).convert('L'), dtype=np.float32)
image2 = np.array(Image.open(sys.argv[2]).convert('L'), dtype=np.float32)
print "Done."

# Step 1 - Find HIPS by computing and thresholding harris response matrix
print "Retrieving harris images"
harris_image1 = harris_matrix_algorithm(image1, 2)  # sigma = 2
harris_image2 = harris_matrix_algorithm(image2, 2)  # sigma = 2

# plot the harris response
pyl.subplot(121)
pyl.imshow(harris_image1, cmap="gray")
pyl.subplot(122)
pyl.imshow(harris_image2, cmap="gray")
pyl.show()

print "Finding interest points"
hips_image1 = get_harris_interest_points(harris_image1, min_distance=10, threshold=0.1)
hips_image2 = get_harris_interest_points(harris_image2, min_distance=10, threshold=0.1)
print "Interest Points (img1): " + str(len(hips_image1))
print "Interest Points (img2): " + str(len(hips_image2))

print "Plotting interest points"
plot_harris_interest_points(image1, hips_image1)
plot_harris_interest_points(image2, hips_image2)

# Step 2 - Form normalised patch descriptor vectors for all HIPScompute feature locations/descriptors
print "Developing descriptors"
descriptor1 = descriptors(image1, hips_image1)
descriptor2 = descriptors(image2, hips_image2)

# Step 3 - Match descriptors -> plot all matching descriptors
print "Matching descriptors"
pairs = match(descriptor1, descriptor2)

# matches = match_stabilised(descriptor1, descriptor2)
plot_matches(image1, image2, hips_image1, hips_image2, pairs)

# Step 4 - Exhaustive RANSAC
row_offset, col_offset, best_matchcount = basic_RANSAC(pairs, hips_image1, hips_image2)
print '--Exhaustive RANSAC--'
print 'Number of agreements (best match count): %d' % best_matchcount
print 'Row offset: %d' % row_offset
print 'Column offset: %d' % col_offset

# Step 5 - Create composite image using RANSAC translation
image1_rgb = Image.open(sys.argv[1])
image2_rgb = Image.open(sys.argv[2])
appendimage(image1_rgb, image2_rgb, int(row_offset), int(col_offset))