"""
tab_nn - tabular neural networks with a FROZEN backbone and last-layer adaptation.

    python tools/nnlab/tab_nn.py [--smoke] [--only <substring>]

Backbones (trained ONCE on all Madrid labels, class-balanced CE, AdamW lr 1e-3,
wd 1e-4, batch 512, seeds 0,1,2):
  mlp   Linear-BN-ReLU-Drop(0.2) x2 (256) -> Linear(64)-ReLU -> Linear(64,64)-ReLU -> Linear(64,4)
  res   ResNet-MLP (Gorishniy et al. 2021): Linear(256), 3 x [x + Lin(Drop(ReLU(Lin(BN(x)))))],
        BN-ReLU -> Linear(64)-ReLU -> Linear(64,4)
  dann  mlp + domain head (Lin(64,64)-ReLU-Lin(64,1)) on the first 64-d features through
        gradient reversal, lambda = 2/(1+exp(-10p))-1, adversary vs UNLABELLED Amsterdam rows
        (as tools/compare_teammate.py train_dann; Madrid and Amsterdam batches go through
        the net separately)
  coral mlp + deep-CORAL loss (weight 1.0) between the penultimate features of a Madrid
        batch and an unlabelled Amsterdam batch
Inputs: rz = 393 feats lab.robustz per city; rg = 393 feats lab.rankgauss per city;
        o60 = official 60 feats lab.standardise per city.
The EMBEDDING is the input of the final Linear(.,4); (W0, b0) are that layer's weights.

Heads per protocol draw (probabilities averaged over the 3 seeds before saving):
  zs    source head unchanged                          (saved once, n=0 rep=0)
  ft    lab.finetune_last_layer(W0, b0, Z_shot, y_shot, lam)   lam chosen on Madrid
  joint the objective of lab.joint_head(Z_madrid, y_madrid, Z_shot, y_shot, shot_weight=30)
        (same weights, C=1) on embeddings standardised with the Madrid embedding mean/sd
        (same affine map for the shots), solved with newton-cholesky warm-started from
        the Madrid-only fit (see head_joint: lab.joint_head's lbfgs does not converge
        in 400 iterations and would cost ~2 h over the run)
  proto prototypes (n_c * shot mean + k * Madrid class mean)/(n_c + k), k=10, embeddings
        centred by their own city's mean and L2-normalised; P = softmax(-20 * sq. distance)
Method names: tab_{backbone}_{input}_{head}. "rg" is run for mlp and dann only (time budget),
so 10 backbone x input combos x 4 heads = 40 methods.

Madrid-only choices (lab.madrid_split, pseudo-target = the 25 % of 2 km blocks furthest
along axis 0; transforms computed separately on the source part and on the pseudo-target,
exactly as they are per city; adversary / CORAL target = the unlabelled pseudo-target):
  * epochs per backbone type from {15, 30, 45}: input rz, seed 0, trained on the source
    part, scored by macro F1 of the decision rule (W.decide, source-part prior) on the
    pseudo-target. Non-adversarial nets: one 45-epoch run evaluated after 15/30/45 epochs
    (identical to separate runs: constant lr, same seed). dann: separate runs, since its
    lambda schedule depends on the total epoch count.
  * lam for ft, ONE value for all backbones and n, from {0.1, 1, 10, 100}: on each
    backbone's selected source-part net (rz, seed 0), draw n in {5, 50} per class from the
    pseudo-target (5 draws, rng(7000 + 100 n + rep)), fine-tune, score the remaining
    pseudo-target with the rule; pick the lam with the best mean over backbones x n x draws.
  Both are cached in CACHE/tab_nn/selection.json and copied into every meta.json.

CACHE LAYOUT (lab.CACHE / "tab_nn", or "tab_nn_SMOKE" for --smoke):
  selection.json                         epochs per backbone, lam, all pseudo-target scores
  sel_{backbone}.npz                     source-part (madrid_split) net at the chosen epochs:
                                         Zs, Zt (source / pseudo-target embeddings), W0, b0
  emb_{backbone}_{input}_s{seed}.npz     one file per (backbone, input, seed); inputs rz,o60
                                         for all backbones, rg for mlp and dann; seeds 0,1,2:
      Zm  (76263, 64) float32  Madrid embeddings (feature-table row order, eval mode)
      Za  (25992, 64) float32  Amsterdam embeddings (feature-table row order)
      W0  (4, 64)     float32  final layer weight;  b0 (4,) final layer bias
      epochs  ()  int          epochs it was trained for
  Logits = Z @ W0.T + b0, columns = classes 1..4.
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import lab                               # noqa: E402
import numpy as np                       # noqa: E402

W = lab.W
T0 = time.time()

BACKBONES = ("mlp", "res", "dann", "coral")
INPUTS = ("rz", "rg", "o60")
# The full run on 3 shared CPU threads is far above 60 min, so per the brief the
# "rg" input is dropped for "res" and "coral" (10 backbone x input combos, 40 methods).
INPUTS_FOR = {"mlp": INPUTS, "dann": INPUTS, "res": ("rz", "o60"), "coral": ("rz", "o60")}
HEADS = ("zs", "ft", "joint", "proto")
SEEDS = (0, 1, 2)
EPOCH_GRID = (15, 30, 45)
LAM_GRID = (0.1, 1.0, 10.0, 100.0)
LAM_NS = (5, 50)
LAM_REPS = 5
SHOT_WEIGHT = 30.0
PROTO_K, PROTO_T = 10.0, 20.0
CORAL_W = 1.0
BATCH, LR, WD = 512, 1e-3, 1e-4


def log(msg):
    print(f"[{time.time() - T0:7.1f}s] {msg}", flush=True)


# --------------------------------------------------------------------------- #
# Networks
# --------------------------------------------------------------------------- #
def build_net(kind: str, d_in: int):
    import torch.nn as nn
    from torch.autograd import Function

    class GradReverse(Function):
        @staticmethod
        def forward(ctx, x, lam):
            ctx.lam = lam
            return x.view_as(x)

        @staticmethod
        def backward(ctx, g):
            return -ctx.lam * g, None

    class ResBlock(nn.Module):
        def __init__(self, h, p=0.2):
            super().__init__()
            self.f = nn.Sequential(nn.BatchNorm1d(h), nn.Linear(h, h), nn.ReLU(), nn.Dropout(p), nn.Linear(h, h))

        def forward(self, x):
            return x + self.f(x)

    class Net(nn.Module):
        def __init__(self):
            super().__init__()
            if kind == "res":
                self.f1 = nn.Sequential(nn.Linear(d_in, 256), ResBlock(256), ResBlock(256), ResBlock(256),
                                        nn.BatchNorm1d(256), nn.ReLU(), nn.Linear(256, 64), nn.ReLU())
                self.f2 = nn.Identity()
            else:
                self.f1 = nn.Sequential(nn.Linear(d_in, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2),
                                        nn.Linear(256, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2),
                                        nn.Linear(256, 64), nn.ReLU())
                self.f2 = nn.Sequential(nn.Linear(64, 64), nn.ReLU())
            self.head = nn.Linear(64, 4)
            self.dom = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 1)) if kind == "dann" else None

        def forward(self, x, lam=0.0):
            f = self.f1(x)
            z = self.f2(f)
            dl = self.dom(GradReverse.apply(f, lam)) if self.dom is not None else None
            return self.head(z), z, dl

    return Net()


def coral_loss(a, b):
    import torch
    d = a.shape[1]

    def cov(x):
        x = x - x.mean(0, keepdim=True)
        return x.T @ x / max(len(x) - 1, 1)
    return ((cov(a) - cov(b)) ** 2).sum() / (4 * d * d)


SMOKE_ROWS = None                        # --smoke: train on a random subset of this many rows


def train(kind, Xs, ys, Xt, epochs, seed, eval_at=(), eval_fn=None):
    """Train a backbone. Xt = unlabelled target rows (dann / coral). If eval_at is
    given, eval_fn(net) is called after those epochs; returns (net, {ep: result})."""
    torch = lab.torch_setup(seed)
    if SMOKE_ROWS and len(Xs) > SMOKE_ROWS:
        sub = np.random.default_rng(seed).choice(len(Xs), SMOKE_ROWS, replace=False)
        Xs, ys = Xs[sub], np.asarray(ys)[sub]
    import torch.nn as nn
    net = build_net(kind, Xs.shape[1])
    opt = torch.optim.AdamW(net.parameters(), lr=LR, weight_decay=WD)
    cw = torch.tensor(lab.balanced_class_weights(ys), dtype=torch.float32)
    ce, bce = nn.CrossEntropyLoss(weight=cw), nn.BCEWithLogitsLoss()
    S, Y = torch.from_numpy(np.ascontiguousarray(Xs, np.float32)), torch.from_numpy(np.asarray(ys) - 1).long()
    T = torch.from_numpy(np.ascontiguousarray(Xt, np.float32))
    g = torch.Generator().manual_seed(1000 + seed)
    steps = int(np.ceil(len(S) / BATCH))
    res = {}
    for ep in range(epochs):
        net.train()
        perm = torch.randperm(len(S), generator=g)
        for k in range(steps):
            i = perm[k * BATCH:(k + 1) * BATCH]
            if len(i) < 2:
                continue
            if kind == "dann":
                p = (ep * steps + k) / (epochs * steps)
                lam = 2 / (1 + np.exp(-10 * p)) - 1
                j = torch.randint(0, len(T), (len(i),), generator=g)
                ls, _, ds = net(S[i], lam)
                _, _, dt = net(T[j], lam)
                loss = ce(ls, Y[i]) + bce(ds.squeeze(1), torch.zeros(len(i))) + bce(dt.squeeze(1), torch.ones(len(j)))
            elif kind == "coral":
                j = torch.randint(0, len(T), (len(i),), generator=g)
                ls, zs, _ = net(S[i])
                _, zt, _ = net(T[j])
                loss = ce(ls, Y[i]) + CORAL_W * coral_loss(zs, zt)
            else:
                ls, _, _ = net(S[i])
                loss = ce(ls, Y[i])
            opt.zero_grad()
            loss.backward()
            opt.step()
        if ep + 1 in eval_at:
            res[ep + 1] = eval_fn(net)
    net.eval()
    return net, res


def embed(net, X):
    import torch
    net.eval()
    out = []
    with torch.no_grad():
        for k in range(0, len(X), 8192):
            out.append(net(torch.from_numpy(np.ascontiguousarray(X[k:k + 8192], np.float32)))[1].numpy())
    return np.vstack(out).astype(np.float32)


def last_layer(net):
    return (net.head.weight.detach().numpy().astype(np.float32).copy(),
            net.head.bias.detach().numpy().astype(np.float32).copy())


# --------------------------------------------------------------------------- #
# Inputs
# --------------------------------------------------------------------------- #
def _by_columns(f, chunk=48):
    """Apply a per-column lab transform in column blocks: identical output (every lab
    transform is column-wise), but bounds the float64 / rank temporaries - lab.rankgauss on
    the full 76k x 393 Madrid matrix alone peaks at ~1.5 GB extra (reviewer fix, RAM cap)."""
    def g(X):
        return np.hstack([f(X[:, k:k + chunk]) for k in range(0, X.shape[1], chunk)])
    return g


TRANSFORM = {"rz": _by_columns(lab.robustz), "rg": _by_columns(lab.rankgauss), "o60": lab.standardise}


def raw_inputs(d, inp):
    return (d["Xm60"], d["Xa60"]) if inp == "o60" else (d["Xm"], d["Xa"])


_INPUT_CACHE = {}


def city_inputs(d, inp):
    """Per-city transformed inputs, kept in memory for the run (~150 MB per 393-feat input)."""
    if inp not in _INPUT_CACHE:
        Xm, Xa = raw_inputs(d, inp)
        f = TRANSFORM[inp]
        _INPUT_CACHE[inp] = (f(Xm), f(Xa))
    return _INPUT_CACHE[inp]


# --------------------------------------------------------------------------- #
# Heads
# --------------------------------------------------------------------------- #
def head_ft(e, Zs, ys, lam):
    Wf, bf = lab.finetune_last_layer(e["W0"], e["b0"], Zs, ys, lam)
    return lab.softmax(e["Za"].astype(np.float64) @ Wf.T + bf)


def joint_init(e):
    """Standardise embeddings with the Madrid mean/sd and fit the Madrid-only
    weighted LR once per seed (warm start for every draw's joint head)."""
    from sklearn.linear_model import LogisticRegression
    if "_Zm_std" not in e:
        mu, sd = e["Zm"].mean(0), e["Zm"].std(0)
        sd = np.where(sd > 1e-8, sd, 1.0)
        e["_Zm_std"] = ((e["Zm"] - mu) / sd).astype(np.float64)
        e["_Za_std"] = ((e["Za"] - mu) / sd).astype(np.float64)
        w = lab.balanced_class_weights(e["ym"])[e["ym"] - 1]
        e["_m0"] = LogisticRegression(C=1.0, solver="newton-cholesky", max_iter=100).fit(
            e["_Zm_std"], e["ym"], sample_weight=w)
    return e


def head_joint(e, shot, ys):
    """Same objective as lab.joint_head(Zm, ym, Zs, ys, shot_weight=30, C=1)
    (identical sample weights), solved with sklearn's newton-cholesky solver
    warm-started from the Madrid-only fit: lab.joint_head's lbfgs stops at
    max_iter=400 unconverged on these embeddings and takes ~9 s per fit; this
    converges in 2-4 Newton steps (~2 s) and lands closer to the true optimum."""
    from sklearn.linear_model import LogisticRegression
    e = joint_init(e)
    Zm, Za = e["_Zm_std"], e["_Za_std"]
    X = np.vstack([Zm, Za[shot]])
    y = np.r_[np.asarray(e["ym"]), np.asarray(ys)]
    w = np.r_[np.ones(len(e["ym"])), np.full(len(ys), SHOT_WEIGHT)]
    w = w * lab.balanced_class_weights(y)[y - 1]
    m = LogisticRegression(C=1.0, solver="newton-cholesky", max_iter=100, warm_start=True)
    m.coef_, m.intercept_ = e["_m0"].coef_.copy(), e["_m0"].intercept_.copy()
    m.fit(X, y, sample_weight=w)
    return m.predict_proba(Za)


def _unit(Z):
    Z = Z - Z.mean(0)
    return Z / np.clip(np.linalg.norm(Z, axis=1, keepdims=True), 1e-12, None)


def head_proto(e, shot, ys):
    if "_Um" not in e:                       # computed once per seed (setdefault would recompute every call)
        e["_Um"], e["_Ua"] = _unit(e["Zm"].astype(np.float64)), _unit(e["Za"].astype(np.float64))
    Um, Ua = e["_Um"], e["_Ua"]
    protos = []
    for c in lab.CLASSES:
        mm = Um[e["ym"] == c].mean(0)
        sc = Ua[shot[ys == c]]
        nc = len(sc)
        protos.append((nc * sc.mean(0) + PROTO_K * mm) / (nc + PROTO_K) if nc else mm)
    Pr = np.stack(protos)
    d2 = (Ua ** 2).sum(1)[:, None] - 2 * Ua @ Pr.T + (Pr ** 2).sum(1)[None]
    return lab.softmax(-PROTO_T * d2)


# --------------------------------------------------------------------------- #
# Madrid-only selection
# --------------------------------------------------------------------------- #
def selection(d, cache, smoke):
    path = cache / "selection.json"
    sel = json.loads(path.read_text()) if path.exists() else {"epochs": {}, "epoch_scores": {}}
    grid = (1, 2) if smoke else EPOCH_GRID
    src, tgt = lab.madrid_split(d["pm"])
    Xs, Xt = TRANSFORM["rz"](d["Xm"][src]), TRANSFORM["rz"](d["Xm"][tgt])
    ys, yt = d["ym"][src], d["ym"][tgt]
    prior = np.array([(ys == c).mean() for c in lab.CLASSES])

    def f1(net):
        P = lab.softmax(embed(net, Xt).astype(np.float64) @ last_layer(net)[0].T.astype(np.float64)
                        + last_layer(net)[1])
        return W.macro_f1(yt, W.decide(P, prior))

    emb = {}
    for kind in BACKBONES:
        ep_path = cache / f"sel_{kind}.npz"
        if kind in sel["epochs"] and ep_path.exists():
            z = np.load(ep_path)
            emb[kind] = dict(Zs=z["Zs"], Zt=z["Zt"], W0=z["W0"], b0=z["b0"])
            continue
        t0 = time.time()
        nets = {}
        if kind == "dann":
            scores = {}
            for E in grid:
                net, _ = train(kind, Xs, ys, Xt, E, 0)
                scores[E], nets[E] = f1(net), (embed(net, Xs), embed(net, Xt), *last_layer(net))
        else:
            snap = {}

            def ev(net):
                r = f1(net)
                snap[len(snap)] = (embed(net, Xs), embed(net, Xt), *last_layer(net))
                net.train()
                return r
            _, scores = train(kind, Xs, ys, Xt, max(grid), 0, eval_at=grid, eval_fn=ev)
            for k, E in enumerate(grid):
                nets[E] = snap[k]
        best = max(grid, key=lambda E: (scores[E], -E))
        sel["epochs"][kind] = best
        sel["epoch_scores"][kind] = {str(k): v for k, v in scores.items()}
        Zs_, Zt_, W0, b0 = nets[best]
        np.savez(ep_path, Zs=Zs_, Zt=Zt_, W0=W0, b0=b0)
        emb[kind] = dict(Zs=Zs_, Zt=Zt_, W0=W0, b0=b0)
        path.write_text(json.dumps(sel, indent=1))
        log(f"select epochs {kind}: pseudo-target rule F1 {scores} -> {best}  ({time.time() - t0:.0f}s)")

    if "lam" not in sel:
        t0 = time.time()
        reps = 1 if smoke else LAM_REPS
        tab = {str(l): [] for l in LAM_GRID}
        per = {}
        for kind in BACKBONES:
            e = emb[kind]
            for n in LAM_NS:
                for rep in range(reps):
                    rng = np.random.default_rng(7000 + 100 * n + rep)
                    shot = np.concatenate([rng.choice(np.flatnonzero(yt == c), n, replace=False) for c in lab.CLASSES])
                    keep = np.setdiff1d(np.arange(len(yt)), shot)
                    for l in LAM_GRID:
                        Wf, bf = lab.finetune_last_layer(e["W0"], e["b0"], e["Zt"][shot], yt[shot], l)
                        P = lab.softmax(e["Zt"][keep].astype(np.float64) @ Wf.T + bf)
                        s = W.macro_f1(yt[keep], W.decide(P, prior))
                        tab[str(l)].append(s)
                        per.setdefault(f"{kind}_n{n}_lam{l}", []).append(s)
        means = {l: float(np.mean(v)) for l, v in tab.items()}
        sel["lam"] = float(max(LAM_GRID, key=lambda l: (means[str(l)], l)))
        sel["lam_scores"] = means
        sel["lam_scores_detail"] = {k: float(np.mean(v)) for k, v in per.items()}
        path.write_text(json.dumps(sel, indent=1))
        log(f"select lam: {means} -> {sel['lam']}  ({time.time() - t0:.0f}s)")
    return sel


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--smoke", action="store_true")
    ap.add_argument("--only", default="")
    ap.add_argument("--seeds", type=int, default=len(SEEDS),
                    help="use only the first k of seeds 0,1,2 (default 3; method names unchanged, "
                         "recorded in meta) - a time-saving fallback, not the planned run")
    a = ap.parse_args()
    smoke = a.smoke
    global SMOKE_ROWS
    SMOKE_ROWS = 3000 if smoke else None
    sfx = "_SMOKE" if smoke else ""
    cache = lab.CACHE / ("tab_nn" + sfx)
    cache.mkdir(parents=True, exist_ok=True)
    seeds = (0,) if smoke else SEEDS[:a.seeds]
    draws = [(5, 0), (200, 0)] if smoke else list(lab.draws())

    d = lab.load(official60=True)
    log("data loaded")

    def mname(bb, inp, h):
        return f"tab_{bb}_{inp}_{h}{sfx}"

    combos = [(bb, inp) for bb in BACKBONES for inp in INPUTS_FOR[bb]
              if any(a.only in mname(bb, inp, h) for h in HEADS)]
    sel = selection(d, cache, smoke)
    lam = sel["lam"]

    for bb, inp in combos:
        heads = [h for h in HEADS if a.only in mname(bb, inp, h)]
        todo_draws = [(n, r) for n, r in draws
                      if any(not lab.have_pred(mname(bb, inp, h), n, r) for h in heads if h != "zs")]
        need_zs = "zs" in heads and not lab.have_pred(mname(bb, inp, "zs"), 0, 0)
        if not todo_draws and not need_zs:
            log(f"{bb}/{inp}: all saved, skip")
            continue
        epochs = int(sel["epochs"][bb])
        embs, Xin = [], None
        for s in seeds:
            f = cache / f"emb_{bb}_{inp}_s{s}.npz"
            if f.exists():                   # reviewer: never reuse a net trained for another epoch count
                with np.load(f) as z0:
                    stale = int(z0["epochs"]) != epochs
                if stale:
                    log(f"cached {f.name} has a different epoch count than selection.json: retraining")
                    f.unlink()
            if not f.exists():
                if Xin is None:
                    Xin = city_inputs(d, inp)
                t0 = time.time()
                net, _ = train(bb, Xin[0], d["ym"], Xin[1], epochs, s)
                W0, b0 = last_layer(net)
                np.savez(f, Zm=embed(net, Xin[0]), Za=embed(net, Xin[1]), W0=W0, b0=b0, epochs=epochs)
                log(f"train {bb}/{inp} seed {s}: {epochs} epochs  {time.time() - t0:.0f}s")
            z = np.load(f)
            embs.append(dict(Zm=z["Zm"], Za=z["Za"], W0=z["W0"], b0=z["b0"], ym=d["ym"]))
        del Xin
        meta = dict(script="tools/nnlab/tab_nn.py", backbone=bb, input=inp, seeds=list(seeds),
                    epochs=epochs, epoch_choice=sel["epoch_scores"].get(bb), ft_lam=lam,
                    lam_scores=sel.get("lam_scores"), joint_shot_weight=SHOT_WEIGHT,
                    proto_k=PROTO_K, proto_temp=PROTO_T, coral_weight=CORAL_W if bb == "coral" else None,
                    optim=dict(opt="AdamW", lr=LR, wd=WD, batch=BATCH),
                    joint_solver="newton-cholesky, warm start from Madrid-only fit, C=1",
                    selection="lab.madrid_split pseudo-target, rule macro F1, input rz seed 0",
                    smoke=smoke)
        if need_zs:
            P = np.mean([lab.softmax(e["Za"].astype(np.float64) @ e["W0"].T + e["b0"]) for e in embs], 0)
            lab.save_pred(mname(bb, inp, "zs"), 0, 0, P, {**meta, "head": "zs"})
            log(f"{mname(bb, inp, 'zs')} n=0 saved")
        for n, rep in todo_draws:
            t0 = time.time()
            shot = lab.shots_for(d, n, rep)
            ys = lab.shot_labels(d, shot)
            for h in heads:
                if h == "zs" or lab.have_pred(mname(bb, inp, h), n, rep):
                    continue
                th = time.time()
                if h == "ft":
                    P = np.mean([head_ft(e, e["Za"][shot], ys, lam) for e in embs], 0)
                elif h == "joint":
                    P = np.mean([head_joint(e, shot, ys) for e in embs], 0)
                else:
                    P = np.mean([head_proto(e, shot, ys) for e in embs], 0)
                lab.save_pred(mname(bb, inp, h), n, rep, P, {**meta, "head": h})
                log(f"{mname(bb, inp, h)} n={n} rep={rep} {time.time() - th:.1f}s")
        del embs
    log("done")


if __name__ == "__main__":
    main()