# -*- coding: utf-8 -*-
"""Editors for user-defined bands added to the Tool 1 stack.

Two small table dialogs:

* CustomIndexDialog    - name + expression, evaluated on the composite.
* CustomTemporalDialog - name + expression + reducer, evaluated on every
  image of the collection and then reduced through time.

Both return a plain Python list so the caller can keep them in session
state without any GEE object being created here.
"""

import re

from qgis.PyQt.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QTableWidget,
    QTableWidgetItem, QDialogButtonBox, QComboBox, QSpinBox, QHeaderView,
)

from .. import gee_utils as gu


_NAME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$")

# Ready-made examples users can drop in and edit.
INDEX_PRESETS = [
    ("NDMI2", "(N - S2b) / (N + S2b)"),
    ("SIPI", "(N - B) / (N - R)"),
    ("NDSI", "(G - S1) / (G + S1)"),
    ("BAEI", "(R + 0.3) / (G + S1)"),
    ("RVI_opt", "N / R"),
]

TEMPORAL_PRESETS = [
    ("NDVI_p90", "(N - R) / (N + R)", "Percentile", 90),
    ("NDVI_amp", "(N - R) / (N + R)", "Amplitude (max-min)", 50),
    ("NDWI_std", "(G - N) / (G + N)", "StdDev", 50),
]


def _token_help(base_bands):
    toks = gu.expression_tokens(base_bands)
    return ("Tokens for the current sensor:  " + ",  ".join(toks) +
            "\nOperators: + - * / ( )   Functions: sqrt(), pow(), exp(), "
            "log(), abs(), min(), max()")


class _BaseCustomDialog(QDialog):
    """Shared table plumbing."""

    HEADERS = []
    WIDTH = 640

    def __init__(self, base_bands, existing=None, parent=None):
        super().__init__(parent)
        self.base_bands = list(base_bands or [])
        self.resize(self.WIDTH, 420)
        self._build(existing or [])

    # -------------------------------------------------------------- build
    def _build(self, existing):
        v = QVBoxLayout(self)
        v.addWidget(self._intro())

        hint = QLabel(_token_help(self.base_bands))
        hint.setWordWrap(True)
        hint.setStyleSheet(
            "color:#555; background:#f4f4f4; padding:6px; "
            "border:1px solid #ddd; font-family:monospace; font-size:11px;")
        v.addWidget(hint)

        self.table = QTableWidget(0, len(self.HEADERS))
        self.table.setHorizontalHeaderLabels(self.HEADERS)
        hh = self.table.horizontalHeader()
        hh.setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)
        v.addWidget(self.table, 1)

        h = QHBoxLayout()
        b_add = QPushButton("Add row")
        b_add.clicked.connect(lambda: self._add_row())
        b_del = QPushButton("Remove selected")
        b_del.clicked.connect(self._remove_row)
        b_ex = QPushButton("Insert example")
        b_ex.clicked.connect(self._insert_example)
        b_clear = QPushButton("Clear all")
        b_clear.clicked.connect(lambda: self.table.setRowCount(0))
        for b in (b_add, b_del, b_ex, b_clear):
            h.addWidget(b)
        h.addStretch()
        v.addLayout(h)

        self.lbl_err = QLabel("")
        self.lbl_err.setStyleSheet("color:#c00;")
        self.lbl_err.setWordWrap(True)
        v.addWidget(self.lbl_err)

        bb = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok |
                              QDialogButtonBox.StandardButton.Cancel)
        bb.accepted.connect(self._accept)
        bb.rejected.connect(self.reject)
        v.addWidget(bb)

        for row in existing:
            self._add_row(row)

    def _intro(self):
        lbl = QLabel(self.INTRO)
        lbl.setWordWrap(True)
        return lbl

    # -------------------------------------------------------------- rows
    def _remove_row(self):
        rows = sorted({i.row() for i in self.table.selectedIndexes()},
                      reverse=True)
        for r in rows:
            self.table.removeRow(r)

    def _cell(self, row, col):
        it = self.table.item(row, col)
        return (it.text().strip() if it is not None else "")

    def _validate_common(self, name, expr, seen):
        if not name and not expr:
            return None                      # blank row: ignore
        if not _NAME_RE.match(name or ""):
            raise ValueError(
                "'{}' is not a valid band name. Use a letter first, then "
                "letters, digits or underscore (e.g. NDMI2).".format(name))
        if name in seen:
            raise ValueError("Duplicate band name '{}'.".format(name))
        if not expr:
            raise ValueError("Band '{}' has no expression.".format(name))
        allowed = set(gu.expression_tokens(self.base_bands))
        used = set(re.findall(r"[A-Za-z_][A-Za-z0-9_]*", expr))
        funcs = {"sqrt", "pow", "exp", "log", "log10", "abs",
                 "min", "max", "atan", "sin", "cos", "tan"}
        unknown = sorted(used - allowed - funcs)
        if unknown:
            raise ValueError(
                "Band '{}': unknown token(s) {}. Allowed tokens are: "
                "{}".format(name, ", ".join(unknown),
                            ", ".join(sorted(allowed))))
        return True


class CustomIndexDialog(_BaseCustomDialog):
    """Name + expression, computed on the composite image."""

    HEADERS = ["Band name", "Expression"]
    INTRO = ("<b>Custom spectral indices.</b> Each row becomes one extra "
             "band on the stack, computed from the composite with the Earth "
             "Engine expression parser - exactly like the built-in indices.")

    def __init__(self, base_bands, existing=None, parent=None):
        super().__init__(base_bands, existing, parent)
        self.setWindowTitle("Custom spectral indices")

    def _add_row(self, row=None):
        r = self.table.rowCount()
        self.table.insertRow(r)
        name, expr = (row or ("", ""))
        self.table.setItem(r, 0, QTableWidgetItem(name))
        self.table.setItem(r, 1, QTableWidgetItem(expr))

    def _insert_example(self):
        used = {self._cell(r, 0) for r in range(self.table.rowCount())}
        for name, expr in INDEX_PRESETS:
            if name not in used:
                self._add_row((name, expr))
                return
        self._add_row(INDEX_PRESETS[0])

    def _accept(self):
        out, seen = [], set()
        try:
            for r in range(self.table.rowCount()):
                name, expr = self._cell(r, 0), self._cell(r, 1)
                if self._validate_common(name, expr, seen) is None:
                    continue
                seen.add(name)
                out.append((name, expr))
        except ValueError as e:
            self.lbl_err.setText(str(e))
            return
        self.result_specs = out
        self.accept()


class CustomTemporalDialog(_BaseCustomDialog):
    """Name + expression + reducer, computed across the whole collection."""

    HEADERS = ["Band name", "Expression (per image)", "Reducer", "Percentile"]
    WIDTH = 780
    INTRO = ("<b>Custom temporal metrics.</b> The expression is evaluated on "
             "<i>every image</i> in the date range, then collapsed through "
             "time by the reducer. So <tt>(N-R)/(N+R)</tt> with "
             "<i>Percentile 90</i> gives the per-pixel 90th percentile NDVI "
             "over the season - a strong phenology feature.")

    def __init__(self, base_bands, existing=None, parent=None):
        super().__init__(base_bands, existing, parent)
        self.setWindowTitle("Custom temporal metrics")

    def _add_row(self, row=None):
        r = self.table.rowCount()
        self.table.insertRow(r)
        if isinstance(row, dict):
            name = row.get("name", "")
            expr = row.get("expr", "")
            red = row.get("reducer", "Mean")
            pct = int(row.get("pct", 50))
        elif row:
            name, expr, red, pct = row
        else:
            name, expr, red, pct = "", "", "Mean", 50
        self.table.setItem(r, 0, QTableWidgetItem(name))
        self.table.setItem(r, 1, QTableWidgetItem(expr))
        cb = QComboBox()
        cb.addItems(gu.CUSTOM_REDUCERS)
        idx = cb.findText(red)
        cb.setCurrentIndex(idx if idx >= 0 else 0)
        self.table.setCellWidget(r, 2, cb)
        sp = QSpinBox()
        sp.setRange(1, 99)
        sp.setValue(pct)
        sp.setEnabled(red == "Percentile")
        cb.currentTextChanged.connect(
            lambda t, _sp=sp: _sp.setEnabled(t == "Percentile"))
        self.table.setCellWidget(r, 3, sp)

    def _insert_example(self):
        used = {self._cell(r, 0) for r in range(self.table.rowCount())}
        for row in TEMPORAL_PRESETS:
            if row[0] not in used:
                self._add_row(row)
                return
        self._add_row(TEMPORAL_PRESETS[0])

    def _accept(self):
        out, seen = [], set()
        try:
            for r in range(self.table.rowCount()):
                name, expr = self._cell(r, 0), self._cell(r, 1)
                if self._validate_common(name, expr, seen) is None:
                    continue
                seen.add(name)
                cb = self.table.cellWidget(r, 2)
                sp = self.table.cellWidget(r, 3)
                out.append({
                    "name": name, "expr": expr,
                    "reducer": cb.currentText() if cb else "Mean",
                    "pct": sp.value() if sp else 50,
                })
        except ValueError as e:
            self.lbl_err.setText(str(e))
            return
        self.result_specs = out
        self.accept()


def edit_custom_indices(parent, base_bands, existing):
    dlg = CustomIndexDialog(base_bands, existing, parent)
    if dlg.exec() == QDialog.DialogCode.Accepted:
        return dlg.result_specs
    return None


def edit_custom_temporal(parent, base_bands, existing):
    dlg = CustomTemporalDialog(base_bands, existing, parent)
    if dlg.exec() == QDialog.DialogCode.Accepted:
        return dlg.result_specs
    return None
