#!/usr/bin/env python3
"""Downscale forecast Tx from the 1.5 deg EC grid to the 0.25 deg Vietnam grid.

Methods, in increasing order of sophistication:
  bilinear  plain bilinear interpolation, no correction        (absolute floor)
  qm        bilinear + empirical quantile mapping per (point, forecast week)
  bcsd      textbook BCSD: bias-correct at COARSE scale against coarsened obs,
            convert to anomaly, interpolate the anomaly, add the FINE obs
            climatology back (this is what "spatial disaggregation" means)

Validation is k-fold by year: each fold holds out 4 whole years, and every
statistic a method needs -- quantile maps, climatologies, thresholds -- is rebuilt
from that fold's training years only.

Usage:
  downscale.py --method bilinear,qm,bcsd [--folds 5] [--out outputs/downscale/...]
"""
import argparse
import json
import os
import sys

import numpy as np
import xarray as xr

ROOT = "/home/tanpv/Progs/claude/HW_EC_10"
COARSE = f"{ROOT}/outputs/proxy/tx_ec_1.5_proxy.nc"
OBS = f"{ROOT}/outputs/OBS_4EC15/tx_obs_aligned.nc"
STATIC = f"{ROOT}/outputs/Cli/static_fields_0.25.nc"

WEEKS = {"W1": (0, 11), "W2": (11, 18), "W3": (18, 25),
         "W4": (25, 32), "W5": (32, 39), "W6": (39, 46)}
YEARS = list(range(2003, 2023))


# ----------------------------------------------------------------- interpolation
def bilinear_matrix(clat, clon, flat, flon):
    """Sparse-ish (nfine, ncoarse) bilinear weight matrix, edge-clamped."""
    order = np.argsort(clat)
    cl = clat[order]
    W = np.zeros((len(flat) * len(flon), len(clat) * len(clon)), dtype="float32")

    jy = np.clip(np.searchsorted(cl, flat) - 1, 0, len(cl) - 2)
    ty = (flat - cl[jy]) / (cl[jy + 1] - cl[jy])
    ty = np.clip(ty, 0, 1)
    ix = np.clip(np.searchsorted(clon, flon) - 1, 0, len(clon) - 2)
    tx = (flon - clon[ix]) / (clon[ix + 1] - clon[ix])
    tx = np.clip(tx, 0, 1)

    for a, j in enumerate(jy):
        j0, j1 = order[j], order[j + 1]
        for b, i in enumerate(ix):
            r = a * len(flon) + b
            wy, wx = ty[a], tx[b]
            W[r, j0 * len(clon) + i] = (1 - wy) * (1 - wx)
            W[r, j0 * len(clon) + i + 1] = (1 - wy) * wx
            W[r, j1 * len(clon) + i] = wy * (1 - wx)
            W[r, j1 * len(clon) + i + 1] = wy * wx
    return W


def fill_coarse_nan(cx, clat, clon):
    """Nearest-neighbour fill of coarse cells that are empty for every field.

    Needed only for the development proxy, whose 1.5 deg cells at the edge of the
    domain have no 0.25 deg source points. The real MARS 1.5 deg data fills every
    cell, so this is a no-op there. It matters because NaN * 0 = NaN inside the
    interpolation matmul: a single empty cell would otherwise poison every output.
    """
    dead = np.isnan(cx).all(axis=0)
    if not dead.any():
        return cx, 0
    la, lo = np.meshgrid(clat, clon, indexing="ij")
    la, lo = la.ravel(), lo.ravel()
    good = np.where(~dead)[0]
    for k in np.where(dead)[0]:
        d = (la[good] - la[k]) ** 2 + (lo[good] - lo[k]) ** 2
        cx[:, k] = cx[:, good[d.argmin()]]
    return cx, int(dead.sum())


def coarsen_matrix(clat, clon, flat, flon, land=None):
    """(ncoarse, nfine) averaging matrix: which fine points fall in each coarse cell.

    `land` restricts the average to land points. That restriction is essential:
    averaging observations over a coarse cell that is mostly sea would mix land
    temperatures with missing values and produce a meaningless "coarse observation".
    Rows for cells containing no land point are returned as all-zero and must be
    excluded by the caller.
    """
    jl = np.abs(flat[:, None] - clat[None, :]).argmin(axis=1)
    il = np.abs(flon[:, None] - clon[None, :]).argmin(axis=1)
    M = np.zeros((len(clat) * len(clon), len(flat) * len(flon)), dtype="float32")
    for a in range(len(flat)):
        for b in range(len(flon)):
            r = a * len(flon) + b
            if land is not None and not land[r]:
                continue
            M[jl[a] * len(clon) + il[b], r] = 1.0
    M /= np.maximum(M.sum(axis=1, keepdims=True), 1)
    return M


# ------------------------------------------------------------------- corrections
def quantile_map(train_src, train_dst, x, nq=101):
    """Map x through the empirical CDF of train_src onto that of train_dst."""
    qs = np.linspace(0, 100, nq)
    src_q = np.nanpercentile(train_src, qs)
    dst_q = np.nanpercentile(train_dst, qs)
    src_q = np.maximum.accumulate(src_q)          # guard against ties
    return np.interp(x, src_q, dst_q)


def apply_qm(pred, obs, train, test, land_idx):
    """Per (land point, week) quantile mapping fitted on the training years."""
    out = pred.copy()
    for _, (l0, l1) in WEEKS.items():
        ptr = pred[train][:, :, l0:l1]            # (ntr, ninit, nlead, npts)
        otr = obs[train][:, :, l0:l1]
        pte = pred[test][:, :, l0:l1]
        for p in land_idx:
            s = ptr[..., p].ravel()
            d = otr[..., p].ravel()
            d = d[np.isfinite(d)]
            s = s[np.isfinite(s)]
            if len(d) < 50 or len(s) < 50:
                continue
            out[test, :, l0:l1, p] = quantile_map(s, d, pte[..., p])
    return out


# -------------------------------------------------------------------- evaluation
def metrics(pred, obs):
    """ME / MAE / RMSE / CORR over the finite entries of two matching slices."""
    p, o = pred.ravel(), obs.ravel()
    ok = np.isfinite(p) & np.isfinite(o)
    p, o = p[ok], o[ok]
    if len(p) < 10:
        return dict(n=int(len(p)), ME=np.nan, MAE=np.nan, RMSE=np.nan, CORR=np.nan)
    d = p - o
    return dict(n=int(len(p)), ME=float(d.mean()), MAE=float(np.abs(d).mean()),
                RMSE=float(np.sqrt((d ** 2).mean())),
                CORR=float(np.corrcoef(p, o)[0, 1]))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--method", default="bilinear,qm,bcsd")
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--coarse", default=COARSE)
    ap.add_argument("--out", default=f"{ROOT}/outputs/downscale/baselines.json")
    args = ap.parse_args()
    methods = args.method.split(",")

    c = xr.open_dataset(args.coarse)
    o = xr.open_dataset(OBS)
    st = xr.open_dataset(STATIC)

    clat, clon = c.lat.values, c.lon.values
    flat, flon = o.lat.values, o.lon.values
    nyear, ninit, nlead = c.sizes["year"], c.sizes["init"], c.sizes["lead"]
    npts = len(flat) * len(flon)
    print(f"coarse {len(clat)}x{len(clon)} -> fine {len(flat)}x{len(flon)}  "
          f"({nyear}y x {ninit}init x {nlead}lead)", flush=True)

    land = (np.isfinite(st.vn.values) & (st.vn.values > 0)).ravel()
    land_idx = np.where(land)[0]
    print(f"land points: {land.sum()}", flush=True)

    W = bilinear_matrix(clat, clon, flat, flon)
    cx = c.tx_ec.values.reshape(-1, len(clat) * len(clon))
    cx, ndead = fill_coarse_nan(cx, clat, clon)
    if ndead:
        print(f"filled {ndead} always-empty coarse cells (proxy edge artefact)",
              flush=True)
    bil = (cx @ W.T).reshape(nyear, ninit, nlead, npts).astype("float32")
    del cx
    obs = o.tx_obs.values.reshape(nyear, ninit, nlead, npts)
    print(f"bilinear done, arrays {bil.nbytes/1e9:.2f} GB each", flush=True)

    # k-fold by year: contiguous blocks of 4 years
    per = nyear // args.folds
    folds = [list(range(k * per, (k + 1) * per)) for k in range(args.folds)]

    results = {}
    for m in methods:
        pred = bil.copy() if m != "bilinear" else bil
        if m == "qm":
            for f in folds:
                tr = [i for i in range(nyear) if i not in f]
                pred = apply_qm(pred, obs, tr, f, land_idx)
                print(f"  qm fold {f[0]//per+1}/{args.folds} done", flush=True)
        elif m == "bcsd":
            # Textbook BCSD:
            #   BC  quantile-map the coarse forecast onto coarse observations
            #   SD  express the corrected field as an anomaly from the COARSE OBS
            #       climatology, interpolate the anomaly, add the FINE OBS climatology
            # Both climatologies come from the training years of the fold.
            Mc = coarsen_matrix(clat, clon, flat, flon, land=land)
            has_land = Mc.sum(axis=1) > 0
            cells = np.where(has_land)[0]
            print(f"  bcsd: {len(cells)} of {Mc.shape[0]} coarse cells contain land",
                  flush=True)

            obs_c = np.einsum("tp,cp->tc", np.nan_to_num(obs.reshape(-1, npts)),
                              Mc).reshape(nyear, ninit, nlead, -1)
            obs_c[..., ~has_land] = np.nan
            cf = c.tx_ec.values.reshape(nyear, ninit, nlead, -1).copy()
            cf, _ = fill_coarse_nan(cf.reshape(-1, cf.shape[-1]), clat, clon)
            cf = cf.reshape(nyear, ninit, nlead, -1)

            for f in folds:
                tr = np.array([i for i in range(nyear) if i not in f])
                corr = cf.copy()
                for _, (l0, l1) in WEEKS.items():
                    for k in cells:
                        s_ = cf[tr][:, :, l0:l1, k].ravel()
                        d_ = obs_c[tr][:, :, l0:l1, k].ravel()
                        s_, d_ = s_[np.isfinite(s_)], d_[np.isfinite(d_)]
                        if len(s_) < 50 or len(d_) < 50:
                            continue
                        for yy in (tr, np.array(f)):        # correct train AND test
                            corr[yy[:, None, None], np.arange(ninit)[None, :, None],
                                 np.arange(l0, l1)[None, None, :], k] = quantile_map(
                                s_, d_, cf[yy][:, :, l0:l1, k]).reshape(
                                    len(yy), ninit, l1 - l0)

                cclim = np.nanmean(obs_c[tr], axis=(0, 1))            # coarse OBS clim
                fclim = np.nanmean(obs[tr], axis=(0, 1))              # fine OBS clim
                anom = corr[f] - cclim[None, None, :, :]
                anom = np.nan_to_num(anom)                # ocean-only cells -> 0 anomaly
                pred[f] = (anom.reshape(-1, anom.shape[-1]) @ W.T
                           ).reshape(len(f), ninit, nlead, npts) + fclim[None, None]
                print(f"  bcsd fold {f[0]//per+1}/{args.folds} done", flush=True)

        # keep the prediction so event scores can be recomputed without refitting
        pdir = os.path.join(os.path.dirname(args.out), "pred")
        os.makedirs(pdir, exist_ok=True)
        np.save(os.path.join(pdir, f"pred_{m}.npy"),
                pred[..., land_idx].astype("float32"))

        res = {}
        for wname, (l0, l1) in WEEKS.items():
            res[wname] = metrics(pred[:, :, l0:l1][..., land_idx],
                                 obs[:, :, l0:l1][..., land_idx])
        res["ALL"] = metrics(pred[..., land_idx], obs[..., land_idx])
        results[m] = res
        print(f"{m}: " + "  ".join(
            f"{w} RMSE={res[w]['RMSE']:.2f}" for w in list(WEEKS) + ["ALL"]), flush=True)

    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    with open(args.out, "w") as fh:
        json.dump(results, fh, indent=2)
    print(f"wrote {args.out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())