"""exercise_0c.py

Exercise 0-C. Regression kriging and what a covariate buys.

Predicts mean annual temperature three ways from identical points, by
ordinary kriging, by a linear trend on elevation and by regression kriging,
and scores each by leave one out cross validation. Elevation off the stations
is itself kriged from the same stations, so every cross validated gain is
measured where elevation was observed.

    python exercise_0c.py
"""
from __future__ import annotations

import sys
from pathlib import Path

import matplotlib
import numpy as np
import pandas as pd

matplotlib.use("Agg")
import matplotlib.pyplot as plt                       # noqa: E402

sys.path.insert(0, str(Path(__file__).resolve().parent))
import common as C                                    # noqa: E402
from kriging import (empirical_variogram, fit_variogram, grid,  # noqa: E402
                     loocv, ordinary_kriging, project_albers)

VAR, COV = "tmean_C", "elev_m"
NOISE_LEVELS = [0, 50, 100, 200, 400, 800]     # meters of covariate error
STANDARD_LAPSE = 6.5                            # degrees C per 1000 m


def load():
    d = (C.load_azmet().dropna(subset=[VAR, COV, "lat", "lon"])
         .reset_index(drop=True))
    x, y = project_albers(d.lon.to_numpy(), d.lat.to_numpy())
    return d, x, y, d[VAR].to_numpy(float), d[COV].to_numpy(float)


def loocv_regression(s, z):
    """Leave one out error for the trend alone, with no kriging at all."""
    n, pred = len(z), np.empty(len(z))
    for i in range(n):
        m = np.ones(n, bool)
        m[i] = False
        b, a = np.polyfit(s[m], z[m], 1)
        pred[i] = a + b * s[i]
    err = pred - z
    return {"rmse": float(np.sqrt((err ** 2).mean())),
            "mae": float(np.abs(err).mean()), "pred": pred, "error": err}


def main():
    C.ensure_dirs()
    d, x, y, z, s = load()
    C.banner("exercise 0-C: regression kriging and what a covariate buys",
             "AZMet daily observations, 2024, keyless",
             "azmet_2024_annual.csv",
             {"stations": len(d), "primary": VAR, "covariate": COV})
    out = {"stations": int(len(d))}

    # ---- 1. establish the relationship ---------------------------------
    slope, intercept = np.polyfit(s, z, 1)
    r, _ = C.corr(s, z)
    reg = loocv_regression(s, z)
    lapse = -slope * 1000
    print("\n[1] THE RELATIONSHIP")
    print(f"    lapse rate      {lapse:.3f} C per 1000 m")
    print(f"    intercept       {intercept:.3f} C")
    print(f"    pearson r       {r:+.4f}   r squared {r*r:.4f}")
    print(f"    LOOCV, trend    rmse {reg['rmse']:.3f} C  "
          f"mae {reg['mae']:.3f} C")
    print(f"    standard atmospheric lapse rate is near {STANDARD_LAPSE} C "
          f"per 1000 m, so the observed value sits shallower by "
          f"{STANDARD_LAPSE - lapse:.2f}")
    out["relationship"] = {
        "lapse_C_per_1000m": round(lapse, 3),
        "intercept_C": round(float(intercept), 3),
        "pearson_r": round(r, 4), "r_squared": round(r * r, 4),
        "trend_only_rmse": round(reg["rmse"], 3),
        "trend_only_mae": round(reg["mae"], 3),
        "standard_lapse": STANDARD_LAPSE,
        "shortfall": round(STANDARD_LAPSE - lapse, 2)}

    # ---- 2. krige three ways -------------------------------------------
    cen, val, cnt, lag, npairs = empirical_variogram(x, y, z, n_bins=12)
    p_ok = fit_variogram(cen, val, cnt)
    ok = loocv(x, y, z, p_ok)

    # Residual variogram, which is what regression kriging actually models.
    resid = z - (intercept + slope * s)
    cenr, valr, cntr, _, _ = empirical_variogram(x, y, resid, n_bins=12)
    p_rk = fit_variogram(cenr, valr, cntr)
    rk = loocv(x, y, z, p_rk, secondary=s)

    print("\n[2] THREE WAYS TO THE SAME SURFACE")
    print(f"    {'method':<44}{'rmse':>8}{'mae':>8}")
    rows = [("ordinary kriging, location only", ok["rmse"], ok["mae"]),
            ("trend on elevation, no kriging", reg["rmse"], reg["mae"]),
            ("regression kriging, elevation known everywhere",
             rk["rmse"], rk["mae"])]
    for name, a, b in rows:
        print(f"    {name:<44}{a:>8.3f}{b:>8.3f}")
    out["methods"] = [{"method": n, "rmse": round(a, 3), "mae": round(b, 3)}
                      for n, a, b in rows]
    out["variogram_primary"] = {"nugget": round(p_ok["nugget"], 3),
                                "sill": round(p_ok["sill"], 3),
                                "range_km": round(p_ok["rng"] / 1000, 1)}
    out["variogram_residual"] = {"nugget": round(p_rk["nugget"], 3),
                                 "sill": round(p_rk["sill"], 3),
                                 "range_km": round(p_rk["rng"] / 1000, 1)}

    # ---- 3. why a station sampled covariate adds little -----------------
    gain_station = ok["rmse"] - reg["rmse"]
    gain_full = ok["rmse"] - rk["rmse"]
    print("\n[3] WHAT THE COVARIATE BOUGHT")
    print(f"    trend alone against ordinary kriging   "
          f"{gain_station:+.3f} C rmse")
    print(f"    regression kriging against ordinary    {gain_full:+.3f} C rmse")
    print(f"    residual sill is {p_rk['sill']:.3f} against a primary sill of "
          f"{p_ok['sill']:.3f}")
    print("    The covariate removed most of the structure the primary "
          "variogram was modeling.")
    out["gain_trend_only"] = round(gain_station, 3)
    out["gain_regression_kriging"] = round(gain_full, 3)

    # ---- 4. degrade the covariate ---------------------------------------
    print("\n[4] DEGRADING THE COVARIATE")
    print(f"    {'noise sd (m)':>13}{'rmse':>9}{'still helping':>16}")
    rng_ = np.random.default_rng(C.SEED)
    out["noise"] = []
    for nz in NOISE_LEVELS:
        errs = []
        for _ in range(10 if nz else 1):
            sn = s + rng_.normal(0, nz, len(s)) if nz else s
            errs.append(loocv(x, y, z, p_rk, secondary=sn)["rmse"])
        e = float(np.median(errs))
        helping = "yes" if e < ok["rmse"] else "no"
        print(f"    {nz:>13d}{e:>9.3f}{helping:>16}")
        out["noise"].append({"noise_sd_m": nz, "rmse": round(e, 3),
                             "still_helping": helping})

    # ---- figures ---------------------------------------------------------
    src = (f"Source: Arizona Meteorological Network, 2024 daily observations. "
           f"{len(d)} stations. EPSG:5070.")
    fig, ax = plt.subplots(figsize=(7.5, 5.5))
    ax.scatter(s, z, s=42, c="#1f3864", zorder=3)
    xs = np.linspace(s.min(), s.max(), 50)
    ax.plot(xs, intercept + slope * xs, color="#b2182b", lw=2,
            label=f"fit, {lapse:.2f} C per 1000 m")
    ax.plot(xs, z.mean() + (-STANDARD_LAPSE / 1000) * (xs - s.mean()),
            ls="--", color="#666", lw=1.5,
            label=f"standard lapse, {STANDARD_LAPSE} C per 1000 m")
    worst = int(np.abs(reg["error"]).argmax())
    ax.annotate(d.name.iloc[worst], (s[worst], z[worst]),
                textcoords="offset points", xytext=(8, -4), fontsize=9,
                color="#b2182b")
    ax.set_xlabel("station elevation (m)")
    ax.set_ylabel("mean annual air temperature (degrees C)")
    ax.set_title("Temperature against elevation", loc="left",
                 fontweight="bold")
    ax.legend(frameon=False, fontsize=9)
    ax.text(0, -0.13, src, transform=ax.transAxes, fontsize=7, color="#666")
    fig.savefig(C.FIG / "exercise_0c_lapse.png", dpi=150, bbox_inches="tight")
    plt.close(fig)

    XX, YY, gx, gy = grid(x, y)
    # Elevation is unknown off-station, so the continuous covariate is itself
    # interpolated. That is the honest version of "known everywhere" and its
    # own error is part of what the exercise measures.
    cen_s, val_s, cnt_s, _, _ = empirical_variogram(x, y, s, n_bins=12)
    p_s = fit_variogram(cen_s, val_s, cnt_s)
    s_grid, _ = ordinary_kriging(x, y, s, XX.ravel(), YY.ravel(), p_s)
    surf_ok, _ = ordinary_kriging(x, y, z, XX.ravel(), YY.ravel(), p_ok)
    surf_rk, _ = ordinary_kriging(x, y, z, XX.ravel(), YY.ravel(), p_rk,
                                  secondary=s, secondary_i=s_grid)

    fig, axes = plt.subplots(1, 3, figsize=(17, 5.2))
    vmin = min(surf_ok.min(), surf_rk.min())
    vmax = max(surf_ok.max(), surf_rk.max())
    for ax, dat, t in ((axes[0], surf_ok,
                        f"Ordinary kriging  rmse {ok['rmse']:.2f}"),
                       (axes[1], surf_rk,
                        f"Regression kriging  rmse {rk['rmse']:.2f}"),
                       (axes[2], surf_rk - surf_ok, "Difference")):
        kw = dict(cmap="RdYlBu_r", vmin=vmin, vmax=vmax)
        if t == "Difference":
            kw = dict(cmap="PuOr", vmin=-4, vmax=4)
        im = ax.pcolormesh(gx / 1000, gy / 1000, dat.reshape(XX.shape),
                           shading="auto", **kw)
        ax.scatter(x / 1000, y / 1000, s=10, c="k", zorder=3)
        ax.set_aspect("equal"), ax.set_xticks([]), ax.set_yticks([])
        ax.set_title(t, loc="left", fontweight="bold", fontsize=11)
        fig.colorbar(im, ax=ax, shrink=0.8)
    axes[0].text(0, -0.06, src, transform=axes[0].transAxes, fontsize=7,
                 color="#666")
    fig.suptitle("What a covariate known everywhere changes",
                 fontsize=13, fontweight="bold")
    fig.tight_layout()
    fig.savefig(C.FIG / "exercise_0c_surfaces.png", dpi=140,
                bbox_inches="tight")
    plt.close(fig)

    tbl = d[["station", "name", "county", COV, "lat", "lon", VAR]].copy()
    tbl["ok_loocv_error"] = ok["error"]
    tbl["trend_loocv_error"] = reg["error"]
    tbl["rk_loocv_error"] = rk["error"]
    C.write_handoff("exercise_0c", tbl, {
        "trend_form": "linear in elevation",
        "residual_variogram_model": p_rk["model"],
        "continuous_covariate": "elevation interpolated by ordinary kriging "
                                "from the same stations, so it carries its "
                                "own error",
        "search_neighborhood": "all stations, no maximum",
        "stations_outside_state": "none excluded",
        "grid_resolution": f"{XX.shape[1]} by {XX.shape[0]}",
        "color_ramp": "RdYlBu_r, unjustified",
    }, [
        "Elevation off-station was interpolated from the same 31 stations. A "
        "real terrain model would be far denser. Does that change the "
        "conclusion, and in which direction?",
        "Is a linear trend on elevation right, or does the relationship bend "
        "at the top of the range where the network has few stations?",
        "The largest residual station has a site description in the published "
        "metadata. Does it explain the residual?",
        "Should the surface be drawn where no station lies within the "
        "variogram range?",
    ])
    C.save_checkvalues("exercise_0c", out)
    print(f"\nwrote {C.FIG / 'exercise_0c_lapse.png'}")
    print(f"wrote {C.FIG / 'exercise_0c_surfaces.png'}")


if __name__ == "__main__":
    main()
