#!/usr/bin/env python3
"""Extract the 16 CFS flux surface fields over the SE-Asia window, per init.

Reading is delegated to grib_read.read_batch, which places records by (valid
time, variable, level) rather than by position and re-reads any valid time the
batched pass returned short. That matters because the archive holds truncated
files and files missing a record, and wgrib2 abandons the rest of a concatenated
stream when it meets one.

Output: outputs/Flux6h/<YYYY>/flx6h_<YYYYMMDD>.npy
        float32, shape (184, 16, nlat_w, nlon_w) = (step, variable, lat, lon)
        step s is valid at init + 6*(s+1) hours; latitude ascending.
"""

from __future__ import annotations

import argparse
import os
import datetime as dt
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")


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


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

    files = C.flux_files(init)
    have = set(os.listdir(C.flux_dir(init)))
    missing = [f for f in files if os.path.basename(f) not in have]
    if missing:
        raise RuntimeError("%s: %d of %d flux 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), C.FLUX_MATCH,
                                        C.FLUX_ORDER, BOX, ny, nx)
    want = C.NSTEP * C.NFVAR
    if nfilled < want:
        print("  %s: %d/%d records missing after retry (%d valid times re-read)"
              % (init, want - nfilled, want, nretry), flush=True)

    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, help="debug: first N inits only")
    a = ap.parse_args()

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

    fail_log = os.path.join(C.LOGS, "exflx_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:                     # one bad init must not
            nfail += 1                               # take the chunk down
            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()