#!/usr/bin/env python3
"""Robust window extraction from a batch of GRIB2 files.

The CFS archive contains files that are TRUNCATED (a 22 MB pgb file cut to 11 MB,
holding 18 of its 40 records) and files that are complete but simply LACK a record
(a full-size pgb file with no PRMSL). Both were found on 2026-09-02.

Two consequences drive this design:

  * A fixed record count cannot be assumed. Records are placed by IDENTITY -
    (valid time, variable, level) read from a custom wgrib2 inventory - and
    anything absent is left as NaN. Counting records and trusting their order is
    what silently scrambles a variable axis.
  * wgrib2 cannot resynchronise after a truncated member of a concatenated
    stream: it stops, and every later file in the batch is lost without an error.
    So the batch result is checked for completeness and any valid time that came
    back short is re-read from its own file, where truncation costs only the
    records that are actually missing.

Extraction is two passes: `-small_grib` cuts the window into a ~170 KB temporary
GRIB, then the inventory and the raw floats are read from that. Cutting first
keeps the temporaries tiny and costs no more than decoding the full grid.
"""

from __future__ import annotations

import os
import subprocess
import tempfile

import numpy as np

WGRIB2 = "/work/apps/gnu_4.8.5/grads/2.2.0/bin/wgrib2"
ENV1 = dict(os.environ, OMP_NUM_THREADS="1")


def _pair(files, match, box, tmpdir):
    """Cut the window from `files` and return (inventory lines, raw float32)."""
    sg = os.path.join(tmpdir, "s.grb2")
    inv = os.path.join(tmpdir, "s.inv")
    binf = os.path.join(tmpdir, "s.bin")
    for p in (sg, inv, binf):
        if os.path.exists(p):
            os.remove(p)

    cat = subprocess.Popen(["cat"] + list(files), stdout=subprocess.PIPE)
    subprocess.run([WGRIB2, "-", "-match", match, "-order", "we:sn",
                    "-inv", "/dev/null", "-small_grib", box[0], box[1], sg],
                   stdin=cat.stdout, stdout=subprocess.DEVNULL,
                   stderr=subprocess.DEVNULL, env=ENV1)
    cat.stdout.close()
    cat.wait()
    if not os.path.exists(sg) or os.path.getsize(sg) == 0:
        return [], np.empty(0, dtype="<f4")

    subprocess.run([WGRIB2, sg, "-inv", inv, "-vt", "-var", "-lev",
                    "-no_header", "-bin", binf],
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, env=ENV1)
    with open(inv) as fh:
        lines = [ln for ln in fh.read().splitlines() if ln.strip()]
    raw = np.fromfile(binf, dtype="<f4")
    return lines, raw


def _place(lines, raw, npts, vt_index, var_index, arr, got):
    """Write each record into arr[step, var] by identity; mark it in `got`.

    `got` is the record ledger, and it is what decides whether a valid time needs
    re-reading. Do NOT infer that from NaN: NaN is also a legitimate undefined
    value (soil moisture over the sea), and treating it as a missing record makes
    every step look short and doubles the run time.
    """
    for r, ln in enumerate(lines):
        if (r + 1) * npts > raw.size:
            break                                    # binary ended early
        f = ln.split(":")
        vt, var, lev = f[2][3:], f[3], ":".join(f[4:])
        s = vt_index.get(vt)
        v = var_index.get((var, lev))
        if s is None or v is None:
            continue
        arr[s, v] = raw[r * npts:(r + 1) * npts].reshape(arr.shape[2], arr.shape[3])
        got[s, v] = True


def read_batch(files, vts, match, order, box, ny, nx, undef=9.9e20):
    """Extract `order` records for every valid time in `vts` from `files`.

    files[i] is the file holding valid time vts[i]. Returns (arr, nfilled,
    nsolo) with arr shaped (len(vts), len(order), ny, nx), NaN wherever a record
    could not be read, and nsolo the number of files that had to be read alone.
    """
    nstep, nvar, npts = len(vts), len(order), ny * nx
    arr = np.full((nstep, nvar, ny, nx), np.nan, dtype=np.float32)
    vt_index = {v: i for i, v in enumerate(vts)}
    var_index = {k: i for i, k in enumerate(order)}

    got = np.zeros((nstep, nvar), dtype=bool)
    nsolo = 0
    with tempfile.TemporaryDirectory(prefix="gribrd_") as td:
        # Batch everything, then re-batch whatever is still outstanding. A pass
        # recovers every file up to the first unreadable one; that one blocker is
        # then read on its own and dropped, and the next pass picks up from there.
        # Clean inits cost one wgrib2 call; each bad file costs two extra.
        remaining = list(range(nstep))
        for _ in range(nstep + 1):
            if not remaining:
                break
            lines, raw = _pair([files[i] for i in remaining], match, box, td)
            _place(lines, raw, npts, vt_index, var_index, arr, got)
            still = [i for i in remaining if not got[i].all()]
            if still and still[0] == remaining[0]:      # no progress on the head
                i = remaining[0]
                lines, raw = _pair([files[i]], match, box, td)
                _place(lines, raw, npts, vt_index, var_index, arr, got)
                nsolo += 1
                still = [k for k in remaining[1:] if not got[k].all()]
            remaining = still

    arr[arr > undef] = np.nan
    return arr, int(got.sum()), nsolo