from PyQt5.QtWidgets import (
    QLabel,
    QRadioButton,
    QApplication,
    QPushButton,
    QVBoxLayout,
    QHBoxLayout,
    QShortcut,
    QLineEdit,
    QAction,
    QWidget,
)
from celldetective.gui.interactive_timeseries_viewer import InteractiveEventViewer
from PyQt5.QtCore import Qt, QSize, QThread, pyqtSignal
from PyQt5.QtGui import QKeySequence, QCloseEvent

from superqt import (
    QLabeledDoubleRangeSlider,
    QSearchableComboBox,
    QLabeledSlider,
)
from celldetective.utils.image_loaders import (
    load_frames,
    _get_img_num_per_channel,
)
from celldetective.utils.experiment import (
    get_experiment_metadata,
    get_experiment_labels,
)
from celldetective.gui.gui_utils import (
    color_from_status,
    color_from_class,
)
from celldetective.gui.base.figure_canvas import FigureCanvas
import numpy as np
from superqt.fonticon import icon
from fonticon_mdi6 import MDI6
import matplotlib.pyplot as plt
from tqdm import tqdm
import gc
from matplotlib.animation import FuncAnimation
from matplotlib.cm import tab10
from typing import Optional, Tuple, Any
from celldetective.gui.base_annotator import BaseAnnotator


class StackLoaderThread(QThread):
    progress = pyqtSignal(int)
    status_update = pyqtSignal(str)
    finished = pyqtSignal()

    def __init__(self, annotator: QWidget) -> None:
        """
        Initialize the loader thread.

        Parameters
        ----------
        annotator : QWidget
            The annotator instance.
        """
        super().__init__()
        self.annotator = annotator
        self._is_cancelled = False

    def stop(self):
        """Stop the thread."""
        self._is_cancelled = True

    def run(self):
        """Run the thread."""

        def callback(progress: int, status: str = "") -> bool:
            """
            Callback to update progress.

            Parameters
            ----------
            progress : int
                Progress percentage.
            status : str, optional
                Status message. Default is "".

            Returns
            -------
            bool
                True if should continue, False if cancelled.
            """
            if self._is_cancelled:
                return False
            self.progress.emit(progress)
            if status:
                self.status_update.emit(status)
            return True

        try:
            self.annotator.prepare_stack(progress_callback=callback)
            if not self._is_cancelled:
                self.finished.emit()
        except Exception as e:
            print(f"Error in loader thread: {e}")
            self.finished.emit()


class EventAnnotator(BaseAnnotator):
    """
    UI to set tracking parameters for bTrack.

    Parameters
    ----------
    *args
        Variable length argument list.
    lazy_load : bool, optional
        If True, do not start loading the stack immediately. Default is False.
    **kwargs
        Arbitrary keyword arguments.
    """

    def __init__(self, *args, lazy_load=False, **kwargs):
        """
        Initialize the annotator.

        Parameters
        ----------
        *args
            Variable length argument list.
        lazy_load : bool, optional
            Whether to load images lazily. Default is False.
        **kwargs
            Arbitrary keyword arguments.
        """
        super().__init__(*args, **kwargs)
        self.setWindowTitle("Signal annotator")

        # default params
        self.class_name = "class"
        self.time_name = "t0"
        self.status_name = "status"
        self._loader_thread = None

        # self.locate_stack()
        if not self.proceed:
            self.close()
        else:
            if not lazy_load:
                self._start_threaded_loading()

    def _start_threaded_loading(self):
        """Start loading the stack in a background thread with progress dialog."""
        from celldetective.gui.base.components import CelldetectiveProgressDialog

        self._progress_dialog = CelldetectiveProgressDialog(
            "Loading stack...", "Cancel", 0, 100, self
        )
        self._progress_dialog.setWindowTitle("Loading")
        self._progress_dialog.setMinimumDuration(0)
        self._progress_dialog.setValue(0)

        self._loader_thread = StackLoaderThread(self)
        self._loader_thread.progress.connect(self._on_load_progress)
        self._loader_thread.status_update.connect(self._on_load_status)
        self._loader_thread.finished.connect(self._on_load_finished)
        self._progress_dialog.canceled.connect(self._on_load_canceled)

        self._loader_thread.start()

    def _on_load_progress(self, value: int) -> None:
        """
        Update progress dialog.

        Parameters
        ----------
        value : int
            Progress value (0-100).
        """
        if hasattr(self, "_progress_dialog") and self._progress_dialog:
            self._progress_dialog.setValue(value)

    def _on_load_status(self, status: str) -> None:
        """
        Update progress dialog label.

        Parameters
        ----------
        status : str
            Status message to display.
        """
        if hasattr(self, "_progress_dialog") and self._progress_dialog:
            self._progress_dialog.setLabelText(status)

    def _on_load_canceled(self):
        """Handle cancel button click."""
        if self._loader_thread:
            self._loader_thread.stop()
            self._loader_thread.wait()
        self.close()

    def _on_load_finished(self):
        """Called when loading completes."""
        if hasattr(self, "_progress_dialog") and self._progress_dialog:
            self._progress_dialog.close()
            self._progress_dialog = None
        self._loader_thread = None
        self.finalize_init()

    def finalize_init(self):
        """Finalize initialization after loading stack."""
        self.frame_lbl = QLabel("frame: ")
        self.looped_animation()
        self.init_event_buttons()
        self.populate_window()

        self.outliers_check.hide()
        if hasattr(self, "contrast_slider"):
            self.im.set_clim(
                self.contrast_slider.value()[0], self.contrast_slider.value()[1]
            )

    def init_event_buttons(self):
        """Initialize event buttons."""

        self.event_btn = QRadioButton("event")
        self.event_btn.setStyleSheet(self.button_style_sheet_2)
        self.event_btn.toggled.connect(self.enable_time_of_interest)

        self.no_event_btn = QRadioButton("no event")
        self.no_event_btn.setStyleSheet(self.button_style_sheet_2)
        self.no_event_btn.toggled.connect(self.enable_time_of_interest)

        self.else_btn = QRadioButton("else")
        self.else_btn.setStyleSheet(self.button_style_sheet_2)
        self.else_btn.toggled.connect(self.enable_time_of_interest)

        self.suppr_btn = QRadioButton("remove")
        self.suppr_btn.setToolTip(
            "Mark for deletion. Upon saving, the cell\nwill be removed from the tables."
        )
        self.suppr_btn.setStyleSheet(self.button_style_sheet_2)
        self.suppr_btn.toggled.connect(self.enable_time_of_interest)

        self.time_of_interest_label = QLabel("time of interest: ")
        self.time_of_interest_le = QLineEdit()

    def populate_options_layout(self):
        """Populate the options layout."""

        # clear options hbox
        for i in reversed(range(self.options_hbox.count())):
            self.options_hbox.itemAt(i).widget().setParent(None)

        options_layout = QVBoxLayout()
        # add new widgets

        btn_hbox = QHBoxLayout()
        btn_hbox.addWidget(self.event_btn, 25, alignment=Qt.AlignCenter)
        btn_hbox.addWidget(self.no_event_btn, 25, alignment=Qt.AlignCenter)
        btn_hbox.addWidget(self.else_btn, 25, alignment=Qt.AlignCenter)
        btn_hbox.addWidget(self.suppr_btn, 25, alignment=Qt.AlignCenter)

        time_option_hbox = QHBoxLayout()
        time_option_hbox.setContentsMargins(0, 5, 100, 10)
        time_option_hbox.addWidget(self.time_of_interest_label, 10)
        time_option_hbox.addWidget(self.time_of_interest_le, 15)
        time_option_hbox.addWidget(QLabel(""), 75)

        options_layout.addLayout(btn_hbox)
        options_layout.addLayout(time_option_hbox)

        self.options_hbox.addLayout(options_layout)

        self.annotation_btns_to_hide = [
            self.event_btn,
            self.no_event_btn,
            self.else_btn,
            self.time_of_interest_label,
            self.time_of_interest_le,
            self.suppr_btn,
        ]
        self.hide_annotation_buttons()

    def populate_window(self):
        """
        Create the multibox design.

        """

        super().populate_window()
        self.populate_options_layout()

        self.del_shortcut = QShortcut(Qt.Key_Delete, self)  # QKeySequence("s")
        self.del_shortcut.activated.connect(self.shortcut_suppr)
        self.del_shortcut.setEnabled(False)

        self.no_event_shortcut = QShortcut(QKeySequence("n"), self)  # QKeySequence("s")
        self.no_event_shortcut.activated.connect(self.shortcut_no_event)
        self.no_event_shortcut.setEnabled(False)

        # Right side
        animation_buttons_box = QHBoxLayout()
        animation_buttons_box.addWidget(self.frame_lbl, 20, alignment=Qt.AlignLeft)

        self.first_frame_btn = QPushButton()
        self.first_frame_btn.clicked.connect(self.set_first_frame)
        self.first_short = QShortcut(QKeySequence("f"), self)
        self.first_short.activated.connect(self.set_first_frame)
        self.first_frame_btn.setIcon(icon(MDI6.page_first, color="black"))
        self.first_frame_btn.setStyleSheet(self.button_select_all)
        self.first_frame_btn.setFixedSize(QSize(60, 60))
        self.first_frame_btn.setIconSize(QSize(30, 30))

        self.last_frame_btn = QPushButton()
        self.last_frame_btn.clicked.connect(self.set_last_frame)
        self.last_short = QShortcut(QKeySequence("l"), self)
        self.last_short.activated.connect(self.set_last_frame)
        self.last_frame_btn.setIcon(icon(MDI6.page_last, color="black"))
        self.last_frame_btn.setStyleSheet(self.button_select_all)
        self.last_frame_btn.setFixedSize(QSize(60, 60))
        self.last_frame_btn.setIconSize(QSize(30, 30))

        self.stop_btn = QPushButton()
        self.stop_btn.clicked.connect(self.stop)
        self.stop_btn.setIcon(icon(MDI6.stop, color="black"))
        self.stop_btn.setStyleSheet(self.button_select_all)
        self.stop_btn.setFixedSize(QSize(60, 60))
        self.stop_btn.setIconSize(QSize(30, 30))

        self.start_btn = QPushButton()
        self.start_btn.clicked.connect(self.start)
        self.start_btn.setIcon(icon(MDI6.play, color="black"))
        self.start_btn.setFixedSize(QSize(60, 60))
        self.start_btn.setStyleSheet(self.button_select_all)
        self.start_btn.setIconSize(QSize(30, 30))
        self.start_btn.hide()

        self.toggle_short = QShortcut(Qt.Key_Space, self)
        self.toggle_short.activated.connect(self.toggle_animation)

        self.speed_slider = QLabeledSlider(Qt.Horizontal)
        self.speed_slider.setRange(1, 60)
        # Convert initial interval (ms) to FPS
        initial_fps = int(1000 / max(1, self.anim_interval))
        self.speed_slider.setValue(initial_fps)
        self.speed_slider.valueChanged.connect(self.update_speed)
        self.speed_slider.setFixedWidth(200)
        self.speed_slider.setToolTip("Adjust animation framerate (FPS)")

        speed_layout = QHBoxLayout()
        speed_layout.addWidget(QLabel("Framerate: "))
        speed_layout.addWidget(self.speed_slider)

        animation_buttons_box.addLayout(speed_layout, 20)

        animation_buttons_box.addWidget(
            self.first_frame_btn, 5, alignment=Qt.AlignRight
        )

        self.prev_frame_btn = QPushButton()
        self.prev_frame_btn.clicked.connect(self.prev_frame)
        self.prev_frame_btn.setIcon(icon(MDI6.chevron_left, color="black"))
        self.prev_frame_btn.setStyleSheet(self.button_select_all)
        self.prev_frame_btn.setFixedSize(QSize(60, 60))
        self.prev_frame_btn.setIconSize(QSize(30, 30))
        self.prev_frame_btn.setEnabled(
            False
        )  # Disabled by default (assuming auto-start or initial state)
        # Actually start is hidden initially, and stop shown?
        # In populate_window: start_btn.hide() (line 235), stop_btn defaults?
        # stop_btn is created.
        # So assumed state is Playing? No, usually it starts, looped_animation matches.
        # But if it starts playing, buttons should be disabled.
        animation_buttons_box.addWidget(self.prev_frame_btn, 5, alignment=Qt.AlignRight)

        animation_buttons_box.addWidget(self.stop_btn, 5, alignment=Qt.AlignRight)
        animation_buttons_box.addWidget(self.start_btn, 5, alignment=Qt.AlignRight)

        self.next_frame_btn = QPushButton()
        self.next_frame_btn.clicked.connect(self.next_frame)
        self.next_frame_btn.setIcon(icon(MDI6.chevron_right, color="black"))
        self.next_frame_btn.setStyleSheet(self.button_select_all)
        self.next_frame_btn.setFixedSize(QSize(60, 60))
        self.next_frame_btn.setIconSize(QSize(30, 30))
        self.next_frame_btn.setEnabled(False)
        animation_buttons_box.addWidget(self.next_frame_btn, 5, alignment=Qt.AlignRight)

        animation_buttons_box.addWidget(self.last_frame_btn, 5, alignment=Qt.AlignRight)

        self.right_panel.addLayout(animation_buttons_box, 5)
        self.right_panel.addWidget(self.fcanvas, 90)

        if not self.rgb_mode:
            contrast_hbox = QHBoxLayout()
            contrast_hbox.setContentsMargins(150, 5, 150, 5)
            self.contrast_slider = QLabeledDoubleRangeSlider()
            self.contrast_slider.setSingleStep(0.001)
            self.contrast_slider.setTickInterval(0.001)
            self.contrast_slider.setOrientation(Qt.Horizontal)
            # Cache percentile values to avoid recomputing on the full stack
            self._stack_p_low = np.nanpercentile(self.stack, 0.001)
            self._stack_p_high = np.nanpercentile(self.stack, 99.999)
            self._stack_p1 = np.nanpercentile(self.stack, 1)
            self._stack_p99 = np.nanpercentile(self.stack, 99.99)
            self.contrast_slider.setRange(self._stack_p_low, self._stack_p_high)
            self.contrast_slider.setValue([self._stack_p1, self._stack_p99])
            self.contrast_slider.valueChanged.connect(self.contrast_slider_action)
            contrast_hbox.addWidget(QLabel("contrast: "))
            contrast_hbox.addWidget(self.contrast_slider, 90)
            self.right_panel.addLayout(contrast_hbox, 5)

        if self.class_choice_cb.currentText() != "":
            self.compute_status_and_colors(0)

        QApplication.processEvents()

        # Add Menu for Interactive Plotter
        menubar = self.menuBar()
        viewMenu = menubar.addMenu("View")

        openPlotterAct = QAction("Interactive Plotter", self)
        openPlotterAct.setShortcut("Ctrl+P")
        openPlotterAct.setStatusTip("Open interactive signal plotter for corrections")
        openPlotterAct.triggered.connect(self.launch_interactive_viewer)
        viewMenu.addAction(openPlotterAct)

    def launch_interactive_viewer(self):
        """Launch the interactive signal viewer."""
        if (
            not hasattr(self, "plotter")
            or self.plotter is None
            or not self.plotter.isVisible()
        ):
            label = None
            if hasattr(self, "class_name") and self.class_name.startswith("class_"):
                label = self.class_name.replace("class_", "")

            # Create with shared DF and callback
            self.plotter = InteractiveEventViewer(
                self.trajectories_path,
                df=self.df_tracks,
                event_label=label,
                callback=self.on_viewer_update,
                parent=self,
            )
        self.plotter.show()
        self.plotter.activateWindow()

    def on_viewer_update(self):
        """Callback from interactive viewer to refresh annotator."""
        self.compute_status_and_colors(0)
        self.update_scatters_only()
        self.fcanvas.canvas.draw_idle()

    def update_scatters_only(self):
        """Update only the scatters without reloading the image."""
        self.status_scatter.set_offsets(self.positions[self.framedata])
        self.status_scatter.set_color(self.colors[self.framedata][:, 1])
        self.class_scatter.set_offsets(self.positions[self.framedata])
        self.class_scatter.set_edgecolor(self.colors[self.framedata][:, 0])

    def compute_status_and_colors(self, i: int) -> None:
        """
        Compute the status and colors of the cells.

        Parameters
        ----------
        i : int
            Current frame index.
        """
        self.class_name = self.class_choice_cb.currentText()
        if self.class_name == "":
            self.class_name = "class"

        suffix = self.class_name.replace("class", "").replace("_", "", 1)
        if suffix != "":
            self.expected_status = "status_" + suffix
            self.expected_time = "t_" + suffix
        else:
            self.expected_status = "status"
            self.expected_time = "t0"
        self.expected_class = self.class_name

        self.time_name = self.expected_time
        self.status_name = self.expected_status

        cols = list(self.df_tracks.columns)

        if (
            self.time_name in cols
            and self.class_name in cols
            and not self.status_name in cols
        ):
            # only create the status column if it does not exist to not erase static classification results
            self.make_status_column()
        elif (
            self.time_name in cols
            and self.class_name in cols
            and self.df_tracks[self.status_name].isnull().all()
        ):
            self.make_status_column()
        elif self.time_name in cols and self.class_name in cols:
            # all good, do nothing
            pass
        else:
            if not self.status_name in self.df_tracks.columns:
                self.df_tracks[self.status_name] = 0
                self.df_tracks["status_color"] = color_from_status(0)
                self.df_tracks["class_color"] = color_from_class(1)

        if not self.class_name in self.df_tracks.columns:
            self.df_tracks[self.class_name] = 1
        if not self.time_name in self.df_tracks.columns:
            self.df_tracks[self.time_name] = -1

        self.df_tracks["status_color"] = [
            color_from_status(i) for i in self.df_tracks[self.status_name].to_numpy()
        ]
        self.df_tracks["class_color"] = [
            color_from_class(i) for i in self.df_tracks[self.class_name].to_numpy()
        ]

        self.extract_scatter_from_trajectories()
        if len(self.selection) > 0:
            self.select_single_cell(self.selection[0][0], self.selection[0][1])

        self.fcanvas.canvas.draw()

    # def close_without_new_class(self):
    # 	self.newClassWidget.close()

    def cancel_selection(self):
        """Cancel the current selection and revert color changes."""
        super().cancel_selection()
        try:
            for k, (t, idx) in enumerate(zip(self.loc_t, self.loc_idx)):
                self.colors[t][idx, 1] = self.previous_color[k][1]
        except Exception as e:
            pass

    def hide_annotation_buttons(self):
        """Hide annotation buttons."""

        for a in self.annotation_btns_to_hide:
            a.hide()
        for b in [self.event_btn, self.no_event_btn, self.else_btn, self.suppr_btn]:
            b.setChecked(False)
        self.time_of_interest_label.setEnabled(False)
        self.time_of_interest_le.setText("")
        self.time_of_interest_le.setEnabled(False)

    def enable_time_of_interest(self):
        """Enable time of interest input if event is selected."""

        if self.event_btn.isChecked():
            self.time_of_interest_label.setEnabled(True)
            self.time_of_interest_le.setEnabled(True)
        else:
            self.time_of_interest_label.setEnabled(False)
            self.time_of_interest_le.setEnabled(False)

    def show_annotation_buttons(self):
        """Show annotation buttons."""

        for a in self.annotation_btns_to_hide:
            a.show()

        cclass = self.df_tracks.loc[
            self.df_tracks["TRACK_ID"] == self.track_of_interest, self.class_name
        ].to_numpy()[0]
        t0 = self.df_tracks.loc[
            self.df_tracks["TRACK_ID"] == self.track_of_interest, self.time_name
        ].to_numpy()[0]

        if cclass == 0:
            self.event_btn.setChecked(True)
            self.time_of_interest_le.setText(str(t0))
        elif cclass == 1:
            self.no_event_btn.setChecked(True)
        elif cclass == 2:
            self.else_btn.setChecked(True)
        elif cclass > 2:
            self.suppr_btn.setChecked(True)

        self.enable_time_of_interest()
        self.correct_btn.setText("submit")

        try:
            self.correct_btn.clicked.disconnect()
        except TypeError:
            pass  # No connections to disconnect
        self.correct_btn.clicked.connect(self.apply_modification)

    def apply_modification(self):
        """Apply modification to the event."""

        t0 = -1
        if self.event_btn.isChecked():
            cclass = 0
            try:
                t0 = float(self.time_of_interest_le.text().replace(",", "."))
                self.line_dt.set_xdata([t0, t0])
                self.cell_fcanvas.canvas.draw_idle()
            except ValueError:
                # Invalid time value entered
                t0 = -1
                cclass = 2
        elif self.no_event_btn.isChecked():
            cclass = 1
        elif self.else_btn.isChecked():
            cclass = 2
        elif self.suppr_btn.isChecked():
            cclass = 42

        self.df_tracks.loc[
            self.df_tracks["TRACK_ID"] == self.track_of_interest, self.class_name
        ] = cclass
        self.df_tracks.loc[
            self.df_tracks["TRACK_ID"] == self.track_of_interest, self.time_name
        ] = t0

        indices = self.df_tracks.loc[
            self.df_tracks["TRACK_ID"] == self.track_of_interest, self.class_name
        ].index
        timeline = self.df_tracks.loc[
            self.df_tracks["TRACK_ID"] == self.track_of_interest, "FRAME"
        ].to_numpy()
        status = np.zeros_like(timeline)
        if t0 > 0:
            status[timeline >= t0] = 1.0
        if cclass == 2:
            status[:] = 2
        if cclass > 2:
            status[:] = 42

        status_color = [color_from_status(s, recently_modified=True) for s in status]
        class_color = [
            color_from_class(cclass, recently_modified=True) for i in range(len(status))
        ]

        # self.df_tracks['status_color'] = [color_from_status(i) for i in self.df_tracks[self.status_name].to_numpy()]
        # self.df_tracks['class_color'] = [color_from_class(i) for i in self.df_tracks[self.class_name].to_numpy()]

        self.df_tracks.loc[indices, self.status_name] = status
        self.df_tracks.loc[indices, "status_color"] = status_color
        self.df_tracks.loc[indices, "class_color"] = class_color

        # self.make_status_column()
        self.extract_scatter_from_trajectories()
        self.give_cell_information()

        try:
            self.correct_btn.clicked.disconnect()
        except TypeError:
            pass  # No connections to disconnect
        self.correct_btn.clicked.connect(self.show_annotation_buttons)

        self.hide_annotation_buttons()
        self.correct_btn.setEnabled(False)
        self.correct_btn.setText("correct")
        self.cancel_btn.setEnabled(False)
        self.del_shortcut.setEnabled(False)
        self.no_event_shortcut.setEnabled(False)

        self.selection.pop(0)

    def make_status_column(self):
        """Create the status column based on class and time."""

        print(
            f"Generating status information for class `{self.class_name}` and time `{self.time_name}`..."
        )
        for tid, group in self.df_tracks.groupby("TRACK_ID"):

            indices = group.index
            t0 = group[self.time_name].to_numpy()[0]
            cclass = group[self.class_name].to_numpy()[0]
            timeline = group["FRAME"].to_numpy()
            status = np.zeros_like(timeline)

            if t0 > 0:
                status[timeline >= t0] = 1.0
            # if cclass == 2:
            # 	status[:] = 1.
            if cclass > 2:
                status[:] = 42

            status_color = [color_from_status(s) for s in status]
            class_color = [color_from_class(cclass) for i in range(len(status))]

            self.df_tracks.loc[indices, self.status_name] = status
            self.df_tracks.loc[indices, "status_color"] = status_color
            self.df_tracks.loc[indices, "class_color"] = class_color

    def generate_signal_choices(self):
        """Generate signal choice combos."""

        self.signal_choice_cb = [QSearchableComboBox() for i in range(self.n_signals)]
        self.signal_choice_label = [
            QLabel(f"signal {i + 1}: ") for i in range(self.n_signals)
        ]
        # self.log_btns = [QPushButton() for i in range(self.n_signals)]

        signals = list(self.df_tracks.columns)

        to_remove = [
            "TRACK_ID",
            "FRAME",
            "x_anim",
            "y_anim",
            "t",
            "state",
            "generation",
            "root",
            "parent",
            "class_id",
            "class",
            "t0",
            "POSITION_X",
            "POSITION_Y",
            "position",
            "well",
            "well_index",
            "well_name",
            "pos_name",
            "index",
            "class_color",
            "status_color",
            "dummy",
            "group_color",
        ]

        meta = get_experiment_metadata(self.exp_dir)
        if meta is not None:
            keys = list(meta.keys())
            to_remove.extend(keys)

        labels = get_experiment_labels(self.exp_dir)
        if labels is not None:
            keys = list(labels.keys())
            to_remove.extend(labels)

        for c in to_remove:
            if c in signals:
                signals.remove(c)

        for i in range(len(self.signal_choice_cb)):
            self.signal_choice_cb[i].addItems(["--"] + signals)
            self.signal_choice_cb[i].setCurrentIndex(i + 1)
            self.signal_choice_cb[i].currentIndexChanged.connect(self.plot_signals)

    def plot_signals(self):
        """Plot the selected signals."""

        range_values = []

        try:
            yvalues = []
            for i in range(len(self.signal_choice_cb)):

                signal_choice = self.signal_choice_cb[i].currentText()
                lbl = signal_choice
                n_cut = 35
                if len(lbl) > n_cut:
                    lbl = lbl[: (n_cut - 3)] + "..."
                self.lines[i].set_label(lbl)

                if signal_choice == "--":
                    self.lines[i].set_xdata([])
                    self.lines[i].set_ydata([])
                else:
                    xdata = self.df_tracks.loc[
                        self.df_tracks["TRACK_ID"] == self.track_of_interest, "FRAME"
                    ].to_numpy()
                    ydata = self.df_tracks.loc[
                        self.df_tracks["TRACK_ID"] == self.track_of_interest,
                        signal_choice,
                    ].to_numpy()

                    range_values.extend(ydata)

                    xdata = xdata[ydata == ydata]  # remove nan
                    ydata = ydata[ydata == ydata]

                    yvalues.extend(ydata)
                    self.lines[i].set_xdata(xdata)
                    self.lines[i].set_ydata(ydata)
                    self.lines[i].set_color(tab10(i / 3.0))

            self.configure_ylims()

            min_val, max_val = self.cell_ax.get_ylim()
            t0 = self.df_tracks.loc[
                self.df_tracks["TRACK_ID"] == self.track_of_interest, self.expected_time
            ].to_numpy()[0]
            self.line_dt.set_xdata([t0, t0])
            self.line_dt.set_ydata([min_val, max_val])

            self.cell_ax.legend(fontsize=8)
            self.cell_fcanvas.canvas.draw()
        except Exception as e:
            print(e)
            pass

        if len(range_values) > 0:
            range_values = np.array(range_values)
            if len(range_values[range_values == range_values]) > 0:
                if len(range_values[range_values > 0]) > 0:
                    self.value_magnitude = np.nanpercentile(range_values, 1)
                else:
                    self.value_magnitude = 1
                self.non_log_ymin = 0.98 * np.nanmin(range_values)
                self.non_log_ymax = np.nanmax(range_values) * 1.02
                if self.cell_ax.get_yscale() == "linear":
                    self.cell_ax.set_ylim(self.non_log_ymin, self.non_log_ymax)
                else:
                    self.cell_ax.set_ylim(self.value_magnitude, self.non_log_ymax)

    def extract_scatter_from_trajectories(self):
        """Extract scatter data from trajectories using efficient groupby."""
        self.positions = []
        self.colors = []
        self.tracks = []

        # Pre-allocate with empty arrays for all frames
        for _ in range(self.len_movie):
            self.positions.append(np.empty((0, 2)))
            self.colors.append(np.empty((0, 2), dtype=object))
            self.tracks.append(np.empty(0))

        # Use groupby for efficient extraction
        grouped = self.df_tracks.groupby("FRAME")
        for frame, group in grouped:
            frame = int(frame)
            if 0 <= frame < self.len_movie:
                self.positions[frame] = group[["x_anim", "y_anim"]].to_numpy()
                self.colors[frame] = group[["class_color", "status_color"]].to_numpy()
                self.tracks[frame] = group["TRACK_ID"].to_numpy()

    def prepare_stack(self, progress_callback: Optional[callable] = None) -> None:
        """
        Load the stack images.

        Parameters
        ----------
        progress_callback : callable, optional
            Callback for progress updates.
        """

        self.img_num_channels = _get_img_num_per_channel(
            self.channels, self.len_movie, self.nbr_channels
        )
        self.stack = []
        disable_tqdm = not len(self.target_channels) > 1

        # Calculate total frames for progress
        total_frames = 0
        for ch in self.target_channels:
            target_ch_name = ch[0]
            indices = self.img_num_channels[
                self.channels[np.where(self.channel_names == target_ch_name)][0]
            ]
            total_frames += len(indices)

        current_frame = 0

        for ch in tqdm(self.target_channels, desc="channel", disable=disable_tqdm):
            target_ch_name = ch[0]
            if self.percentile_mode:
                normalize_kwargs = {"percentiles": (ch[1], ch[2]), "values": None}
            else:
                normalize_kwargs = {"values": (ch[1], ch[2]), "percentiles": None}

            if self.rgb_mode:
                normalize_kwargs.update({"amplification": 255.0, "clip": True})

            chan = []
            indices = self.img_num_channels[
                self.channels[np.where(self.channel_names == target_ch_name)][0]
            ]

            if progress_callback:
                if not progress_callback(
                    int((current_frame / total_frames) * 100),
                    f"Loading channel {target_ch_name}...",
                ):
                    return

            for t in tqdm(range(len(indices)), desc="frame"):

                if progress_callback:
                    if not progress_callback(int((current_frame / total_frames) * 100)):
                        return
                    current_frame += 1

                if self.rgb_mode:
                    f = load_frames(
                        indices[t],
                        self.stack_path,
                        scale=self.fraction,
                        normalize_input=True,
                        normalize_kwargs=normalize_kwargs,
                    )
                    f = f.astype(np.uint8)
                else:
                    f = load_frames(
                        indices[t],
                        self.stack_path,
                        scale=self.fraction,
                        normalize_input=False,
                    )

                chan.append(f[:, :, 0])

            self.stack.append(chan)

        self.stack = np.array(self.stack)
        if self.rgb_mode:
            self.stack = np.moveaxis(self.stack, 0, -1)
        else:
            self.stack = self.stack[0]
            if self.log_option:
                self.stack[np.where(self.stack > 0.0)] = np.log(
                    self.stack[np.where(self.stack > 0.0)]
                )

    def closeEvent(self, event: QCloseEvent) -> None:
        """
        Handle close event.

        Parameters
        ----------
        event : QCloseEvent
            The close event.
        """
        try:
            self.stop()
        except Exception:
            pass

        # Stop and delete animation to break reference cycles
        if hasattr(self, "anim") and self.anim:
            try:
                self.anim.event_source.stop()
            except Exception:
                pass
            del self.anim

        # Close matplotlib figures
        if hasattr(self, "fig"):
            try:
                plt.close(self.fig)
            except Exception:
                pass

        if hasattr(self, "cell_fig"):
            try:
                plt.close(self.cell_fig)
            except Exception:
                pass

        # Delete large objects
        if hasattr(self, "stack"):
            del self.stack

        if hasattr(self, "df_tracks"):
            del self.df_tracks

        gc.collect()
        super().closeEvent(event)

    def animation_generator(self):
        """
        Generator yielding frame indices for the animation,
        starting from the current self.framedata.
        """
        i = self.framedata
        while True:
            yield i
            i += 1
            if i >= self.len_movie:
                i = 0

    def looped_animation(self):
        """
        Load an image.

        """

        self.framedata = 0

        self.fig, self.ax = plt.subplots(tight_layout=True)
        self.fcanvas = FigureCanvas(self.fig, interactive=True)
        self.ax.clear()

        self.im = self.ax.imshow(self.stack[0], cmap="gray", interpolation="none")
        self.status_scatter = self.ax.scatter(
            self.positions[0][:, 0],
            self.positions[0][:, 1],
            marker="x",
            c=self.colors[0][:, 1],
            s=50,
            picker=True,
            pickradius=100,
        )
        self.class_scatter = self.ax.scatter(
            self.positions[0][:, 0],
            self.positions[0][:, 1],
            marker="o",
            facecolors="none",
            edgecolors=self.colors[0][:, 0],
            s=200,
        )

        self.ax.set_xticks([])
        self.ax.set_yticks([])
        self.ax.set_aspect("equal")

        self.fig.set_facecolor("none")  # or 'None'
        self.fig.canvas.setStyleSheet("background-color: black;")

        self.anim = FuncAnimation(
            self.fig,
            self.draw_frame,
            frames=self.animation_generator,  # Use generator to allow seamless restarts
            interval=self.anim_interval,  # in ms
            blit=True,
            cache_frame_data=False,
        )

        self._pick_cid = self.fig.canvas.mpl_connect("pick_event", self.on_scatter_pick)
        self.fcanvas.canvas.draw()

    def select_single_cell(self, index: int, timepoint: int) -> None:
        """
        Select a single cell.

        Parameters
        ----------
        index : int
            The cell index.
        timepoint : int
            The timepoint.
        """

        self.correct_btn.setEnabled(True)
        self.cancel_btn.setEnabled(True)
        self.del_shortcut.setEnabled(True)
        self.no_event_shortcut.setEnabled(True)

        self.track_of_interest = self.tracks[timepoint][index]
        print(f"You selected cell #{self.track_of_interest}...")
        self.give_cell_information()
        self.plot_signals()

        self.loc_t = []
        self.loc_idx = []
        for t in range(len(self.tracks)):
            indices = np.where(self.tracks[t] == self.track_of_interest)[0]
            if len(indices) > 0:
                self.loc_t.append(t)
                self.loc_idx.append(indices[0])

        self.previous_color = []
        for t, idx in zip(self.loc_t, self.loc_idx):
            self.previous_color.append(self.colors[t][idx].copy())
            self.colors[t][idx] = "lime"

    def shortcut_no_event(self):
        """Handle no event shortcut."""
        self.correct_btn.click()
        self.no_event_btn.click()
        self.correct_btn.click()

    def configure_ylims(self):
        """Configure y-axis limits for signal plots."""

        try:
            min_values = []
            max_values = []
            feats = []
            for i in range(len(self.signal_choice_cb)):
                signal = self.signal_choice_cb[i].currentText()
                if signal == "--":
                    continue
                else:
                    maxx = np.nanpercentile(
                        self.df_tracks.loc[:, signal].to_numpy().flatten(), 99
                    )
                    minn = np.nanpercentile(
                        self.df_tracks.loc[:, signal].to_numpy().flatten(), 1
                    )
                    min_values.append(minn)
                    max_values.append(maxx)
                    feats.append(signal)

            smallest_value = np.amin(min_values)
            feat_smallest_value = feats[np.argmin(min_values)]
            min_feat = self.df_tracks[feat_smallest_value].min()
            max_feat = self.df_tracks[feat_smallest_value].max()
            pad_small = (max_feat - min_feat) * 0.05
            if pad_small == 0:
                pad_small = 0.05

            largest_value = np.amax(max_values)
            feat_largest_value = feats[np.argmax(max_values)]
            min_feat = self.df_tracks[feat_largest_value].min()
            max_feat = self.df_tracks[feat_largest_value].max()
            pad_large = (max_feat - min_feat) * 0.05
            if pad_large == 0:
                pad_large = 0.05

            if len(min_values) > 0:
                self.cell_ax.set_ylim(
                    smallest_value - pad_small, largest_value + pad_large
                )
        except Exception as e:
            print(f"L1170 {e=}")
            pass

    def draw_frame(self, framedata: int) -> Tuple[Any, ...]:
        """
        Update plot elements at each timestep of the loop.

        Parameters
        ----------
        framedata : int
            The current frame index.

        Returns
        -------
        tuple
            The updated plot elements.
        """

        self.framedata = framedata
        self.frame_lbl.setText(f"frame: {self.framedata}")
        self.im.set_array(self.stack[self.framedata])
        self.status_scatter.set_offsets(self.positions[self.framedata])
        self.status_scatter.set_color(self.colors[self.framedata][:, 1])

        self.class_scatter.set_offsets(self.positions[self.framedata])
        self.class_scatter.set_edgecolor(self.colors[self.framedata][:, 0])

        return (
            self.im,
            self.status_scatter,
            self.class_scatter,
        )

    def stop(self):
        """Stop animation."""
        # # On stop we disconnect all of our events.
        self.stop_btn.hide()
        self.start_btn.show()
        self.anim.pause()
        self.prev_frame_btn.setEnabled(True)
        self.next_frame_btn.setEnabled(True)
        self.stop_btn.clicked.connect(self.start)

    def start(self):
        """Start animation."""
        self.start_btn.hide()
        self.stop_btn.show()

        self.prev_frame_btn.setEnabled(False)
        self.next_frame_btn.setEnabled(False)

        self.anim.resume()
        self.stop_btn.clicked.connect(self.stop)

    def next_frame(self):
        """Go to likely next frame."""
        self.framedata += 1
        if self.framedata >= self.len_movie:
            self.framedata = 0
        self.draw_frame(self.framedata)
        self.fcanvas.canvas.draw()

    def prev_frame(self):
        """Go to likely previous frame."""
        self.framedata -= 1
        if self.framedata < 0:
            self.framedata = self.len_movie - 1
        self.draw_frame(self.framedata)
        self.fcanvas.canvas.draw()

    def toggle_animation(self):
        """Toggle animation play/pause."""
        if self.stop_btn.isVisible():
            self.stop()
        else:
            self.start()

    def update_speed(self):
        """Update animation speed."""
        fps = self.speed_slider.value()
        # Convert FPS to interval in ms
        # FPS = 1000 / interval_ms => interval_ms = 1000 / FPS
        val = int(1000 / max(1, fps))
        self.anim_interval = val
        print(
            f"DEBUG: Speed slider moved. FPS: {fps} -> Interval: {val} ms. Recreating animation object."
        )

        # Check if animation is allowed to run (Pause button is visible means we are Playing)
        should_play = self.stop_btn.isVisible()

        if hasattr(self, "anim") and self.anim:
            try:
                self.anim.event_source.stop()
            except Exception as e:
                print(f"DEBUG: Error stopping animation: {e}")

        # Recreate animation with new interval
        try:
            # Disconnect the old pick event to avoid accumulating connections
            if hasattr(self, "_pick_cid"):
                try:
                    self.fig.canvas.mpl_disconnect(self._pick_cid)
                except Exception:
                    pass

            self.anim = FuncAnimation(
                self.fig,
                self.draw_frame,
                frames=self.animation_generator,
                interval=self.anim_interval,
                blit=True,
                cache_frame_data=False,
            )

            # Reconnect pick event and store cid
            self._pick_cid = self.fig.canvas.mpl_connect(
                "pick_event", self.on_scatter_pick
            )

            # If we were NOT playing (i.e. Paused), pause the new animation immediately
            if not should_play:
                self.anim.event_source.stop()

        except Exception as e:
            print(f"DEBUG: Error recreating animation: {e}")

    def give_cell_information(self):
        """Display cell information."""

        cell_selected = f"cell: {self.track_of_interest}\n"
        cell_class = f"class: {self.df_tracks.loc[self.df_tracks['TRACK_ID'] == self.track_of_interest, self.class_name].to_numpy()[0]}\n"
        cell_time = f"time of interest: {self.df_tracks.loc[self.df_tracks['TRACK_ID'] == self.track_of_interest, self.time_name].to_numpy()[0]}\n"
        self.cell_info.setText(cell_selected + cell_class + cell_time)

    def save_trajectories(self):
        """Save trajectories to file."""

        if self.normalized_signals:
            self.normalize_features_btn.click()
        if self.selection:
            self.cancel_selection()

        self.df_tracks = self.df_tracks.drop(
            self.df_tracks[self.df_tracks[self.class_name] > 2].index
        )
        self.df_tracks.to_csv(self.trajectories_path, index=False)
        print("Table successfully exported...")
        if self.class_choice_cb.currentText() != "":
            self.compute_status_and_colors(0)
        self.extract_scatter_from_trajectories()

    def set_first_frame(self):
        """Set to the first frame."""
        self.stop()
        self.framedata = 0
        self.draw_frame(self.framedata)
        self.fcanvas.canvas.draw()

    def set_last_frame(self):
        """Set to the last frame."""
        self.stop()
        self.framedata = len(self.stack) - 1
        while len(np.where(self.stack[self.framedata].flatten() == 0)[0]) > 0.99 * len(
            self.stack[self.framedata].flatten()
        ):
            self.framedata -= 1
            if self.framedata < 0:
                self.framedata = 0
                break

        self.draw_frame(self.framedata)
        self.fcanvas.canvas.draw()
