"""
Central plot panel - individual pg.PlotWidget rows on a shared X axis.

Row visibility:
  RF Magnitude   - always visible, no toggle
  RF I/Q         - optional, hidden by default
  Gx / Gy / Gz   - optional, hidden by default
  Gate RF/ADC/TR - optional, hidden by default
  Block dur.     - optional, hidden by default

Toggle buttons in the header bar show/hide entire rows.
Visibility state is preserved across plot_all() calls (e.g. after rerun).
"""
from __future__ import annotations

from typing import List, Optional

import numpy as np
import pyqtgraph as pg
from PySide6.QtCore import Signal, Qt, QPoint
from PySide6.QtGui import QColor, QFont
from PySide6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QLabel, QMenu,
    QApplication, QFileDialog,
)

from seq_interp.src.gui.adapters import (
    BlockRow, gate_to_step, build_block_rows, find_block_at_time,
)
from seq_interp.src.gui.scheme_panel import system_is_dark

# Waveform / overlay colours (used for curve pens and region fills)
_C = {
    "rf_mag":   "#d63031",
    "rf_real":  "#00b894",
    "rf_imag":  "#0984e3",
    "gx":       "#e17055",
    "gy":       "#6c5ce7",
    "gz":       "#00cec9",
    "gate_rf":  "#d63031",
    "gate_adc": "#0984e3",
    "gate_tr":  "#6c5ce7",
    "cl":       "#636e72",
    "start_d":  "#fd79a8",
    "rf_d":     "#fdcb6e",
    "tr_d":     "#74b9ff",
    "bound":    "#b2bec3",
}

# Toggle-button label colours - >=4.5:1 contrast on each background.
_C_TEXT_LIGHT = {
    "rf_mag":   "#b71c1c",
    "rf_real":  "#00695c",
    "rf_imag":  "#01579b",
    "gx":       "#bf360c",
    "gy":       "#4527a0",
    "gz":       "#006064",
    "gate_rf":  "#b71c1c",
    "gate_adc": "#01579b",
    "gate_tr":  "#4527a0",
    "cl":       "#37474f",
}
_C_TEXT_DARK = {
    "rf_mag":   "#ef5350",
    "rf_real":  "#26d4b0",
    "rf_imag":  "#42a5f5",
    "gx":       "#ff8a65",
    "gy":       "#9575cd",
    "gz":       "#26ceca",
    "gate_rf":  "#ef5350",
    "gate_adc": "#42a5f5",
    "gate_tr":  "#9575cd",
    "cl":       "#90a4ae",
}

_MAX_VLINES = 400


class PlotPanel(QWidget):
    """
    All waveforms and sync gates on a shared time axis.
    Emits blockClicked(sync_index) on left-click.
    Emits timeHovered(t_s, channel, value) for status-bar updates.
    """
    blockClicked = Signal(int)
    timeHovered  = Signal(float, str, float)

    def __init__(self, parent: QWidget | None = None):
        super().__init__(parent)

        _dark = system_is_dark()
        pg.setConfigOptions(
            antialias=True,
            background=QColor("#1e1e1e") if _dark else QColor("#ffffff"),
            foreground="w" if _dark else "k",
        )

        root = QVBoxLayout(self)
        root.setContentsMargins(0, 0, 0, 0)
        root.setSpacing(0)

        # -- toggle button bar ---------------------------------------------
        self._btn_bar = QWidget()
        self._btn_bar.setObjectName("PlotBtnBar")
        self._btn_bar.setStyleSheet(
            "#PlotBtnBar { border-bottom: 1px solid palette(mid); }"
        )
        self._btn_lay = QHBoxLayout(self._btn_bar)
        self._btn_lay.setContentsMargins(6, 3, 6, 3)
        self._btn_lay.setSpacing(4)
        lbl = QLabel("Show:")
        lbl.setFont(QFont("Arial", 9))
        self._btn_lay.addWidget(lbl)   # index 0 - stays forever
        self._btn_lay.addStretch()     # index 1 - stays forever; buttons inserted before it
        root.addWidget(self._btn_bar)

        # -- plot container ------------------------------------------------
        self._plot_area = QWidget()
        self._plot_area_lay = QVBoxLayout(self._plot_area)
        self._plot_area_lay.setContentsMargins(0, 0, 0, 0)
        self._plot_area_lay.setSpacing(1)
        root.addWidget(self._plot_area, stretch=1)

        # State
        self._plot_widgets: dict[str, pg.PlotWidget] = {}
        self._plots:        dict[str, pg.PlotItem]   = {}
        self._curves:       dict[str, pg.PlotDataItem] = {}
        self._row_btns:     dict[str, QPushButton]   = {}
        self._block_rows:   List[BlockRow] = []
        self._ref_plot:     Optional[pg.PlotItem] = None
        # Remembered visibility (persists across plot_all calls).
        # For simple rows the key is the row key; for combined rows (gradients,
        # gates) the key is the individual curve key (gx, gy, gz, gate_rf, ...).
        self._visible: set[str] = set()
        # Maps curve_key -> row_key for curves that live in a combined row.
        self._curve_to_row: dict[str, str] = {}
        # Maps row_key -> [curve_keys] for combined rows.
        self._row_curves: dict[str, list[str]] = {}

    # ------------------------------------------------------------------ #
    #  Public API                                                          #
    # ------------------------------------------------------------------ #

    def plot_all(self, seq_data: dict, sync_data: dict) -> None:
        self._clear()
        self._block_rows = build_block_rows(seq_data, sync_data)

        self._build_rf_rows(seq_data)
        self._build_gradients_row(seq_data)
        self._build_gates_row(sync_data)
        self._build_cl_row(sync_data)

        durs      = sync_data["blocks_duration"]
        cumtimes  = np.cumsum([0.0] + list(durs))
        self._draw_boundaries(cumtimes)
        self._draw_delay_regions()
        self._finalize_x_axis()
        self._attach_mouse_events()

        # Apply visibility.
        # rf_mag is always on (fall back to first row if sequence has no RF).
        always_key = "rf_mag" if "rf_mag" in self._plot_widgets else next(
            iter(self._plot_widgets), None
        )
        for key, pw in self._plot_widgets.items():
            if key == always_key:
                pw.setVisible(True)
            elif key in self._row_curves:
                # Combined row: visible when >=1 of its curves is on.
                # Also restore per-curve visibility and button states.
                curves = self._row_curves[key]
                for c_key in curves:
                    on = c_key in self._visible
                    c = self._curves.get(c_key)
                    if c:
                        c.setVisible(on)
                    if c_key in self._row_btns:
                        self._row_btns[c_key].setChecked(on)
                pw.setVisible(any(c in self._visible for c in curves))
            else:
                # Simple row: row key itself tracked in _visible.
                on = key in self._visible
                pw.setVisible(on)
                if key in self._row_btns:
                    self._row_btns[key].setChecked(on)

    def highlight_block(self, sync_index: int) -> None:
        for row in self._block_rows:
            if row.sync_index == sync_index:
                mid  = (row.t_start + row.t_end) / 2.0
                span = max(row.t_end - row.t_start, 1e-6)
                if self._ref_plot:
                    self._ref_plot.setXRange(mid - span * 5, mid + span * 5)
                break

    def fit_all(self) -> None:
        if self._ref_plot:
            self._ref_plot.enableAutoRange(axis="x")
            for p in self._plots.values():
                p.enableAutoRange(axis="y")

    # ------------------------------------------------------------------ #
    #  Row builders                                                        #
    # ------------------------------------------------------------------ #

    def _build_rf_rows(self, seq_data: dict) -> None:
        rf_raw = seq_data.get("rf")
        t_rf   = seq_data.get("t_rf")
        if rf_raw is None or t_rf is None or len(rf_raw) == 0:
            return

        rf  = np.asarray(rf_raw)
        t   = np.asarray(t_rf, dtype=float)
        mag = np.abs(rf)

        # RF Magnitude - always on, no toggle button
        p = self._add_row("rf_mag", "RF Mag", "a.u.", h=90)
        self._curves["rf_mag"] = p.plot(
            t, mag, pen=pg.mkPen(_C["rf_mag"], width=1.5), name="RF Mag"
        )

        # RF I/Q - optional
        p2 = self._add_row("rf_iq", "RF I/Q", "a.u.", h=75)
        self._curves["rf_real"] = p2.plot(
            t, rf.real, pen=pg.mkPen(_C["rf_real"], width=1.2), name="RF Re"
        )
        self._curves["rf_imag"] = p2.plot(
            t, rf.imag, pen=pg.mkPen(_C["rf_imag"], width=1.2), name="RF Im"
        )
        self._add_row_btn("rf_iq", "RF I/Q", "rf_real")

    def _build_gradients_row(self, seq_data: dict) -> None:
        axes = [("gx", "Gx"), ("gy", "Gy"), ("gz", "Gz")]
        present = [
            (key, lbl) for key, lbl in axes
            if key in seq_data and len(seq_data.get(key, []))
        ]
        if not present:
            return
        curve_keys = [key for key, _ in present]
        self._row_curves["gradients"] = curve_keys
        for key in curve_keys:
            self._curve_to_row[key] = "gradients"

        p = self._add_row("gradients", "Gradients", "Hz/m", h=100)
        p.addLegend(offset=(10, 5))
        for key, label in present:
            t = np.asarray(seq_data[f"t_{key}"], dtype=float)
            v = np.asarray(seq_data[key],         dtype=float)
            self._curves[key] = p.plot(
                t, v, pen=pg.mkPen(_C[key], width=1.5), name=label
            )
            self._add_curve_btn(key, label, key)

    def _build_gates_row(self, sync_data: dict) -> None:
        durs  = sync_data["blocks_duration"]
        gates = [
            ("gate_rf",        "Gate RF",  "gate_rf"),
            ("gate_adc",       "Gate ADC", "gate_adc"),
            ("gate_tr_switch", "Gate TR",  "gate_tr"),
        ]
        curve_keys = [dk for dk, _, _ in gates]
        self._row_curves["gates"] = curve_keys
        for key in curve_keys:
            self._curve_to_row[key] = "gates"

        p = self._add_row("gates", "Gates", "", h=80)
        p.setYRange(-0.15, 1.15)
        p.addLegend(offset=(10, 5))
        for data_key, label, color_key in gates:
            t, v  = gate_to_step(sync_data[data_key], durs)
            color = _C[color_key]
            self._curves[data_key] = p.plot(
                t, v,
                pen=pg.mkPen(color, width=1.5),
                fillLevel=0.0,
                brush=pg.mkBrush(color + "40"),
                name=label,
            )
            self._add_curve_btn(data_key, label, color_key)

    def _build_cl_row(self, sync_data: dict) -> None:
        durs     = sync_data["blocks_duration"]
        cumtimes = np.cumsum([0.0] + list(durs))
        t_pts, v_pts = [], []
        for i, d in enumerate(durs):
            t_pts += [cumtimes[i], cumtimes[i + 1]]
            v_pts += [d * 1e6,     d * 1e6]
        p = self._add_row("cl", "Block dur.", "uss", h=65)
        self._curves["cl"] = p.plot(
            np.array(t_pts), np.array(v_pts),
            pen=pg.mkPen(_C["cl"], width=1.2),
        )
        self._add_row_btn("cl", "Dur", "cl")

    # ------------------------------------------------------------------ #
    #  Overlays                                                            #
    # ------------------------------------------------------------------ #

    def _draw_boundaries(self, cumtimes: np.ndarray) -> None:
        n    = len(cumtimes)
        step = max(1, n // _MAX_VLINES)
        pen  = pg.mkPen(_C["bound"] + "80", width=0.5, style=Qt.DotLine)
        for plot in self._plots.values():
            for t in cumtimes[::step]:
                plot.addItem(
                    pg.InfiniteLine(pos=float(t), angle=90, pen=pen, movable=False)
                )

    def _draw_delay_regions(self) -> None:
        color_map = {
            "START": _C["start_d"] + "30",
            "RF":    _C["rf_d"]    + "30",
            "TR":    _C["tr_d"]    + "30",
        }
        for row in self._block_rows:
            if not row.is_delay:
                continue
            color = color_map.get(row.delay_type, "#ffffff30")
            for plot in self._plots.values():
                region = pg.LinearRegionItem(
                    values=[row.t_start, row.t_end],
                    brush=pg.mkBrush(color),
                    pen=pg.mkPen(None),
                    movable=False,
                )
                region.setZValue(-10)
                plot.addItem(region)

    # ------------------------------------------------------------------ #
    #  Mouse interactions                                                  #
    # ------------------------------------------------------------------ #

    def _attach_mouse_events(self) -> None:
        for name, plot in self._plots.items():
            plot.scene().sigMouseClicked.connect(
                lambda ev, p=plot, n=name: self._on_click(ev, p, n)
            )
            plot.scene().sigMouseMoved.connect(
                lambda pos, p=plot, n=name: self._on_hover(pos, p, n)
            )

    def _on_click(self, ev, plot: pg.PlotItem, channel: str) -> None:
        if not plot.sceneBoundingRect().contains(ev.scenePos()):
            return
        if ev.button() == Qt.RightButton:
            self._show_context_menu(ev, plot, channel)
            ev.accept()
            return
        if ev.button() == Qt.LeftButton:
            vb = plot.getViewBox()
            pt = vb.mapSceneToView(ev.scenePos())
            block = find_block_at_time(pt.x(), self._block_rows)
            if block is not None:
                self.blockClicked.emit(block.sync_index)
            ev.accept()

    def _on_hover(self, scene_pos, plot: pg.PlotItem, channel: str) -> None:
        if not plot.sceneBoundingRect().contains(scene_pos):
            return
        vb = plot.getViewBox()
        pt = vb.mapSceneToView(scene_pos)
        self.timeHovered.emit(pt.x(), channel, pt.y())

    def _show_context_menu(self, ev, plot: pg.PlotItem, channel: str) -> None:
        vb    = plot.getViewBox()
        t_val = vb.mapSceneToView(ev.scenePos()).x()

        menu     = QMenu(self)
        a_fit    = menu.addAction("Fit all")
        a_fit.triggered.connect(self.fit_all)
        a_reset  = menu.addAction("Reset zoom")
        a_reset.triggered.connect(lambda: vb.autoRange())
        menu.addSeparator()
        a_grid = menu.addAction("Toggle grid")
        a_grid.triggered.connect(
            lambda: plot.showGrid(
                x=not plot.ctrl.xGridCheck.isChecked(),
                y=not plot.ctrl.yGridCheck.isChecked(),
                alpha=0.3,
            )
        )
        a_copy = menu.addAction(f"Copy time  {t_val * 1e6:.4f} uss")
        a_copy.triggered.connect(
            lambda: QApplication.clipboard().setText(f"{t_val * 1e6:.6f}")
        )
        block = find_block_at_time(t_val, self._block_rows)
        if block is not None:
            menu.addSeparator()
            a_jump = menu.addAction(
                f"Jump to block #{block.sync_index}  [{block.block_type}]"
            )
            a_jump.triggered.connect(lambda: self.blockClicked.emit(block.sync_index))
        menu.addSeparator()
        a_export = menu.addAction("Export plot as PNG...")
        a_export.triggered.connect(lambda: self._export_plot(plot, channel))
        sp = ev.screenPos()
        menu.exec(QPoint(int(sp.x()), int(sp.y())))

    def _export_plot(self, plot: pg.PlotItem, channel: str) -> None:
        path, _ = QFileDialog.getSaveFileName(
            self, f"Export {channel}", f"{channel}.png",
            "PNG images (*.png);;All files (*)",
        )
        if not path:
            return
        pg.exporters.ImageExporter(plot).export(path)

    # ------------------------------------------------------------------ #
    #  Helpers                                                             #
    # ------------------------------------------------------------------ #

    def _add_row(self, key: str, label: str, units: str, h: int) -> pg.PlotItem:
        dark = system_is_dark()
        lbl_color = "#ccccee" if dark else "#333355"

        pw = pg.PlotWidget()
        pw.setMinimumHeight(h)
        pw.setMaximumHeight(h + 40)
        p = pw.getPlotItem()
        p.setLabel("left", label, units=units, color=lbl_color, size="9pt")
        # Bottom axis: numbers visible, label added later by _finalize_x_axis()
        p.showGrid(x=True, y=True, alpha=0.25)
        # Hide bottom axis label/numbers by default — shown only on last row
        p.getAxis("bottom").setStyle(showValues=False)
        p.getAxis("bottom").setLabel("")

        self._plot_widgets[key] = pw
        self._plots[key]        = p
        self._plot_area_lay.addWidget(pw)
        if self._ref_plot is None:
            self._ref_plot = p
        else:
            p.setXLink(self._ref_plot)
        return p

    def _finalize_x_axis(self) -> None:
        """Show X-axis tick values + label only on the last (bottom) row."""
        dark = system_is_dark()
        lbl_color = "#ccccee" if dark else "#333355"
        keys = list(self._plots.keys())
        for i, key in enumerate(keys):
            p = self._plots[key]
            ax = p.getAxis("bottom")
            if i == len(keys) - 1:
                ax.setStyle(showValues=True)
                p.setLabel("bottom", "Время", units="с",
                           color=lbl_color, size="9pt")
            else:
                ax.setStyle(showValues=False)
                p.getAxis("bottom").setLabel("")

    def _add_row_btn(self, key: str, label: str, color_key: str) -> None:
        _text = _C_TEXT_DARK if system_is_dark() else _C_TEXT_LIGHT
        color = _text.get(color_key, _C.get(color_key, "#555555"))
        btn   = QPushButton(label)
        btn.setCheckable(True)
        btn.setChecked(key in self._visible)
        btn.setFixedHeight(20)
        btn.setStyleSheet(
            f"QPushButton {{"
            f"  color: palette(mid); font-size: 11px;"
            f"  padding: 0px 7px;"
            f"  border: 1px solid palette(mid);"
            f"  border-radius: 3px;"
            f"}}"
            f"QPushButton:checked {{"
            f"  color: {color}; font-weight: bold;"
            f"  border-color: {color};"
            f"  background: palette(alternateBase);"
            f"}}"
        )
        btn.toggled.connect(lambda checked, k=key: self._on_row_toggle(k, checked))
        # Insert before the trailing stretch (last item in _btn_lay)
        self._btn_lay.insertWidget(self._btn_lay.count() - 1, btn)
        self._row_btns[key] = btn

    def _on_row_toggle(self, key: str, checked: bool) -> None:
        if checked:
            self._visible.add(key)
        else:
            self._visible.discard(key)
        pw = self._plot_widgets.get(key)
        if pw is not None:
            pw.setVisible(checked)

    def _add_curve_btn(self, curve_key: str, label: str, color_key: str) -> None:
        """Toggle button for one curve inside a combined row."""
        _text = _C_TEXT_DARK if system_is_dark() else _C_TEXT_LIGHT
        color = _text.get(color_key, _C.get(color_key, "#555555"))
        btn   = QPushButton(label)
        btn.setCheckable(True)
        btn.setChecked(curve_key in self._visible)
        btn.setFixedHeight(20)
        btn.setStyleSheet(
            f"QPushButton {{"
            f"  color: palette(mid); font-size: 11px;"
            f"  padding: 0px 7px;"
            f"  border: 1px solid palette(mid);"
            f"  border-radius: 3px;"
            f"}}"
            f"QPushButton:checked {{"
            f"  color: {color}; font-weight: bold;"
            f"  border-color: {color};"
            f"  background: palette(alternateBase);"
            f"}}"
        )
        btn.toggled.connect(
            lambda checked, k=curve_key: self._on_curve_toggle(k, checked)
        )
        self._btn_lay.insertWidget(self._btn_lay.count() - 1, btn)
        self._row_btns[curve_key] = btn

    def _on_curve_toggle(self, curve_key: str, checked: bool) -> None:
        if checked:
            self._visible.add(curve_key)
        else:
            self._visible.discard(curve_key)
        c = self._curves.get(curve_key)
        if c is not None:
            c.setVisible(checked)
        # Show the parent row when >=1 of its curves is on; hide when all off.
        row_key = self._curve_to_row.get(curve_key)
        if row_key:
            any_on = any(k in self._visible for k in self._row_curves.get(row_key, []))
            pw = self._plot_widgets.get(row_key)
            if pw is not None:
                pw.setVisible(any_on)

    def _clear(self) -> None:
        for pw in self._plot_widgets.values():
            self._plot_area_lay.removeWidget(pw)
            pw.deleteLater()
        self._plot_widgets.clear()
        self._plots.clear()
        self._curves.clear()
        self._ref_plot = None
        self._block_rows.clear()
        for btn in self._row_btns.values():
            self._btn_lay.removeWidget(btn)
            btn.deleteLater()
        self._row_btns.clear()
        self._curve_to_row.clear()
        self._row_curves.clear()
