import os
import csv
import traceback
import time

from ij import IJ, WindowManager
from ij.gui import ImageCanvas, ImageWindow, OvalRoi, Overlay
from ij.plugin.frame import RoiManager
from ij.measure import ResultsTable, Measurements
from ij.plugin.filter import ParticleAnalyzer


from java.io import File, IOException
from java.nio.file import Files, StandardCopyOption, Paths
from java.beans import PropertyChangeListener
from java.lang import Runnable, System

from javax.swing import (JFrame, JDialog, JMenuBar, JMenu, JMenuItem, JSplitPane,
                         JPanel, JComboBox, JScrollPane, JOptionPane, JTree, JTable,
                         JButton, JLabel, JFileChooser, ListSelectionModel, BorderFactory,
                         JTextField, JList, JCheckBox, DefaultListModel,
                         SwingWorker, JProgressBar, ProgressMonitor, SwingUtilities)
from javax.swing.table import AbstractTableModel, DefaultTableModel
from javax.swing.tree import DefaultMutableTreeNode, DefaultTreeModel
from javax.swing.event import ListSelectionListener, ListDataListener
from javax.swing.border import EmptyBorder
from javax.swing.filechooser import FileNameExtensionFilter

from java.awt import BorderLayout, FlowLayout, Font, GridLayout, Cursor
from java.awt.event import WindowAdapter, MouseAdapter, KeyListener, FocusAdapter, ActionListener

#==============================================
# Project structure and file managment
#==============================================

class ProjectImage(object):
    """ Simple class to hold info about a single image file """
    def __init__(self, filename, project_path):
        self.filename = filename
        self.full_path = os.path.join(project_path, "Images", filename)

        base_name, _ = os.path.splitext(self.filename)
        self.roi_path = os.path.join(project_path, "ROI_Files", base_name + "_ROIs.zip")
        self.outline_path = os.path.join(project_path, "Final_Cell_Selections", base_name + "_Outlines.zip")

        self.rois = [] # list of dictionaries
        self.status = "Pending ROIs" 
    
    def has_outlines(self):
        """ Check if image has corrosponding cell outline selections file """
        return os.path.exists(self.outline_path)

    def has_roi(self):
        """ Checks if corrosponding ROI file exists """
        return os.path.exists(self.roi_path)
    
    def add_roi(self, roi_data):
        """ Adds an ROI's data to the image"""
        self.rois.append(roi_data)

    def _load_rois_from_zip(self):
        """
        Loads all ROIs directly from the .zip file, making it the source of truth
        for names and individual bregma values.
        """
        if self.has_roi():
            rm = RoiManager(True)
            try:
                rm.open(self.roi_path)
                rois_array = rm.getRoisAsArray()
                self.rois = [] # Clear any old data
                for roi in rois_array:
                    self.rois.append({
                        'roi_name': roi.getName(),
                        'bregma': roi.getProperty("comment") or 'N/A'
                    })
            finally:
                rm.close()

class Project(object):
    """ Class representing a project, holding its structure and data once opened from folder """
    def __init__(self, root_dir):
        self.root_dir = root_dir
        self.name = os.path.basename(os.path.normpath(root_dir))
        self.paths = self._discover_paths()
        self._verify_and_create_dirs()
        self.images = [] # list of ProjectImage objects
        self._load_project_db()
        self._scan_for_new_images()
        self.images.sort(key=self._get_natural_sort_key)

    def _get_natural_sort_key(self, image_object):
        """ correctly sorts filenames by extracting leading number """
        try:
            return int(image_object.filename.split('_')[0])
        except (ValueError, IndexError):
            return float('inf')
        
    def _verify_and_create_dirs(self):
        """ Check for essential project files and creates them if missing"""
        for key, path in self.paths.items():
            if not os.path.exists(path):
                try:
                    # For csv databases
                    if path.endswith(".csv"):
                        headers = []
                        if key == 'roi_db':
                            headers = ['filename', 'roi_name', 'bregma', 'status']
                        elif key == 'image_status_db': 
                            headers = ['filename', 'status']
                        elif key == 'results_db':
                            headers = ['filename', 'roi_name', 'roi_area', 'bregma_value', 'cell_count', 'total_cell_area' ]

                        if headers:
                            with open(path, 'w') as csvfile:
                                writer = csv.writer(csvfile)
                                writer.writerow(headers)
                                IJ.log("Created missing project database: {}".format(path))
                    else:
                        os.makedirs(path)
                        IJ.log("Created missing project directory: {}".format(path))
                except OSError as e:
                    IJ.log("Error creating directory {}: {}".format(path, e))

    def _discover_paths(self):
        """ Creates dict of essential project components """
        return {
            'images': os.path.join(self.root_dir, 'Images'),
            'rois': os.path.join(self.root_dir, 'ROI_Files'),
            'processed': os.path.join(self.root_dir, 'Processed_Images'),
            'probabilities': os.path.join(self.root_dir, 'Ilastik_Probabilites'),
            'cell_outlines': os.path.join(self.root_dir, 'Final_Cell_Selections'),
            'temp': os.path.join(self.root_dir, 'temp'),
            'roi_db': os.path.join(self.root_dir, 'Roi_DB.csv'),
            'image_status_db': os.path.join(self.root_dir, 'Image_Status_DB.csv'),
            'results_db': os.path.join(self.root_dir, 'Results_DB.csv')
        }

    def _load_project_db(self):
        """
        Loads and parses both databases, but ONLY for images that currently exist
        in the Images folder, effectively pruning missing entries.
        """
        images_map = {}
        images_dir = self.paths['images']

        # Load Image Status DB
        status_db_path = self.paths['image_status_db']
        if os.path.exists(status_db_path):
            with open(status_db_path, 'r') as csvfile:
                reader = csv.DictReader(csvfile)
                for row in reader:
                    filename = row['filename']
                    
                    # --- NEW CHECK ---
                    # Verify the image file actually exists before processing its DB entry.
                    image_path = os.path.join(images_dir, filename)
                    if not os.path.exists(image_path):
                        continue # Skip this entry if the image file is missing.
                    # --- END NEW CHECK ---

                    if filename not in images_map:
                        images_map[filename] = ProjectImage(filename, self.root_dir)
                    images_map[filename].status = row.get('status', 'New')

        # Load ROI DB
        roi_db_path = self.paths['roi_db']
        if os.path.exists(roi_db_path):
            with open(roi_db_path, 'r') as csvfile:
                reader = csv.DictReader(csvfile)
                for row in reader:
                    filename = row['filename']

                    # --- NEW CHECK ---
                    # Also check here to catch images that might only be in the ROI DB.
                    image_path = os.path.join(images_dir, filename)
                    if not os.path.exists(image_path):
                        continue # Skip this entry if the image file is missing.
                    # --- END NEW CHECK ---

                    if filename not in images_map:
                        images_map[filename] = ProjectImage(filename, self.root_dir)
                    images_map[filename].add_roi(row)

        # Loop through all loaded images and populate from zip files as before
        for image in images_map.values():
            image._load_rois_from_zip()

        self.images = sorted(images_map.values(), key=lambda img: img.filename)

    def _scan_for_new_images(self):
        """ Scans images folder for any files not already loaded from the DBs. """
        if not os.path.isdir(self.paths['images']):
            return
        
        existing_filenames = {img.filename for img in self.images}
        for f in sorted(os.listdir(self.paths['images'])):
            if f.lower().endswith(('.tif', '.tiff', 'jpg', 'jpeg')) and f not in existing_filenames:
                new_image = ProjectImage(f, self.root_dir)
                new_image.status = "Pending ROIs"
                new_image._load_rois_from_zip() # new images
                self.images.append(new_image)

    def sync_project_db(self):
        """ Master save function that syncs both databases. """
        roi_success = self._sync_roi_db()
        status_success = self._sync_image_status_db()
        return roi_success and status_success

    def _sync_roi_db(self):
        """ Rewrites the Roi_DB.csv (ROI data) from memory. """
        db_path = self.paths['roi_db']
        headers = ['filename', 'roi_name', 'bregma', 'status']
        try:
            with open(db_path, 'wb') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=headers)
                writer.writeheader()
                for image in self.images:
                    if not image.rois:
                        continue # Skip images with no ROIs
                    for roi_data in image.rois:
                        row = {
                            'filename': image.filename,
                            'roi_name': roi_data.get('roi_name', 'N/A'),
                            'bregma': roi_data.get('bregma', 'N/A'),
                            'status': roi_data.get('status', 'Pending')
                        }
                        writer.writerow(row)
            return True
        except IOError as e:
            IJ.log("Error syncing ROI DB: {}".format(e))
            return False

    def _sync_image_status_db(self):
        """ Rewrites the Image_Status_DB.csv from memory. """
        db_path = self.paths['image_status_db']
        headers = ['filename', 'status']
        try:
            with open(db_path, 'wb') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=headers)
                writer.writeheader()
                for image in self.images:
                    writer.writerow({'filename': image.filename, 'status': image.status})
            return True
        except IOError as e:
            IJ.log("Error syncing Image Status DB: {}".format(e))
            return False

#==============================================
# Main GUI Classes
#==============================================

class ProjectManagerGUI(WindowAdapter):
    """ Builds and manages the main GUI, facilitating dialogs and and controling the script """
    def __init__(self):
        self.project = None
        self.unsaved_changes = False
        self.save_proj_item = None
        
        self.frame = JFrame("Project Manager")
        self.frame.setSize(900, 700)
        self.frame.setLayout(BorderLayout())
        self.frame.setDefaultCloseOperation(JFrame.DO_NOTHING_ON_CLOSE)

        self.build_menu()
        self.build_main_panel()
        self.build_status_bar()

        self.frame.addWindowListener(self)

    def show(self):
        self.frame.setLocationRelativeTo(None)
        self.frame.setVisible(True)

    def build_menu(self):
        menu_bar = JMenuBar()
        file_menu = JMenu("File")
        open_proj_item = JMenuItem("Open Project", actionPerformed=self.open_project_action)
        self.save_proj_item = JMenuItem("Save Project", actionPerformed=self.save_project_action, enabled=False)
        exit_item = JMenuItem("Exit", actionPerformed=lambda event: self.frame.dispose())
        file_menu.add(open_proj_item)
        file_menu.add(self.save_proj_item)
        file_menu.addSeparator()
        file_menu.add(exit_item)
        menu_bar.add(file_menu)
        self.frame.setJMenuBar(menu_bar)

    def build_main_panel(self):
        # Project header
        self.project_name_label = JLabel("No Project Loaded")
        self.project_name_label.setFont(Font("SansSerif", Font.BOLD, 16))
        self.project_name_label.setBorder(EmptyBorder(10,10,10,10))
        self.frame.add(self.project_name_label, BorderLayout.NORTH)

        # File Tree
        root_node = DefaultMutableTreeNode("Project")
        self.tree_model = DefaultTreeModel(root_node)
        self.file_tree = JTree(self.tree_model)
        tree_scroll_pane = JScrollPane(self.file_tree)

        right_panel = JPanel(BorderLayout())

        # Image table 
        image_cols = ["Filename", "ROI File", "# ROIs", "Status"]
        self.image_table_model = DefaultTableModel(None, image_cols)
        self.image_table = JTable(self.image_table_model)
        self.image_table.setSelectionMode(ListSelectionModel.MULTIPLE_INTERVAL_SELECTION)
        self.image_table.getSelectionModel().addListSelectionListener(self.on_image_selection)
        image_table_pane = JScrollPane(self.image_table)
        image_table_pane.setBorder(BorderFactory.createTitledBorder("Project Images"))
        
        # ROI detail table
        self.roi_table = JTable()
        roi_table_pane = JScrollPane(self.roi_table)
        roi_table_pane.setBorder(BorderFactory.createTitledBorder("ROI Details (Editable)"))
        
        # Split pane for two tables
        right_split_pane = JSplitPane(JSplitPane.VERTICAL_SPLIT, image_table_pane, roi_table_pane)
        right_split_pane.setDividerLocation(300)
        right_panel.add(right_split_pane, BorderLayout.CENTER)

        # Main split pane for tree and tables
        main_split_pane = JSplitPane(JSplitPane.HORIZONTAL_SPLIT, tree_scroll_pane, right_panel)
        main_split_pane.setDividerLocation(220)
        self.frame.add(main_split_pane, BorderLayout.CENTER)

    def build_status_bar(self):
        control_panel = JPanel(BorderLayout())
        control_panel.setBorder(EmptyBorder(5,5,5,5))

        self.status_label = JLabel("Open a project folder to begin")
        control_panel.add(self.status_label, BorderLayout.CENTER)
        
        button_panel = JPanel(FlowLayout(FlowLayout.RIGHT))

        self.import_button = JButton("Import Images", enabled=False)
        self.select_all_button = JButton("Select All / None")
        self.roi_button = JButton("Define/Edit ROIs", enabled=False)
        self.quant_button = JButton("Run Quantification", enabled=False)
        self.show_results_button = JButton("Show Results", enabled=False)

        button_panel.add(self.import_button)
        button_panel.add(self.select_all_button)
        button_panel.add(self.roi_button)
        button_panel.add(self.quant_button)
        button_panel.add(self.show_results_button)

        control_panel.add(button_panel, BorderLayout.EAST)
        self.frame.add(control_panel, BorderLayout.SOUTH)

        self.import_button.addActionListener(self.import_images_action)
        self.select_all_button.addActionListener(self.toggle_select_all_action)
        self.roi_button.addActionListener(self.open_roi_editor_action)
        self.quant_button.addActionListener(self.open_quantification_dialog_action)
        self.show_results_button.addActionListener(self.show_results_action)

    def set_unsaved_changes(self, state):
        """ Updates UI to show if there are unsaved changes """
        self.unsaved_changes = state
        self.save_proj_item.setEnabled(state)
        title = "Project Manager"
        if state:
            title += " *"
        self.frame.setTitle(title)

    # Event Handlers and actions

    def open_project_action(self, event):
        chooser = JFileChooser()
        chooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY)
        chooser.setDialogTitle("Select Project Directory")
        if chooser.showOpenDialog(self.frame) == JFileChooser.APPROVE_OPTION:
            project_dir = chooser.getSelectedFile().getAbsolutePath()
            self.load_project(project_dir)

    def save_project_action(self, event):
        """ Saves current state of project to csv file"""
        if not (self.project and self.unsaved_changes):
            return True

        # Sync database
        if self.project.sync_project_db():
            self.status_label.setText("Project saved successfully.")
            self.set_unsaved_changes(False)
            return True
        else:
            self.status_label.setText("Error saving project. See Log.")
            return False
        
    def show_results_action(self, event):
        """Launches the ResultsViewer dialog for the selected image."""
        selected_row = self.image_table.getSelectedRow()
        if selected_row == -1: return

        selected_image = self.project.images[selected_row]
        
        # The ResultsViewer will handle checking for the outline file
        viewer = ResultsViewer(self.frame, selected_image)
        viewer.show()

    def on_image_selection(self, event):
        """ 
        Called when the user selects image(s) in the top table.
        It can also be called programmatically by passing event=None to refresh the view.
        """
        # This condition allows the method's logic to run either when a user 
        # selection event has finalized (getValueIsAdjusting is False) or 
        # when the method is called directly without an event.
        if event is None or not event.getValueIsAdjusting():
            selection_count = self.image_table.getSelectedRowCount()

            # Enable/disable action buttons based on how many images are selected
            self.roi_button.setEnabled(selection_count == 1)
            self.quant_button.setEnabled(selection_count > 0)

            if selection_count == 1:
                selected_row = self.image_table.getSelectedRow()
                # Safety check in case the selection is cleared before this code runs
                if selected_row == -1: 
                    return

                selected_image = self.project.images[selected_row]
                self.status_label.setText("Selected: {}".format(selected_image.filename))
                self.show_results_button.setEnabled(selected_image.has_outlines())

                # Populate the bottom ROI details table for the selected image
                editable_model = EditableROIsTableModel(selected_image)
                editable_model.addTableModelListener(lambda e: self.set_unsaved_changes(True))
                self.roi_table.setModel(editable_model)

            elif selection_count > 1:
                self.status_label.setText("Selected: {} images".format(selection_count))
                # Clear the details table when multiple images are selected
                self.roi_table.setModel(EditableROIsTableModel(None)) 
                self.show_results_button.setEnabled(False)

            else: # Corresponds to selection_count == 0
                self.status_label.setText("No Image(s) Selected")
                # Clear the details table when the selection is empty
                self.roi_table.setModel(EditableROIsTableModel(None)) 
                self.show_results_button.setEnabled(False)

    def toggle_select_all_action(self, event):
        """ Selects all rows in the image table if not all are selected or clears selection if all are already selected"""
        row_count = self.image_table.getRowCount()
        if row_count == 0:
            return
        
        selected_count = self.image_table.getSelectedRowCount()

        if selected_count == row_count:
            self.image_table.clearSelection()
        else:
            self.image_table.selectAll()

    def open_roi_editor_action(self, event):
        """ Opens ROI editor window for selected image """
        selected_row = self.image_table.getSelectedRow()
        if selected_row != -1:
            selected_image = self.project.images[selected_row]

            editor = ROIEditor(self, self.project, selected_image)
            editor.show()

    def open_quantification_dialog_action(self, event):
        """ Gathers selected images and opens the quantification settings dialog. """
        selected_rows = self.image_table.getSelectedRows()
        if not selected_rows: return

        selected_images = [self.project.images[row] for row in selected_rows]

        quant_dialog = QuantificationDialog(self.frame, selected_images)
        settings = quant_dialog.show_dialog()

        if settings:
            progress_dialog = ProgressDialog(self.frame, "Processing images...", 100)
            worker = QuantificationWorker(self, self.project, settings, progress_dialog)
            progress_dialog.setVisible(True)
            worker.execute()

    def import_images_action(self, event):
        """Opens a file chooser and starts the background import process."""
        if not self.project:
            return
        
        chooser = JFileChooser()
        chooser.setDialogTitle("Select Images to Import")
        chooser.setMultiSelectionEnabled(True)
        chooser.setFileFilter(FileNameExtensionFilter("Image Files (tif, tiff, jpg, jpeg)", ["tif","tiff","jpg","jpeg"]))

        if chooser.showOpenDialog(self.frame) == JFileChooser.APPROVE_OPTION:
            selected_files = chooser.getSelectedFiles()

            # 1. Create an instance of our new worker class
            worker = ImageImportWorker(self, self.project, selected_files)

            # 2. Create a ProgressMonitor to watch the worker
            progress_monitor = ProgressMonitor(self.frame, "Importing Images", "Starting...", 0, 100)
            progress_monitor.setMillisToDecideToPopup(100) # Show the dialog quickly

            # 3. Link the worker's progress changes to the monitor's display
            class ProgressListener(PropertyChangeListener):
                def propertyChange(self, evt):
                    prop = evt.getPropertyName()
                    if "progress" == prop:
                        progress_monitor.setProgress(evt.getNewValue())
                    elif "note" == prop:
                        progress_monitor.setNote(evt.getNewValue())
                    
                    if progress_monitor.isCanceled():
                        worker.cancel(True)
            
            worker.addPropertyChangeListener(ProgressListener())

            # 4. Start the background task
            worker.execute()


    def windowClosing(self, event):
        """ Called when user attempts to close window, intercepts and prompts to save changes """
        if self.unsaved_changes:
            title = "Unsaved Changes"
            message = "You have unsaved changes. Would you like to save before closing?"

            # show dialog
            result = JOptionPane.showConfirmDialog(self.frame, message, title, JOptionPane.YES_NO_CANCEL_OPTION)

            if result == JOptionPane.YES_OPTION:
                if self.save_project_action(None):
                    self.frame.dispose()
                # If save fails, do nothing

            elif result == JOptionPane.NO_OPTION:
                self.frame.dispose()

            # if cancel, do nothing

        else: # no unsaved changes
            self.frame.dispose()

    # UI update logic
    def load_project(self, project_dir):
        """ Loads a project's data and update entire UI"""
        self.status_label.setText("Loading Project {}".format(project_dir))
        try:
            self.project = Project(project_dir)
            self.update_ui_for_project()

            self.import_button.setEnabled(True)

            self.status_label.setText("Sucessfully loaded project: {}".format(self.project.name))
            self.set_unsaved_changes(False)
        except Exception as e:
            self.status_label.setText("Error Loading Project. See Log for details")
            IJ.log("--- ERROR while loading project ---")
            IJ.log(traceback.format_exc())
            IJ.log("-----------------------------------")

    def update_ui_for_project(self):
        """ Populates the UI componenets with the current project's data """
        if not self.project:
            return
        
        # Update name
        self.project_name_label.setText("Project: " + self.project.name)
        
        # Image table
        while self.image_table_model.getRowCount() > 0:
            self.image_table_model.removeRow(0)
        
        for img in self.project.images:
            roi_file_status = "Yes" if img.has_roi() else "No"
            self.image_table_model.addRow([
                img.filename,
                roi_file_status,
                len(img.rois),
                img.status
            ])

        # update file tree 
        root_node = DefaultMutableTreeNode(self.project.name)
        for name, path in self.project.paths.items():
            # show directorys and key files
            if os.path.isdir(path) or name.endswith('_db'):
                node = DefaultMutableTreeNode(os.path.basename(path))
                root_node.add(node)

        self.tree_model.setRoot(root_node)

    def update_view_for_image(self, updated_image):
        """
        Finds and updates a single image's row in the JTable instead of
        reloading the entire UI.
        """
        for i in range(self.image_table_model.getRowCount()):
            # Find the row corresponding to our image
            if self.image_table_model.getValueAt(i, 0) == updated_image.filename:
                # Update the values in the table model
                self.image_table_model.setValueAt("Yes" if updated_image.has_roi() else "No", i, 1)
                self.image_table_model.setValueAt(len(updated_image.rois), i, 2)
                self.image_table_model.setValueAt(updated_image.status, i, 3)
                
                # Refresh the ROI details table as well
                self.on_image_selection(None) # Pass a dummy event or refactor to take an index
                break

class ImageImportWorker(SwingWorker):
    """
    Handles the image import process on a background thread to keep the GUI responsive,
    and reports progress updates that can be displayed by a progress bar.
    """
    def __init__(self, parent_gui, project, selected_files):
        super(ImageImportWorker, self).__init__()

        self.parent_gui = parent_gui
        self.project = project
        self.selected_files = selected_files
        self.newly_added_count = 0
        self.skipped_files = []

    def doInBackground(self):
        """This is where the long-running work happens."""
        images_dir = self.project.paths['images']
        total_files = len(self.selected_files)

        for i, source_file in enumerate(self.selected_files):
            # Check if the user has clicked the "Cancel" button on the progress monitor
            if self.isCancelled():
                break

            # Update the note on the progress monitor to show the current file
            self.firePropertyChange("note", "", "Copying {}...".format(source_file.getName()))
            
            dest_file = File(images_dir, source_file.getName())

            if dest_file.exists():
                self.skipped_files.append(source_file.getName())
                continue # Skip existing files

            try:
                Files.copy(source_file.toPath(), dest_file.toPath(), StandardCopyOption.REPLACE_EXISTING)
                
                # Update the project data structure in memory
                new_image = ProjectImage(dest_file.getName(), self.project.root_dir)
                new_image.status = "Pending ROIs"
                self.project.images.append(new_image)
                self.newly_added_count += 1
            except Exception as e:
                # Proper error handling should be added here if needed
                IJ.log("Failed to import '{}': {}".format(source_file.getName(), e))

            # Report the percentage complete
            progress = int(100.0 * (i + 1) / total_files)
            self.super__setProgress(progress)
        
        return self.newly_added_count

    def done(self):
        """This runs on the GUI thread after doInBackground is finished."""
        try:
            # The get() method retrieves the result and also raises any exceptions
            # that occurred during the background task.
            count = self.get()
            
            if count > 0:
                self.parent_gui.status_label.setText("Successfully imported {} new images.".format(count))
                self.parent_gui.update_ui_for_project()
                self.parent_gui.set_unsaved_changes(True)
            
            if self.skipped_files:
                IJ.log("Skipped {} existing files.".format(len(self.skipped_files)))

        except Exception as e:
            error_msg = "An error occurred during import: {}".format(e)
            IJ.log(error_msg)
            JOptionPane.showMessageDialog(self.parent_gui.frame, error_msg, "Import Error", JOptionPane.ERROR_MESSAGE)

class ROIEditor(WindowAdapter):
    """ 
    Creates a JFrame with tools for creating, modifying, and managing ROIs 
    for a single image using a "live update" model.
    """
    def __init__(self, parent_gui, project, project_image):
        self.parent_gui = parent_gui
        self.project = project
        self.image_obj = project_image
        self.win = None
        self.unsaved_changes = False 
        self.updating_fields = False  # Flag to prevent recursive updates
        self.last_selected_index = -1  # Track the last selected ROI index

        # Open Image and create canvas and imagewindow to hold it
        self.imp = IJ.openImage(self.image_obj.full_path)
        if not self.imp:
            IJ.error("Failed to open image:", self.image_obj.full_path)
            return
        self.imp.show()
        self.win = self.imp.getWindow()

        # Open a local ROI manager instance
        self.rm = RoiManager(True) 
        self.rm.reset()

        if self.image_obj.has_roi():
            self.rm.runCommand("Open", self.image_obj.roi_path)
            self.rm.runCommand("Show All")

        # Build GUI
        self.base_title = "ROI Editor: " + self.image_obj.filename
        self.frame = JDialog(self.win, self.base_title, False)
        self.frame.setSize(350, 650)  # Made taller to accommodate new button
        self.frame.addWindowListener(self)
        self.frame.setLayout(BorderLayout(5,5))

        # ROI list - Initialize components first, then populate
        self.roi_list_model = DefaultListModel()
        self.roi_list = JList(self.roi_list_model)
        self.update_roi_list_from_manager()  # Now safe to call
        self.roi_list.setSelectionMode(ListSelectionModel.SINGLE_SELECTION)
        self.roi_list.addListSelectionListener(self._on_roi_select)

        list_pane = JScrollPane(self.roi_list)
        list_pane.setBorder(BorderFactory.createTitledBorder("ROIs"))

        # Edit Panel
        edit_panel = JPanel(GridLayout(0, 2, 5, 5))
        edit_panel.setBorder(BorderFactory.createTitledBorder("Edit Selected ROI"))
        self.roi_name_field = JTextField()
        self.bregma_field = JTextField()
        edit_panel.add(JLabel("ROI Name:"))
        edit_panel.add(self.roi_name_field)
        edit_panel.add(JLabel("Bregma Value:"))
        edit_panel.add(self.bregma_field)

        self.show_all_checkbox = JCheckBox("Show All ROIs", True)
        self.show_all_checkbox.addActionListener(self._toggle_show_all)
        edit_panel.add(self.show_all_checkbox)
        
        # Event listeners for live updates with debouncing
        class TextFieldUpdater(FocusAdapter, ActionListener):
            def __init__(self, editor):
                self.editor = editor
            def actionPerformed(self, e):
                self.editor._apply_text_field_changes()
            def focusLost(self, e):
                self.editor._apply_text_field_changes()
        
        updater = TextFieldUpdater(self)
        self.roi_name_field.addActionListener(updater)
        self.roi_name_field.addFocusListener(updater)
        self.bregma_field.addActionListener(updater)
        self.bregma_field.addFocusListener(updater)
        
        # Button panel
        button_panel = JPanel(GridLayout(0, 1, 10, 10))
        button_panel.setBorder(BorderFactory.createEmptyBorder(10, 10, 10, 10))

        create_button = JButton("Create New From Selection", actionPerformed=self._create_new_roi)
        delete_button = JButton("Delete Selected ROI", actionPerformed=self._delete_selected_roi)
        
        # NEW: Save Current ROI button
        self.save_current_button = JButton("Save Current ROI Changes", actionPerformed=self._save_current_roi)
        self.save_current_button.setEnabled(False)  # Disabled until ROI is selected
        
        self.ready_checkbox = JCheckBox("Mark as Ready for Quantification")
        is_ready = (self.image_obj.status == "Ready to Quantify")
        self.ready_checkbox.setSelected(is_ready)
        self.ready_checkbox.addActionListener(self._toggle_ready_status)
        
        save_button = JButton("Save All ROIs & Close", actionPerformed=self._save_and_close)

        button_panel.add(create_button)
        button_panel.add(delete_button)
        button_panel.add(self.save_current_button)  # NEW button added here
        button_panel.add(self.ready_checkbox)
        button_panel.add(save_button)

        # Main layout
        south_panel = JPanel(BorderLayout())
        south_panel.add(edit_panel, BorderLayout.NORTH)
        south_panel.add(button_panel, BorderLayout.CENTER)

        self.frame.add(list_pane, BorderLayout.CENTER)
        self.frame.add(south_panel, BorderLayout.SOUTH)

    def show(self):
        if not self.frame or not self.win:
            return
        img_win_x = self.win.getX()
        img_win_width = self.win.getWidth()
        img_win_y = self.win.getY()
        self.frame.setLocation(img_win_x + img_win_width, img_win_y)
        self.frame.setVisible(True)

    def update_roi_list_from_manager(self):
        """Syncs the JList with the IJ ROI manager."""
        current_selection = self.roi_list.getSelectedIndex()
        self.roi_list_model.clear()
        
        rois = self.rm.getRoisAsArray()
        for i, roi in enumerate(rois):
            roi_name = roi.getName() or "Untitled"
            display_text = "{}. {}".format(i + 1, roi_name)
            self.roi_list_model.addElement(display_text)
        
        # Restore selection more carefully
        if current_selection != -1 and current_selection < self.roi_list_model.getSize():
            SwingUtilities.invokeLater(lambda: self.roi_list.setSelectedIndex(current_selection))

    def _set_unsaved_changes(self, state):
        """Updates the UI to show if there are unsaved changes."""
        self.unsaved_changes = state
        title = self.base_title
        if state:
            title += " *"
        self.frame.setTitle(title)

    def _toggle_show_all(self, event):
        """Toggles visibility of all ROIs on the image."""
        try:
            if event.getSource().isSelected():
                self.rm.runCommand("Show All")
            else:
                self.rm.runCommand("Show None")
            
            # Force a refresh of the image display
            if self.imp:
                self.imp.updateAndDraw()
        except Exception as e:
            IJ.log("Error toggling ROI visibility: " + str(e))

    def _on_roi_select(self, event):
        """Populates text fields when a new ROI is selected."""
        if not event.getValueIsAdjusting() and not self.updating_fields:
            selected_index = self.roi_list.getSelectedIndex()
            
            # Before switching to a new ROI, save any changes to the previously selected one
            if (self.last_selected_index != -1 and 
                self.last_selected_index != selected_index and 
                self.last_selected_index < self.rm.getCount()):
                self._sync_current_roi_from_image()
            
            if selected_index != -1 and selected_index < self.rm.getCount():
                try:
                    self.updating_fields = True  # Prevent recursive updates
                    
                    # Load and display the selected ROI
                    self._refresh_roi_display(selected_index)
                    
                    # Update text fields with the selected ROI's data
                    selected_roi = self.rm.getRoi(selected_index)
                    if selected_roi:
                        self.roi_name_field.setText(selected_roi.getName() or "")
                        bregma_val = selected_roi.getProperty("comment") or ''
                        self.bregma_field.setText(str(bregma_val))
                    
                    # Enable the save current button
                    self.save_current_button.setEnabled(True)
                    self.last_selected_index = selected_index
                    
                except Exception as e:
                    IJ.log("Error in ROI selection: " + str(e))
                finally:
                    self.updating_fields = False
            else:
                # Clear fields if no valid selection
                if not self.updating_fields:
                    self.roi_name_field.setText("")
                    self.bregma_field.setText("")
                    self.save_current_button.setEnabled(False)
                    self.last_selected_index = -1

    def _refresh_roi_display(self, selected_index):
        """Properly loads and displays the selected ROI from the manager."""
        try:
            # Clear any existing selection on the image first
            self.imp.deleteRoi()
            
            # Get the ROI from the manager and set it on the image
            selected_roi = self.rm.getRoi(selected_index)
            if selected_roi:
                # Clone the ROI to avoid modifying the original in the manager
                display_roi = selected_roi.clone()
                self.imp.setRoi(display_roi)
            
            # Handle overlay display
            if self.show_all_checkbox.isSelected():
                self.rm.runCommand("Show All")
            else:
                self.rm.runCommand("Show None")
            
            # Force display update
            self.imp.updateAndDraw()
            
        except Exception as e:
            IJ.log("Error refreshing ROI display: " + str(e))

    def _sync_current_roi_from_image(self):
        """
        NEW METHOD: Captures any geometric changes made to the ROI on the image
        and updates the corresponding ROI in the manager.
        """
        if self.last_selected_index == -1 or self.last_selected_index >= self.rm.getCount():
            return
            
        try:
            # Get the current ROI from the image (which may have been modified by dragging)
            current_image_roi = self.imp.getRoi()
            
            if current_image_roi:
                # Get the original ROI from the manager
                original_manager_roi = self.rm.getRoi(self.last_selected_index)
                
                # Create a new ROI with the updated geometry but preserve the metadata
                updated_roi = current_image_roi.clone()
                if original_manager_roi:
                    # Preserve the name from the original
                    updated_roi.setName(original_manager_roi.getName())
                    
                    # Copy specific properties we know about (safer approach)
                    comment = original_manager_roi.getProperty("comment")
                    if comment is not None:
                        updated_roi.setProperty("comment", comment)
                
                # Replace the ROI in the manager
                self.rm.setRoi(updated_roi, self.last_selected_index)
                self._set_unsaved_changes(True)
                
        except Exception as e:
            IJ.log("Error syncing ROI from image: " + str(e))

    def _save_current_roi(self, event):
        """
        NEW METHOD: Explicitly saves changes to the currently selected ROI.
        """
        selected_index = self.roi_list.getSelectedIndex()
        if selected_index == -1:
            JOptionPane.showMessageDialog(self.frame, 
                "Please select an ROI from the list first.", 
                "No ROI Selected", 
                JOptionPane.WARNING_MESSAGE)
            return
            
        try:
            # First apply any text field changes
            self._apply_text_field_changes()
            
            # Then sync any geometric changes from the image
            self._sync_current_roi_from_image()
            
            # Update the list display to reflect any name changes
            self.update_roi_list_from_manager()
            
            # Re-select the same ROI to maintain selection
            if selected_index < self.roi_list_model.getSize():
                self.roi_list.setSelectedIndex(selected_index)
            
            # Show confirmation
            roi_name = self.rm.getRoi(selected_index).getName() or "Untitled"
                
        except Exception as e:
            IJ.log("Error saving current ROI: " + str(e))
            JOptionPane.showMessageDialog(self.frame, 
                "Error saving ROI: " + str(e), 
                "Save Error", 
                JOptionPane.ERROR_MESSAGE)

    def _apply_text_field_changes(self):
        """Applies text field changes with better error handling."""
        if self.updating_fields:  # Prevent recursive calls
            return
            
        selected_index = self.roi_list.getSelectedIndex()
        if selected_index == -1 or selected_index >= self.rm.getCount():
            return

        try:
            selected_roi = self.rm.getRoi(selected_index)
            if not selected_roi:
                return
                
            new_name = self.roi_name_field.getText().strip()
            new_bregma = self.bregma_field.getText().strip()
            
            # Get current values for comparison
            current_name = selected_roi.getName() or ""
            current_bregma = str(selected_roi.getProperty("comment") or "")
            
            # Check if anything actually changed
            if current_name == new_name and current_bregma == new_bregma:
                return
                
            # Validate bregma value if not empty
            if new_bregma:
                try:
                    float(new_bregma)  # Test if it's a valid number
                except ValueError:
                    JOptionPane.showMessageDialog(self.frame, 
                        "Bregma value must be a number or empty.", 
                        "Invalid Input", 
                        JOptionPane.WARNING_MESSAGE)
                    # Reset to original value
                    self.bregma_field.setText(current_bregma)
                    return
            
            # Apply the changes
            self.updating_fields = True
            
            # Update ROI properties
            selected_roi.setProperty("comment", new_bregma if new_bregma else None)
            if new_name != current_name:
                self.rm.rename(selected_index, new_name)
            
            self._set_unsaved_changes(True)
                
        except Exception as e:
            IJ.log("Error applying text field changes: " + str(e))
            JOptionPane.showMessageDialog(self.frame, 
                "Error updating ROI: " + str(e), 
                "Update Error", 
                JOptionPane.ERROR_MESSAGE)
        finally:
            self.updating_fields = False

    def _create_new_roi(self, event):
        """Creates a new ROI from the current image selection using the name from the text field."""
        current_roi = self.imp.getRoi()
        if not current_roi:
            JOptionPane.showMessageDialog(self.frame, 
                "Please draw a selection on the image first.", 
                "No Selection Found", 
                JOptionPane.WARNING_MESSAGE)
            return
        
        # Get the name from the ROI name text field
        new_name = self.roi_name_field.getText().strip()
        if not new_name:
            JOptionPane.showMessageDialog(self.frame, 
                "Please enter a name in the 'ROI Name' field first.", 
                "No Name Provided", 
                JOptionPane.WARNING_MESSAGE)
            return

        try:
            # Clone the ROI to avoid modifying the original
            roi_clone = current_roi.clone()
            roi_clone.setName(new_name)
            
            # Also set the bregma value if provided
            bregma_value = self.bregma_field.getText().strip()
            if bregma_value:
                try:
                    float(bregma_value)  # Validate it's a number
                    roi_clone.setProperty("comment", bregma_value)
                except ValueError:
                    JOptionPane.showMessageDialog(self.frame, 
                        "Invalid bregma value. Using empty value instead.", 
                        "Invalid Bregma", 
                        JOptionPane.WARNING_MESSAGE)
            
            # Add to ROI manager
            self.rm.addRoi(roi_clone)
            
            # Update UI
            self.update_roi_list_from_manager()
            new_index = self.rm.getCount() - 1
            self.roi_list.setSelectedIndex(new_index)
            self._set_unsaved_changes(True)
            
            # Clear the image selection
            self.imp.deleteRoi()
            
            # Clear the text fields for the next ROI
            self.roi_name_field.setText("")
            self.bregma_field.setText("")
            
        except Exception as e:
            IJ.log("Error creating new ROI: " + str(e))
            JOptionPane.showMessageDialog(self.frame, 
                "Error creating ROI: " + str(e), 
                "Creation Error", 
                JOptionPane.ERROR_MESSAGE)

    def _delete_selected_roi(self, event):
        """Deletes the selected ROI."""
        selected_index = self.roi_list.getSelectedIndex()
        if selected_index == -1:
            JOptionPane.showMessageDialog(self.frame, 
                "Please select an ROI from the list to delete.", 
                "No ROI Selected", 
                JOptionPane.WARNING_MESSAGE)
            return

        # Confirm deletion
        roi_name = self.rm.getRoi(selected_index).getName() or "Untitled"
        result = JOptionPane.showConfirmDialog(self.frame, 
            "Delete ROI '{}'?".format(roi_name), 
            "Confirm Deletion", 
            JOptionPane.YES_NO_OPTION)
        
        if result != JOptionPane.YES_OPTION:
            return

        try:
            # Delete from ROI manager
            self.rm.select(selected_index)
            self.rm.runCommand("Delete")
            
            # Reset tracking since we deleted the selected ROI
            self.last_selected_index = -1
            
            # Update UI
            self.update_roi_list_from_manager()
            self.roi_name_field.setText("")
            self.bregma_field.setText("")
            self.save_current_button.setEnabled(False)
            self._set_unsaved_changes(True)
            
            # Clear image selection
            self.imp.deleteRoi()
            
        except Exception as e:
            IJ.log("Error deleting ROI: " + str(e))
            JOptionPane.showMessageDialog(self.frame, 
                "Error deleting ROI: " + str(e), 
                "Deletion Error", 
                JOptionPane.ERROR_MESSAGE)

    def _toggle_ready_status(self, event):
        """Updates the image's status in memory when the checkbox is toggled."""
        if self.ready_checkbox.isSelected():
            self.image_obj.status = "Ready to Quantify"
        else:
            self.image_obj.status = "Pending ROIs"
        self._set_unsaved_changes(True)
        
    def _save_all_changes(self):
        """Save with better error handling and validation."""
        try:
            # Before saving, sync any changes to the currently selected ROI
            if self.last_selected_index != -1:
                self._sync_current_roi_from_image()
            
            # Validate all ROI names are not empty (duplicate names are allowed)
            rois_from_manager = self.rm.getRoisAsArray()
            
            for i, roi in enumerate(rois_from_manager):
                name = roi.getName()
                if not name or name.strip() == "":
                    JOptionPane.showMessageDialog(self.frame, 
                        "ROI #{} has no name. Please name all ROIs before saving.".format(i + 1), 
                        "Validation Error", 
                        JOptionPane.WARNING_MESSAGE)
                    return False
            
            # Update project data structure
            new_rois_list = []
            for roi in rois_from_manager:
                new_rois_list.append({
                    'roi_name': roi.getName(),
                    'bregma': roi.getProperty("comment") or 'N/A',
                })
            self.image_obj.rois = new_rois_list

            # Save to .zip file
            if not os.path.exists(os.path.dirname(self.image_obj.roi_path)):
                os.makedirs(os.path.dirname(self.image_obj.roi_path))
                
            self.rm.runCommand("Save", self.image_obj.roi_path)
            return True
            
        except Exception as e:
            IJ.log("Error saving ROIs: " + str(e))
            IJ.log(traceback.format_exc())
            JOptionPane.showMessageDialog(self.frame, 
                "Failed to save ROI data: " + str(e), 
                "Save Error", 
                JOptionPane.ERROR_MESSAGE)
            return False

    def _save_and_close(self, event=None):
        """Saves all changes and closes the editor."""
        # Validate current field contents before saving
        if not self.updating_fields:
            self._apply_text_field_changes()
        
        # Save ROIs to .zip and update in-memory object
        if not self._save_all_changes():
            return

        # Save project databases
        if not self.project.sync_project_db():
            JOptionPane.showMessageDialog(self.frame, 
                "Could not save project databases. See log for details.", 
                "Database Sync Failed", 
                JOptionPane.ERROR_MESSAGE)
            return

        # Update main GUI
        self.parent_gui.update_view_for_image(self.image_obj)
        self.parent_gui.set_unsaved_changes(True)

        self._set_unsaved_changes(False)
        self.cleanup()

    def cleanup(self):
        """Closes the image, ROI Manager, and disposes the frame."""
        try:
            if self.imp:
                self.imp.changes = False  # Prevent save dialog
                self.imp.close()
            if self.rm:
                self.rm.close()
            if self.frame:
                self.frame.dispose()
        except Exception as e:
            IJ.log("Error during cleanup: " + str(e))

    def windowClosing(self, event):
        """Handles the window 'X' button."""
        if self.unsaved_changes:
            title = "Unsaved ROI Changes"
            message = "You have unsaved changes. Would you like to save before closing?"
            result = JOptionPane.showConfirmDialog(self.frame, message, title, 
                JOptionPane.YES_NO_CANCEL_OPTION, JOptionPane.QUESTION_MESSAGE)

            if result == JOptionPane.YES_OPTION:
                self._save_and_close(None)
            elif result == JOptionPane.NO_OPTION:
                self.cleanup()
            # If CANCEL, do nothing
        else:
            self.cleanup()

class EditableROIsTableModel(AbstractTableModel):
    """ Helper class to creat custom table model that allows editing of ROI details table"""
    def __init__(self, project_image):
        self.image = project_image
        self.headers = ["ROI Name", "Bregma", "Status"]
        self.data = self.image.rois if self.image else []
        self.header_map = {'roi_name': 0, 'bregma': 1, 'status': 2}

    def getRowCount(self):
        return len(self.data)
    
    def getColumnCount(self):
        return len(self.headers)
    
    def getValueAt(self, rowIndex, columnIndex):
        key = self.headers[columnIndex].lower().replace(" ", "_")
        return self.data[rowIndex].get(key, "")
    
    def getColumnName(self, columnIndex):
        return self.headers[columnIndex]
    
    def isCellEditable(self, rowIndex, columnIndex):
        return True

    def setValueAt(self, aValue, rowIndex, columnIndex):
        key = self.headers[columnIndex].lower().replace(" ", "_")
        self.data[rowIndex][key] = aValue
        # Updates data in projectImage directly
        self.fireTableCellUpdated(rowIndex, columnIndex)


#==============================================
# Quantification dialog class
#==============================================

class QuantificationDialog(JDialog):
    """
    modal dialog to configure setting for a batch quantification process.
    Returns selected settings to be passed to the worker class.
    """
    def __init__(self, parent_frame, selected_images):
        super(QuantificationDialog, self).__init__(parent_frame, "Quantification Setting", True)

        self.selected_images = selected_images
        self.settings = None
        self.available_models = self._get_models()

        # Main panel
        main_panel = JPanel(BorderLayout(10,10))
        main_panel.setBorder(EmptyBorder(15,15,15,15))
        self.add(main_panel)

        # Info label
        info_text = "Ready to process {} selected images.".format(len(self.selected_images))
        info_label = JLabel(info_text)
        main_panel.add(info_label, BorderLayout.NORTH)

        # Settings panel
        settings_panel = JPanel(GridLayout(0,2,10,10))
        settings_panel.setBorder(BorderFactory.createTitledBorder("Processing Options"))

        # workflow selection
        workflows = ["cFosDAB+ Detection (Generic Model)", "cFosDAB+ Detection (region specific model)"]
        settings_panel.add(JLabel("Choose Your Quantification Type: "))
        self.workflow_combo = JComboBox(workflows)
        settings_panel.add(self.workflow_combo)

        # Verbose images or no
        settings_panel.add(JLabel("Display Options: "))
        self.show_images_checkbox = JCheckBox("Show images during processing", False)
        settings_panel.add(self.show_images_checkbox)

        main_panel.add(settings_panel, BorderLayout.CENTER)

        # Bottom button panel
        button_panel = JPanel(FlowLayout(FlowLayout.RIGHT))
        run_button = JButton("Run", actionPerformed=self._run_action)
        cancel_button = JButton("Cancel", actionPerformed=self._cancel_action)
        button_panel.add(run_button)
        button_panel.add(cancel_button)
        main_panel.add(button_panel, BorderLayout.SOUTH)

        self.pack()

    def _run_action(self, event):
        """ Gathers settings into dictionary and closes dialog """
        selected_workflow = self.workflow_combo.getSelectedItem()

        if selected_workflow == "cFosDAB+ Detection (Generic Model)": 
            self.settings = {
                'images': self.selected_images,
                'pixel_classifier': self.available_models['PIXEL_cFosDAB_TiffIO_Generic'],
                'object_classifier': self.available_models['OBJECT_cFosDAB_TiffIO_Generic'],  
                'show_images': self.show_images_checkbox.isSelected()
                }
        elif selected_workflow == "cFosDAB+ Detection (region specific model)":
            IJ.error("NOT IMPLEMENTED", "Havnent made this yet. use the generic model.")
        
        self.dispose()

    def _cancel_action(self,event):
        """ Leaves settings=None and closes dialog"""
        self.settings = None
        self.dispose()

    def show_dialog(self):
        """ Public method called by the GUI """
        self.setLocationRelativeTo(self.getParent())
        self.setVisible(True)
        return self.settings
    
    def _get_models(self):
        """
        Finds models in a dedicated folder inside Fiji's 'lib' directory.
        This works by locating the core ImageJ .jar file
        to determine the Fiji root directory, regardless of how
        the application was launched.
        """
        from java.net import URLDecoder
        from java.lang import System

        MODELS_FOLDER_NAME = "cell-quantifier-toolkit-models"
        models = {}
        
        try:
            class_loader = IJ.getClassLoader()
            if class_loader is None:
                raise IOError("Could not get ImageJ ClassLoader.")

            resource_url = class_loader.getResource("IJ_Props.txt")
            if resource_url is None:
                raise IOError("Could not find core resource 'IJ_Props.txt'. Is Fiji installed correctly?")

            url_str = URLDecoder.decode(resource_url.toString(), "UTF-8")
            path_part = url_str.split("!")[0].replace("jar:file:", "")

            if System.getProperty("os.name").lower().startswith("windows") and path_part.startswith("/"):
                path_part = path_part[1:]

            jar_file = File(path_part)
            fiji_root_file = jar_file.getParentFile().getParentFile()
            fiji_root = fiji_root_file.getAbsolutePath()
           
            models_dir = os.path.join(fiji_root, "lib", MODELS_FOLDER_NAME)

            if os.path.isdir(models_dir):
                for f in os.listdir(models_dir):
                    if f.lower().endswith('.ilp'):
                        display_name = os.path.splitext(f)[0]
                        full_path = os.path.join(models_dir, f)
                        models[display_name] = full_path
            else:
                IJ.log("Model directory not found. Please create it at: " + models_dir)

        except Exception as e:
            IJ.log("Error discovering models: " + str(e))
            IJ.log(traceback.format_exc())

        return models

class ProgressDialog(JDialog):
    """ A simple, non-modal dialog to display a progress bar. """
    def __init__(self, parent_frame, title, max_value):
        super(ProgressDialog, self).__init__(parent_frame, title, False)
        self.progress_bar = JProgressBar(0, max_value)
        self.progress_bar.setStringPainted(True)
        self.add(self.progress_bar)
        self.pack()
        self.setSize(400, 80)
        self.setLocationRelativeTo(parent_frame)

#==============================================
# Processor Classes
#==============================================

class QuantificationWorker(SwingWorker):
    """ Processor Classs facilitating image quantification on a background thread given settings from the dialog """
    def __init__(self, parent_gui, project, settings, progress_dialog):
        super(QuantificationWorker, self).__init__()
        self.parent_gui = parent_gui
        self.project = project
        self.settings = settings
        self.progress_dialog = progress_dialog
        self.all_results = []

    def doInBackground(self):
        """
        Processes each ROI individually after loading all ROIs from the zip file.
        Uses an index to create unique temporary filenames, preventing overwrites.
        """
        # --- Helper class for updating the progress bar on the GUI thread ---
        class UpdateProgressBarTask(Runnable):
            def __init__(self, dialog, value):
                self.dialog = dialog
                self.value = value
            def run(self):
                self.dialog.progress_bar.setValue(self.value)

        # --- Main processing logic ---
        images_to_process = self.settings['images']

        # Set status to "Processing" at the beginning
        for image_obj in images_to_process:
            image_obj.status = "Processing"
        
        # Immediately save and refresh the UI to show the "Processing" status
        self.project._sync_image_status_db()
        SwingUtilities.invokeLater(self.parent_gui.update_ui_for_project)
        
        # Calculate total number of individual ROIs for the progress bar
        total_rois_to_process = 0
        for img in images_to_process:
            if img.has_roi():
                rm_temp = RoiManager(True)
                rm_temp.open(img.roi_path)
                total_rois_to_process += rm_temp.getCount()
                rm_temp.close()

        if total_rois_to_process == 0: 
            return "No ROIs to process."
        roi_counter = 0

        for image_obj in images_to_process:
            try:    
                all_image_outlines = []
                if self.isCancelled(): 
                    break
                
                if not image_obj.has_roi(): 
                    continue

                imp_original = IJ.openImage(image_obj.full_path)
                if not imp_original:
                    IJ.log("ERROR: Failed to open original image: " + image_obj.full_path)
                    continue
                
                # 1. Load ALL ROIs from the .zip file ONCE per image.
                rm = RoiManager(True)
                rm.open(image_obj.roi_path)
                all_rois_for_image = rm.getRoisAsArray()
                rm.close()

                # 2. Loop through the loaded ROIs using enumerate to get a unique index 'i'
                for i, roi in enumerate(all_rois_for_image):
                    if self.isCancelled(): 
                        break
                    
                    temp_cropped_path = None
                    try:
                        # Read the bregma value directly from the ROI object's property
                        bregma_val_str = roi.getProperty("comment")
                        try:
                            bregma_val = float(bregma_val_str) if bregma_val_str else 0.0
                        except (ValueError, TypeError):
                            bregma_val = 0.0

                        # Get bounding box coordinates for offsetting results later
                        roi_x = roi.getBounds().x
                        roi_y = roi.getBounds().y

                        # Create a duplicate for cropping to preserve the original image
                        imp_cropped = imp_original.duplicate()
                        imp_cropped.setRoi(roi)
                        IJ.run(imp_cropped, "Crop", "")
                        
                        # 3. Add the unique index 'i' to the base_name to prevent file overwriting
                        base_name = "{}_{}_{}".format(os.path.splitext(image_obj.filename)[0], roi.getName(), i)
                        
                        temp_cropped_path = os.path.join(self.project.paths['temp'], base_name + "_cropped.tif")
                        prob_map_path = os.path.join(self.project.paths['probabilities'], base_name)
                        IJ.saveAs(imp_cropped, "Tiff", temp_cropped_path)

                        imp_cropped.show()

                        # Run external processing (e.g., ilastik)
                        result_imp = self._run_ilastik_classification(roi, temp_cropped_path, image_obj.filename, prob_map_path)

                        if not self.settings.get('show_images', True):
                            if imp_cropped and imp_cropped.isVisible():
                                imp_cropped.close()

                        # Analyze the results in Fiji
                        analysis = self._analyze_results(result_imp, roi, roi_x, roi_y)

                        if not self.settings.get('show_images', True):
                            if result_imp:
                                result_imp.changes = False
                                result_imp.close()

                        if analysis['outlines']:
                            all_image_outlines.extend(analysis['outlines'])

                        # Collect the result for this single ROI piece
                        single_roi_result = {
                            'filename': image_obj.filename,
                            'roi_name': roi.getName(),
                            'roi_area': roi.getStatistics().area,
                            'bregma_value': bregma_val,
                            'cell_count': analysis['count'],
                            'total_cell_area': analysis['total area']
                        }
                        self.all_results.append(single_roi_result)

                    except Exception as e:
                        IJ.log("ERROR processing ROI #{} ('{}') in '{}': {}".format(i, roi.getName(), image_obj.filename, e))
                        IJ.log(traceback.format_exc())
                        continue 

                    finally:
                        # Clean up temporary cropped file
                        if temp_cropped_path and os.path.exists(temp_cropped_path):
                            try:
                                os.remove(temp_cropped_path)
                            except Exception as ex:
                                IJ.log("Warning: Could not delete temporary file " + temp_cropped_path)

                        if not self.settings.get('show_images', True):
                            self._cleanup_stray_windows()
                        
                        # Update progress
                        roi_counter += 1
                        progress = int(100.0 * roi_counter / total_rois_to_process)
                        update_task = UpdateProgressBarTask(self.progress_dialog, progress)
                        SwingUtilities.invokeLater(update_task)
                
                # After processing all ROIs for an image, save the collected cell outlines
                if all_image_outlines:
                    outline_rm = RoiManager(True)
                    for outline_roi in all_image_outlines:
                        outline_rm.addRoi(outline_roi)
                    outline_rm.runCommand("Save", image_obj.outline_path)
                    outline_rm.close()
                    IJ.log("Saved {} cell outlines for {}.".format(len(all_image_outlines), image_obj.filename))

                # Close the original image window if it's not meant to be shown
                if not self.settings.get('show_images', True) and imp_original and imp_original.isVisible():
                    imp_original.close()

                image_obj.status = "Completed" # Mark for final update

            except Exception as e:
                IJ.log("ERROR processing '{}': {}".format(image_obj.filename, e))
                image_obj.status = "Failed" # Mark as failed
                continue # Move to the next image

            finally:
                IJ.run("Collect Garbage", "")
                System.gc()

                self._cleanup_stray_windows()  

        return "Quantification completed successfully for {} ROIs.".format(roi_counter)
                
    
    def _run_ilastik_classification(self, roi, temp_cropped_path, img_name, prob_map_path):
        """
        Runs the full Ilastik workflow, correctly resuming from intermediate steps
        and handling the 'show images' setting by keeping required images open but hidden.
        """
        pixel_imp = None  # Define here for access in finally block
        try:
            pixel_classifier = self.settings['pixel_classifier']
            object_classifier = self.settings['object_classifier']
    
            pixel_prob_path = prob_map_path + "_probabilities.tif"
            object_prob_path = prob_map_path + "_objects.tif"

            # Case 1: The final object classification file already exists.
            if os.path.exists(object_prob_path):
                IJ.log("Found existing object file, skipping Ilastik processing for: " + os.path.basename(object_prob_path))
                result_imp = IJ.openImage(object_prob_path)
                if self.settings.get('show_images', True):
                    result_imp.show()
                return result_imp

            # Case 2: The intermediate pixel probability file exists, but the final one does not.
            elif os.path.exists(pixel_prob_path):
                IJ.log("Found existing probability map, running Object Classification only for: " + os.path.basename(pixel_prob_path))
                # CRITICAL: Open the existing probability map, as the next step depends on it.
                pixel_imp = IJ.openImage(pixel_prob_path)
                if not self.settings.get('show_images', True):
                    pixel_imp.hide() # Keep it open but invisible

                # Run only the Object Classification step
                object_macro_cmd = 'run("Run Object Classification Prediction", "projectfilename=[{}] rawinputimage=[{}] inputproborsegimage=[{}] secondinputtype=Probabilities ");'.format(object_classifier, temp_cropped_path, pixel_prob_path)
                IJ.runMacro(object_macro_cmd)
                object_imp = IJ.getImage()
                if not object_imp or (pixel_imp and object_imp.getID() == pixel_imp.getID()):
                    raise Exception("Object classification did not produce a new result image.")
                
                IJ.saveAs(object_imp, "Tiff", object_prob_path)
                if not self.settings.get('show_images', True):
                    object_imp.hide()

                IJ.run("Collect Garbage", "")
                System.gc()

                return object_imp

            # Case 3: Neither file exists. Run the full workflow.
            else:
                # Run Pixel Classification
                pixel_macro_cmd = 'run("Run Pixel Classification Prediction", "projectfilename=[{}] inputimage=[{}] pixelclassificationtype=Probabilities");'.format(pixel_classifier, temp_cropped_path)
                IJ.runMacro(pixel_macro_cmd)
                pixel_imp = IJ.getImage()
                if not pixel_imp:
                    raise Exception("No probability map was generated by the Ilastik pixel classifier.")

                IJ.saveAs(pixel_imp, "Tiff", pixel_prob_path)

                # Keep the image open but hide it for the next step.
                if not self.settings.get('show_images', True):
                    pixel_imp.hide()

                IJ.run("Collect Garbage", "")
                System.gc()

                # Run Object Classification
                object_macro_cmd = 'run("Run Object Classification Prediction", "projectfilename=[{}] rawinputimage=[{}] inputproborsegimage=[{}] secondinputtype=Probabilities ");'.format(object_classifier, temp_cropped_path, pixel_prob_path)
                IJ.runMacro(object_macro_cmd)
                object_imp = IJ.getImage()
                if not object_imp or (pixel_imp and object_imp.getID() == pixel_imp.getID()):
                    raise Exception("Object classification did not produce a new result image.")
                
                IJ.saveAs(object_imp, "Tiff", object_prob_path)
                if self.settings.get('show_images', True):
                    object_imp.show()

                IJ.run("Collect Garbage", "")
                System.gc()
                return object_imp

        except Exception as e:
            IJ.log("Ilastik processing failed: " + str(e))
            raise e
        finally:
            # Final cleanup of any lingering intermediate windows
            if pixel_imp:
                pixel_imp.changes = False
                pixel_imp.close() 
                

    def _analyze_results(self, result_imp, roi, offset_x, offset_y):
        """
        Final processing and analysis of ilastik output in Fiji.
        This version includes a thresholding step to create the required
        binary image for the Watershed command, resolving the error.
        """
        # --- START: MANUAL MASKING AND BINARIZATION ---

        # 1. Create a perfect black-and-white mask from the user's ROI.
        width = result_imp.getWidth()
        height = result_imp.getHeight()
        mask_title = "mask_" + str(System.nanoTime())
        mask_imp = IJ.createImage(mask_title, "8-bit black", width, height, 1)
        
        roi_clone_for_masking = roi.clone()
        roi_clone_for_masking.setLocation(0, 0)
        mask_imp.setRoi(roi_clone_for_masking)
        IJ.run(mask_imp, "Fill", "slice")
        mask_imp.deleteRoi()

        # 2. Use the Image Calculator's "AND" operation to apply the ROI mask
        # to the original Ilastik label image.
        from ij.plugin import ImageCalculator
        ic = ImageCalculator()
        ic.run("AND", result_imp, mask_imp)

        mask_imp.changes = False
        mask_imp.close()

        # The Watershed command requires a binary input. We select all labeled
        # pixels (values 1 and up) and convert them to a single mask.
        IJ.setThreshold(result_imp, 1, 65535) 
        IJ.run(result_imp, "Convert to Mask", "") 

        # 4. Now, run Watershed on the proper binary image.
        IJ.run(result_imp, "Watershed", "")
        
        rm = RoiManager(True)
        rt = ResultsTable()

        # Configure and run the ParticleAnalyzer
        options = ParticleAnalyzer.SHOW_OUTLINES | ParticleAnalyzer.EXCLUDE_EDGE_PARTICLES
        measurements = Measurements.AREA
        pa = ParticleAnalyzer(options, measurements, rt, 20, float('inf'), 0.0, 1.0)
        pa.setRoiManager(rm)
        pa.analyze(result_imp)

        # Get stats safely from our local results table.
        count = rt.getCounter()
        total_area = 0
        if count > 0:
            area_col_index = rt.getColumnIndex("Area")
            if area_col_index != -1:
                area_col = rt.getColumn(area_col_index)
                if area_col is not None:
                    total_area = sum(area_col)

        # Get the particle outlines
        particle_outlines_relative = rm.getRoisAsArray()
        rm.reset()
        rm.close()
        
        result_imp.changes = False
        result_imp.close()

        if particle_outlines_relative is None:
            particle_outlines_relative = []

        # Translate outlines to the full image coordinates
        particle_outlines_absolute = []
        for outline in particle_outlines_relative:
            current_bounds = outline.getBounds()
            outline.setLocation(current_bounds.x + offset_x, current_bounds.y + offset_y)
            particle_outlines_absolute.append(outline)

        analysis = {
            'count': count,
            'total area': total_area,
            'outlines': particle_outlines_absolute
        }
        
        return analysis
    
    def _cleanup_stray_windows(self):
        """Aggressively find and close any stray temporary image windows."""
        # Get a list of all currently open image windows
        image_ids = WindowManager.getIDList()
        if not image_ids:
            return
        
        # Keywords found in the titles of temporary windows
        temp_keywords = ["_cropped", "_probabilities", "_objects", "mask of"]

        # Iterate over a copy of the list, as closing images can modify it
        for img_id in list(image_ids):
            img = WindowManager.getImage(img_id)
            if not img:
                continue
            
            title = img.getTitle().lower()
            
            # If the window title contains any of our keywords, close it
            if any(keyword in title for keyword in temp_keywords):
                img.changes = False  # Prevent "Save changes?" dialog
                img.close()
    
    def done(self):
        """ Runs on GUI thread after background work is finished. """
        try:
            if self.all_results:
                aggregated_results = {}
                # This dictionary will hold temporary sums for averaging
                bregma_data = {}

                for result in self.all_results:
                    key = (result['filename'], result['roi_name'])
                    if key not in aggregated_results:
                        aggregated_results[key] = result.copy()
                        # Initialize sum and count for averaging Bregma
                        bregma_data[key] = {'sum': result['bregma_value'], 'count': 1}
                    else:
                        # Sum the quantitative values
                        aggregated_results[key]['roi_area'] += result['roi_area']
                        aggregated_results[key]['cell_count'] += result['cell_count']
                        aggregated_results[key]['total_cell_area'] += result['total_cell_area']
                        # Add to sum and increment count for averaging
                        bregma_data[key]['sum'] += result['bregma_value']
                        bregma_data[key]['count'] += 1
                
                # Now, calculate the average Bregma for each group
                for key, data in aggregated_results.items():
                    bregma_sum = bregma_data[key]['sum']
                    bregma_count = bregma_data[key]['count']
                    # Calculate average and format to 3 decimal places, avoid division by zero
                    average_bregma = (bregma_sum / bregma_count) if bregma_count > 0 else 0
                    aggregated_results[key]['bregma_value'] = "{:.3f}".format(average_bregma)

                final_results_list = aggregated_results.values()
                
                # Now write the FINAL, aggregated list to the CSV
                results_db_path = self.project.paths['results_db']
                headers = ['filename', 'roi_name', 'roi_area', 'bregma_value', 'cell_count', 'total_cell_area']
                file_exists = os.path.isfile(results_db_path)
                with open(results_db_path, 'ab') as csvfile:
                    writer = csv.DictWriter(csvfile, fieldnames=headers)
                    if not file_exists or os.path.getsize(results_db_path) == 0:
                        writer.writeheader()
                    writer.writerows(final_results_list)
            
            # Show final status message
            final_message = self.get()
            JOptionPane.showMessageDialog(self.progress_dialog, final_message, "Status", JOptionPane.INFORMATION_MESSAGE)

        except Exception as e:
            # This will catch errors from the background thread
            IJ.log(traceback.format_exc())
            JOptionPane.showMessageDialog(self.progress_dialog, "An error occurred during processing:\n" + str(e), "Error", JOptionPane.ERROR_MESSAGE)
            for image in self.settings['images']:
                if image.status == "Processing":
                    image.status = "Failed"
        finally:
            self.progress_dialog.dispose()

            image_ids = WindowManager.getIDList()
            if image_ids:
                # Iterate over a copy of the list, as closing images modifies the original list.
                for img_id in list(image_ids):
                    img = WindowManager.getImage(img_id)
                    if img:
                        img.changes = False
                        img.close()

            # Save the final "Completed" or "Failed" statuses and refresh the UI
            self.project.sync_project_db()
            self.parent_gui.update_ui_for_project()

#==============================================
# Results Viewer class
#==============================================
#==============================================
# Results Viewer Dialog Class
#==============================================

class ResultsViewer(WindowAdapter):
    """
    A self-contained dialog for viewing an image with toggleable overlays
    for analysis ROIs and quantified cell outlines.
    """
    def __init__(self, parent_frame, project_image):
        self.image_obj = project_image
        self.imp = IJ.openImage(self.image_obj.full_path)
        if not self.imp:
            IJ.error("Failed to open image: " + self.image_obj.full_path)
            return
        self.imp.show()

        self.image_window = self.imp.getWindow()

        # Load both sets of ROIs into memory
        self.analysis_rois = self._load_rois_from_zip(self.image_obj.roi_path)
        self.outline_rois = self._load_rois_from_zip(self.image_obj.outline_path)

        # Build the control dialog
        self.dialog = JDialog(self.image_window, "Results Viewer: " + self.image_obj.filename, False)
        self.dialog.setSize(300, 150)
        self.dialog.addWindowListener(self)

        self.image_window.addWindowListener(ImageWindowListener(self.dialog))
        
        panel = JPanel(GridLayout(0, 1, 10, 10))
        panel.setBorder(EmptyBorder(10, 10, 10, 10))

        # Create checkboxes
        self.analysis_checkbox = JCheckBox("Show Analysis ROIs", True)
        self.outlines_checkbox = JCheckBox("Show Cell Outlines", True)

        # Enable checkboxes only if their corresponding ROIs were found
        self.analysis_checkbox.setEnabled(bool(self.analysis_rois))
        self.outlines_checkbox.setEnabled(bool(self.outline_rois))

        # Add a single action listener to both
        action_listener = self._update_overlay
        self.analysis_checkbox.addActionListener(action_listener)
        self.outlines_checkbox.addActionListener(action_listener)

        panel.add(self.analysis_checkbox)
        panel.add(self.outlines_checkbox)
        self.dialog.add(panel)

        # Initial display
        self._update_overlay()

    def _load_rois_from_zip(self, zip_path):
        """Helper function to load all ROIs from a zip file into a list."""
        if not os.path.exists(zip_path):
            return []
        rm = RoiManager(True)
        rm.open(zip_path)
        rois = rm.getRoisAsArray()
        rm.close()
        return list(rois)

    def _update_overlay(self, event=None):
        """Builds and applies a new overlay based on checkbox states."""
        overlay = Overlay()

        if self.analysis_checkbox.isSelected() and self.analysis_rois:
            for roi in self.analysis_rois:
                overlay.add(roi)
        
        if self.outlines_checkbox.isSelected() and self.outline_rois:
            for roi in self.outline_rois:
                overlay.add(roi)
        
        self.imp.setOverlay(overlay)
        self.imp.updateAndDraw()

    def show(self):
        """Positions and shows the dialog."""
        if not self.dialog: return
        # Position control dialog next to the image window
        self.dialog.setLocation(self.imp.getWindow().getX() + self.imp.getWindow().getWidth(), self.imp.getWindow().getY())
        self.dialog.setVisible(True)

    def windowClosing(self, event):
        """Cleans up when the dialog is closed."""
        if self.imp:
            self.imp.close()

class ImageWindowListener(WindowAdapter):
    """A listener that closes the control dialog when its image window is closed."""
    def __init__(self, viewer_dialog):
        self.viewer_dialog = viewer_dialog

    def windowClosing(self, event):
        # When the image window is closed by the user,
        # programmatically close and dispose of our control dialog.
        if self.viewer_dialog:
            self.viewer_dialog.dispose()

#==============================================
# Program entry point
#==============================================
if __name__ == '__main__':
    def create_and_show_gui():
        gui = ProjectManagerGUI()
        gui.show()

    SwingUtilities.invokeLater(create_and_show_gui)
