# -*- coding: utf-8 -*-

import sys
import numpy as np
from scipy.ndimage import filters
from PIL import Image
import matplotlib.pylab as pyl
from skimage.measure import ransac
from skimage.transform import ProjectiveTransform
from sphinx.addnodes import desc


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)

    print det_M / tr_M

    # return response matrix R
    return det_M / tr_M


def get_harris_interest_points(harris_im, min_distance=10, threshold=0.2):
    """ 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)
                 -> points extracted with min_distance > wid """

    descriptor = []
    for coords in hips:
        patch = img[coords[0] - wid:coords[0] + wid + 1, coords[1] - wid:coords[1] + wid + 1].flatten()
        patch = (patch - np.mean(patch)) / np.linalg.norm(patch)
        descriptor.append(patch)

    return descriptor


def match(descriptor1, descriptor2, threshold=0.95):
    m1 = np.asarray(descriptor1, dtype=np.float32)
    m2 = np.asarray(descriptor2, dtype=np.float32).T

    # print max(descriptor1)
    # print descriptor1
    R_12 = np.dot(m1, m2)
    print "Max is: " + str(m1.max())
    print "Max is: " + str(m2.max())
    print "Max is: " + str(R_12.max())
    img = Image.fromarray(R_12 * 255)
    pairs = []
    for r in range(R_12.shape[0]):
        row_max = R_12[r][0]
        max_index = (r, 0)

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

        # R_12[max_index[0]][max_index[1]] = 255

    img2 = Image.fromarray(R_12 * 255)
    pyl.subplot(121)
    pyl.imshow(img)
    pyl.subplot(122)
    pyl.imshow(R_12)
    pyl.show()
    print R_12.shape
    return R_12, pairs


def plot_matches(image1, image2, hips_image1, hips_image2, matches):
    """ 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]
    matches = []
    # print hips_image1[0]
    # print R_12.shape[1]
    # print R_12.shape[0]
    # for i in range(R_12.shape[0]):
    #     for j in range(R_12.shape[1]):
    #         if R_12[i][j] != 0:
    #             matches.append(i)
    #             matches.append(j)

    print matches
    # 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)
    # image3 = np.vstack((image3, image3))
    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(0, len(matches), 2):
    for i, m in enumerate(matches):
        # print "------------I-------------"
        # print i
        # 1 is X - plotting on x axis, corresponds to columns (remember offset second column value)
        # 2 is Y - plotting on y axis, corresponds to rows (no offset)
        # 3 is style
        if m > 0:
            pyl.plot([hips_image1[i][1], hips_image2[m][1] + column1], [hips_image1[i][0], hips_image2[m][0]], 'c')
            # pyl.plot([hips_image1[matches[i]][1], hips_image2[matches[i + 1]][1] + column1], [hips_image1[matches[i]][0], hips_image2[matches[i + 1]][0]], 'c')
    pyl.axis('off')
    pyl.show()


def RANSAC(R_12, descriptor1, descriptor2):
    """ Step 4 """

    matches = []
    # print hips_image1[0]
    print R_12.shape[1]
    print R_12.shape[0]
    for i in range(R_12.shape[0]):
        for j in range(R_12.shape[1]):
            if R_12[i][j] != 0:
                matches.append(i)
                matches.append(j)

    offset = np.zeros((len(matches), 2))
    for i in range(0, len(matches), 2):
        ndx = matches[i]
        ndx2 = matches[i + 1]
        print ndx
        print ndx2
        offset[i, 0] = descriptor1[ndx][0] - descriptor2[ndx2][0]  # row
        offset[i, 1] = descriptor1[ndx][1] - descriptor2[ndx2][1]  # column

    row_offset, col_offset = offset[0, 0], offset[0, 1]  # select offset at random
    best_matchcount = 0

    for i in range(len(offset)):
        matchcount = 1  # counter for agreements
        offset_i0 = offset[i, 0]
        offset_i1 = offset[i, 1]
        if ((offset_i0 - row_offset) ** 2 + (offset_i1 - col_offset) ** 2) >= 1.6 ** 2:  # compare with WHAT
            sum_row_offset, sum_col_offset = offset_i0, offset_i1
            for j in range(len(matches)):
                if j != 1:
                    offset_j0 = offset[j, 0]
                    offset_j1 = offset[j, 1]
                    if ((offset_i0 - offset_j0) ** 2 + (offset_i1 - offset_j1) ** 2) < 1.6 ** 2:  # compare with WHAT
                        sum_row_offset += offset_j0
                        sum_col_offset += offset_j1
                        matchcount += 1
                if matchcount != best_matchcount:
                    row_offset = sum_row_offset / matchcount
                    col_offset = sum_col_offset / matchcount
                    best_matchcount = matchcount

    return row_offset, col_offset, best_matchcount


def appendimage(image1, image2, row_offset, col_offset):
    """ Step 5 -- note image1, image2 are np arrays -> convert to image"""

    x2, y2 = image2.shape
    x3 = int(x2 + row_offset)
    y3 = int(y2 + col_offset)
    canvas = np.zeros((x3, y3))
    canvas = Image.fromarray(canvas)
    image1 = Image.fromarray(image1)
    image2 = Image.fromarray(image2)
    canvas.paste(image1, (0, col_offset))  # (0,220)
    canvas.paste(image2, (row_offset, 0))  # (290,0)
    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)

image1 /= 255
image2 /= 255
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
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)

#
# print "-----descriptors-------------"
# print descriptor1
# print descriptor2

# im1 = Image.fromarray(descriptor1)
# im2 = Image.fromarray(descriptor2)
# pyl.subplot(121)
# pyl.imshow(im1)
# pyl.subplot(122)
# pyl.imshow(im2)

# Step 3 - Match descriptors -> plot all matching descriptors
print "Matching descriptors"
R_12 = match(descriptor1, descriptor2)
print "Len desc1 %d" % len(descriptor1)
print R_12
# print R_12.shape
# matches = match_stabilised(descriptor1, descriptor2)
plot_matches(image1, image2, hips_image1, hips_image2, R_12)

# Step 4 - Exhaustive RANSAC
row_offset, col_offset, best_matchcount = RANSAC(R_12, descriptor1, descriptor2)
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
appendimage(image1, image2, int(row_offset), int(col_offset))