"""Extract 40 CFS pgb pressure-level fields over the SE-Asia window, per init.

Reading is delegated to grib_read.read_batch (see there for why records are
placed by identity and short valid times re-read). Two valid times per lead day are kept
(12Z of day D and 00Z of day D+1) rather than all four: free-atmosphere fields
have a weak diurnal cycle compared with their synoptic variability, and the cost
here is dominated by reading 22 MB GRIB files, half of them across the NAS mount.

Output: outputs/Pgb2x/<YYYY>/pgb2x_<YYYYMMDD>.npy
        float32, shape (46, 2, 40, nlat_w, nlon_w) = (lead, time-of-day, var, lat, lon)
        time index 0 = 12Z of the lead day, 1 = 00Z of the next day.
"""

from __future__ import annotations

import argparse
import datetime as dt
import os
import sys
import time

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C
import grib_read as G

BOX = ("95:120", "0:30")

PGB_MATCH = (
    ":(HGT|TMP|UGRD|VGRD):(1000|925|850|700|500|200) mb:"
    "|:SPFH:(1000|925|850|700|500) mb:"
    "|:RH:(925|850|700|500) mb:"
    "|:VVEL:(925|850|700|500) mb:"
    "|:PRMSL:mean sea level:"
    "|:PWAT:entire atmosphere"
    "|:(CAPE|CIN):surface:"
)

# Record order as written in the GRIB files (verified 2026-09-02).
PGB_ORDER = (
    [("HGT", "200 mb"), ("TMP", "200 mb"), ("UGRD", "200 mb"), ("VGRD", "200 mb")]
    + [(v, "%d mb" % p) for p in (500, 700, 850, 925)
       for v in ("HGT", "TMP", "RH", "SPFH", "VVEL", "UGRD", "VGRD")]
    + [("HGT", "1000 mb"), ("TMP", "1000 mb"), ("SPFH", "1000 mb"),
       ("UGRD", "1000 mb"), ("VGRD", "1000 mb")]
    + [("CAPE", "surface"), ("CIN", "surface"), ("PRMSL", "mean sea level")]
)
PGB_NAMES = ["%s%s" % (v.lower(), l.split()[0] if l.endswith("mb") else "")
             for v, l in PGB_ORDER]
NPVAR = len(PGB_ORDER)
NTOD = 2                       # 12Z of the lead day, 00Z of the next day
PGB2X = os.path.join(C.OUT, "Pgb2x")
PGB_GRID = os.path.join(C.META, "pgb_grid.npz")


def pgb_grid():
    z = np.load(PGB_GRID)
    return dict(nlon=int(z["nlon"]), nlat=int(z["nlat"]), i0=int(z["i0"]),
                i1=int(z["i1"]), j0=int(z["j0"]), j1=int(z["j1"]),
                wlon=z["wlon"], wlat=z["wlat"])


def pgb_path(init):
    return os.path.join(PGB2X, init[:4], "pgb2x_%s.npy" % init)


def pgb_files(init):
    """(46 x 2) valid times: 12Z of each lead day and 00Z of the following day."""
    d0 = dt.datetime(int(init[:4]), int(init[4:6]), int(init[6:8]))
    base = C.pgb_dir(init)
    out = []
    for k in range(C.NLT):
        for h in (12, 24):
            v = d0 + dt.timedelta(days=k, hours=h)
            out.append(os.path.join(base, "pgbf%s.grb2" % v.strftime("%Y%m%d%H")))
    return out


def valid_times(init):
    d0 = dt.datetime(int(init[:4]), int(init[4:6]), int(init[6:8]))
    return [(d0 + dt.timedelta(days=k, hours=h)).strftime("%Y%m%d%H")
            for k in range(C.NLT) for h in (12, 24)]


def extract_one(init, grid, overwrite=False):
    out = pgb_path(init)
    if os.path.exists(out) and not overwrite:
        return "skip", 0
    os.makedirs(os.path.dirname(out), exist_ok=True)

    files = pgb_files(init)
    have = set(os.listdir(C.pgb_dir(init)))
    missing = [f for f in files if os.path.basename(f) not in have]
    if missing:
        raise RuntimeError("%s: %d of %d pgb files absent from disk"
                           % (init, len(missing), len(files)))

    ny, nx = len(grid["wlat"]), len(grid["wlon"])
    arr, nfilled, nretry = G.read_batch(files, valid_times(init), PGB_MATCH,
                                        PGB_ORDER, BOX, ny, nx)
    want = len(files) * NPVAR
    if nfilled < want:
        print("  %s: %d/%d records missing after retry (%d valid times re-read)"
              % (init, want - nfilled, want, nretry), flush=True)

    arr = arr.reshape(C.NLT, NTOD, NPVAR, ny, nx)
    tmp = out + ".tmp%d.npy" % os.getpid()
    np.save(tmp, arr)
    os.replace(tmp, out)
    return "done", want - nfilled


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--chunk", type=int, default=0)
    ap.add_argument("--nchunk", type=int, default=1)
    ap.add_argument("--overwrite", action="store_true")
    ap.add_argument("--limit", type=int, default=0)
    a = ap.parse_args()

    inits = C.both_inits()
    if a.limit:
        inits = inits[:a.limit]
    mine = inits[a.chunk::a.nchunk]
    grid = pgb_grid()

    fail_log = os.path.join(C.LOGS, "expgb_failed_%d.txt" % a.chunk)
    t0 = time.time()
    ndone = nskip = nfail = ngap = 0
    for n, init in enumerate(mine, 1):
        try:
            st, gaps = extract_one(init, grid, a.overwrite)
        except Exception as exc:
            nfail += 1
            with open(fail_log, "a") as fh:
                fh.write("%s\t%s\n" % (init, exc))
            print("  FAILED %s: %s" % (init, exc), flush=True)
            continue
        ngap += gaps
        ndone += st == "done"
        nskip += st == "skip"
        if n % 20 == 0 or n == len(mine):
            el = time.time() - t0
            print("[chunk %d] %d/%d done=%d skip=%d fail=%d gaps=%d %.1f s %.2f s/init"
                  % (a.chunk, n, len(mine), ndone, nskip, nfail, ngap, el, el / n),
                  flush=True)
    print("[chunk %d] FINISHED %d inits in %.1f s"
          % (a.chunk, len(mine), time.time() - t0), flush=True)


if __name__ == "__main__":
    main()