"""
in-situ TEM Toolbox - TEM raw data conversion

Assembles of functions related to read and process raw dm4, dm3 or FEI TIA or Ceta datas.

Example:
    
    import insitu_EMRaw as EM  
    IO.f2tif(path,1)

Created on Aug 8 2024
@author: Meng Li
"""
import numpy as np
import tqdm
import ncempy
import tifffile
import os

def img_2_8bit(img):
    """
    Convert image to Uint8bit with custom min and max intensity: vmin to 0, vmax=255

    """
    vmax = int(np.median(img)*2)
    vmin = int(np.min(img))
    if (vmax-vmin)==0:
        scale =1 
    else:
        scale = 255/(vmax-vmin)
    imgOut=(img*scale).astype(np.uint8)
    return imgOut

def getScaleBar(barLenPix,scale):
    """
    Function to convert truescale into str with correct format
    """
    truescale= barLenPix/scale
    
    # Find the nearest power of 10 to the truescale
    nearest_power = 10 ** np.floor(np.log10(truescale))

    # Determine the most appropriate scale bar value
    if truescale >= 5 * nearest_power:
        scalebar = 5 * nearest_power
    elif truescale >= 2 * nearest_power:
        scalebar = 2 * nearest_power
    else:
        scalebar = nearest_power

    # Calculate the scale in pixels
    scalepix = int(scalebar * scale)

    # Determine the appropriate unit for the scale bar text
    if scalebar >= 1000:
        text = f"{int(scalebar/1000)} um"
    else:
        text = f"{int(scalebar)} nm"
    return scalepix, text

def AddScalePix(img,barLenPix =250, scale = 34.5, px=0.02,py=0.96, barthick=5,lw=1,stroke=True,tscale=1,textThick=4):
    """
    Function to add scale to image
    Input:
        img: greyscale image
        scaleLen: desired scale length
        scale: scale of the image: px/nm: 1000kX- 49.5; 700kX - 34.5
        (px,py): position of the scale bar: r(0-1)
        lw: boarder of the text
        stroke=1: add stroke; stroke=0 No stroke
 
    """
    import cv2
    h,w=img.shape

    scalepix, text = getScaleBar(barLenPix, scale)

    # w_scale=int(scale*barLen)
    x1=int(w*px)
    y1=int(h*py)
    x2=x1+scalepix
    y2=y1+barthick


    #find bcakground color and make the reverse as scalebar color
    img_bg=img[y1:y2,x1:x2]
    mean_intensity=np.mean(img_bg)

    if mean_intensity<128:
        color = 255
    else:
        color = 0

    bcolor=255-color

    # text = str(barLen)+' nm'
    font = cv2.FONT_HERSHEY_SIMPLEX#CV_FONT_HERSHEY_SIMPLEX normal size sans-serif font
    if stroke == True:
        cv2.putText(img,text,(x1,y1-8), font, tscale, bcolor,int(textThick+lw), cv2.LINE_AA) #Stroke  
        cv2.rectangle(img,(x1,y1),(x2,y2),bcolor,2)#boarder color

        
    cv2.rectangle(img,(x1,y1),(x2,y2),color,-1)
    
    cv2.putText(img,text,(x1,y1-8), font,tscale, color, textThick, cv2.LINE_AA)


def DM2Tiff(path,addscale = True,scalebar_pix=200, tscale=1.2, textThick=4):
    """
    Function to convert dm4 file to tiff. 
    If the input is .dm4  stacked image, convert the stack to tiff stack.
    If the input is .dm4 single image, convert the image to tiff image.
    If the input is a folder with .dm4 single images, convert the folder into tiff stack.
    """
    import ncempy
    import tifffile



    if path.endswith(('.dm4', '.dm3')):
        
        dm = ncempy.read(path)
        pathout =path[:-4]+'.tif'
     
        if len(dm['data'].shape)==2:
            #single image
            print("convert .dm3/4 file to tiff image")
            #for single image
            frame=dm['data']
            img=img_2_8bit(frame)
    
            if addscale==True: 
            
                scale = 1/(dm['pixelSize'][0]*1e3)
    
                AddScalePix(img,barLenPix =scalebar_pix, scale = scale, px=0.02,py=0.98, barthick=5,lw=1,stroke=True,tscale=tscale,textThick=textThick)
            
            tifffile.imwrite(pathout,img, append=False)
        else:
        #image stacks
            print("convert .dm4 stack to tiff stack")
            imgs=dm['data']
            nFrames,w,h =dm['data'].shape

                 
            with tifffile.TiffWriter(pathout,bigtiff=True) as tif:
                for i in tqdm.tqdm(range(nFrames)):
                    img=imgs[i]
                    # img = cv2.convertScaleAbs(img)
                    # img = (img/ 65535.0 * 255).astype(np.uint8)
                    img=img_2_8bit(img)
    
                    # img =cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
                    if addscale==True: 
                        scale = 1/(dm['pixelSize'][0]*1e3)
                        AddScalePix(img,barLenPix =scalebar_pix, scale = scale, px=0.02,py=0.98,barthick=5,lw=1,stroke=True,tscale=tscale,textThick=textThick)
                
                    tif.write(img, contiguous=True)

    else: 
        files = [os.path.join(path, f) for f in os.listdir(path) if f.endswith(('.dm4', '.dm3'))]
        #convert each file to tiff stack
        print('Convert folder with dm3/4 to tiff stack')
        nFrames=len(files)
        pathout=path+'.tif'
        with tifffile.TiffWriter(pathout,bigtiff=True) as tif:
            for file_path in tqdm.tqdm(files):
                
                dm = ncempy.read(file_path)
                frame=dm['data']
                img=img_2_8bit(frame)
        
                if addscale==True: 
                
                    scale = 1/(dm['pixelSize'][0]*1e3)
        
                    AddScalePix(img,barLenPix =scalebar_pix, scale = scale, px=0.02,py=0.98, barthick=5,lw=1,stroke=True,tscale=tscale,textThick=textThick)
                
                tif.write(img, contiguous=True)
    print("Convertion done!")


def GMS2folder(root_dir):

    import os
    import shutil
    import re
    import tqdm
    """
    Function to convert GMS saved in situ data into one folder named by frame name

    input: path to root_dir : the folder before Hour_00
    output: dm4 images saved in one folder
    """

    # Define the target directory
    target_dir = root_dir + '_output'
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)
    
    # Initialize a list to hold file paths and their corresponding timestamps
    images = []
    
    # Regular expression to extract hour, minute, second, and frame number from filenames
    pattern = re.compile(r"Hour_(\d+)_Minute_(\d+)_Second_(\d+)_Frame_(\d+)\.dm4")
    
    print("Converting GMS datasets into one folder...")
    
    # Walk through the directory structure and count frames per second
    fps = 0
    for root, dirs, files in os.walk(root_dir):
        for file in files:
            if file.endswith(".dm4"):  # Check if the file is an image file
                match = pattern.search(file)
                if match:
                    hour, minute, second, frame = match.groups()
                    # Create a timestamp tuple for sorting
                    timestamp = (int(hour), int(minute), int(second), int(frame))
                    # Add the file path and timestamp to the list
                    images.append((os.path.join(root, file), timestamp))
                    # Count frames in the Second_00 folder
                    if int(second) == 0:
                        fps += 1
    
    print(f"Frames per second (fps) detected: {fps}")
    
    # Sort the images by their timestamps
    images.sort(key=lambda x: x[1])
    
    # Copy and rename images to the target directory in the correct order
    for index, (filepath, timestamp) in enumerate(images, start=0):
        new_filename = f"F{index:06d}.dm4"
        new_filepath = os.path.join(target_dir, new_filename)
        shutil.copy(filepath, new_filepath)
    
    print(f"Total number of images: {len(images)}")
    print(f"All images have been saved to: {target_dir}")
    return target_dir,fps

def DMfolder2tiff(folder_path):
    """
    Function to convert entire folder of dm files into tiff
    If the dm file is single image, convert to tiff single image
    If the dm file is stack, convert to tiff stack
    """

    print("Batch convert folder into tiff or movie: ")
    print("----------------------------------------")
    # List all files in the directory
    import os
    files = os.listdir(folder_path)

    # Filter files with .ser extension
    dm_files = [file for file in files if file.endswith(('.dm4', '.dm3'))]

    # Process each .ser file
    for dm_file in dm_files:
        print("convert ", dm_file)
            # Check file size
        
        file_path = os.path.join(folder_path, dm_file)
        file_size = os.path.getsize(file_path)
        if file_size == 0:
            print(f"Skipping empty file: {dm_file}")
            continue
        DM2Tiff(file_path,scalebar_pix=250,tscale=1.2, textThick=4)
    print("----------------------------------------")
    print("Convertion done! ")


# def getScaleBar(barLenPix,scale):
#     """
#     Function to convert truescale into str with correct format
#     """
#     truescale= barLenPix/scale
    
#     # Find the nearest power of 10 to the truescale
#     nearest_power = 10 ** np.floor(np.log10(truescale))

#     # Determine the most appropriate scale bar value
#     if truescale >= 5 * nearest_power:
#         scalebar = 5 * nearest_power
#     elif truescale >= 2 * nearest_power:
#         scalebar = 2 * nearest_power
#     else:
#         scalebar = nearest_power

#     # Calculate the scale in pixels
#     scalepix = int(scalebar * scale)

#     # Determine the appropriate unit for the scale bar text
#     if scalebar >= 1000:
#         text = f"{int(scalebar/1000)} um"
#     else:
#         text = f"{int(scalebar)} nm"
#     return scalepix, text

# def AddScalePix(img,barLenPix =250, scale = 34.5, px=0.02,py=0.96, barthick=5,lw=1,stroke=True,tscale=1,textThick=4):
#     """
#     Function to add scale to image
#     Input:
#         img: greyscale image
#         scaleLen: desired scale length
#         scale: scale of the image: px/nm: 1000kX- 49.5; 700kX - 34.5
#         (px,py): position of the scale bar: r(0-1)
#         lw: boarder of the text
#         stroke=1: add stroke; stroke=0 No stroke
#         tscale: scale of the text
#         textThick: thickness of the text in pxs
 
#     """
#     import cv2
#     h,w=img.shape

#     scalepix, text = getScaleBar(barLenPix, scale)

#     # w_scale=int(scale*barLen)
#     x1=int(w*px)
#     y1=int(h*py)
#     x2=x1+scalepix
#     y2=y1+barthick


#     #find bcakground color and make the reverse as scalebar color
#     img_bg=img[y1:y2,x1:x2]
#     mean_intensity=np.mean(img_bg)

#     if mean_intensity<128:
#         color = 255
#     else:
#         color = 0

    
#     bcolor=255-color


    
#     # text = str(barLen)+' nm'
#     font = cv2.FONT_HERSHEY_SIMPLEX#CV_FONT_HERSHEY_SIMPLEX normal size sans-serif font
#     if stroke == True:
#         cv2.putText(img,text,(x1,y1-8), font, tscale, bcolor,int(textThick+lw), cv2.LINE_AA) #Stroke  
#         cv2.rectangle(img,(x1,y1),(x2,y2),bcolor,2)#boarder color

        
#     cv2.rectangle(img,(x1,y1),(x2,y2),color,-1)
    
#     cv2.putText(img,text,(x1,y1-8), font,tscale, color, textThick, cv2.LINE_AA)
# #Save image to tiff or tiff stack

def convertser(path,addscale = True,scalebar_pct=20,fps=10):
    """
    Convert ser file to image, stack to tiff stack and movie
    input: 
        path: .ser file path
        addscale: whether scale bar is added
        scalebar_pix: roughly how many pixels the scale bar takes
        scalebar_pct: roughly how much percnet the scale bar takes on the image
        fps: for .ser stack, the fps of the original data
        tscale: scale of the text for the scale bar
        textThick: thickness of the text for the scale bar
    Output:
        single .ser image: Tiff image
        .ser stack: Tiff stack  and avi movie
    """
    import ncempy
    import tifffile
    import tqdm
    import cv2

    import os
    os.environ["IMAGEIO_FFMPEG_EXE"] = "/opt/homebrew/bin/ffmpeg"
    import insituTEM.insitu_IO as IO

    # print("convert .ser file to tiff image")
    ser = ncempy.read(path)
    pathout =path[:-4]+'.tif'
 

    if len(ser['data'].shape)==2:
        #single image
        # print("convert .ser file to tiff image")
        #for single image
        img=ser['data']
        img =cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        # img = cv2.convertScaleAbs(img)
        # img = (img/ 65535.0 * 255).astype(np.uint8)

        if addscale==True: 
            w= ser['data'].shape[0]
            scalebar_pix = round(w/100) *scalebar_pct
            tscale=round(w/1024)
            textThick=round(w/1024)*2
            
            scale = 1/(ser['pixelSize'][0]*1e9)
       
            # scalebar= int(scalebar_pix*scale/100)*100
            
            AddScalePix(img,barLenPix =scalebar_pix, scale = scale, px=0.02,py=0.98, barthick=5,lw=1,stroke=True,tscale=tscale,textThick=textThick)
        
        tifffile.imwrite(pathout,img, append=False)

    else:
        #image stacks
        # print("convert .ser file to tiff stack")
        imgs=ser['data']
        nFrames,w,h =ser['data'].shape
        scalebar_pix = round(w/100) *scalebar_pct
        # fps = int(input('Input desired output fps:'))
        pathout2 =path[:-4]+'_'+str(fps)+'.avi'    
        codec  = cv2.VideoWriter_fourcc(*'MJPG')
        video_out = cv2.VideoWriter(pathout2, codec , fps, (w,  h),isColor=False)

        with tifffile.TiffWriter(pathout,bigtiff=True) as tif:
            for i in tqdm.tqdm(range(nFrames)):
                img=imgs[i]
                # img = cv2.convertScaleAbs(img)
                # img = (img/ 65535.0 * 255).astype(np.uint8)
                img =cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
                if addscale==True: 
                    scale = 1/(ser['pixelSize'][0]*1e9)
                    AddScalePix(img,barLenPix =scalebar_pix, scale = scale, px=0.02,py=0.98,barthick=5,lw=1,stroke=True,tscale=tscale,textThick=textThick)
            
                tif.write(img, contiguous=True)
                video_out.write(img)




def SerFolder2Tiff(folder_path,fps=15):
    print("Batch convert folder into tiff or movie: ")
    print("----------------------------------------")
    # List all files in the directory
    import os
    files = os.listdir(folder_path)
    
    # Filter files with .ser extension
    ser_files = [file for file in files if file.endswith('.ser')]
    
    # Process each .ser file
    for ser_file in ser_files:
        print("convert ", ser_file)
            # Check file size
        
        file_path = os.path.join(folder_path, ser_file)
        file_size = os.path.getsize(file_path)
        if file_size == 0:
            print(f"Skipping empty file: {ser_file}")
            continue
            
        # Try to convert the SER file, and skip if an error occurs
        try:
            # convertser(file_path, fps=fps, scalebar_pix=scalebar_pix, tscale=tscale, textThick=textThick)
            convertser(file_path, fps=fps, scalebar_pct=20)
    
        except Exception as e:
            print(f"Error converting {ser_file}: {e}")
            continue  # Skip to the next file
        # convertser(file_path,fps=10,scalebar_pix=scalebar_pix,tscale=tscale, textThick=textThick)
    
    print("----------------------------------------")
    print("Convertion done! ")