"""
This module contains the functions necessary to extract a thumbprint for a given image. When run in main, 
plots the thumbprint generated by sys.argv[1]
"""
import skimage.filters as skfill
import skimage.morphology as skmorph
import skimage.measure as skmeasure
import skimage.exposure as ske
import matplotlib.pyplot as plt
import numpy as np
import cv2 as cv
import heapq
import sys


def sin2d(u, v, ii, jj, N, M):
    """
    Generates a sine wave structuring element to be used in convolving.
    
    Params: u: int. factor of the vertical period of waves (higher value, more waves)
            v: int. factor of the horizontal period of waves (higher value, more waves)
            ii: NxM 2D numpy array. Horizontal gradient where col0 = 0 and coln = coln-1 + 1
            jj: NxM 2D numpy array. Vertical gradient where col0 = 0 and coln = coln-1 + 1
            N: int. # rows
            M: int. # cols
            
    Output: 2D numpy array.  Our sine wave structure element
    """
    selem = np.sin(u * ii / (2 * np.pi) + v * jj / (2 * np.pi))
    return selem * 255 / (N * M)


def high_freq_detection(imgray):
    """
    Given an image, finds the high spatial frequency areas of the image by convolving high 
    frequency sine waves with the image in question.  High frequency areas are thus brighter
    and low frequency areas are darker.
    
    Params: imgray: 2D numpy array.  B&W image in question
    
    Output: 2D numpy array with high spatial frequency regions outlined
    """
    N = 200
    M = 200
    i = np.linspace(0, N - 1, N)
    j = np.linspace(0, M - 1, M)
    jj, ii = np.meshgrid(j, i)

    # Looked at gabor filter but doesn't work because of the gray region outside of the peak sign region.
    # Here we use sign waves of varying frequencies and directions to cature as much of the thumb print
    # as possible.
    
    # Creating our filters
    fill_1 = sin2d(2, 2, ii, jj, N, M)
    fill_2 = sin2d(1, 2, ii, jj, N, M)
    fill_3 = sin2d(2, 1, ii, jj, N, M)
    fill_4 = sin2d(-2, 1, ii, jj, N, M)
    fill_5 = sin2d(-2, 2, ii, jj, N, M)
    fill_6 = sin2d(-2, 3, ii, jj, N, M)
    
    # Creating a blank image which we will combine each of our convolved images with
    new_img = np.zeros(imgray.shape)
    conv_1 = cv.filter2D(imgray, -1, fill_1)
    conv_2 = cv.filter2D(imgray, -1, fill_2)
    conv_3 = cv.filter2D(imgray, -1, fill_3)
    conv_4 = cv.filter2D(imgray, -1, fill_4)
    conv_5 = cv.filter2D(imgray, -1, fill_5)
    conv_6 = cv.filter2D(imgray, -1, fill_6)
    
    # new_img capures all the high frequency regions located by each output image
    new_img += conv_1 + conv_2 + conv_3 + conv_4 + conv_5 + conv_6

    return new_img


def thresh_series_mask(imgray):
    """
    Generates a mask to be used on original image to crop out the thumb region.
    
    Params: imgray: 2D numpy array.  B&W image in question
    
    Output: 2D numpy array.  Binary mask for B&W image in question
    """
    # first equilize the histogram to emphasize the area around the thumb
    imgray_eq = ske.equalize_hist(imgray)
    
    # median filter to perserve the edges while removing spatial frequency
    median_img = skfill.median(imgray_eq, selem=skmorph.disk(21))
    
    # some blocks sizes work better than others for different images,
    # we try a wide range of them to get the best qualities of each
    block_sizes = [21, 41, 51, 61, 81]
    masks = []
    for block_size in block_sizes:
        
        # some larger some smaller threshold blocks
        local_thresh = skfill.threshold_local(median_img, block_size)
        binary_local = median_img > local_thresh
        
        # detect regions
        labels = skmeasure.label(binary_local)
        hight_freq_img = high_freq_detection(imgray)
        
        # take 15 largest area regions detected in our thumb
        label_props = skmeasure.regionprops(labels)
        areas = [label_props[i]["area"] for i in range(len(label_props))]
        max_areas = heapq.nlargest(15, areas)
        max_indicies = [areas.index(max_areas[i]) + 1 for i in range(len(max_areas))]
        
        # find the largest/highest frequency region
        max_i = 0
        max_count = 0
        for i in range(len(max_indicies)):
            label_vals = labels == max_indicies[i]
            
            # finding our region using our high_freq_detection() output
            current_count = (label_vals.astype(int) * hight_freq_img).sum()
            if current_count > max_count:
                max_count = current_count
                max_i = max_indicies[i]

        mask = labels == max_i
        masks.append(mask)

    mask = np.zeros(masks[0].shape).astype(bool)
    
    # combine all masks from each of our thresholds with bitwise or
    for m in masks:
        mask = mask | m
    
    # clean up the mask
    mask = skmorph.binary_closing(mask, selem=skmorph.disk(30))
    return skmorph.binary_erosion(mask, selem=skmorph.disk(10))


def extract_print(mask, imgray):
    """
    Gets our thumbprint, given a mask generated by thresh_series_mask(imgray) and
    imgray, our image in question.
    
    Params: imgray: 2D numpy array.  B&W image in question
            mask: 2D numpy array.  Binary image mask
    
    Output: 2D numpy array.  Our final thumbprint
    """
    imgray[mask != 1] = 0
    
    # on the cropped image, run localized threshold to extract the print
    block_size = 41
    local_thresh = skfill.threshold_local(imgray, block_size)
    binary_local = imgray > local_thresh
    
    # make background white
    binary_local[mask == 0] = 1
    
    # denoise
    f_print = skfill.median(binary_local, selem=skmorph.disk(3))

    return f_print


if __name__ == "__main__":
    imgray = cv.imread(sys.argv[1], cv.IMREAD_GRAYSCALE)
    f_print = extract_print(thresh_series_mask(imgray), imgray)
