#!/usr/bin/env python3
"""Stage 2, table 2 rows 43-62: the ten leading EOF modes of H500 and T850.

The wide-domain (80-150E, -15-45N, 1 deg) forecast fields carry the large-scale
circulation that local surface variables cannot express.  For each variable:

  1. read every (init, lt) field, keeping steps 0..45 so they line up with
     EC_Combined, and flip latitude to ascending;
  2. remove the model climatology, computed per (lead time, day of year) on the
     same 31-day window as everywhere else, so both the seasonal cycle and the
     model drift with lead time are taken out;
  3. weight by sqrt(cos(lat)) so equal-area is respected;
  4. build the space-space covariance over the TRAINING YEARS ONLY and take its
     leading eigenvectors -- the validation years never touch the basis;
  5. project every sample, training and validation alike, onto the first ten
     modes and standardise each PC by its training standard deviation.

Output: pcs (ninit, 46, 10) per variable, plus the patterns and explained
variance for later inspection.
"""

from __future__ import annotations

import argparse
import os
import sys

import numpy as np
from joblib import Parallel, delayed
from netCDF4 import Dataset

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import build_cube as B
from calc_climatology import bin_by_doy, moments

WIDE = {"H500": ("gh", os.path.join(C.DATA, "EC", "H500"), "H500_cf_%s.nc"),
        "T850": ("t", os.path.join(C.DATA, "EC", "T850"), "T850_cf_%s.nc")}
NMODE = 10
EOF_DIR = os.path.join(C.OUT, "EOF")


def _read_year(year, varname, folder, pattern):
    """(ninit_year, 46, nlat, nlon) with latitude flipped to ascending."""
    out = None
    for i, init in enumerate(C.init_dates(year)):
        path = os.path.join(folder, str(year), pattern % init.isoformat())
        with Dataset(path) as ds:
            a = np.ma.filled(ds[varname][:C.NLT], np.nan).astype("f4")[:, ::-1, :]
        if out is None:
            out = np.empty((len(C.init_dates(year)),) + a.shape, dtype="f4")
        out[i] = a
    return out


def load_field(name, jobs):
    varname, folder, pattern = WIDE[name]
    parts = Parallel(n_jobs=jobs, verbose=5)(
        delayed(_read_year)(y, varname, folder, pattern) for y in C.YEARS)
    field = np.concatenate(parts, axis=0)
    with Dataset(os.path.join(folder, str(C.YEARS[0]),
                              pattern % C.init_dates(C.YEARS[0])[0].isoformat())) as ds:
        lat = np.asarray(ds["latitude"][:])[::-1]
        lon = np.asarray(ds["longitude"][:])
    return field, lat, lon


def anomalies(field, doy, keep=None):
    """Subtract the per (lead time, day of year) climatology in place.

    `keep` is the training-year mask.  The climatology removed here is a fitted
    quantity like any other, so it is accumulated over the training years only and
    then applied to every sample -- the same rule the EOF basis itself follows.
    """
    ninit, nlt, nlat, nlon = field.shape
    ng = nlat * nlon
    flat = field.reshape(ninit, nlt, ng)
    for k in range(nlt):
        col, d = flat[:, k, :], doy[:, k]
        if keep is not None:
            s0, s1, s2 = bin_by_doy(col[keep], d[keep], ng)
        else:
            s0, s1, s2 = bin_by_doy(col, d, ng)
        mean, _ = moments(s0, s1, s2)
        flat[:, k, :] -= mean[d - 1]
    return flat


def eof_basis(x_train, nmode, block=4096):
    """Leading eigenvectors of the space-space covariance of x_train (ns, ng).

    The covariance is accumulated in float64 over row blocks so a full float64
    copy of the sample matrix is never materialised.
    """
    ns, ng = x_train.shape
    cov = np.zeros((ng, ng), dtype="f8")
    for i in range(0, ns, block):
        b = x_train[i:i + block].astype("f8")
        cov += b.T @ b
    cov /= max(ns - 1, 1)
    vals, vecs = np.linalg.eigh(cov)
    order = np.argsort(vals)[::-1][:nmode]
    return vecs[:, order].astype("f4"), vals[order], float(np.trace(cov))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--jobs", type=int, default=20)
    ap.add_argument("--vars", nargs="*", default=list(WIDE))
    args = ap.parse_args()

    meta = B.load_meta()
    doy = meta["doy"]
    init_year = meta["init_year"]
    ninit = int(meta["ninit"][0])
    os.makedirs(EOF_DIR, exist_ok=True)

    train_init = np.isin(init_year, C.TRAIN_YEARS)
    print("[split] %d train / %d valid initialisations"
          % (train_init.sum(), (~train_init).sum()), flush=True)

    for name in args.vars:
        print("[%s] reading wide-domain fields ..." % name, flush=True)
        field, lat, lon = load_field(name, args.jobs)
        assert field.shape[0] == ninit, (field.shape, ninit)
        nlt, nlat, nlon = field.shape[1:]
        ng = nlat * nlon
        print("[%s] %s = %.2f GB" % (name, field.shape, field.nbytes / 1e9), flush=True)

        flat = anomalies(field, doy, train_init)           # view on `field`
        w = np.sqrt(np.cos(np.deg2rad(lat))).astype("f4")
        flat *= np.repeat(w, nlon)[None, None, :]

        x = flat.reshape(ninit * nlt, ng)
        train_mask = np.repeat(train_init, nlt)
        vecs, vals, total = eof_basis(x[train_mask], NMODE)
        frac = vals / total
        print("[%s] explained variance: %s (sum %.3f)"
              % (name, " ".join("%.3f" % f for f in frac), frac.sum()), flush=True)

        pcs = (x @ vecs).reshape(ninit, nlt, NMODE)
        sd = pcs.reshape(-1, NMODE)[train_mask].std(axis=0)
        pcs /= np.maximum(sd, 1e-9)

        np.savez(os.path.join(EOF_DIR, "eof_%s.npz" % name),
                 pcs=pcs.astype("f4"),
                 patterns=vecs.reshape(nlat, nlon, NMODE),
                 eigenvalues=vals, explained=frac, lat=lat, lon=lon,
                 train_years=np.array(C.TRAIN_YEARS),
                 names=np.array(["%s%02d" % (name.lower(), i + 1) for i in range(NMODE)]))
        del field, flat, x, pcs
        print("[%s] saved" % name, flush=True)

    print("[done] -> %s" % EOF_DIR, flush=True)


if __name__ == "__main__":
    main()