import os  # Importing the os module to interact with the operating system
import sys  # Importing the sys module to access system-specific parameters and functions
from ij import IJ, WindowManager, Prefs  # Importing classes and methods from the ImageJ library
from ij.gui import NonBlockingGenericDialog  # Importing NonBlockingGenericDialog from ImageJ's GUI module
from java.io import File  # Importing the File class from Java's io package
from javax.swing import JOptionPane, JFileChooser  # Importing JOptionPane and JFileChooser from Java's Swing package
from fiji.plugin.trackmate import Model, Settings, TrackMate, SelectionModel, Logger  # Importing TrackMate-related classes from Fiji
from fiji.plugin.trackmate.detection import DogDetectorFactory  # Importing DogDetectorFactory for spot detection in TrackMate
from fiji.plugin.trackmate.tracking.jaqaman import SparseLAPTrackerFactory  # Importing SparseLAPTrackerFactory for tracking in TrackMate
from fiji.plugin.trackmate.features import FeatureFilter  # Importing FeatureFilter from TrackMate
from fiji.plugin.trackmate.features.track import TrackIndexAnalyzer  # Importing TrackIndexAnalyzer from TrackMate
from fiji.plugin.trackmate.gui.displaysettings import DisplaySettingsIO  # Importing DisplaySettingsIO for TrackMate display settings
from fiji.plugin.trackmate.visualization.table import AllSpotsTableView  # Importing AllSpotsTableView for displaying spots in a table
from fiji.plugin.trackmate.visualization.hyperstack import HyperStackDisplayer  # Importing HyperStackDisplayer for visualizing tracks in a hyperstack
from loci.plugins import BF  # Importing Bio-Formats from the loci.plugins package
from loci.plugins.in import ImporterOptions  # Importing ImporterOptions from loci.plugins.in

import csv  # Importing the csv module for handling CSV files
import glob  # Importing the glob module for file pattern matching
import math  # Importing the math module for mathematical functions

# Constants
DEBUG = False  # Debug mode flag
SAVE_SPOTS = False  # Flag to save detected spots
OUTPUT_DIR = os.path.join(os.path.expanduser("~"), "Desktop", "Output")  # Output directory for saving files
IMAGE_EXTENSIONS = ('.tif', '.tiff', '.jpg', '.jpeg', '.png', '.czi', '.nd2', '.lif', '.lsm', '.ome.tif', '.ome.tiff')  # Supported image file extensions
SAVE_INDIVIDUAL_CSVs=False

# Ensure UTF-8 encoding
reload(sys)  # Reload the sys module to change the default encoding
sys.setdefaultencoding('utf-8')  # Set the default encoding to UTF-8

# Template for TrackMate settings
TRACKMATE_SETTINGS_TEMPLATE = {
	'StageRepositioning.TrackMate.Subpixel_Localization': True,  # Enable subpixel localization
	'StageRepositioning.TrackMate.Spot_Diameter': 4.0,  # Spot diameter in microns
	'StageRepositioning.TrackMate.Threshold_Value': 20.904,  # Threshold value for spot detection
	'StageRepositioning.TrackMate.Median_Filtering': False,  # Enable/disable median filtering
	'StageRepositioning.TrackMate.BatchMode': False  # Batch mode flag
}

def log_message(message):
	"""Logs a message if debug mode is enabled."""
	if DEBUG:  # Check if debug mode is enabled
		IJ.log(message)  # Log the message using ImageJ's logging system

def read_preferences(settings):
	"""Reads the user preferences for TrackMate settings."""
	log_message("Reading Preferences...")  # Log the start of reading preferences
	output = {}  # Initialize an empty dictionary to store the preferences
	for key, default_value in settings.items():  # Iterate over the settings
		value = Prefs.get(key, str(default_value))  # Get the preference value
		if isinstance(default_value, bool):  # Convert to appropriate type
			value = bool(int(value))
		elif isinstance(default_value, float):
			value = float(value)
		elif isinstance(default_value, list):
			value = value.split(",")
		elif isinstance(default_value, str):
			value = str(value)
		output[key] = value  # Store the converted value in the output dictionary
	log_message("Reading Preferences: Done")  # Log the completion of reading preferences
	return output  # Return the preferences

def save_preferences(settings):
	"""Saves the user preferences for TrackMate settings."""
	log_message("Saving Preferences...")  # Log the start of saving preferences
	for key, value in settings.items():  # Iterate over the settings
		if isinstance(value, list):  # Convert to appropriate format for saving
			value = ",".join(map(str, value))
		elif isinstance(value, bool):
			value = int(value)
		Prefs.set(key, value)  # Save the preference value
	Prefs.savePreferences()  # Save all preferences
	log_message("Saving Preferences: Done")  # Log the completion of saving preferences

def get_image_info(imp):
	"""Retrieves and returns information about the image."""
	log_message("Getting Image Info...")  # Log the start of getting image info
	file_info = imp.getOriginalFileInfo()  # Get the file information of the image
	filename = file_info.fileName  # Get the filename
	basename, extension = os.path.splitext(filename)  # Split the filename into base name and extension
	input_dir = file_info.directory  # Get the directory of the file
	input_file_path = os.path.join(input_dir, filename)  # Get the full file path
	image_name = imp.getTitle()  # Get the title of the image
	width = imp.getWidth()  # Get the width of the image
	height = imp.getHeight()  # Get the height of the image
	nb_channels = imp.getNChannels()  # Get the number of channels in the image
	nb_slices = imp.getNSlices()  # Get the number of slices in the image
	nb_timepoints = imp.getNFrames()  # Get the number of timepoints in the image
	bit_depth = imp.getBitDepth()  # Get the bit depth of the image
	current_channel = imp.getChannel()  # Get the current channel being displayed
	current_slice = imp.getSlice()  # Get the current slice being displayed
	current_frame = imp.getFrame()  # Get the current frame being displayed
	calibration = imp.getCalibration()  # Get the calibration of the image
	pixel_width = calibration.pixelWidth  # Get the pixel width
	pixel_height = calibration.pixelHeight  # Get the pixel height
	pixel_depth = calibration.pixelDepth  # Get the pixel depth
	space_unit = calibration.getUnit()  # Get the unit of space
	time_unit = calibration.getTimeUnit()  # Get the unit of time
	frame_interval = calibration.frameInterval  # Get the frame interval
	
	# Store all image information in a dictionary
	image_info = {
		'InputFilePath': input_file_path,
		'InputDir': input_dir,
		'Filename': filename,
		'Basename': basename,
		'Extension': extension,
		'ImageName': image_name,
		'Width': width,
		'Height': height,
		'NbChannels': nb_channels,
		'NbSlices': nb_slices,
		'NbTimepoints': nb_timepoints,
		'BitDepth': bit_depth,
		'CurrentChannel': current_channel,
		'CurrentSlice': current_slice,
		'CurrentFrame': current_frame,
		'Calibration': calibration,
		'PixelWidth': pixel_width,
		'PixelHeight': pixel_height,
		'PixelDepth': pixel_depth,
		'SpaceUnit': space_unit,
		'TimeUnit': time_unit,
		'FrameInterval': frame_interval
	}
	
	log_message("Getting Image Info: Done.")  # Log the completion of getting image info
	return image_info  # Return the image information

def generate_unique_filepath(directory, basename, suffix, extension):
	"""Generates a unique file path to avoid overwriting existing files."""
	log_message("Generating Unique Filepath...")  # Log the start of generating unique filepath
	file_counter = 1  # Initialize a counter for file numbering
	while True:  # Loop until a unique filepath is found
		filename = "{}_{}-{:03d}{}".format(basename, suffix, file_counter, extension)  # Generate a filename
		filepath = os.path.join(directory, filename)  # Generate the full file path
		if not os.path.exists(filepath):  # Check if the file path exists
			log_message("Generating Unique Filepath: Done.")  # Log the completion of generating unique filepath
			return filepath  # Return the unique file path
		file_counter += 1  # Increment the counter if the file path exists

def select_folder(default_path):
	"""Opens a dialog for the user to select a folder."""
	log_message("Selecting Folder...")  # Log the start of selecting folder
	chooser = JFileChooser(default_path)  # Create a file chooser with the default path
	chooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY)  # Set the file chooser to select directories only
	chooser.setDialogTitle("Choose a directory containing the images to process")  # Set the dialog title
	return_value = chooser.showOpenDialog(None)  # Show the open dialog
	if return_value == JFileChooser.APPROVE_OPTION:  # Check if the user approved the selection
		input_dir = chooser.getSelectedFile().getAbsolutePath()  # Get the selected directory
	log_message("Selecting Folder: Done.")  # Log the completion of selecting folder
	return input_dir  # Return the selected directory

def open_image_bioformats(filepath):
	"""Opens an image using Bio-Formats."""
	log_message("Importing with Bioformats...")  # Log the start of importing with Bioformats
	options = ImporterOptions()  # Create importer options
	options.setId(filepath)  # Set the file path in the importer options
	try:
		imps = BF.openImagePlus(options)  # Open the image using Bio-Formats
		if imps and len(imps) > 0:  # Check if images were found
			log_message("Importing with Bioformats: Done.")  # Log the completion of importing with Bioformats
			return imps[0]  # Return the first image
		else:
			log_message("Importing with Bioformats: No Images found in the file: " + filepath + ".")  # Log if no images were found
			return None  # Return None if no images were found
	except Exception as e:
		log_message("Importing with Bioformats: Error opening file" + str(e) + ".")  # Log the error if an exception occurred
		return None  # Return None if an exception occurred

def run_trackmate_predetection(imp, trackmate_settings, save_spots):
	"""Runs the pre-detection step using TrackMate to detect spots in the image."""
	log_message("Running PreDetection...")  # Log the start of pre-detection
	image_info = get_image_info(imp)  # Get the image information
	
	max_quality_threshold_values = []  # Initialize a list to store max quality threshold values
	nb_detected_spots = 0  # Initialize the number of detected spots
	
	imp.setDisplayMode(IJ.COLOR)  # Set the display mode to color
	imp.setC(image_info['CurrentChannel'])  # Set the current channel
	imp.updateAndDraw()  # Update and draw the image
	model = Model()  # Create a new TrackMate model
	trackmate_settings_obj = Settings(imp)  # Create a new TrackMate settings object
	model.setPhysicalUnits(image_info['SpaceUnit'], image_info['TimeUnit'])  # Set the physical units in the model
	trackmate_settings_obj.detectorFactory = DogDetectorFactory()  # Set the detector factory
	trackmate_settings_obj.detectorSettings = {
		'DO_SUBPIXEL_LOCALIZATION': trackmate_settings['StageRepositioning.TrackMate.Subpixel_Localization'],  # Set subpixel localization
		'RADIUS': trackmate_settings['StageRepositioning.TrackMate.Spot_Diameter'] / 2,  # Set the spot radius
		'TARGET_CHANNEL': image_info['CurrentChannel'],  # Set the target channel
		'THRESHOLD': trackmate_settings['StageRepositioning.TrackMate.Threshold_Value'],  # Set the threshold value
		'DO_MEDIAN_FILTERING': trackmate_settings['StageRepositioning.TrackMate.Median_Filtering']  # Set median filtering
	}

	trackmate_settings_obj.trackerFactory = SparseLAPTrackerFactory()  # Set the tracker factory
	trackmate_settings_obj.trackerSettings = trackmate_settings_obj.trackerFactory.getDefaultSettings()  # Get the default tracker settings
	trackmate_settings_obj.trackerSettings['ALLOW_TRACK_SPLITTING'] = True  # Allow track splitting
	trackmate_settings_obj.trackerSettings['ALLOW_TRACK_MERGING'] = True  # Allow track merging
	trackmate_settings_obj.addAllAnalyzers()  # Add all analyzers

	trackmate = TrackMate(model, trackmate_settings_obj)  # Create a new TrackMate instance
	
	ok1 = trackmate.checkInput()  # Check the input
	ok2 = trackmate.process()  # Process the input
	if not ok1 or not ok2:  # Check if the input check or processing failed
		nb_detected_spots = 0  # Set the number of detected spots to 0
		max_quality_threshold_value = 100  # Set the max quality threshold value to 100
		return nb_detected_spots, max_quality_threshold_value  # Return the number of detected spots and max quality threshold value

	selection_model = SelectionModel(model)  # Create a new selection model
	ds = DisplaySettingsIO.readUserDefault()  # Read the user default display settings
	ds.setSpotDisplayRadius(0.9)  # Set the spot display radius
	ds.setSpotDisplayedAsRoi(True)  # Set the spots to be displayed as ROI
	ds.setSpotShowName(True)  # Set the spots to show names
	ds.setSpotTransparencyAlpha(0.7)  # Set the spot transparency alpha
	ds.setSpotFilled(True)  # Set the spots to be filled
	displayer = HyperStackDisplayer(model, selection_model, imp, ds)  # Create a new hyperstack displayer
	displayer.render()  # Render the displayer
	displayer.refresh()  # Refresh the displayer

	nb_detected_spots = int(model.getSpots().getNSpots(False))  # Get the number of detected spots
	
	max_quality = 0  # Initialize the max quality
	if nb_detected_spots > 0:  # Check if there are detected spots
		for spot in model.getSpots().iterable(True):  # Iterate over the detected spots
			spot_quality = spot.getFeature('QUALITY')  # Get the quality of the spot
			if max_quality is None or spot_quality > max_quality:  # Check if the current spot quality is higher than the max quality
				max_quality = spot_quality  # Update the max quality
		max_quality_threshold_values.append(max_quality)  # Append the max quality to the threshold values
		
		if save_spots:  # Check if saving spots is enabled
			all_spots_table = AllSpotsTableView(model, selection_model, ds, filename)  # Create a new all spots table view
			output_csv_path = generate_unique_filepath(OUTPUT_DIR, basename, Suffix="Spots-", Extension=".csv")  # Generate a unique file path for the CSV
			all_spots_table.exportToCsv(output_csv_path)  # Export the spots to CSV
		
	if max_quality_threshold_values:  # Check if there are max quality threshold values
		max_quality_threshold_value = max(max_quality_threshold_values)  # Set the max quality threshold value to the max of the threshold values
	else:
		max_quality_threshold_value = 100  # Set the max quality threshold value to 100
	log_message("Running PreDetection: Done.")  # Log the completion of pre-detection
		
	return nb_detected_spots, max_quality_threshold_value  # Return the number of detected spots and max quality threshold value

def display_stage_repositioning_dialog(imp, trackmate_settings, nb_detected_spots, test_detection, max_quality_threshold_value, dialog_counter, SAVE_INDIVIDUAL_CSVs):
	"""Displays a dialog for the user to adjust TrackMate settings for stage repositioning."""
	global DEBUG
	log_message("Displaying Stage Repositioning Dialog...")  # Log a message indicating that the dialog is being displayed.
	image_info = get_image_info(imp)  # Retrieve image information.
	gd = NonBlockingGenericDialog("Stage Repositioning with Batch TrackMate")  # Create a non-blocking dialog titled "Stage Repositioning with Batch TrackMate".
	gd.addMessage("=== Detection Settings ===")  # Add a message to the dialog for detection settings.
	gd.addCheckbox("Enable Subpixel Localization", trackmate_settings['StageRepositioning.TrackMate.Subpixel_Localization'])  # Add a checkbox for subpixel localization.
	gd.addNumericField("Spot Diameter (microns):", trackmate_settings['StageRepositioning.TrackMate.Spot_Diameter'], 2, 5, "")  # Add a numeric field for spot diameter.
	gd.addSlider("Threshold Value:", 0, max_quality_threshold_value, trackmate_settings['StageRepositioning.TrackMate.Threshold_Value'])  # Add a slider for threshold value.
	gd.addCheckbox("Apply Median Filtering", trackmate_settings['StageRepositioning.TrackMate.Median_Filtering'])  # Add a checkbox for median filtering.
	gd.addCheckbox("Batch Mode", trackmate_settings['StageRepositioning.TrackMate.BatchMode'])  # Add a checkbox for batch mode.

	# Handle the number of detected spots and provide feedback to the user.
	if nb_detected_spots == False:
		nb_detected_spots = 0
	if nb_detected_spots > image_info['NbTimepoints']:
		gd.addMessage("Nb of detected spots too high. Increase the threshold.")  # Add a message if detected spots are too high.
	elif nb_detected_spots < image_info['NbTimepoints']:
		gd.addMessage("No of detected spots too low. Decrease the threshold.")  # Add a message if detected spots are too low.
	elif nb_detected_spots == image_info['NbTimepoints']:
		gd.addMessage("Exactly one spot detected per Frame. You can proceed.")  # Add a message if detected spots are correct.
	
	gd.addMessage("Nb of detected spots: " + str(nb_detected_spots) + ". Nb of Frame: " + str(image_info['NbTimepoints']))  # Display the number of detected spots and frames.
	gd.addCheckbox("Test Detection", test_detection)  # Add a checkbox for test detection.
	gd.addCheckbox("Save individual CSVs", SAVE_INDIVIDUAL_CSVs)
	gd.addCheckbox("Chatty Mode", DEBUG)  # Add a checkbox for chatty mode (debug mode).
	gd.showDialog()  # Show the dialog to the user.

	if gd.wasOKed():
		user_click = "OK"  # Set user click status to "OK".
		trackmate_settings_user = {}
		trackmate_settings_user['StageRepositioning.TrackMate.Subpixel_Localization'] = gd.getNextBoolean()  # Get user input for subpixel localization.
		trackmate_settings_user['StageRepositioning.TrackMate.Spot_Diameter'] = gd.getNextNumber()  # Get user input for spot diameter.
		trackmate_settings_user['StageRepositioning.TrackMate.Threshold_Value'] = gd.getNextNumber()  # Get user input for threshold value.
		trackmate_settings_user['StageRepositioning.TrackMate.Median_Filtering'] = gd.getNextBoolean()  # Get user input for median filtering.
		trackmate_settings_user['StageRepositioning.TrackMate.BatchMode'] = gd.getNextBoolean()  # Get user input for batch mode.
		test_detection = gd.getNextBoolean()  # Get user input for test detection.
		SAVE_INDIVIDUAL_CSVs=gd.getNextBoolean()
		DEBUG = gd.getNextBoolean()  # Get user input for debug mode.
		save_preferences(trackmate_settings_user)  # Save the user preferences.
		dialog_counter += 1  # Increment the dialog counter.
	elif gd.wasCanceled():
		user_click = "Cancel"  # Set user click status to "Cancel".
	
	log_message("Displaying Stage Repositioning Dialog: Done.")  # Log a message indicating that the dialog has been displayed.
	return trackmate_settings_user, user_click, nb_detected_spots, image_info['NbTimepoints'], test_detection, dialog_counter, SAVE_INDIVIDUAL_CSVs  # Return the updated settings and status.

def run_trackmate_detection(imp, trackmate_settings_stored, SAVE_INDIVIDUAL_CSVs):
	"""Runs the TrackMate detection process on the image."""
	log_message("Running TrackMate Detection...")  # Log a message indicating that TrackMate detection is starting.
	image_info = get_image_info(imp)  # Retrieve image information.
	model = Model()  # Create a new TrackMate model.
	trackmate_settings_obj = Settings(imp)  # Create a new TrackMate settings object.
	trackmate_settings_obj.detectorFactory = DogDetectorFactory()  # Set the detector factory to DogDetectorFactory.
	model.setPhysicalUnits(image_info['SpaceUnit'], image_info['TimeUnit'])  # Set the physical units for the model.
	trackmate_settings_obj.detectorSettings = {
		'DO_SUBPIXEL_LOCALIZATION': trackmate_settings_stored['StageRepositioning.TrackMate.Subpixel_Localization'],  # Set subpixel localization.
		'RADIUS': trackmate_settings_stored['StageRepositioning.TrackMate.Spot_Diameter'] / 2,  # Set the radius for spot detection.
		'TARGET_CHANNEL': image_info['CurrentChannel'],  # Set the target channel.
		'THRESHOLD': trackmate_settings_stored['StageRepositioning.TrackMate.Threshold_Value'],  # Set the threshold value.
		'DO_MEDIAN_FILTERING': trackmate_settings_stored['StageRepositioning.TrackMate.Median_Filtering']  # Set median filtering.
	}

	trackmate_settings_obj.trackerFactory = SparseLAPTrackerFactory()  # Set the tracker factory to SparseLAPTrackerFactory.
	trackmate_settings_obj.trackerSettings = trackmate_settings_obj.trackerFactory.getDefaultSettings()  # Get the default tracker settings.
	trackmate_settings_obj.trackerSettings['ALLOW_TRACK_SPLITTING'] = True  # Enable track splitting.
	trackmate_settings_obj.trackerSettings['ALLOW_TRACK_MERGING'] = True  # Enable track merging.
	trackmate_settings_obj.addAllAnalyzers()  # Add all analyzers to the settings.

	trackmate = TrackMate(model, trackmate_settings_obj)  # Create a new TrackMate instance.
	if not trackmate.checkInput() or not trackmate.process():
		return 0, 100  # Return if TrackMate fails.

	selection_model = SelectionModel(model)  # Create a new selection model.
	ds = DisplaySettingsIO.readUserDefault()  # Read the user default display settings.
	ds.setSpotDisplayRadius(0.9)  # Set the spot display radius.
	ds.setSpotDisplayedAsRoi(True)  # Set spots to be displayed as ROI.
	ds.setSpotShowName(True)  # Show spot names.
	ds.setSpotTransparencyAlpha(0.7)  # Set spot transparency.
	ds.setSpotFilled(True)  # Set spots to be filled.
	displayer = HyperStackDisplayer(model, selection_model, imp, ds)  # Create a new hyperstack displayer.
	displayer.render()  # Render the display.
	displayer.refresh()  # Refresh the display.

	nb_detected_spots = model.getSpots().getNSpots(False)  # Get the number of detected spots.
	
	spots_data = []  
	if nb_detected_spots > 0:
		for spot in model.getSpots().iterable(True):  # Iterate over the detected spots
			spot_data = {
				'SPOT_ID': spot.ID(),
				'QUALITY': spot.getFeature('QUALITY'),
				'POSITION_X': spot.getFeature('POSITION_X'),
				'POSITION_Y': spot.getFeature('POSITION_Y'),
				'POSITION_Z': spot.getFeature('POSITION_Z'),
				'POSITION_T': spot.getFeature('POSITION_T'),
				'FRAME': spot.getFeature('FRAME'),
				'RADIUS': spot.getFeature('RADIUS'),
				'VISIBILITY': spot.getFeature('VISIBILITY'),
				'MANUAL_SPOT_COLOR': spot.getFeature('MANUAL_SPOT_COLOR'),
				'MEAN_INTENSITY_CH1': spot.getFeature('MEAN_INTENSITY_CH1'),
				'MEDIAN_INTENSITY_CH1': spot.getFeature('MEDIAN_INTENSITY_CH1'),
				'MIN_INTENSITY_CH1': spot.getFeature('MIN_INTENSITY_CH1'),
				'MAX_INTENSITY_CH1': spot.getFeature('MAX_INTENSITY_CH1'),
				'TOTAL_INTENSITY_CH1': spot.getFeature('TOTAL_INTENSITY_CH1'),
				'STD_INTENSITY_CH1': spot.getFeature('STD_INTENSITY_CH1'),
				'CONTRAST_CH1': spot.getFeature('CONTRAST_CH1'),
				'SNR_CH1': spot.getFeature('SNR_CH1'),
				'SourceFile': image_info['Filename']
			}
			spots_data.append(spot_data)
	if SAVE_INDIVIDUAL_CSVs:
		all_spots_table = AllSpotsTableView(model, selection_model, ds, image_info['Filename'])  # Create a table view of all spots.
	   	output_csv_path = generate_unique_filepath(OUTPUT_DIR, image_info['Basename'], "Stage-Repositioning", ".csv")  # Generate a unique file path for the CSV.
	   	all_spots_table.exportToCsv(output_csv_path)  # Export the spots table to CSV.
	
	log_message("Running TrackMate Detection: Done.")  # Log a message indicating that TrackMate detection is done.
	return nb_detected_spots, image_info['NbTimepoints'], spots_data   # Return the number of detected spots and the number of timepoints.

def process_image(imp, SAVE_INDIVIDUAL_CSVs):
	"""Processes the image using TrackMate with user-defined settings."""
	image_name = imp.getTitle()  # Get the title of the image.
	log_message("Processing Image:" + str(image_name) + "...")  # Log a message indicating that the image is being processed.
	test_detection = False  # Initialize test detection as False.
	dialog_counter = 0  # Initialize the dialog counter.
	user_click = None  # Initialize user click status as None.
	all_spots_data = []  # List to accumulate spots data
	while True:
		trackmate_settings_stored = read_preferences(TRACKMATE_SETTINGS_TEMPLATE)  # Read the stored TrackMate settings.
		nb_detected_spots, max_quality_threshold_value = run_trackmate_predetection(imp, trackmate_settings_stored, SAVE_SPOTS)  # Run pre-detection to get the number of detected spots and max quality threshold value.
		trackmate_settings_user, user_click, nb_detected_spots, nb_timepoints, test_detection, dialog_counter, SAVE_INDIVIDUAL_CSVs = display_stage_repositioning_dialog(imp, trackmate_settings_stored, nb_detected_spots, test_detection, max_quality_threshold_value, dialog_counter, SAVE_INDIVIDUAL_CSVs)  # Display the stage repositioning dialog.
		save_preferences(trackmate_settings_user)  # Save the user-defined TrackMate settings.
		
		trackmate_settings_stored_filtered = {key: value for key, value in trackmate_settings_stored.items() if key != 'StageRepositioning.TrackMate.BatchMode'}  # Filter out batch mode from stored settings.
		trackmate_settings_user_filtered = {key: value for key, value in trackmate_settings_user.items() if key != 'StageRepositioning.TrackMate.BatchMode'}  # Filter out batch mode from user settings.
		
		# Check conditions to break the loop or handle user cancellation.
		if user_click == "OK" and nb_detected_spots == nb_timepoints and trackmate_settings_stored_filtered == trackmate_settings_user_filtered and not test_detection:
			break
		elif user_click == "Cancel":
			log_message("Processing Image:" + str(image_name) + ": User canceled Operation.")  # Log a message indicating user cancellation.
			if image_from_folder:
				imp.close()  # Close the image if it is from a folder.
			sys.exit(1)  # Exit the program.
	
	nb_detected_spots, nb_timepoints,spots_data  = run_trackmate_detection(imp, trackmate_settings_stored, SAVE_INDIVIDUAL_CSVs)  # Run TrackMate detection with the stored settings.
	all_spots_data.extend(spots_data)  # Accumulate spots data
	return trackmate_settings_user, nb_timepoints,SAVE_INDIVIDUAL_CSVs, all_spots_data  # Return the user-defined settings and number of timepoints.
	
def process_image_batch(imp, close, SAVE_INDIVIDUAL_CSVs):
	"""Processes the image in batch mode using TrackMate with stored settings."""
	image_name = imp.getTitle()  # Get the title of the image.
	log_message("Processing in Batch Image:" + str(image_name) + "...")  # Log a message indicating the start of batch processing for the image.
	trackmate_settings_stored = read_preferences(TRACKMATE_SETTINGS_TEMPLATE)  # Read the stored TrackMate settings.
	nb_detected_spots, nb_timepoints, spots_data  = run_trackmate_detection(imp, trackmate_settings_stored, SAVE_INDIVIDUAL_CSVs)  # Run TrackMate detection with the stored settings.
	if close:
		imp.close()  # If the close flag is True, close the image.
	log_message("Processing in Batch Image:" + str(image_name) + ": Done.")  # Log a message indicating that batch processing is done for the image.
	return nb_detected_spots, nb_timepoints, spots_data  # Return the number of detected spots and the number of timepoints.

def clean_filename(filename):
	"""Cleans the filename by removing specific substrings."""
	log_message("Cleaning Filename "+str(filename)+" ...")
	filename = filename.replace(" - ", "_")
	filename = filename.replace(".czi", "")
	filename = filename.replace("_spots", "")
	filename = filename.replace(".csv", "")
	
	log_message("Cleaning Filename "+str(filename)+": Done")
	return filename

def is_float(value):
	"""Checks if a value can be converted to a float."""
	try:
		float(value)
		return True
	except ValueError:
		return False

def merge_and_process_spots_data(all_spots_data, output_csv_path):
	"""Merges, processes, and orders spots data from a list of dictionaries."""
	log_message("Merging and Processing Spots Data...")  # Log start

	if not all_spots_data:
		log_message("No spots data available to merge.")
		return

	# Define the desired column order, including placeholders for the known columns
	column_order = [
		'SPOT_ID', 'QUALITY', 'POSITION_X', 'POSITION_Y', 'POSITION_Z', 'POSITION_T', 'FRAME',
		'RADIUS', 'VISIBILITY', 'MANUAL_SPOT_COLOR', 'MEAN_INTENSITY_CH1', 'MEDIAN_INTENSITY_CH1',
		'MIN_INTENSITY_CH1', 'MAX_INTENSITY_CH1', 'TOTAL_INTENSITY_CH1', 'STD_INTENSITY_CH1',
		'CONTRAST_CH1', 'SNR_CH1', 'SourceFile', 'Time (sec)', 'Time (min)', 'X (nm)', 'Y (nm)', 'Z (nm)',
		'Displacement 3D (nm)', 'Displacement 2D (nm)'
	]

	# Ensure all data has required fields and dynamically add 'Variable-xxx' keys to each spot
	max_variables = 0
	for spot_data in all_spots_data:
		source_file = spot_data.get("SourceFile", "")
		# Clean the source file name and split into variables
		cleaned_filename = clean_filename(source_file)
		spot_data['SourceFile'] = cleaned_filename
		
		# Split the cleaned filename by underscores and add variables to the spot_data
		filename_parts = cleaned_filename.split('_')
		for i, part in enumerate(filename_parts):
			variable_key = 'Variable-%03d' % (i+1)  # String formatting for Python 2
			spot_data[variable_key] = part
		
		max_variables = max(max_variables, len(filename_parts))  # Track the maximum number of variables

	# Add 'Variable-xxx' to the column order
	for i in range(max_variables):
		column_order.append('Variable-%03d' % (i+1))

	# Group data by SourceFile
	grouped_data = {}
	for spot_data in all_spots_data:
		source_file = spot_data.get("SourceFile", "")
		cleaned_filename = clean_filename(source_file)
		spot_data['SourceFile'] = cleaned_filename
		if cleaned_filename not in grouped_data:
			grouped_data[cleaned_filename] = []
		grouped_data[cleaned_filename].append(spot_data)

	# Process and compute differences and displacements per SourceFile
	processed_data = []
	for source_file, spots in grouped_data.items():
		# Sort spots by Frame
		sorted_spots = sorted(spots, key=lambda x: int(x.get("FRAME", 0)))

		# Initialize reference and previous positions
		ref_x = ref_y = ref_z = None
		prev_x = prev_y = prev_z = None

		for i, spot in enumerate(sorted_spots):
			# Set reference positions only for the first frame (FRAME == 0)
			if i == 0 and int(spot.get("FRAME", -1)) == 0:
				if is_float(spot["POSITION_X"]):
					ref_x = float(spot["POSITION_X"])
				if is_float(spot["POSITION_Y"]):
					ref_y = float(spot["POSITION_Y"])
				if is_float(spot["POSITION_Z"]):
					ref_z = float(spot["POSITION_Z"])

			# Compute relative positions in nm (from the reference)
			x = float(spot["POSITION_X"]) if is_float(spot["POSITION_X"]) else None
			y = float(spot["POSITION_Y"]) if is_float(spot["POSITION_Y"]) else None
			z = float(spot["POSITION_Z"]) if is_float(spot["POSITION_Z"]) else None

			spot["X (nm)"] = (x - ref_x) * 1000 if x is not None and ref_x is not None else 0
			spot["Y (nm)"] = (y - ref_y) * 1000 if y is not None and ref_y is not None else 0
			spot["Z (nm)"] = (z - ref_z) * 1000 if z is not None and ref_z is not None else 0

			# Compute differences relative to the previous frame
			x_diff = y_diff = z_diff = 0  # Default differences are 0
			if prev_x is not None and x is not None:
				x_diff = (x - prev_x) * 1000  # Convert to nm
			if prev_y is not None and y is not None:
				y_diff = (y - prev_y) * 1000  # Convert to nm
			if prev_z is not None and z is not None:
				z_diff = (z - prev_z) * 1000  # Convert to nm

			# Compute 2D and 3D displacements
			spot["Displacement 2D (nm)"] = math.sqrt(x_diff ** 2 + y_diff ** 2)
			spot["Displacement 3D (nm)"] = math.sqrt(x_diff ** 2 + y_diff ** 2 + z_diff ** 2)
			spot['Time (sec)'] = float(spot["POSITION_T"]) if is_float(spot["POSITION_T"]) else None
			spot['Time (min)'] = spot['Time (sec)'] / 60

			# Update previous positions for the next iteration
			prev_x, prev_y, prev_z = x, y, z

			processed_data.append(spot)

	# Sort all processed data by SourceFile and Frame for final output
	final_data = sorted(processed_data, key=lambda x: (x.get("SourceFile", ""), int(x.get("FRAME", 0))))

	# Write the processed data to CSV (Python 2 compatible)
	with open(output_csv_path, mode='wb') as output_file:
		log_message("Writing Data to CSV...")
		writer = csv.DictWriter(output_file, fieldnames=column_order)
		writer.writeheader()
		for row in final_data:
			# Convert all values to UTF-8 if they are unicode
			writer.writerow({k: v.encode('utf-8') if isinstance(v, unicode) else v for k, v in row.items()})

	log_message("Merging and Processing Spots Data: Done.")




### Main Script here

# Create output directory on the user's desktop if it doesn't already exist
if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)  # Create the output directory for saving files

# Initialize variables
image_from_folder = False  # Flag to track if the images are coming from a folder
open_images = WindowManager.getImageTitles()  # Get titles of all open images in ImageJ
all_spots_data = []  # Initialize an empty list to store all spot data

# Check if there are any open images in ImageJ
if open_images:
    # Loop through each open image and process it
    for index, image_title in enumerate(open_images):
        imp = WindowManager.getImage(image_title)  # Get the ImagePlus object for the current image
        image = WindowManager.getFrame(imp.getTitle())  # Get the image's frame
        image.toFront()  # Bring the image window to the front
        image_name = imp.getTitle()  # Get the title of the current image
        
        # For the first image, run the full processing pipeline
        if index == 0:
            trackmate_settings_user, nb_timepoints, SAVE_INDIVIDUAL_CSVs, spots_data = process_image(imp, SAVE_INDIVIDUAL_CSVs)  # Process the image with user settings
            IJ.log("Success processing " + str(image_name) + ".")  # Log success message
            all_spots_data.extend(spots_data)  # Add the spot data from this image to the list
        else:
            # For subsequent images, check if batch mode is enabled
            if trackmate_settings_user['StageRepositioning.TrackMate.BatchMode']:
                nb_detected_spots, nb_timepoints, spots_data = process_image_batch(imp, False, SAVE_INDIVIDUAL_CSVs)  # Process the image in batch mode
                if nb_detected_spots == nb_timepoints:  # Check if the number of detected spots matches the number of timepoints
                    IJ.log("Success in batch processing " + str(image_name) + ".")  # Log success message
                    all_spots_data.extend(spots_data)  # Add the spot data from this image to the list
                else:
                    log_message("Detection failed for " + str(image_name) + ".\nDetected spots " + str(nb_detected_spots) + ". Nb of Frame: " + str(nb_timepoints))  # Log failure message
                    trackmate_settings_user, nb_timepoints, SAVE_INDIVIDUAL_CSVs, spots_data = process_image(imp, SAVE_INDIVIDUAL_CSVs)  # Reprocess the image
                    all_spots_data.extend(spots_data)  # Add the spot data from this image to the list
                    IJ.log("Success processing " + str(image_name) + ".")  # Log success after reprocessing
            else:
                trackmate_settings_user, nb_timepoints, SAVE_INDIVIDUAL_CSVs, spots_data = process_image(imp, SAVE_INDIVIDUAL_CSVs)  # Process the image if batch mode is not enabled
                all_spots_data.extend(spots_data)  # Add the spot data from this image to the list
                IJ.log("Success processing " + str(image_name) + ".")  # Log success message
else:
    # If no images are open, prompt the user to select a folder containing images
    folder_input_path = select_folder(default_path=os.path.join(os.path.expanduser("~"), "Desktop"))  # Open folder selection dialog with default path to Desktop
    if folder_input_path:
        # Get a list of valid image files in the selected folder
        image_files = [
            f for f in os.listdir(folder_input_path)
            if f.lower().endswith(IMAGE_EXTENSIONS)  # Only include files with valid image extensions
        ]
        
        if not image_files:
            # Show error message if no valid image files were found in the selected folder
            JOptionPane.showMessageDialog(None, "No valid image files found in the selected folder.", "Error", JOptionPane.ERROR_MESSAGE)
            sys.exit(1)  # Exit the program
       
        image_from_folder = True  # Set flag indicating images are coming from a folder
        # Process each image in the folder
        for index, image_file in enumerate(image_files):
            imp = open_image_bioformats(os.path.join(folder_input_path, image_file))  # Open image using Bio-Formats
            if imp is None:
                continue  # Skip the image if it could not be opened
            image_name = imp.getTitle()  # Get the image title
            
            # For the first image, run the full processing pipeline
            if index == 0:
                trackmate_settings_user, nb_timepoints, SAVE_INDIVIDUAL_CSVs, spots_data = process_image(imp, SAVE_INDIVIDUAL_CSVs)  # Process the image with user settings
                all_spots_data.extend(spots_data)  # Add the spot data from this image to the list
                IJ.log("Success processing " + str(image_name) + ".")  # Log success message
                imp.close()  # Close the image after processing
            else:
                # For subsequent images, check if batch mode is enabled
                if trackmate_settings_user['StageRepositioning.TrackMate.BatchMode']:
                    nb_detected_spots, nb_timepoints, spots_data = process_image_batch(imp, False, SAVE_INDIVIDUAL_CSVs)  # Process the image in batch mode
                    if nb_detected_spots == nb_timepoints:  # Check if the number of detected spots matches the number of timepoints
                        all_spots_data.extend(spots_data)  # Add the spot data from this image to the list
                        IJ.log("Success batch processing " + str(image_name) + ".")  # Log success message for batch processing
                        imp.close()  # Close the image after processing
                    else:
                        IJ.log("Detection failed for " + image_name + ".\nDetected spots " + str(nb_detected_spots) + ". Nb of Frame: " + str(nb_timepoints))  # Log failure message
                        trackmate_settings_user, nb_timepoints, SAVE_INDIVIDUAL_CSVs, spots_data = process_image(imp, SAVE_INDIVIDUAL_CSVs)  # Reprocess the image
                        all_spots_data.extend(spots_data)  # Add the spot data from this image to the list
                        IJ.log("Success processing " + str(image_name) + ".")  # Log success after reprocessing
                        imp.close()  # Close the image after reprocessing
                else:
                    trackmate_settings_user, nb_timepoints, SAVE_INDIVIDUAL_CSVs, spots_data = process_image(imp, SAVE_INDIVIDUAL_CSVs)  # Process the image if batch mode is not enabled
                    all_spots_data.extend(spots_data)  # Add the spot data from this image to the list
                    IJ.log("Success processing " + str(image_name) + ".")  # Log success message
                    imp.close()  # Close the image after processing
    else:
        # Show error message if no folder is selected by the user
        JOptionPane.showMessageDialog(None, "No folder selected.", "Error", JOptionPane.ERROR_MESSAGE)
        sys.exit(1)  # Exit the program

# Determine the number of processed images
if open_images:
    Nb_Processed_Images = len(open_images)  # Count processed images if images were open in ImageJ
elif image_files:
    Nb_Processed_Images = len(image_files)  # Count processed images if images were from a folder
else:
    Nb_Processed_Images = 1  # Default to 1 if no images were processed

# Log the success message indicating the number of processed images
MessageProcessing = "Stage Repositioning with Batch TrackMate completed.\n" + format(Nb_Processed_Images) + " images have been processed successfully."
if SAVE_INDIVIDUAL_CSVs:
    MessageProcessing = MessageProcessing + "\nFiles are saved in " + str(OUTPUT_DIR)
IJ.log(MessageProcessing)  

# Generate a unique filepath for merging the CSV files
merged_data_output_path = generate_unique_filepath(OUTPUT_DIR, basename="Stage-Repositioning", suffix="Merged-Data", extension=".csv")
IJ.log("Merging and Processing CSV Files...")
merge_and_process_spots_data(all_spots_data, merged_data_output_path)  # Merge and process the collected spot data

# Notify the user of successful completion
MessageMerging = "Merging and processing completed successfully.\nMerged Data saved as {}".format(merged_data_output_path)
IJ.log(MessageMerging)
MessageFinal = MessageProcessing + "\n" + MessageMerging  # Combine messages for final notification
JOptionPane.showMessageDialog(None, MessageFinal, "Stage Repositioning with Batch TrackMate", JOptionPane.INFORMATION_MESSAGE)

# Load and install necessary libraries at the beginning
rm(list = ls())

# Check and install required packages if they are not already installed
if (!require(dplyr)) install.packages("dplyr", dependencies = TRUE)
if (!require(stringr)) install.packages("stringr", dependencies = TRUE)
if (!require(ggplot2)) install.packages("ggplot2", dependencies = TRUE)
if (!require(tcltk)) install.packages("tcltk", dependencies = TRUE)

# Load libraries
library(dplyr)
library(stringr)
library(ggplot2)
library(tcltk)

# Set default input and output directories
default_input_dir <- file.path(Sys.getenv("USERPROFILE"), "Desktop")  # Default to "Output" on Desktop
InputFolder <- tclvalue(tkchooseDirectory(initialdir = default_input_dir, title = "Select a folder containing CSV files"))

#InputFolder <- default_input_dir  # Use default folder

# Specify the Output folder (this is a fixed folder on the Desktop)
OutputFolder <- file.path(Sys.getenv("USERPROFILE"), "Desktop", "Output")
if (!dir.exists(OutputFolder)) dir.create(OutputFolder, recursive = TRUE)

# List all CSV files in the folder
csv_files <- list.files(path = InputFolder, pattern = "\\.csv$", full.names = TRUE)
if (length(csv_files) == 0) stop("No CSV files found in the selected directory.")

header <- names(read.csv(csv_files[1], nrows = 1))

# Function to clean the filename
clean_filename <- function(filename) {
  filename_parts <- strsplit(filename, " - ")[[1]]
  cleaned_filename <- sub("\\.czi", "", filename_parts[1])
  cleaned_filename <- sub("_spots", "", cleaned_filename)
  cleaned_filename <- sub("\\.csv", "", cleaned_filename)
  return(cleaned_filename)
}

# Read and merge all CSV files
merged_data <- csv_files %>%
  lapply(function(file) {
    data <- read.csv(file, skip = 4, header = FALSE)
    colnames(data) <- header
    data <- data %>% arrange(FRAME)
    
    # Clean and add source file info
    filename <- basename(file)
    data$SourceFile <- clean_filename(filename)
    
    # Extract variables from the filename
    filename_parts <- strsplit(clean_filename(filename), "_")[[1]]
    for (i in seq_along(filename_parts)) {
      variable_name <- paste("Variable-", sprintf("%03d", i), sep = "")
      data[[variable_name]] <- filename_parts[i]
    }
    
    # Add time columns if available
    if ("POSITION_T" %in% colnames(data)) {
      data$`Time (sec)` <- round(data$POSITION_T, 0)
      data$`Time (min)` <- round(data$`Time (sec)` / 60, 2)
    }
    
    # Calculate displacement columns (X, Y, Z, 3D, 2D)
    if ("POSITION_X" %in% colnames(data)) {
      first_value <- data$POSITION_X[data$FRAME == 0][[1]]
      data$`X (nm)` <- (data$POSITION_X - first_value) * 1000
    }
    if ("POSITION_Y" %in% colnames(data)) {
      first_value <- data$POSITION_Y[data$FRAME == 0][[1]]
      data$`Y (nm)` <- (data$POSITION_Y - first_value) * 1000
    }
    if ("POSITION_Z" %in% colnames(data)) {
      first_value <- data$POSITION_Z[data$FRAME == 0][[1]]
      data$`Z (nm)` <- (data$POSITION_Z - first_value) * 1000
    }
    
    # Calculate displacement (3D and 2D)
    if (all(c("X (nm)", "Y (nm)", "Z (nm)") %in% colnames(data))) {
      data$`Displacement 3D (nm)` <- sqrt(
        diff(c(0, data$`X (nm)`))^2 + 
          diff(c(0, data$`Y (nm)`))^2 + 
          diff(c(0, data$`Z (nm)`))^2
      )
      data$`Displacement 3D (nm)`[1] <- 0
    }
    if (all(c("X (nm)", "Y (nm)") %in% colnames(data))) {
      data$`Displacement 2D (nm)` <- sqrt(
        diff(c(0, data$`X (nm)`))^2 + 
          diff(c(0, data$`Y (nm)`))^2
      )
      data$`Displacement 2D (nm)`[1] <- 0
    }
    
    return(data)
  }) %>%
  bind_rows()

# Save the merged data to a new CSV file
output_file <- "Stage-Repositioning_Merged-Data.csv"
write.csv(merged_data, file = file.path(OutputFolder, output_file), row.names = FALSE)

cat("All CSV files have been merged and saved to", output_file, "\n")