"""Scoring for the C. elegans graph-dynamics benchmark (L2, Atanas et al. 2023, 38 animals, 89 identified neurons).

Task: for every animal, predict standardized activity h = 4 frames (2.4 s) ahead. A submission gives, per animal, an
array P of shape (T, 89) where P[t] is the prediction of X[t + 4]. Each animal is scored on held-out data: train on the
animals of the other four folds of the frozen split, predict the animals of the fifth.

Per animal: pooled R² = 1 - SSE / SST over valid (t, node) targets, SST around each node's mean in that animal. Entries are
scored only where the target is valid and the node's own input at time t was valid (the rule every reference model
follows), so all submissions are scored on the same entries. Summary: mean over animals with a 2,000-draw animal
bootstrap; comparisons with the reference models are paired by animal.
"""
from __future__ import annotations

import io
import json
from pathlib import Path

import numpy as np

DATA = Path(__file__).resolve().parent / "data"
H = 4
REFERENCE_LABELS = {
    "gnn_connectome": "GNN, connectome edges",
    "gnn_shuffled": "GNN, degree-preserving shuffles (mean of 3)",
    "gnn_geometry": "GNN, nearest physical neighbours",
    "gnn_none": "GNN, no edges",
    "gnn_full": "GNN, fully connected",
    "persistence": "Persistence (predict no change)",
}


def load_benchmark():
    z = np.load(DATA / "benchmark.npz")
    animals = sorted(k[2:] for k in z.files if k.startswith("X/"))
    X = {a: z[f"X/{a}"].astype(float) for a in animals}
    M = {a: z[f"M/{a}"] for a in animals}
    meta = json.loads((DATA / "benchmark_meta.json").read_text())
    ref = json.loads((DATA / "reference_scores.json").read_text())
    shuf = [k for k in ref if k.startswith("gnn_shuffled_")]
    ref["gnn_shuffled"] = {a: float(np.mean([ref[s][a] for s in shuf])) for a in animals}
    return animals, X, M, meta, ref


def score_animal(X, M, P, h=H):
    """Pooled R² for one animal. Same arithmetic as wormlab.gd.bench.forecast.score_animal."""
    P = np.where(M, P, np.nan)
    T = len(X)
    y, yh = X[h:], P[: T - h]
    ok = M[h:] & np.isfinite(yh)
    if ok.sum() < 50:
        return None
    mu = np.array([y[ok[:, k], k].mean() if ok[:, k].any() else 0.0 for k in range(y.shape[1])])
    sse = float(((y - yh) ** 2)[ok].sum())
    sst = float(((y - mu) ** 2)[ok].sum())
    return 1 - sse / sst


def boot_mean(vals, n=2000, seed=0):
    v = np.asarray(vals, float)
    rng = np.random.default_rng(seed)
    b = rng.choice(v, (n, len(v))).mean(1)
    return float(v.mean()), float(np.percentile(b, 2.5)), float(np.percentile(b, 97.5))


def read_submission(raw: bytes, animals, X):
    """Accept an .npz whose keys are animal ids (optionally prefixed 'P/'), each a (T, 89) float array."""
    z = np.load(io.BytesIO(raw), allow_pickle=False)
    sub = {(k[2:] if k.startswith("P/") else k): z[k] for k in z.files}
    missing = [a for a in animals if a not in sub]
    if missing:
        raise ValueError(f"{len(missing)} of {len(animals)} animals missing, e.g. {missing[:3]}. All 38 are required.")
    extra = sorted(set(sub) - set(animals))
    if extra:
        raise ValueError(f"Unknown keys: {extra[:3]}")
    out = {}
    for a in animals:
        P = np.asarray(sub[a], float)
        if P.shape != X[a].shape:
            raise ValueError(f"{a}: shape {P.shape}, expected {X[a].shape}")
        out[a] = P
    return out


def score_submission(raw: bytes, bench=None):
    animals, X, M, meta, ref = bench or load_benchmark()
    sub = read_submission(raw, animals, X)
    per = {}
    for a in animals:
        r = score_animal(X[a], M[a], sub[a])
        if r is None or not np.isfinite(r):
            raise ValueError(f"{a}: fewer than 50 scorable entries or non-finite score")
        per[a] = r
    you = [per[a] for a in animals]
    rows = [dict(model="Your submission", **dict(zip(("r2", "lo", "hi"), boot_mean(you))))]
    for k, lab in REFERENCE_LABELS.items():
        r = [ref[k][a] for a in animals]
        m, lo, hi = boot_mean(r)
        d, dlo, dhi = boot_mean(np.array(you) - np.array(r))
        rows.append(dict(model=lab, r2=m, lo=lo, hi=hi, diff=d, diff_lo=dlo, diff_hi=dhi))
    return dict(rows=rows, per_animal=per, fold_of=meta["fold_of"])