#!/usr/bin/env python3
"""Stage 2 sections 5.2-5.3: z-scores and the assembled AI input store.

Writes into $ROOT/outputs/AI_Inp:

  ec_z.f4     (ninit, 46, 40, npt)  z-score of every EC variable of table 1
  obs_z.f4    (ninit, 46,  3, npt)  z-score of tasmax / tasmin / ln(pr+1)
  derived.f4  (ninit, 46,  2, npt)  sv_mod, ev_mod   (raw, not standardised)
  target.f4   (ninit, 46,  3, npt)  tx_obs (deg C), sev_obs, evt_obs
  static.npz                        lead, doy_cos, doy_sin, lat, lon, region,
                                    Tx_thres_obs and Tx_thres_mod per sample/point

The flat sample index of section 5.3 is `init_index * 46 + lt`, so the arrays
already have the 20 x 104 x 46 block layout the docx describes; the observation
arrays are kept strictly as targets and are never offered to a model as input.
"""

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
import build_cube as B
import calc_climatology as K

AI_DIR = os.path.join(C.OUT, "AI_Inp")
EC_Z = os.path.join(AI_DIR, "ec_z.f4")
OBS_Z = os.path.join(AI_DIR, "obs_z.f4")
DERIVED = os.path.join(AI_DIR, "derived.f4")
TARGET = os.path.join(AI_DIR, "target.f4")
STATIC = os.path.join(AI_DIR, "static.npz")

DERIVED_NAMES = ["sv_mod", "ev_mod"]
TARGET_NAMES = ["tx_obs", "sev_obs", "evt_obs"]
STD_FLOOR = 1e-3


def _z_one_var(j, ninit, npt, doy, kind):
    """z-score of one variable across all lead times; returns (j, array)."""
    nv = len(B.EC_VARS); nov = len(B.OBS_VARS)
    ec, obs = B.open_cubes(ninit, npt, mode="r")
    if kind == "ec":
        cube = ec
        mean = np.memmap(K.EC_MEAN, "f4", "r", shape=(nv, C.NLT, 365, npt))[j]
        std = np.memmap(K.EC_STD, "f4", "r", shape=(nv, C.NLT, 365, npt))[j]
        get_clim = lambda k: (mean[k][doy[:, k] - 1], std[k][doy[:, k] - 1])
    else:
        cube = obs
        mean = np.memmap(K.OBS_MEAN, "f4", "r", shape=(nov, 365, npt))[j]
        std = np.memmap(K.OBS_STD, "f4", "r", shape=(nov, 365, npt))[j]
        get_clim = lambda k: (mean[doy[:, k] - 1], std[doy[:, k] - 1])

    out = np.empty((ninit, C.NLT, npt), dtype="f4")
    for k in range(C.NLT):
        m, s = get_clim(k)
        out[:, k, :] = (np.asarray(cube[:, k, j, :]) - m) / np.maximum(s, STD_FLOOR)
    return j, out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--jobs", type=int, default=20)
    ap.add_argument("--skip-z", action="store_true",
                    help="reuse ec_z/obs_z on disk and rebuild only the parts that "
                         "depend on Tx_thres (derived, target, static)")
    args = ap.parse_args()

    grid = C.Grid()
    meta = B.load_meta()
    ninit = int(meta["ninit"][0]); npt = int(meta["npt"][0])
    doy = meta["doy"]; init_year = meta["init_year"]
    nv, nov = len(B.EC_VARS), len(B.OBS_VARS)
    os.makedirs(AI_DIR, exist_ok=True)

    if args.skip_z:
        print("[z] skipped, reusing %s and %s" % (EC_Z, OBS_Z), flush=True)
    else:
        print("[z] EC variables ...", flush=True)
        z = np.memmap(EC_Z, "f4", "w+", shape=(ninit, C.NLT, nv, npt))
        for j, arr in Parallel(n_jobs=args.jobs, verbose=5,
                               return_as="generator_unordered")(
                delayed(_z_one_var)(j, ninit, npt, doy, "ec") for j in range(nv)):
            z[:, :, j, :] = arr
        z.flush(); del z

        print("[z] observation variables ...", flush=True)
        zo = np.memmap(OBS_Z, "f4", "w+", shape=(ninit, C.NLT, nov, npt))
        for j, arr in Parallel(n_jobs=min(args.jobs, nov), verbose=5,
                               return_as="generator_unordered")(
                delayed(_z_one_var)(j, ninit, npt, doy, "obs") for j in range(nov)):
            zo[:, :, j, :] = arr
        zo.flush(); del zo

    print("[sev] severity and event flags ...", flush=True)
    with Dataset(os.path.join(C.OUT, "Cli", "Qp_tx_obs.nc")) as ds:
        thr_obs = grid.to_points(np.ma.filled(ds["Tx_thres_obs"][:], np.nan).astype("f4"))
    with Dataset(os.path.join(C.OUT, "Cli", "Qp_tx_mod.nc")) as ds:
        thr_mod = grid.to_points(np.ma.filled(ds["Tx_thres_mod"][:], np.nan).astype("f4"))

    ec, obs = B.open_cubes(ninit, npt, mode="r")
    jm = B.EC_VARS.index("mx2t6"); jo = B.OBS_VARS.index("tasmax")
    der = np.memmap(DERIVED, "f4", "w+", shape=(ninit, C.NLT, 2, npt))
    tgt = np.memmap(TARGET, "f4", "w+", shape=(ninit, C.NLT, 3, npt))
    th_o_s = np.empty((ninit, C.NLT, npt), dtype="f4")
    th_m_s = np.empty((ninit, C.NLT, npt), dtype="f4")
    for k in range(C.NLT):
        d = doy[:, k] - 1
        tm = np.asarray(ec[:, k, jm, :]); to = np.asarray(obs[:, k, jo, :])
        thm = thr_mod[k][d]; tho = thr_obs[d]
        th_m_s[:, k, :] = thm; th_o_s[:, k, :] = tho
        sm = tm - thm; so = to - tho
        der[:, k, 0, :] = sm
        der[:, k, 1, :] = (sm >= 0).astype("f4")
        tgt[:, k, 0, :] = to
        tgt[:, k, 1, :] = so
        tgt[:, k, 2, :] = (so >= 0).astype("f4")
    der.flush(); tgt.flush()

    lead = (np.arange(C.NLT, dtype="f4") / 45.0)
    np.savez(STATIC,
             ec_vars=np.array(B.EC_VARS), obs_vars=np.array(B.OBS_VARS),
             derived_names=np.array(DERIVED_NAMES), target_names=np.array(TARGET_NAMES),
             lead=lead,
             doy=doy.astype("i4"),
             doy_cos=np.cos(2 * np.pi * doy / 365.0).astype("f4"),
             doy_sin=np.sin(2 * np.pi * doy / 365.0).astype("f4"),
             init_year=init_year, init_dates=meta["init_dates"],
             lat=grid.lat[grid.iy].astype("f4"), lon=grid.lon[grid.ix].astype("f4"),
             region=grid.reg.astype("i4"),
             thr_obs_doy=thr_obs, thr_mod_lt_doy=thr_mod,
             ninit=np.array([ninit]), npt=np.array([npt]))

    ev = np.asarray(tgt[:, :, 2, :])
    print("[stats] observed EHD base rate = %.4f | model = %.4f"
          % (np.nanmean(ev), np.nanmean(np.asarray(der[:, :, 1, :]))), flush=True)
    print("[done] -> %s" % AI_DIR, flush=True)


if __name__ == "__main__":
    main()