#Code to compute radial-max profiles of the FFT of live-view image in GMS. 
#Run the code with the live-view image front-most in GMS. 
#To stop calculation, close the live view image. 
#Profile is computed as often as possible.
#Lines of code between #XXXXXXXX... lines are specific to computing a radial-max profile
	#All other lines of code are general, and can be re-used to produce other kinds of profiles from a live-view image

#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 time
import sys
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

median = 1 #set to 1 to apply median filter to the FFT prior to profile creation (slows calculation, especially for large input images)
initial_result_image_width = 20 #how many profiles can be displayed in the intial result window (window is automatically expanded as needed)
profile_result_length_ratio = 2 #Set to some integer 2^N, N=>0. Smaller N will make calculation slower. Default: 4
profile_angular_sampling_resolution = 512 #set how many samples are taken around the circumference of the radial profile

class CListen(DM.Py_ScriptObject):

	#Initialization Function
	def __init__(self, img):
		self.name = "ImageO"
		self.i = 0
		self.j = True #this is only for IS player testing
		#get the original image and assign it to self.imgref
		self.imgref = img
		#get the shape and calibration of the original image
		(input_sizex, input_sizey) = img.GetNumArray().shape
		origin, self.scale_orig, self.unit_stringo = self.imgref.GetDimensionCalibration(0, 0)
		#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
		#Set the size and calibration of the result image
		(self.scale,self.unit_string) = self.calibrate_result(self.scale_orig,input_sizex,profile_result_length_ratio,self.unit_stringo)
		r_img_size=self.profile_length(input_sizex,profile_result_length_ratio)
		#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
		#create empty set for result images
		self.result_images = {}
		#create 1st result image and set calibration
		self.result_images[self.name] = DM.CreateImage(np.zeros((r_img_size,initial_result_image_width)))
		self.result_images[self.name].SetDimensionCalibration(1,0,self.scale,self.unit_string,0)
		self.result_images[self.name].ShowImage()		
		#get numpy array from result image
		self.result_array = self.result_images[self.name].GetNumArray()
		DM.Py_ScriptObject.__init__(self)
		self.stop = 0
	
	#Function to end Image Listener
	def __del__(self):
		DM.Py_ScriptObject.__del__(self)
		
	#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX	
	#Function to Determine Result Profile Length
	def profile_length(self,sx,sf):
		return int(sx/sf)
	#Function to Set Result Image Calibration
	def calibrate_result(self,scale,sx,sf,ustr):
		diff_scale = sf/scale/2/sx
		unit_str = ustr+"-1"
		return(diff_scale,unit_str)
	#Funtion to convert cartesian-coordinate image to polar-coordinate image
	def topolar(self, 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(self, image_o, profile_ang_res,length_ratio,do_median):	
		(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]
		#compute FFT
		fft_im = np.absolute(scipy.fftpack.fftshift(np.fft.fft2(image_o)))
		#Median-Filter FFT to remove single-pixel outliers
		if do_median: fft_im_median = scipy.ndimage.median_filter(fft_im, size=3)
		else: fft_im_median = fft_im
		#determine profile size
		sx, sy = fft_im.shape
		profile_size = int(sx/length_ratio)
		#convert FFT image to polar coordinates
		polar_im = self.topolar(fft_im_median, profile_size, profile_ang_res, order=1)
		#compute radial mean and maximum profiles
		radial_max=np.amax(polar_im,1)
		radial_mean=np.mean(polar_im,1)
		#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
	#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX	

	
	#This function is run each time the image changes
	def HandleDataChangedEvent(self, flags, image):
		if not self.stop:
			if self.j: #this is only for IS player testing
				#start timing
				start=time.perf_counter()
				self.j = not self.j #this is only for IS player testing
				(result_sizey, result_sizex) = self.result_array.shape
				#if the result image is nearly full, make it 2x larger
				if self.i > result_sizex-2:
					#create a new numpy array 2x larger
					self.result_array_temp = np.append(self.result_array, np.zeros_like(self.result_array), axis=1)
					#close the old results image in DM
					DM.DeleteImage(self.result_images[self.name])
					print("ImageDeleted")
					#create a new results image and calibrate it
					self.name="Image{0}".format(self.i)
					self.result_images[self.name] = DM.CreateImage(np.copy((self.result_array_temp)))
					self.result_images[self.name].SetDimensionCalibration(1,0,self.scale,self.unit_string,0)
					print("ImageCreated")
					#display new result image in DM
					self.result_images[self.name].ShowImage()	
					#get numpy array from new result image
					self.result_array = (self.result_images[self.name].GetNumArray())
				#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
				#compute radial FFT profile from the image, and place this profile into results image
				self.result_array[:,self.i] = self.FFT_radial_profile(self.imgref.GetNumArray(),profile_angular_sampling_resolution,profile_result_length_ratio, median)
				#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
				self.result_images[self.name].UpdateImage()
				#end timing and output time to process this frame
				end=time.perf_counter()
				print("Processed Image "+str(self.i)+"  Processing Time= "+str(end-start))
				self.i = self.i+1	
			else: 
				self.j = not self.j #this is only for IS player testing
	#Function to end script if source image window is closed
	def HandleWindowClosedEvent(self, event_flags, window):
		print("Window Closed")
		self.stop = 1
		self.__del__()
		#del self.imgref
		print("Script Ended")		

#Main Code Starts Here
#Get front image in GMS		
img1 = DM.GetFrontImage()
#Get the image window, so we can check if it gets closed
imageDoc = DM.GetFrontImageDocument()
imDocWin = imageDoc.GetWindow()
#initiate the image listener
listener = CListen(img1)
#check if the source window closes
listener.WindowHandleWindowClosedEvent(imDocWin, 'pythonplugin')
#check if the source image changes
listener.ImageHandleDataChangedEvent(img1, 'pythonplugin')