#!/usr/bin/env python
"""Stage 02 — reconcile gridded vs station observed Tx before choosing a calibration target.

Quantifies gridded-vs-station disagreement in daily Tx and in derived EHD frequency, per
region, over the common gridded period (2002-2023). Nearest-cell matching; 29 Feb dropped
(gridded calendar is 365_day). EHD uses the diagnostic gridded Tx90p climatology from
stage 01 (min(Tx90p,35)); this is a reconciliation diagnostic, not a CV skill number.

Outputs:
  results/tables/obs_reconciliation_station.csv   (per station)
  results/tables/obs_reconciliation_region.csv    (per region)
  results/figures/obs_reconciliation.png
  docs/OBS_RECONCILIATION.md                       (decision: target vs validation)
"""
from __future__ import annotations
import os, sys, logging
import numpy as np
import pandas as pd
import xarray as xr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from ehd.config import load_config
from ehd import stations as stn
from ehd import climatology as clim

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("stage02")

REGION_NAMES = {1: "R1 Tay Bac", 2: "R2 Dong Bac", 3: "R3 DB Bac Bo", 4: "R4 Bac Trung Bo",
                5: "R5 Nam Trung Bo", 6: "R6 Tay Nguyen", 7: "R7 Nam Bo"}


def doy_of(idx: pd.DatetimeIndex) -> np.ndarray:
    """Day-of-year 1..365 on a no-leap convention: after 28 Feb, subtract leap offset."""
    doy = np.asarray(idx.dayofyear).copy()
    isleap = np.asarray(idx.is_leap_year)
    after_feb29 = np.asarray(idx.month) > 2
    doy[isleap & after_feb29] -= 1
    return doy


def main():
    cfg = load_config()
    y0, y1 = cfg["obs"]["gridded_years"]
    abs_c = float(cfg["index"]["abs_degC"])

    da = clim.load_gridded_obs(cfg["paths"]["gridded_obs"], cfg["obs"]["gridded_var"],
                               float(cfg["obs"]["gridded_fill"]))
    tx90p = xr.open_dataarray(os.path.join(cfg["paths"]["results"], "diagnostics",
                                           "tx90p_climatology.nc"))
    thr_min = clim.ehd_threshold(tx90p, abs_c, "min")

    lat = da.lat.values; lon = da.lon.values
    # gridded time as pandas + doy for threshold lookup
    gt = pd.to_datetime([f"{t.year}-{t.month:02d}-{t.day:02d}" for t in da.time.values])
    gdoy = np.arange(len(gt)) % 365 + 1  # 365_day contiguous
    gobs = da.values                     # (time, lat, lon)

    rows = []
    for p in stn.list_tx_stations(cfg["paths"]["station_dir"]):
        s = stn.load_station(p)
        ser = s["series"]
        ser = ser[(ser.index.year >= y0) & (ser.index.year <= y1)]
        ser = ser[~((ser.index.month == 2) & (ser.index.day == 29))]
        if ser.size < 365:
            continue
        iy = int(np.argmin(np.abs(lat - s["lat"])))
        ix = int(np.argmin(np.abs(lon - s["lon"])))
        gcell = pd.Series(gobs[:, iy, ix], index=gt)
        common = ser.index.intersection(gcell.index)
        if len(common) < 365:
            continue
        st_v = ser.loc[common].values
        gr_v = gcell.loc[common].values
        m = np.isfinite(st_v) & np.isfinite(gr_v)
        if m.sum() < 365:
            continue
        st_v, gr_v = st_v[m], gr_v[m]
        cdoy = doy_of(common[m])
        thr = thr_min.values[cdoy - 1, iy, ix]
        st_ehd = (st_v >= thr); gr_ehd = (gr_v >= thr)
        bias = float(np.mean(st_v - gr_v))
        rows.append(dict(name=s["name"], region=s["region"], lon=s["lon"], lat=s["lat"],
                         elev=s["elev"], n=int(m.sum()),
                         bias_stn_minus_grid=round(bias, 2),
                         mae=round(float(np.mean(np.abs(st_v - gr_v))), 2),
                         rmse=round(float(np.sqrt(np.mean((st_v - gr_v) ** 2))), 2),
                         corr=round(float(np.corrcoef(st_v, gr_v)[0, 1]), 3),
                         ehd_freq_stn=round(float(st_ehd.mean()), 4),
                         ehd_freq_grid=round(float(gr_ehd.mean()), 4)))
    df = pd.DataFrame(rows)
    res = cfg["paths"]["results"]
    df.to_csv(os.path.join(res, "tables", "obs_reconciliation_station.csv"), index=False)
    log.info("matched %d stations", len(df))

    reg = df.groupby("region").agg(
        n_stn=("name", "count"), elev_mean=("elev", "mean"),
        bias=("bias_stn_minus_grid", "mean"), mae=("mae", "mean"),
        rmse=("rmse", "mean"), corr=("corr", "mean"),
        ehd_freq_stn=("ehd_freq_stn", "mean"), ehd_freq_grid=("ehd_freq_grid", "mean"),
    ).round(3).reset_index()
    reg["region_name"] = reg.region.map(REGION_NAMES)
    reg["ehd_freq_diff"] = (reg.ehd_freq_stn - reg.ehd_freq_grid).round(4)
    reg.to_csv(os.path.join(res, "tables", "obs_reconciliation_region.csv"), index=False)
    log.info("per-region reconciliation:\n%s", reg.to_string(index=False))

    # figure: (a) bias vs elevation, (b) region MAE, (c) EHD freq station vs grid
    fig, ax = plt.subplots(1, 3, figsize=(16, 4.5))
    sc = ax[0].scatter(df.elev, df.bias_stn_minus_grid, c=df.region, cmap="tab10", s=18)
    ax[0].axhline(0, color="k", lw=0.6); ax[0].set_xlabel("station elevation (m)")
    ax[0].set_ylabel("Tx bias station - gridded (degC)")
    ax[0].set_title("(a) station-gridded bias vs elevation")
    ax[1].bar(reg.region.astype(str), reg.mae, color="#0072B2")
    ax[1].set_xlabel("region"); ax[1].set_ylabel("mean |station-grid| Tx (degC)")
    ax[1].set_title("(b) daily Tx MAE by region")
    ax[2].scatter(reg.ehd_freq_grid, reg.ehd_freq_stn, c=reg.region, cmap="tab10", s=60)
    mx = float(max(reg.ehd_freq_grid.max(), reg.ehd_freq_stn.max())) * 1.1
    ax[2].plot([0, mx], [0, mx], "k--", lw=0.6)
    for _, r in reg.iterrows():
        ax[2].annotate(str(r.region), (r.ehd_freq_grid, r.ehd_freq_stn), fontsize=8)
    ax[2].set_xlabel("EHD freq gridded"); ax[2].set_ylabel("EHD freq station")
    ax[2].set_title("(c) EHD frequency: station vs gridded")
    fig.tight_layout()
    fig.savefig(os.path.join(res, "figures", "obs_reconciliation.png"), dpi=130)
    log.info("wrote obs_reconciliation.png")

    _write_decision_note(cfg, df, reg)


def _write_decision_note(cfg, df, reg):
    note = f"""# Observation reconciliation — gridded vs station Tx

Compared {len(df)} station Tx records against the nearest 0.25 deg gridded cell over
{cfg['obs']['gridded_years'][0]}-{cfg['obs']['gridded_years'][1]} (common days only; 29 Feb
dropped; EHD via diagnostic min(Tx90p,35) climatology). Full numbers in
`results/tables/obs_reconciliation_*.csv`.

## Per-region summary (station minus gridded)

{reg[['region','region_name','n_stn','elev_mean','bias','mae','rmse','corr','ehd_freq_stn','ehd_freq_grid','ehd_freq_diff']].to_string(index=False)}

## Decision

- **Calibration target = gridded Tx** (`Tx_EC_0.25_2002-2023.nc`). The forecast product is
  gridded over all 541 land cells and every lead block; only the gridded obs provides the
  spatially complete, temporally aligned target the calibration and verification need.
- **Station Tx = independent validation.** Stations are used to audit the gridded product at
  point scale, especially where the two disagree.
- **Where they disagree, station is weighted for point realism, gridded for spatial
  coverage.** The disagreement is dominated by a **general warm offset**: stations read on
  average ~1.5-2.2 degC warmer than their 0.25 deg cell in every region (figure panel a),
  with elevation-dependent scatter (a few high peaks, ~1500 m, run colder). Day-to-day
  correlation stays high (0.83-0.98), so this is a systematic representativeness/offset gap,
  not a timing error. Because the shared threshold is derived from the cooler gridded field,
  the warmer stations exceed it far more often, so **station EHD frequency (0.31-0.59) is 2-5x
  the gridded (~0.10)** (panel c). Consequence for the product: skill and thresholds are fit
  and verified in the gridded world; the station comparison bounds how far point-scale EHD
  risk may exceed the gridded product, and is reported as a limitation.

## Documented data-quality handling

- Missing values (`-99.0`) are dropped, never imputed. Per-station missing fractions are in
  `results/tables/station_metadata.csv`; stations with <365 matched days are excluded here.
- Station periods and lengths vary; matching is done on the intersecting dates only.
- Obvious inhomogeneities (station relocation, instrument change) are not corrected; flagged
  as a limitation for point validation.
"""
    with open(os.path.join(cfg["paths"]["root"], "docs", "OBS_RECONCILIATION.md"), "w") as f:
        f.write(note)
    log.info("wrote docs/OBS_RECONCILIATION.md")


if __name__ == "__main__":
    main()