#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Figure 10: perimeter-area scaling and the check of the Moore-VS perimeter against three other methods.

Input: ../Data/Fig10_perimeter_v11_mlj.csv (one row per validation record; P_km_v11 is the released Moore-VS perimeter, perim_mlj_cells_km the Crofton perimeter of MorphoLibJ 1.6.5 on the same cells).
Output: ../Figures/Fig10_Perimeter_v11_lat.png and .pdf.

Run from this folder; ../Data holds the inputs and the figure is written to ../Figures.
"""

from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec


mpl.rcParams.update({
    "font.family":      "Helvetica",
    "font.size":        10,
    "axes.labelsize":   11,
    "legend.fontsize":  9,
    "xtick.labelsize":  9,
    "ytick.labelsize":  9,
    "mathtext.fontset": "cm",
    "axes.linewidth":   0.8,
})


HERE = Path(__file__).resolve().parent
INPUT_CSV = HERE.parent / "Data" / "Fig10_perimeter_v11_mlj.csv"
ASPECT_MAX = 1.2
SPAN_MAX = 1.0          # % (MorphoLibJ one pixel size per image)

OUT_DIR = HERE.parent / "Figures"


# Method labels as in the manuscript
ALT_METHODS = [
    # Ordered top→bottom by median Δ (control → cross-check → reference)
    ("GIS staircase",                   "perim_vector_raw_km",      "#E8723A"),
    ("Marching squares",                "perim_vector_subpixel_km", "#3E78B2"),
    ("Crofton, MorphoLibJ (reference)", "perim_mlj_cells_km",       "#A9B537"),
]
NEAR_SQUARE_ONLY = {"perim_vector_raw_km", "perim_vector_subpixel_km"}


def build_long(df):
    """One row per mask (outlet and sub-mid, already long-form); in (b) marching squares only where the pixels are
    close to square (mean aspect <= ASPECT_MAX), MorphoLibJ where one pixel size fits the mask (span <= SPAN_MAX)."""
    out = pd.DataFrame({"A_km2": df["A_km2"], "P_km": df["P_km_v11"]})
    near_square = df["aspect"] <= ASPECT_MAX
    for _, key, _ in ALT_METHODS:
        out[key] = df[key].where(near_square) if key in NEAR_SQUARE_ONLY else df[key]
    out["perim_mlj_cells_km"] = out["perim_mlj_cells_km"].where(df["mlj_span_pct"] <= SPAN_MAX)
    return out.reset_index(drop=True)


def main():
    print(f"Reading {INPUT_CSV}")
    df = pd.read_csv(INPUT_CSV)
    long_df = build_long(df)
    n = len(long_df)
    print(f"  long-form rows = {n:,}")

    A = long_df["A_km2"].to_numpy()
    P = long_df["P_km"].to_numpy()

    logA = np.log10(A); logP = np.log10(P)
    alpha, beta = np.polyfit(logA, logP, 1)
    r2 = 1.0 - np.sum((logP - (alpha * logA + beta)) ** 2) / \
               np.sum((logP - logP.mean()) ** 2)
    print(f"  fit: P = {10**beta:.3f} · A^{alpha:.4f}    R² = {r2:.4f}")

    A_lo = 10 ** np.floor(np.log10(A.min()))
    A_hi = 10 ** np.ceil (np.log10(A.max()))
    A_ref  = np.logspace(np.log10(A_lo), np.log10(A_hi), 200)
    P_circ = 2.0 * np.sqrt(np.pi * A_ref)
    P_fit  = 10.0 ** (alpha * np.log10(A_ref) + beta)

    n_txt = f"~{round(n/1000)}k" if n >= 1000 else f"{n:,}"

    # ---------------------------------------------------------------
    fig = plt.figure(figsize=(6.4, 10.2))
    gs  = GridSpec(4, 1, height_ratios=[2.3, 1.0, 0.42, 1.25], hspace=0.08)
    ax_top = fig.add_subplot(gs[0, 0])
    ax_bot = fig.add_subplot(gs[1, 0], sharex=ax_top)
    ax_lat = fig.add_subplot(gs[3, 0])

    # (a) =========================================================== P-A
    ax_top.scatter(A, P, s=1.2, alpha=0.22, color="#6E6E6E",
                   edgecolors="none", rasterized=True)

    ax_top.plot(A_ref, P_circ, "k--", lw=1.0, alpha=0.75,
                label=r"Euclidean circle:  $P=2\sqrt{\pi}\,A^{0.5}$")

    ax_top.plot(A_ref, P_fit, color="#C0392B", lw=2.0, alpha=0.95,
                label=(fr"Moore–VS fit ($N\approx{n_txt}$):  "
                       fr"$P={10**beta:.2f}\,A^{{{alpha:.3f}}}$"
                       fr"  ($R^2={r2:.3f}$)"))

    ax_top.set_xscale("log")
    ax_top.set_yscale("log")
    # tight P range: ignore extreme outliers (a few ~0.1 km records on the
    # smallest sub-mid tiles) and pad lightly
    P_lo = np.percentile(P, 0.5) * 0.7
    P_hi = np.percentile(P, 99.9) * 1.5
    ax_top.set_ylim(P_lo, P_hi)
    ax_top.set_ylabel(r"Basin perimeter  $P$  (km)")
    ax_top.legend(loc="lower right", frameon=True, framealpha=0.95,
                  edgecolor="0.7")
    ax_top.grid(True, which="both", lw=0.3, alpha=0.3)
    ax_top.tick_params(axis="x", labelbottom=False)
    ax_top.text(-0.13, 1.01, "(a)", transform=ax_top.transAxes,
                fontsize=12, fontweight="bold", va="bottom")

    # (b) =========================================================== Δ%
    ax_bot.axhline( 0.0, color="k",   lw=0.9, alpha=0.7)
    ax_bot.axhline( 5.0, color="0.5", lw=0.6, ls="--", alpha=0.6)
    ax_bot.axhline(-5.0, color="0.5", lw=0.6, ls="--", alpha=0.6)

    stat_lines = []  # (color, text)
    for label, key, color in ALT_METHODS:
        y_alt = long_df[key].to_numpy()
        m = (~np.isnan(y_alt)) & (y_alt > 0)
        if m.sum() == 0:
            continue
        d = 100.0 * (y_alt[m] - P[m]) / P[m]
        ax_bot.scatter(A[m], d, s=0.8, alpha=0.35, color=color,
                       edgecolors="none", rasterized=True)
        stat_lines.append((
            color,
            f"{label}:  median $\\Delta = {np.nanmedian(d):+.2f}\\%$",
        ))

    # White backdrop + three colour-coded stat lines, anchored in DATA
    # coords to sit in the visual gap between the sub-pixel cluster (~+5 %)
    # and the Vector-raw cluster (~+27 %).
    from matplotlib.patches import FancyBboxPatch
    # bbox footprint, in data coords for the y axis (Δ %) and axes-frac
    # for the x axis (flush right of the axis)
    y_hi_data, y_lo_data = 24.0, 7.0                       # Δ ∈ [7, 24] %
    y_hi_frac = (y_hi_data - (-8)) / (38 - (-8))           # → axes frac
    y_lo_frac = (y_lo_data - (-8)) / (38 - (-8))
    bg_w       = 0.62
    bg_x       = 1.0 - bg_w - 0.012                        # flush right
    bg_y       = y_lo_frac
    bg_h       = y_hi_frac - y_lo_frac

    ax_bot.add_patch(FancyBboxPatch(
        (bg_x, bg_y), bg_w, bg_h,
        boxstyle="round,pad=0.004,rounding_size=0.008",
        facecolor="white", edgecolor="0.25", lw=0.9, alpha=0.95,
        transform=ax_bot.transAxes, zorder=4,
    ))

    # centre the 3 lines vertically in the box
    n_lines  = len(stat_lines)
    dy_frac  = 0.095
    y_top    = bg_y + bg_h - 0.030
    for i, (color, txt) in enumerate(stat_lines):
        ax_bot.text(bg_x + 0.015, y_top - i * dy_frac, txt,
                    transform=ax_bot.transAxes, ha="left", va="top",
                    fontsize=8.5, color=color, zorder=6)

    ax_bot.set_xscale("log")
    ax_bot.set_xlim(A_lo, A_hi)
    ax_bot.set_ylim(-8, 38)
    ax_bot.set_xlabel(r"Basin area  $A$  (km$^2$)")
    ax_bot.set_ylabel(r"$\Delta$ vs. Moore–VS  (%)")
    ax_bot.grid(True, which="both", lw=0.3, alpha=0.3)
    ax_bot.text(-0.13, 1.02, "(b)", transform=ax_bot.transAxes,
                fontsize=12, fontweight="bold", va="bottom")

    # (c) =========================================================== Δ% by latitude
    lat = df["latitude"].abs().to_numpy()
    ok = (df["mlj_span_pct"] <= SPAN_MAX).to_numpy()
    edges = np.arange(0.0, 90.0, 5.0)
    mid = 0.5 * (edges[:-1] + edges[1:])
    which = np.digitize(lat, edges) - 1

    def by_band(values, q):
        out = np.full(mid.size, np.nan)
        for i in range(mid.size):
            v = values[ok & (which == i)]
            v = v[np.isfinite(v)]
            if v.size >= 30:
                out[i] = np.percentile(v, q)
        return out

    same = 100.0 * (df["P_km_v11"] / df["perim_mlj_cells_km"] - 1.0).to_numpy()
    ax_lat.axhline(0.0, color="k", lw=0.9, alpha=0.7)
    ax_lat.fill_between(mid, by_band(same, 10), by_band(same, 90), color="#C0392B", alpha=0.18, lw=0)
    ax_lat.plot(mid, by_band(same, 50), color="#C0392B", lw=2.0, ls="-",
                label="Moore–VS vs MorphoLibJ on the same cells (median, 10–90 %)")
    ax_lat.legend(loc="upper left", frameon=True, framealpha=0.95, edgecolor="0.7", fontsize=8.5)
    ax_lat.set_xlim(0, 85)
    ax_lat.set_xticks([0, 20, 40, 60, 80])
    ax_lat.set_xticklabels(["0°", "20°", "40°", "60°", "80°"])
    ax_lat.set_ylim(-4, 4)
    ax_lat.set_xlabel("Latitude (north or south)")
    ax_lat.set_ylabel(r"$\Delta$ vs. MorphoLibJ  (%)")
    ax_lat.grid(True, which="both", lw=0.3, alpha=0.3)
    top = ax_lat.secondary_xaxis("top")
    a = 6378137.0; e2 = (1 / 298.257223563) * (2 - 1 / 298.257223563)
    def aspect(phi):
        s = np.sin(np.radians(phi)); return (1 - e2) / (1 - e2 * s * s) / np.cos(np.radians(phi))
    ticks = [0.0, 45.0, 60.0, 70.5, 80.0]
    top.set_xticks(ticks)
    top.set_xticklabels([f"{aspect(x):.1f}:1" for x in ticks], fontsize=8)
    top.set_xlabel("Pixel height : width", fontsize=8.5)
    ax_lat.text(-0.13, 1.02, "(c)", transform=ax_lat.transAxes, fontsize=12, fontweight="bold", va="bottom")

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    for ext in ("png", "pdf"):
        p = OUT_DIR / f"Fig10_Perimeter_v11_lat.{ext}"
        if p.exists():
            raise SystemExit(f"{p} is there already")
        fig.savefig(p, dpi=300, bbox_inches="tight")
        print(f"  wrote {p}")


if __name__ == "__main__":
    main()
