#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Shared drawing code of Figures 8 and 9 (per-pixel attributes of the example basin 8595).

Input: ../Data/ex8595_*.tif, ex8595_record.json. Imported by Fig08_Terrain_v11.py and Fig09_Planform_v11.py.
"""

from __future__ import annotations

from pathlib import Path
from dataclasses import dataclass

import numpy as np
import geopandas as gpd
import rasterio
from rasterio.features import shapes as raster_shapes
from shapely.geometry import shape as shp_shape
from shapely.ops import unary_union
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec


HERE    = Path(__file__).resolve().parent
EX_DIR  = HERE.parent / "Data"
FIG_DIR = HERE.parent / "Figures"
FIG_DIR.mkdir(parents=True, exist_ok=True)

FOCUS_BSN_ID = 8595
ATTRIBUTES = ["mean_elv", "min_elv", "max_elv", "elv_std", "mean_slp", "slp_std", "perimeter_km", "basin_length_km",
              "basin_width_km", "convexity", "relief", "hypsometric_integral", "elongation_ratio", "gravelius",
              "circularity", "lemniscate_ratio"]
UPA_MIN_KM2  = 10.0


# ----------------------------------------------------------------------
def apply_style():
    mpl.rcParams.update({
        "font.family":       "Helvetica",
        "font.size":         8,
        "mathtext.fontset":  "custom",
        "mathtext.rm":       "Helvetica",
        "mathtext.it":       "Helvetica:italic",
        "axes.linewidth":    0.6,
        "xtick.major.width": 0.5,
        "ytick.major.width": 0.5,
        "xtick.labelsize":   6.5,
        "ytick.labelsize":   6.5,
        "axes.titlesize":    8,
        "axes.titlepad":     2.5,
    })


# ----------------------------------------------------------------------
def _sample_raster(tif: Path, xs, ys):
    with rasterio.open(tif) as s:
        arr = s.read(1)
        tr  = s.transform
        nd  = s.nodata
    rr, cc = rasterio.transform.rowcol(tr, xs, ys)
    vals = arr[np.array(rr), np.array(cc)].astype(np.float32)
    if nd is not None:
        vals = np.where(vals == nd, np.nan, vals)
    return vals


def _polygonise_focus(bsn_tif: Path, focus_id: int):
    with rasterio.open(bsn_tif) as s:
        arr = s.read(1)
        tr  = s.transform
    mask = (arr == focus_id).astype(np.uint8)
    polys = [
        shp_shape(geom)
        for geom, val in raster_shapes(
            mask, mask=mask.astype(bool), transform=tr)
        if int(val) == 1
    ]
    if not polys:
        raise RuntimeError(f"No polygon found for bsn == {focus_id}")
    merged = unary_union(polys)
    return gpd.GeoDataFrame(geometry=[merged], crs="EPSG:4326")


# ----------------------------------------------------------------------
def _read(name):
    with rasterio.open(EX_DIR / f"ex8595_{name}.tif") as src:
        arr = src.read(1).astype(np.float64)
        nd = src.nodata
        tr = src.transform
        bounds = src.bounds
    if nd is not None:
        arr = np.where(arr == nd, np.nan, arr)
    return arr, tr, bounds


def load_context():
    """Per-pixel GeoDataFrame (pixel centres, every pixel with UPA >= 10 km^2),
    basin polygon, outlet and window."""
    import json
    rec = json.load(open(EX_DIR / "ex8595_record.json"))
    bsn, tr, bounds = _read("bsn")
    upa, _, _ = _read("upa")
    gdf_basin = _polygonise_focus(EX_DIR / "ex8595_bsn.tif", FOCUS_BSN_ID)

    orow, ocol = rec["outlet_row_in_window"], rec["outlet_col_in_window"]
    ox, oy = rasterio.transform.xy(tr, orow, ocol)
    focus_all = bsn == FOCUS_BSN_ID
    assert np.nanargmax(np.where(focus_all, upa, -1)) == orow * bsn.shape[1] + ocol, "outlet is not the largest UPA"
    outlet_upa = float(upa[orow, ocol])

    keep = (bsn > 0) & (upa >= UPA_MIN_KM2)
    rr, cc = np.nonzero(keep)
    xs, ys = rasterio.transform.xy(tr, rr, cc)
    data = {"bsn": bsn[rr, cc], "upa": upa[rr, cc]}
    for name in ATTRIBUTES:
        data[name] = _read(name)[0][rr, cc]
    gdf_all = gpd.GeoDataFrame(data, geometry=gpd.points_from_xy(xs, ys), crs="EPSG:4326")

    mask_focus = (gdf_all["bsn"] == FOCUS_BSN_ID)
    gdf_focus  = gdf_all[ mask_focus].reset_index(drop=True)
    gdf_ctx    = gdf_all[~mask_focus].reset_index(drop=True)

    # Marker sizes for focus (log10(UPA))
    S_MIN, S_MAX = 0.8, 4.5
    log_upa      = np.log10(gdf_focus["upa"].values)
    log_min      = np.log10(UPA_MIN_KM2)
    log_max      = np.log10(gdf_focus["upa"].max())
    ms_focus     = S_MIN + (S_MAX - S_MIN) * (log_upa - log_min) / (log_max - log_min)
    s_focus      = ms_focus ** 2

    order        = np.argsort(gdf_focus["upa"].values)
    gdf_focus    = gdf_focus.iloc[order].reset_index(drop=True)
    s_focus      = s_focus[order]

    S_CTX = 0.5
    s_ctx = S_CTX ** 2

    # bbox = the raster window (basin 8595's bounding box)
    xmin, ymin, xmax, ymax = bounds.left, bounds.bottom, bounds.right, bounds.top
    pad_x = (xmax - xmin) * 0.02
    pad_y = (ymax - ymin) * 0.05
    bbox  = (xmin - pad_x, xmax + pad_x, ymin - pad_y, ymax + pad_y)

    print(f"focus basin (bsn={FOCUS_BSN_ID}): {len(gdf_focus)} pts (UPA>={UPA_MIN_KM2} km^2)")
    print(f"context                          : {len(gdf_ctx)} pts")
    print(f"outlet xy (UPA={outlet_upa:.2f} km²): ({ox:.4f}, {oy:.4f})")
    print(f"bbox window                      : lon [{bbox[0]:.4f}, {bbox[1]:.4f}]  lat [{bbox[2]:.4f}, {bbox[3]:.4f}]")

    return {
        "gdf_focus": gdf_focus,
        "s_focus":   s_focus,
        "gdf_ctx":   gdf_ctx,
        "s_ctx":     s_ctx,
        "outlet_xy": (ox, oy),
        "gdf_basin": gdf_basin,
        "bbox":      bbox,
    }


# ----------------------------------------------------------------------
@dataclass
class PanelSpec:
    col:   str
    lab:   str
    title: str
    cmap:  str


def plot_4x2(panels: list[PanelSpec], ctx: dict, stem: str,
             essd_style: bool = True):
    assert len(panels) == 8
    apply_style_essd() if essd_style else apply_style()

    gdf_focus = ctx["gdf_focus"]
    gdf_ctx_  = ctx["gdf_ctx"]
    s_focus   = ctx["s_focus"]
    s_ctx     = ctx["s_ctx"]
    ox, oy    = ctx["outlet_xy"]
    gdf_basin = ctx["gdf_basin"]
    xmin, xmax, ymin, ymax = ctx["bbox"]

    fig = plt.figure(figsize=(5.4, 9.0))
    gs = GridSpec(
        nrows=8, ncols=2, figure=fig,
        height_ratios=[1.0, 0.07] * 4,
        left=0.08, right=0.995, top=0.97, bottom=0.04,
        wspace=0.10, hspace=0.42,
    )

    for i, sp in enumerate(panels):
        row_pair = i // 2
        col_idx  = i %  2
        ax  = fig.add_subplot(gs[row_pair * 2,     col_idx])
        cax = fig.add_subplot(gs[row_pair * 2 + 1, col_idx])

        vmin = float(np.nanpercentile(gdf_focus[sp.col], 2))
        vmax = float(np.nanpercentile(gdf_focus[sp.col], 98))

        ax.scatter(
            gdf_ctx_.geometry.x, gdf_ctx_.geometry.y,
            c=gdf_ctx_[sp.col], cmap=sp.cmap,
            s=s_ctx, alpha=0.5,
            vmin=vmin, vmax=vmax,
            linewidths=0, zorder=2,
        )
        ax.scatter(
            gdf_focus.geometry.x, gdf_focus.geometry.y,
            c=gdf_focus[sp.col], cmap=sp.cmap,
            s=s_focus,
            vmin=vmin, vmax=vmax,
            linewidths=0, zorder=3,
        )

        gdf_basin.boundary.plot(ax=ax, edgecolor="black", linewidth=0.6, zorder=4)
        ax.plot(
            ox, oy, marker="v",
            markersize=7, markerfacecolor="#E53935",
            markeredgecolor="black", markeredgewidth=0.9,
            linestyle="", zorder=6,
        )

        ax.set_aspect("equal")
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(ymin, ymax)
        ax.set_xticks([round(xmin + 0.02, 2),
                       round((xmin + xmax) / 2, 2),
                       round(xmax - 0.02, 2)])
        ax.set_yticks([round(ymin + 0.02, 2),
                       round((ymin + ymax) / 2, 2),
                       round(ymax - 0.02, 2)])
        ax.tick_params(direction="out", length=2, pad=1.2)
        ax.set_title(f"({sp.lab}) {sp.title}", loc="center")

        sm = plt.cm.ScalarMappable(
            cmap=sp.cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax)
        )
        sm.set_array([])
        cb = fig.colorbar(sm, cax=cax, orientation="horizontal")
        cb.ax.tick_params(labelsize=6, direction="out", length=1.5, pad=0.8)
        cb.outline.set_linewidth(0.3)

        if col_idx == 0:
            ax.set_ylabel("Lat [°N]", fontsize=7)
        else:
            ax.set_yticklabels([])
        if row_pair == 3:
            cax.set_xlabel("Lon [°E]", fontsize=7, labelpad=10)

    out_stem = FIG_DIR / stem
    fig.savefig(f"{out_stem}.png", dpi=300, bbox_inches="tight",
                facecolor="white")
    print(f"  wrote {out_stem}.png")
    plt.close(fig)


# ----------------------------------------------------------------------
def apply_style_essd():
    """ESSD-ready style: Helvetica + cm mathtext + 9 pt (matches Figs 4-7 v_)."""
    mpl.rcParams.update({
        "font.family":         "Helvetica",
        "font.size":            9,
        "axes.labelsize":       9,
        "axes.titlesize":       9,
        "axes.titleweight":     "bold",
        "xtick.labelsize":      7,
        "ytick.labelsize":      7,
        "mathtext.fontset":   "cm",
        "axes.linewidth":     0.5,
    })


def plot_2x4(panels: list[PanelSpec], ctx: dict, stem: str,
             figsize=(7.4, 4.0)):
    """ESSD 2-column layout: 2 rows × 4 cols of panels.

    The example basin is ~2:1 wide so each cell fits the basin tightly.
    Per-cell layout:
        +---------------------------+
        | (a) attribute name [unit] |  <- title (top)
        |   ............basin..     |  <- scatter
        | [horizontal colorbar  ]   |  <- compact cbar (bottom)
        +---------------------------+
    Lat/Lon ticks suppressed (same basin in all 8 cells, redundant).
    """
    assert len(panels) == 8
    apply_style_essd()

    gdf_focus = ctx["gdf_focus"]
    gdf_ctx_  = ctx["gdf_ctx"]
    s_focus   = ctx["s_focus"]
    s_ctx     = ctx["s_ctx"]
    ox, oy    = ctx["outlet_xy"]
    gdf_basin = ctx["gdf_basin"]
    xmin, xmax, ymin, ymax = ctx["bbox"]

    fig = plt.figure(figsize=figsize)
    gs = GridSpec(
        nrows=4, ncols=4, figure=fig,
        height_ratios=[1.0, 0.07, 1.0, 0.07],
        left=0.015, right=0.99, top=0.94, bottom=0.04,
        wspace=0.06, hspace=0.30,
    )

    for i, sp in enumerate(panels):
        row_pair = i // 4         # 0 (top) or 1 (bottom)
        col_idx  = i %  4
        ax  = fig.add_subplot(gs[row_pair * 2,     col_idx])
        cax = fig.add_subplot(gs[row_pair * 2 + 1, col_idx])

        vmin = float(np.nanpercentile(gdf_focus[sp.col], 2))
        vmax = float(np.nanpercentile(gdf_focus[sp.col], 98))

        ax.scatter(
            gdf_ctx_.geometry.x, gdf_ctx_.geometry.y,
            c=gdf_ctx_[sp.col], cmap=sp.cmap,
            s=s_ctx, alpha=0.5,
            vmin=vmin, vmax=vmax,
            linewidths=0, zorder=2,
        )
        ax.scatter(
            gdf_focus.geometry.x, gdf_focus.geometry.y,
            c=gdf_focus[sp.col], cmap=sp.cmap,
            s=s_focus,
            vmin=vmin, vmax=vmax,
            linewidths=0, zorder=3,
        )

        gdf_basin.boundary.plot(ax=ax, edgecolor="black",
                                linewidth=0.5, zorder=4)
        ax.plot(
            ox, oy, marker="v",
            markersize=5, markerfacecolor="#E53935",
            markeredgecolor="black", markeredgewidth=0.6,
            linestyle="", zorder=6,
        )

        ax.set_aspect("equal")
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(ymin, ymax)
        ax.set_xticks([]); ax.set_yticks([])
        for spn in ax.spines.values():
            spn.set_linewidth(0.4)
            spn.set_color("#999999")
        ax.set_title(f"({sp.lab}) {sp.title}",
                     loc="left", pad=2.5, fontsize=8.5)

        sm = plt.cm.ScalarMappable(
            cmap=sp.cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax))
        sm.set_array([])
        cb = fig.colorbar(sm, cax=cax, orientation="horizontal")
        cb.ax.tick_params(labelsize=6.5, direction="out",
                          length=1.5, pad=0.8)
        cb.outline.set_linewidth(0.3)
        # Limit ticks to 3-4 to avoid label crowding
        from matplotlib.ticker import MaxNLocator
        cb.locator = MaxNLocator(nbins=4)
        cb.update_ticks()

    out_stem = FIG_DIR / stem
    fig.savefig(f"{out_stem}.png", dpi=300, bbox_inches="tight",
                facecolor="white")
    print(f"  wrote {out_stem}.png")
    plt.close(fig)


# ----------------------------------------------------------------------
@dataclass
class ContextSpec:
    """Per-pixel raster shown as panel (a) anchor before the 8 derived
    attributes.  Use imshow on the raster (full per-pixel detail), masked
    to FOCUS_BSN_ID."""
    raster:    str       # filename in EX_DIR (e.g. "ex8595_dem.tif")
    title:     str       # panel title (without "(a)" prefix)
    cmap:      str
    log_scale: bool = False  # use LogNorm (e.g. for UPA)


def plot_3x3_with_context(panels: list[PanelSpec],
                          ctx: dict,
                          context: ContextSpec,
                          stem: str,
                          figsize=(7.0, 5.5)):
    """ESSD 3×3 layout: panel (a) is a per-pixel raster context (DEM or
    UPA), panels (b)-(i) are the 8 derived attribute scatters.  Cbar is
    inset inside each panel (lower-right corner) so the basin gets the
    most visual area and the cbar reads as a subordinate element.
    Lat/Lon ticks suppressed (same basin in all 9 cells).
    """
    from matplotlib.colors import LogNorm
    from matplotlib.ticker import MaxNLocator, LogLocator, LogFormatterSciNotation

    assert len(panels) == 8, "panels must list the 8 derived attributes"
    apply_style_essd()

    gdf_focus = ctx["gdf_focus"]
    gdf_ctx_  = ctx["gdf_ctx"]
    s_focus   = ctx["s_focus"]
    s_ctx     = ctx["s_ctx"]
    ox, oy    = ctx["outlet_xy"]
    gdf_basin = ctx["gdf_basin"]
    xmin, xmax, ymin, ymax = ctx["bbox"]

    # ---- Read context raster, mask to focus basin --------------------
    with rasterio.open(EX_DIR / "ex8595_bsn.tif") as s:
        bsn_arr = s.read(1)
    mask_focus = (bsn_arr == FOCUS_BSN_ID)
    with rasterio.open(EX_DIR / context.raster) as s:
        ctx_arr  = s.read(1).astype(np.float64)
        ctx_extent = (s.bounds.left, s.bounds.right,
                      s.bounds.bottom, s.bounds.top)
        ctx_nd = s.nodata
    valid = mask_focus & np.isfinite(ctx_arr)
    if ctx_nd is not None:
        valid &= (ctx_arr != ctx_nd)
    if context.log_scale:
        valid &= (ctx_arr > 0)
    ctx_disp = np.ma.masked_where(~valid, ctx_arr)
    c_vmin = float(np.nanpercentile(ctx_arr[valid], 2))
    c_vmax = float(np.nanpercentile(ctx_arr[valid], 98))
    ctx_norm = (LogNorm(vmin=max(c_vmin, 1e-3), vmax=c_vmax)
                if context.log_scale
                else plt.Normalize(vmin=c_vmin, vmax=c_vmax))

    # ---- Layout ------------------------------------------------------
    fig = plt.figure(figsize=figsize)
    gs = GridSpec(
        nrows=3, ncols=3, figure=fig,
        left=0.015, right=0.995, top=0.96, bottom=0.06,
        wspace=0.02, hspace=0.25,
    )

    def style_axes(ax):
        ax.set_aspect("equal")
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(ymin, ymax)
        ax.set_xticks([]); ax.set_yticks([])
        for spn in ax.spines.values():
            spn.set_linewidth(0.4); spn.set_color("#999999")

    def add_outlet_and_outline(ax):
        gdf_basin.boundary.plot(ax=ax, edgecolor="black",
                                linewidth=0.5, zorder=4)
        ax.plot(ox, oy, marker="v",
                markersize=5, markerfacecolor="#E53935",
                markeredgecolor="black", markeredgewidth=0.6,
                linestyle="", zorder=6)

    def below_cbar(ax, sm, n_ticks=5, log=False):
        # Standard cbar OUTSIDE panel, just below ax — full panel width
        # (~90 %), readable thickness, 7 pt ticks.  The cbar is a
        # first-class element (readers map colour → value here), so it
        # gets normal proportions.
        cax = ax.inset_axes([0.05, -0.18, 0.90, 0.045])
        cb = fig.colorbar(sm, cax=cax, orientation="horizontal")
        cb.ax.tick_params(labelsize=7, direction="out",
                          length=1.8, pad=1.2)
        cb.outline.set_linewidth(0.4)
        if log:
            cb.locator   = LogLocator(base=10)
            cb.formatter = LogFormatterSciNotation(base=10)
        else:
            cb.locator   = MaxNLocator(nbins=n_ticks)
        cb.update_ticks()
        return cb

    # ---- Panel (a): context raster (imshow) --------------------------
    ax_a = fig.add_subplot(gs[0, 0])
    im = ax_a.imshow(ctx_disp, extent=ctx_extent, cmap=context.cmap,
                     origin="upper", interpolation="nearest",
                     norm=ctx_norm, zorder=2)
    add_outlet_and_outline(ax_a)
    style_axes(ax_a)
    ax_a.set_title(f"(a) {context.title}", loc="left",
                   pad=2.5, fontsize=8.5)
    below_cbar(ax_a, im, n_ticks=5, log=context.log_scale)

    # ---- Panels (b)-(i): 8 derived attribute scatters ----------------
    LETTERS = "bcdefghi"
    for k, sp in enumerate(panels):
        idx = k + 1            # cell index (skip (0,0) which is context)
        row = idx // 3
        col = idx %  3
        ax  = fig.add_subplot(gs[row, col])

        vmin = float(np.nanpercentile(gdf_focus[sp.col], 2))
        vmax = float(np.nanpercentile(gdf_focus[sp.col], 98))

        ax.scatter(
            gdf_ctx_.geometry.x, gdf_ctx_.geometry.y,
            c=gdf_ctx_[sp.col], cmap=sp.cmap,
            s=s_ctx, alpha=0.5,
            vmin=vmin, vmax=vmax,
            linewidths=0, zorder=2,
        )
        ax.scatter(
            gdf_focus.geometry.x, gdf_focus.geometry.y,
            c=gdf_focus[sp.col], cmap=sp.cmap,
            s=s_focus,
            vmin=vmin, vmax=vmax,
            linewidths=0, zorder=3,
        )
        add_outlet_and_outline(ax)
        style_axes(ax)
        ax.set_title(f"({LETTERS[k]}) {sp.title}",
                     loc="left", pad=2.5, fontsize=8.5)

        sm = plt.cm.ScalarMappable(
            cmap=sp.cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax))
        sm.set_array([])
        below_cbar(ax, sm, n_ticks=5)

    out_stem = FIG_DIR / stem
    for ext in ("png", "pdf"):
        if Path(f"{out_stem}.{ext}").exists():
            raise SystemExit(f"{out_stem}.{ext} is there already")
        fig.savefig(f"{out_stem}.{ext}", dpi=300, bbox_inches="tight",
                    facecolor="white")
        print(f"  wrote {out_stem}.{ext}")
    plt.close(fig)