import os
import re
import numpy as np
from qtpy import QtWidgets, QtCore
from matplotlib.figure import Figure
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
from matplotlib.widgets import RectangleSelector
from matplotlib.ticker import MaxNLocator, FormatStrFormatter

from ..ds_ramspec import Spectrum


class SequenceController(object):
    def __init__(self, model, widget):
        self.model = model
        self.widget = widget
        self.base_spectrum_ctrl = None
        self.plot_ctrl = None

        self._seq_canvas = None
        self._seq_ax = None
        self._seq_line = None
        self._seq_x = None
        self._seq_y = None

        self._selector_1d = None
        self._roi_artist_1d = None

        self._chi_files = []
        self._file_numbers = []
        self._chi_cache = {}
        self._roi_1d = None

        self._build_canvas()
        self._connect_channel()

    def set_helpers(self, base_spectrum_ctrl=None, base_ptn_ctrl=None, plot_ctrl=None):
        self.base_spectrum_ctrl = base_spectrum_ctrl or base_ptn_ctrl
        self.plot_ctrl = plot_ctrl

    def _build_canvas(self):
        if not hasattr(self.widget, "verticalLayout_SeqCanvas"):
            return
        self._seq_fig = Figure()
        self._seq_fig.subplots_adjust(left=0.09, right=0.985, top=0.96, bottom=0.14)
        self._seq_ax = self._seq_fig.add_subplot(111)
        self._seq_canvas = FigureCanvasQTAgg(self._seq_fig)
        self.widget.verticalLayout_SeqCanvas.addWidget(self._seq_canvas, 1)
        self._seq_canvas.mpl_connect("button_press_event", self._on_seq_click)

    def _connect_channel(self):
        if not hasattr(self.widget, "tabWidget"):
            return
        self.widget.tabWidget.currentChanged.connect(self._on_main_tab_changed)

        self.widget.pushButton_SeqLoadChi.clicked.connect(self._load_chi_files)
        self.widget.pushButton_SeqSetRoi.clicked.connect(self._arm_roi_selection)
        self.widget.pushButton_SeqClearRoi.clicked.connect(self._clear_roi)
        self.widget.pushButton_SeqCompute.clicked.connect(self._compute_sequence)
        self.widget.pushButton_SeqExportImage.clicked.connect(self._export_image)
        self.widget.pushButton_SeqExportNpy.clicked.connect(self._export_npy)

    def _on_main_tab_changed(self, _idx):
        if not hasattr(self.widget, "tab_Seq"):
            return
        if self.widget.tabWidget.currentWidget() != self.widget.tab_Seq:
            self.deactivate_interactions()
            self._clear_roi_overlays()
        else:
            self.refresh_roi_overlays()

    def deactivate_interactions(self):
        self._disable_roi_selectors()
        if hasattr(self.widget, "pushButton_SeqSetRoi"):
            self.widget.pushButton_SeqSetRoi.setChecked(False)

    def is_roi_selection_active(self):
        sel_1d_active = (self._selector_1d is not None) and \
            bool(getattr(self._selector_1d, "active", False))
        return bool(sel_1d_active)

    def _set_status(self, msg):
        if hasattr(self.widget, "label_SeqStatus"):
            self.widget.label_SeqStatus.setText(str(msg))

    def _set_loaded_count(self):
        if hasattr(self.widget, "label_SeqLoaded"):
            self.widget.label_SeqLoaded.setText(f"Loaded: {len(self._chi_files)}")

    def _load_chi_files(self):
        files, _ = QtWidgets.QFileDialog.getOpenFileNames(
            self.widget,
            "Select spectra for sequence",
            self.model.chi_path,
            "Spectra (*.spe *.SPE *.chi)",
        )
        if not files:
            return

        self._chi_files = sorted(list(files), key=self._filename_sort_key)
        self._file_numbers = self._derive_file_numbers(self._chi_files)
        self._chi_cache = {}
        self._roi_1d = None
        self._seq_x = None
        self._seq_y = None

        self._set_loaded_count()
        self._set_status(f"Loaded {len(self._chi_files)} spectra.")
        self._preview_first_file()
        self._set_default_1d_full_range_roi()
        if self._roi_1d is not None:
            self._compute_sequence()
        else:
            self._draw_sequence()

    def _filename_sort_key(self, filename):
        name = os.path.splitext(os.path.basename(filename))[0]
        name_lower = name.lower()
        nums = re.findall(r"(\d+)", name_lower)
        if nums:
            return (0, name_lower[:name_lower.rfind(nums[-1])], int(nums[-1]))
        return (1, name_lower)

    def _derive_file_numbers(self, files):
        out = []
        fallback = False
        for i, f in enumerate(files):
            name = os.path.splitext(os.path.basename(f))[0].lower()
            nums = re.findall(r"(\d+)", name)
            if nums:
                out.append(int(nums[-1]))
            else:
                fallback = True
                out.append(i + 1)
        if fallback:
            return list(range(1, len(files) + 1))
        return out

    def _preview_first_file(self):
        if not self._chi_files:
            return
        self._load_file_to_main_plot(0)

    def _load_file_to_main_plot(self, idx):
        if self.base_spectrum_ctrl is None:
            return
        if (idx < 0) or (idx >= len(self._chi_files)):
            return
        try:
            self.base_spectrum_ctrl._load_a_new_pattern(self._chi_files[idx])
            if self.plot_ctrl is not None:
                self.plot_ctrl.update()
            self._schedule_overlay_refresh()
        except Exception as exc:
            self._set_status(f"Failed to preview file: {exc}")

    def _schedule_overlay_refresh(self):
        self.refresh_roi_overlays()
        QtCore.QTimer.singleShot(0, self.refresh_roi_overlays)
        QtCore.QTimer.singleShot(80, self.refresh_roi_overlays)

    def _set_default_1d_full_range_roi(self):
        if not self._chi_files:
            return
        try:
            x, __ = self._load_bgsub_xy_if_requested(self._chi_files[0])
            if x.size == 0:
                return
            xmin = float(np.nanmin(x))
            xmax = float(np.nanmax(x))
            if (not np.isfinite(xmin)) or (not np.isfinite(xmax)) or (xmax <= xmin):
                return
            self._roi_1d = (xmin, xmax)
            self.widget.lineEdit_SeqRoiSummary.setText(f"Shift [{xmin:.3f}, {xmax:.3f}]")
        except Exception:
            pass

    def _arm_roi_selection(self):
        if self.widget.tabWidget.currentWidget() != self.widget.tab_Seq:
            self._set_status("Open Seq tab first.")
            return
        self._disable_roi_selectors()
        self.widget.pushButton_SeqSetRoi.setChecked(True)
        self._selector_1d = RectangleSelector(
            self.widget.mpl.canvas.ax_pattern,
            self._on_roi_1d_selected,
            useblit=True,
            button=[1],
            interactive=False,
            drag_from_anywhere=False,
        )
        self._set_status("Draw ROI on the spectrum plot.")

    def _disable_roi_selectors(self):
        if self._selector_1d is not None:
            try:
                self._selector_1d.set_active(False)
            except Exception:
                pass
            self._selector_1d = None

    def _clear_roi(self):
        self._roi_1d = None
        self.widget.lineEdit_SeqRoiSummary.setText("")
        self._set_status("ROI cleared.")
        self.deactivate_interactions()
        self._clear_roi_overlays()

    def _on_roi_1d_selected(self, eclick, erelease):
        if (eclick.xdata is None) or (erelease.xdata is None):
            return
        xmin = min(float(eclick.xdata), float(erelease.xdata))
        xmax = max(float(eclick.xdata), float(erelease.xdata))
        if xmax <= xmin:
            return
        self._roi_1d = (xmin, xmax)
        self.widget.lineEdit_SeqRoiSummary.setText(f"Shift [{xmin:.3f}, {xmax:.3f}]")
        self._set_status("ROI selected.")
        self.deactivate_interactions()
        self.refresh_roi_overlays()
        self._compute_sequence()

    def _load_processed_xy(self, spectrum_path):
        spectrum = Spectrum(spectrum_path)
        if str(spectrum_path).lower().endswith(".spe"):
            spectrum.apply_excitation_wavelength(
                float(self.widget.doubleSpinBox_SetWavelength.value()))
            if hasattr(self.widget, "spinBox_CCDRowMin") and hasattr(self.widget, "spinBox_CCDRowMax"):
                spectrum.set_spe_row_roi(
                    int(self.widget.spinBox_CCDRowMin.value()),
                    int(self.widget.spinBox_CCDRowMax.value()))
        use_bgsub = bool(
            getattr(self.widget, "checkBox_BgSub", None) and
            self.widget.checkBox_BgSub.isChecked()
        )
        fit_areas = []
        table = getattr(self.widget, "tableWidget_BackgroundConstraints", None)
        if table is not None:
            for row in range(table.rowCount()):
                item_min = table.item(row, 0)
                item_max = table.item(row, 1)
                if item_min is None or item_max is None:
                    continue
                try:
                    xmin = float(item_min.text())
                    xmax = float(item_max.text())
                except Exception:
                    continue
                if xmax < xmin:
                    xmin, xmax = xmax, xmin
                fit_areas.append([xmin, xmax])
        if use_bgsub:
            x_raw, y_raw = spectrum.get_raw()
            x_raw = np.asarray(x_raw, dtype=float)
            if x_raw.size == 0:
                return x_raw, np.asarray([], dtype=float)
            roi_min = float(self.widget.doubleSpinBox_Background_ROI_min.value())
            roi_max = float(self.widget.doubleSpinBox_Background_ROI_max.value())
            roi_min = max(float(np.nanmin(x_raw)), roi_min)
            roi_max = min(float(np.nanmax(x_raw)), roi_max)
            if roi_max <= roi_min:
                roi_min = float(np.nanmin(x_raw))
                roi_max = float(np.nanmax(x_raw))
            y_fit = np.asarray(y_raw, dtype=float)
            if (self.plot_ctrl is not None) and \
                    bool(getattr(self.plot_ctrl, "_smoothing_active", lambda: False)()):
                __, y_fit = self.plot_ctrl._get_smoothed_pattern_xy(x_raw, y_fit)
            spectrum.get_chbg(
                [roi_min, roi_max],
                [int(self.widget.spinBox_BGParam1.value())],
                yshift=0,
                fit_areas=fit_areas,
                y_source=y_fit,
            )
            x, y = spectrum.get_bgsub()
        else:
            x, y = spectrum.get_raw()
            x = np.asarray(x, dtype=float)
            y = np.asarray(y, dtype=float)
            if (self.plot_ctrl is not None) and \
                    bool(getattr(self.plot_ctrl, "_smoothing_active", lambda: False)()):
                x, y = self.plot_ctrl._get_smoothed_pattern_xy(x, y)
            return x, y
        return np.asarray(x, dtype=float), np.asarray(y, dtype=float)

    def _compute_sequence(self):
        if not self._chi_files:
            self._set_status("Load SPE files first.")
            return
        if self._roi_1d is None:
            self._set_status("Select ROI first.")
            return

        values = np.full(len(self._chi_files), np.nan, dtype=float)
        failures = []
        for i, chi_path in enumerate(self._chi_files):
            try:
                x, y = self._load_processed_xy(chi_path)
                xmin, xmax = self._roi_1d
                m = (x >= xmin) & (x <= xmax)
                if not np.any(m):
                    values[i] = np.nan
                else:
                    values[i] = float(np.nansum(y[m]))
            except Exception as exc:
                failures.append((chi_path, str(exc)))
                values[i] = np.nan

        self._seq_x = np.asarray(self._file_numbers, dtype=float)
        self._seq_y = values
        self._draw_sequence()

        if failures:
            first_path, first_err = failures[0]
            first_name = os.path.basename(first_path)
            self._set_status(
                f"Sequence computed with {len(failures)} failures. "
                f"First: {first_name} ({first_err})")
        else:
            self._set_status("Sequence computed from processed spectrum intensity in ROI.")
        self._schedule_overlay_refresh()

    def _current_title_text(self):
        if self._roi_1d is not None:
            xmin, xmax = self._roi_1d
            return f"Raman Shift [{xmin:.3f}, {xmax:.3f}]"
        return "Raman Shift"

    def _draw_sequence(self):
        if self._seq_ax is None:
            return
        self._seq_ax.clear()
        if (self._seq_x is None) or (self._seq_y is None) or (self._seq_x.size == 0):
            self._seq_ax.set_axis_off()
            self._seq_canvas.draw_idle()
            return
        self._seq_ax.set_axis_on()
        self._seq_line, = self._seq_ax.plot(
            self._seq_x,
            self._seq_y,
            linestyle="-",
            marker="o",
            markersize=5,
            linewidth=1.3,
            color="#1f77b4",
            markerfacecolor="#1f77b4",
            markeredgecolor="white",
            markeredgewidth=0.6,
        )
        self._seq_ax.set_title(self._current_title_text())
        self._seq_ax.set_xlabel("File number")
        self._seq_ax.set_ylabel("Integrated intensity")
        self._seq_ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        self._seq_ax.xaxis.set_major_formatter(FormatStrFormatter("%d"))
        self._seq_ax.grid(True, alpha=0.22, linewidth=0.6)
        self._seq_fig.tight_layout()
        self._seq_canvas.draw_idle()

    def _on_seq_click(self, event):
        if (self._seq_x is None) or (self._seq_y is None):
            return
        if event.inaxes != self._seq_ax:
            return
        if event.xdata is None:
            return
        finite = np.isfinite(self._seq_x) & np.isfinite(self._seq_y)
        if not np.any(finite):
            return
        x = self._seq_x[finite]
        y = self._seq_y[finite]
        idx_map = np.where(finite)[0]

        # Nearest point in display units for robust click selection.
        click_px = self._seq_ax.transData.transform((event.xdata, event.ydata if event.ydata is not None else 0.0))
        pts_px = self._seq_ax.transData.transform(np.column_stack([x, y]))
        dist2 = np.sum((pts_px - click_px) ** 2, axis=1)
        k = int(np.argmin(dist2))
        if dist2[k] > 20.0 ** 2:
            return

        file_idx = int(idx_map[k])
        self._load_file_to_main_plot(file_idx)
        file_num = int(self._file_numbers[file_idx]) if file_idx < len(self._file_numbers) else file_idx + 1
        self._set_status(f"Selected sequence point: file number {file_num}")

    def _is_seq_tab_active(self):
        return hasattr(self.widget, "tab_Seq") and \
            (self.widget.tabWidget.currentWidget() == self.widget.tab_Seq)

    def _should_show_roi_overlay(self):
        if (self._seq_y is None) or (self._seq_x is None):
            return False
        if not self._is_seq_tab_active():
            return False
        return self._roi_1d is not None

    def _clear_roi_overlays(self):
        changed = False
        if self._roi_artist_1d is not None:
            try:
                self._roi_artist_1d.remove()
            except Exception:
                pass
            self._roi_artist_1d = None
            changed = True
        if changed and hasattr(self.widget, "mpl") and hasattr(self.widget.mpl, "canvas"):
            try:
                self.widget.mpl.canvas.draw_idle()
            except Exception:
                pass

    def refresh_roi_overlays(self):
        self._clear_roi_overlays()
        if not self._should_show_roi_overlay():
            return
        try:
            if self._roi_1d is not None:
                xmin, xmax = self._roi_1d
                ax = self.widget.mpl.canvas.ax_pattern
                self._roi_artist_1d = ax.axvspan(
                    xmin, xmax, ymin=0.0, ymax=1.0,
                    facecolor="red", edgecolor="red", alpha=0.2, linewidth=1.2
                )
            self.widget.mpl.canvas.draw_idle()
        except Exception:
            self._roi_artist_1d = None

    def _export_image(self):
        if (self._seq_x is None) or (self._seq_y is None):
            self._set_status("No sequence data to export.")
            return
        base = os.path.join(self.model.chi_path, "sequence.png")
        filen, _ = QtWidgets.QFileDialog.getSaveFileName(
            self.widget,
            "Export sequence image",
            base,
            "Image files (*.png *.pdf)",
        )
        if not filen:
            return
        self._seq_fig.savefig(filen, dpi=300, bbox_inches="tight")
        self._set_status("Image exported.")

    def _export_npy(self):
        if (self._seq_x is None) or (self._seq_y is None):
            self._set_status("No sequence data to export.")
            return
        root = self.model.chi_path if str(getattr(self.model, "chi_path", "")).strip() else os.getcwd()
        out_dir = os.path.join(root, "sequence_py")
        os.makedirs(out_dir, exist_ok=True)

        npy_path = os.path.join(out_dir, "sequence.npy")
        py_path = os.path.join(out_dir, "plot_sequence.py")

        data = np.column_stack([self._seq_x, self._seq_y])
        np.save(npy_path, data)
        title = self._current_title_text()

        script = (
            "import numpy as np\n"
            "import matplotlib.pyplot as plt\n\n"
            "from matplotlib.ticker import MaxNLocator, FormatStrFormatter\n\n"
            "def main():\n"
            "    data = np.load('sequence.npy')\n"
            "    file_number = data[:, 0]\n"
            "    intensity = data[:, 1]\n"
            "    fig, ax = plt.subplots(figsize=(7, 4.5), facecolor='white')\n"
            "    ax.set_facecolor('white')\n"
            "    ax.plot(\n"
            "        file_number,\n"
            "        intensity,\n"
            "        linestyle='-',\n"
            "        marker='o',\n"
            "        markersize=5,\n"
            "        linewidth=1.3,\n"
            "        color='#1f77b4',\n"
            "        markerfacecolor='#1f77b4',\n"
            "        markeredgecolor='white',\n"
            "        markeredgewidth=0.6,\n"
            "    )\n"
            f"    ax.set_title({title!r})\n"
            "    ax.set_xlabel('File number')\n"
            "    ax.set_ylabel('Integrated intensity')\n"
            "    ax.xaxis.set_major_locator(MaxNLocator(integer=True))\n"
            "    ax.xaxis.set_major_formatter(FormatStrFormatter('%d'))\n"
            "    ax.grid(True, alpha=0.22, linewidth=0.6)\n"
            "    fig.tight_layout()\n"
            "    plt.show()\n\n"
            "if __name__ == '__main__':\n"
            "    main()\n"
        )
        with open(py_path, "w", encoding="utf-8") as fh:
            fh.write(script)

        self._set_status("Exported sequence_py (sequence.npy + plot_sequence.py).")
