import skimage as sk
import skimage.transform as skt
import scipy.ndimage as ndimage
import numpy as np
from itertools import product
import math

# filters an image to give it a linear gradient along the row direction (axis 0)
def linear_gradient(im):
    im_border = np.zeros((im.shape[0]+2, im.shape[1]+2))
    im_border[1:im.shape[0]+1, 1:im.shape[1]+1] = im
    im_filter = 0.5*(im_border - np.roll(im_border, -2, axis=1))
    im_filter = np.abs(im_filter[0:im.shape[0], 0:im.shape[1]])
    return im_filter

# def sobel_gradient(im):
#     im_border = np.zeros((im.shape[0]+2, im.shape[1]+2))
#     im_border[1:im.shape[0]+1, 1:im.shape[1]+1] = im
    
#     sobel_x = np.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
#     sobel_y = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
#     filter_dx = ndimage.convolve(im_border, sobel_x)
#     filter_dy = ndimage.convolve(im_border, sobel_y)
#     im_filter = np.square(filter_dx) + np.square(filter_dy)
#     im_filter = np.sqrt(im_filter) / 8.0
#     return im_filter

# uses the L2 norm to calculate the smallest displacement
def L2(im1, im2, dp, shift=(0, 0)):
    m = im2.shape[0]
    n = im2.shape[1]
    error = [sum(sum((im1[max(0,i):min(m, m+i),max(0,j):min(n,n+j)] - im2[max(0,-i):min(m,m-i), max(0,-j):min(n,n-j)])**2)) for i, j in product(range(-dp+shift[0], dp+shift[0]), range(-dp+shift[1], dp+shift[1]))]
    min_index = np.argmin(error)
    result = (min_index // (dp*2) -dp+shift[0], min_index % (dp*2) - dp + shift[1])
    return result

# uses the NCC to calculate the smallest displacement
def NCC(im1, im2, dp):
    m = im2.shape[0]
    n = im2.shape[1]
    error = np.empty(4*dp*dp)
    for i in range(-dp, dp):
        for j in range(-dp, dp):
            v1 = im1[max(0,i):min(m, m+i),max(0,j):min(n,n+j)].flatten()
            v2 = im2[max(0,-i):min(m,m-i), max(0,-j):min(n,n-j)].flatten()
            error[(i+dp)*dp*2 + (j+dp)] = sum(v1*v2)/(np.linalg.norm(v1)*np.linalg.norm(v2))
    min_index = np.argmin(error)
    result = (min_index // (dp*2) -dp, min_index % (dp*2) - dp)
    return result

# crops photo by certain offsets
def crop(im, top, bottom, left, right):
    m = im.shape[0]
    n = im.shape[1]
    im_cropped = im[top:m-bottom, left:n-right]
    return im_cropped

# shifts image over by a vector, puts zeros in the empty space
def shift(v, im):
    m = im.shape[0]
    n = im.shape[1]
    dp = max(abs(v[0]), abs(v[1])) + 1
    im_shifted = np.zeros((m+dp*2, n+dp*2))
    im_shifted[dp:(m+dp), dp:(n+dp)] = im
    im_shifted = im_shifted[dp-v[0]:m+dp-v[0], dp-v[1]:n+dp-v[1]]
    return im_shifted

# combines three r, g, b matrices into one 3D m x n x 3 matrix
def combine(r, g, b):
    r = r.reshape((r.shape[0], r.shape[1], 1))
    g = g.reshape((g.shape[0],g.shape[1],1))
    b = b.reshape((b.shape[0],b.shape[1],1))
    return np.concatenate((r, g, b), axis=2)

# takes two images a metric, and a value for number of levels in image pyramid and aligns them
def align(im1, im2, metric=L2, image_pyramid=0):
    dp = max(im1.shape[0], im1.shape[1], im2.shape[0], im2.shape[1]) // 20
    if image_pyramid:
        corr_v = (0, 0)
        for i in range(image_pyramid, -1, -1): 
            im1_resize = skt.rescale(im1, 0.5**i)
            im2_resize = skt.rescale(im2, 0.5**i)
            if (im1_resize.shape[1] < 512):
                dp = 25
            elif (im1_resize.shape[1] < 1024):
                dp = 15
            elif (im1_resize.shape[1] < 2048):
                dp = 5
            else:
                dp = 1
            corr_v = metric(im1_resize, im2_resize, dp, (2*corr_v[0], 2*corr_v[1]))
        return corr_v
    return metric(im1, im2, dp)

# automatically crops an image by starting at the edges for each color channel and checking how much each pixel row or column differs from each other
def auto_crop(im):
    r = im[:,:, 0]
    r.reshape((im.shape[0], im.shape[1]))

    top = 0
    bottom = 0
    left = 0
    right = 0
    while True:
        if np.std(r[top]) > 0.1:
            break;
        top += 1
    while True:
        if np.std(r[r.shape[0]-1-bottom]) > 0.1:
            break;
        bottom += 1
    while True:
        if np.std(r[:,left]) > 0.1:
            break;
        left += 1
    while True:
        if np.std(r[:,r.shape[1]-1-right]) > 0.1:
            break
        right += 1
    print(top, bottom, left, right)
    im_crop = crop(im, top, bottom, left, right)

    g = im_crop[:,:, 1]
    g.reshape((im_crop.shape[0], im_crop.shape[1]))

    top = 0
    bottom = 0
    left = 0
    right = 0
    while True:
        if np.std(g[top]) > 0.1:
            break;
        top += 1
    while True:
        if np.std(g[g.shape[0]-1-bottom]) > 0.1:
            break;
        bottom += 1
    while True:
        if np.std(g[:,left]) > 0.1:
            break;
        left += 1
    while True:
        if np.std(g[:,g.shape[1]-1-right]) > 0.1:
            break;
        right += 1
    im_crop = crop(im_crop, top, bottom, left, right)

    b = im_crop[:,:, 2]
    b.reshape((im_crop.shape[0], im_crop.shape[1]))

    top = 0
    bottom = 0
    left = 0
    right = 0
    while True:
        if np.std(b[top]) > 0.1:
            break;
        top += 1
    while True:
        if np.std(b[b.shape[0]-1-bottom]) > 0.1:
            break;
        bottom += 1
    while True:
        if np.std(b[:,left]) > 0.1:
            break;
        left += 1
    while True:
        if np.std(b[:,b.shape[1]-1-right]) > 0.1:
            break;
        right += 1
    im_crop = crop(im_crop, top, bottom, left, right)
    
    return im_crop