#Code to compute radial-max profiles of the FFT of IS Video Datasets in GMS. 
#Lines of code between #XXXXXXXX... lines are specific to computing a radial-max profile over time
	#All other lines of code are general, and can be re-used to produce other kinds of profiles over time from IS Video Datasets

#Requires Scipy. 
#To install packages like scipy, see instructions in GMS Help:Python:Installation and Configuration:Additional Packages

#Code written by Ben Miller. Last Updated 2020-06

import numpy as np
import os
import sys
import time
if (DM.IsScriptOnMainThread() == False):
	print( ' MatplotLib and scipy scripts require to be run on the main thread.',
		'\n Uncheck the "Execute on Background Thread"',
		'checkbox at the bottom of the Script Window' )
	exit()
import scipy
from scipy import ndimage
from scipy import signal
from scipy import fftpack
from scipy.ndimage.interpolation import geometric_transform
import tkinter as tk
import tkinter.filedialog as tkfd


#User-Set Parameters
Profile_Resolution = 200 

#Function to Get List of All Files in In-Situ Dataset
def BrowseforFileList():
	# Let User Select the IS Dataset Directory
	sys.argv.extend(['-a', ' '])
	root = tk.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)

	# 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]
	listOfFiles.sort()
	return (listOfFiles,newdir)

#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
# Function to Process the Image Data in Each Image 
def processimage(numpy_data, profile_res):
	#Funtion to convert cartesian-coordinate image to polar-coordinate image
	def topolar(img,  r_size, theta_size, order=1):
		sx, sy = img.shape
		max_radius = int(sx/2)
		#define transform
		def transform(coords):	
			theta = 2.0*np.pi*coords[1] / (theta_size - 1.)
			radius = max_radius * coords[0] / r_size
			i = int(sx/2) - radius*np.sin(theta)
			j = radius*np.cos(theta) + int(sx/2)
			return i,j
		#perform transform
		polar = geometric_transform(img, transform, output_shape=(r_size,theta_size), order=order,mode='constant',cval=1.0,prefilter=False)	
		return polar

	#Function to calculate radial profile of FFT from image
	def FFT_radial_profile(image_o, profile_res):
		(sx,sy)= image_o.shape
		if sx>sy:
			image_o = image_o[sx//2-sy//2:sx//2-sy//2+sy,:]
		elif sx<sy:
			image_o = image_o[:,sy//2-sx//2:sy//2-sx//2+sx]
		sizefraction=2
		#compute FFT
		fft_im = np.absolute(scipy.fftpack.fftshift(scipy.fftpack.fft2(image_o)))
		#Median-Filter FFT to remove single-pixel outliers
		#fft_im_median=scipy.ndimage.median_filter(fft_im, size=3)
		fft_im_median=fft_im
		#determine profile size
		sx, sy = fft_im.shape
		profile_size = int(sx/sizefraction)
		#convert FFT image to polar coordinates
		polar_im = topolar(fft_im_median, profile_size, profile_res, order=1)
		#compute radial mean and maximum profiles
		labels = np.mgrid[1:profile_size+1,1:profile_res+1]
		index = np.arange(1,profile_size+1)
		radial_max = ndimage.measurements.maximum(polar_im, labels=labels[0,:,:],index=index)
		radial_mean = ndimage.measurements.mean(polar_im, labels=labels[0,:,:],index=index)
		#median-filter the radial mean profile to smooth this further
		radial_mean_median = scipy.signal.medfilt(radial_mean)
		#radial profile is radial-max minus radial-mean
		radial_profile = np.atleast_2d(radial_max-radial_mean_median)	
		return radial_profile
		
	processed_data = FFT_radial_profile(numpy_data, profile_res)
	return processed_data
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

#Browse for IS Dataset and Get List of Files in Dataset
(listOfFiles,newdir) = BrowseforFileList()
#Main loop over all image files
i=0
N=len(listOfFiles)
start=time.perf_counter()
for file in listOfFiles:
	image = DM.OpenImage(file)
	imagedata = image.GetNumArray()
	processedimagedata = processimage(imagedata, Profile_Resolution)
	if i==0:
		sx, sy = processedimagedata.shape
		result_data = np.zeros((sy,len(listOfFiles)))
		result_data[:,i]  = processedimagedata
		result_image = DM.CreateImage(result_data)
		result_data = result_image.GetNumArray()
		result_image.ShowImage()
	else:
		result_data[:,i]  = processedimagedata
		if(i%10==1): 
			result_image.UpdateImage()
			print("Processed Image %s of %s" %(i,N))
	DM.DeleteImage(image)
	i=i+1
end=time.perf_counter()
print('Processed %s Images in %s seconds' %(N,end-start))