#!/usr/bin/env python3
"""Aggregate the 6-hourly flux extraction to daily values on the 46 lead days.

AXIS LAYOUT -- write it down and check it, a wrong reshape here does not raise,
it just produces plausible-looking wrong numbers:

    outputs/Cube/flux_daily.f4   float32  (NINIT, NLT, NFVAR, NY, NX)
    outputs/Cube/tmax_B.f4       float32  (NINIT, NLT, NY, NX)

NINIT follows outputs/meta/init_both.txt in order; NLT = 46 lead days, lead k is
valid on init + k days; NFVAR = 16 flux variables in cfs_common.FLUX_VARS order;
NY x NX = 32 x 27, latitude ascending.

Day alignment (developing.txt 0.6). The 6-hourly step s is valid at init+6(s+1) h.
    A (default, in flux_daily.f4): lead k uses steps 4k..4k+3
        = buckets ending 06Z, 12Z, 18Z of day D and 00Z of D+1 = 07 LT D -> 07 LT D+1
    B (in tmax_B.f4): lead k uses steps 4k-1..4k+2
        = buckets ending 00Z, 06Z, 12Z, 18Z of D = 07 LT D-1 -> 07 LT D
    Lead 0 has no step -1, so B at lead 0 falls back to the three available buckets.
Which one is right is decided against observations by pick_alignment.py.
"""

from __future__ import annotations

import argparse
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

CUBE = os.path.join(C.OUT, "Cube")
DAILY = os.path.join(CUBE, "flux_daily.f4")
TMAXB = os.path.join(CUBE, "tmax_B.f4")
META = os.path.join(CUBE, "daily_meta.npz")

OPS = {"mean": np.nanmean, "max": np.nanmax, "min": np.nanmin}


def shapes():
    g = C.flux_grid()
    ny, nx = len(g["wlat"]), len(g["wlon"])
    n = len(C.both_inits())
    return n, C.NLT, C.NFVAR, ny, nx


def create():
    os.makedirs(CUBE, exist_ok=True)
    n, nlt, nv, ny, nx = shapes()
    np.memmap(DAILY, dtype="float32", mode="w+", shape=(n, nlt, nv, ny, nx)).flush()
    np.memmap(TMAXB, dtype="float32", mode="w+", shape=(n, nlt, ny, nx)).flush()
    g = C.flux_grid()
    np.savez(META, inits=np.array(C.both_inits()), vars=np.array(C.FLUX_NAMES),
             lat=g["wlat"], lon=g["wlon"], shape=np.array([n, nlt, nv, ny, nx]))
    print("created %s  (%d, %d, %d, %d, %d)" % (DAILY, n, nlt, nv, ny, nx))


def run(chunk, nchunk):
    n, nlt, nv, ny, nx = shapes()
    inits = C.both_inits()
    daily = np.memmap(DAILY, dtype="float32", mode="r+", shape=(n, nlt, nv, ny, nx))
    tmb = np.memmap(TMAXB, dtype="float32", mode="r+", shape=(n, nlt, ny, nx))
    ktmax = C.FLUX_NAMES.index("tmax")

    idx = list(range(chunk, n, nchunk))
    t0 = time.time()
    for c, ii in enumerate(idx, 1):
        a = np.load(C.flux6h_path(inits[ii]))            # (184, 16, ny, nx)
        six = a.reshape(nlt, 4, nv, ny, nx)              # alignment A grouping
        with np.errstate(invalid="ignore"):
            for v, (_, op) in enumerate(C.FLUX_VARS):
                daily[ii, :, v] = OPS[op](six[:, :, v], axis=1)
            # alignment B: steps 4k-1..4k+2; lead 0 keeps only its three buckets
            tx = a[:, ktmax]                             # (184, ny, nx)
            b = np.empty((nlt, ny, nx), dtype=np.float32)
            b[0] = np.nanmax(tx[0:3], axis=0)
            for k in range(1, nlt):
                b[k] = np.nanmax(tx[4 * k - 1:4 * k + 3], axis=0)
            tmb[ii] = b
        if c % 50 == 0 or c == len(idx):
            el = time.time() - t0
            print("[chunk %d] %d/%d  %.1f s  %.2f s/init"
                  % (chunk, c, len(idx), el, el / c), flush=True)
    daily.flush(); tmb.flush()


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


if __name__ == "__main__":
    main()