#!/usr/bin/env python3
"""Stage 2 section 5.1: model and observed climatology on a 31-day moving window.

For every EC variable, lead time and day of year:
    mean_ec (40, 46, 365, npt)   std_ec (40, 46, 365, npt)
For the observations the climatology is a function of day of year only (an
observed field does not depend on the forecast lead time), pooled over every
sample that is valid on that day:
    mean_obs (3, 365, npt)       std_obs (3, 365, npt)

cp, tp and pr already carry the ln(x+1) transform applied in build_cube.py.

The window sums are built by first binning every sample onto its own day of year
and then running a circular 31-day sum over those bins, which is exact and costs
O(365) per series instead of O(365 x window).
"""

from __future__ import annotations

import argparse
import os
import sys

import numpy as np
from joblib import Parallel, delayed

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

CLI_DIR = os.path.join(C.OUT, "Cli")
EC_MEAN = os.path.join(CLI_DIR, "clim_ec_mean.f4")
EC_STD = os.path.join(CLI_DIR, "clim_ec_std.f4")
OBS_MEAN = os.path.join(CLI_DIR, "clim_obs_mean.f4")
OBS_STD = os.path.join(CLI_DIR, "clim_obs_std.f4")
CLIM_META = os.path.join(CLI_DIR, "clim_meta.npz")

STD_FLOOR = 1e-6


def circular_window_sum(a, half=C.WIN_HALF):
    """Sum over a centred, circular window along axis 0 of length 365."""
    n = a.shape[0]
    pad = np.concatenate([a[n - half:], a, a[:half]], axis=0)
    cs = np.cumsum(pad, axis=0)
    cs = np.concatenate([np.zeros((1,) + a.shape[1:], dtype=cs.dtype), cs], axis=0)
    width = 2 * half + 1
    return cs[width:width + n] - cs[:n]


def bin_by_doy(values, doy, npt, ndoy=365):
    """Per-day sums of x, x^2 and the valid count.  values (nsample, npt)."""
    ok = np.isfinite(values)
    x = np.where(ok, values, 0.0).astype("f8")
    flat = (doy[:, None] - 1) * npt + np.arange(npt)[None, :]
    flat = np.broadcast_to(flat, values.shape).ravel()
    size = ndoy * npt
    s0 = np.bincount(flat, weights=ok.ravel().astype("f8"), minlength=size)
    s1 = np.bincount(flat, weights=x.ravel(), minlength=size)
    s2 = np.bincount(flat, weights=(x * x).ravel(), minlength=size)
    return (s0.reshape(ndoy, npt), s1.reshape(ndoy, npt), s2.reshape(ndoy, npt))


def moments(s0, s1, s2):
    """Window mean and standard deviation from binned sums."""
    n = circular_window_sum(s0)
    a = circular_window_sum(s1)
    b = circular_window_sum(s2)
    with np.errstate(invalid="ignore", divide="ignore"):
        mean = np.where(n > 0, a / np.maximum(n, 1), np.nan)
        var = np.where(n > 1, b / np.maximum(n, 1) - mean * mean, np.nan)
    std = np.sqrt(np.maximum(var, 0.0))
    return mean.astype("f4"), std.astype("f4")


def _ec_one_var(j, ninit, npt, doy, keep=None):
    """Climatology of EC variable j for all 46 lead times.

    `keep` selects the initialisations the statistics are accumulated over.  It is
    the training-year mask: the mean and standard deviation that standardise a
    feature are themselves fitted quantities, so letting the validation years into
    them leaks, mildly, through the scaling.
    """
    ec, _ = B.open_cubes(ninit, npt, mode="r")
    mean = np.empty((C.NLT, 365, npt), dtype="f4")
    std = np.empty((C.NLT, 365, npt), dtype="f4")
    for k in range(C.NLT):
        col = np.asarray(ec[:, k, j, :])
        d = doy[:, k]
        if keep is not None:
            col, d = col[keep], d[keep]
        s0, s1, s2 = bin_by_doy(col, d, npt)
        mean[k], std[k] = moments(s0, s1, s2)
    del ec
    return j, mean, std


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--jobs", type=int, default=20)
    ap.add_argument("--all-years", action="store_true",
                    help="accumulate the climatology over all 20 years (the old "
                         "behaviour); the default is training years only")
    args = ap.parse_args()

    meta = B.load_meta()
    ninit = int(meta["ninit"][0]); npt = int(meta["npt"][0])
    doy = meta["doy"]                                    # (ninit, 46)
    nv = len(B.EC_VARS); nov = len(B.OBS_VARS)
    os.makedirs(CLI_DIR, exist_ok=True)

    keep = None if args.all_years else np.isin(meta["init_year"], C.TRAIN_YEARS)
    print("[split] climatology fitted on %s"
          % ("all %d initialisations" % ninit if keep is None
             else "%d training initialisations of %d" % (keep.sum(), ninit)),
          flush=True)

    print("[ec] %d variables x %d lead times ..." % (nv, C.NLT), flush=True)
    m = np.memmap(EC_MEAN, dtype="f4", mode="w+", shape=(nv, C.NLT, 365, npt))
    s = np.memmap(EC_STD, dtype="f4", mode="w+", shape=(nv, C.NLT, 365, npt))
    for j, mean, std in Parallel(n_jobs=args.jobs, verbose=5, return_as="generator_unordered")(
            delayed(_ec_one_var)(j, ninit, npt, doy, keep) for j in range(nv)):
        m[j] = mean; s[j] = std
    m.flush(); s.flush(); del m, s

    print("[obs] %d variables (day of year only) ..." % nov, flush=True)
    _, obs = B.open_cubes(ninit, npt, mode="r")
    om = np.empty((nov, 365, npt), dtype="f4")
    os_ = np.empty((nov, 365, npt), dtype="f4")
    sel = slice(None) if keep is None else keep
    flat_doy = doy[sel].ravel()
    for j in range(nov):
        vals = np.asarray(obs[:, :, j, :])[sel].reshape(-1, npt)
        s0, s1, s2 = bin_by_doy(vals, flat_doy, npt)
        om[j], os_[j] = moments(s0, s1, s2)
    np.memmap(OBS_MEAN, dtype="f4", mode="w+", shape=om.shape)[:] = om
    np.memmap(OBS_STD, dtype="f4", mode="w+", shape=os_.shape)[:] = os_

    np.savez(CLIM_META,
             ec_vars=np.array(B.EC_VARS), obs_vars=np.array(B.OBS_VARS),
             ec_shape=np.array([nv, C.NLT, 365, npt]),
             obs_shape=np.array([nov, 365, npt]),
             window=np.array([2 * C.WIN_HALF + 1]), std_floor=np.array([STD_FLOOR]),
             train_only=np.array([keep is not None]),
             train_years=np.array(C.TRAIN_YEARS))
    print("[done] -> %s" % CLI_DIR, flush=True)


if __name__ == "__main__":
    main()