#!/usr/bin/env python3
"""Build the EC_0.25_SD cube in the exact ../HW build_cube.py format.

This is the ONLY place the EC side differs from ../HW. Everything downstream
(calc_climatology, calc_thresholds, calc_eof, build_ai_inputs, ai_data.Store,
train_national, operational_*, final_verify) then runs on the standard cube.

EC_0.25_SD = the 40 EC_VARS on the 0.25 deg grid, 541 land points, where:
  * mx2t6  = AI/statistically downscaled forecast Tx (0.25 deg, 541 pts) -- the
             defining product; taken from --tx-ds (default: real BCSD).
  * mn2t6  = 1.5 deg daily Tn (timselmin of the four 6-hourly mn2t6), bilinearly
             interpolated to the 541 land points.
  * the other 38 vars = raw 1.5 deg fields (ec15_predictors.nc, already
             de-accumulated) plus the derived fields (dtd2m, srn, bowen, uv10,
             thk_*, ivt_*, dt_*), computed with the exact PrePro/ope1_process
             formulas at 1.5 deg, then bilinearly interpolated to 541 pts.

Because every EC variable except mx2t6 is z-scored downstream against a
climatology recomputed FROM THIS SAME CUBE, unit conventions and the tiny
non-commutativity of "derive-then-interpolate" vs "interpolate-then-derive"
wash out. mx2t6 is in deg C (the downscaled product is obs-calibrated), which is
what sv_mod = mx2t6 - Tx_thres_mod requires.

Output (into C.OUT/Cube, i.e. outputs_sd/Cube): ec_raw.f4, obs_raw.f4, meta.npz
-- byte-compatible with build_cube.open_cubes / load_meta.
"""
from __future__ import annotations

import argparse
import glob
import os
import sys
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import xarray as xr
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   # EC_VARS, OBS_VARS, LOG_VARS, cube_shapes, paths

HERE_DATA = "/home/tanpv/Progs/claude/HW_EC_10"
EC15 = os.path.join(HERE_DATA, "outputs/real/ec15_predictors.nc")
TXDS_DEFAULT = os.path.join(HERE_DATA, "outputs/downscale_real/pred/pred_bcsd.npy")
TXTM_ROOT = os.path.join(HERE_DATA, "data/nc/txtm")
HDATES = "/home/tanpv/Data/QG.25.196/hdate_2023.txt"

EC15_LON = 100.0 + 1.5 * np.arange(10)          # 100.0 .. 113.5
EC15_LAT = 24.0 - 1.5 * np.arange(13)           # 24.0 .. 6.0 (descending)
NLEAD = C.NLT
NSTEP = 4


# --------------------------------------------------------------------------- #
#  bilinear weights: 1.5 deg grid -> the 541 land points, edge-clamped
# --------------------------------------------------------------------------- #
def bilinear_weights(plat, plon):
    """(npt, 13*10) weights mapping the flattened 1.5 deg grid onto points."""
    order = np.argsort(EC15_LAT)                 # ascending copy for searchsorted
    cl = EC15_LAT[order]
    npt = plat.size
    W = np.zeros((npt, EC15_LAT.size * EC15_LON.size), dtype="f4")
    jy = np.clip(np.searchsorted(cl, plat) - 1, 0, cl.size - 2)
    ty = np.clip((plat - cl[jy]) / (cl[jy + 1] - cl[jy]), 0, 1)
    ix = np.clip(np.searchsorted(EC15_LON, plon) - 1, 0, EC15_LON.size - 2)
    tx = np.clip((plon - EC15_LON[ix]) / (EC15_LON[ix + 1] - EC15_LON[ix]), 0, 1)
    nlon = EC15_LON.size
    for a in range(npt):
        j0, j1 = order[jy[a]], order[jy[a] + 1]
        i, wy, wx = ix[a], ty[a], tx[a]
        W[a, j0 * nlon + i]     = (1 - wy) * (1 - wx)
        W[a, j0 * nlon + i + 1] = (1 - wy) * wx
        W[a, j1 * nlon + i]     = wy * (1 - wx)
        W[a, j1 * nlon + i + 1] = wy * wx
    return W


_FILL_MAP = None   # (130,) nearest-valid index, lazily built per NaN pattern


def _nearest_valid_map(bad):
    """For each flat 1.5 deg cell, index of the nearest cell that is never bad."""
    la, lo = np.meshgrid(EC15_LAT, EC15_LON, indexing="ij")
    la, lo = la.ravel(), lo.ravel()
    good = np.where(~bad)[0]
    m = np.arange(bad.size)
    for k in np.where(bad)[0]:
        d = (la[good] - la[k]) ** 2 + (lo[good] - lo[k]) ** 2
        m[k] = good[d.argmin()]
    return m


def fill_coarse(field15):
    """Nearest-neighbour fill of cells that are NaN anywhere in time (sea cells of
    land-only fields such as soil moisture/temperature). Without it a single NaN
    neighbour poisons the whole bilinear stencil (NaN*0 = NaN) and every land
    point downstream comes out NaN."""
    flat = field15.reshape(-1, EC15_LAT.size * EC15_LON.size)   # (N, 130)
    bad = ~np.isfinite(flat).all(axis=0)
    if not bad.any():
        return field15
    m = _nearest_valid_map(bad)
    flat = flat[:, m]
    return flat.reshape(field15.shape)


def interp_field(field15, W):
    """(..., 13, 10) -> (..., npt) via the precomputed weight matrix."""
    field15 = fill_coarse(field15)
    flat = field15.reshape(field15.shape[:-2] + (-1,))
    return flat @ W.T


# --------------------------------------------------------------------------- #
#  derived variables at 1.5 deg, following PrePro/ope1_process exactly
# --------------------------------------------------------------------------- #
def build_derived(v):
    """v: dict name -> (Y,I,L,13,10). Adds the derived EC_VARS in place, returns v."""
    v["dtd2m"] = v["t2m"] - v["d2m"]
    v["srn"] = v["ssr"] + v["str"]
    with np.errstate(invalid="ignore", divide="ignore"):
        bw = v["sshf"] / v["slhf"]
    bw[~np.isfinite(bw)] = np.nan          # slhf ~ 0 -> undefined, drop from clim
    v["bowen"] = bw
    v["uv10"] = np.hypot(v["u10"], v["v10"])
    # thickness = gh(upper=lower pressure) - gh(lower=higher pressure)
    v["thk_1000_500"] = v["gh500"] - v["gh1000"]
    v["thk_1000_850"] = v["gh850"] - v["gh1000"]
    v["thk_850_700"]  = v["gh700"] - v["gh850"]
    v["thk_925_500"]  = v["gh500"] - v["gh925"]
    v["thk_925_700"]  = v["gh700"] - v["gh925"]
    # lapse: dt_<a>_<b> = t(a) - t(b)
    v["dt_850_500"] = v["t850"] - v["t500"]
    v["dt_925_850"] = v["t925"] - v["t850"]
    v["dt_700_500"] = v["t700"] - v["t500"]
    # IVT: terrain-aware trapezoidal integral of q*u, q*v over common levels
    ivt_u, ivt_v = compute_ivt(v)
    v["ivt_u"], v["ivt_v"] = ivt_u, ivt_v
    v["ivt"] = np.hypot(ivt_u, ivt_v)
    return v


def compute_ivt(v):
    """IVT on the common q/u/v levels {925,850,700,500}, sp-aware bottom layer."""
    G = 9.80665
    levs = [925, 850, 700, 500]                  # descending altitude order below
    # sort by pressure descending (surface -> aloft) => 925,850,700,500 already is
    p = np.array(levs, dtype="f8") * 100.0
    q = np.stack([v[f"q{L}"] for L in levs], axis=0)   # (nlev,Y,I,L,y,x)
    u = np.stack([v[f"u{L}"] for L in levs], axis=0)
    w = np.stack([v[f"v{L}"] for L in levs], axis=0)
    sp = v["sp"]                                       # (Y,I,L,y,x), Pa
    qu, qv = q * u, q * w
    nlay = len(levs) - 1
    p_lo, p_hi = p[:-1], p[1:]
    dp = p_lo - p_hi
    q_m = 0.5 * (qu[:-1] + qu[1:])
    v_m = 0.5 * (qv[:-1] + qv[1:])
    ivt_u = np.zeros_like(sp, dtype="f8")
    ivt_v = np.zeros_like(sp, dtype="f8")
    for k in range(nlay):
        dpe = np.where(sp >= p_lo[k], dp[k],
                       np.where(sp <= p_hi[k], 0.0, sp - p_hi[k]))
        ivt_u += q_m[k] * dpe
        ivt_v += v_m[k] * dpe
    return (ivt_u / G).astype("f4"), (ivt_v / G).astype("f4")


# --------------------------------------------------------------------------- #
#  1.5 deg daily Tn (mn2t6) from the txtm group
# --------------------------------------------------------------------------- #
def _load_mn_year(a):
    root, year, inits = a
    out = np.empty((len(inits), NLEAD, EC15_LAT.size, EC15_LON.size), dtype="f4")
    for k, mmdd in enumerate(inits):
        hits = glob.glob(os.path.join(root, str(year), f"*{year}-{mmdd}*.nc"))
        if not hits:
            raise FileNotFoundError(f"txtm {year}-{mmdd}")
        with xr.open_dataset(hits[0]) as d:
            arr = d["mn2t6"].values.astype("f4")[:NLEAD * NSTEP]
            daily = arr.reshape(NLEAD, NSTEP, arr.shape[1], arr.shape[2]).min(axis=1)
            out[k] = daily - 273.15
    return year, out


def build_mn15(inits, jobs):
    data = np.empty((len(C.YEARS), len(inits), NLEAD, EC15_LAT.size, EC15_LON.size),
                    dtype="f4")
    tasks = [(TXTM_ROOT, y, inits) for y in C.YEARS]
    with ProcessPoolExecutor(max_workers=jobs) as pool:
        for year, arr in pool.map(_load_mn_year, tasks):
            data[C.YEARS.index(year)] = arr
            print(f"  mn2t6 {year} done", flush=True)
    return data


# --------------------------------------------------------------------------- #
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tx-ds", default=TXDS_DEFAULT,
                    help="downscaled forecast Tx .npy (Y,I,L,npt), = mx2t6")
    ap.add_argument("--jobs", type=int, default=10)
    args = ap.parse_args()

    grid = C.Grid()
    npt = grid.npt
    plat = grid.lat[grid.iy].astype("f8")
    plon = grid.lon[grid.ix].astype("f8")
    W = bilinear_weights(plat, plon)
    print(f"[grid] {npt} land points, bilinear weights {W.shape}", flush=True)

    # ---- canonical init ordering (identical to ../HW build_cube) ---------- #
    import datetime as dt
    canon = {y: C.init_dates(y) for y in C.YEARS}
    per_year = [len(canon[y]) for y in C.YEARS]
    offsets = np.concatenate([[0], np.cumsum(per_year)[:-1]]).astype(int)
    ninit = int(sum(per_year))
    print(f"[init] {ninit} = {per_year[0]}/yr x {len(C.YEARS)} yr", flush=True)

    # ---- 1.5 deg predictors + their init coord --------------------------- #
    ds = xr.open_dataset(EC15)
    inits15 = [str(x) for x in ds["init"].values]     # mmdd order of the arrays
    # map: for each canonical (year, i) -> index into the 1.5 deg 'init' axis
    mmdd_to_j = {m: j for j, m in enumerate(inits15)}
    v15 = {name: ds[name].values.astype("f4") for name in ds.data_vars}  # (Y,I,L,13,10)
    ds.close()
    # sanity: canonical mmdd must all exist in the 1.5 deg set, same count
    for y in C.YEARS:
        assert len(canon[y]) == len(inits15), (y, len(canon[y]), len(inits15))
    print(f"[ec15] {len(v15)} raw vars, init axis {len(inits15)}", flush=True)

    # ---- derived vars, then the 1.5 deg 40-var stack (minus mx2t6) -------- #
    build_derived(v15)
    mn15 = build_mn15(inits15, args.jobs)             # (Y, I15, L, 13, 10)
    v15["mn2t6"] = mn15

    # downscaled Tx already on the 541 points
    txds = np.load(args.tx_ds)                        # (Y, I, L, npt)
    assert txds.shape[-1] == npt, txds.shape
    print(f"[tx-ds] {os.path.basename(args.tx_ds)} {txds.shape}", flush=True)

    # canonical (year,i) -> 1.5 deg init-axis index, precomputed once
    j15_of = np.empty(ninit, dtype=int)
    for yi, y in enumerate(C.YEARS):
        for i, init in enumerate(canon[y]):
            j15_of[offsets[yi] + i] = mmdd_to_j[init.isoformat()[5:]]

    # ---- assemble EC cube IN RAM (fill contiguously per var), write once. --
    # Strided writes straight into a lustre memmap are pathologically slow; a
    # single sequential .tofile of the finished array is orders faster.
    se, so = B.cube_shapes(ninit, npt)
    os.makedirs(B.CUBE_DIR, exist_ok=True)
    print(f"[cube] assembling ec {se} = {np.prod(se)*4/1e9:.2f} GB in RAM", flush=True)
    ecram = np.empty(se, dtype="f4")                  # (ninit,46,40,npt)

    # source flat row into (Y*I15) for each canonical cube row
    yi_of = np.repeat(np.arange(len(C.YEARS)), per_year)
    src_row = yi_of * len(inits15) + j15_of           # (ninit,)

    jmx = B.EC_VARS.index("mx2t6")
    for j, name in enumerate(B.EC_VARS):
        if name == "mx2t6":
            continue
        pts = interp_field(v15.pop(name), W).astype("f4")   # (Y,I15,L,npt); free src
        if name in B.LOG_VARS:
            pts = np.log1p(np.maximum(pts, 0.0))
        pts = pts.reshape(-1, C.NLT, npt)             # (Y*I15, L, npt)
        ecram[:, :, j, :] = pts[src_row]
    # mx2t6 = downscaled Tx, canonical (Y,I) order == txds (Y,I) order
    ecram[:, :, jmx, :] = txds.reshape(ninit, C.NLT, npt)
    print("[cube] writing ec_raw.f4 ...", flush=True)
    ecram.tofile(B.EC_CUBE); del ecram
    print("[cube] ec_raw written", flush=True)

    # ---- obs_raw IN RAM then one write (identical obs source as ../HW) ----- #
    obsram = np.empty(so, dtype="f4")
    for yi, y in enumerate(C.YEARS):
        for i, init in enumerate(canon[y]):
            with Dataset(C.obs_file(init)) as d:
                for j, vv in enumerate(B.OBS_VARS):
                    a = np.ma.filled(d[vv][:], np.nan).astype("f4")
                    a[~np.isfinite(a)] = np.nan
                    if vv in B.LOG_VARS:
                        a = np.log1p(np.maximum(a, 0.0))
                    obsram[offsets[yi] + i, :, j, :] = grid.to_points(a)
        print(f"  obs {y} done", flush=True)
    obsram.tofile(B.OBS_CUBE); del obsram

    # ---- meta.npz (identical construction to build_cube) ----------------- #
    init_dates, init_year, doy = [], [], []
    for y in C.YEARS:
        for init in canon[y]:
            init_dates.append(init.isoformat())
            init_year.append(y)
            doy.append([C.doy_noleap(d) for d in C.valid_days(init)])
    np.savez(B.META,
             ec_vars=np.array(B.EC_VARS), obs_vars=np.array(B.OBS_VARS),
             init_dates=np.array(init_dates), init_year=np.array(init_year, "i4"),
             doy=np.array(doy, "i4"), lead=np.arange(C.NLT, dtype="i4"),
             years=np.array(C.YEARS, "i4"),
             log_vars=np.array(sorted(B.LOG_VARS)),
             npt=np.array([npt]), ninit=np.array([ninit]))
    print(f"[done] -> {B.CUBE_DIR}", flush=True)


if __name__ == "__main__":
    main()