#!/usr/bin/env python3
"""Stage 2 step 1: pack EC_Combined and OBS_4EC into two dense memmap cubes.

    ec_raw   (ninit, 46, 40, npt) float32   the 40 EC variables of table 1
    obs_raw  (ninit, 46,  3, npt) float32   tasmax, tasmin, pr

`ninit` is 2080 = 20 years x 104 initialisations and the flat sample index of
section 5.3 is simply `init_index * 46 + lt`, so the cube already has the block
structure the docx asks for (20 year blocks, 104 sub-blocks, 46 lead times).

Only the 541 land points of the mask are stored.  Everything downstream
(climatology, z-scores, EOF join, model training) reads these two files instead
of re-opening 2080 netCDF files, which is what dominates the run time.
"""

from __future__ import annotations

import argparse
import os
import sys

import numpy as np
from joblib import Parallel, delayed
from netCDF4 import Dataset

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C

EC_VARS = ["cape", "tcw", "tcc", "sm20", "sm100", "t2m", "d2m", "skt", "st20",
           "st100", "dtd2m", "sshf", "slhf", "ssrd", "strd", "ssr", "str", "ttr",
           "cp", "tp", "srn", "bowen", "msl", "u10", "v10", "uv10",
           "thk_1000_500", "thk_1000_850", "thk_850_700", "thk_925_500",
           "thk_925_700", "ivt_u", "ivt_v", "ivt", "dt_850_500", "dt_925_850",
           "dt_700_500", "w", "mx2t6", "mn2t6"]
OBS_VARS = ["tasmax", "tasmin", "pr"]

# Variables that get the ln(x+1) transform before any statistic is taken.
LOG_VARS = {"cp", "tp", "pr"}

CUBE_DIR = os.path.join(C.OUT, "Cube")
EC_CUBE = os.path.join(CUBE_DIR, "ec_raw.f4")
OBS_CUBE = os.path.join(CUBE_DIR, "obs_raw.f4")
META = os.path.join(CUBE_DIR, "meta.npz")


def cube_shapes(ninit, npt):
    return ((ninit, C.NLT, len(EC_VARS), npt), (ninit, C.NLT, len(OBS_VARS), npt))


def open_cubes(ninit, npt, mode="r"):
    se, so = cube_shapes(ninit, npt)
    return (np.memmap(EC_CUBE, dtype="f4", mode=mode, shape=se),
            np.memmap(OBS_CUBE, dtype="f4", mode=mode, shape=so))


def load_meta():
    z = np.load(META, allow_pickle=True)
    return dict(z)


def _fill_year(year, offset, ninit, grid):
    """Write the 104 initialisations of `year` into their slice of both cubes."""
    ec, obs = open_cubes(ninit, grid.npt, mode="r+")
    inits = C.init_dates(year)
    for i, init in enumerate(inits):
        row = offset + i
        with Dataset(C.ec_file(init)) as ds:
            for j, v in enumerate(EC_VARS):
                a = np.ma.filled(ds[v][:], np.nan).astype("f4")
                a[~np.isfinite(a)] = np.nan
                if v in LOG_VARS:
                    a = np.log1p(np.maximum(a, 0.0))
                ec[row, :, j, :] = grid.to_points(a)
        with Dataset(C.obs_file(init)) as ds:
            for j, v in enumerate(OBS_VARS):
                a = np.ma.filled(ds[v][:], np.nan).astype("f4")
                a[~np.isfinite(a)] = np.nan
                if v in LOG_VARS:
                    a = np.log1p(np.maximum(a, 0.0))
                obs[row, :, j, :] = grid.to_points(a)
    ec.flush(); obs.flush()
    return year, len(inits)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--jobs", type=int, default=20)
    ap.add_argument("--years", type=int, nargs="*", default=C.YEARS)
    args = ap.parse_args()

    grid = C.Grid()
    os.makedirs(CUBE_DIR, exist_ok=True)

    per_year = [len(C.init_dates(y)) for y in args.years]
    offsets = np.concatenate([[0], np.cumsum(per_year)[:-1]])
    ninit = int(sum(per_year))
    se, so = cube_shapes(ninit, grid.npt)
    print("[cube] ec %s = %.2f GB | obs %s = %.2f GB" %
          (se, np.prod(se) * 4 / 1e9, so, np.prod(so) * 4 / 1e9), flush=True)

    # allocate on disk
    np.memmap(EC_CUBE, dtype="f4", mode="w+", shape=se).flush()
    np.memmap(OBS_CUBE, dtype="f4", mode="w+", shape=so).flush()

    Parallel(n_jobs=args.jobs, verbose=5)(
        delayed(_fill_year)(y, int(o), ninit, grid)
        for y, o in zip(args.years, offsets))

    # per-sample metadata
    init_dates, init_year, doy = [], [], []
    for y in args.years:
        for init in C.init_dates(y):
            init_dates.append(init.isoformat())
            init_year.append(y)
            doy.append([C.doy_noleap(d) for d in C.valid_days(init)])
    np.savez(META,
             ec_vars=np.array(EC_VARS), obs_vars=np.array(OBS_VARS),
             init_dates=np.array(init_dates), init_year=np.array(init_year, "i4"),
             doy=np.array(doy, "i4"),                       # (ninit, 46)
             lead=np.arange(C.NLT, dtype="i4"),
             years=np.array(args.years, "i4"),
             log_vars=np.array(sorted(LOG_VARS)),
             npt=np.array([grid.npt]), ninit=np.array([ninit]))
    print("[done] %s" % CUBE_DIR, flush=True)


if __name__ == "__main__":
    main()