#!/usr/bin/env python3
"""Point-wise upper-air predictors from the pgb extraction, aggregated by week.

The form matters more than the variable list. Large-scale fields entered as
NATIONAL scalars (EOF principal components) are known to be worthless for this
problem: one number broadcast to every point can only shift the national level,
and the national level is not what is failing. The same fields entered as
POINT-WISE derived quantities -- thickness, layer thermal gradients, moisture
transport, omega, mean sea level pressure -- are what carried the large jump in
skill. So everything here is evaluated at the grid point.

Input : outputs/Pgb2x/<YYYY>/pgb2x_<init>.npy  (46, 2, 40, 31, 26)
Output: outputs/Cube/pgb_weekly.npz  X (ninit, 6, nderived, 541), float32
"""

from __future__ import annotations

import argparse
import multiprocessing as mp
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C
import build_obs as B
import extract_pgb as P

OUTF = os.environ.get("PGB_WEEKLY",
                      os.path.join(C.OUT, "Cube", "pgb_weekly.npz"))
G0 = 9.80665

DERIVED = [
    "thk_1000_500", "thk_1000_850", "thk_850_700", "thk_925_500", "thk_925_700",
    "dt_850_500", "dt_925_850", "dt_700_500",
    "ivt_u", "ivt_v", "ivt",
    "w500", "w700", "w850",
    "prmsl", "cape", "cin",
    "t850", "t700", "h500", "h200",
    "rh850", "rh700", "rh500",
    "shear_u", "shear_v", "wspd200",
]


def _idx(name):
    return P.PGB_NAMES.index(name)


def derive(a):
    """a: (nlead, ntod, nvar, ny, nx) -> (nlead, nderived, ny, nx), tod averaged."""
    m = np.nanmean(a, axis=1)                       # (nlead, nvar, ny, nx)
    g = lambda n: m[:, _idx(n)]
    out = {}
    out["thk_1000_500"] = g("hgt500") - g("hgt1000")
    out["thk_1000_850"] = g("hgt850") - g("hgt1000")
    out["thk_850_700"] = g("hgt700") - g("hgt850")
    out["thk_925_500"] = g("hgt500") - g("hgt925")
    out["thk_925_700"] = g("hgt700") - g("hgt925")
    out["dt_850_500"] = g("tmp850") - g("tmp500")
    out["dt_925_850"] = g("tmp925") - g("tmp850")
    out["dt_700_500"] = g("tmp700") - g("tmp500")
    # crude vertically integrated moisture transport over 1000-500 hPa
    lev = [1000, 925, 850, 700, 500]
    dp = np.array([75., 112.5, 187.5, 275., 100.]) * 100.0     # Pa, midpoint slabs
    iu = sum(g("spfh%d" % p) * g("ugrd%d" % p) * d for p, d in zip(lev, dp)) / G0
    iv = sum(g("spfh%d" % p) * g("vgrd%d" % p) * d for p, d in zip(lev, dp)) / G0
    out["ivt_u"], out["ivt_v"], out["ivt"] = iu, iv, np.hypot(iu, iv)
    out["w500"], out["w700"], out["w850"] = g("vvel500"), g("vvel700"), g("vvel850")
    out["prmsl"], out["cape"], out["cin"] = g("prmsl"), g("cape"), g("cin")
    out["t850"], out["t700"] = g("tmp850"), g("tmp700")
    out["h500"], out["h200"] = g("hgt500"), g("hgt200")
    out["rh850"], out["rh700"], out["rh500"] = g("rh850"), g("rh700"), g("rh500")
    out["shear_u"] = g("ugrd200") - g("ugrd850")
    out["shear_v"] = g("vgrd200") - g("vgrd850")
    out["wspd200"] = np.hypot(g("ugrd200"), g("vgrd200"))
    return np.stack([out[k] for k in DERIVED], axis=1)


_G = {}


def _init(pj, pi):
    _G["pj"], _G["pi"] = pj, pi


def _one(init):
    """Weekly means for one init, or None when that init has no pgb file.

    Runs in a worker process: the loop over 3086 inits is the whole cost here and
    the inits are independent, so it is spread over the node's cores. Each task
    reads one 12 MB .npy from /work, so this scales with cores, not with the NAS.
    """
    f = P.pgb_path(init)
    if not os.path.exists(f):
        return None
    d = derive(np.load(f))[:, :, _G["pj"], _G["pi"]]          # (46, nd, npt)
    out = np.full((len(C.WEEKS), d.shape[1], len(_G["pj"])), np.nan, dtype="float32")
    for w, (_, idx) in enumerate(C.WEEKS):
        with np.errstate(invalid="ignore"):
            out[w] = np.nanmean(d[idx], axis=0)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--nproc", type=int,
                    default=int(os.environ.get("SLURM_CPUS_PER_TASK", 8)))
    a = ap.parse_args()

    inits = C.both_inits()
    z = np.load(B.TX_NPZ)
    olat, olon = z["lat"][z["jj"]], z["lon"][z["ii"]]
    gp = P.pgb_grid()
    pj = np.abs(olat[:, None] - gp["wlat"][None, :]).argmin(1)
    pi = np.abs(olon[:, None] - gp["wlon"][None, :]).argmin(1)
    npt, nw, nd = len(pj), len(C.WEEKS), len(DERIVED)

    X = np.full((len(inits), nw, nd, npt), np.nan, dtype="float32")
    have = np.zeros(len(inits), dtype=bool)
    print("gop tuan %d init tren %d tien trinh" % (len(inits), a.nproc), flush=True)
    with mp.Pool(a.nproc, initializer=_init, initargs=(pj, pi)) as pool:
        for i, r in enumerate(pool.imap(_one, inits, chunksize=8)):
            if r is not None:
                X[i] = r
                have[i] = True
            if (i + 1) % 500 == 0:
                print("  %d/%d (have %d)" % (i + 1, len(inits), int(have.sum())),
                      flush=True)
    np.savez(OUTF, X=X, have=have, names=np.array(DERIVED), inits=np.array(inits))
    print("saved %s  inits with pgb: %d/%d" % (OUTF, int(have.sum()), len(inits)))


if __name__ == "__main__":
    main()