# --------------------------
# User Prompt Toggle
# --------------------------
# Set this to False to skip the initial dialog and reuse previous settings
show_initial_prompt = True  # Change this as needed

import os
import sys
import json
from ij import IJ, WindowManager
from java.io import File
from ij.gui import GenericDialog
from fiji.plugin.trackmate import Model, Settings, TrackMate, SelectionModel, Logger
from fiji.plugin.trackmate.detection import DogDetectorFactory
from fiji.plugin.trackmate.tracking.jaqaman import SparseLAPTrackerFactory
from fiji.plugin.trackmate.features import FeatureFilter
from fiji.plugin.trackmate.features.track import TrackIndexAnalyzer
from fiji.plugin.trackmate.gui.displaysettings import DisplaySettingsIO
from fiji.plugin.trackmate.gui.displaysettings.DisplaySettings import TrackMateObject
from fiji.plugin.trackmate.visualization.table import TrackTableView
from fiji.plugin.trackmate.visualization.hyperstack import HyperStackDisplayer

# Ensure UTF-8 encoding
reload(sys)
sys.setdefaultencoding('utf-8')

# Path to the config file that will store last user settings
config_file_path = os.path.join(os.path.expanduser("~"), "Desktop", "trackmate_config.json")

# Default settings in case there's no previous config file
default_settings = {
    'subpixel_localization': True,
    'spot_diameter': 4.0,  # Default 0.5 microns
    'threshold_value': 20.904,  # Default threshold value
    'apply_median_filtering': False
}

# Function to load settings from the config file
def load_settings():
    if os.path.exists(config_file_path):
        with open(config_file_path, 'r') as f:
            return json.load(f)
    else:
        return default_settings

# Function to save settings to the config file
def save_settings(settings):
    with open(config_file_path, 'w') as f:
        json.dump(settings, f)

# Load previous settings (or use default if no file exists)
user_settings = load_settings()

# Create output directory on the user's desktop
output_dir = os.path.join(os.path.expanduser("~"), "Desktop", "Output")
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# ----------------------------
# Create the model object
# ----------------------------
model = Model()
model.setLogger(Logger.IJ_LOGGER)  # Send messages to the ImageJ log window

# ------------------------
# Create input dialog
# ------------------------

# Create the dialog box for user input
gd = GenericDialog("TrackMate Parameters")

# Add fields to the dialog, prefill with last saved values
gd.addCheckbox("Enable Subpixel Localization?", user_settings['subpixel_localization'])  # Default: from saved settings
gd.addNumericField("Spot Diameter (microns):", user_settings['spot_diameter'], 2)  # Default diameter: from saved settings
gd.addSlider("Threshold Value:", 0, 255, user_settings['threshold_value'])  # Slider for threshold value (0-255 range, default: from saved settings)
gd.addCheckbox("Apply Median Filtering?", user_settings['apply_median_filtering'])  # Default: from saved settings
gd.addCheckbox("Process All Open Images?", False)  # Add a checkbox to select processing mode

# Show the dialog
gd.showDialog()

# Check if the user canceled the dialog
if gd.wasCanceled():
    sys.exit("User canceled the operation.")

# Get user inputs from the dialog
subpixel_localization = gd.getNextBoolean()  # Whether to enable subpixel localization
spot_diameter = gd.getNextNumber()  # Spot diameter in microns
threshold_value = gd.getNextNumber()  # Threshold value from the slider
apply_median_filtering = gd.getNextBoolean()  # Whether to apply median filtering
process_all_images = gd.getNextBoolean()  # Whether to process all open images

# Save the new settings to the configuration file
user_settings = {
    'subpixel_localization': subpixel_localization,
    'spot_diameter': spot_diameter,
    'threshold_value': threshold_value,
    'apply_median_filtering': apply_median_filtering
}
save_settings(user_settings)

# ------------------------
# Prepare settings object
# ------------------------

# If processing all images
if process_all_images:
    open_images = WindowManager.getImageTitles()
    total_images = len(open_images)
    
    if total_images == 0:
        IJ.log("No images are open!")
        sys.exit("No images are open!")

    for idx, title in enumerate(open_images):
        imp = WindowManager.getImage(title)

        # Ensure the image is not null
        if imp is None:
            continue

        # Get the image title for naming output files
        filename = imp.getTitle()

        # Configure TrackMate for this image
        settings = Settings(imp)
        settings.detectorFactory = DogDetectorFactory()
        settings.detectorSettings = {
            'DO_SUBPIXEL_LOCALIZATION': subpixel_localization,  # Set subpixel localization
            'RADIUS': spot_diameter / 2,  # Convert diameter to radius for detector
            'TARGET_CHANNEL': 1,  # Target channel (customize if needed)
            'THRESHOLD': threshold_value,  # User-defined threshold value
            'DO_MEDIAN_FILTERING': apply_median_filtering,  # Apply median filtering if selected
        }

        # Configure tracker
        settings.trackerFactory = SparseLAPTrackerFactory()
        settings.trackerSettings = settings.trackerFactory.getDefaultSettings()
        settings.trackerSettings['ALLOW_TRACK_SPLITTING'] = True
        settings.trackerSettings['ALLOW_TRACK_MERGING'] = True

        # Add all known feature analyzers to compute track statistics
        settings.addAllAnalyzers()

        # Configure track filters
        track_filter = FeatureFilter('TRACK_DISPLACEMENT', 10, False)  # Filter tracks with displacement >10 pixels
        settings.addTrackFilter(track_filter)

        # -------------------
        # Instantiate plugin
        # -------------------
        trackmate = TrackMate(model, settings)

        # --------
        # Process
        # --------
        if not trackmate.checkInput():
            IJ.log("Error checking input: " + trackmate.getErrorMessage())
            continue
        if not trackmate.process():
            IJ.log("Error during processing: " + trackmate.getErrorMessage())
            continue

        # ----------------
        # Display results
        # ----------------
        selection_model = SelectionModel(model)

        # Read the default display settings
        ds = DisplaySettingsIO.readUserDefault()
        ds.setTrackColorBy(TrackMateObject.TRACKS, TrackIndexAnalyzer.TRACK_INDEX)
        ds.setSpotColorBy(TrackMateObject.TRACKS, TrackIndexAnalyzer.TRACK_INDEX)

        # Display tracks and spots on the image
        displayer = HyperStackDisplayer(model, selection_model, imp, ds)
        displayer.render()
        displayer.refresh()

        # -----------------
        # Export results
        # -----------------
        # Export spot table as CSV
        spot_table = TrackTableView.createSpotTable(model, ds)
        output_csv_path = os.path.join(output_dir, filename + "_spots.csv")
        spot_table.exportToCsv(File(output_csv_path))

        # Close the image after processing
        imp.close()

        # Update progress bar
        IJ.showProgress(idx + 1, total_images)

else:
    # If processing only the active image, leave it open after processing
    imp = WindowManager.getCurrentImage()

    if imp is None:
        sys.exit("No image is currently open.")

    filename = imp.getTitle()

    # Configure TrackMate for this image
    settings = Settings(imp)
    settings.detectorFactory = DogDetectorFactory()
    settings.detectorSettings = {
        'DO_SUBPIXEL_LOCALIZATION': subpixel_localization,  # Set subpixel localization
        'RADIUS': spot_diameter / 2,  # Convert diameter to radius for detector
        'TARGET_CHANNEL': 1,  # Target channel (customize if needed)
        'THRESHOLD': threshold_value,  # User-defined threshold value
        'DO_MEDIAN_FILTERING': apply_median_filtering,  # Apply median filtering if selected
    }

    # Configure tracker
    settings.trackerFactory = SparseLAPTrackerFactory()
    settings.trackerSettings = settings.trackerFactory.getDefaultSettings()
    settings.trackerSettings['ALLOW_TRACK_SPLITTING'] = True
    settings.trackerSettings['ALLOW_TRACK_MERGING'] = True

    # Add all known feature analyzers to compute track statistics
    settings.addAllAnalyzers()

    # Configure track filters
    track_filter = FeatureFilter('TRACK_DISPLACEMENT', 10, False)  # Filter tracks with displacement >10 pixels
    settings.addTrackFilter(track_filter)

    # -------------------
    # Instantiate plugin
    # -------------------
    trackmate = TrackMate(model, settings)

    # --------
    # Process
    # --------
    if not trackmate.checkInput():
        IJ.log("Error checking input: " + trackmate.getErrorMessage())
    if not trackmate.process():
        IJ.log("Error during processing: " + trackmate.getErrorMessage())

    # ----------------
    # Display results
    # ----------------
    selection_model = SelectionModel(model)

    # Read the default display settings
    ds = DisplaySettingsIO.readUserDefault()
    ds.setTrackColorBy(TrackMateObject.TRACKS, TrackIndexAnalyzer.TRACK_INDEX)
    ds.setSpotColorBy(TrackMateObject.TRACKS, TrackIndexAnalyzer.TRACK_INDEX)

    # Display tracks and spots on the image
    displayer = HyperStackDisplayer(model, selection_model, imp, ds)
    displayer.render()
    displayer.refresh()

    # -----------------
    # Export results
    # -----------------
    # Export spot table as CSV
    spot_table = TrackTableView.createSpotTable(model, ds)
    output_csv_path = os.path.join(output_dir, filename + "_spots.csv")
    spot_table.exportToCsv(File(output_csv_path))

# Show completion message
IJ.log("Processing Complete!")

# 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(corrplot)) install.packages("corrplot", dependencies = TRUE)

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

# Set default input and output directories
default_input_dir <- file.path(Sys.getenv("USERPROFILE"), "Desktop", "Output")  # Default to "Output" on Desktop
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 <- "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")

# Load merged data for visualization (or analysis)