"""
in-situ TEM Toolbox - Diffraction pattern analysis

Assembles of functions related to process FFT and diffraction patterns

Example:
    
    import insitu_Diff as Diff  
    Diff.f2tif(path,1)

Created on Aug 8 2024
@author: Meng Li
"""
import numpy as np
import cv2
# from skimage.feature import blob_dog
import matplotlib.pyplot as plt
# from scipy import ndimage
import math

# import pandas as pd

def next_power_of_2(x):
    return 1 if x == 0 else 2**(x - 1).bit_length()

def pad_to_nearest_power_of_2(image):
    import cv2
    # Get original dimensions
    h, w = image.shape
    size=max(h,w)

    # Calculate the nearest power of 2 dimensions
    new_size = next_power_of_2(size)

    # Calculate padding amounts
    top_pad = (new_size - h) // 2
    bottom_pad = new_size - h - top_pad
    left_pad = (new_size - w) // 2
    right_pad = new_size - w - left_pad

    # Apply padding
    padded_image = cv2.copyMakeBorder(
        image,
        top=top_pad, bottom=bottom_pad,
        left=left_pad, right=right_pad,
        borderType=cv2.BORDER_CONSTANT,
        value=0  # Black color
    )

    return padded_image

def getFFT(im):
    """
    Function to get FFT from grayscale image, auto pad the image to keep the image size  2^n
    Input: grayscale image (np array)
    Output: grayscale FFT with 2^n size, 32 bit
    """

    pad_im=pad_to_nearest_power_of_2(im)


    fft_image = np.fft.fft2(pad_im)
    fft_shifted = np.fft.fftshift(fft_image)
    fft = 20 * np.log(np.abs(fft_shifted))
    return fft

def FFTstack(path):
    """
    Function to get FFT of the stack
    input: path of the image stack
    output: FFT stack with 2^n size 8 bit images
    """
    import tifffile
    import tqdm
    import cv2

    stack = tifffile.imread(path)
    nFrames,h,w = stack.shape #make sure the order of h and w!
    pathout=path[:-4]+'_FFT.tif'

    print("------------------------------------------")
    print('Calculating FFT of each frame')
    with tifffile.TiffWriter(pathout,bigtiff=True) as tif:
        for i in tqdm.tqdm(range(nFrames)):
            fr=stack[i]
            fft=getFFT(fr)
            #convert the 32 bit image to 8 bit
            normalized_image = cv2.normalize(fft, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)
            # Convert the normalized image to 8-bit
            fft_8bit = np.uint8(normalized_image)
            tif.write(fft_8bit, contiguous=True)

def FFT2Polar(fft):
    """
    Function to convert FFT to polar image, y= 180˚, x= Rho
    Input: grayscale fft image (np array)
    """
    import cv2
    #function to convert fft image to polar coordinates from the center from 0-180 deg
    center = (fft.shape[1] // 2, fft.shape[0] // 2)
    polar_fft = cv2.linearPolar(fft, center, center[0], cv2.WARP_FILL_OUTLIERS)
    # polar_fft= polar_fft[:,:center[1]]
    #resize the image to realistic size
    resized_polar_fft = cv2.resize(polar_fft, (center[0], 360))
        # plt.imshow(resized_polar_fft)
    
    p_fft1=resized_polar_fft[:180,:]
    p_fft2=resized_polar_fft[180:,:]
    p_fft=p_fft1+p_fft2
    return p_fft

def FFTidentify(fft,pathout,scale=1,sigmaMax=2.5,sigmaMin=2.4,thrs=0.96,sigma_ratio=10):
    """
    Function to identify DPs from FFT 
    input: 
        fft: image
        scale: scale of the original image : 1 px  = XXX nm
        parameters for blobs detection: sigmaMax, sigmaMin, thrs, sigma_ratio
        pathout: path for output image (.tif) and csv

    output:
        idnetified image
        indexed csv



    """
    import math
    import matplotlib.pyplot as plt
    from scipy import ndimage
    from skimage.feature import blob_dog, blob_log, blob_doh
    from math import sqrt
    from insituTEM import insitu_IO as IO
    
    h,w = fft.shape
    x_c=int(w/2)
    y_c=int(h/2)
    r=int(x_c/2)   
    if w>320:
        fft_crop=fft[x_c-r:x_c+r,y_c-r:y_c+r]
        image= ndimage.maximum_filter(fft_crop, size=2, mode='wrap')
    else:
        fft_crop = fft
        image= ndimage.maximum_filter(fft_crop, size=3, mode='wrap')




    blobs_dog = blob_dog(image, max_sigma=sigmaMax,min_sigma=sigmaMin, threshold=thrs,sigma_ratio=sigma_ratio,exclude_border=True)
    blobs_dog[:, 2] = blobs_dog[:, 2] * sqrt(2)

    title = 'DP detection'

    fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True)
    ax = axes.ravel()

    ax[0].set_title('FFT_')
    ax[0].imshow(fft_crop)
    # ax[0].set_axis_off()
    ax[1].set_title(title)
    ax[1].imshow(image)


    plt.tight_layout()
    
   
    data=np.zeros((len(blobs_dog),6))
    data[:,0]=blobs_dog[:,1]
    data[:,1]=blobs_dog[:,0]

    for j in range(len(blobs_dog)):
        y=data[j,1]-r
        x=data[j,0]-r

        data[j,2]=math.sqrt(x**2+y**2)
        # if data[j,2] < 10:
        #     data[j,2]=0
        #     theta = 0
        # else:
        #     theta=math.atan(y/x)
        theta=math.atan(y/x)
        data[j,3]=math.degrees(theta)

    data[:,4]=data[:,2]*scale
    data[:,5]=1/data[:,4]

    plt.scatter(data[:,0],data[:,1],marker='+',color='red')
    for i, (x_i, y_i) in enumerate(zip(data[:,0],data[:,1])):
        plt.annotate(str(i+1), (x_i, y_i), textcoords="offset points", xytext=(0,10), ha='center',color='white')
    pathoutcsv=pathout[:-4]+'.csv'
    IO.writeCSV(pathoutcsv,data,fmt='%1.3f')
    plt.savefig(pathout)

# def getFFT(im):
#     """
#     Function to calculate FFT of the image
#     input: im- np array of image data
#     output: fft image
#     """
#     # Get the original dimensions of the image
#     h, w = im.shape

#     size = min(h,w) #in case h w don't match
#     im_crop=im[:size,:size] #crop image 
#     fft_image = np.fft.fft2(im_crop)
#     fft_shifted = np.fft.fftshift(fft_image)
#     fft = 20 * np.log(np.abs(fft_shifted) + 1e-10)  # Add small constant to avoid log(0)
    
#     return fft

def spotFFT(fft, sigmaMax=2.5, sigmaMin=2.45, sigma_ratio=20, thrs=0.98, R_tol=2, fftfilter=3, tolerance=5, display=True):
    """
    Function to get diffraction spot positions in FFT
    Input: fft- FFT image
            sigmaMax, sigmaMin, sigma_ratio, thrs - Parameters for blob detection
            R_tol - tolerance of center spot
            fftfilter- filtersize to denoise FFT to make diffraction spots more visible
            tolerance - tolerance for symmetry matching
    
    Output: 
        Data of all detected spots:    X(px)   Y(px) 
        Image labeling the spots
     """
    from skimage.feature import blob_dog
    import matplotlib.pyplot as plt
    from scipy import ndimage
    import math
  
    import numpy as np
    import pandas as pd

    h, w = fft.shape
    size = max(h, w)
    scale_fft = 1 / (size)
    
    x_c, y_c = int(w / 2), int(h / 2)
    image = ndimage.maximum_filter(fft, size=fftfilter, mode='wrap')
    blobs_dog = blob_dog(image, max_sigma=sigmaMax, min_sigma=sigmaMin, threshold=thrs, sigma_ratio=sigma_ratio, exclude_border=True)
    blobs_dog[:, 2] *= math.sqrt(2)
    

    
    data = []
    blob_set = set()
    
    for j in range(len(blobs_dog)):
        y, x = blobs_dog[j, 1] - y_c, blobs_dog[j, 0] - x_c
        distance = math.sqrt(x**2 + y**2)
        if distance < R_tol:
            continue
        
        data.append([blobs_dog[j, 1], blobs_dog[j, 0]])
        blob_set.add((blobs_dog[j, 1], blobs_dog[j, 0]))
    
    # Ensure symmetry
    for x, y in data:
        sym_x = 2 * x_c - x
        sym_y = 2 * y_c - y
        
        if not any(math.isclose(sym_x, bx, abs_tol=tolerance) and math.isclose(sym_y, by, abs_tol=tolerance) for bx, by in blob_set):
            data.append([sym_x, sym_y])
            blob_set.add((sym_x, sym_y))
    
    data = np.array(data)

    if display == True:
        fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True)
        ax = axes.ravel()
        ax[0].set_title('FFT')
        ax[0].imshow(fft)
        ax[1].set_title('DP detection')
        ax[1].imshow(image)
        plt.tight_layout()
        plt.scatter(data[:, 0], data[:, 1], marker='+', color='red')
        for i, (x_i, y_i) in enumerate(zip(data[:, 0], data[:, 1])):
            plt.annotate(str(i+1), (x_i, y_i), textcoords="offset points", xytext=(0,10), ha='center', color='white')
        
    return data

def DP_mask(fft,circle_radius=10, center_rad=1,sigmaMax=2.5, sigmaMin=2.45, sigma_ratio=20, thrs=0.98, R_tol=2, fftfilter=3, tolerance=5):
    """
    Create a mask image with circles drawn at detected diffraction points.
    
    Parameters:
        image_shape (tuple): Shape of the image (height, width).
        points (np.ndarray): Array of detected points [[x1, y1], [x2, y2], ...].
        circle_radius (int): Radius of the circles to draw.
        center_rad: radius of center circle
        sigma...: parameters for detect DP
    
    Returns:
        np.ndarray: Mask image with drawn circles.
    """

    mask = np.zeros(fft.shape, dtype=np.uint8)
    points= spotFFT(fft,sigmaMax,sigmaMin,sigma_ratio,thrs,R_tol,fftfilter,tolerance,display=False)
    for x, y in points:
        cv2.circle(mask, (int(x), int(y)), circle_radius, 255, thickness=-1)

    if center_rad != 0:
        center=int(fft.shape[0]/2)
        cv2.circle(mask, (center, center), center_rad, 255, thickness=-1)
    return mask

def maskedIFFT(image,mask,mix_ratio=0.3):
    """
    generate masked image from FFT
    Input: original image; mask; mix_ratio for IFFT filtering
    Output: FFT_image,  IFFT image, IFFT filterd image
    Example: fft_image, ifft_image,filtered_image=maskedIFFT(image,mask,mix_ratio)
    """
    h,w=image.shape

    if h !=w:
        d=min(h,w)
        image=image[0:d,0:d]

    # Perform FFT on the image
    fft = cv2.dft(np.float32(image), flags=cv2.DFT_COMPLEX_OUTPUT)
    fft_shifted = np.fft.fftshift(fft)
    fft_image=np.log(1 + cv2.magnitude(fft_shifted[:, :, 0], fft_shifted[:, :, 1])) #using fft_image to apply mask results in empty image
    
    # Ensure the mask image has the same size as the FFT
    mask_image = cv2.resize(mask,(fft_shifted.shape[1], fft_shifted.shape[0]))

    # Apply the mask to the FFT
    masked_fft_shifted = cv2.bitwise_and(fft_shifted, fft_shifted, mask=mask_image)

    # Perform IFFT on the masked FFT
    masked_fft = np.fft.ifftshift(masked_fft_shifted)
    masked_image = cv2.idft(masked_fft)
    masked_image = cv2.magnitude(masked_image[:, :, 0], masked_image[:, :, 1])
    
    #convert to 8bit images
    fft_image=cv2.normalize(fft_image, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
    ifft_image=cv2.normalize(masked_image, None, 0, 128, cv2.NORM_MINMAX, cv2.CV_8U) #cv2.normalize(src, dst, alpha, beta, norm_type, dtype)
    
    # # Normalize the pixel values to the range of 0-255
    ifft_image_8 = cv2.normalize(ifft_image, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
    

    # Add IFFT to original image with mix_ratio
    masked_image_normalized = cv2.normalize(ifft_image, None, 0, int(255*mix_ratio), cv2.NORM_MINMAX, cv2.CV_8U)
    image_normalized = cv2.normalize(image, None, 0, int(255*(1-mix_ratio*0.5)), cv2.NORM_MINMAX, cv2.CV_8U)
    
    filtered_image=cv2.add(image_normalized,masked_image_normalized)


    return fft_image, ifft_image, filtered_image


def measureFFT(fft,path,scale,j=0,sigmaMax=2.5,sigmaMin=2.4,thrs=0.97,R_tol=2,fftfilter=2,savefig=True,displayfig=True):
    """
    Function to measure diffraction spots in FFT
    Input: fft- FFT image
            path: path to save data
            sigmaMax, sigmaMin, thrs - Parameters for blob detection
            R_tol - tolerance of center spot
            
    Output: 
        Data of all detected spots: Point   X(px)   Y(px)  distance(px)  angle(˚)  1/d (1/nm)  d (nm)
        Image labeling the spots

    """
    from skimage.feature import blob_dog
    import matplotlib.pyplot as plt
    from scipy import ndimage
    import math
  
    import numpy as np
    import pandas as pd

    h,w = fft.shape
    size=max(h,w)
    scale_fft=1/(scale*size)
    
    x_c=int(w/2)
    y_c=int(h/2)
    center=x_c/2 #center of cropped_fft
    
    fft_crop=fft[x_c-int(center):x_c+int(center),y_c-int(center):y_c+int(center)]
    image= ndimage.maximum_filter(fft_crop, size=fftfilter, mode='wrap')
    blobs_dog = blob_dog(image, max_sigma=sigmaMax,min_sigma=sigmaMin, threshold=thrs,sigma_ratio=10,exclude_border=True)
    blobs_dog[:, 2] = blobs_dog[:, 2] * math.sqrt(2)
    
    
    title = 'DP detection'
    title1= 'FFT_'+str(j)
    fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True)
    ax = axes.ravel()
    
    ax[0].set_title(title1)
    ax[0].imshow(fft_crop)
    ax[1].set_title(title)
    ax[1].imshow(image)
    
    plt.tight_layout()
    
    pathout=path[:-4]+'_'+str(j)+'_FFT.csv'
    pathout_img=path[:-4]+'_'+str(j)+'_DP.tif'
    
    
    data=np.zeros((len(blobs_dog)-1,8))
    
    
    idx=0
    for j in range(len(blobs_dog)):
    
    
        y=blobs_dog[j,1]-center
        x=blobs_dog[j,0]-center
    
        
        distance=math.sqrt(x**2+y**2)#px distance to center
        if distance<R_tol: #skip center point
            continue

        theta=math.atan2(x,y)#atan gives +-90˚ values, atan2 gives 180
        angle=0-math.degrees(theta)
        data[idx,:7] = [j , blobs_dog[j, 1], blobs_dog[j, 0], distance, angle, distance * scale_fft, 1 / (distance * scale_fft)]
        idx += 1
    
    
    plt.scatter(data[:,1],data[:,2],marker='+',color='red')
    for i, (x_i, y_i) in enumerate(zip(data[:,1],data[:,2])):
        plt.annotate(str(i+1), (x_i, y_i), textcoords="offset points", xytext=(0,10), ha='center',color='white')
    
    
    data[:,7]=data[:,4]-data[0,4]
    
    columns= ["Point","X(px)","Y(px)","distance(px)","angle to X(˚)","1/d (1/nm)","d (nm)","angle to P1(˚)"]
    # Set pandas global display format for floats
    pd.options.display.float_format = '{:.3f}'.format
    # Create a DataFrame
    df = pd.DataFrame(data, columns=columns)
    df.to_csv(pathout,index=False)
    if savefig:
        plt.savefig(pathout_img,bbox_inches='tight', pad_inches=0.1,dpi=150)
    if displayfig:
        plt.show()

    return df




def FFT2Polar(fft):
    """
    Function to convert fft image to polar image
    input: fft pimage 
    output: polar FFt image
    """
    import cv2
    #function to convert fft image to polar coordinates from the center from 0-180 deg
    center = (fft.shape[1] // 2, fft.shape[0] // 2)
    polar_fft = cv2.linearPolar(fft, center, center[0], cv2.WARP_FILL_OUTLIERS)
    # polar_fft= polar_fft[:,:center[1]]
    #resize the image to realistic size
    resized_polar_fft = cv2.resize(polar_fft, (center[0], 360))
        # plt.imshow(resized_polar_fft)
    
    p_fft1=resized_polar_fft[:180,:]
    p_fft2=resized_polar_fft[180:,:]
    p_fft=p_fft1+p_fft2
    return p_fft

def maskedFFT(image,mask):
    """
    generate masked image from FFT
    Input: original image; mask
    Output: FFT_image, filtered IFFT image
    """
    h,w=image.shape

    if h !=w:
        d=min(h,w)
        image=image[0:d,0:d]

    # Perform FFT on the image
    fft = cv2.dft(np.float32(image), flags=cv2.DFT_COMPLEX_OUTPUT)
    fft_shifted = np.fft.fftshift(fft)
    fft_image=np.log(1 + cv2.magnitude(fft_shifted[:, :, 0], fft_shifted[:, :, 1]))
    

    # Ensure the mask image has the same size as the FFT
    mask_image = cv2.resize(mask,(fft_shifted.shape[1], fft_shifted.shape[0]))

    # Apply the mask to the FFT
    masked_fft_shifted = cv2.bitwise_and(fft_shifted, fft_shifted, mask=mask_image)


    # Perform inverse FFT on the masked FFT
    masked_fft = np.fft.ifftshift(masked_fft_shifted)
    masked_image = cv2.idft(masked_fft)
    masked_image = cv2.magnitude(masked_image[:, :, 0], masked_image[:, :, 1])
    
    #convert to 8bit images
    fft_image=cv2.normalize(fft_image, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
    masked_image=cv2.normalize(masked_image, None, 0, 128, cv2.NORM_MINMAX, cv2.CV_8U) #cv2.normalize(src, dst, alpha, beta, norm_type, dtype)
    


    return fft_image, masked_image


# def FFTstack(stackpath,ndfilersize=3,fps=15):
#     import tqdm
#     from scipy import ndimage
#     import tifffile
#     from insituTEM import insitu_IO as IO
#     """
#     Apply FFT to the entire stack. 
#     """
#     pathout=stackpath[:-4]+'_FFT.tif'
#     stack =tifffile.imread(stackpath)
    
#     nFrames, h,w = stack.shape
#     print("------------------------------------------")
#     print("Getting FFT to the stack:")   
#     with tifffile.TiffWriter(pathout,bigtiff=True) as tif:
#         for i in tqdm.tqdm(range(nFrames)):
#             fr=stack[i]
#             fft=Diff.getFFT(fr)
#             if ndfilersize!=0:
#                 fft = ndimage.maximum_filter(fft, size=ndfilersize, mode='wrap')
#             fft_image=cv2.normalize(fft, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
#             tif.write(fft_image, contiguous=True)
#     IO.tiff2avi(pathout,fps=15)

#     print("Conversion done!")