# -*- coding: utf-8 -*-
"""
RiverBankErosionAccretion_V2 - QGIS plugin for quantifying riverbank
erosion and accretion.

Workflow:
  1) Load two waterbody polygon layers (Year-1, Year-2)
  2) Select river morphological type → algorithm adapts automatically
  3) Optional: Riparian buffer masking, MMU filter, braiding index,
     connected component labelling, regime classification
  4) Generate analysis grid and compute per-cell change
  5) Aggregate results by river stretches + Top-N ranking
  6) Optional transect-based rate-of-change analysis
  7) Export results as shapefiles and CSV reports
"""

import os

import numpy as np
import geopandas as gpd

from qgis.PyQt.QtWidgets import (
    QAction, QDialog, QLabel, QLineEdit, QPushButton,
    QGridLayout, QFileDialog, QMessageBox, QComboBox,
    QGroupBox, QCheckBox, QSpinBox, QDoubleSpinBox,
    QVBoxLayout, QHBoxLayout, QFormLayout,
)
from qgis.PyQt.QtGui import QIcon
from qgis.core import QgsCoordinateReferenceSystem
from qgis.gui import QgsMapLayerComboBox, QgsProjectionSelectionWidget

from . import change_analysis

PLUGIN_MENU = "&RiverBank Erosion and Accretion Analyzer (v2)"
CELL_SIZE_M = 30.0
TOP_N = 15

RIVER_TYPES = [
    "Meandering",
    "Braided",
    "Deltaic / Distributary",
    "Straight",
    "Anastomosing",
]


class RiverBankErosionAccretion_V2:

    def __init__(self, iface):
        self.iface = iface
        self.plugin_dir = os.path.dirname(__file__)
        self.action = None
        self.dlg = None
        self.stretch_m = 3000
        self.cell_size_m = CELL_SIZE_M
        self.top_n = TOP_N
        self.buffer_m = 500
        self.mmu_cells = 3
        self.min_polygon_area_ha = 5
        self.years_interval = 20

    def initGui(self):
        icon_path = os.path.join(self.plugin_dir, "icon.png")
        icon = QIcon(icon_path) if os.path.exists(icon_path) else QIcon()
        title = "RiverBank Erosion and Accretion Analyzer (v2)"
        self.action = QAction(icon, title, self.iface.mainWindow())
        self.action.triggered.connect(self.run)
        self.iface.addToolBarIcon(self.action)
        self.iface.addPluginToMenu(PLUGIN_MENU, self.action)

    def unload(self):
        if self.action:
            self.iface.removeToolBarIcon(self.action)
            self.iface.removePluginMenu(PLUGIN_MENU, self.action)

    # ---------------- Dialog builder ----------------
    def build_ui(self):
        self.dlg = QDialog(self.iface.mainWindow())
        self.dlg.setWindowTitle(
            "RiverBank Erosion & Accretion Analyzer (v2)"
        )
        self.dlg.resize(520, 680)

        main_layout = QVBoxLayout(self.dlg)

        self._build_input_section()
        main_layout.addWidget(self.input_group)

        self._build_river_section()
        main_layout.addWidget(self.river_group)

        self._build_output_section()
        main_layout.addWidget(self.output_group)

        self._build_analysis_params()
        main_layout.addWidget(self.params_group)

        main_layout.addLayout(self._build_button_layout())

    def _build_input_section(self):
        self.input_group = QGroupBox("Input Layers")
        layout = QFormLayout(self.input_group)
        self.comboYear1 = QgsMapLayerComboBox()
        self.comboYear2 = QgsMapLayerComboBox()
        for cb in (self.comboYear1, self.comboYear2):
            try:
                cb.setFilters(cb.GeometryType.PolygonLayer)
            except Exception:
                pass
        layout.addRow("Year 1 Waterbody:", self.comboYear1)
        layout.addRow("Year 2 Waterbody:", self.comboYear2)

    def _build_river_section(self):
        self.river_group = QGroupBox("River Settings")
        layout = QFormLayout(self.river_group)

        self.comboRiverType = QComboBox()
        self.comboRiverType.addItems(RIVER_TYPES)
        self.comboRiverType.currentIndexChanged.connect(
            self._on_river_type_changed
        )
        layout.addRow("River morphological type:", self.comboRiverType)

        self.spinBufferDist = QDoubleSpinBox()
        self.spinBufferDist.setRange(100, 5000)
        self.spinBufferDist.setValue(self.buffer_m)
        self.spinBufferDist.setSuffix(" m")
        self.spinBufferDist.setToolTip(
            "Buffer distance for riparian corridor mask. "
            "Cells outside this corridor are excluded."
        )
        layout.addRow("Riparian buffer distance:", self.spinBufferDist)

        self.spinYearsInterval = QSpinBox()
        self.spinYearsInterval.setRange(1, 100)
        self.spinYearsInterval.setValue(self.years_interval)
        self.spinYearsInterval.setSuffix(" years")
        self.spinYearsInterval.setToolTip(
            "Number of years between Year 1 and Year 2."
        )
        layout.addRow("Years between epochs:", self.spinYearsInterval)

        self.spinMinPolygonArea = QDoubleSpinBox()
        self.spinMinPolygonArea.setRange(0, 100)
        self.spinMinPolygonArea.setValue(self.min_polygon_area_ha)
        self.spinMinPolygonArea.setSuffix(" ha")
        self.spinMinPolygonArea.setToolTip(
            "Minimum polygon area (ha) for deltaic connected "
            "component labelling. Features below this are excluded."
        )
        layout.addRow("Min. polygon area (CCL):", self.spinMinPolygonArea)

    def _on_river_type_changed(self, idx):
        rtype = RIVER_TYPES[idx]
        if rtype == "Braided":
            self.spinMinPolygonArea.setEnabled(False)
            self.spinBufferDist.setToolTip(
                "Buffer distance for braided rivers. "
                "Use larger values (>1000m) for wide braidplains."
            )
            self.spinBufferDist.setValue(1000)
        elif rtype == "Deltaic / Distributary":
            self.spinMinPolygonArea.setEnabled(True)
            self.spinMinPolygonArea.setToolTip(
                "Features below this area are excluded as "
                "non-riverine (agricultural ponds, tidal creeks)."
            )
            self.spinBufferDist.setValue(500)
        else:
            self.spinMinPolygonArea.setEnabled(False)
            self.spinBufferDist.setValue(500)

    def _build_output_section(self):
        self.output_group = QGroupBox("Output Settings")
        layout = QGridLayout(self.output_group)

        layout.addWidget(QLabel("Projected CRS:"), 0, 0)
        self.projTarget = QgsProjectionSelectionWidget()
        self.projTarget.setToolTip(
            "Select a projected CRS (equal-area preferred)."
        )
        layout.addWidget(self.projTarget, 0, 1)

        layout.addWidget(QLabel("Output folder:"), 1, 0)
        self.editOutDir = QLineEdit()
        self.editOutDir.setReadOnly(True)
        layout.addWidget(self.editOutDir, 1, 1)

        self.btnBrowseOut = QPushButton("Browse\u2026")
        layout.addWidget(self.btnBrowseOut, 2, 1)
        self.btnBrowseOut.clicked.connect(self._pick_outdir)

    def _build_analysis_params(self):
        self.params_group = QGroupBox("Analysis Parameters")
        layout = QGridLayout(self.params_group)

        layout.addWidget(QLabel("Grid cell size (m):"), 0, 0)
        self.spinCellSize = QDoubleSpinBox()
        self.spinCellSize.setRange(5, 500)
        self.spinCellSize.setValue(self.cell_size_m)
        self.spinCellSize.setSuffix(" m")
        layout.addWidget(self.spinCellSize, 0, 1)

        layout.addWidget(QLabel("Stretch distance:"), 1, 0)
        self.comboStretch = QComboBox()
        self.comboStretch.addItems(
            ["90 m", "3 km", "6 km", "9 km"]
        )
        layout.addWidget(self.comboStretch, 1, 1)

        layout.addWidget(QLabel("Top-N stretches:"), 2, 0)
        self.spinTopN = QSpinBox()
        self.spinTopN.setRange(5, 100)
        self.spinTopN.setValue(self.top_n)
        layout.addWidget(self.spinTopN, 2, 1)

        layout.addWidget(QLabel("MMU threshold (cells):"), 3, 0)
        self.spinMMU = QSpinBox()
        self.spinMMU.setRange(0, 20)
        self.spinMMU.setValue(self.mmu_cells)
        self.spinMMU.setSuffix(" cells")
        self.spinMMU.setToolTip(
            "Minimum contiguous cells for erosion/accretion patches. "
            "Patches smaller than this are removed. 0 = disabled."
        )
        layout.addWidget(self.spinMMU, 3, 1)

        self.chkTransect = QCheckBox(
            "Enable transect-based analysis (requires centerline)"
        )
        layout.addWidget(self.chkTransect, 4, 0, 1, 2)

        layout.addWidget(QLabel("Transect spacing (m):"), 5, 0)
        self.spinTransectSpacing = QDoubleSpinBox()
        self.spinTransectSpacing.setRange(20, 1000)
        self.spinTransectSpacing.setValue(100)
        self.spinTransectSpacing.setSuffix(" m")
        self.spinTransectSpacing.setEnabled(False)
        layout.addWidget(self.spinTransectSpacing, 5, 1)

        self.chkTransect.toggled.connect(
            self.spinTransectSpacing.setEnabled
        )

    def _build_button_layout(self):
        self.btnStart = QPushButton("Start Analysis")
        self.btnStart.setMinimumHeight(40)
        layout = QHBoxLayout()
        layout.addStretch()
        layout.addWidget(self.btnStart)
        self.btnStart.clicked.connect(self._start)
        return layout

    def run(self):
        if self.dlg is None:
            self.build_ui()
        self.dlg.show()

    # ---------------- UI Actions ----------------

    def _pick_outdir(self):
        folder = QFileDialog.getExistingDirectory(
            self.iface.mainWindow(), "Select Output Folder"
        )
        if folder:
            self.editOutDir.setText(folder)

    def _start(self):
        outdir = self.editOutDir.text().strip()
        if not outdir:
            QMessageBox.warning(
                self.dlg, "Missing Output Folder",
                "Please choose an output folder."
            )
            return

        crs = self.projTarget.crs()
        if not crs.isValid():
            crs = QgsCoordinateReferenceSystem("EPSG:6933")
        target_crs = crs.authid()

        self.cell_size_m = self.spinCellSize.value()
        self.top_n = self.spinTopN.value()
        self.buffer_m = self.spinBufferDist.value()
        self.mmu_cells = self.spinMMU.value()
        self.min_polygon_area_ha = self.spinMinPolygonArea.value()
        self.years_interval = self.spinYearsInterval.value()

        label = self.comboStretch.currentText().strip().lower()
        self.stretch_m = {
            "90 m": 90, "3 km": 3000,
            "6 km": 6000, "9 km": 9000,
        }.get(label, 3000)

        river_type = RIVER_TYPES[self.comboRiverType.currentIndex()]

        try:
            self._run_polygon_pipeline(outdir, target_crs, river_type)
            QMessageBox.information(
                self.dlg, "Finished",
                f"Processing complete.\nResults saved to:\n{outdir}"
            )
        except Exception as e:
            import traceback
            QMessageBox.critical(
                self.dlg, "Error",
                f"Processing failed:\n{e}\n\n"
                f"{traceback.format_exc()}"
            )

    def _run_polygon_pipeline(self, out_dir, target_crs, river_type):
        lyr1 = self.comboYear1.currentLayer()
        lyr2 = self.comboYear2.currentLayer()

        if not lyr1 or not lyr2:
            raise RuntimeError(
                "Please select both Year-1 and Year-2 polygon layers."
            )

        y1_gdf = gpd.read_file(lyr1.source())
        y2_gdf = gpd.read_file(lyr2.source())

        if target_crs:
            y1_gdf = y1_gdf.to_crs(target_crs)
            y2_gdf = y2_gdf.to_crs(target_crs)

        y1_gdf.to_file(os.path.join(out_dir, "Waterbody_Y1.shp"))
        y2_gdf.to_file(os.path.join(out_dir, "Waterbody_Y2.shp"))

        # Connected component labeling for deltaic regime
        if river_type == "Deltaic / Distributary" and self.min_polygon_area_ha > 0:
            y1_gdf = change_analysis.connected_component_label(
                y1_gdf, min_area_ha=self.min_polygon_area_ha
            )
            y2_gdf = change_analysis.connected_component_label(
                y2_gdf, min_area_ha=self.min_polygon_area_ha
            )
            if 'channel_id' in y1_gdf.columns and 'channel_id' in y2_gdf.columns:
                y1_gdf.to_file(os.path.join(out_dir, "Waterbody_Y1_labelled.shp"))
                y2_gdf.to_file(os.path.join(out_dir, "Waterbody_Y2_labelled.shp"))

        self._run_change_analysis(
            y1_gdf, y2_gdf, out_dir, target_crs, river_type
        )

    def _run_change_analysis(
        self, y1_gdf, y2_gdf, out_dir, target_crs, river_type
    ):
        # Compute riparian corridor
        corridor = change_analysis.compute_riparian_corridor(
            y1_gdf, y2_gdf, self.buffer_m
        )
        if corridor is not None:
            corridor_gdf = gpd.GeoSeries([corridor], crs=target_crs)
            corridor_gdf.to_file(
                os.path.join(out_dir, "Riparian_Corridor.shp")
            )

        # Fishnet: inside corridor if available, else union bounds
        if corridor is not None and not corridor.is_empty:
            bounds = corridor.bounds
        else:
            union_geom = y1_gdf.geometry.unary_union.union(
                y2_gdf.geometry.unary_union
            )
            bounds = gpd.GeoSeries(
                [union_geom], crs=y1_gdf.crs
            ).total_bounds

        fishnet = change_analysis.create_fishnet(
            bounds, self.cell_size_m, target_crs
        )

        # Clip fishnet to corridor
        if corridor is not None and not corridor.is_empty:
            fishnet = fishnet[fishnet.intersects(corridor)].copy()

        fishnet['grid_id'] = np.arange(1, len(fishnet) + 1)
        fishnet.to_file(os.path.join(out_dir, "Fishnet_Analysis.shp"))

        # Per-cell change computation
        per_grid = change_analysis.compute_per_grid_change(
            fishnet, y1_gdf, y2_gdf, out_dir, target_crs
        )

        # MMU filter
        if self.mmu_cells > 0:
            per_grid = change_analysis.apply_mmu_filter(
                per_grid, min_cells=self.mmu_cells, out_dir=out_dir
            )

        # Stretch aggregation
        stretch = change_analysis.aggregate_by_stretch(
            per_grid, self.stretch_m, self.cell_size_m,
            out_dir, top_n=self.top_n
        )

        # Braiding index for braided regime
        if river_type == "Braided":
            stretch = change_analysis.compute_braiding_index(
                stretch, y1_gdf, y2_gdf, self.stretch_m
            )

        # Regime classification
        if corridor is not None and not corridor.is_empty:
            regime = change_analysis.classify_stretch_regime(
                stretch, y1_gdf, y2_gdf, self.cell_size_m, self.stretch_m,
                user_river_type=river_type
            )
            if regime is not None:
                regime.to_file(
                    os.path.join(out_dir, "Stretch_Regime_Class.shp")
                )

        # Transect analysis
        transect_enabled = self.chkTransect.isChecked()
        if transect_enabled:
            if river_type == "Deltaic / Distributary":
                import warnings
                warnings.warn(
                    "Transect analysis disabled for deltaic regime: "
                    "no single centerline exists for distributary systems."
                )
            else:
                try:
                    merged_rows = (
                        list(y1_gdf[['geometry']].itertuples(
                            index=False, name=None
                        ))
                        + list(y2_gdf[['geometry']].itertuples(
                            index=False, name=None
                        ))
                    )
                    merged_gdf = gpd.GeoDataFrame(
                        merged_rows, columns=['geometry'],
                        crs=y1_gdf.crs
                    )
                    centerline = change_analysis.extract_centerline(
                        merged_gdf, cell_size=self.cell_size_m
                    )
                    if centerline is not None:
                        centerline.to_file(
                            os.path.join(out_dir, "Centerline.shp")
                        )
                        spacing = self.spinTransectSpacing.value()
                        transects = change_analysis.generate_transects(
                            centerline, spacing=spacing
                        )
                        if transects is not None:
                            transects.to_file(
                                os.path.join(out_dir, "Transects.shp")
                            )
                            rates = change_analysis.compute_transect_rates(
                                transects, y1_gdf, y2_gdf,
                                years_interval=self.years_interval,
                                out_dir=out_dir,
                            )
                            # Grid-Transect consistency
                            change_analysis.compute_consistency_index(
                                per_grid, rates, transects,
                                self.stretch_m, self.cell_size_m,
                                out_dir=out_dir
                            )
                except Exception as exc:
                    import warnings
                    warnings.warn(
                        f"Transect analysis failed: {exc}"
                    )
