#########################
# Inspired by hackathon code avaialble at:
# https://imagej.net/tutorials/analyze-frap-movies
#
# This code uses the double normalization method outlined in:
# l. Phair RD, Gorski SA, Misteli T. (2004).
# Measurement of dynamic protein binding to chromatin
# in vivo, using photobleaching microscopy. 
# Methods Enzymol. 375:393-414
# doi: 10.1016/s0076-6879(03)75025-3. PMID: 14870680.
#
# July 31, 2025
# Jeff Hardin, Univ, of Wisconsin-Madison
# jdhardin@wisc.edu
#########################

import java.awt.Color as Color
import java.awt.Frame as Frame
import java.awt.Window as Window
from java.io import File as File
from java.lang import System as System
from ij import WindowManager as WindowManager
from ij.plugin.frame import RoiManager as RoiManager
from ij.process import ImageStatistics as ImageStatistics
from ij.measure import Measurements as Measurements
from ij import IJ as IJ
from ij.measure import CurveFitter as CurveFitter
from ij.gui import Plot as Plot
from ij.gui import PlotWindow as PlotWindow
from ij.measure import ResultsTable as ResultsTable
from ij.gui import PlotWindow as ImageWindow
from ij.gui import GenericDialog
from ij.text import TextWindow
from ij import ImagePlus as ImagePlus
from ij.io import FileInfo as FileInfo
import math
import os.path
from os import path, mkdir
import csv

def doFRAP():
	#Check that there is an image wondow open.
	if (WindowManager.getImageCount() == 0):
		IJ.showMessage("No image open","No image windows open.")
		return None
		
	# Get ROIs
	roi_manager = RoiManager.getInstance()
	# If the ROI manager wasn't open, open it and get instance
	if (roi_manager is None):
		IJ.run("ROI Manager...","")
		roi_manager = RoiManager.getInstance()
    	
	roi_list    = roi_manager.getRoisAsArray()

	if len(roi_list) != 3:
		IJ.showMessage("ROIs needed","This routine requires exactly three ROIs (1 = bleach, 2 = unbleached signal,\n3 = background). Add these ROIs to the ROI Manager and try again.")
		return None

	# We assume first one is FRAP roi, the 2nd one is normalizing roi.
	roi_FRAP    = roi_list[0]
	roi_norm    = roi_list[1]
	roi_BG		= roi_list[2]


	# Check that this is a stack with more than two slices
	current_imp  = WindowManager.getCurrentImage()
	if (current_imp.getStackSize()<2):
		IJ.showMessage("FRAP Analyzer", "This command requires a stack.")
		return None
	stack        = current_imp.getImageStack()
	calibration  = current_imp.getCalibration()
	
	# Get window
	myWindow = current_imp.getWindow()

	#############################################
	# Specify up to what frame to fit and plot. Right now set to total slices.
	n_slices = stack.size()

	############################################
	# Get parameters
	gd = GenericDialog("FRAP Analysis Settings")
	gd.addChoice("Curve fitting method:", ["Single", "Double"], "Single")
	
	#Next provides option to force fit the curve through the origin.
	#In theory recovery curves should do so; other FRAP curve fit routines omit this.
	#NOTE: Forcing this can change the fit parameters quite markedly in some cases,
	#so caveat emptor!
	gd.addCheckbox("Force curve fit through origin for single exponential", True)
	gd.addCheckbox("Stretch intensities to full range (0-1)", True)		
	gd.addNumericField("Time interval (sec): ",1,1)
	gd.addCheckbox("Create image with ROIs", True)
	#Saving manually gets tiresome after a while.
	#Next allows autocreation of a subdirectory and things to save there.
	#If the subdirectory exists it will ask permission to overwrite files therein.
	gd.addMessage("Autosave settings:")
	#Saves analysis results as text file
	gd.addCheckbox("Autosave results", True)
	gd.addCheckbox("Close windows after autosaving", True)
	gd.showDialog()
	
	if gd.wasOKed():
		#IJ.log("\\Clear")
		#IJ.log("FRAP Settings")
		fitMethod = gd.getNextChoice()
		throughOrigin = gd.getNextBoolean()
		stretchValues = gd.getNextBoolean()
		timeScale = gd.getNextNumber()
		createROIs = gd.getNextBoolean()
		autoSave = gd.getNextBoolean()
		closeWindows = gd.getNextBoolean()
	else: return None
	
	# Collect intensity values

	# Create empty lists of numbers
	If = []  # Frap values
	In = []  # Norm values
	Ib = []  # Bacground values

	# Loop over each slice of the stack
	for i in range(0, n_slices):
 
		# Get the current slice 
		ip = stack.getProcessor(i+1)
 
		# Put the FRAP ROI on it
		ip.setRoi(roi_FRAP)
 
		# Make a measurement in it
		stats = ImageStatistics.getStatistics(ip, Measurements.MEAN, calibration)
		mean  = stats.mean
 
		# Store the measurement in the list
		If.append(mean)

		# Do the same for non-FRAPed area
		ip.setRoi(roi_norm)
		stats = ImageStatistics.getStatistics(ip, Measurements.MEAN, calibration)
		mean = stats.mean
		In.append( mean  )

		# And now for BG area
		ip.setRoi(roi_BG)
		stats = ImageStatistics.getStatistics(ip, Measurements.MEAN, calibration)
		mean = stats.mean
		Ib.append( mean  )

 #####################################
 	# Gather image parameters
 	# Orignal code calaculated frame info from calibration data, which some images don't have
 	# Commented out next two lines below but retained for reference to original code
	#frame_interval = calibration.frameInterval
	#time_units = calibration.getTimeUnit()
	frame_interval = timeScale
	time_units = "sec"
	myTextWindow = TextWindow("FRAP Results - " + current_imp.getTitle(),"File name: " + current_imp.getTitle()+ "\n",600,300)
	myTextWindow.append("Time interval: " + str(timeScale) + " " + time_units + "\n")

	# Find minimal intensity value to find bleach frame
	min_intensity = min( If )
	bleach_frame = If.index( min_intensity )
 	myTextWindow.append('FRAP frame is ' + str(bleach_frame+1))

	# Compute mean pre-bleach intensity for each ROI
	mean_If = 0.0
	mean_In = 0.0
	mean_BG = 0.0
	#loop until the bleach time
	for i in range(bleach_frame):
		mean_If = mean_If + If[i]
		mean_In = mean_In + In[i]
		mean_BG = mean_BG + Ib[i]
	mean_If = mean_If / bleach_frame
	mean_In = mean_In / bleach_frame
	mean_BG = mean_BG / bleach_frame
 
	# Calculate normalized curve
	normalized_curve = []
	for i in range(n_slices):
		normalized_curve.append( ((If[i] - mean_BG) * (mean_In - mean_BG)) / ((In[i] - mean_BG) * (mean_If - mean_BG)) )
	x = [i * frame_interval for i in range( n_slices ) ] 
	y = normalized_curve

	xtofit = [ i * frame_interval for i in range( n_slices - bleach_frame ) ]

	if (stretchValues):
		#Now stretch to range 0 to 1
		yMax = 0.0
		yMin = 0.0
		yMin = min (normalized_curve)
		yMax = max (normalized_curve)
	
		for i in range(n_slices):
			normalized_curve[i] = (normalized_curve[i] - yMin)/(yMax - yMin)

	ytofit = normalized_curve[ bleach_frame : n_slices ]
	yMin = min (normalized_curve)
	yMax = max (normalized_curve)
 
	# Fitter
	fitter = CurveFitter(xtofit, ytofit)
	if fitMethod == "Single":
		if throughOrigin: 
			fitter.doFit(CurveFitter.EXP_RECOVERY_NOOFFSET)
		else:
			fitter.doFit(CurveFitter.EXP_RECOVERY)
	else:
		eqn = "y = a*(1-exp(-b*x)) +c*(1-exp(-d*x)) + e"
		params = fitter.doCustomFit(eqn, None, False)

	myTextWindow.append("Fit FRAP curve by " + fitter.getFormula() )
	param_values = fitter.getParams()

	# plot raw intensities of the three ROIs
	x2 = [i * frame_interval for i in range( n_slices ) ] 
	x3 = [i * frame_interval for i in range( n_slices ) ] 
	plot_raw = Plot(current_imp.getTitle() + " - raw", "Time ("+time_units+")", "Intensity")
	plot_raw.setLineWidth(2)
	# Sets minumum y to 80% of lowest value so that background is visible; 
	# use next line if yMin of 0 preferred
	plot_raw.setLimits(0, max(x), 0.8 * min(If), max(If) )
	#plot_raw.setLimits(0, max(x), 0, max(If) )
	plot_raw.setColor(Color.BLACK)
	plot_raw.addPoints(x, If, Plot.LINE)
	plot_raw_window = plot_raw.show()
	plot_raw.setColor(Color.RED)
	plot_raw.addPoints(x2, In, Plot.LINE)
	plot_raw.setColor(Color.BLUE)
	plot_raw.addPoints(x3, Ib,Plot.LINE)
	plot_raw.setColor(Color.BLACK)
	legendString = "Bleach\nNon-bleach\nBackground"
	plot_raw.setLegend(legendString,Plot.AUTO_POSITION)
	plot_raw.update()

	# Overlay fit curve, with oversampling (for plot)
	xfit = [ (t / 10.0  + bleach_frame) * frame_interval for t in range(10 * len(xtofit) ) ]
	yfit = []
	for xt in xfit:
		yfit.append( fitter.f( fitter.getParams(), xt - xfit[0]) )

	# plot normalized curve without fit
	plot = Plot(current_imp.getTitle() + " - FRAP", "Time ("+time_units+')', "NU")
	#plot.setLimits(0, max(x), 0, 1.0 )
	plot.setLimits(0, max(x), 0, yMax )
	plot.setLineWidth(2)
	plot.setColor(Color.BLACK)
	plot.addPoints(x, y, Plot.LINE)
	plot_window =  plot.show()

	# plot normalized curve offset to x = 0 being the bleach frame
	plot = Plot(current_imp.getTitle() + " - FRAP offset", "Time ("+time_units+')', "NU")
	#plot.setLimits(0, max(x)-bleach_frame, 0, 1.0 )
	plot.setLimits(0, max(x)-bleach_frame, 0, yMax )
	plot.setLineWidth(2)
	plot.setColor(Color.BLACK)
	
	#clone x array; can't use assignment operator, which doesn't clone the list
	xOffset = list(x)
	for i in range(len(xOffset)):
		xOffset[i] = xOffset[i] - bleach_frame

	plot.addPoints(xOffset[bleach_frame:], y[bleach_frame:], Plot.LINE)
	plot_offset =  plot.show()

	# plot normalized curve with fit
	plot_fit = Plot(current_imp.getTitle() + " - FRAP fit", "Time ("+time_units+')', "NU")
	#plot_fit.setLimits(0, max(x), 0, 1.0 )
	plot_fit.setLimits(0, max(x), 0, yMax )
	plot_fit.setLineWidth(2)
	
	plot_fit.setColor(Color.BLACK)
	plot_fit.addPoints(x, y, Plot.LINE)
	plot_fit.setColor(Color.RED)
	plot_fit.addPoints(xfit, yfit, Plot.LINE)

	plot_fit.setColor(Color.black);
	plot_fit_window =  plot_fit.show()
	legendString = "FRAP\nFit"
	plot_fit.setLegend(legendString,Plot.AUTO_POSITION)
	plot_fit.update()

	if createROIs:
		#draw duplicate image and color ROIs
		WindowManager.setCurrentWindow(myWindow)
		current_imp.setSlice(bleach_frame + 1)
		IJ.run("Select All")
		str6 = current_imp.getTitle() + ' - ROIs'
		IJ.run("Duplicate...", "title='" + str6 + "'")
		IJ.run("RGB Color")
		dupWindow = WindowManager.getCurrentWindow()
		dupImp = dupWindow.getImagePlus()
		dupImp.setRoi(roi_FRAP)
		IJ.run("Colors...", "foreground=red")
		IJ.run("Draw")
		dupImp.setRoi(roi_norm)
		IJ.run("Colors...", "foreground=cyan")
		IJ.run("Draw")
		dupImp.setRoi(roi_BG)
		IJ.run("Colors...", "foreground=yellow")
		IJ.run("Draw")
		IJ.run("Line Width...", "line=1")
		IJ.run("Colors...", "foreground=black")
		dupImp.setRoi(0,0,0,0)
		WindowManager.setCurrentWindow(myWindow)
		current_imp.setRoi(0,0,0,0)

	# Output FRAP parameters
	if fitMethod == "Single":
		if throughOrigin:
			myTextWindow.append("Fit constrained through origin: TRUE")
		else:
			myTextWindow.append("Fit constrained through origin: FALSE")
		thalf = math.log(2) / param_values[1]
		mobile_fraction = param_values[0]
		str1 = ('Half-life = %.2f ' + time_units) % thalf
		myTextWindow.append( str1 )
		str2 = "Mobile fraction = %.1f %%" % (100 * mobile_fraction)
		myTextWindow.append( str2 )
		myTextWindow.append( "" )
		myTextWindow.append("*******Details********" + fitter.getResultString() )

	else:
		thalf1 = math.log(2) / param_values[1]
		thalf2 = math.log(2) / param_values[3]
		mobile_fraction = param_values[0] + param_values[2]
		str1 = ('Half-life #1= %.2f ' + time_units) % thalf1
		myTextWindow.append( str1 )
		str2 = ('Half-life #2= %.2f ' + time_units) % thalf2
		myTextWindow.append( str2 )
		str3 = "Mobile fraction = %.1f %%" % (100 * mobile_fraction)
		myTextWindow.append( str3 )
		myTextWindow.append( "" )
		myTextWindow.append("*******Details********" + fitter.getResultString() )


	if autoSave:
		# Automatically save relevant files to new subdirectory
		# Get directory in which image resides
		# The following doesn't work!
		# myInfo = current_imp.getFileInfo()
		# The following does...
		myDir = current_imp.getOriginalFileInfo().directory

		# Create new subdirectory using the file name
		subDir = myWindow.title + "-FRAP"
		rootName = myWindow.title
		mySubDir = myDir + subDir

		# Get new directory to pass to the rest of the macro
		# See if this directory already exists and present dialog if it does
		IJ.selectWindow(rootName)
		dirExists = False
		if path.isdir(mySubDir):
			gd2 = GenericDialog("Warning")
			gd2.addMessage("Warning: Directory exists. Overwrite?")
			#if user selects "Cancel" in following then cease execution
			gd2.showDialog()
			dirExists = True
			myTextWindow.append("Directory already existed...overwrote old data")

		# Create new subdirectory if it doesn't exist
		if not dirExists: 
			mkdir(mySubDir)
		
		# Create text file with analysis results
		IJ.selectWindow("FRAP Results - " + current_imp.getTitle())
		IJ.saveAs("Text", mySubDir + System.getProperty("file.separator") + rootName + " - results.txt")

		# Save ROIs as ROI set
		roi_manager.select(0)
		roi_manager.runCommand("Set Color", "red")
		roi_manager.runCommand("Rename", "bleach")
		roi_manager.select(1)
		roi_manager.runCommand("Set Color", "cyan")
		roi_manager.runCommand("Rename", "non-bleach")
		roi_manager.select(2)
		roi_manager.runCommand("Set Color", "yellow")
		roi_manager.runCommand("Rename", "background")
		roi_manager.runCommand("Save", mySubDir + System.getProperty("file.separator") + rootName + " - RoiSet.zip")
		IJ.selectWindow(rootName)
		IJ.run(current_imp, "Select None", "")

		# Save image with ROIs if this was generated
		if createROIs:
			IJ.selectWindow(rootName + " - ROIs")
			IJ.saveAs("Tiff", str(mySubDir) + System.getProperty("file.separator") + rootName + " - ROIs.tif")
			# This save as call renames the window! Need to rename below if automatically closing it...
						
		# Save data as .csv file
		# Save FRAP curve
		myResults = plot_window.getResultsTable()
		myResults.saveAs(str(mySubDir) + System.getProperty("file.separator") + rootName + " - FRAP.csv")
		
		# Save FRAP fit data
		myResults2 = plot_fit_window.getResultsTable()
		#myResults2.setValue​(java.lang.String column, int row, java.lang.String value)
		myResults2.saveAs(str(mySubDir) + System.getProperty("file.separator") + rootName + " - FRAP fit.csv")

		# Save FRAP data offset to origin
		myResults3 = plot_offset.getResultsTable()
		myResults3.saveAs(str(mySubDir) + System.getProperty("file.separator") + rootName + " - FRAP offset.csv")

		# Save raw intensity plot
		#Use CSV routines in Jython to use custom column names
		# open the file first (if its not there, newly created)
		fullpath = str(mySubDir) + System.getProperty("file.separator") + rootName + " - raw.csv"
		#print("Path for raw data:")
		#print(fullpath)
		f = open(fullpath, 'wb')
 		# create csv writer
		writer = csv.writer(f)
		row = ["time","bleach","non-bleach","background"]
		writer.writerow(row)
		# for loop to write each row
		for i in range(n_slices):
			row = [x[i], If[i], In[i], Ib[i]]
			writer.writerow(row)
		# close the file. 
		f.close()

		if (autoSave & closeWindows):
			plot_fit_window.close()
			plot_window.close()
			plot_offset.close()
			plot_raw_window.close()
			IJ.selectWindow(rootName + " - ROIs.tif")
			IJ.run("Close")

	
#main
doFRAP()