"""Assemble the 104 reforecast GRIBs (per group) into per-variable logical cubes.

Input  : data/raw/ecmwf_s2s_1p5deg/s2s_cf_<group>_<dom>_mv2023-MM-DD_hy2003-2022.grib
Output : data/processed/cube/<group>/<group>_mv2023-MM-DD.nc   (one per model-version date)

Each intermediate NetCDF holds every variable of the group with dims
    (hindcast_year[20], lead[46], [level], lat, lon)
standardized to DAILY lead days LT0..LT45 on the native coarse grid:
  * mx2t6 -> daily-max Tx on the LOCAL (UTC+7) calendar day (daily_max_from_6h);
  * daily/instantaneous fields -> the 24-h value whose valid 00Z maps to each
    local day (lead day k uses the step at k*24 h).
Latitude is flipped to ascending; temperature fields (t, 2t, mx2t6) -> degC.

The full logical cube per variable is the 104 files stacked on a new
`model_version_date` axis -> (model_version_date[104], hindcast_year[20],
lead[46], [level], lat, lon). Stack lazily with:
    xr.open_mfdataset("data/processed/cube/<group>/<group>_mv2023-*.nc",
                      concat_dim="model_version_date", combine="nested")

Resumable: skips model-version dates whose output already exists. Safe to
stop/re-run.  Usage:
    python3 scripts/08_assemble_cube.py --group sfc_tx --dates 1   # smoke test
    python3 scripts/08_assemble_cube.py --group all --dates all    # full run
"""
import argparse
import os
import subprocess
import sys
import tempfile
import time

import numpy as np
import xarray as xr

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from data.grib_preprocess import S2SReforecastPreprocessor, daily_max_from_6h  # noqa: E402
from utils.timealign import local_calendar_day  # noqa: E402

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
RAW = os.path.join(ROOT, "data", "raw", "ecmwf_s2s_1p5deg")
OUT = os.path.join(ROOT, "data", "processed", "cube")
HDATES = os.path.join(ROOT, "hdate_2023.txt")
MV_YEAR = "2023"
TZ = 7
N_LEAD = 46          # LT0..LT45 (local calendar days)
TEMP_VARS = {"t", "2t", "mx2t6"}   # K -> degC

# group -> (filename domain tag, kind)
#   tx6h  : mx2t6 6-hourly -> daily-max Tx on local day
#   daily : instantaneous daily fields (24-hourly steps)
#   avg   : 24-h average fields (step ranges L_L+24), e.g. daily-mean 2t
GROUPS = {
    "sfc_tx":  ("loc", "tx6h"),
    "sfc_dyn": ("loc", "daily"),
    "sfc_t2m": ("loc", "avg"),
    "pl_hi":   ("ls",  "daily"),
}


def log(*a):
    print(*a, flush=True)


def _mktemp(suffix=""):
    """mkstemp but close the returned fd immediately (mkstemp leaves it open)."""
    fd, path = tempfile.mkstemp(suffix=suffix)
    os.close(fd)
    return path


def read_hdates():
    return [l.strip() for l in open(HDATES) if l.strip()]


def grib_path(group, mmdd):
    dom = GROUPS[group][0]
    return os.path.join(RAW, f"s2s_cf_{group}_{dom}_mv{MV_YEAR}-{mmdd}_hy2003-2022.grib")


def _cdo_nc(grib, nc):
    r = subprocess.run(["cdo", "-s", "-f", "nc", "copy", grib, nc],
                       capture_output=True, text=True, timeout=600)
    if r.returncode != 0 or not os.path.exists(nc):
        raise RuntimeError(f"cdo failed: {r.stderr[:400]}")


def _split_by_init(grib):
    """Split a 20-hindcast-year GRIB into one temp GRIB per dataDate (init)."""
    out = subprocess.run(["grib_get", "-p", "dataDate", grib],
                         capture_output=True, text=True, timeout=300).stdout.split()
    dds = sorted({d for d in out if d.isdigit() and len(d) == 8})
    parts = {}
    for dd in dds:
        tmp = _mktemp(f"_{dd}.grib")
        subprocess.run(["grib_copy", "-w", f"dataDate={dd}", grib, tmp],
                       capture_output=True, text=True, timeout=300)
        if os.path.exists(tmp) and os.path.getsize(tmp) > 0:
            parts[dd] = tmp
    return parts


def _lead_index(valid_times, init_date):
    """Local-day lead index (int) for each valid time: local_day - init_date."""
    lday = local_calendar_day(valid_times, TZ)
    return (lday - init_date.astype("datetime64[D]")).astype("timedelta64[D]").astype(int)


def assemble_sfc_tx(grib):
    """mx2t6 (6-hourly) -> daily-max Tx per hindcast year, lead LT0..LT45."""
    pp = S2SReforecastPreprocessor(tz_hours=TZ)
    fields = pp.load_canonical_by_init(grib, to_degC=True)   # 20 fields, degC
    f0 = fields[0]
    lat, lon = f0.lat, f0.lon
    years, arrs = [], []
    for f in fields:
        days, txday, cov = daily_max_from_6h(f, tz_hours=TZ)  # (nday,lat,lon)
        k = (days - f.init_time.astype("datetime64[D]")).astype("timedelta64[D]").astype(int)
        buf = np.full((N_LEAD, len(lat), len(lon)), np.nan, "float32")
        for j, kk in enumerate(k):
            if 0 <= kk < N_LEAD:
                buf[kk] = txday[j].astype("float32")
        years.append(int(str(f.init_time)[:4]))
        arrs.append(buf)
    data = np.stack(arrs, 0)                                 # (year,lead,lat,lon)
    ds = xr.Dataset(
        {"mx2t6": (("hindcast_year", "lead", "lat", "lon"), data)},
        coords={"hindcast_year": np.array(years, "int16"),
                "lead": np.arange(N_LEAD, dtype="int16"), "lat": lat, "lon": lon},
    )
    ds["mx2t6"].attrs.update(units="degC", long_name="daily-max Tx (local UTC+7 day)")
    return ds, str(f0.model_version_date)


def _shortnames(grib):
    out = subprocess.run(["grib_get", "-p", "shortName", grib],
                         capture_output=True, text=True, timeout=120).stdout.split()
    return sorted(set(out))


def _load_init_daily(init_grib):
    """Convert one init GRIB to an xr.Dataset of DAILY fields. Fast path: a single
    cdo copy of the whole init (works when every variable shares the clean daily
    step set, e.g. pl_hi). Fallback: convert per shortName and drop any variable
    whose series has < N_LEAD steps (e.g. 2t, archived only as a single 0-24h
    average). Returns (dataset|None, {dropped_var: reason})."""
    nc = _mktemp(".nc")
    try:
        try:
            _cdo_nc(init_grib, nc)
            return xr.open_dataset(nc).load(), {}
        except RuntimeError:
            pass                                   # ragged steps -> per-var path
    finally:
        if os.path.exists(nc):
            os.remove(nc)

    das, coords, dropped = {}, {}, {}
    for nm in _shortnames(init_grib):
        sub = _mktemp(f"_{nm}.grib")
        nc = _mktemp(".nc")
        try:
            subprocess.run(["grib_copy", "-w", f"shortName={nm}", init_grib, sub],
                           capture_output=True, text=True, timeout=120)
            _cdo_nc(sub, nc)
            # detach fully into numpy inside the context manager so no backend
            # file handle survives (avoids an fd leak across 20x104 inits)
            with xr.open_dataset(nc) as d:
                if d.sizes.get("time", 0) < N_LEAD:
                    dropped[nm] = f"{d.sizes.get('time', 0)} step(s)"
                    continue
                da = d[[x for x in d.data_vars][0]]
                das[nm] = (da.dims, da.values.copy())
                for c in da.dims:
                    if c not in coords and c in d.coords:
                        coords[c] = d[c].values.copy()
        finally:
            for p in (sub, nc):
                if os.path.exists(p):
                    os.remove(p)
    if not das:
        return None, dropped
    ds = xr.Dataset({nm: dv for nm, dv in das.items()}, coords=coords)
    return ds, dropped


def assemble_daily(grib):
    """Instantaneous daily fields (sfc_dyn / pl_hi) -> lead LT0..LT45."""
    parts = _split_by_init(grib)
    per_year = {}        # var -> list of (year, array)
    lat = lon = plev = None
    dropped_any = {}
    try:
        for dd in sorted(parts):
            ds, dropped = _load_init_daily(parts[dd])
            dropped_any.update(dropped)
            if ds is None:
                continue
            if ds["lat"].values[0] > ds["lat"].values[-1]:   # -> ascending
                ds = ds.isel(lat=slice(None, None, -1))
            init = np.datetime64(f"{dd[:4]}-{dd[4:6]}-{dd[6:8]}")
            k = _lead_index(ds["time"].values, init)          # (time,)
            keep = (k >= 0) & (k < N_LEAD)
            lat = ds["lat"].values
            lon = ds["lon"].values
            has_lev = "plev" in ds.dims
            if has_lev:
                plev = (ds["plev"].values / 100.0).astype("int16")  # Pa->hPa
            year = int(dd[:4])
            for v in ds.data_vars:
                da = ds[v]
                shp = (N_LEAD, len(plev), len(lat), len(lon)) if has_lev \
                    else (N_LEAD, len(lat), len(lon))
                buf = np.full(shp, np.nan, "float32")
                vals = da.values.astype("float32")
                if v in TEMP_VARS:
                    vals = vals - 273.15
                for ti, kk in zip(np.nonzero(keep)[0], k[keep]):
                    buf[kk] = vals[ti]
                per_year.setdefault(v, []).append((year, buf))
            ds.close()
    finally:
        for p in parts.values():
            if os.path.exists(p):
                os.remove(p)

    if not per_year:
        raise RuntimeError(f"no usable daily variables (dropped={dropped_any})")
    if dropped_any:
        log(f"    dropped vars: {dropped_any}")
    sample = next(iter(per_year.values()))
    years = [y for y, _ in sorted(sample)]
    has_lev = per_year[next(iter(per_year))][0][1].ndim == 4
    data_vars, mv = {}, None
    for v, lst in per_year.items():
        stack = np.stack([a for _, a in sorted(lst)], 0)       # (year,lead,[lev,]lat,lon)
        dims = ("hindcast_year", "lead", "level", "lat", "lon") if has_lev \
            else ("hindcast_year", "lead", "lat", "lon")
        data_vars[v] = (dims, stack)
    coords = {"hindcast_year": np.array(years, "int16"),
              "lead": np.arange(N_LEAD, dtype="int16"), "lat": lat, "lon": lon}
    if has_lev:
        coords["level"] = plev
    ds = xr.Dataset(data_vars, coords=coords)
    for v in ds.data_vars:
        if v in TEMP_VARS:
            ds[v].attrs["units"] = "degC"
    if dropped_any:
        ds.attrs["dropped_vars"] = "; ".join(f"{k}:{v}" for k, v in dropped_any.items())
    mv = f"{MV_YEAR}-" + os.path.basename(grib).split("mv2023-")[1][:5]
    return ds, mv


def assemble_t2m(grib):
    """Daily-MEAN 2t (24-h average windows 0_24..1080_1104) -> lead LT0..LT45.
    Each init has 46 avg records; cdo labels them by end-time, so record order
    (sorted ascending) == lead day 0..45 (the k-th forecast day's mean)."""
    parts = _split_by_init(grib)
    lat = lon = None
    by_year = {}
    try:
        for dd in sorted(parts):
            ds, _ = _load_init_daily(parts[dd])
            ds = ds.sortby("time")
            if ds["lat"].values[0] > ds["lat"].values[-1]:
                ds = ds.isel(lat=slice(None, None, -1))
            lat, lon = ds["lat"].values, ds["lon"].values
            v = [x for x in ds.data_vars][0]
            vals = ds[v].squeeze().values.astype("float32") - 273.15   # K->degC
            buf = np.full((N_LEAD, len(lat), len(lon)), np.nan, "float32")
            n = min(N_LEAD, vals.shape[0])
            buf[:n] = vals[:n]
            by_year[int(dd[:4])] = buf
            ds.close()
    finally:
        for p in parts.values():
            if os.path.exists(p):
                os.remove(p)
    years = sorted(by_year)
    data = np.stack([by_year[y] for y in years], 0)
    ds = xr.Dataset(
        {"2t": (("hindcast_year", "lead", "lat", "lon"), data)},
        coords={"hindcast_year": np.array(years, "int16"),
                "lead": np.arange(N_LEAD, dtype="int16"), "lat": lat, "lon": lon},
    )
    ds["2t"].attrs.update(units="degC", long_name="daily-mean 2 m temperature (0-24h fcst-day avg)")
    mv = f"{MV_YEAR}-" + os.path.basename(grib).split("mv2023-")[1][:5]
    return ds, mv


def process(group, mmdd):
    grib = grib_path(group, mmdd)
    if not (os.path.exists(grib) and os.path.getsize(grib) > 0):
        raise FileNotFoundError(grib)
    kind = GROUPS[group][1]
    if kind == "tx6h":
        ds, mv = assemble_sfc_tx(grib)
    elif kind == "avg":
        ds, mv = assemble_t2m(grib)
    else:
        ds, mv = assemble_daily(grib)
    ds.attrs.update(model_version_date=mv, group=group,
                    source=os.path.basename(grib), tz="UTC+7 local day",
                    lead="LT0..LT45 (local calendar days)")
    return ds


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--group", default="all",
                    choices=list(GROUPS) + ["all"])
    ap.add_argument("--dates", default="1", help="N (first N hdates) or 'all'")
    ap.add_argument("--force", action="store_true", help="rebuild existing outputs")
    a = ap.parse_args()
    groups = list(GROUPS) if a.group == "all" else [a.group]
    dates = read_hdates()
    todo = dates if a.dates == "all" else dates[:int(a.dates)]

    for group in groups:
        gout = os.path.join(OUT, group)
        os.makedirs(gout, exist_ok=True)
        log(f"=== group={group}: {len(todo)} date(s) ===")
        n_ok = n_skip = n_fail = 0
        for i, mmdd in enumerate(todo, 1):
            dst = os.path.join(gout, f"{group}_mv{MV_YEAR}-{mmdd}.nc")
            if os.path.exists(dst) and not a.force:
                n_skip += 1
                continue
            t0 = time.time()
            try:
                ds = process(group, mmdd)
                enc = {v: {"zlib": True, "complevel": 4} for v in ds.data_vars}
                tmp = dst + ".tmp"
                ds.to_netcdf(tmp, encoding=enc)
                ds.close()
                os.replace(tmp, dst)
                sz = os.path.getsize(dst) / 1e6
                log(f"[{i}/{len(todo)}] {group} {mmdd} OK  {sz:.1f} MB  {time.time()-t0:.0f}s")
                n_ok += 1
            except Exception as e:
                log(f"[{i}/{len(todo)}] {group} {mmdd} ERROR {type(e).__name__}: {e}")
                n_fail += 1
        log(f"--- {group}: ok={n_ok} skip={n_skip} fail={n_fail} ---")


if __name__ == "__main__":
    main()