# -*- coding: utf-8 -*-
"""Tool 2: Create sample points for training / validation.

Sample size: Cochran formula n = z^2 * p * (1 - p) / e^2.
Three methods:
  1. Simple random.
  2. Stratified by ESA WorldCover on GEE (recommended) - no local
     raster needed; each point also gets a WorldCover "hint".
  3. Stratified by a local reference raster (GDAL file, not a GEE
     tile layer).
"""

import math
import random
import secrets
_train_random = secrets.SystemRandom()  # cryptographically secure, avoids Bandit S311
import traceback

from qgis.PyQt.QtCore import Qt, QVariant
from qgis.PyQt.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QFormLayout, QGroupBox, QComboBox,
    QSpinBox, QDoubleSpinBox, QCheckBox, QPushButton, QLabel, QMessageBox,
    QApplication, QScrollArea, QWidget,
)
from qgis.PyQt.QtGui import QColor
from qgis.gui import QgsMapLayerComboBox
from qgis.core import (
    QgsMapLayerProxyModel, QgsProject, QgsField, QgsVectorLayer, QgsFeature,
    QgsGeometry, QgsPointXY, QgsCategorizedSymbolRenderer, QgsRendererCategory,
    QgsMarkerSymbol,
)

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

M_RANDOM, M_WC, M_LOCAL = 0, 1, 2


class SamplingDialog(QDialog):
    def __init__(self, iface, parent=None, state=None):
        super().__init__(parent)
        self.iface = iface
        self.state = state if state is not None else {}
        self.setWindowTitle("2. Create Sample Points")
        self.setMinimumWidth(460)
        self.resize(480, 620)
        self._build_ui()
        self._recalc()

    def _build_ui(self):
        root = QVBoxLayout(self)
        scroll = QScrollArea(); scroll.setWidgetResizable(True)
        content = QWidget()
        lay = QVBoxLayout(content)

        gb0 = QGroupBox("AOI polygon (your shapefile / GPKG / AOI layer)")
        f0 = QFormLayout(gb0)
        self.cb_aoi = QgsMapLayerComboBox()
        self.cb_aoi.setFilters(QgsMapLayerProxyModel.Filter.PolygonLayer)
        f0.addRow("Polygon layer:", self.cb_aoi)
        self.chk_sel = QCheckBox("Use only selected feature(s)")
        f0.addRow("", self.chk_sel)
        lay.addWidget(gb0)

        gb1 = QGroupBox("Sample size (Cochran formula)")
        f1 = QFormLayout(gb1)
        self.sp_p = QDoubleSpinBox(); self.sp_p.setRange(0.5, 0.99)
        self.sp_p.setSingleStep(0.05); self.sp_p.setValue(0.85)
        self.sp_e = QDoubleSpinBox(); self.sp_e.setRange(0.01, 0.20)
        self.sp_e.setSingleStep(0.01); self.sp_e.setValue(0.05)
        self.sp_z = QDoubleSpinBox(); self.sp_z.setRange(1.0, 3.0)
        self.sp_z.setSingleStep(0.01); self.sp_z.setValue(1.96)
        self.sp_extra = QSpinBox(); self.sp_extra.setRange(0, 100)
        self.sp_extra.setValue(10); self.sp_extra.setSuffix(" %")
        for w in (self.sp_p, self.sp_e, self.sp_z):
            w.valueChanged.connect(self._recalc)
        self.sp_extra.valueChanged.connect(self._recalc)
        f1.addRow("Expected accuracy p:", self.sp_p)
        f1.addRow("Margin of error e:", self.sp_e)
        f1.addRow("Confidence z:", self.sp_z)
        f1.addRow("Extra buffer:", self.sp_extra)
        self.lbl_n = QLabel("n = ...")
        f1.addRow("Computed n:", self.lbl_n)
        self.sp_n = QSpinBox(); self.sp_n.setRange(10, 100000)
        f1.addRow("Points to create:", self.sp_n)
        lay.addWidget(gb1)

        gb2 = QGroupBox("Method")
        f2 = QFormLayout(gb2)
        self.cb_method = QComboBox()
        self.cb_method.addItems([
            "Simple random",
            "Stratified by ESA WorldCover (GEE, recommended)",
            "Stratified by local reference raster (GDAL file)",
        ])
        self.cb_method.setCurrentIndex(M_WC)          # recommended default
        self.cb_method.currentIndexChanged.connect(
            lambda i: self.cb_ref.setEnabled(i == M_LOCAL))
        f2.addRow("Sampling:", self.cb_method)
        self.cb_alloc = QComboBox()
        self.cb_alloc.addItems([
            "Equal per class (best for TRAINING data)",
            "Proportional to class area (best for VALIDATION, "
            "Olofsson 2014)",
        ])
        self.cb_alloc.setCurrentIndex(1)              # proportional default
        f2.addRow("Allocation:", self.cb_alloc)
        self.sp_minclass = QSpinBox()
        self.sp_minclass.setRange(10, 200); self.sp_minclass.setValue(30)
        self.sp_minclass.setToolTip(
            "Minimum points per class. Olofsson et al. (2014) recommend "
            "50-100 per class for validation samples.")
        f2.addRow("Min per class:", self.sp_minclass)
        self.cb_ref = QgsMapLayerComboBox()
        self.cb_ref.setFilters(QgsMapLayerProxyModel.Filter.RasterLayer)
        self.cb_ref.setEnabled(False)
        self.cb_ref.setToolTip(
            "A classified land cover raster (GeoTIFF), e.g. ESA WorldCover "
            "or a previous LULC map. GEE/XYZ tile layers are hidden "
            "because QGIS cannot read pixel values from them.")
        self._refresh_ref_filter()
        QgsProject.instance().layersAdded.connect(self._refresh_ref_filter)
        QgsProject.instance().layersRemoved.connect(self._refresh_ref_filter)
        f2.addRow("Reference raster:", self.cb_ref)
        self.sp_mindist = QDoubleSpinBox()
        self.sp_mindist.setRange(0, 100000); self.sp_mindist.setValue(0)
        self.sp_mindist.setToolTip(
            "Minimum distance between points, AOI layer CRS units. "
            "Not used by the WorldCover method.")
        f2.addRow("Min. distance:", self.sp_mindist)
        self.sp_train = QSpinBox(); self.sp_train.setRange(10, 90)
        self.sp_train.setValue(70); self.sp_train.setSuffix(" % training")
        f2.addRow("Train / valid split:", self.sp_train)
        lay.addWidget(gb2)

        lay.addStretch(1)
        scroll.setWidget(content)
        root.addWidget(scroll, 1)

        hbtn = QHBoxLayout()
        self.btn_run = QPushButton("Create sample points layer")
        self.btn_run.setMinimumHeight(32)
        self.btn_run.setStyleSheet("font-weight: bold;")
        self.btn_run.clicked.connect(self.run)
        self.btn_reset = QPushButton("Reset")
        self.btn_reset.clicked.connect(self.reset_defaults)
        self.btn_close = QPushButton("Close")
        self.btn_close.clicked.connect(self.close)
        hbtn.addWidget(help_button("sampling", self))
        hbtn.addWidget(self.btn_run, 1)
        hbtn.addSpacing(20)
        hbtn.addWidget(self.btn_reset); hbtn.addWidget(self.btn_close)
        root.addLayout(hbtn)
        self.lbl_status = QLabel(""); self.lbl_status.setWordWrap(True)
        root.addWidget(self.lbl_status)

    def reset_defaults(self):
        self.sp_p.setValue(0.85); self.sp_e.setValue(0.05)
        self.sp_z.setValue(1.96); self.sp_extra.setValue(10)
        self.cb_method.setCurrentIndex(M_WC)
        self.cb_alloc.setCurrentIndex(1)
        self.sp_minclass.setValue(30)
        self.sp_mindist.setValue(0)
        self.sp_train.setValue(70)
        self.chk_sel.setChecked(False)
        self.lbl_status.setText("Reset to defaults.")
        self._recalc()

    def _refresh_ref_filter(self, *_):
        from qgis.core import QgsRasterLayer
        try:
            bad = [lyr for lyr in QgsProject.instance().mapLayers().values()
                   if isinstance(lyr, QgsRasterLayer)
                   and lyr.providerType() != "gdal"]
            self.cb_ref.setExceptedLayerList(bad)
        except (AttributeError, TypeError, RuntimeError):  # QGIS core
            pass

    def _recalc(self):
        p, e, z = self.sp_p.value(), self.sp_e.value(), self.sp_z.value()
        n = (z * z * p * (1 - p)) / (e * e)
        n = int(round(n * (1 + self.sp_extra.value() / 100.0)))
        self.lbl_n.setText("n = {}  (Cochran + buffer)".format(n))
        self.sp_n.setValue(n)

    # ------------------------------------------------------------ AOI helpers
    def _aoi_layer_and_geom(self):
        lyr = self.cb_aoi.currentLayer()
        if lyr is None:
            raise RuntimeError("Choose an AOI polygon layer "
                               "(your shapefile works here).")
        feats = (lyr.selectedFeatures()
                 if self.chk_sel.isChecked() and lyr.selectedFeatureCount() > 0
                 else list(lyr.getFeatures()))
        if not feats:
            raise RuntimeError("The AOI layer has no features / no selection.")
        geom = QgsGeometry(feats[0].geometry())
        for ft in feats[1:]:
            geom = geom.combine(ft.geometry())
        return lyr, geom

    # ------------------------------------------------------------ run
    def run(self):
        try:
            QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
            self.lbl_status.setText("Working ...")
            QApplication.processEvents()
            method = self.cb_method.currentIndex()
            n_target = self.sp_n.value()

            if method == M_WC:
                pts = self._gee_worldcover_stratified(n_target)
            elif method == M_LOCAL:
                pts = self._local_raster_stratified(n_target)
            else:
                pts = self._simple_random(n_target)

            out = self._finalize(pts)
            QgsProject.instance().addMapLayer(out)
            gu.raise_points_above(("sample_points",))
            self.lbl_status.setText(
                "Created '{}' with {} points ({}% train / {}% valid).\n"
                "Open Tool 3 to label them. Save the layer as GPKG so "
                "your labels are safe.".format(
                    out.name(), out.featureCount(),
                    self.sp_train.value(), 100 - self.sp_train.value()))
        except Exception as e:
            self.lbl_status.setText("Error: {}".format(e))
            QMessageBox.critical(self, "Error",
                                 "{}\n\n{}".format(e, traceback.format_exc()))
        finally:
            QApplication.restoreOverrideCursor()

    # ------------------------------------------------------------ methods
    def _simple_random(self, n_target):
        import processing
        lyr, _ = self._aoi_layer_and_geom()
        src = lyr
        if self.chk_sel.isChecked() and lyr.selectedFeatureCount() > 0:
            src = processing.run("native:saveselectedfeatures",
                                 {"INPUT": lyr, "OUTPUT": "memory:"})["OUTPUT"]
        pts = processing.run("native:randompointsinpolygons", {
            "INPUT": src, "POINTS_NUMBER": n_target,
            "MIN_DISTANCE": self.sp_mindist.value(),
            "OUTPUT": "memory:"})["OUTPUT"]
        return pts

    def _gee_worldcover_stratified(self, n_target):
        """Server-side stratified sample on ESA WorldCover. Equal
        allocation per class present in the AOI, min 20 per class."""
        ee = gu.init_ee()
        lyr, geom = self._aoi_layer_and_geom()
        region = gu.qgis_geom_to_ee(geom, lyr.crs())
        wc = gu.worldcover_image()

        area_km2 = region.area(1).getInfo() / 1e6
        scale = 10 if area_km2 <= 2000 else 100

        hist = wc.reduceRegion(
            reducer=ee.Reducer.frequencyHistogram(), geometry=region,
            scale=max(scale, 50), maxPixels=1e13, bestEffort=True
        ).get("Map").getInfo() or {}
        hist = {int(float(k)): v for k, v in hist.items()}
        classes = sorted(hist)
        if not classes:
            raise RuntimeError("WorldCover has no pixels inside this AOI.")
        total = float(sum(hist.values()))
        min_pc = self.sp_minclass.value()
        if self.cb_alloc.currentIndex() == 1:      # proportional
            pts = [max(min_pc, int(round(n_target * hist[c] / total)))
                   for c in classes]
        else:                                       # equal
            per = max(min_pc, n_target // len(classes))
            pts = [per for _ in classes]

        sample = wc.stratifiedSample(
            numPoints=min_pc, classBand="Map", region=region,
            scale=scale, seed=42, geometries=True,
            classValues=classes, classPoints=pts)
        feats = sample.getInfo().get("features", [])
        if not feats:
            raise RuntimeError("GEE returned no sample points.")

        out = QgsVectorLayer("Point?crs=EPSG:4326", "tmp", "memory")
        pr = out.dataProvider()
        pr.addAttributes([QgsField("wc_code", QVariant.Int),
                          QgsField("wc_hint", QVariant.String)])
        out.updateFields()
        qfeats = []
        for f in feats:
            x, y = f["geometry"]["coordinates"]
            code = int(f["properties"].get("Map", 0))
            qf = QgsFeature(out.fields())
            qf.setGeometry(QgsGeometry.fromPointXY(QgsPointXY(x, y)))
            qf["wc_code"] = code
            qf["wc_hint"] = gu.WC_LABELS.get(code, str(code))
            qfeats.append(qf)
        random.shuffle(qfeats)
        pr.addFeatures(qfeats)
        return out

    def _local_raster_stratified(self, n_target):
        import processing
        ref = self.cb_ref.currentLayer()
        if ref is None:
            raise RuntimeError("Choose a reference raster.")
        if ref.providerType() != "gdal":
            raise RuntimeError(
                "'{}' is a tile/streaming layer (provider: {}).\n"
                "QGIS cannot read pixel values from GEE/XYZ layers.\n"
                "Use the WorldCover (GEE) method instead, or download a "
                "real GeoTIFF first.".format(ref.name(), ref.providerType()))

        lyr, _ = self._aoi_layer_and_geom()
        src = lyr
        if self.chk_sel.isChecked() and lyr.selectedFeatureCount() > 0:
            src = processing.run("native:saveselectedfeatures",
                                 {"INPUT": lyr, "OUTPUT": "memory:"})["OUTPUT"]
        pts = processing.run("native:randompointsinpolygons", {
            "INPUT": src, "POINTS_NUMBER": n_target * 5,
            "MIN_DISTANCE": self.sp_mindist.value(),
            "OUTPUT": "memory:"})["OUTPUT"]
        # reproject points to the raster CRS before sampling
        if pts.crs() != ref.crs():
            pts = processing.run("native:reprojectlayer", {
                "INPUT": pts, "TARGET_CRS": ref.crs(),
                "OUTPUT": "memory:"})["OUTPUT"]
        pts = processing.run("native:rastersampling", {
            "INPUT": pts, "RASTERCOPY": ref,
            "COLUMN_PREFIX": "ref_", "OUTPUT": "memory:"})["OUTPUT"]

        groups = {}
        for f in pts.getFeatures():
            v = f["ref_1"]
            try:
                fv = float(v)
                if not math.isfinite(fv):   # NoData is often inf / nan
                    continue
                v = int(fv)
            except (TypeError, ValueError):
                continue
            groups.setdefault(v, []).append(f)
        total = sum(len(v) for v in groups.values())
        if total == 0:
            raise RuntimeError(
                "No raster values found under the points. Check that the "
                "raster really covers the AOI and has valid pixel values.")
        chosen = []
        min_pc = self.sp_minclass.value()
        for cls, feats in groups.items():
            share = max(min_pc, int(round(n_target * len(feats) / total)))
            random.shuffle(feats)
            chosen.extend(feats[:share])
        random.shuffle(chosen)
        out = QgsVectorLayer("Point?crs=" + pts.crs().authid(),
                             "tmp", "memory")
        out.dataProvider().addAttributes(pts.fields())
        out.updateFields()
        out.dataProvider().addFeatures(chosen)
        return out

    # ------------------------------------------------------------ finalize
    def _finalize(self, pts):
        pr = pts.dataProvider()
        existing = [f.name() for f in pts.fields()]
        new = [QgsField(n, t) for n, t in (
            ("point_id", QVariant.Int), ("class_code", QVariant.Int),
            ("class_name", QVariant.String), ("lu_code", QVariant.Int),
            ("lu_name", QVariant.String), ("use", QVariant.String))
            if n not in existing]
        if new:
            pr.addAttributes(new)
            pts.updateFields()
        train_frac = self.sp_train.value() / 100.0
        idx = {n: pts.fields().indexOf(n)
               for n in ("point_id", "class_code", "class_name",
                         "lu_code", "lu_name", "use")}
        changes = {}
        for i, f in enumerate(pts.getFeatures(), start=1):
            changes[f.id()] = {
                idx["point_id"]: i, idx["class_code"]: 0,
                idx["class_name"]: "", idx["lu_code"]: 0, idx["lu_name"]: "",
                idx["use"]: "train" if _train_random.random() < train_frac
                else "valid"}
        pr.changeAttributeValues(changes)
        pts.setName("sample_points")

        scheme = self.state.get("scheme") or "lulc"
        cats = [QgsRendererCategory(0, self._marker("#888888"), "Unlabeled")]
        for code, name, color in gu.scheme_classes(scheme):
            cats.append(QgsRendererCategory(code, self._marker(color), name))
        pts.setRenderer(QgsCategorizedSymbolRenderer("class_code", cats))
        return pts

    @staticmethod
    def _marker(color):
        s = QgsMarkerSymbol.createSimple(
            {"name": "circle", "size": "2.4",
             "outline_color": "black", "outline_width": "0.2"})
        s.setColor(QColor(color))
        return s
