#!/usr/bin/env python3
"""Weekly flux fields kept ON THE MODEL GRID, as the source of large-scale features.

weekly_feats.npz already carries the 16 variables sampled AT each observation
point. That is the local view and it is all the pooled experiments have used. The
synoptic setting a heat wave sits in -- the ridge, the gradient across the domain,
the contrast between the point and its surroundings -- is not in it at all, and
cannot be recovered from a point value.

So the same weekly aggregation is kept on the full 32x27 window instead of being
sampled at points, and exp3.py derives box means, domain means and gradients from
it. Deliberately NOT stored as EOF principal components: the EC project showed a
national scalar broadcast to every point is worthless, because it can only move
the national level and the national level is not what fails.

Output: outputs/Cube/weekly_grid.f4  float32 (ninit, 6, 16, 32, 27) ~ 1.0 GB
"""

from __future__ import annotations

import argparse
import multiprocessing as mp
import os
import sys

import numpy as np

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

GRIDF = os.path.join(C.OUT, "Cube", "weekly_grid.f4")
ITMAX = C.FLUX_NAMES.index("tmax")
ITMIN = C.FLUX_NAMES.index("tmin")


def _shape():
    n, nlt, nv, ny, nx = D.shapes()
    return n, len(C.WEEKS), nv, ny, nx


def _one(i):
    n, nlt, nv, ny, nx = D.shapes()
    daily = np.memmap(D.DAILY, dtype="float32", mode="r", shape=(n, nlt, nv, ny, nx))
    out = np.memmap(GRIDF, dtype="float32", mode="r+", shape=_shape())
    cube = np.asarray(daily[i])                                  # (46, nv, ny, nx)
    for w, (_, idx) in enumerate(C.WEEKS):
        sub = cube[idx]
        with np.errstate(invalid="ignore"):
            out[i, w] = np.nanmean(sub, axis=0)
            out[i, w, ITMAX] = np.nanmax(sub[:, ITMAX], axis=0)
            out[i, w, ITMIN] = np.nanmin(sub[:, ITMIN], axis=0)
    del out
    return i


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--nproc", type=int,
                    default=int(os.environ.get("SLURM_CPUS_PER_TASK", 8)))
    a = ap.parse_args()
    sh = _shape()
    print("dung %s  shape %s  (%.2f GB)"
          % (GRIDF, sh, np.prod(sh) * 4 / 1e9), flush=True)
    np.memmap(GRIDF, dtype="float32", mode="w+", shape=sh).flush()
    with mp.Pool(a.nproc) as pool:
        for k, _ in enumerate(pool.imap_unordered(_one, range(sh[0]), chunksize=8)):
            if (k + 1) % 500 == 0:
                print("  %d/%d" % (k + 1, sh[0]), flush=True)
    print("xong")


if __name__ == "__main__":
    main()