# Importing the necessary modules
import numpy as np
from scipy import misc
from scipy import ndimage,misc
from scipy.ndimage.filters import convolve
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import os

def gaussian_smoothing(size, sigma=1):
    '''
    A way to get rid of the noise on the image, is by applying Gaussian blur to smooth it. 
    To do so, image convolution technique is applied with a Gaussian Kernel of any size
    The kernel size depends on the expected blurring effect. Basically, the smallest the kernel, 
    the less visible is the blur.
    '''
    size = int(size) // 2
    x, y = np.mgrid[-size:size+1, -size:size+1]
    normal = 1 / (2.0 * np.pi * sigma**2)
    g =  np.exp(-(x**2 + y**2) / (2.0*sigma**2)) * normal
    return g

def rgb2gray(image):
    '''
    This is a way in which we convert the colored RGB Image to Grayscale
    '''    
    r, g, b = image[:,:,0], image[:,:,1], image[:,:,2]
    gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
    return gray

def visualize(img, format_image='gray'):
    '''
    This is a function which helps in visualizing the image generated by the hysterisis function
    using matplotlib.
    '''
    plt.figure(figsize=(8, 8))
    if img.shape[0] == 3:
            img = img.transpose(1,2,0)    
    plt.imshow(img,cmap=format_image)
    plt.show()
    
def gaussian_kernel_image(image,g):
    '''
    This function helps us to apply the filter to our image and get back the desired image results
    '''
    gaussianBlur = convolve(image, g)    
    return gaussianBlur

def sobel_filter(img):
    '''
    The Gradient calculation step detects the edge intensity and direction 
    by calculating the gradient of the image using edge detection operators.
    When the image is smoothed, the derivatives Ix and Iy w.r.t. x and y are calculated. 
    Then, the magnitude gradient and the slope(slope) of the gradient is calculated
    Here Kh and Kv are the vertical and horizontal filters
    '''
    Kh = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], np.float32)
    Kv = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], np.float32)

    Ix = ndimage.filters.convolve(img, Kh)
    Iy = ndimage.filters.convolve(img, Kv)

    gradient = np.hypot(Ix, Iy)
    gradient = gradient / gradient.max() * 255
    slope = np.arctan2(Iy, Ix)
    return (gradient, slope)

def non_max_supression(image,slope):
    '''
    Ideally, the final image should have thin edges. Thus, we must perform non-maximum 
    suppression to thin out the edges.The principle is simple: 
    the algorithm goes through all the points on the gradient intensity matrix and 
    finds the pixels with the maximum value in the edge directions.
    n this algorithm we propose additional penalties to produce more compact edges
    and thus become less sensitive to the threshold of NMS
    The gradient intensity matrixis taken along 4 different angles : 22.5, 67.5, 112.5 and 157.5.
    The steps are :
    - Create a matrix initialized to 0 of the same size of the original gradient intensity matrix;
    - Identify the edge direction based on the angle value from the angle matrix;
    - Check if the pixel in the same direction has a higher intensity than the pixel that is currently processed;
    - Return the image processed with the non-max suppression algorithm.
    '''
    R, C = image.shape
    nms_image = np.zeros((R,C), dtype=np.int32)
    angle = slope * 180. / np.pi
    angle[angle < 0] += 180

    
    for i in range(1,R-1):
        for j in range(1,C-1):
            try:
                q = 255
                r = 255
                
               #angle 0
                if (0 <= angle[i,j] < 22.5) or (157.5 <= angle[i,j] <= 180):
                    q = image[i, j+1]
                    r = image[i, j-1]
                #angle 45
                elif (22.5 <= angle[i,j] < 67.5):
                    q = image[i+1, j-1]
                    r = image[i-1, j+1]
                #angle 90
                elif (67.5 <= angle[i,j] < 112.5):
                    q = image[i+1, j]
                    r = image[i-1, j]
                #angle 135
                elif (112.5 <= angle[i,j] < 157.5):
                    q = image[i-1, j-1]
                    r = image[i+1, j+1]

                if (image[i,j] >= q) and (image[i,j] >= r):
                    nms_image[i,j] = image[i,j]
                else:
                    nms_image[i,j] = 0

            except IndexError as e:
                pass
    
    return nms_image

def threshold(img,ht,lt,wp,sp):
    '''
    The double threshold step aims at identifying 3 kinds of pixels: strong, weak, and non-relevant:
    - Strong pixels are pixels that have an intensity so high that we are sure they contribute to the final edge.
    - Weak pixels are pixels that have an intensity value that is not enough to be considered as strong ones, 
      but yet not small enough to be considered as non-relevant for the edge detection.
    - Other pixels are considered as non-relevant for the edge.
    
    Similarly :
    - High threshold is used to identify the strong pixels (intensity higher than the high threshold)
    - Low threshold is used to identify the non-relevant pixels (intensity lower than the low threshold)
    - All pixels having intensity between both thresholds are flagged as weak and the Hysteresis mechanism (next step) 
      will help us identify the ones that could be considered as strong and the ones that are considered as non-relevant.
    '''
    highThreshold = img.max() * ht;
    lowThreshold = highThreshold * lt;

    M, N = img.shape
    res = np.zeros((M,N), dtype=np.int32)

    weak = np.int32(wp)
    strong = np.int32(sp)

    strong_i, strong_j = np.where(img >= highThreshold)
    zeros_i, zeros_j = np.where(img < lowThreshold)

    weak_i, weak_j = np.where((img <= highThreshold) & (img >= lowThreshold))

    res[strong_i, strong_j] = strong
    res[weak_i, weak_j] = weak

    return res

def hysteresis(image,wp,sp):
    '''
    Based on the threshold results, the hysteresis consists of transforming weak pixels into strong ones, 
    if and only if at least one of the pixels around the one being processed is a strong one.Weak pixels 
    connected to strong pixels are kept, while the weak pixels connected to weak pixels are discarded.
    This is done to determine the edges which are actually weak but are real edges
    '''
    
    R, C = image.shape
    weak = wp
    strong = sp
    for i in range(1, R-1):
        for j in range(1, C-1):
            if (image[i,j] == weak):
                try:
                    if ((image[i+1, j-1] == strong) or (image[i+1, j] == strong) or (image[i+1, j+1] == strong)
                        or (image[i, j-1] == strong) or (image[i, j+1] == strong)
                        or (image[i-1, j-1] == strong) or (image[i-1, j] == strong) or (image[i-1, j+1] == strong)):
                        image[i, j] = strong
                    else:
                        image[i, j] = 0
                except IndexError as e:
                    pass

    return image

def image_edge_detect(path, sigma=1, kernel_size=5, weakp=70, strongp=225, low_t=0.07, high_t=0.18):
    '''
    This is like a feed forward function which feeds in the original throgh all the functions 
    and returns back the image obtained.
    '''
    gauss_kernel = gaussian_smoothing(kernel_size, sigma)
    img=mpimg.imread(path)
    image=rgb2gray(img)
    smooth_image = convolve(image, gauss_kernel)
    
    gradient, slope = sobel_filter(smooth_image)
    nms_image = non_max_supression(gradient, slope)
    
    thresh_image = threshold(nms_image, high_t, low_t, weakp, strongp)
    
    final_image = hysteresis(thresh_image, weakp, strongp)   
    return final_image

def folder_edge_detect(path, sigma=1, kernel_size=5, weakp=70, strongp=225, low_t=0.07, high_t=0.18):
    '''
    A similar function as above. This iterates through images in a folder instead of asingle image
    '''
    final=[]
    for img in os.listdir(path):
        img=mpimg.imread(os.path.join(path,img))
        canny_img=image_edge_detect(img,sigma,kernel_size,weakp,strongp,low_t,high_t)
        final.append(canny_img)        
    return final       
    






    
    
    
    



