import os
import warnings
import numpy as np
from sklearn.cluster import KMeans
from scipy.interpolate import interp1d
from scipy.integrate import trapz
from scipy import optimize
from pspylib.tiff.reader import TiffReader

def cal_histogram_range(flatten_image):
    split_hists, split_edges = cluster_kmeans(flatten_image)
    # change unit micro -> pico
    scale = 1000000
    # popt_gauss = (amplitude, center, sigma)
    _, popt_gauss_1 = one_peak_gaussian_fit(split_edges[0][:-1]*scale, split_hists[0])
    _, popt_gauss_2 = one_peak_gaussian_fit(split_edges[1][:-1]*scale, split_hists[1])
    #change unit pico -> micro -> nano
    popt_gauss_1 = list(popt_gauss_1)
    popt_gauss_2 = list(popt_gauss_2)
    popt_gauss_1[1] = popt_gauss_1[1]/scale * 1000
    popt_gauss_2[1] = popt_gauss_2[1]/scale * 1000
    histogram_range = np.abs((popt_gauss_1[1]-popt_gauss_2[1]))
    return histogram_range, popt_gauss_1, popt_gauss_2

def cal_histogram_range_with_infimum(flatten_image):
    split_hists, split_edges = cluster_kmeans(flatten_image)
    # change unit Pico -> Micro
    scale = 1000000
    ## add dummy infimum value: len(hist)+1 = len(edges)
    split_hists[0] = np.insert(split_hists[0],0,0)
    split_hists[1] = np.insert(split_hists[1],0,0)
    # popt_gauss = (amplitude, center, sigma)
    _, popt_gauss_1 = one_peak_gaussian_fit(split_edges[0]*scale, split_hists[0])
    _, popt_gauss_2 = one_peak_gaussian_fit(split_edges[1]*scale, split_hists[1])
    #change unit Micro -> Pico -> nano
    popt_gauss_1 = list(popt_gauss_1)
    popt_gauss_2 = list(popt_gauss_2)
    popt_gauss_1[1] = popt_gauss_1[1]/scale * 1000
    popt_gauss_2[1] = popt_gauss_2[1]/scale * 1000
    histogram_range = np.abs((popt_gauss_1[1]-popt_gauss_2[1]))
    return histogram_range, popt_gauss_1, popt_gauss_2

def region_histogram(region):
    bins = __optimal_bins(region)
    hist,edges = np.histogram(region,bins)
    return hist, edges

def __optimal_bins(gray_image):
    vmin = np.min(gray_image)
    vmax = np.max(gray_image)
    # cvt Scale (Pico -> Nano)
    bound = int((vmax - vmin)* 1000)
    if bound < 10:
        bound = 128
    count = gray_image.size
    sqrt_count = int(np.sqrt(count))

    if (sqrt_count < bound) & (count < bound * 50):
        bins = sqrt_count                                                  
    else: bins = bound
    if bins > 128:
        bins = 128
    return bins

def cluster_kmeans(gray,n_cluster:int=2):
    data  = np.reshape(gray.copy(), (-1,1))
    kmeans = KMeans(n_clusters=n_cluster,n_init=10 ,tol=1e-6,random_state=0)
    #kmeans = KMeans(n_clusters=n_cluster,n_init='auto',tol=1e-4)
    kmeans.fit_predict(data)
    cluster_id = kmeans.fit_predict(data)
    _ = kmeans.cluster_centers_
    split_data_hist = []
    split_data_edges = []
    for id in np.unique(cluster_id):
        subset = data[cluster_id==id]
        hist, edges = region_histogram(subset)
        split_data_hist.append(hist)
        split_data_edges.append(edges)
    return split_data_hist, split_data_edges

def one_peak_gaussian_fit(x,y):
    y, denominator = __norm_gaussian(x,y)
    popt_gauss = __one_peak_gaussian_fit(x, y)
    fit_data = __one_peak_gaussian(x, *popt_gauss)
    return fit_data*denominator, popt_gauss

def __one_peak_gaussian_fit(x_ori,y_ori):
    """Do not use pico & Nano scale
        fit_data = __one_peak_gaussian(x_axis,popt_gauss[0],popt_gauss[1],popt_gauss[2])
        fit_data = __one_peak_gaussian(x_axis,*popt_gauss)
    """
    x = x_ori.copy()
    y = y_ori.copy() 
    x, y = __interpolate(x, y)
    x,y = __clip_negative(x, y)

    amp = y.max()
    cen = x[np.where(y == amp)][0]
    sigma = np.std(y)
    warnings.filterwarnings("error", category=UserWarning)
    try:
        popt_gauss, _ = optimize.curve_fit(__one_peak_gaussian, x, y, p0=[amp, cen, sigma])
    except UserWarning as e:
        popt_gauss = (amp,cen,sigma)
    except RuntimeError:
        popt_gauss = (amp,cen,sigma)
    return popt_gauss

def __norm_gaussian(x,y):
    # integral = 1
    y = np.abs(y)
    eps = 7/3. - 4/3. -1 
    denominator = np.abs(trapz(x,y)) + eps
    return y / denominator, denominator

def __one_peak_gaussian(x, amp,cen,sigma):
    eps = 7/.3 - 4/.3 - 1
    return amp *( 1 / (eps + sigma * (np.sqrt(2 * np.pi)))) * (np.exp((-1.0 / 2.0) * (((x-cen) / (sigma + eps))**2)))

def __clip_negative(x,y):
    clip_index = np.where(y > 0)
    return x[clip_index], y[clip_index]

def __interpolate(x,y, num:int=10000):
    fun = interp1d(x,y,kind='cubic')
    x_new = np.linspace(x.min(), x.max(),num=num,endpoint=True)
    y_new = fun(x_new)
    return x_new, y_new

def __one_peak_gaussian(x, amp,cen,sigma):
    eps = 7/.3 - 4/.3 - 1
    return amp *( 1 / (eps + sigma * (np.sqrt(2 * np.pi)))) * (np.exp((-1.0 / 2.0) * (((x-cen) / (sigma + eps))**2)))

if __name__ == "__main__":
    samples_path = r"C:\Park Systems\SmartScan\samples"
    tiff_path = os.path.join(samples_path, "Image", "Cheese.tiff")
    tiff = TiffReader(tiff_path)
    Zdata = tiff.data.scanData.ZData
    header = tiff.data.scanHeader.scanHeader
    dshape = (int(header['height'][0]), int(header['width'][0]))
    tiff_image = np.reshape(Zdata,dshape)
    tiff_image = np.flipud(tiff_image)

    result_1, _, _ = cal_histogram_range(tiff_image)
    print(f'region histogram range = {result_1:.3f} nm')
    result_2, _, _ =cal_histogram_range_with_infimum(tiff_image)
    print(f'region histogram(with infimum) range = {result_2:.3f} nm')