#!/usr/bin/env python3
"""Raw-CFS skill on the WEEKLY windows, the frame the forecast is actually judged in.

Day-to-day pairing of a 46-day forecast is the wrong frame: phase error alone
destroys it, and the collapse of daily anomaly correlation past ~10 days says
nothing about whether the week is skilful. The operational quantity is the week,
and the predictor is max{daily Tx} over the week -- the same statistic the event
is defined from.

Windows follow the project convention: W1 = lead 0-10, W2 = 11-17, ... W6 = 39-45.

Anomalies are taken against the MODEL's own climatology of the weekly maximum at
(week, day of year, point), fitted on training years only. Using the observed
climatology instead would charge the model's drift to the noise term and
understate its skill -- which is exactly what the daily diagnostic did.
"""

from __future__ import annotations

import argparse
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_obs as B
import build_daily as D


def auc(score, label):
    """Rank-based AUC; ties get average ranks."""
    s, y = np.asarray(score, float), np.asarray(label, bool)
    ok = np.isfinite(s)
    s, y = s[ok], y[ok]
    n1, n0 = int(y.sum()), int((~y).sum())
    if n1 == 0 or n0 == 0:
        return np.nan
    r = np.empty(len(s))
    order = np.argsort(s, kind="mergesort")
    sr = s[order]
    i = 0
    while i < len(sr):
        j = i
        while j + 1 < len(sr) and sr[j + 1] == sr[i]:
            j += 1
        r[order[i:j + 1]] = 0.5 * (i + j) + 1
        i = j + 1
    return (r[y].sum() - n1 * (n1 + 1) / 2.0) / (n1 * n0)


def build():
    """Weekly max Tx from model and observations, plus the EHDE label."""
    n, nlt, nv, ny, nx = D.shapes()
    inits = C.both_inits()
    ktmax = C.FLUX_NAMES.index("tmax")
    daily = np.memmap(D.DAILY, dtype="float32", mode="r", shape=(n, nlt, nv, ny, nx))

    z = np.load(B.TX_NPZ)
    obs, oord, odoy = z["tx"], z["ordinal"], z["doy"]
    row = {int(o): r for r, o in enumerate(oord)}
    ehd, _ = B.ehd_of("train")
    olat, olon = z["lat"][z["jj"]], z["lon"][z["ii"]]
    g = C.flux_grid()
    pj = np.abs(olat[:, None] - g["wlat"][None, :]).argmin(1)
    pi = np.abs(olon[:, None] - g["wlon"][None, :]).argmin(1)

    npt, nw = len(pj), len(C.WEEKS)
    M = np.full((n, nw, npt), np.nan, dtype="float32")   # model weekly max Tx
    O = np.full((n, nw, npt), np.nan, dtype="float32")   # observed weekly max Tx
    E = np.zeros((n, nw, npt), dtype=bool)               # observed EHDE
    doy = np.zeros((n, nw), dtype="i4")                  # doy of window centre
    yrs = np.array([int(s[:4]) for s in inits])

    for i, s in enumerate(inits):
        tx = daily[i, :, ktmax][:, pj, pi] - 273.15      # (46, npt)
        base = C.to_date(s).toordinal()
        for w, (_, idx) in enumerate(C.WEEKS):
            rows = [row.get(base + int(k), -1) for k in idx]
            if min(rows) < 0:
                continue
            M[i, w] = np.nanmax(tx[idx], axis=0)
            O[i, w] = obs[rows].max(axis=0)
            E[i, w] = ehd[rows].any(axis=0)
            doy[i, w] = odoy[rows[len(rows) // 2]]
        if (i + 1) % 500 == 0:
            print("  %d/%d" % (i + 1, n), flush=True)
    return inits, yrs, M, O, E, doy


def climz(X, doy, yrs, train):
    """z-score of X against its own (week, doy, point) climatology, train-only.

    The degenerate-sd guard is RELATIVE to each variable's own spread. An absolute
    floor is a units bug waiting to happen: prate is kg m-2 s-1 with sd ~1e-5, so
    a 1e-3 floor silently voids the whole variable and the design matrix simply
    comes back empty.
    """
    n, nw, npt = X.shape
    Z = np.full_like(X, np.nan)
    for w in range(nw):
        d = doy[:, w]
        bydoy = [np.where((d == q) & train)[0] for q in range(1, 366)]
        mu = np.full((366, npt), np.nan, dtype="float32")
        sd = np.full((366, npt), np.nan, dtype="float32")
        for q in range(1, 366):
            idx = np.concatenate([bydoy[(q - 1 + o) % 365] for o in
                                  range(-C.WIN_HALF, C.WIN_HALF + 1)])
            if len(idx) < 30:
                continue
            mu[q] = np.nanmean(X[idx, w], axis=0)
            sd[q] = np.nanstd(X[idx, w], axis=0)
        gsd = np.nanstd(X[train, w])
        floor = max(1e-6 * float(gsd), 1e-30) if np.isfinite(gsd) else 1e-30
        sd = np.where(sd > floor, sd, np.nan)
        ok = d > 0
        Z[ok, w] = (X[ok, w] - mu[d[ok]]) / sd[d[ok]]
    return Z


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache", default=os.path.join(C.OUT, "Cube", "weekly.npz"))
    ap.add_argument("--rebuild", action="store_true")
    a = ap.parse_args()

    if a.rebuild or not os.path.exists(a.cache):
        inits, yrs, M, O, E, doy = build()
        np.savez(a.cache, inits=np.array(inits), yrs=yrs, M=M, O=O, E=E, doy=doy)
    z = np.load(a.cache)
    yrs, M, O, E, doy = z["yrs"], z["M"], z["O"], z["E"], z["doy"]

    train = np.isin(yrs, C.TRAIN_YEARS)
    valid = np.isin(yrs, C.VALID_YEARS)
    ZM = climz(M, doy, yrs, train)
    ZO = climz(O, doy, yrs, train)

    print("\nRAW CFS SKILL ON WEEKLY MAX Tx -- validation years %s, %d inits, %d points"
          % (C.VALID_YEARS, int(valid.sum()), M.shape[2]))
    print("%-4s %8s %8s %8s %8s %8s %9s %9s"
          % ("week", "ME", "MAE", "RMSE", "CORR", "AUC", "base", "AUC_req"))
    from signal_ceiling import far_to_pofd, auc_needed
    for w, name in enumerate(C.WEEK_NAMES):
        m, o = M[valid, w].ravel(), O[valid, w].ravel()
        zm, zo = ZM[valid, w].ravel(), ZO[valid, w].ravel()
        e = E[valid, w].ravel()
        ok = np.isfinite(m) & np.isfinite(o) & np.isfinite(zm) & np.isfinite(zo)
        br = float(e[ok].mean())
        req = auc_needed(0.60, min(far_to_pofd(0.60, 0.20, br), 0.999))
        print("%-4s %8.3f %8.3f %8.3f %8.4f %8.4f %9.4f %9.4f"
              % (name, float((m[ok] - o[ok]).mean()),
                 float(np.abs(m[ok] - o[ok]).mean()),
                 float(np.sqrt(((m[ok] - o[ok]) ** 2).mean())),
                 float(np.corrcoef(zm[ok], zo[ok])[0, 1]),
                 auc(zm[ok], e[ok]), br, req))

    print("\nAUC by region (validation years)")
    reg = np.load(B.TX_NPZ)["reg"]
    print("%-20s %s" % ("region", "  ".join("%6s" % w for w in C.WEEK_NAMES)))
    for r in range(1, 8):
        k = reg == r
        out = []
        for w in range(len(C.WEEKS)):
            zm, e = ZM[valid, w][:, k].ravel(), E[valid, w][:, k].ravel()
            ok = np.isfinite(zm)
            out.append(auc(zm[ok], e[ok]))
        print("%-20s %s" % (C.REGION_NAMES[r], "  ".join("%6.3f" % v for v in out)))


if __name__ == "__main__":
    main()