#!/usr/bin/env python3
"""Build the 46-variable 1.5 deg predictor cube from the REAL ECDS download.

Drop-in replacement for the proxy `ec15_predictors.nc` (build_ec_proxy_all.py):
same dims, same coords, same 46 variable names, so every downstream script
(ehd_full.py etc.) runs unchanged after swapping the input path.

The real GRIB->NetCDF files are already on the native 1.5 deg MARS grid
(lat 24..6, lon 100..113.5, 13x10), so NO spatial averaging is needed -- we read
the cells directly.  Only the time convention has to match the proxy:

  srf_daily : 46 daily means, one per lead day               -> value[:46]
  srf_24h   : 47 points at 0,24,..,1104 h.
              accumulated fields (fluxes/radiation/precip) are cumulative from
              the start -> lead day d = value[d+1]-value[d]  (diff of first 47);
              instantaneous fields (msl,sp,u10,v10)          -> value[:46]
  plev      : 47 instantaneous points                        -> value[:46]

Variable/level set is chosen to reproduce EXACTLY the 46 proxy names:
  gh {1000,925,850,700,500}  q {1000,925,850,700,500}
  t/u/v {925,850,700,500}    w -> level 500, stored as bare name 'w'
Levels 300/200 (present in the real plev files) are intentionally dropped so the
name set matches the proxy; add them later if a model wants upper levels.
"""
import argparse
import glob
import os
import sys
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import xarray as xr

YEARS = list(range(2003, 2023))
NLEAD = 46
NC = "/home/tanpv/Progs/claude/HW_EC_10/data/nc"
OUT = "/home/tanpv/Progs/claude/HW_EC_10/outputs/real/ec15_predictors.nc"

EC15_LON = 100.0 + 1.5 * np.arange(10)
EC15_LAT = 24.0 - 1.5 * np.arange(13)

SURF_DAILY = ["t2m", "d2m", "skt", "cape", "tcw", "tcc",
              "sm20", "sm100", "st20", "st100"]
SURF_24H_INST = ["msl", "sp", "u10", "v10"]
SURF_24H_ACC = ["sshf", "slhf", "ssrd", "strd", "ssr", "str", "ttr", "cp", "tp"]
NONNEG = {"cp", "tp", "ssrd", "ssr"}
PLEV = {"gh": [1000, 925, 850, 700, 500],
        "q":  [1000, 925, 850, 700, 500],
        "t":  [925, 850, 700, 500],
        "u":  [925, 850, 700, 500],
        "v":  [925, 850, 700, 500],
        "w":  [500]}                      # single level, stored as 'w'


def _open(group, year, mmdd):
    f = os.path.join(NC, group, str(year), f"{group}_cf_{year}-{mmdd}.nc")
    if not os.path.exists(f):
        hits = glob.glob(os.path.join(NC, group, str(year), f"*{year}-{mmdd}*.nc"))
        if not hits:
            raise FileNotFoundError(f)
        f = hits[0]
    return xr.open_dataset(f)


def _align(da):
    """Return (time,13,10) float32 on the exact EC15 grid order."""
    da = da.reindex(latitude=EC15_LAT, longitude=EC15_LON, method="nearest",
                    tolerance=0.01)
    return da.values.astype("float32")


def load_year(year, inits):
    res = {}
    for k, mmdd in enumerate(inits):
        vals = {}
        with _open("srf_daily", year, mmdd) as d:
            for v in SURF_DAILY:
                vals[v] = _align(d[v])[:NLEAD]
        with _open("srf_24h", year, mmdd) as d:
            for v in SURF_24H_INST:
                vals[v] = _align(d[v])[:NLEAD]
            for v in SURF_24H_ACC:
                x = _align(d[v])[:NLEAD + 1]
                x = np.diff(x, axis=0)                 # de-accumulate to daily
                if v in NONNEG:
                    np.clip(x, 0, None, out=x)
                vals[v] = x
        with _open("plev", year, mmdd) as d:
            for v, levs in PLEV.items():
                for L in levs:
                    name = v if v == "w" else f"{v}{L}"
                    vals[name] = _align(d[v].sel(level=L))[:NLEAD]
        for name, arr in vals.items():
            if name not in res:
                res[name] = np.empty((len(inits), NLEAD, 13, 10), dtype="float32")
            res[name][k] = arr
    return year, res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--jobs", type=int, default=10)
    ap.add_argument("--out", default=OUT)
    ap.add_argument("--hdates", default="/home/tanpv/Data/QG.25.196/hdate_2023.txt")
    args = ap.parse_args()

    with open(args.hdates) as fh:
        inits = sorted({t for t in fh.read().split() if len(t) == 5 and t[2] == "-"})
    print(f"{len(inits)} init dates, {len(YEARS)} years", flush=True)

    data = {}
    tasks = [(y, inits) for y in YEARS]
    with ProcessPoolExecutor(max_workers=args.jobs) as pool:
        for year, res in pool.map(load_year, *zip(*tasks)):
            for name, arr in res.items():
                if name not in data:
                    data[name] = np.empty((len(YEARS),) + arr.shape, dtype="float32")
                data[name][YEARS.index(year)] = arr
            print(f"  {year} done", flush=True)

    ds = xr.Dataset(
        {k: (("year", "init", "lead", "lat", "lon"), v) for k, v in data.items()},
        coords={"year": YEARS, "init": inits, "lead": np.arange(NLEAD),
                "lat": EC15_LAT, "lon": EC15_LON})
    ds.attrs["note"] = ("1.5 deg predictors from the real ECDS S2S download "
                        "(model 2023, control fc), native MARS grid, no regridding.")
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    ds.to_netcdf(args.out, encoding={k: {"zlib": True, "complevel": 4} for k in data})
    print(f"wrote {args.out} ({os.path.getsize(args.out)/1e9:.2f} GB), "
          f"{len(data)} variables: {sorted(data)}")
    return 0


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