# -*- coding: utf-8 -*-
"""

in-situ TEM Toolbox - PreProcess
By: Meng Li and Michael
Assembles of functions related to movie preprocess 
Example:
    
    import insitu_Preprocess as PP
    PP.blur_score(image)

Created on Tue Jun 11 10:25:25 2019


"""
import numpy as np
import cv2
from scipy import signal
from PIL import Image 
import matplotlib.pyplot as plt
import tifffile
import tqdm

def estimate_dup(A,B,threshold=1):
    """
    Function to detect if two images are duplicated
    """
    D=np.subtract(A,B)
    score=np.mean(D)#mean works better than sum due to noises
    return bool(score<threshold)


def estimate_blank(image,threshold=50):
    """
    funciton to detect blank frames
    returns 1 if is blank
    """
    score = np.sum(image[10:15,:])
    
    return bool(score < threshold)

def blur_score(image):
    """
    funciton to get blur_score of image
    returns blur score
    """
    if image.ndim == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    blur_map = cv2.Laplacian(image, cv2.CV_64F)
    score = np.var(blur_map)
    return score

def median_blur (score_list, i,thrs=150,wid=21):
    """
    funciton to detect if frame i is blur
    input: score_list: 1D array of scores; i: frame number; thrs: tolerance of blur threshold; wid： median filter width
    returns 1 if is blur
    """
    medthreshold=signal.medfilt(score_list,wid)
    #The threshold should be moved because there is no need to recreate the threshold everytime
    score = score_list[i]
    medthres = medthreshold[i]
    return medthres,bool(abs(score-medthres)>=thrs)

    
def detect_blur_fft(image, size=20, thresh=15,vis=False):
    """
    Function to detect if one frame is blurred using FFT
    Example:
    score,is_blur = PP.detect_blur_fft(image, size=20, thresh=15,vis=False)
    Input:
    image: grayscale image, default size: 128*128 px
size: filter size to get the background of FFT
    thresh: threshold value for blur scores, smaller score, more blurry
    Output:
    mean: mean value of the frame
    is_blur: whether the frame is blur
    """
    from insituTEM import insitu_DP as DP
    # grab the dimensions of the image and use the dimensions to derive the center (x, y)-coordinates
    # h, w = image.shape
  

    magnitude=DP.getFFT(image)
    h,w = magnitude.shape
    cX, cY = w //2, h//2
    bar=int(h/100)
    center=bar*10
        #use only the center of the FFT
    crop_h,crop_w=(h-cX)//2, (w-cY)//2
    

    # compute the background of FFT to make the diffraaction spots more obvious
    bg = cv2.blur(magnitude,(size,size))

    enhanced =cv2.subtract(magnitude, bg)
    #remove noise by mean
    # enhanced=enhanced[crop_h:(crop_h+cY),crop_w:(crop_w+cX)]

    #remove white noise from the substracted image
    enhanced[enhanced<20]=0
    
    # Check the data type and convert if necessary
    if enhanced.dtype != np.uint8:
        # Normalize the image to the range [0, 255] if it has a different range
        enhanced = cv2.normalize(enhanced, None, 0, 255, cv2.NORM_MINMAX)
        # Convert the image to uint8
        enhanced = enhanced.astype(np.uint8)
    
    # Apply median filter
    enhanced = cv2.medianBlur(enhanced, 3)

    # enhanced = cv2.medianBlur(enhanced, 3)
    # #remove center spot
    enhanced[(cY - center): (cY + center), cX - center:cX + center] =0
    enhanced[(cY - bar): (cY + bar), :]=0
    enhanced[:, (cX - bar): (cX + bar)]=0
    # enhanced[


    croppedFFT=enhanced[crop_h:(crop_h+cY),crop_w:(crop_w+cX)]

    if vis:
        (fig, ax) = plt.subplots(1, 2,figsize=(10, 5) )
        ax[0].imshow(image, cmap="gray")
        ax[0].set_title("Input")
        # ax[0].set_xticks([])
        # ax[0].set_yticks([])
        # display the magnitude image
        ax[1].imshow(croppedFFT,cmap="viridis")
        ax[1].set_title("Filtered FFT")
        # ax[1].set_xticks([])
        # ax[1].set_yticks([])
        
        # show our plots
        plt.show()
    # mean = np.mean(enhanced)*1000
    mean = np.mean(croppedFFT)*1000
#     print(mean)
    # the image will be considered "blurry" if the mean value of the magnitudes is less than the threshold value
    return (mean,mean<thresh)
        
def tiff_crop(moviepath):
    '''  
    function to crop a tiff stack
    input: path
    output: cropped tiff stack
    '''
    from skimage import io
    from tqdm import tqdm
    import tifffile
    moviepathout = moviepath[:-4] + "_Crop.Tif"
    im = cv2.imread(moviepath)
    #only reads the first image in the Tiff stack
    # Select ROI
    r = cv2.selectROI(im, False, False)
    imCrop = im[int(r[1]):int(r[1]+r[3]), int(r[0]):int(r[0]+r[2])]
    cv2.imshow("Image", imCrop)
    #Displays the cropped image
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    
    #Read the Tiff file properly
    im = io.imread(moviepath)
    nframes, height, width = im.shape
    iterable = range(nframes)
    
    for num in tqdm(iterable): 
        # Crop image
        imCrop = im[num, int(r[1]):int(r[1]+r[3]), int(r[0]):int(r[0]+r[2])]
     
        # Display cropped image
        
        #save cropped image
        if num == 0:                    
            tifffile.imsave(moviepathout, imCrop, append=False)
        else:
            tifffile.imsave(moviepathout, imCrop, append=True)
            
def rotateFr(im,angle=0,expand=True,bgColor=150):
    """
    funciton to rotate frame clockwise by angle degree with filled color
    input: im-greyscale image frame
    returns rotated image
    """
#    from PIL import Image
#    angle = 73
    a=360-angle
    img = Image. fromarray(im)
    fr=img.rotate(a,expand=expand,fillcolor=bgColor)
    imgRt = np.array(fr) 
    return imgRt

def removeDP(img,position):
    """
    funciton to  remove deadpoint in greyscale image using automated color based on surounding pixels

    input: im-greyscale tiff stack
    position:position of DP[x1 y1 w h], or DPs[[x1 y1 w1 h1],[x2 y2,w2,h2],...]
    returns rotated image
    """
    #remove deadpoint in greyscale image using automated color based on surounding pixels
    #position:[x1 y1 w h]
    
    
    x1=position[0]
    y1=position[1]
    w=position[2]
    h=position[3]

    x2=x1+w
    y2=y1+h
    c1=np.mean(img[y1-1,x1:x2])
    c2=np.mean(img[y2+1,x1:x2])
#    c3=np.mean(img[x1-1,y1:y2])
#    c4=np.mean(img[x2+1,y1:y2])
#    color=int((c1+c2+c3+c4)/4)
    color = int((c1+c2)/2)
#    color = 177
    cv2.rectangle(img,(x1,y1),(x2,y2),color,-1)  
    
def removeDPs(img,DP):
    """"
    Function to remove multiple DPs in one image
    DP[[x1 y1 w1 h1],[x2 y2,w2,h2],...]
    """
    nDP=len(DP)
    if nDP==1:
        removeDP(img,DP)
    else:
        for j in range(nDP):
            position=DP[j]
            removeDP(img,position)


def BlurDetection(path,wid=31):
    """
    Function to remove blurry frames based on the blur score
    Input: tif file path
    Output: deblurred tif stack
    """
    from insituTEM import insitu_IO as IO

    indexout=path[:-4]+'_DB-index.csv'
    DBout=path[:-4]+'_DB.csv'
    DBout2=path[:-4]+'_Blur.csv'

    stack = tifffile.imread(path)
    nFrames,h,w = stack.shape #make sure the order of h and w!
    index = []#index of remaining frames
    score = []

    print("------------------------------------------")
    print("Calculating blur scores!")  

    for i in tqdm.tqdm(range(nFrames)):
        fr=stack[i]
        s = blur_score(fr)
        score.append(s)
        index.append(i)
        
    length = len(score)
    Blur=np.zeros((length,3))#[i,blur_score,blur_th]
    Blur[:,0]=index
    Blur[:,1]=score            

    plt.plot(score,'gray')
    plt.show()     

    print("------------------------------------------")
    print("Check blur detection!")  
  
    thrs=int(input( 'Input blur threshold(default: 70): ')) 
    while thrs!=0:
        index2=[]

        for i in tqdm.tqdm(range(len(score))):
            medthres,is_blur=median_blur (score, i,thrs,wid=wid)#wid must be odd
    #        DB[index[i],3]=is_blur
            Blur[i,2]=medthres
            if is_blur == 0:
                index2.append(index[i])
                
        plt.plot(Blur[:,0],Blur[:,1],'gray',Blur[:,0],Blur[:,2]-thrs,'r',Blur[:,0],Blur[:,2]+thrs,'r')
        plt.show()
        thrs=int(input( 'Press 0 if you are satisfied, press corrected threshold if you want to change:  ')) 
    
    IO.writeCSV(DBout2,Blur)  
        
    result= np.zeros((len(index2),2)) 
    for i in range(len(index2)):
        result[i,0]=i
        result[i,1]=index2[i]
        
    IO.writeCSV(indexout,result)

    print("------------------------------------------")
    print("Save blur removed file") 
    pathout2=path[:-4]+'_DBlur.tif'

    with tifffile.TiffWriter(pathout2,bigtiff=True) as tif:
        for j in tqdm.tqdm(range(len(index2))):
            i=int(result[j,1])
            fr=stack[i]

            tif.write(fr, contiguous=True)
    print("------------------------------------------")
    print("Blur removed!") 
    print("==========================================")


def movie_preprocess(path,remove_DP=True,DP=[[0, 0, 1, 1],[0,0,2,2]],Crop=True,cropboarder=[0,0,0,0],Trim=False,start_time_mmss="00:00",end_time_mmss="00:00",grayscale=True,exp_time=0):
    """
    Function to preprocess the raw TEM movie to tiff
    Input:
        path: file path of the movie
        remove DP: if there are constant dead point in the movie to remove, 
            if so, input the DP value. 
        Crop: if crop the movie to remove the boarders
            if so, input the cropboarder: [left, top, right, bottom]
        Trim: if trim the movie length
            if so, input the start and end time in mm:ss
        grayscale: if convert the movie to grayscale
        exp_time: if remove the duplicated frame based on the real exposure time. 
            if no need to remove duplicates, put exp_time = 0
    """
    # print("=========================================")
    print("ETEM movie preprocessor start!")

    import moviepy.editor as mp
    import moviepy.video.fx.all as vfx
    from insituTEM import insitu_IO as IO


    video = mp.VideoFileClip(path)
    if Trim == True:
        # Convert mm:ss to seconds
        start_time_seconds = sum(x * int(t) for x, t in zip([60, 1], start_time_mmss.split(":")))
        end_time_seconds = sum(x * int(t) for x, t in zip([60, 1], end_time_mmss.split(":")))

        video = video.subclip(start_time_seconds, end_time_seconds)# cut the clip between t=0 and 28 secs.

    fps = int(video.fps)
    # fps= 15
    w = int(video.w)
    h = int(video.h)
    nFrames = int(fps*video.duration)

    print("fps=", fps, "   w =",w, "   h=", h, "duration= ",video.duration)
    print("Total number of Frames=", nFrames)



    #Preview first frame
    print("------------------------------------------")
    print("Preview the processed frame: press 0 if you want to continue: ")

    if Crop == True:
        video =  video.fx(vfx.crop, x1=cropboarder[0], y1=cropboarder[1], x2=w-cropboarder[2], y2=h-cropboarder[3])#crop movie snapshot
    
    fr = video.get_frame(mp.cvsecs(1/fps))
    
    if grayscale==True: 
        fr = cv2.cvtColor(fr, cv2.COLOR_BGR2GRAY)

    if remove_DP==True:
        removeDPs(fr,DP)

    IO.showFr(fr)
    prompt = input("Input 0 if you want to continue, other keys to break :")

    if prompt != "0":
        print("Stopped by user. Modify the cropboarder values and restart.")
        return
    else: 
        #continue to apply to the full movie
        #remove duplicate frames with known fps conversion
        if exp_time==0:
            dup=1
        else:
            dup= fps*exp_time

        nFramesOut=int(nFrames/dup)
        print ('nFramesOut= ', nFramesOut)
        print('nFramesIn= ',nFrames)


        pathout=path[:-4]+'_CutDB.tif'

        print("------------------------------------------")
        print("Saving duplicate removed processed frames~")   
        

        with tifffile.TiffWriter(pathout,bigtiff=True) as tif:
            for j in tqdm.tqdm(range(nFramesOut)):
                i=j*dup
                fr = video.get_frame(mp.cvsecs(i/fps))
                if grayscale==True: 
                    fr = cv2.cvtColor(fr, cv2.COLOR_BGR2GRAY)
                if remove_DP==True:
                    removeDPs(fr,DP) 
                tif.write(fr, contiguous=True)

        print("Conversion done!")