# -*- coding: utf-8 -*-
"""
Created on Thu Jun  6 10:21:54 2019

@author: Student
"""

import numpy as np
from matplotlib import pyplot as plt
from skimage import io 
from scipy import ndimage as nd
from scipy import interpolate
from scipy import linalg
from skimage import color
import time


class LineBuilder:
    def __init__(self, line):
        self.line = line
        self.xs = list(line.get_xdata())
        self.ys = list(line.get_ydata())
        self.cid = line.figure.canvas.mpl_connect('button_press_event', self)
        self.cid = line.figure.canvas.mpl_connect('key_press_event', self)
        self.cid = line.figure.canvas.mpl_connect('scroll_event', self)

    def __call__(self, event):
        if(event.name=='key_press_event'):
            self.xs.append(int(self.xs[0]))
            self.ys.append(int(self.ys[0]))
            self.line.set_data(self.xs, self.ys)
            self.line.figure.canvas.draw()
            self.line.figure.canvas.mpl_disconnect(self.cid)
            print("xs = ", self.xs)
            print("ys = ", self.ys)
        if (event.name == 'button_press_event'):
            self.xs.append(int(event.xdata))
            self.ys.append(int(event.ydata))
            self.line.set_data(self.xs, self.ys)
            self.line.figure.canvas.draw()
            print("xs = ", self.xs)
            print("ys = ", self.ys)
        if(event.name=='scroll_event'):
            self.xs.pop()
            self.ys.pop()
            self.line.set_data(self.xs, self.ys)
            print("xs = ", self.xs)
            print("ys = ", self.ys)
        

def click_points(image):
    fig = plt.figure()
    ax = fig.add_subplot(111)            
    ax.imshow(image,cmap='gray')
    ax.set_title('click')
    line, = ax.plot([], [])
    linebuilder = LineBuilder(line)
    plt.show()
    plt.pause(7)
    xs=np.asarray(linebuilder.xs)
    ys=np.asarray(linebuilder.ys)
    return xs, ys

            
def reinterpolate_contours(xs, ys, dmin, dmax):
    
    n_xs = xs
    n_ys = ys
    grid_x, grid_y = np.meshgrid(n_xs, n_ys)
    a = n_ys.shape[0]
    b = n_xs.shape[0]
    
    for i in range(n_xs.shape[0]):
        if np.sqrt((n_xs[i+1] - n_ys[i+1])**2 + ((n_xs[i] - n_ys[i])**2)) > dmax:
            for j in range(grid_x.shape[1]):
                I1 = b-n_xs[j]/b - grid_x[0,0] + n_xs[j]/b * grid_x[0,b]
                I2 = b-n_xs[j]/b - grid_x[a,0] + n_xs[j]/b* grid_x[a,b]
                I3 = a-n_ys[j]/a * I1 + n_ys[j]/a * I2
            grid_y[j], grid_x[j] = I3
    
    return n_xs, n_ys




def calculate_gradient_at_contour_points(image, sigma, xs, ys):
    fimage = nd.gaussian_filter(image, sigma)
    gx, gy = np.gradient(fimage)
    mag = (gx**2+gy**2)**0.5
    mag = (mag - mag.min()/(mag.max()-mag.min()))
    xs_gradient, ys_gradient = np.gradient(mag)
    
    return xs_gradient, ys_gradient

def fx(xs, ys):
       
    xs[ xs < 0 ] = 0.
    ys[ ys < 0 ] = 0.

    xs[ xs > image.shape[1]-1 ] = image.shape[1]-1
    ys[ ys > image.shape[0]-1 ] = image.shape[0]-1

    return xs_gradient[ (y.round().astype(int), x.round().astype(int)) ]

def fy(x, y):
        
    x[ x < 0 ] = 0.
    y[ y < 0 ] = 0.
    
    x[ x > image.shape[1]-1 ] = image.shape[1]-1
    y[ y > image.shape[0]-1 ] = image.shape[0]-1
            
    return ys_gradient[ (y.round().astype(int), x.round().astype(int)) ]

        


def generate_coefficient_matrix(N, alpha, beta):
    def create_matrix_A(a, b, N):
        row = np.r_[
        -2*a - 6*b, 
        a + 4*b,
        -b,
        np.zeros(N-5),
        -b,
        a + 4*b
    ]
    A = np.zeros((N,N))
    for i in range(N):
        A[i] = np.roll(row, i)
        return A

def run():
    

    image = io.imread('xray.jpg')
    image = color.rgb2gray(image)
    
    plt.show()
    xs,ys = click_points(image)
    time.sleep(1)
    n_xs, n_ys = reinterpolate_contours(xs, ys, 5, 1)
    xs_gradient, ys_gradient = calculate_gradient_at_contour_points(image, 1, xs, ys)
    

    plt.figure(2)
    plt.imshow(image)
    
if __name__ == "__main__":
	run()