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

# 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

# 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):
    """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("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.
        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  # Return the updated settings and status.

def run_trackmate_detection(imp, trackmate_settings_stored, save_spots):
    """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.
    if nb_detected_spots > 0:
        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']  # Return the number of detected spots and the number of timepoints.

def process_image(imp):
    """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.
    
    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 = display_stage_repositioning_dialog(imp, trackmate_settings_stored, nb_detected_spots, test_detection, max_quality_threshold_value, dialog_counter)  # 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 = run_trackmate_detection(imp, trackmate_settings_stored, SAVE_SPOTS)  # Run TrackMate detection with the stored settings.
    return trackmate_settings_user, nb_timepoints  # Return the user-defined settings and number of timepoints.
    
def process_image_batch(imp, close):
    """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 = run_trackmate_detection(imp, trackmate_settings_stored, SAVE_SPOTS)  # 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  # Return the number of detected spots and the number of timepoints.

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

# Initialize variables
image_from_folder = False  # Flag to indicate if images are from a folder.
open_images = WindowManager.getImageTitles()  # Get the titles of all open images in ImageJ.

if open_images:
    # If images are already open in ImageJ, process them
    for index, image_title in enumerate(open_images):
        imp = WindowManager.getImage(image_title)  # Get the ImagePlus object for the given image title.
        image = WindowManager.getFrame(imp.getTitle())  # Get the frame for the image.
        image.toFront()  # Bring the image frame to the front.
        image_name = imp.getTitle()  # Get the title of the image.
        
        # For the first image, run the full process
        if index == 0:
            trackmate_settings_user, nb_timepoints = process_image(imp)  # Process the image and get user-defined settings.
            IJ.log("Success processing " + str(image_name) + ".")  # Log a success message.
        else:
            # For subsequent images, check if batch mode is enabled
            if trackmate_settings_user['StageRepositioning.TrackMate.BatchMode']:
                nb_detected_spots, nb_timepoints = process_image_batch(imp, close=False)  # Process the image in batch mode.
                if nb_detected_spots == nb_timepoints:
                    IJ.log("Success in batch processing " + str(image_name) + ".")  # Log a success message for batch processing.
                else:
                    log_message("Detection failed for " + str(image_name) + ".\nDetected spots " + str(nb_detected_spots) + ". Nb of Frame: " + str(nb_timepoints))  # Log a failure message.
                    trackmate_settings_user, nb_timepoints = process_image(imp)  # Reprocess the image if detection failed.
                    IJ.log("Success processing " + str(image_name) + ".")  # Log a success message after reprocessing.
            else:
                trackmate_settings_user, nb_timepoints = process_image(imp)  # Process the image if batch mode is not enabled.
                IJ.log("Success processing " + str(image_name) + ".")  # Log a success message.
    JOptionPane.showMessageDialog(None, "Stage Repositioning with Batch TrackMate completed.\n{} images have been processed successfully.\nFiles are saved in {}".format(len(image_files), OUTPUT_DIR), "Stage Repositioning with Batch TrackMate", JOptionPane.INFORMATION_MESSAGE)  # Show a message dialog indicating the completion of the process.
    IJ.log("{} images have been processed successfully.\nFiles are saved in {}".format(len(open_images), OUTPUT_DIR))  # Log the completion message.

else:
    # If no images are open, prompt the user to select a folder
    folder_input_path = select_folder(default_path=os.path.join(os.path.expanduser("~"), "Desktop"))  # Prompt the user to select a folder.
    if folder_input_path:
        image_files = [
            f for f in os.listdir(folder_input_path)
            if f.lower().endswith(IMAGE_EXTENSIONS)
        ]  # Get a list of valid image files in the selected folder.
        if not image_files:
            JOptionPane.showMessageDialog(None, "No valid image files found in the selected folder.", "Error", JOptionPane.ERROR_MESSAGE)  # Show an error message if no valid images are found.
            sys.exit(1)  # Exit the program.
        
        image_from_folder = True  # Set the flag to indicate images are from a folder.
        for index, image_file in enumerate(image_files):
            imp = open_image_bioformats(os.path.join(folder_input_path, image_file))  # Open the image using Bio-Formats.
            if imp is None:
                continue  # Skip the image if it could not be opened.
            image_name = imp.getTitle()  # Get the title of the image.
            
            # For the first image, run the full process
            if index == 0:
                trackmate_settings_user, nb_timepoints = process_image(imp)  # Process the image and get user-defined settings.
                IJ.log("Success processing " + str(image_name) + ".")  # Log a success message.
                imp.close()  # Close the image.
            else:
                # For subsequent images, check if batch mode is enabled
                if trackmate_settings_user['StageRepositioning.TrackMate.BatchMode']:
                    nb_detected_spots, nb_timepoints = process_image_batch(imp, close=False)  # Process the image in batch mode.
                    if nb_detected_spots == nb_timepoints:
                        IJ.log("Success batch processing " + str(image_name) + ".")  # Log a success message for batch processing.
                        imp.close()  # Close the image.
                    else:
                        IJ.log("Detection failed for " + image_name + ".\nDetected spots " + str(nb_detected_spots) + ". Nb of Frame: " + str(nb_timepoints))  # Log a failure message.
                        trackmate_settings_user, nb_timepoints = process_image(imp)  # Reprocess the image if detection failed.
                        IJ.log("Success processing " + str(image_name) + ".")  # Log a success message after reprocessing.
                        imp.close()  # Close the image.
                else:
                    trackmate_settings_user, nb_timepoints = process_image(imp)  # Process the image if batch mode is not enabled.
                    IJ.log("Success processing " + str(image_name) + ".")  # Log a success message.
                    imp.close()  # Close the image.
        JOptionPane.showMessageDialog(None, "Stage Repositioning with Batch TrackMate completed.\n{} images have been processed successfully.\nFiles are saved in {}".format(len(image_files), OUTPUT_DIR), "Stage Repositioning with Batch TrackMate", JOptionPane.INFORMATION_MESSAGE)  # Show a message dialog indicating the completion of the process.
        IJ.log("{} images have been processed successfully.\nFiles are saved in {}".format(len(image_files), OUTPUT_DIR))  # Log the completion message.
    
    else:
        JOptionPane.showMessageDialog(None, "No folder selected.", "Error", JOptionPane.ERROR_MESSAGE)  # Show an error message if no folder is selected.
        sys.exit(1)  # Exit the program.

# 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")