"""
nnlab - shared harness for the overnight model search (2026-11-11).

Every model family writes, for each protocol draw (n, rep), ONE array of class
probabilities over ALL 25,992 Amsterdam pixels (row order = the feature table):

    lab.save_pred("my_method", n, rep, P)      # P: (25992, 4), columns = class 1..4

and nothing else. Scoring is done here, identically for everybody:

    python tools/nnlab/lab.py score my_method [other_method ...]
    python tools/nnlab/lab.py board

RULES every module must keep (reviewers check these):
  * Amsterdam LABELS may be read only through lab.shot_labels(d, idx) for the
    indices returned by lab.shots_for(d, n, rep). Never touch d["ya"] directly.
  * Unlabelled Amsterdam features / coordinates may be used freely (transductive).
  * Hyper-parameters are fixed in advance or chosen on Madrid (lab.madrid_split
    gives a spatial pseudo-target). Never chosen by looking at an Amsterdam score.
  * The held-out organisers' data is never opened.

The draws are the protocol's: seed 42, rng(42 + 1000 n + rep), n in
(5, 25, 50, 100, 200). Shots are removed from the scored set.
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent.parent
# ROOT itself holds winners.py when these scripts run from the published bundle
for _p in (ROOT, HERE, ROOT / "src", ROOT / "tools"):
    if str(_p) not in sys.path:
        sys.path.insert(0, str(_p))
THREADS = int(os.environ.get("NNLAB_THREADS", "3"))
os.environ.setdefault("OMP_NUM_THREADS", str(THREADS))
os.environ.setdefault("MKL_NUM_THREADS", str(THREADS))

import numpy as np                      # noqa: E402
import pandas as pd                     # noqa: E402

import winners as W                     # noqa: E402  (verified port of the protocol + rule)

PROC = ROOT / "data" / "processed"
OUT = ROOT / "reports" / "nnlab"
PRED = OUT / "preds"
SCORES = OUT / "scores"
CACHE = OUT / "cache"                   # embeddings, time-series tensors, ...
for _d in (PRED, SCORES, CACHE):
    _d.mkdir(parents=True, exist_ok=True)

SEED = 42
N_LIST = (5, 25, 50, 100, 200)
REPS = int(os.environ.get("NNLAB_REPS", "5"))   # reps 0..4 by default; NNLAB_REPS=12 adds the confirmation draws
CLASSES = np.array([1, 2, 3, 4])
PID = ["px_key", "py_key"]
BLOCK = 67                              # 67 px = 2 km
SMOOTH_R, SMOOTH_ALPHA = 6, 0.5         # teammate's choice, made on Madrid


def torch_setup(seed: int = 0):
    import torch
    torch.set_num_threads(THREADS)
    torch.manual_seed(seed)
    np.random.seed(seed)
    return torch


# --------------------------------------------------------------------------- #
# Data
# --------------------------------------------------------------------------- #
def load(official60: bool = True) -> dict:
    """Madrid + Amsterdam tables. Xm/Xa: our 393 features (float32, raw scale);
    Xm60/Xa60: the organisers' official 60 (first valid observation per year);
    pm/pa: integer pixel coordinates (30 m units); ym: Madrid labels;
    prior: Madrid class shares. Amsterdam labels are kept under a private key."""
    feats = W.feature_names()
    mad = pd.read_parquet(PROC / "madrid_features.parquet")
    ams = pd.read_parquet(PROC / "amsterdam_features.parquet")
    d = dict(
        feats=feats,
        Xm=mad[feats].to_numpy(np.float32), Xa=ams[feats].to_numpy(np.float32),
        ym=mad["age_class"].to_numpy().astype(int),
        pm=mad[PID].to_numpy().astype(np.int64), pa=ams[PID].to_numpy().astype(np.int64),
        _ya=ams["age_class"].to_numpy().astype(int),
    )
    d["prior"] = np.array([(d["ym"] == c).mean() for c in CLASSES])
    if official60:
        import config as C
        f60 = list(C.FEATURE_COLS_60)
        for city, ref, key in (("madrid", mad, "Xm60"), ("amsterdam", ams, "Xa60")):
            o = pd.read_parquet(PROC / f"{city}_official60_first.parquet")
            o = ref[PID].merge(o, on=PID, how="left")
            assert len(o) == len(ref) and o[f60].notna().all().all()
            d[key] = o[f60].to_numpy(np.float32)
        d["feats60"] = f60
    return d


def shots_for(d: dict, n: int, rep: int) -> np.ndarray:
    """Row indices of the n-per-class labelled Amsterdam pixels of draw (n, rep)."""
    return W.shots_for(d["_ya"], n, rep)


def shot_labels(d: dict, idx: np.ndarray) -> np.ndarray:
    """Labels (1..4) of shot rows. The ONLY sanctioned way to read Amsterdam labels."""
    return d["_ya"][np.asarray(idx)]


def draws(reps: int = REPS, n_list=N_LIST):
    for n in n_list:
        for r in range(reps):
            yield n, r


def balanced_class_weights(y) -> np.ndarray:
    return W.balanced_class_weights(y)


# per-city input transforms (label-free) ------------------------------------- #
def robustz(X: np.ndarray, clip: float | None = 10.0) -> np.ndarray:
    """Median/IQR scaling with this matrix's own statistics (call once per city)."""
    X = np.asarray(X, np.float64)
    med = np.median(X, 0)
    q75, q25 = np.percentile(X, [75, 25], 0)
    iqr = np.where((q75 - q25) > 1e-12, q75 - q25, 1.0)
    Z = (X - med) / iqr
    if clip:
        Z = np.clip(Z, -clip, clip)
    return Z.astype(np.float32)


def standardise(X: np.ndarray) -> np.ndarray:
    X = np.asarray(X, np.float64)
    sd = X.std(0)
    sd[sd == 0] = 1
    return ((X - X.mean(0)) / sd).astype(np.float32)


def rankgauss(X: np.ndarray) -> np.ndarray:
    """Per-column rank -> standard normal quantile, within this matrix (one city)."""
    from scipy.special import ndtri
    from scipy.stats import rankdata
    X = np.asarray(X, np.float64)
    R = rankdata(X, axis=0, method="average")
    return ndtri((R - 0.5) / len(X)).astype(np.float32)


# Madrid-only model selection ------------------------------------------------- #
def block_ids(p: np.ndarray, block: int = BLOCK) -> np.ndarray:
    b = p // block
    return (b[:, 0] * 100_003 + b[:, 1]).astype(np.int64)


def madrid_split(pm: np.ndarray, frac: float = 0.25, axis: int = 0, high: bool = True):
    """A spatial pseudo-target inside Madrid: the `frac` of 2 km blocks furthest
    along one axis. Returns (source_idx, pseudo_target_idx). Use it to choose
    epochs / regularisation without touching Amsterdam labels."""
    b = pm[:, axis] // BLOCK
    cut = np.quantile(b, 1 - frac) if high else np.quantile(b, frac)
    tgt = b > cut if high else b < cut
    return np.flatnonzero(~tgt), np.flatnonzero(tgt)


# --------------------------------------------------------------------------- #
# Heads on a frozen backbone
# --------------------------------------------------------------------------- #
def softmax(L: np.ndarray) -> np.ndarray:
    L = L - L.max(1, keepdims=True)
    E = np.exp(L)
    return E / E.sum(1, keepdims=True)


def finetune_last_layer(W0, b0, Zs, ys, lam: float = 1.0, class_balanced: bool = True,
                        max_iter: int = 200):
    """True last-layer fine-tuning: start from the Madrid-trained last layer
    (W0: (4, k), b0: (4,)) and minimise, on the labelled shots only,
        sum_i w_i CE(softmax(W z_i + b), y_i)  +  lam * (||W - W0||^2 + ||b - b0||^2).
    The CE is a SUM, so the pull toward the source weights fades as n grows.
    Returns (W, b). Probabilities: softmax(Z @ W.T + b)."""
    torch = torch_setup(0)
    Z = torch.tensor(np.asarray(Zs), dtype=torch.float64)
    y = torch.tensor(np.asarray(ys) - 1, dtype=torch.long)
    w = torch.ones(len(y), dtype=torch.float64)
    if class_balanced:
        w = torch.tensor(balanced_class_weights(np.asarray(ys))[np.asarray(ys) - 1], dtype=torch.float64)
    W0t = torch.tensor(np.asarray(W0), dtype=torch.float64)
    b0t = torch.tensor(np.asarray(b0), dtype=torch.float64)
    Wt, bt = W0t.clone().requires_grad_(True), b0t.clone().requires_grad_(True)
    opt = torch.optim.LBFGS([Wt, bt], max_iter=max_iter, line_search_fn="strong_wolfe")

    def closure():
        opt.zero_grad()
        ce = torch.nn.functional.cross_entropy(Z @ Wt.T + bt, y, reduction="none")
        loss = (w * ce).sum() + lam * (((Wt - W0t) ** 2).sum() + ((bt - b0t) ** 2).sum())
        loss.backward()
        return loss

    opt.step(closure)
    return Wt.detach().numpy(), bt.detach().numpy()


def joint_head(Zm, ym, Zs, ys, shot_weight: float = 30.0, C: float = 1.0, max_iter: int = 400):
    """Last layer re-fitted on frozen embeddings of Madrid + shots (shots
    weighted `shot_weight`x, classes balanced) - the frozen-backbone analogue
    of the retrained LightGBM. Returns a fitted sklearn LogisticRegression
    (predict_proba columns = classes 1..4)."""
    from sklearn.linear_model import LogisticRegression
    X = np.vstack([Zm, Zs])
    y = np.r_[np.asarray(ym), np.asarray(ys)]
    w = np.r_[np.ones(len(ym)), np.full(len(ys), shot_weight)]
    w = w * balanced_class_weights(y)[y - 1]
    return LogisticRegression(C=C, max_iter=max_iter).fit(X, y, sample_weight=w)


# --------------------------------------------------------------------------- #
# Prediction store
# --------------------------------------------------------------------------- #
def pred_path(method: str, n: int, rep: int) -> Path:
    return PRED / method / f"n{n}_r{rep}.npy"


def have_pred(method: str, n: int, rep: int) -> bool:
    return pred_path(method, n, rep).exists()


def save_pred(method: str, n: int, rep: int, P: np.ndarray, meta: dict | None = None):
    """P: (25992, 4) probabilities for ALL Amsterdam pixels, feature-table row
    order. A method that ignores the shots saves once with n=0, rep=0."""
    P = np.asarray(P, np.float32)
    assert P.shape == (25992, 4) and np.isfinite(P).all(), P.shape
    p = pred_path(method, n, rep)
    p.parent.mkdir(parents=True, exist_ok=True)
    np.save(p, P)
    if meta is not None:
        (p.parent / "meta.json").write_text(json.dumps(meta, indent=1, default=str))


def load_pred(method: str, n: int, rep: int) -> np.ndarray:
    p = pred_path(method, n, rep)
    if not p.exists():
        p = pred_path(method, 0, 0)       # shot-independent method
    return np.load(p).astype(np.float64)


# --------------------------------------------------------------------------- #
# Scoring (the only place Amsterdam labels are compared with predictions)
# --------------------------------------------------------------------------- #
_ADJ = {}


def adjacency(coords: np.ndarray, r: float = SMOOTH_R):
    """Sparse 0/1 matrix: pixels within r (Euclidean, pixel units), self included."""
    from scipy.sparse import coo_matrix, identity
    from scipy.spatial import cKDTree
    key = (coords.shape[0], float(r), int(coords[:, 0].sum()), int(coords[:, 1].sum()))
    if key not in _ADJ:
        pr = cKDTree(coords.astype(float)).query_pairs(r, output_type="ndarray")
        n = len(coords)
        A = coo_matrix((np.ones(len(pr)), (pr[:, 0], pr[:, 1])), shape=(n, n)).tocsr()
        _ADJ[key] = (A + A.T + identity(n, format="csr")).tocsr()
    return _ADJ[key]


def smooth(P: np.ndarray, coords: np.ndarray, r: float = SMOOTH_R, alpha: float = SMOOTH_ALPHA,
           keep: np.ndarray | None = None) -> np.ndarray:
    """Teammate's neighbour smoothing: alpha * own + (1 - alpha) * mean over
    pixels within r (self included). With `keep`, only those rows take part
    (P must then already be P_all[keep]) - shots are never neighbours."""
    A = adjacency(coords, r)
    if keep is not None:
        A = A[keep][:, keep]
    deg = np.asarray(A.sum(1)).ravel()
    return alpha * P + (1 - alpha) * (A @ P) / deg[:, None]


def half_A(pa: np.ndarray) -> np.ndarray:
    """Checkerboard of 2 km blocks. Half A is for comparing variants; half B is
    looked at only to confirm the final picks."""
    b = pa // BLOCK
    return ((b[:, 0] + b[:, 1]) % 2) == 0


def score_method(d: dict, method: str, reps: int = REPS, n_list=N_LIST, verbose: bool = True) -> dict:
    ya, pa, prior = d["_ya"], d["pa"], d["prior"]
    inA = half_A(pa)
    rows = []
    for n, rep in draws(reps, n_list):
        if not (have_pred(method, n, rep) or have_pred(method, 0, 0)):
            continue
        P_all = load_pred(method, n, rep)
        shot = shots_for(d, n, rep)
        keep = np.setdiff1d(np.arange(len(ya)), shot)
        P, y = P_all[keep], ya[keep]
        lab = {
            "argmax": CLASSES[P.argmax(1)],
            "rule": W.decide(P, prior),
            "rule_smooth": W.decide(smooth(P, pa, keep=keep), prior),
        }
        row = {"n": n, "rep": rep}
        for k, l in lab.items():
            row[k] = W.macro_f1(y, l)
            row[k + "_A"] = W.macro_f1(y[inA[keep]], l[inA[keep]])
            row[k + "_B"] = W.macro_f1(y[~inA[keep]], l[~inA[keep]])
        rows.append(row)
    if not rows:
        return {}
    df = pd.DataFrame(rows)
    summ = df.groupby("n").mean(numeric_only=True).drop(columns="rep")
    std = df.groupby("n").std(numeric_only=True).drop(columns="rep")
    out = {"method": method, "draws": rows, "n_draws": int(len(rows)),
           "mean": {c: {int(n): float(v) for n, v in summ[c].items()} for c in summ.columns},
           "std": {c: {int(n): float(v) for n, v in std[c].items()} for c in std.columns}}
    (SCORES / f"{method}.json").write_text(json.dumps(out, indent=1))
    if verbose:
        for k in ("argmax", "rule", "rule_smooth"):
            m = summ[k]
            print(f"{method:<34s} {k:<12s} " + "  ".join(f"n{n}={m[n]:.4f}" for n in m.index)
                  + f"  avg={m.mean():.4f}  peak={m.max():.4f}  ({len(rows)} draws)", flush=True)
    return out


def board(metric_suffix: str = "") -> pd.DataFrame:
    """All scored methods, one row per (method, decision). suffix '' / '_A' / '_B'."""
    rows = []
    for f in sorted(SCORES.glob("*.json")):
        s = json.loads(f.read_text())
        for k in ("argmax", "rule", "rule_smooth"):
            m = s["mean"].get(k + metric_suffix)
            if not m:
                continue
            r = {"method": s["method"], "decision": k, **{f"n{n}": m.get(str(n), np.nan) for n in N_LIST}}
            v = [r[f"n{n}"] for n in N_LIST]
            r["avg"], r["peak"], r["draws"] = float(np.nanmean(v)), float(np.nanmax(v)), s["n_draws"]
            rows.append(r)
    df = pd.DataFrame(rows).sort_values("avg", ascending=False)
    df.to_csv(OUT / f"leaderboard{metric_suffix}.csv", index=False)
    return df


if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else "board"
    if cmd == "score":
        d = load(official60=False)
        names = sys.argv[2:] or sorted(p.name for p in PRED.iterdir() if p.is_dir())
        t0 = time.time()
        for m in names:
            score_method(d, m)
        print(f"scored {len(names)} methods in {time.time() - t0:.0f}s")
    else:
        suffix = sys.argv[2] if len(sys.argv) > 2 else ""
        pd.set_option("display.width", 250)
        print(board(suffix).round(4).to_string(index=False))