import time
import sys
import glob
import os
import cv2
from skimage import morphology, measure, restoration
from skimage import util
from skimage import color
from skimage.filters import scharr, gaussian, threshold_otsu, threshold_sauvola
import numpy as np
from scipy import ndimage
import matplotlib.pyplot as plt

PROC_VERSION = 101
EDGE_THRESH = 10
AREA_FILTER = 5
DECONV = True
DECONV_ITER = 9
DECONV_METHOD = 'LR'
DECONV_MASK_WEIGHT = 0.5
ESTIMATE_SHARPNESS = True
SMALL_FLOAT_VAL = 0.0001
BW_BLUR_RADIUS = 1
BAYER_PATTERN = cv2.COLOR_BAYER_RG2RGB

# convert image to 8 bit with or without autoscaling
def convert_to_8bit(img, auto_scale=True):
    # Convert to 8 bit and autoscale
    if auto_scale:

        result = np.float32(img) - np.min(img)
        if np.max(img) != 0:
            result = result / np.max(img)

        img_8bit = np.uint8(255 * result)
    else:
        img_8bit = np.unit8(img)

    return img_8bit

# import raw image
def import_image(abs_path, filename, raw=True, color=False, bayer_pattern=BAYER_PATTERN):
    # Load and convert image as needed
    img_c = cv2.imread(os.path.join(abs_path, filename), cv2.IMREAD_UNCHANGED)
    if raw:
        if color:
            img_c = cv2.cvtColor(img_c, bayer_pattern)

    return img_c

def find_particles(filename,save_to_disk=False,min_edge=EDGE_THRESH,min_area=AREA_FILTER,show=False,color=False):

    gray = import_image(os.path.dirname(filename),os.path.basename(filename),raw=True,color=color)

    gray = convert_to_8bit(gray)

    #if show:
    #    plt.imshow(gray)
    #    plt.show()
    # remove single-p

    # edge-based segmentation and region filling to define the object
    edges_mag = scharr(gray)
    #edges_med = np.median(edges_mag[edges_mag > 0])
    #edges_thresh = min_edge*edges_med
    #edges = edges_mag >= edges_thresh
    
    thresh = threshold_otsu(edges_mag)
    edges = edges_mag > thresh*.5
    
    if show:
    #    plt.imshow(edges_mag)
    #    plt.show()
        plt.imshow(edges)
        plt.show()
    # Post process
    #edges = morphology.erosion(edges, morphology.disk(1))
    # define the binary image for further operations
    bw_img = edges

    # Compute morphological descriptors
    label_img = morphology.label(bw_img.astype('uint16'), neighbors=8, background=0)
    props = measure.regionprops(label_img, gray)

    color_img = cv2.cvtColor(gray,cv2.COLOR_GRAY2RGB)

    particle_count = 0
    pdiam = []
    for p in props:
        if p['area'] > min_area and p['eccentricity'] < 0.6:
            particle_count += 1
            cv2.circle(color_img,
                       (int(p['centroid'][1]),
                        int(p['centroid'][0])),
                        int(p['major_axis_length']/2.0),
                        (0, 255, 0),
                        2)
            #print(p['major_axis_length']/2.0)
            majal = p['major_axis_length'] - 1
            minal = p['minor_axis_length'] - 1
            esd = np.sqrt(majal*minal)/2
            
            pdiam.append(esd)
    #print(filename + "," + str(particle_count))

    if show:
        plt.imshow(color_img)
        plt.show()
        plt.hist(1000*np.array(pdiam)*.04, color = 'blue', edgecolor = 'black',
         bins = range(10,1000,10))
        plt.show()

    if save_to_disk:
        output = filename.rstrip('\\')+'.proc.tif'
        print('Saving to: ' + output)
        cv2.imwrite(output,label_img.astype('uint16'))


    return particle_count
    
if __name__=="__main__":

    if len(sys.argv) != 2:
        print('Please supply a directory')
        exit()

    #scales = ['x1', 'x2', 'x4', 'x8', 'x16', 'x32', 'x64', 'x128']

    total_focus_files = glob.glob(os.path.join(sys.argv[1],"*_TotalFocus.png"))
    for im in total_focus_files:
        find_particles(im,
            save_to_disk=False,
            min_edge=EDGE_THRESH,
            min_area=AREA_FILTER,
            show=True,
            color=False)