"""
This script can be used to process in-situ camera datasets (or in-situ STEM) in DigitalMicrograph
The script supports custom processing of each frame by modifying the contents of the custom_frame_processing function below.


This script uses a module called ISDatasetProcessing 
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    
"""
from benmillerscripts import ISDatasetProcessing as IDP
import numpy as np

# User-set parameters   XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
IDP.num_frames_avg = 1 #Set this parameter >1 to perform an exponentially  
    #weighted moving average on the data over time. 
IDP.ProcessedName = "FFT of " #Use this parameter to set the prefix of the name of the 
    #resulting datasets, files. This will be followed by the original dataset name.
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX


    # Available Processing Functions from IDP module to use inside custom_frame_processing
    #IDP.bin2D(numpy_array_2D,binning=(2,2))
    #IDP.FFT(array, complex=false) 
    #You can of course also use any numpy functions, or other packages that have been 
        #imported to write your own frame processing code



def custom_frame_processing(numpy_array_2D): 
    #XXXXXXXXXX Change the code within these lines to modify the processing XXXXXXXXXXX
    #result = IDP.bin2D(numpy_array_2D,binning=(2,2))
    #result = np.mean(numpy_array_2D, axis=0)
    result = IDP.FFT(numpy_array_2D, complex=True)
    #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
    return result, len(result.shape) #Function must return both processed data and its shape


IDP.processimage = custom_frame_processing #This replaces the default processing 
    #from the module with the custom frame processing defined in this script
IDP.ProcessFrontISDataset()