import numpy as np
import scipy
import scipy.signal as sg
import scipy.ndimage as ndimage
import scipy.ndimage.filters as filters
import cv2
import os
import math
from matplotlib import pyplot as plt


output_path = './results'

def main(img, label):

    smooth_5, smooth_10 = gaussian_smooth(img)
    gradient_5, magnitude_5, R_5 = sobel_edge_detection(smooth_5)
    gradient_10, magnitude_10, R_10 = sobel_edge_detection(smooth_10)

    R_10[R_10 < 50] = 0

    nms_3 = nms(R_10, 3)
    nms_30 = nms(R_10, 30)


    output_3 = np.copy(img)
    output_30 = np.copy(img)

    coordinates_3 = np.argwhere(nms_3 > 0)
    coordinates_30 = np.argwhere(nms_30 > 0)

    for coordinate in coordinates_3:
        cv2.circle(output_3, (coordinate[1], coordinate[0]), 5, (0, 0, 200), -1)

    for coordinate in coordinates_30:
        cv2.circle(output_30, (coordinate[1], coordinate[0]), 5, (0, 0, 200), -1)


    if (label == 'original'):

        saveimg('kernel_5.jpg', smooth_5)
        saveimg('kernel_10.jpg', smooth_10)

        # savepltimg('gradient_5.jpg', gradient_5, 'hsv')
        gmap_5 = hsv2rgb(gradient_5, magnitude_5, 16)
        saveimg('gradient_5.jpg', gmap_5)
        savepltimg('magnitude_5.jpg', magnitude_5, 'gray')
        # savepltimg('gradient_10.jpg', gradient_10, 'hsv')
        gmap_10 = hsv2rgb(gradient_10, magnitude_10, 8)
        saveimg('gradient_10.jpg', gmap_10)
        savepltimg('magnitude_10.jpg', magnitude_10, 'gray')

        # savepltimg('tensor_10.jpg', R_10, 'gray')
        savepltimg('nms_3.jpg', nms_3, 'gray')
        savepltimg('nms_30.jpg', nms_30, 'gray')


    print(label, nms_3.sum())
    saveimg('%s_3.jpg' % label, output_3)
    print(label, nms_30.sum())
    saveimg('%s_30.jpg' % label, output_30)




def gaussian_smooth(img):
    row, col = np.meshgrid(np.arange(-2, 3), np.arange(-2, 3))
    kernel_5 = (1 / (2 * np.pi * 25)) * np.exp(-(row**2 + col**2) / (2 * 25))
    kernel_5 /= kernel_5.sum()

    R = img[:, :, 0]
    G = img[:, :, 1]
    B = img[:, :, 2]


    smooth_5R = sg.convolve2d(R, kernel_5, boundary='symm', mode='same')
    smooth_5G = sg.convolve2d(G, kernel_5, boundary='symm', mode='same')
    smooth_5B = sg.convolve2d(B, kernel_5, boundary='symm', mode='same')
    smooth_5 = np.stack((smooth_5R, smooth_5G, smooth_5B), axis=2)

    row, col = np.meshgrid(np.arange(-5, 5), np.arange(-5, 5))
    kernel_10 = (1 / (2 * np.pi * 25)) * np.exp(-(row**2 + col**2) / (2 * 25))
    kernel_10 /= kernel_10.sum()

    smooth_10R = sg.convolve2d(R, kernel_10, boundary='symm', mode='same')
    smooth_10G = sg.convolve2d(G, kernel_10, boundary='symm', mode='same')
    smooth_10B = sg.convolve2d(B, kernel_10, boundary='symm', mode='same')
    smooth_10 = np.stack((smooth_10R, smooth_10G, smooth_10B), axis=2)


    return smooth_5, smooth_10

def sobel_edge_detection(img):

    gray = cv2.cvtColor(np.float32(img), cv2.COLOR_BGR2GRAY)
    sobel_operator_x = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]) / 8
    sobel_operator_y = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]]) / 8
    Gx = sg.convolve2d(gray, sobel_operator_x, boundary='symm', mode='same')
    Gy = sg.convolve2d(gray, sobel_operator_y, boundary='symm', mode='same')
    magnitude = np.sqrt(Gx**2 + Gy**2)
    gradient = np.arctan2(Gx, Gy)
    R_map = structure_tensor(gray, Gx, Gy)

    return gradient, magnitude, R_map

def structure_tensor(img, gradient_x, gradient_y):
    row, col = np.meshgrid(np.arange(-5, 5), np.arange(-5, 5))
    kernel_10 = (1 / (2 * np.pi * 25)) * np.exp(-(row**2 + col**2) / (2 * 25))
    kernel_10 /= kernel_10.sum()

    Ixx = np.square(gradient_x)
    Ixy = np.multiply(gradient_x, gradient_y)
    Iyy = np.square(gradient_y)

    Axx = sg.convolve2d(Ixx, kernel_10, boundary='symm', mode='same')
    Axy = sg.convolve2d(Ixy, kernel_10, boundary='symm', mode='same')
    Ayy = sg.convolve2d(Iyy, kernel_10, boundary='symm', mode='same')

    det = np.multiply(Axx, Ayy) - np.square(Axy)
    trace = Axx + Ayy
    k = 0.05

    return det - k * (trace**2)


def nms(R_map, window_size):
    data_max = filters.maximum_filter(R_map, window_size)
    diff = np.absolute(data_max - R_map)
    # np.savetxt('diff_%d.csv' % window_size, diff, delimiter=',')
    # np.savetxt('Rmap_%d.csv' % window_size, R_map, delimiter=',')
    # np.savetxt('Dmax_%d.csv' % window_size, data_max, delimiter=',')
    output = np.zeros(R_map.shape, dtype=np.uint8)
    output[np.logical_and(diff < 1e-3, R_map > 50)] = 1
    # print(output.sum())
    return output

def saveimg(filename, ref):
    filepath = os.path.join(output_path, filename)
    cv2.imwrite(filepath, ref)

def savepltimg(filename, ref, colormap=None):
    filepath = os.path.join(output_path, filename)
    plt.figure()
    plt.imshow(ref, cmap=colormap, alpha=1.0)
    plt.savefig(filepath)

def hsv2rgb(gradient, magnitude, threshold):

    norm = plt.Normalize(-np.pi, np.pi)
    # print(norm(gradient))
    output = cv2.applyColorMap((norm(gradient) * 255).astype(np.uint8), cv2.COLORMAP_HSV)
    output[magnitude < threshold] = [0, 0, 0]

    return output

if __name__ == '__main__':
    original = cv2.imread('original.jpg')

    rotate = ndimage.rotate(original, 45, reshape=False)
    scale = cv2.resize(original, None, fx=0.5, fy=0.5)


    if not os.path.isdir(output_path):
        os.mkdir(output_path)

    main(original, 'original')
    main(rotate, 'rotate')
    main(scale, 'scale')