from pathlib import Path
from dataclasses import dataclass
import numpy as np
from PyQt5 import QtWidgets, QtCore
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas, NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure
from matplotlib.widgets import SpanSelector

from hdsemg_shared.fileio.file_io import EMGFile, Grid
from hdsemg_pipe._log.log_config import logger
from hdsemg_pipe.ui_elements.theme import Colors, Spacing, BorderRadius, Fonts, Styles

@dataclass
class GridData:
    """Helper to pair an EMGFile with one of its Grids."""
    emgfile: EMGFile
    grid: Grid

def _normalize_single(x):
    if isinstance(x, str):
        return x.lower()
    if isinstance(x, (list, tuple, np.ndarray)):
        return " ".join(_normalize_single(xx) for xx in x)
    return str(x).lower()

def normalize(desc):
    if isinstance(desc, np.ndarray):
        return np.array([_normalize_single(item) for item in desc])
    return _normalize_single(desc)

class CropRoiDialog(QtWidgets.QDialog):
    def __init__(self, file_paths, parent=None):
        super().__init__(parent)
        logger.info("Initializing Crop ROI Dialog for %d files", len(file_paths))

        self.file_paths = file_paths
        self.grid_items: list[GridData] = []
        self.selected_thresholds = (0, 0)
        self.reference_signal_map = {}
        self.threshold_lines = []
        self.span_selector = None
        self.lower_threshold = 0
        self.upper_threshold = 0
        self.click_mode = False  # Track if we're in click mode
        self.first_click_pos = None  # Store position of first click

        self.load_files()
        self.init_ui()

    def load_files(self):
        """Load each file via EMGFile and collect its Grids."""
        for fp in self.file_paths:
            try:
                logger.info("Loading file: %s", fp)
                emg = EMGFile.load(fp)
                for grid in emg.grids:
                    self.grid_items.append(GridData(emgfile=emg, grid=grid))
                logger.debug("→ %d grids from %s", len(emg.grids), Path(fp).name)
            except Exception as e:
                logger.error("Failed to load %s: %s", fp, e, exc_info=True)
                QtWidgets.QMessageBox.warning(self, "Loading Error", f"Failed to load {fp}:\n{e}")

        logger.info("Total grids loaded: %d", len(self.grid_items))

    def init_ui(self):
        """Initialize modern UI with GitHub-style design."""
        self.setWindowTitle("Crop Region of Interest (ROI)")
        self.resize(1400, 800)

        # Set dialog background
        self.setStyleSheet(f"""
            QDialog {{
                background-color: {Colors.BG_SECONDARY};
            }}
        """)

        # Build ref-signal map now that grid_items exists
        self.reference_signal_map = self.build_reference_signal_map()

        # Main layout
        main_layout = QtWidgets.QVBoxLayout(self)
        main_layout.setContentsMargins(Spacing.LG, Spacing.LG, Spacing.LG, Spacing.LG)
        main_layout.setSpacing(Spacing.LG)

        # Header with instructions
        header = QtWidgets.QLabel("Select Region of Interest")
        header.setStyleSheet(f"""
            QLabel {{
                color: {Colors.TEXT_PRIMARY};
                font-size: {Fonts.SIZE_XXL};
                font-weight: {Fonts.WEIGHT_BOLD};
                margin-bottom: {Spacing.SM}px;
            }}
        """)

        instruction = QtWidgets.QLabel(
            "🖱️ Drag to select region  •  🖱️ Click twice to set start/end points  •  🔍 Use toolbar to zoom/pan"
        )
        instruction.setStyleSheet(f"""
            QLabel {{
                color: {Colors.TEXT_SECONDARY};
                font-size: {Fonts.SIZE_BASE};
                margin-bottom: {Spacing.MD}px;
            }}
        """)
        instruction.setWordWrap(True)

        main_layout.addWidget(header)
        main_layout.addWidget(instruction)

        # Content area (plot + sidebar)
        content_layout = QtWidgets.QHBoxLayout()
        content_layout.setSpacing(Spacing.LG)

        # --- Plot area ---
        plot_container = QtWidgets.QFrame()
        plot_container.setStyleSheet(f"""
            QFrame {{
                background-color: {Colors.BG_PRIMARY};
                border: 1px solid {Colors.BORDER_DEFAULT};
                border-radius: {BorderRadius.LG};
                padding: {Spacing.MD}px;
            }}
        """)
        plot_layout = QtWidgets.QVBoxLayout(plot_container)
        plot_layout.setContentsMargins(Spacing.SM, Spacing.SM, Spacing.SM, Spacing.SM)

        self.figure = Figure(facecolor=Colors.BG_PRIMARY)
        self.canvas = FigureCanvas(self.figure)
        self.ax = self.figure.add_subplot(111)
        self.ax.set_facecolor(Colors.BG_PRIMARY)

        # Add matplotlib navigation toolbar for zoom/pan
        self.toolbar = NavigationToolbar(self.canvas, self)
        self.toolbar.setStyleSheet(f"""
            QToolBar {{
                background-color: {Colors.BG_SECONDARY};
                border: 1px solid {Colors.BORDER_DEFAULT};
                border-radius: {BorderRadius.SM};
                padding: {Spacing.XS}px;
                spacing: {Spacing.XS}px;
            }}
            QToolButton {{
                background-color: transparent;
                border: 1px solid transparent;
                border-radius: {BorderRadius.SM};
                padding: {Spacing.XS}px;
            }}
            QToolButton:hover {{
                background-color: {Colors.GRAY_100};
                border-color: {Colors.BORDER_DEFAULT};
            }}
        """)

        plot_layout.addWidget(self.toolbar)
        plot_layout.addWidget(self.canvas)

        # ROI info display
        roi_info_layout = QtWidgets.QHBoxLayout()
        roi_info_layout.setSpacing(Spacing.MD)

        self.roi_start_label = QtWidgets.QLabel("Start: 0")
        self.roi_end_label = QtWidgets.QLabel("End: 0")
        self.roi_duration_label = QtWidgets.QLabel("Duration: 0 samples")

        for label in [self.roi_start_label, self.roi_end_label, self.roi_duration_label]:
            label.setStyleSheet(f"""
                QLabel {{
                    color: {Colors.TEXT_PRIMARY};
                    font-size: {Fonts.SIZE_SM};
                    font-weight: {Fonts.WEIGHT_MEDIUM};
                    background-color: {Colors.GRAY_100};
                    border: 1px solid {Colors.BORDER_DEFAULT};
                    border-radius: {BorderRadius.SM};
                    padding: {Spacing.XS}px {Spacing.SM}px;
                }}
            """)

        roi_info_layout.addWidget(self.roi_start_label)
        roi_info_layout.addWidget(self.roi_end_label)
        roi_info_layout.addWidget(self.roi_duration_label)
        roi_info_layout.addStretch()

        plot_layout.addLayout(roi_info_layout)

        content_layout.addWidget(plot_container, stretch=3)

        # --- Control panel ---
        scroll = QtWidgets.QScrollArea()
        scroll.setWidgetResizable(True)
        scroll.setMaximumWidth(350)
        scroll.setStyleSheet(f"""
            QScrollArea {{
                border: none;
                background-color: transparent;
            }}
        """)

        panel = QtWidgets.QWidget()
        panel.setStyleSheet(f"background-color: transparent;")
        vbox = QtWidgets.QVBoxLayout(panel)
        vbox.setSpacing(Spacing.SM)
        vbox.setContentsMargins(0, 0, 0, 0)

        # Reference signals section
        ref_header = QtWidgets.QLabel("Reference Signals")
        ref_header.setStyleSheet(f"""
            QLabel {{
                color: {Colors.TEXT_PRIMARY};
                font-size: {Fonts.SIZE_LG};
                font-weight: {Fonts.WEIGHT_SEMIBOLD};
                margin-bottom: {Spacing.SM}px;
            }}
        """)
        vbox.addWidget(ref_header)

        ref_hint = QtWidgets.QLabel("Select signals to display in the plot:")
        ref_hint.setStyleSheet(f"""
            QLabel {{
                color: {Colors.TEXT_SECONDARY};
                font-size: {Fonts.SIZE_SM};
                margin-bottom: {Spacing.SM}px;
            }}
        """)
        vbox.addWidget(ref_hint)

        self.checkboxes = {}

        for gd in self.grid_items:
            key = gd.grid.grid_key
            uid = gd.grid.grid_uid
            box = QtWidgets.QGroupBox(f"Grid: {key}")
            box.setStyleSheet(Styles.groupbox())
            box_layout = QtWidgets.QVBoxLayout()
            box_layout.setSpacing(Spacing.XS)
            self.checkboxes[uid] = []

            ref_descs = self.reference_signal_map[uid]["ref_descriptions"]
            ref_indices = gd.grid.ref_indices  # Actual indices in the data array

            for idx, desc in enumerate(ref_descs):
                cb = QtWidgets.QCheckBox(f"Ref {idx} – {desc}")

                # Determine if this checkbox should be checked by default
                # Check if this reference signal corresponds to requested_path or performed_path
                actual_ref_idx = ref_indices[idx] if idx < len(ref_indices) else None
                should_check = False

                if gd.grid.requested_path_idx is not None and actual_ref_idx == gd.grid.requested_path_idx:
                    should_check = True
                elif gd.grid.performed_path_idx is not None and actual_ref_idx == gd.grid.performed_path_idx:
                    should_check = True
                elif gd.grid.requested_path_idx is None and gd.grid.performed_path_idx is None and idx == 0:
                    # Fallback: if no path indices are set, check the first one
                    should_check = True

                cb.setChecked(should_check)
                cb.stateChanged.connect(self.update_plot)
                cb.setStyleSheet(f"""
                    QCheckBox {{
                        color: {Colors.TEXT_PRIMARY};
                        font-size: {Fonts.SIZE_BASE};
                        spacing: {Spacing.SM}px;
                    }}
                    QCheckBox::indicator {{
                        width: 16px;
                        height: 16px;
                        border-radius: {BorderRadius.SM};
                        border: 1px solid {Colors.BORDER_DEFAULT};
                    }}
                    QCheckBox::indicator:checked {{
                        background-color: {Colors.BLUE_600};
                        border-color: {Colors.BLUE_600};
                    }}
                """)
                box_layout.addWidget(cb)
                self.checkboxes[uid].append(cb)

            box.setLayout(box_layout)
            vbox.addWidget(box)

        vbox.addStretch(1)

        # Action buttons
        btn_layout = QtWidgets.QVBoxLayout()
        btn_layout.setSpacing(Spacing.SM)

        reset_btn = QtWidgets.QPushButton("Reset Selection")
        reset_btn.setStyleSheet(Styles.button_secondary())
        reset_btn.clicked.connect(self.reset_selection)
        btn_layout.addWidget(reset_btn)

        ok_btn = QtWidgets.QPushButton("Apply & Close")
        ok_btn.setStyleSheet(Styles.button_primary())
        ok_btn.clicked.connect(self.on_ok_pressed)
        btn_layout.addWidget(ok_btn)

        vbox.addLayout(btn_layout)

        scroll.setWidget(panel)
        content_layout.addWidget(scroll, stretch=1)

        main_layout.addLayout(content_layout)

        # Initialize ROI bounds
        lo, hi = self.compute_data_xrange()
        self.lower_threshold = lo
        self.upper_threshold = hi

        # Setup interactive span selector for drag mode
        self.span_selector = SpanSelector(
            self.ax,
            self.on_span_select,
            'horizontal',
            useblit=True,
            props=dict(alpha=0.3, facecolor=Colors.BLUE_500),
            interactive=True,
            drag_from_anywhere=True
        )

        # Connect click event for click mode
        self.canvas.mpl_connect('button_press_event', self.on_click)

        self.update_plot()
        self.update_roi_info()

    def compute_data_xrange(self):
        maxlen = max((gd.emgfile.data.shape[0] for gd in self.grid_items), default=0)
        return (0, maxlen - 1 if maxlen>0 else 0)

    def on_click(self, event):
        """Handle click events for two-click selection mode."""
        # Only process clicks inside the plot area and with left mouse button
        if event.inaxes != self.ax or event.button != 1:
            return

        # Check if we're in a toolbar mode (zoom/pan)
        if self.toolbar.mode != '':
            return

        x_pos = int(event.xdata)

        if self.first_click_pos is None:
            # First click - set start position
            self.first_click_pos = x_pos
            logger.debug(f"First click at position: {x_pos}")

            # Draw a vertical line at the first click position
            self.draw_threshold_lines()

        else:
            # Second click - set end position and complete selection
            second_pos = x_pos
            self.lower_threshold = min(self.first_click_pos, second_pos)
            self.upper_threshold = max(self.first_click_pos, second_pos)
            logger.debug(f"Second click at position: {second_pos}, ROI: {self.lower_threshold}-{self.upper_threshold}")

            # Reset for next selection
            self.first_click_pos = None

            # Update visualization
            self.update_roi_info()
            self.draw_threshold_lines()

    def on_span_select(self, xmin, xmax):
        """Called when user selects a region with the span selector (drag mode)."""
        # Reset click mode when dragging
        self.first_click_pos = None

        self.lower_threshold = int(xmin)
        self.upper_threshold = int(xmax)
        self.update_roi_info()
        self.draw_threshold_lines()
        logger.debug(f"ROI selected via drag: {self.lower_threshold} - {self.upper_threshold}")

    def update_roi_info(self):
        """Update the ROI info labels."""
        duration = self.upper_threshold - self.lower_threshold
        self.roi_start_label.setText(f"Start: {self.lower_threshold}")
        self.roi_end_label.setText(f"End: {self.upper_threshold}")
        self.roi_duration_label.setText(f"Duration: {duration} samples")

    def draw_threshold_lines(self):
        """Draw vertical lines and shaded region for the current ROI selection."""
        # Remove old lines
        for line in self.threshold_lines:
            try:
                line.remove()
            except:
                pass
        self.threshold_lines.clear()

        # Draw based on selection state
        if self.first_click_pos is not None:
            # First click made - draw single line
            line = self.ax.axvline(
                self.first_click_pos,
                color=Colors.BLUE_600,
                linestyle='--',
                linewidth=2,
                label='Start'
            )
            self.threshold_lines.append(line)

        elif self.lower_threshold != 0 or self.upper_threshold != 0:
            # Complete selection - draw both lines and shaded region
            line1 = self.ax.axvline(
                self.lower_threshold,
                color=Colors.GREEN_600,
                linestyle='-',
                linewidth=2,
                alpha=0.7
            )
            line2 = self.ax.axvline(
                self.upper_threshold,
                color=Colors.GREEN_600,
                linestyle='-',
                linewidth=2,
                alpha=0.7
            )
            # Add shaded region
            span = self.ax.axvspan(
                self.lower_threshold,
                self.upper_threshold,
                alpha=0.15,
                color=Colors.GREEN_500
            )
            self.threshold_lines.extend([line1, line2, span])

        self.canvas.draw_idle()

    def reset_selection(self):
        """Reset the ROI selection to full range."""
        lo, hi = self.compute_data_xrange()
        self.lower_threshold = lo
        self.upper_threshold = hi
        self.first_click_pos = None  # Reset click mode
        self.update_roi_info()

        # Remove all threshold lines
        for line in self.threshold_lines:
            try:
                line.remove()
            except:
                pass
        self.threshold_lines.clear()

        # Remove the old span selector completely
        if self.span_selector:
            self.span_selector.set_visible(False)
            self.span_selector = None

        # Redraw the plot to remove all highlights
        self.update_plot()
        logger.info("ROI selection reset to full range")

    def on_ok_pressed(self):
        """Apply the selected ROI and close dialog."""
        self.selected_thresholds = (self.lower_threshold, self.upper_threshold)
        logger.info("User selected x-range: %s", self.selected_thresholds)
        self.accept()

    def update_plot(self):
        """Update the plot with selected reference signals."""
        self.ax.clear()
        self.ax.set_facecolor(Colors.BG_PRIMARY)

        # Plot selected signals
        for gd in self.grid_items:
            uid = gd.grid.grid_uid
            ref_data = self.reference_signal_map[uid]["ref_signals"]
            for idx, cb in enumerate(self.checkboxes[uid]):
                if cb.isChecked():
                    self.ax.plot(
                        ref_data[:, idx],
                        label=f"{gd.grid.grid_key}-Ref{idx}",
                        linewidth=1.5
                    )

        # Style the plot
        self.ax.set_xlabel("Sample Index", fontsize=12)
        self.ax.set_ylabel("Amplitude", fontsize=12)
        self.ax.legend(loc='upper right', framealpha=0.9)
        self.ax.grid(True, alpha=0.3, linestyle='--')

        # Re-create span selector after plot update
        if self.span_selector:
            self.span_selector = None
        self.span_selector = SpanSelector(
            self.ax,
            self.on_span_select,
            'horizontal',
            useblit=True,
            props=dict(alpha=0.3, facecolor=Colors.BLUE_500),
            interactive=True,
            drag_from_anywhere=True
        )

        self.canvas.draw_idle()

    def build_reference_signal_map(self):
        """
        Map each grid_uid → {
            'ref_signals': 2D np.array (samples×nRefs),
            'ref_descriptions': list[str]
        }
        """
        mp = {}
        for gd in self.grid_items:
            uid = gd.grid.grid_uid
            data = gd.emgfile.data
            desc = gd.emgfile.description
            idxs = gd.grid.ref_indices or []
            if not idxs:
                # fallback to first EMG channel
                idxs = [gd.grid.emg_indices[0]] if gd.grid.emg_indices else [0]

            ref_descs = [desc[i] for i in idxs]
            # normalize to str
            ref_descs = [normalize(rd) for rd in ref_descs]
            ref_data = data[:, idxs]

            mp[uid] = {
                "ref_signals": ref_data,
                "ref_descriptions": ref_descs
            }
        return mp
