# -*- coding: utf-8 -*-
from .geomaticape_algorithm import GeomaticapeAlgorithm
"""
Estadistica zonal de raster (mono o multiespectral)
====================================================
Calcula estadisticas zonales (count, min, max, mean, median, sum, std,
range, majority, minority) de los pixeles de un raster que caen dentro
de cada poligono.

Formatos de salida:
  0 - Excel (.xlsx)          : tabla sin geometria
  1 - CSV   (.csv)           : tabla sin geometria
  2 - Shapefile (.shp)       : capa vectorial con estadisticas como atributos
  3 - GeoPackage (.gpkg)     : capa vectorial con estadisticas como atributos

Autor : Geomatica Ambiental - https://www.geomatica.pe
Plugin: GeoRemote Tools v1.15
Grupo : Geoprocesamiento
"""

import os
import csv
import math

import numpy as np

from qgis.core import (
    QgsProcessingParameterVectorLayer,
    QgsProcessingParameterRasterLayer,
    QgsProcessingParameterBand,
    QgsProcessingParameterEnum,
    QgsProcessingParameterField,
    QgsProcessingParameterFileDestination,
    QgsProcessingParameterFeatureSink,
    QgsProcessingException,
    QgsProcessing,
    QgsFeatureSink,
    QgsFeature,
    QgsFields,
    QgsField,
    QgsWkbTypes,
    QgsVectorFileWriter,
    QgsProject,
)
from qgis.PyQt.QtCore import QVariant
from qgis import processing
from osgeo import gdal, ogr, osr


# ---------------------------------------------------------------------------
# Constantes
# ---------------------------------------------------------------------------
STATS_LIST = [
    "count", "min", "max", "mean", "median",
    "sum", "std", "range", "majority", "minority"
]

OUT_FORMAT_OPTIONS = [
    "Excel (.xlsx)",
    "CSV (.csv)",
    "Shapefile (.shp)",
    "GeoPackage (.gpkg)",
]

FMT_XLSX = 0
FMT_CSV  = 1
FMT_SHP  = 2
FMT_GPKG = 3


# ---------------------------------------------------------------------------
# Calculo de estadisticas
# ---------------------------------------------------------------------------
def _calc_stats(arr_valid, stats_seleccionadas):
    """Devuelve dict {stat: valor} sobre arr_valid (1-D, sin NaN/NoData)."""
    out = {}
    if arr_valid.size == 0:
        for s in stats_seleccionadas:
            out[s] = None
        return out
    for s in stats_seleccionadas:
        if   s == "count":   out[s] = int(arr_valid.size)
        elif s == "min":     out[s] = float(np.min(arr_valid))
        elif s == "max":     out[s] = float(np.max(arr_valid))
        elif s == "mean":    out[s] = float(np.mean(arr_valid))
        elif s == "median":  out[s] = float(np.median(arr_valid))
        elif s == "sum":     out[s] = float(np.sum(arr_valid))
        elif s == "std":     out[s] = float(np.std(arr_valid))
        elif s == "range":   out[s] = float(np.max(arr_valid) - np.min(arr_valid))
        elif s in ("majority", "minority"):
            try:
                vals_int = np.round(arr_valid).astype(np.int64)
                uniq, cnt = np.unique(vals_int, return_counts=True)
                if uniq.size == 0:
                    out[s] = None
                elif s == "majority":
                    out[s] = int(uniq[np.argmax(cnt)])
                else:
                    out[s] = int(uniq[np.argmin(cnt)])
            except Exception:
                out[s] = None
    return out


# ---------------------------------------------------------------------------
# Clase principal
# ---------------------------------------------------------------------------
class EstadisticaZonalRaster(GeomaticapeAlgorithm):
    _algorithm_name = "estadistica_zonal_raster"
    _icon_name      = "zonal_raster.png"

    INPUT_VECTOR = "INPUT_VECTOR"
    INPUT_RASTER = "INPUT_RASTER"
    BAND         = "BAND"
    ID_FIELD     = "ID_FIELD"
    STATS        = "STATS"
    OUT_FORMAT   = "OUT_FORMAT"
    OUTPUT       = "OUTPUT"

    def displayName(self):
        return self.tr("Estadistica zonal raster (Excel/CSV/SHP/GPKG)")

    def group(self):
        return self.tr("Geoprocesamiento")

    def groupId(self):
        return "geomaticape_geoprocesamiento"

    def shortHelpString(self):
        return """
<h3>Estadistica zonal raster</h3>
<b>Autor:</b> Geomatica Ambiental<br>
<b>Plugin:</b> GeoRemote Tools<br>
<b>Version:</b> 1.15<br><br>

<b>Descripcion:</b><br>
Calcula estadisticas de los pixeles del raster que caen dentro de cada
poligono. Soporta raster mono o multiespectral (selecciona la banda).<br><br>

<b>Estadisticas disponibles:</b>
<ul>
  <li><b>count</b>   — numero de pixeles validos</li>
  <li><b>min / max</b> — valores extremos</li>
  <li><b>mean</b>    — promedio</li>
  <li><b>median</b>  — mediana</li>
  <li><b>sum</b>     — suma total</li>
  <li><b>std</b>     — desviacion estandar</li>
  <li><b>range</b>   — max - min</li>
  <li><b>majority</b> — valor mas frecuente (entero)</li>
  <li><b>minority</b> — valor menos frecuente (entero)</li>
</ul>

<b>Formatos de salida:</b>
<ul>
  <li><b>Excel (.xlsx)</b>    — tabla sin geometria</li>
  <li><b>CSV (.csv)</b>       — tabla sin geometria</li>
  <li><b>Shapefile (.shp)</b> — capa vectorial con estadisticas como atributos</li>
  <li><b>GeoPackage (.gpkg)</b> — capa vectorial con estadisticas como atributos</li>
</ul>

Los formatos <b>Shapefile y GeoPackage</b> conservan la geometria original
de los poligonos junto con todas las estadisticas calculadas, permitiendo
visualizacion directa en QGIS.<br><br>

<b>Web:</b> https://www.geomatica.pe/
"""

    def initAlgorithm(self, config=None):
        self.addParameter(QgsProcessingParameterVectorLayer(
            self.INPUT_VECTOR,
            self.tr("Capa de poligonos"),
            types=[QgsProcessing.SourceType.TypeVectorPolygon]
        ))
        self.addParameter(QgsProcessingParameterRasterLayer(
            self.INPUT_RASTER,
            self.tr("Capa raster (mono o multiespectral)")
        ))
        self.addParameter(QgsProcessingParameterBand(
            self.BAND,
            self.tr("Banda a analizar (si es multiespectral)"),
            parentLayerParameterName=self.INPUT_RASTER,
            defaultValue=1,
            optional=False
        ))
        self.addParameter(QgsProcessingParameterField(
            self.ID_FIELD,
            self.tr("Campo identificador del poligono (opcional)"),
            parentLayerParameterName=self.INPUT_VECTOR,
            optional=True
        ))
        self.addParameter(QgsProcessingParameterEnum(
            self.STATS,
            self.tr("Estadisticas a calcular (todas por defecto)"),
            options=STATS_LIST,
            allowMultiple=True,
            defaultValue=list(range(len(STATS_LIST)))
        ))
        self.addParameter(QgsProcessingParameterEnum(
            self.OUT_FORMAT,
            self.tr("Formato de salida"),
            options=OUT_FORMAT_OPTIONS,
            defaultValue=FMT_XLSX,
            allowMultiple=False
        ))
        self.addParameter(QgsProcessingParameterFileDestination(
            self.OUTPUT,
            self.tr("Archivo de resultados"),
            fileFilter=(
                "Excel (*.xlsx);;"
                "CSV (*.csv);;"
                "Shapefile (*.shp);;"
                "GeoPackage (*.gpkg)"
            )
        ))

    # -----------------------------------------------------------------------
    def processAlgorithm(self, parameters, context, feedback):
        v_layer  = self.parameterAsVectorLayer(parameters, self.INPUT_VECTOR, context)
        r_layer  = self.parameterAsRasterLayer(parameters, self.INPUT_RASTER, context)
        band_n   = self.parameterAsInt(parameters, self.BAND, context)
        id_field = self.parameterAsString(parameters, self.ID_FIELD, context)
        stats_idx = self.parameterAsEnums(parameters, self.STATS, context)
        out_fmt  = self.parameterAsEnum(parameters, self.OUT_FORMAT, context)
        out_path = self.parameterAsFileOutput(parameters, self.OUTPUT, context)

        if v_layer is None:
            raise QgsProcessingException("No se pudo cargar la capa de poligonos.")
        if r_layer is None:
            raise QgsProcessingException("No se pudo cargar la capa raster.")

        stats_seleccionadas = (
            [STATS_LIST[i] for i in stats_idx] if stats_idx else STATS_LIST[:]
        )

        # Ajustar extension de salida segun formato seleccionado
        ext_map = {FMT_XLSX: ".xlsx", FMT_CSV: ".csv",
                   FMT_SHP: ".shp", FMT_GPKG: ".gpkg"}
        ext_req = ext_map[out_fmt]
        if not out_path.lower().endswith(ext_req):
            out_path = os.path.splitext(out_path)[0] + ext_req

        feedback.pushInfo("=" * 55)
        feedback.pushInfo("Estadistica zonal raster  [GeoRemote Tools]")
        feedback.pushInfo(f"  Vector  : {v_layer.name()} ({v_layer.featureCount()} poligonos)")
        feedback.pushInfo(f"  Raster  : {r_layer.name()} (banda {band_n})")
        feedback.pushInfo(f"  Stats   : {', '.join(stats_seleccionadas)}")
        feedback.pushInfo(f"  Formato : {OUT_FORMAT_OPTIONS[out_fmt]}")
        feedback.pushInfo(f"  Salida  : {out_path}")
        feedback.pushInfo("=" * 55)

        # --- Abrir raster con GDAL ---
        raster_path = r_layer.source()
        ds_r = gdal.Open(raster_path, gdal.GA_ReadOnly)
        if ds_r is None:
            raise QgsProcessingException("GDAL no pudo abrir el raster.")
        if band_n < 1 or band_n > ds_r.RasterCount:
            raise QgsProcessingException(
                f"Banda {band_n} fuera de rango (raster tiene {ds_r.RasterCount}).")
        gt    = ds_r.GetGeoTransform()
        proj  = ds_r.GetProjection()
        rcols = ds_r.RasterXSize
        rrows = ds_r.RasterYSize
        band  = ds_r.GetRasterBand(band_n)
        nodata = band.GetNoDataValue()

        # --- Reproyectar vector si CRS difiere ---
        crs_vec = v_layer.crs()
        srs_r   = osr.SpatialReference()
        if proj:
            srs_r.ImportFromWkt(proj)
        crs_r_authid = None
        if srs_r.GetAuthorityCode(None):
            crs_r_authid = f"{srs_r.GetAuthorityName(None)}:{srs_r.GetAuthorityCode(None)}"

        vector_to_use = v_layer
        if crs_r_authid and crs_vec.authid() != crs_r_authid:
            feedback.pushInfo(f"Reproyectando vector {crs_vec.authid()} -> {crs_r_authid}")
            res = processing.run(
                "native:reprojectlayer",
                {"INPUT": v_layer, "TARGET_CRS": crs_r_authid, "OUTPUT": "memory:"},
                context=context, feedback=feedback
            )
            vector_to_use = res["OUTPUT"]

        # --- Iterar features ---
        feats  = list(vector_to_use.getFeatures())
        total  = len(feats)
        registros = []   # lista de (feat_id, attr_dict, geom_wkt, orig_attrs)
        fname  = id_field if id_field else ""

        for i, feat in enumerate(feats):
            geom = feat.geometry()
            if geom is None or geom.isEmpty():
                continue
            bbox = geom.boundingBox()

            inv_gt = gdal.InvGeoTransform(gt)
            if inv_gt is None:
                raise QgsProcessingException("GeoTransform invalido del raster.")
            px_min, py_min = gdal.ApplyGeoTransform(inv_gt, bbox.xMinimum(), bbox.yMaximum())
            px_max, py_max = gdal.ApplyGeoTransform(inv_gt, bbox.xMaximum(), bbox.yMinimum())

            xoff  = max(0,     int(math.floor(min(px_min, px_max))))
            yoff  = max(0,     int(math.floor(min(py_min, py_max))))
            xend  = min(rcols, int(math.ceil(max(px_min, px_max))))
            yend  = min(rrows, int(math.ceil(max(py_min, py_max))))
            xsize = xend - xoff
            ysize = yend - yoff
            if xsize <= 0 or ysize <= 0:
                continue

            sub_minx = gt[0] + xoff * gt[1] + yoff * gt[2]
            sub_maxy = gt[3] + xoff * gt[4] + yoff * gt[5]
            sub_gt   = (sub_minx, gt[1], gt[2], sub_maxy, gt[4], gt[5])

            arr = band.ReadAsArray(xoff, yoff, xsize, ysize)
            if arr is None:
                continue
            arr = arr.astype(np.float64)

            # Rasterizar geometria como mascara
            drv   = ogr.GetDriverByName("Memory")
            ds_v  = drv.CreateDataSource("mem")
            srs   = osr.SpatialReference()
            if proj:
                srs.ImportFromWkt(proj)
            lyr   = ds_v.CreateLayer("p", srs=srs, geom_type=ogr.wkbPolygon)
            f_ogr = ogr.Feature(lyr.GetLayerDefn())
            wkb   = bytes(geom.asWkb())
            f_ogr.SetGeometry(ogr.CreateGeometryFromWkb(wkb))
            lyr.CreateFeature(f_ogr)
            f_ogr = None

            mask_ds = gdal.GetDriverByName("MEM").Create("", xsize, ysize, 1, gdal.GDT_Byte)
            mask_ds.SetGeoTransform(sub_gt)
            if proj:
                mask_ds.SetProjection(proj)
            gdal.RasterizeLayer(mask_ds, [1], lyr, burn_values=[1])
            mask    = mask_ds.GetRasterBand(1).ReadAsArray()
            mask_ds = None
            ds_v    = None

            valid_mask = (mask == 1)
            if nodata is not None:
                valid_mask &= (arr != nodata)
            valid_mask &= np.isfinite(arr)

            arr_valid = arr[valid_mask]
            stat_vals = _calc_stats(arr_valid, stats_seleccionadas)

            # Construir fila de atributos
            row = {}
            row["fid"] = feat.id()
            if fname:
                try:
                    row[fname] = feat[fname]
                except Exception:
                    row[fname] = None
            for s in stats_seleccionadas:
                row[s] = stat_vals.get(s)

            # Guardar geometria (WKT) y atributos originales para formatos vectoriales
            registros.append({
                "row":       row,
                "geom_wkt":  geom.asWkt(),
                "orig_attrs": {f.name(): feat[f.name()]
                               for f in vector_to_use.fields()},
            })

            feedback.setProgress(int((i + 1) * 100 / max(total, 1)))
            if feedback.isCanceled():
                break

        ds_r = None

        if not registros:
            raise QgsProcessingException(
                "No se generaron registros (revisa CRS y extension).")

        cols = list(registros[0]["row"].keys())
        for rec in registros:
            for c in cols:
                rec["row"].setdefault(c, None)

        # --- Escribir salida ---
        if out_fmt == FMT_XLSX:
            self._escribir_excel(out_path, [r["row"] for r in registros], cols)
        elif out_fmt == FMT_CSV:
            self._escribir_csv(out_path, [r["row"] for r in registros], cols)
        elif out_fmt == FMT_SHP:
            self._escribir_vector(
                out_path, registros, cols,
                vector_to_use, stats_seleccionadas, fname, "ESRI Shapefile", feedback
            )
        elif out_fmt == FMT_GPKG:
            self._escribir_vector(
                out_path, registros, cols,
                vector_to_use, stats_seleccionadas, fname, "GPKG", feedback
            )

        feedback.pushInfo(f"OK - {len(registros)} registros escritos en: {out_path}")
        return {self.OUTPUT: out_path}

    # -----------------------------------------------------------------------
    # Escritores
    # -----------------------------------------------------------------------

    def _escribir_excel(self, ruta, registros, cols):
        try:
            from openpyxl import Workbook
            from openpyxl.styles import Font, PatternFill, Alignment
            from openpyxl.utils import get_column_letter
        except ImportError:
            raise QgsProcessingException(
                "Para exportar a Excel instala openpyxl:\n"
                "  python -m pip install openpyxl"
            )
        wb = Workbook()
        ws = wb.active
        ws.title = "zonal_stats"

        # Encabezado con estilo
        header_fill = PatternFill("solid", fgColor="2E75B6")
        header_font = Font(bold=True, color="FFFFFF")
        ws.append(cols)
        for cell in ws[1]:
            cell.fill = header_fill
            cell.font = header_font
            cell.alignment = Alignment(horizontal="center")

        for r in registros:
            ws.append([r.get(c) for c in cols])

        # Ajustar ancho de columnas
        for col_idx, col_name in enumerate(cols, 1):
            max_len = max(len(str(col_name)),
                          max((len(str(r.get(col_name, "") or "")) for r in registros), default=0))
            ws.column_dimensions[get_column_letter(col_idx)].width = min(max_len + 4, 30)

        wb.save(ruta)

    def _escribir_csv(self, ruta, registros, cols):
        with open(ruta, "w", newline="", encoding="utf-8") as f:
            w = csv.DictWriter(f, fieldnames=cols)
            w.writeheader()
            for r in registros:
                w.writerow(r)

    def _escribir_vector(self, ruta, registros, stat_cols,
                         source_layer, stats_seleccionadas, fname, driver_name, feedback):
        """
        Escribe un archivo vectorial (SHP o GPKG) con la geometria original
        de los poligonos y las estadisticas como atributos adicionales.
        """
        from qgis.core import QgsGeometry

        # ── 1. Eliminar archivo preexistente ─────────────────────────────────
        # QgsVectorFileWriter falla si el archivo ya existe
        if os.path.exists(ruta):
            try:
                os.remove(ruta)
                feedback.pushInfo("Archivo existente eliminado: " + ruta)
            except Exception as e:
                feedback.pushWarning("No se pudo eliminar archivo existente: " + str(e))

        # Para Shapefile eliminar archivos auxiliares (.dbf, .shx, .prj, etc.)
        if ruta.lower().endswith(".shp"):
            base = ruta[:-4]
            for ext in (".dbf", ".shx", ".prj", ".cpg", ".qpj", ".qix", ".sbn", ".sbx"):
                sidecar = base + ext
                if os.path.exists(sidecar):
                    try:
                        os.remove(sidecar)
                    except Exception as e:
                        feedback.pushWarning(f"No se pudo eliminar {sidecar}: {e}")

        # ── 2. Definir campos de salida ──────────────────────────────────────
        orig_fields = source_layer.fields()
        out_fields  = QgsFields()

        for f in orig_fields:
            out_fields.append(QgsField(f.name(), f.type()))

        int_stats = {"count", "majority", "minority"}
        for s in stats_seleccionadas:
            vtype = QVariant.LongLong if s in int_stats else QVariant.Double
            out_fields.append(QgsField("zs_" + s, vtype))

        crs           = source_layer.crs()
        transform_ctx = QgsProject.instance().transformContext()
        geom_type     = QgsWkbTypes.Type.MultiPolygon

        # ── 3. Crear el writer ───────────────────────────────────────────────
        writer = None

        # Metodo A: QgsVectorFileWriter.create() — disponible en QGIS 3.x
        # Puede devolver el writer directamente O una tupla (writer, error_msg)
        if hasattr(QgsVectorFileWriter, "create"):
            try:
                opts              = QgsVectorFileWriter.SaveVectorOptions()
                opts.driverName   = driver_name
                opts.fileEncoding = "UTF-8"
                result = QgsVectorFileWriter.create(
                    ruta, out_fields, geom_type, crs, transform_ctx, opts
                )
                if isinstance(result, tuple):
                    writer    = result[0]
                    error_msg = result[1] if len(result) > 1 else ""
                    if error_msg:
                        feedback.pushWarning("create(): " + error_msg)
                else:
                    # Devuelve el writer directamente
                    writer = result
                    error_msg = ""

                # Verificar que el writer es valido
                if writer is not None and hasattr(writer, "hasError") and writer.hasError():
                    feedback.pushWarning(
                        "create() writer con error: " +
                        (writer.errorMessage() if hasattr(writer, "errorMessage") else "?")
                    )
                    writer = None
            except Exception as e:
                feedback.pushWarning("create() fallo: " + str(e))
                writer = None

        # Metodo B: Constructor legacy — funciona en todas las versiones QGIS 3.x
        if writer is None:
            try:
                writer = QgsVectorFileWriter(
                    ruta, "UTF-8", out_fields, geom_type, crs, driver_name
                )
                if hasattr(writer, "hasError") and writer.hasError():
                    err = writer.errorMessage() if hasattr(writer, "errorMessage") else "error"
                    raise QgsProcessingException(
                        "QgsVectorFileWriter (constructor): " + err
                    )
            except QgsProcessingException:
                raise
            except Exception as e:
                raise QgsProcessingException(
                    "No se pudo crear el archivo vectorial (" + driver_name + "): " + str(e)
                )

        if writer is None:
            raise QgsProcessingException(
                "No se pudo crear el archivo vectorial: " + driver_name
            )

        # ── 4. Escribir features ─────────────────────────────────────────────
        for rec in registros:
            row      = rec["row"]
            geom_wkt = rec["geom_wkt"]
            orig     = rec["orig_attrs"]

            feat = QgsFeature(out_fields)
            feat.setGeometry(QgsGeometry.fromWkt(geom_wkt))

            for f in orig_fields:
                feat.setAttribute(f.name(), orig.get(f.name()))

            for s in stats_seleccionadas:
                feat.setAttribute("zs_" + s, row.get(s))

            writer.addFeature(feat)

        del writer
        feedback.pushInfo(
            "Capa vectorial creada (" + driver_name + "): " + ruta
        )

    def run(self):
        processing.execAlgorithmDialog(self)
