# -*- coding: utf-8 -*-
"""Signatures tool: mean spectral signatures per class and Jeffries-Matusita
separability between class pairs.

Signatures come from the GEE stack sampled at the labeled points, so the
numbers match exactly what the classifier sees. J-M distance ranges 0
(identical, inseparable) to 2 (perfectly separable). Pairs < 1.0 are the
classes most likely to be confused.
"""

import math
import traceback

from qgis.PyQt.QtCore import Qt
from qgis.PyQt.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QFormLayout, QGroupBox, QComboBox,
    QPushButton, QLabel, QTextEdit, QMessageBox, QApplication, QFileDialog,
    QCheckBox,
)
from qgis.gui import QgsMapLayerComboBox
from qgis.core import QgsMapLayerProxyModel, QgsCoordinateTransform, QgsProject

from .. import gee_utils as gu
from .help_texts import help_button


class SignaturesDialog(QDialog):
    def __init__(self, iface, state, parent=None):
        super().__init__(parent)
        self.iface = iface
        self.state = state
        self.sig = None          # {class_code: {band: mean}}
        self.cov = None          # {class_code: covariance matrix}
        self.bands = []
        self.setWindowTitle("Signatures & Separability")
        self.setMinimumSize(640, 560)
        self._build_ui()

    def _build_ui(self):
        lay = QVBoxLayout(self)
        gb = QGroupBox("Inputs")
        f = QFormLayout(gb)
        self.cb_pts = QgsMapLayerComboBox()
        self.cb_pts.setFilters(QgsMapLayerProxyModel.Filter.PointLayer)
        f.addRow("Labeled points layer:", self.cb_pts)
        self.cb_level = QComboBox()
        self.cb_level.addItems(["Level 1 (class_code)",
                                "Level 2 (lu_code)"])
        f.addRow("Class level:", self.cb_level)
        self.lbl_scheme = QLabel("")
        self.lbl_scheme.setStyleSheet("color:#0a6; font-weight:bold;")
        self.lbl_scheme.setWordWrap(True)
        f.addRow("Active scheme:", self.lbl_scheme)
        self.chk_indices_only = QCheckBox(
            "Plot only indices in [-1, 1] range (cleaner curves)")
        self.chk_indices_only.setChecked(True)
        f.addRow("", self.chk_indices_only)
        lay.addWidget(gb)

        h = QHBoxLayout()
        self.btn_run = QPushButton("Compute signatures & separability")
        self.btn_run.clicked.connect(self.run)
        self.btn_plot = QPushButton("Show signature plot")
        self.btn_plot.clicked.connect(self.show_plot)
        self.btn_plot.setEnabled(False)
        h.addWidget(help_button("signatures", self))
        h.addWidget(self.btn_run); h.addWidget(self.btn_plot)
        lay.addLayout(h)

        self.txt = QTextEdit(); self.txt.setReadOnly(True)
        self.txt.setStyleSheet("font-family: Consolas, monospace;")
        lay.addWidget(self.txt)

        h2 = QHBoxLayout()
        self.btn_save = QPushButton("Save report (.txt)")
        self.btn_save.clicked.connect(self.save_report)
        self.btn_close = QPushButton("Close")
        self.btn_close.clicked.connect(self.close)
        h2.addStretch()
        h2.addWidget(self.btn_save); h2.addWidget(self.btn_close)
        lay.addLayout(h2)
        self.lbl_note = QLabel(
            "Needs the GEE stack from Tool 1 (or a loaded asset).")
        self.lbl_note.setWordWrap(True)
        lay.addWidget(self.lbl_note)
        self._sync_scheme()

    # ------------------------------------------------------------ helpers
    def _scheme(self):
        return (self.state.get("scheme") or "lulc")

    def showEvent(self, e):
        super().showEvent(e)
        self._sync_scheme()

    def _sync_scheme(self):
        scheme = self._scheme()
        names = {"lulc": "LULC (7 broad classes)",
                 "worldcover": "ESA WorldCover 2021 (11 classes)",
                 "custom": "Custom classes"}
        self.lbl_scheme.setText(names.get(scheme, scheme))
        cur = self.cb_level.currentIndex()
        self.cb_level.blockSignals(True)
        self.cb_level.clear()
        if scheme == "lulc":
            self.cb_level.addItems(["Level 1 (class_code)",
                                    "Level 2 (lu_code)"])
            self.cb_level.setEnabled(True)
            self.cb_level.setCurrentIndex(0 if cur not in (0, 1) else cur)
        else:
            self.cb_level.addItem(names.get(scheme, scheme))
            self.cb_level.setEnabled(False)
            self.cb_level.setCurrentIndex(0)
        self.cb_level.blockSignals(False)

    def _field(self):
        if self._scheme() == "lulc" and self.cb_level.currentIndex() == 1:
            return "lu_code"
        return "class_code"

    def _names(self):
        if self._scheme() == "lulc" and self.cb_level.currentIndex() == 1:
            return {k: v[2] for k, v in gu.LU2003.items()}
        return {c: n for c, n, _ in gu.scheme_classes(self._scheme())}

    def _points_fc(self, ee, layer, field):
        ct = QgsCoordinateTransform(layer.crs(), gu.WGS84,
                                    QgsProject.instance())
        feats = []
        names = [f.name() for f in layer.fields()]
        if field not in names:
            raise RuntimeError(
                "The points have no '{}' field. Label them in Tool 3 "
                "first.".format(field))
        for f in layer.getFeatures():
            code = int(f[field] or 0)
            if code <= 0:
                continue
            pt = ct.transform(f.geometry().asPoint())
            feats.append(ee.Feature(ee.Geometry.Point([pt.x(), pt.y()]),
                                    {"cls": code}))
        return ee.FeatureCollection(feats), len(feats)

    # ------------------------------------------------------------ run
    def run(self):
        stack = self.state.get("stack")
        if stack is None:
            QMessageBox.warning(self, "No stack",
                                "Build the stack in Tool 1 first.")
            return
        pts = self.cb_pts.currentLayer()
        if pts is None:
            QMessageBox.warning(self, "Points", "Choose the points layer.")
            return
        try:
            QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
            self.txt.setPlainText("Sampling stack at points on GEE ...")
            QApplication.processEvents()
            ee = gu.init_ee()
            field = self._field()
            fc, n = self._points_fc(ee, pts, field)
            if n < 10:
                raise RuntimeError("Too few labeled points ({}).".format(n))
            scale = int(self.state.get("scale") or 10)
            stack.bandNames()
            samples = stack.sampleRegions(
                collection=fc, properties=["cls"], scale=scale,
                tileScale=4, geometries=False)

            # pull sampled values once, compute stats locally (needs mean+cov)
            data = gu._retry(lambda: samples.getInfo())
            self._compute(data, field)
            self.btn_plot.setEnabled(True)
            self._write_report(n, scale)
        except Exception as e:
            self.txt.setPlainText("Error: {}\n\n{}".format(
                e, traceback.format_exc()))
        finally:
            QApplication.restoreOverrideCursor()

    def _compute(self, data, field):
        feats = data.get("features", [])
        band_names = [b for b in
                      (feats[0]["properties"].keys() if feats else [])
                      if b != "cls"]
        band_names.sort()
        by_class = {}
        for f in feats:
            pr = f["properties"]
            c = int(pr.get("cls", 0))
            if c <= 0:
                continue
            row = [pr.get(b) for b in band_names]
            if any(v is None for v in row):
                continue
            by_class.setdefault(c, []).append(row)

        self.bands = band_names
        self.sig = {}
        self.cov = {}
        self.counts = {}
        for c, rows in by_class.items():
            m = len(rows)
            k = len(band_names)
            means = [sum(r[j] for r in rows) / m for j in range(k)]
            self.sig[c] = dict(zip(band_names, means))
            self.counts[c] = m
            # covariance (unbiased); guard tiny samples
            cov = [[0.0] * k for _ in range(k)]
            if m > 1:
                for r in rows:
                    d = [r[j] - means[j] for j in range(k)]
                    for a in range(k):
                        for b in range(k):
                            cov[a][b] += d[a] * d[b]
                for a in range(k):
                    for b in range(k):
                        cov[a][b] /= (m - 1)
            # regularize diagonal for invertibility
            for a in range(k):
                cov[a][a] += 1e-6
            self.cov[c] = cov

    # ------------------------------------------------------------ J-M
    @staticmethod
    def _jm(m1, S1, m2, S2):
        import numpy as np
        m1 = np.array(m1); m2 = np.array(m2)
        S1 = np.array(S1); S2 = np.array(S2)
        Sm = (S1 + S2) / 2.0
        try:
            inv = np.linalg.inv(Sm)
        except (AttributeError, TypeError, ValueError):  # numpy/math
            inv = np.linalg.pinv(Sm)
        d = (m1 - m2)
        a = 0.125 * d.T @ inv @ d
        det_m = np.linalg.det(Sm)
        det1 = np.linalg.det(S1); det2 = np.linalg.det(S2)
        denom = math.sqrt(abs(det1 * det2)) or 1e-12
        b = 0.5 * math.log(abs(det_m) / denom) if det_m > 0 else 0.0
        B = a + b
        return 2.0 * (1.0 - math.exp(-B))

    # ------------------------------------------------------------ report
    def _write_report(self, n, scale):
        names = self._names()
        codes = sorted(self.sig)
        lines = []
        lines.append("SPECTRAL SIGNATURES & SEPARABILITY")
        lines.append("=" * 60)
        lines.append("Points used: {}   Scale: {} m   Bands: {}".format(
            n, scale, len(self.bands)))
        lines.append("Points per class: " + ", ".join(
            "{} ({})".format(names.get(c, c), self.counts[c])
            for c in codes))
        lines.append("")
        lines.append("Mean signatures (first bands shown):")
        show = self.bands[:12]
        header = "{:<14}".format("Class") + "".join(
            "{:>8}".format(b[:7]) for b in show)
        lines.append(header)
        for c in codes:
            row = "{:<14}".format(str(names.get(c, c))[:13])
            for b in show:
                row += "{:>8}".format(round(self.sig[c][b], 3))
            lines.append(row)
        lines.append("")
        lines.append("Jeffries-Matusita separability "
                     "(0 = identical, 2 = perfectly separable):")
        lines.append("Pairs below 1.0 are likely to be confused.")
        lines.append("")

        m = {c: [self.sig[c][b] for b in self.bands] for c in codes}
        worst = []
        header = "{:<14}".format("") + "".join(
            "{:>10}".format(str(names.get(c, c))[:9]) for c in codes)
        lines.append(header)
        for i in codes:
            row = "{:<14}".format(str(names.get(i, i))[:13])
            for j in codes:
                if i == j:
                    row += "{:>10}".format("-")
                else:
                    v = self._jm(m[i], self.cov[i], m[j], self.cov[j])
                    row += "{:>10.3f}".format(v)
                    if i < j:
                        worst.append((v, names.get(i, i), names.get(j, j)))
            lines.append(row)
        lines.append("")
        worst.sort()
        lines.append("Least separable pairs:")
        for v, a, b in worst[:6]:
            flag = "  <-- CONFUSED" if v < 1.0 else (
                "  (weak)" if v < 1.5 else "")
            lines.append("  {:<14} vs {:<14} : {:.3f}{}".format(
                str(a)[:13], str(b)[:13], v, flag))
        self.report_text = "\n".join(lines)
        self.txt.setPlainText(self.report_text)

    # ------------------------------------------------------------ plot
    def show_plot(self):
        """Open a LIVE, interactive vector plot (pan/zoom/toggle classes),
        not a saved image. Uses matplotlib's Qt canvas inside a window."""
        if not self.sig:
            return
        try:
            import matplotlib
            matplotlib.use("QtAgg")
            from matplotlib.backends.backend_qtagg import (
                FigureCanvasQTAgg as FigureCanvas,
                NavigationToolbar2QT as NavToolbar)
            from matplotlib.figure import Figure
        except (ImportError, ModuleNotFoundError, RuntimeError):  # GEE/Import
            try:
                import matplotlib
                matplotlib.use("Qt5Agg")
                from matplotlib.backends.backend_qt5agg import (
                    FigureCanvasQTAgg as FigureCanvas,
                    NavigationToolbar2QT as NavToolbar)
                from matplotlib.figure import Figure
            except (ImportError, ModuleNotFoundError, RuntimeError):  # GEE/Import
                QMessageBox.information(
                    self, "matplotlib missing",
                    "Install matplotlib in the QGIS Python environment:\n"
                    "  pip install matplotlib\n"
                    "(OSGeo4W Shell, then restart QGIS)")
                return

        names = self._names()
        colmap = {c: h for c, _, h in gu.scheme_classes(self._scheme())}
        if self.chk_indices_only.isChecked():
            bands = [b for b in self.bands
                     if all(-1.2 <= self.sig[c][b] <= 1.2 for c in self.sig)]
        else:
            bands = list(self.bands)
        if not bands:
            bands = list(self.bands)

        win = QDialog(self)
        win.setWindowTitle("Spectral signatures (live)")
        win.resize(1000, 620)
        v = QVBoxLayout(win)
        fig = Figure(figsize=(10, 5.5), tight_layout=True)
        canvas = FigureCanvas(fig)
        toolbar = NavToolbar(canvas, win)
        v.addWidget(toolbar)
        v.addWidget(canvas, 1)

        ax = fig.add_subplot(111)
        x = list(range(len(bands)))
        lines = {}
        for c in sorted(self.sig):
            y = [self.sig[c][b] for b in bands]
            if self._scheme() == "lulc" and self.cb_level.currentIndex() == 1:
                l1 = gu.LU2003.get(c, (0,))[0]
            else:
                l1 = c

            (ln,) = ax.plot(x, y, marker="o", ms=4, lw=1.8,
                            label="{} {}".format(c, names.get(c, "")),
                            color=colmap.get(l1))
            lines[c] = ln
        ax.set_xticks(x)
        ax.set_xticklabels(bands, rotation=60, ha="right", fontsize=8)
        ax.set_ylabel("Mean value")
        ax.set_title("Spectral signatures per class "
                     "(click a legend entry to show/hide)")
        ax.grid(True, alpha=0.3)
        leg = ax.legend(fontsize=8, ncol=2)

        # clickable legend: toggle a class on/off live
        handles = leg.get_lines()
        for h in handles:
            h.set_picker(True)
            h.set_pickradius(6)
        mapping = dict(zip(handles, lines.values()))

        def _on_pick(event):
            ln = mapping.get(event.artist)
            if ln is None:
                return
            vis = not ln.get_visible()
            ln.set_visible(vis)
            event.artist.set_alpha(1.0 if vis else 0.25)
            canvas.draw_idle()
        canvas.mpl_connect("pick_event", _on_pick)

        h = QHBoxLayout()
        b_svg = QPushButton("Export vector (SVG/PDF) ...")

        def _export():
            path, _ = QFileDialog.getSaveFileName(
                win, "Export plot", "signatures.svg",
                "SVG (*.svg);;PDF (*.pdf);;PNG (*.png)")
            if path:
                fig.savefig(path)
        b_svg.clicked.connect(_export)
        b_close = QPushButton("Close")
        b_close.clicked.connect(win.close)
        h.addWidget(b_svg); h.addStretch(); h.addWidget(b_close)
        v.addLayout(h)

        self._plot_win = win        # keep a reference so it is not GC'd
        win.show()

    def save_report(self):
        if not getattr(self, "report_text", ""):
            return
        path, _ = QFileDialog.getSaveFileName(
            self, "Save report", "separability_report.txt", "Text (*.txt)")
        if path:
            with open(path, "w", encoding="utf-8") as f:
                f.write(self.report_text)
