'''Code to compute radial-max profiles of IS Video Datasets in DigitalMicrograph
    The data can be either 
        HRTEM images (an FFT will be computed by this script) 
        Diffraction patterns
Run the script with the in-situ dataset window front-most in the software

This script produces both a 2D display of the radial profiles over time (with  
time on the horizontal axis), and an in-situ dataset which is automatically 
synched with the original in-situ dataset when either is played back with the 
In-Situ Player

Requires Scipy. 
To install packages like scipy, see instructions in DigitalMicrograph Help:Python:Installation and Configuration:Additional Packages
Code written by Ben Miller. Last Updated Feb 2024
'''
import numpy as np
import os
import sys
import time
if not DM.IsScriptOnMainThread(): print('Scipy scripts cannot be run on Background Thread.'); exit()
import scipy
from scipy import ndimage
from scipy import fftpack
from scipy.ndimage import map_coordinates
from tkinter import *
sys.argv.extend(['-a', ' '])
import tkinter.filedialog as tkfd
from numpy.lib.stride_tricks import as_strided

#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
#User-Set Parameters
Diff_not_FFT = False #Set to True to calculate profiles from a live diffraction pattern, rather than from the FFT of an HRTEM image.
mask_center_lines = True #Set to True to mask the vertical and horizontal center lines (often useful for FFTs)
ask_about_diffraction = True #Set to True to always ask to confirm whether to process diffraction patterns or HRTEM images
binx = 2  #Set to an integer number >=1 to bin the diffraction pattern or diffractogram prior to computing the radial profile. Smaller N will make the calculation slower. 
maskP=5 #Set Percentage of Radial Profile (from center) to be Ignored 
num_frames_avg = 1 # Set to value >1 to apply an exponentially weighted moving average prior to computing the radial profiles 
    #Note that this is applied to HRTEM images before computing an FFT, so if there is significant drift this may produce poor results.
ask_for_centering = True #Set to True to ask the user to verify that an ROI on the image is centered before proceeeding with computation.
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX


change = True
if ask_about_diffraction: 
    if Diff_not_FFT: 
        message = "Do you want to process a series of diffraction patterns? \nSelect OK to continue or cancel to instead process HRTEM images."
    else:    
        message = "Do you want to process a series of HRTEM images? \nSelect OK to continue or cancel to instead process diffraction patterns."
    change = not DM.OkCancelDialog(message)
    
if (change and ask_about_diffraction): Diff_not_FFT = not Diff_not_FFT 

profile_result_length_ratio = 1 #Not working now... keep this equal to 1 
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

#Function to Browse for a Folder and Get a List of All Files in In-Situ Dataset
def BrowseForFileList():
    # Let User Select the IS Dataset Directory
    sys.argv.extend(['-a', ' '])
    root = Tk()
    root.withdraw() #use to hide tkinter window
    currdir = os.getcwd()
    dirname = tkfd.askdirectory(parent=root, initialdir=currdir, title='Please select the IS Dataset Root Directory')
    if len(dirname) > 0:
        print("\nOriginal IS DataSet Directory: %s" % dirname)
        newdir=dirname[:3] + 'DMScript Edited Datasets/' + dirname[3:]
        os.chdir(dirname)
    else:
        root.destroy()
        print("User Canceled File Dialog")
        exit()
    # Get the list of all files in directory tree at given path
    listOfFiles = list()
    for (dirpath, dirnames, filenames) in os.walk(dirname):
        listOfFiles += [os.path.join(dirpath, file) for file in filenames if file.endswith('.dm4')]
    listOfFiles.sort()
    root.destroy()
    return (listOfFiles,dirname,newdir)

#Function to Get a List of All Files in In-Situ Dataset
def Get_IS_File_List(dirname):
    # Get the list of all files in directory tree at given path
    listOfFiles = list()
    for (dirpath, dirnames, filenames) in os.walk(dirname):
        listOfFiles += [os.path.join(dirpath, file) for file in filenames if file.endswith('.dm4')]
        listOfFiles.sort()
    return listOfFiles

#Function to enable the user to interact with DigitalMicrograph, and then resume the script when they are ready
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()
    
#Function to find an ROI on an image, or to add one if none exists. 
def find_ROI(image):
        imageDisplay = image.GetImageDisplay(0)
        numROIs = imageDisplay.CountROIs()
        id = None
        for n in range(numROIs):
            roi = imageDisplay.GetROI(0) 
            if roi.IsRectangle():
                roi.SetVolatile(False)
                roi.SetResizable(False)  
                id = roi.GetID()
                break
        if id is None:
            #If No ROI is found, create one that covers the whole image. 
            print("\nRectangular ROI not found... using whole image")
            data_shape = image.GetNumArray().shape
            roi=DM.NewROI()
            roi.SetRectangle(0, 0, data_shape[0], data_shape[1])
            imageDisplay.AddROI(roi)
            roi.SetVolatile(False)
            roi.SetResizable(False) 
            id = roi.GetID()
        return id
        
#Function to Set Result Image Calibration
def calibrate_result(scale,sx,sf,ustr,diff_not_fft=False):
    if diff_not_fft:
        diff_scale = scale*sf*binx
        unit_str = ustr
    else:
        diff_scale = binx*sf/scale/2/sx
        unit_str = ustr+"-1"
    return(diff_scale,unit_str)
    
#Function to Determine Result Profile Length
def profile_length(sx,sf):
    return int(sx/sf/binx/2)
    
#Function to bin an image using only numpy
def strided_binning2D(array,binning=(1,1)):
    bin=np.flip(binning)
    nh = (array.shape[0]-bin[0])//bin[0]+1
    nw = (array.shape[1]-bin[1])//bin[1]+1
    strides = (array.strides[0],array.strides[1],array.strides[0]*bin[0],array.strides[1]*bin[1])
    shape = (bin[0],bin[1],nh,nw)
    virtual_datacube = as_strided(array,shape=shape,strides=strides)
    result = np.sum(virtual_datacube,axis=(0,1))
    return result
    
#Function to transform cartesian image to polar image
def topolar(f,m,profile_res, coords, order=1, mask_percent=10):
        g = map_coordinates(f, coords, order=order)
        g = g.reshape(profile_res,m//2)
        mask_n = int((m*mask_percent/100)/2)
        g[:,0:mask_n]=0
        return g
        
#Function to produce coordinates to transform cartesian image to polar image
def topolar_init(data,m, profile_res, center,order=1, mask_percent=10):
    def rphi2xy(r, phi,center):
        x=r*np.cos(phi)+center[0]
        y=r*np.sin(phi)+center[1]
        return x, y
    #center = m//2
    rmax = m//2
    phimax = 2*np.pi
    rs, phis = np.meshgrid(np.linspace(0, rmax, m//2),np.linspace(0, phimax, profile_res),sparse=True)
    xs, ys = rphi2xy(rs, phis,center)
    xs, ys = xs.reshape(-1), ys.reshape(-1)
    coords = np.vstack((ys, xs))
    array = topolar(data,m,profile_res,coords,order=order,mask_percent=mask_percent)
    return (coords, array)
    
#Function to process FFT data and return a profile and polar image
def processimage(numpy_data, profile_res,coords,maskP, FFT=False):
    
    #Funtion to convert cartesian-coordinate image to polar-coordinate image
    def angular_profile(Diffim, profile_res, exclude_rad=10):
        sx,sy = Diffim.shape
        FFT_size = sx//2
        if(sx!=sy):
            print("Input image for radial profile should be square")
        
        polar_im = topolar(Diffim,sx,profile_res,coords, order=1, mask_percent=maskP)
        radial_max = np.max(polar_im,axis=0)
        angular_max = np.max(polar_im,axis=1)
        radial_mean = np.mean(polar_im,axis=0)
        radial_max_minus_mean = radial_max - radial_mean
        return (radial_max_minus_mean,polar_im)

    if FFT: 
        numpy_data = np.absolute(scipy.fftpack.fftshift(np.fft.fft2(numpy_data)))
        numpy_data = strided_binning2D(numpy_data,binning=(binx,binx))
    sx, sy = numpy_data.shape
    if mask_center_lines: 
            numpy_data[:,sy//2-1:sy//2-1+1] = 0
            numpy_data[sx//2-1:sx//2-1+1,:] = 0
    numpy_data = ndimage.median_filter(numpy_data,size=3)
    profile,polar_im = angular_profile(numpy_data, profile_res, exclude_rad=(max(1,numpy_data.shape[0]//256)*10))
    #p_length = len(profile)//2
    #filtered_profile = scipy.ndimage.gaussian_filter1d(profile[0:p_length]+profile[p_length:],sigma = 2)
    filtered_profile = ndimage.gaussian_filter1d(profile,sigma = 1)
    return filtered_profile
    
#Function to add a Circular ROI to an image
#(Currently unused)
def AddCircleROI(image,center,radii):
    im_disp = image.GetImageDisplay(0)
    for radius in radii:
        t = center[1]-radius
        b = center[1]+radius
        l = center[0]-radius
        r = center[0]+radius
        my_roi=DM.NewROI()
        my_roi.SetOval(t,l,b,r)
        im_disp.AddROI(my_roi)
        
#Function to compute an exponentially weighted moving average from an existing average and a new input frame
def MyExpMovingAvg(frame,old_avg,num_frames_sum):
    if num_frames_sum <=1:
        return(frame)
    else:
        persistence = (num_frames_sum-1)/(num_frames_sum+1)
        new_avg = persistence*old_avg + (1-persistence)*frame
        equivalent_sum = new_avg*num_frames_sum
        return(new_avg)
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

#XXXXXXXX Getting IS Dataset data XXXXXXXX
num_frames = 10
image0 = 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?")
data_0 = image0.GetNumArray()
data_type = data_0.dtype
raw_filepath = filepath[:-4]+".raw"
folder = os.path.dirname(filepath)
#Determine the IS dataset data structure based on what is saved to disk
if os.path.isfile(raw_filepath): 
    IS_type = "Director-Raw"
    raw_length = os.path.getsize(raw_filepath)
    pixels_per_frame = np.prod(data_0.shape)
    bytes_per_frame = data_0.nbytes
    (b,num_frames) = image0.GetTagGroup().GetTagAsUInt32("In-situ:Recorded:# Frames")
elif os.path.exists(os.path.join(folder,"Hour_00")): 
    IS_type = "Director-Hybrid"
    (b,num_frames) = image0.GetTagGroup().GetTagAsUInt32("In-situ:Recorded:# Frames")
    listOfFiles = Get_IS_File_List(os.path.join(folder,"Hour_00"))
else: 
    IS_type = "H-m-s"
    (listOfFiles,dirname,newdir) = BrowseForFileList()
    num_frames = len(listOfFiles)
print("IS Data Type is: "+IS_type)



#XXXXXXXX Initialize Radial Profile XXXXXXXX 
#Get the data from the region within an ROI     
roi = DM.GetROIFromID(find_ROI(image0))
if Diff_not_FFT:
    roi.SetVolatile(True)
    roi.SetResizable(True)
    if ask_for_centering: DM_interaction_wait_dialog("Move the ROI in your dataset so it is exactly centered on the diffraction pattern center, and includes as much of the pattern as you want to analyze.", title_str='DM_Interaction_Dialog',button1_str='Continue',button2_str='Cancel' )
    roi = DM.GetROIFromID(find_ROI(image0))
val, val2, val3, val4 = roi.GetRectangle()
data = image0.GetNumArray()[int(val):int(val3),int(val2):int(val4)]
#get the shape and calibration of the original image
(input_sizex, input_sizey) = data.shape
origin, x_scale_orig, scale_unit_orig = image0.GetDimensionCalibration(1, 0)
if scale_unit_orig == b'\xb5m': scale_unit_orig = 'um' #scale unit of microns causes problems for python in DM
#Set the size and calibration of the result image
(scale,unit_string) = calibrate_result(x_scale_orig,input_sizex,profile_result_length_ratio,scale_unit_orig,diff_not_fft = Diff_not_FFT)
r_img_size=profile_length(input_sizex,profile_result_length_ratio)
Profile_Resolution = r_img_size
binnned_data = strided_binning2D(data,binning=(binx,binx))
if Diff_not_FFT: data = binnned_data
center = np.asarray(binnned_data.shape)//2
coords = topolar_init(binnned_data,binnned_data.shape[0],Profile_Resolution,center,mask_percent=maskP)[0]

#create 2D result image and set calibration
result_2D = DM.CreateImage(np.zeros((r_img_size,num_frames)))
result_2D.SetDimensionCalibration(1,0,scale,unit_string,0)
result_2D.ShowImage()
#Set the image name which will be displayed in the image window's title bar
result_2D.SetName("IS Radial Max Profiles of "+image0.GetName())

#XXXXXXXX Apply radial profile to every frame XXXXXXXX 
j=0
raw_out_path = os.path.join(folder,result_2D.GetName())+"_1D.raw"
if os.path.exists(raw_out_path):
    print("Over-writing existing processed data")
    os.remove(raw_out_path)
with open(raw_out_path, "ab") as raw_file:
    for i in range(num_frames):
        if IS_type == "Director-Raw":
            read_start = i*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)
                print(raw_filepath)
                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)
            frame = array.reshape(data_0.shape)
        if IS_type == "Director-Hybrid":
            image = DM.OpenImage(listOfFiles[i])
            frame = image.GetNumArray()
        if IS_type == "H-m-s": 
            image = DM.OpenImage(listOfFiles[i])
            frame = image.GetNumArray()
        # every 10 frames, update the user on current progress
        if i%10 == 0: 
            print("Processing Frame: %s of %s" %(i,num_frames))
            result_2D.UpdateImage()
            DM.DoEvents()
        #val, val2, val3, val4 = roi.GetRectangle() this can be uncommented if you want the ROI to be moveable during processing
        old_avg = data
        data = frame[int(val):int(val3),int(val2):int(val4)]
        if Diff_not_FFT: data = strided_binning2D(data,binning=(binx,binx))
        data = MyExpMovingAvg(data,old_avg,num_frames_avg)
        processedimagedata = processimage(data, Profile_Resolution,coords,maskP, FFT=(not Diff_not_FFT) )
        result_2D.GetNumArray()[:,i] = processedimagedata
        raw_file.write(processedimagedata.tobytes())
        if IS_type != "Director-Raw":
            del image
    

#XXXXXXXX Save and display results XXXXXXXX 
#Save 2D Image (already displayed)
result_2D.UpdateImage()
result_2D.SaveAsGatan(os.path.join(folder,result_2D.GetName()))
#Save 1D IS Dataset
result_1D = DM.CreateImage(processedimagedata)
result_1D.SetDimensionCalibration(0,0,scale,unit_string,0)
result_1D.SetName("IS Radial Max Profiles of "+image0.GetName()+"_1D")
result_1D.ShowImage()
#Set tags needed to sync data using In-Situ Player
dm  = "InSituTags_CopyRecordTags(FindImageByID("+str(image0.GetID())+"),FindImageByID("+str(result_1D.GetID())+"))"
dm += "\n FindImageByID("+str(result_1D.GetID())+").IMDSetFunctionInSituDirector()"
dm += "\n InSituTags_SetRawDataInfo(FindImageByID("+str(result_1D.GetID())+"),FindImageByID("+str(result_1D.GetID())+"))"
DM.ExecuteScriptString(dm)
DM.DoEvents()
path_1D = os.path.join(folder,result_1D.GetName())+".dm4"
#result_1D.SaveAsGatan(os.path.join(folder,result_1D.GetName()))
result_1D.SaveAsGatan(path_1D)
print("Saved 2D and 1D 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_1D)
print("Opening 1D Director Image")
DM.DoEvents()
#doc = DM.NewImageDocumentFromFile(os.path.join(folder,result_1D.GetName())+".dm4")
doc = DM.NewImageDocumentFromFile(path_1D)
doc.Show()
print("Done")