'''
This script produces maps of crystalline regions from an in-situ 4D STEM dataset in DigitalMicrograph
Run the script with the in-situ 4D STEM dataset window front-most in the software
Then place a picker-tool on the 4D dataset so a diffration pattern is shown. 
Next place a circular ROI around the region in the center that should be 
ignored when finding the brightest pixel in each diffraction pattern. 

The ROI must be circular, not oval, or an error will be returned.

This requires a module FFTArrayAnalysis, which requires scipy, skimage, matplotlib, and tqdm 
The module is part of a package which can be installed with pip:
    pip install BenMillerScripts
Code written by Ben Miller. Last Updated July 2024
'''
import os, sys
from datetime import datetime
start_time = datetime.now()
import tkinter as tk
from tkinter import *
sys.argv.extend(['-a', ' '])
import tkinter.filedialog as tkfd
import numpy as np
from benmillerscripts import FFTArrayAnalysis as FAA
from pathlib import PureWindowsPath
import scipy
import tqdm
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
#User-Set Parameters
FAA.maskC_width = 0 #Pixels to mask on the center vertical and horizontal
map_var = 'theta'   #Variable to Map in the colormap ('theta' or 'radius')
                    #theta maps the angle of the brightest FFT spot
                    #radius maps the spacing of the brightest FFT spot
scale = 1           #Optional Scaling factor (should be a positive integer)
                    #This scales the data up in real space, but down in diffraction space

FAA.pre_filter = False      #Optionally filter the 4DSTEM dataset with a median filter
gaussian_blur = True        #Optionally filter the 4DSTEM dataset with a gaussian blur 
FAA.median_size = (1,1,3,3)
showcube = False            #Optionally show the datacube that is analyzed (after cropping, scaling, and masking)
show_grey_maps = False      #Optionally show the 2 maps that are used to generate the color map
subtract_mean_bg = False    #Optionally subtract the mean diffraction pattern from every pattern
draw_color_scale = True    #Optionally produce a circular color scale

find_RGB_scale_max = False  #Find the approximate max over the whole dataset (this will process a subset of the frames)
RGB_scale_max = 535000      #The maximum found using the find_RGB_scale_max option 
numframes = None            #number of in-situ frames to process. Set to None to get the total number from the director image and process the whole dataset
step = 1                    # set to >1 to process only a fraction of the frames
GUI_Progress_Bar = True
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    
#Use TQDM??
if GUI_Progress_Bar: 
    from  tqdm.gui import tqdm
    import matplotlib.pyplot as plt
    
#If we want to find a reasonable RGB scale maximum to use for this script, we can just process a subset of the data    
if find_RGB_scale_max: 
    step = step*5
    RGB_scale_max = None
    if not DM.OkCancelDialog("Finding Colorscale Max is on (find_RGB_scale_max = True)... the script will only process 1 out of "+str(step)+" frames."): exit()

def GetSelectedFiles():
    """
    Function to Get List of File Locations for Files Selected with a File Dialog

    Returns:
        listOfFiles        (list of strings) list of files selected by user
    """
    sys.argv.extend(['-a', ' '])
    root = tk.Tk()
    root.withdraw() #use to hide tkinter window
    currdir = os.getcwd()
    listOfFiles = list(tkfd.askopenfilenames(parent=root,title='Choose DM4 File'))
    if len(listOfFiles)>0:
        os.chdir(os.path.dirname(listOfFiles[0]))
    root.destroy()
    return listOfFiles

def GetNthFrame(director_image_filepath,director_image=None, N=0):
    """
    Function to get the Nth frame from an in-situ dataset saved with a 
    director-image and raw-file(s) format. 
    """
    global j
    try:
        raw_filepath = director_image_filepath[:-4]+".raw"
        raw_length = os.path.getsize(raw_filepath)
        if j != 0:
            raw_filepath = os.path.splitext(raw_filepath)[0]+".raw_"+str(j)
    except:
        print("No raw file found for: " +director_image_filepath)
    if director_image is None:
        director_image = DM.OpenImage(director_image_filepath)
    director_data = director_image.GetNumArray()
    data_type = director_data.dtype
    pixels_per_frame = np.abs(np.prod(director_data.shape,dtype='uint'))

    bytes_per_frame = director_data.nbytes
    read_start = N*bytes_per_frame - j*raw_length
    if read_start + bytes_per_frame > raw_length:
        j=j+1
        count_1 = int((raw_length-read_start)/bytes_per_frame*pixels_per_frame)
        count_2 = pixels_per_frame-count_1
        array_1 = np.fromfile(raw_filepath, dtype=data_type, count=count_1, sep='', offset=read_start)
        raw_filepath = os.path.splitext(raw_filepath)[0]+".raw_"+str(j)
        array_2 = np.fromfile(raw_filepath, dtype=data_type, count=count_2, sep='', offset=0)
        array = np.concatenate([array_1,array_2])
    else: 
        array = np.fromfile(raw_filepath, dtype=data_type, count=pixels_per_frame, sep='', offset=read_start)
    array = array.reshape(director_data.shape)
    return array
    


def get_pickerROI_num():
    '''
    Function to get a picker-tool index number from the front-most data in GMS
    returns: ROI_num    integer index for an region of interest (ROI) in a GMS image
    '''
    #DM Script Embedded in this Python Code
    #The script gets the picker tool number from the front-most image
    # and puts it in a global tag that Python can access
    dm='imagedocument imdoc = GetFrontImageDocument()' +'\n'
    dm+='imageDisplay disp = imdoc.ImageDocumentGetImageModeDisplay()' +'\n'
    dm+='number nR = disp.ImageDisplayCountROIs()' +'\n'
    dm+='if(nR==0) Result("No Picker Tool Present on 4D STEM Dataset")' +'\n'
    dm+='if(nR>1) Result("Warning, Multiple Picker Tools Present on 4D STEM Dataset, Using Number "+disp.ImageDisplayGetROI(0).ROIGetLabel())' +'\n'
    dm+='number ROI_num = disp.ImageDisplayGetROI(0).ROIGetLabel().val()' +'\n'
    dm+='GetPersistentTagGroup().TagGroupSetTagAsLong("Python_temp:out:ROInum",ROI_num)' +'\n'
    #Run DM Script
    try: DM.ExecuteScriptString(dm)
    except: print("No Picker Tool Found on Front-Most Image"); exit()
    #Get # from Tags
    (b,ROI_num) = DM.GetPersistentTagGroup().GetTagAsUInt32('Python_temp:out:ROInum')
    return ROI_num
    
def get_picker_data_and_ROI(image):
    '''
    Function to get a diffraction pattern from the linked picker-tool data and a rectangular ROI within it
    Accepts: image                         the original DM image containing a 4DSTEM dataset- this should have a picker-tool placed on it, be front-most in GMS,
    returns: (top,left,bottom,right)    the coordinates (in pixels) of the rectangular ROI position
             diff_data                    numpy array with the picker tool data (a diffraction pattern)
    '''
    #make sure this image is front-most
    image.ShowImage() 
    #check that this is a 4D dataset
    if image.GetNumDimensions() != 4: DM.OkDialog("Front-Most Image is not a 4D Dataset...\n\nAborting Script") ; exit()
    #Get picker tool diffraction data
    Diff_Name = "("+str(get_pickerROI_num())+") Diffraction of "+image.GetName()
    diff_img = DM.FindImageByName( Diff_Name )
    diff_data = diff_img.GetNumArray()
    sy,sx = diff_data.shape
    if ( diff_img is None ):
        print( 'No image "', Diff_Name, '" was found' )
    #Get ROI coordinates
    imageDisplay = diff_img.GetImageDisplay(0) 
    roi = imageDisplay.GetROI(0)
    try: 
        if roi.IsRectangle:
            top,left,bottom,right = np.array(roi.GetRectangle(),dtype='int32')
            x = (right+left)//2
            y = (bottom+top)//2
            r = (right-left)/2
        if roi.IsCircle():
            x,y,r = np.array(roi.GetCircle(),dtype='int32')
    except: print("Problem Reading ROI on the Picker Tool Image... Script Aborted. Please Place a Circular ROI around the Central Spot"); exit()
    del diff_img, image
    return (x,y,r),diff_data
    
def DM_interaction_wait_dialog(instructions_str, title_str='DM_Interaction_Dialog',button1_str='Continue',button2_str='Cancel' ):
    '''
    Function to Create a tkinter dialog that enables the user to interact with GMS
    while the script is paused, and then resume the script after they are done
    Accepts:
        instructions_str    string giving instructions to the user about what actions are expected in GMS
        title_str            optional string for the dialog window title
        button1_str            optional string to change the continue button label text
        button1_str            optional string to change the cancel button label text
    Returns: Nothing
    '''
    root = Tk()
    root.attributes('-topmost', 'true')
    root.title(title_str)
    def go_on():
        root.destroy()
    def cancel():
        print("Script Aborted by User")
        root.destroy()
        sys.exit()
    Label(root, text=instructions_str).grid(row=0, sticky=W)
    Button(root, text=button1_str, command=go_on).grid(column=0, row=1, sticky=W, pady=4)
    Button(root, text=button2_str, command=cancel).grid(column=1, row=1, sticky=W, pady=4)
    mainloop()

def process_4D_data(image_4D,raw_datacube=None):
    """
    Function to process a 4D datacube, outputing only an RGB array
    This expects a DM image, not a numpy array, but a numpy array can be passed in
    with the optional raw_datacube parameter
    
    """
    if raw_datacube is None:
        raw_datacube = image_4D.GetNumArray()
    (x,y,r), picker_diff= get_picker_data_and_ROI(image_4D)
    sx,sy = picker_diff.shape
    center = (x,y)
    FAA.maskP = r*2/sx*100
    try:
        if r == 0: 
            raise ValueError("ROI must be a circle (with radius defined in control panel on the left of DM) , not an oval.")
    except:
        DM.OkDialog("ROI must be a circle (with radius defined in control panel on the left of DM) , not an oval.")
        if GUI_Progress_Bar: plt.close('all')
    name = image_4D.GetName()
    origin, x_scale, scale_unit =  image_4D.GetDimensionCalibration(2, 0)

    #Optionally Scale Data (up in real space, down in diff space)
    if scale > 1:
        data_s = FAA.STEMx_Crop_to_Center_Square(raw_datacube, center)
        datasquare = np.copy(FAA.Upscale_4D_STEM(data_s, scale))
    else: 
        datasquare = np.copy(FAA.STEMx_Crop_to_Center_Square(raw_datacube, center))

    #Optionally Subtract Mean Pattern from Data
    if subtract_mean_bg: 
        avg_diffraction = np.mean(datasquare,axis=(0,1))
        data_BGSub = datasquare-avg_diffraction
        data_BGSub[data_BGSub<0] = 0
        data_input = data_BGSub
    else: 
        data_input = datasquare
    sigma =(0,0,2,2)
    if gaussian_blur: data_input = scipy.ndimage.gaussian_filter(data_input, sigma,truncate=2)
    #Process Dataset
    (RGB_im,direction_image,intensity_image,spacing_image,diffractogram_max) = FAA.STEMx_Process_Cube(image_4D,map_var, im_data = data_input, show_cube = showcube,RGB_scale_max=RGB_scale_max)
    DM_RGB = FAA.ShowRGB(RGB_im, name)
    DM.DoEvents()
    return RGB_im

def find_colorscale_max(showplot=True):
    '''
    Function to use values stored by the FAA module to determine the best
    colorscale maximum to use during future runs of this script
    '''
    if showplot:
        high_im = DM.CreateImage(np.asarray(FAA.high_lims))
        high_im.ShowImage()
        del high_im
    #Output the suggested maximum
    high_limit = np.max(np.asarray(FAA.high_lims))
    print("Max spot intensity over dataset is %s" %high_limit)
    print("Use this value for the RGB_scale_max parameter")
    
def mean_max_4D(image_4D,raw_datacube=None):
    if raw_datacube is None:
        raw_datacube = image_4D.GetNumArray()
    mean = np.mean(raw_datacube,axis =(0,1))
    max = np.max(raw_datacube,axis =(0,1))
    return mean, max

director_image = DM.GetFrontImage()
#Get the file path of the 4D STEM dataset open in DM (this needs a one-line DM script)
dm='GetPersistentTagGroup().TagGroupSetTagAsString("Python_temp:out:FrontFileLocation",GetFrontImageDocument().ImageDocumentGetCurrentFile())'
DM.ExecuteScriptString(dm)
(b,filepath) = DM.GetPersistentTagGroup().GetTagAsString('Python_temp:out:FrontFileLocation')
if os.path.isfile(filepath):
    print(filepath)
else: 
    print("Front image could not be found on disk. Has it been saved?")
#director_file = GetSelectedFiles()[0]
director_file = filepath
folder,file = os.path.split(director_file)
output_file = os.path.join(folder,"IS Processing",file)
output_file2 = os.path.join(folder,"IS Processing Mean-Max",file)
if not os.path.exists(os.path.dirname(output_file)):
        os.makedirs(os.path.dirname(output_file))
if not os.path.exists(os.path.dirname(output_file2)):
        os.makedirs(os.path.dirname(output_file2))
        
DM_interaction_wait_dialog("Put a Picker tool on the 4D Dataset and then an oval ROI around the central region of the diffration pattern to be masked out")
if numframes is None:
    returnVal, numframes = director_image.GetTagGroup().GetTagAsLong('In-situ:Recorded:# Frames')

raw_out_path = PureWindowsPath(os.path.join(folder,output_file)).with_suffix(".raw")
if os.path.exists(raw_out_path):
        message = str(raw_out_path) + " Already exists\n Overwrite existing processed data?"
        if not DM.OkCancelDialog(message): exit()
        print("Over-writing existing processed data")
        os.remove(raw_out_path)

j=0
with open(raw_out_path, "ab") as raw_file:
    for i in tqdm(range(0,numframes,step)):
        if i>0: DM.DeleteImage(result)
        print("Processing frame %s of %s " %(i,numframes))
        data = GetNthFrame(director_file,director_image=director_image,N=i)
        RGB_im = process_4D_data(director_image,raw_datacube=data)
        result = DM.GetFrontImage()
        DM.DoEvents()
        raw_file.write(RGB_im.tobytes())
        

result_2D = DM.GetFrontImage()
dm  = "InSituTags_CopyRecordTags(FindImageByID("+str(director_image.GetID())+"),FindImageByID("+str(result_2D.GetID())+"))"
dm += "\n FindImageByID("+str(result_2D.GetID())+").IMDSetFunctionInSituDirector()"
dm += "\n InSituTags_SetRawDataInfo(FindImageByID("+str(result_2D.GetID())+"),FindImageByID("+str(result_2D.GetID())+"))"
DM.ExecuteScriptString(dm)
DM.DoEvents()
path_2D = str(raw_out_path.with_suffix(".dm4"))
print(path_2D)
DM.GetFrontImage().SaveImage(path_2D)
print("Saved 2D Results")
#Must close and re-open in-situ director image (using an image document) to get the In-Situ Player to sync it
DM.DeleteImage(result_2D)
print("Opening 2D Director Image\n")
DM.DoEvents()
doc = DM.NewImageDocumentFromFile(path_2D)
doc.Show()
if GUI_Progress_Bar: plt.close('all')
find_colorscale_max(showplot=find_RGB_scale_max)

#Optionally Produce Circular Color Scale
if draw_color_scale: FAA.DrawColorScale()
print("\nStart time was %s" %start_time.time())
print("End time was %s" %datetime.now().time())
print("Processing took %s minutes" %((datetime.now()-start_time ).total_seconds()/60))
print("\nDone")