#!/usr/bin/env python3
"""AI downscaling of forecast Tx from 1.5 deg to the 0.25 deg Vietnam grid.

The model learns the RESIDUAL  r = Tx_obs - bilinear(Tx_ec)  rather than Tx itself,
so the interpolation keeps the large-scale signal and the learner only has to supply
the correction the coarse grid cannot represent. Predictions are bilinear + r_hat.

Features (per fine point, init, lead):
  bil          bilinearly interpolated coarse Tx
  bil_anom     bil minus its own training climatology for that (week, doy)
  clim_delta   obs climatology - bil climatology: the systematic correction
  dz           orog - orog_coarse, the elevation the coarse model misses
  orog, orog_std, landfrac, lat, lon
  lead, doy_sin, doy_cos
  c00..c33     the 3x3 coarse Tx patch around the point (local gradient / advection)

Every climatology is rebuilt from the training years of the fold, so nothing about a
test year reaches the model. Validation is k-fold by whole years.

Usage:
  downscale_ai.py [--folds 5] [--subsample 4000000] [--out ...]
"""
import argparse
import json
import os
import sys

import numpy as np
import xarray as xr
import lightgbm as lgb

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)]
NDOY = 365
HALF = 15                       # 31-day climatology window

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from downscale import bilinear_matrix, fill_coarse_nan, metrics   # noqa: E402


def window_climatology(values, doy, npts):
    """31-day circular-window mean of `values` binned by day of year.

    values: (nsamp, npts); doy: (nsamp,) 1..365. Returns (365, npts).
    Binning once and then smoothing the sums is O(N) -- gathering the window
    explicitly for each of the 365 days would be orders of magnitude slower.
    """
    s = np.zeros((NDOY, npts), dtype="float64")
    c = np.zeros((NDOY, npts), dtype="float64")
    v = np.nan_to_num(values, nan=0.0)
    ok = np.isfinite(values).astype("float64")
    np.add.at(s, doy - 1, v)
    np.add.at(c, doy - 1, ok)
    ker = np.ones(2 * HALF + 1)
    pad = lambda a: np.concatenate([a[-HALF:], a, a[:HALF]], axis=0)   # noqa: E731
    ss = np.apply_along_axis(lambda x: np.convolve(x, ker, "valid"), 0, pad(s))
    cc = np.apply_along_axis(lambda x: np.convolve(x, ker, "valid"), 0, pad(c))
    with np.errstate(invalid="ignore", divide="ignore"):
        out = ss / cc
    out[cc == 0] = np.nan
    return out.astype("float32")


def coarse_patch(cx, clat, clon, flat, flon):
    """For every fine point, the 3x3 coarse cells centred on its nearest cell."""
    jl = np.abs(flat[:, None] - clat[None, :]).argmin(axis=1)
    il = np.abs(flon[:, None] - clon[None, :]).argmin(axis=1)
    nlc, nc = len(clat), len(clon)
    idx = np.empty((len(flat) * len(flon), 9), dtype="int32")
    for a in range(len(flat)):
        for b in range(len(flon)):
            r = a * len(flon) + b
            k = 0
            for dj in (-1, 0, 1):
                for di in (-1, 0, 1):
                    j = np.clip(jl[a] + dj, 0, nlc - 1)
                    i = np.clip(il[b] + di, 0, nc - 1)
                    idx[r, k] = j * nc + i
                    k += 1
    return idx


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--subsample", type=int, default=4_000_000)
    ap.add_argument("--coarse", default=COARSE)
    ap.add_argument("--out", default=f"{ROOT}/outputs/downscale/ai_lgbm.json")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    rng = np.random.default_rng(args.seed)

    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)

    land = (np.isfinite(st.vn.values) & (st.vn.values > 0)).ravel()
    land_idx = np.where(land)[0]
    nland = len(land_idx)
    print(f"coarse {len(clat)}x{len(clon)} -> fine {len(flat)}x{len(flon)}, "
          f"{nland} land points", flush=True)

    W = bilinear_matrix(clat, clon, flat, flon)
    cxf = c.tx_ec.values.reshape(-1, len(clat) * len(clon))
    cxf, ndead = fill_coarse_nan(cxf, clat, clon)
    if ndead:
        print(f"filled {ndead} empty coarse cells (proxy edge artefact)", flush=True)
    bil = (cxf @ W.T).reshape(nyear, ninit, nlead, npts).astype("float32")
    obs = o.tx_obs.values.reshape(nyear, ninit, nlead, npts)
    doy = o.doy.values                                    # (nyear, ninit, nlead)

    # static features, land points only
    stat = np.stack([st.dz.values.ravel(), st.orog.values.ravel(),
                     st.orog_std.values.ravel(), st.landfrac.values.ravel(),
                     st.lat2d.values.ravel(), st.lon2d.values.ravel()],
                    axis=1)[land_idx].astype("float32")

    pidx = coarse_patch(cxf, clat, clon, flat, flon)[land_idx]      # (nland, 9)
    cx4 = cxf.reshape(nyear, ninit, nlead, -1)

    bil_l = bil[..., land_idx]
    obs_l = obs[..., land_idx]
    del bil, obs, cxf

    per = nyear // args.folds
    folds = [list(range(k * per, (k + 1) * per)) for k in range(args.folds)]
    pred_l = np.full_like(bil_l, np.nan)

    FEATNAMES = (["bil", "bil_anom", "clim_delta", "dz", "orog", "orog_std",
                  "landfrac", "lat", "lon", "lead", "doy_sin", "doy_cos"]
                 + [f"c{k}" for k in range(9)])

    gains = np.zeros(len(FEATNAMES))
    for fi, te in enumerate(folds, 1):
        tr = np.array([i for i in range(nyear) if i not in te])
        te = np.array(te)

        # ---- climatologies from training years only
        wk_of_lead = np.zeros(nlead, dtype=int)
        for w, (_, l0, l1) in enumerate(WEEKS):
            wk_of_lead[l0:l1] = w
        clim_obs = np.empty((len(WEEKS), NDOY, nland), dtype="float32")
        clim_bil = np.empty((len(WEEKS), NDOY, nland), dtype="float32")
        for w, (_, l0, l1) in enumerate(WEEKS):
            d = doy[tr][:, :, l0:l1].ravel()
            clim_obs[w] = window_climatology(
                obs_l[tr][:, :, l0:l1].reshape(-1, nland), d, nland)
            clim_bil[w] = window_climatology(
                bil_l[tr][:, :, l0:l1].reshape(-1, nland), d, nland)
        print(f"fold {fi}: climatologies built", flush=True)

        # Rows are (year, init, lead) triples; every land point of a chosen row is
        # used. Sampling whole rows instead of scattered cells keeps the spatial
        # structure intact and, more importantly, lets the feature matrix be built
        # one chunk at a time -- materialising all 41M training rows at once would
        # need ~3.5 GB.
        def build_rows(years, rows):
            """Feature matrix for the given flat (year,init,lead) row indices."""
            b = bil_l[years].reshape(-1, nland)[rows]
            ob = obs_l[years].reshape(-1, nland)[rows]
            dd = doy[years].reshape(-1)[rows]
            ll = np.tile(np.arange(nlead), len(years) * ninit)[rows]
            wk = wk_of_lead[ll]
            cb = clim_bil[wk, dd - 1]
            co = clim_obs[wk, dd - 1]
            patch = cx4[years].reshape(-1, cx4.shape[-1])[rows][:, pidx]

            n = len(rows) * nland
            X = np.empty((n, len(FEATNAMES)), dtype="float32")
            X[:, 0] = b.ravel()
            X[:, 1] = (b - cb).ravel()
            X[:, 2] = (co - cb).ravel()
            for k in range(6):
                X[:, 3 + k] = np.tile(stat[:, k], len(rows))
            X[:, 9] = np.repeat(ll / 45.0, nland)
            X[:, 10] = np.repeat(np.sin(2 * np.pi * dd / NDOY), nland)
            X[:, 11] = np.repeat(np.cos(2 * np.pi * dd / NDOY), nland)
            for k in range(9):
                X[:, 12 + k] = patch[:, :, k].ravel()
            return X, (ob - b).ravel(), b.ravel()

        ntime_tr = len(tr) * ninit * nlead
        nrows_want = max(1, args.subsample // nland)
        sel_rows = (rng.choice(ntime_tr, min(nrows_want, ntime_tr), replace=False)
                    if nrows_want < ntime_tr else np.arange(ntime_tr))
        Xtr, ytr, _ = build_rows(tr, np.sort(sel_rows))

        idx = np.where(np.isfinite(ytr) & np.isfinite(Xtr).all(axis=1))[0]
        print(f"fold {fi}: {len(sel_rows):,} of {ntime_tr:,} time rows -> "
              f"{len(idx):,} valid training samples", flush=True)

        model = lgb.LGBMRegressor(
            n_estimators=700, learning_rate=0.05, num_leaves=127,
            min_child_samples=200, subsample=0.8, subsample_freq=1,
            colsample_bytree=0.8, reg_lambda=1.0, n_jobs=8, verbose=-1)
        model.fit(Xtr[idx], ytr[idx], feature_name=FEATNAMES)
        del Xtr, ytr

        g = model.booster_.feature_importance(importance_type="gain")
        gains += g / g.sum()

        # predict in chunks of rows so the test matrix never has to fit whole
        ntime_te = len(te) * ninit * nlead
        out = np.empty((ntime_te, nland), dtype="float32")
        CH = 2000
        for a in range(0, ntime_te, CH):
            rows = np.arange(a, min(a + CH, ntime_te))
            Xte, _, bte = build_rows(te, rows)
            okte = np.isfinite(Xte).all(axis=1)
            r = np.full(len(Xte), np.nan, dtype="float32")
            if okte.any():
                r[okte] = model.predict(Xte[okte])
            out[a:a + len(rows)] = (bte + r).reshape(len(rows), nland)
        pred_l[te] = out.reshape(len(te), ninit, nlead, nland)
        del out
        print(f"fold {fi}/{args.folds} done", flush=True)

    pdir = os.path.join(os.path.dirname(args.out), "pred")
    os.makedirs(pdir, exist_ok=True)

    results = {}
    for tag, arr in [("bilinear", bil_l), ("lgbm", pred_l)]:
        np.save(os.path.join(pdir, f"pred_{tag}.npy"), arr.astype("float32"))
        res = {}
        for wname, l0, l1 in WEEKS:
            res[wname] = metrics(arr[:, :, l0:l1], obs_l[:, :, l0:l1])
        res["ALL"] = metrics(arr, obs_l)
        results[tag] = res
        print(f"{tag}: " + "  ".join(
            f"{w} RMSE={res[w]['RMSE']:.2f}" for w, _, _ in WEEKS) +
            f"  ALL RMSE={res['ALL']['RMSE']:.2f}", flush=True)

    gains /= args.folds
    order = np.argsort(-gains)
    results["feature_gain"] = {FEATNAMES[i]: float(gains[i]) for i in order}
    print("\nfeature gain (mean over folds):")
    for i in order[:12]:
        print(f"  {FEATNAMES[i]:12s} {gains[i]:.4f}")

    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())