"""
metric.py - METRIC LEARNING family: an embedding trained on Madrid so that distance
means "same age class", then adapted per protocol draw with the n Amsterdam shots.

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

Encoders (MLP 393-256-256-64, BN/ReLU/Dropout 0.2, AdamW lr 1e-3 wd 1e-4, 2 seeds):
    supcon   supervised contrastive (T = 0.1), class-balanced batches of 512
    proto    prototypical episodes (4-way, k in {5, 25}, 50 queries/class, scale 10),
             cross-block: support from one random half of the 2 km blocks, queries from the other
    protoda  proto + deep-CORAL (weight 1) to an unlabelled Amsterdam batch every step
    cos      Baseline++ cosine classifier (scale 16), class-balanced batches
Heads (encoder frozen, CL2N per city, seeds' probabilities averaged):
    pn, maha, tpn, knn, ftproj   ->  methods met_{encoder}_{head}

Training length per encoder is chosen ONCE on Madrid (lab.madrid_split pseudo-target,
25-shot prototype macro F1, candidates evaluated as checkpoints of one constant-LR run),
then the encoder is retrained on all Madrid. Amsterdam labels: only lab.shot_labels.
"""
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

T0 = time.time()


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


ap = argparse.ArgumentParser()
ap.add_argument("--smoke", action="store_true")
ap.add_argument("--only", default=None)
args = ap.parse_args()
SMOKE = args.smoke
SFX = "_SMOKE" if SMOKE else ""

ENCODERS = ["supcon", "proto", "protoda", "cos"]
HEADS = ["pn", "maha", "tpn", "knn", "ftproj"]

# configuration ---------------------------------------------------------------- #
# candidate training lengths (optimizer steps; one step = one batch / one episode)
# (sized to the measured ~0.09-0.2 s/step on the shared CPU so the full run fits ~60 min;
#  1500 steps x 512 = ~10 Madrid epochs)
CAND = {"supcon": [300, 700, 1500], "cos": [300, 700, 1500],
        "proto": [400, 800, 1600], "protoda": [400, 800, 1600]}
if SMOKE:
    CAND = {k: [15, 30] for k in CAND}
SEEDS = [0] if SMOKE else [0, 1]
FT_SEEDS = 1                         # ftproj uses only the first seed's encoder (time budget)
SEL_SEED = 100                       # seed of the Madrid-only selection run
SEL_DRAWS = 2 if SMOKE else 5        # pseudo-target 25-shot draws per checkpoint
DRAWS = [(5, 0), (200, 0)] if SMOKE else list(lab.draws())
BATCH = 512
TEMP = 0.1
PROTO_SCALE = 10.0
PROTO_K = (5, 25)
PROTO_Q = 50
CORAL_W = 1.0
DA_BATCH = 300                       # Amsterdam batch for CORAL ~ size of a Madrid episode
COS_SCALE = 16.0
LR, WD = 1e-3, 1e-4
# heads
K_MIX, S_PN = 5.0, 20.0
SHRINK, RIDGE = 50.0, 1e-3
TPN_IT, TPN_CAP = 10, 5.0
KNN_K = 15
FT_STEPS, FT_LR, FT_L2SP = 100, 1e-3, 1e-2

MC = lab.CACHE / "metric"
MC.mkdir(parents=True, exist_ok=True)


def mname(enc, head):
    return f"met_{enc}_{head}{SFX}"


WANT = {e: [h for h in HEADS if args.only is None or args.only in f"met_{e}_{h}"] for e in ENCODERS}
WANT = {e: hs for e, hs in WANT.items() if hs}

# data ------------------------------------------------------------------------- #
d = lab.load(official60=False)
ym, pm = d["ym"], d["pm"]
Xm = lab.robustz(d["Xm"])
Xa = lab.robustz(d["Xa"])
del d["Xa"]                         # raw Madrid kept for the selection split
nM, nA = len(Xm), len(Xa)
log(f"loaded Madrid {Xm.shape} Amsterdam {Xa.shape}; encoders {list(WANT)}")

torch = lab.torch_setup(0)
nn = torch.nn
F = torch.nn.functional


# --------------------------------------------------------------------------- #
# Encoder + losses
# --------------------------------------------------------------------------- #
class Enc(nn.Module):
    def __init__(self, din=393):
        super().__init__()
        self.body = nn.Sequential(
            nn.Linear(din, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(256, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2))
        self.proj = nn.Linear(256, 64)

    def forward(self, x):
        return self.proj(self.body(x))


def supcon_loss(z, y, T=TEMP):
    z = F.normalize(z, dim=1)
    sim = z @ z.T / T
    n = len(y)
    eye = torch.eye(n, dtype=torch.bool)
    sim = sim.masked_fill(eye, -1e9)
    logp = sim - torch.logsumexp(sim, 1, keepdim=True)
    pos = (y[:, None] == y[None, :]) & ~eye
    return -((logp * pos).sum(1) / pos.sum(1).clamp(min=1)).mean()


def proto_loss(zs, ys, zq, yq, scale=PROTO_SCALE, ncls=4):
    zs, zq = F.normalize(zs, dim=1), F.normalize(zq, dim=1)
    P = torch.stack([zs[ys == c].mean(0) for c in range(ncls)])
    d2 = ((zq[:, None, :] - P[None]) ** 2).sum(-1)
    return F.cross_entropy(-scale * d2, yq)


def coral(a, b):
    def cov(x):
        x = x - x.mean(0, keepdim=True)
        return x.T @ x / (len(x) - 1)
    dd = a.shape[1]
    return ((cov(a) - cov(b)) ** 2).sum() / (4 * dd * dd)


class Sampler:
    """Class-balanced batches and cross-block episodes over one labelled pool."""

    def __init__(self, y, p, rng):
        self.y, self.rng = y, rng
        self.cls = [np.flatnonzero(y == c) for c in lab.CLASSES]
        b = lab.block_ids(p)
        _, self.code = np.unique(b, return_inverse=True)
        self.nb = self.code.max() + 1
        self.ccode = [self.code[ix] for ix in self.cls]

    def batch(self, size=BATCH):
        m = size // 4
        return np.concatenate([ix[self.rng.integers(0, len(ix), m)] for ix in self.cls])

    def episode(self, k=None):
        """Cross-block episode. Sampling is with replacement inside pools of thousands of
        pixels (duplicates are rare; it avoids a full permutation per class per step)."""
        k = k or int(self.rng.choice(PROTO_K))
        half = self.rng.random(self.nb) < 0.5
        S, Q = [], []
        for ix, cc in zip(self.cls, self.ccode):
            m = half[cc]
            a, b = ix[m], ix[~m]
            if len(a) < k or len(b) < PROTO_Q:     # degenerate split: fall back to any pixel
                a = b = ix
            S.append(a[self.rng.integers(0, len(a), k)])
            Q.append(b[self.rng.integers(0, len(b), PROTO_Q)])
        ys = np.repeat(np.arange(4), k)
        yq = np.repeat(np.arange(4), PROTO_Q)
        return np.concatenate(S), ys, np.concatenate(Q), yq


@torch.no_grad()
def embed_H(model, X, bs=8192):
    model.eval()
    out = [model.body(torch.from_numpy(X[i:i + bs])).numpy() for i in range(0, len(X), bs)]
    return np.vstack(out).astype(np.float32)


def train_encoder(kind, X, y, p, Xt, lengths, seed, eval_fn=None):
    """Constant-LR training; returns (model, {length: eval score}) with checkpoints at
    `lengths` (the model returned is the one at max(lengths))."""
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    model = Enc(X.shape[1])
    params = list(model.parameters())
    Wc = None
    if kind == "cos":
        Wc = nn.Parameter(torch.randn(4, 64) * 0.1)
        params.append(Wc)
    opt = torch.optim.AdamW(params, lr=LR, weight_decay=WD)
    smp = Sampler(y, p, rng)
    Xt_t = torch.from_numpy(X)
    yt = torch.from_numpy(y - 1)
    scores = {}
    t0 = time.time()
    for step in range(1, max(lengths) + 1):
        model.train()
        if kind in ("supcon", "cos"):
            ix = smp.batch()
            z = model(Xt_t[ix])
            if kind == "supcon":
                loss = supcon_loss(z, yt[ix])
            else:
                logits = COS_SCALE * F.normalize(z, dim=1) @ F.normalize(Wc, dim=1).T
                loss = F.cross_entropy(logits, yt[ix])
        else:
            S, ys, Q, yq = smp.episode()
            z = model(Xt_t[np.concatenate([S, Q])])
            zs, zq = z[:len(S)], z[len(S):]
            loss = proto_loss(zs, torch.from_numpy(ys), zq, torch.from_numpy(yq))
            if kind == "protoda":
                it = rng.integers(0, len(Xt), DA_BATCH)
                zt = model(torch.from_numpy(Xt[it]))
                loss = loss + CORAL_W * coral(z, zt)
        opt.zero_grad()
        loss.backward()
        opt.step()
        if step in lengths:
            if eval_fn is not None:
                scores[step] = eval_fn(model)
                log(f"  {kind} seed{seed} step {step}: loss {loss.item():.4f} "
                    f"pseudo-target F1 {scores[step]:.4f} ({time.time() - t0:.0f}s)")
            else:
                log(f"  {kind} seed{seed} step {step}: loss {loss.item():.4f} ({time.time() - t0:.0f}s)")
    return model, scores


# --------------------------------------------------------------------------- #
# Embedding utilities
# --------------------------------------------------------------------------- #
def l2n(Z):
    return Z / np.maximum(np.linalg.norm(Z, axis=1, keepdims=True), 1e-12)


def cl2n(Z):
    """SimpleShot CL2N with this city's own mean (label-free)."""
    E = l2n(np.asarray(Z, np.float64))
    return l2n(E - E.mean(0))


def macro_f1(y, yhat):
    from sklearn.metrics import f1_score
    return float(f1_score(y, yhat, average="macro"))


def proto_f1_eval(Xpt, ypt, ndraw=SEL_DRAWS, shots=25):
    """Pseudo-target score: 25-shot pure prototype classifier on CL2N embeddings."""
    rs = np.random.default_rng(7)
    draws_ = []
    for _ in range(ndraw):
        sh = np.concatenate([rs.choice(np.flatnonzero(ypt == c), shots, replace=False) for c in lab.CLASSES])
        draws_.append(sh)

    def fn(model):
        H = embed_H(model, Xpt)
        with torch.no_grad():
            Z = model.proj(torch.from_numpy(H)).numpy()
        E = cl2n(Z)
        out = []
        for sh in draws_:
            P = np.stack([E[sh][ypt[sh] == c].mean(0) for c in lab.CLASSES])
            keep = np.setdiff1d(np.arange(len(E)), sh)
            pred = lab.CLASSES[(-((E[keep, None, :] - P[None]) ** 2).sum(-1)).argmax(1)]
            out.append(macro_f1(ypt[keep], pred))
        return float(np.mean(out))
    return fn


def select_length(kind):
    p = MC / f"{kind}_select{SFX}.json"
    if p.exists():
        return json.loads(p.read_text())
    src, pt = lab.madrid_split(pm)
    Xs, Xpt = lab.robustz(d["Xm"][src]), lab.robustz(d["Xm"][pt])
    log(f"{kind}: selecting length on Madrid split (src {len(src)}, pseudo-target {len(pt)}) "
        f"from {CAND[kind]}")
    # protoda: the unlabelled 'target' of the selection run is the pseudo-target itself
    _, sc = train_encoder(kind, Xs, ym[src], pm[src], Xpt, CAND[kind], SEL_SEED,
                          eval_fn=proto_f1_eval(Xpt, ym[pt]))
    best = max(sc, key=lambda s: (sc[s], -s))
    res = {"candidates": CAND[kind], "pseudo_target_f1": {str(k): v for k, v in sc.items()},
           "chosen": int(best), "seed": SEL_SEED, "split": "lab.madrid_split(pm) default",
           "da_target_in_selection": "pseudo-target (unlabelled)" if kind == "protoda" else None}
    p.write_text(json.dumps(res, indent=1))
    log(f"{kind}: chose {best} steps")
    return res


def encoder_cache(kind, seed, length):
    p = MC / f"{kind}_s{seed}_L{length}{SFX}.npz"
    if p.exists():
        z = np.load(p)
        return {k: z[k] for k in z.files}
    log(f"{kind}: training seed {seed} on all Madrid for {length} steps")
    model, _ = train_encoder(kind, Xm, ym, pm, Xa, [length], seed)
    out = dict(Hm=embed_H(model, Xm), Ha=embed_H(model, Xa),
               W=model.proj.weight.detach().numpy().copy(), b=model.proj.bias.detach().numpy().copy())
    np.savez(p, **out)
    return out


# --------------------------------------------------------------------------- #
# Heads
# --------------------------------------------------------------------------- #
def class_means(E, y):
    return np.stack([E[y == c].mean(0) for c in lab.CLASSES])


def within_cov(E, y):
    R = np.vstack([E[y == c] - E[y == c].mean(0) for c in lab.CLASSES])
    return R.T @ R / max(len(R) - 4, 1)


def pn_protos(Ea, sh, ysh, mu_m):
    P = []
    for j, c in enumerate(lab.CLASSES):
        S = Ea[sh[ysh == c]]
        P.append((len(S) * S.mean(0) + K_MIX * mu_m[j]) / (len(S) + K_MIX))
    return np.stack(P)


def sqd(E, P):
    return (E ** 2).sum(1)[:, None] - 2 * E @ P.T + (P ** 2).sum(1)[None]


def head_pn(Ea, sh, ysh, mu_m):
    return lab.softmax(-S_PN * sqd(Ea, pn_protos(Ea, sh, ysh, mu_m)))


def head_maha(Ea, sh, ysh, mu_m, S_m):
    from scipy.linalg import cho_factor, cho_solve
    P = pn_protos(Ea, sh, ysh, mu_m)
    nt = len(sh)
    S = (nt * within_cov(Ea[sh], ysh) + SHRINK * S_m) / (nt + SHRINK) + RIDGE * np.eye(Ea.shape[1])
    cf = cho_factor(S)
    D = np.empty((len(Ea), 4))
    for j in range(4):
        R = Ea - P[j]
        D[:, j] = (R * cho_solve(cf, R.T).T).sum(1)
    return lab.softmax(-0.5 * D)


def head_tpn(Ea, sh, ysh, mu_m):
    P = pn_protos(Ea, sh, ysh, mu_m)
    unl = np.setdiff1d(np.arange(len(Ea)), sh)
    Eu = Ea[unl]
    sums = np.stack([Ea[sh[ysh == c]].sum(0) for c in lab.CLASSES])
    nc = np.array([(ysh == c).sum() for c in lab.CLASSES], float)
    for _ in range(TPN_IT):
        R = lab.softmax(-S_PN * sqd(Eu, P))           # soft assignments of the unlabelled pool
        mass = R.sum(0)
        w = np.minimum(1.0, TPN_CAP * nc / np.maximum(mass, 1e-12))
        P = (sums + w[:, None] * (R.T @ Eu)) / (nc + w * mass)[:, None]
    return lab.softmax(-S_PN * sqd(Ea, P))


def head_knn(Ea, sh, ysh, mu_m, n):
    if n < 25:
        return head_pn(Ea, sh, ysh, mu_m)
    Es = Ea[sh]
    C = Ea @ Es.T                                       # cosine (unit vectors)
    k = min(KNN_K, len(sh))
    nb = np.argpartition(-C, k - 1, axis=1)[:, :k]
    cs = np.take_along_axis(C, nb, 1)
    w = 1.0 / (1.0 - cs + 1e-3)                         # inverse cosine distance
    lab_nb = ysh[nb] - 1
    P = np.zeros((len(Ea), 4))
    for j in range(4):
        P[:, j] = (w * (lab_nb == j)).sum(1)
    P += 1e-3
    return P / P.sum(1, keepdims=True)


def head_ftproj(enc, Hm, Ha, W0, b0, sh, ysh, n, seed):
    """Fine-tune ONLY the projection Linear(256,64) on cached frozen activations."""
    torch.manual_seed(seed + 1000 * n)
    rng = np.random.default_rng(seed + 1000 * n)
    Wt = torch.tensor(W0, requires_grad=True)
    bt = torch.tensor(b0, requires_grad=True)
    W0t, b0t = torch.tensor(W0), torch.tensor(b0)
    opt = torch.optim.Adam([Wt, bt], lr=FT_LR)
    Hs = torch.from_numpy(Ha[sh])
    Hm_t = torch.from_numpy(Hm)
    smp = MAD_SMP
    smp.rng = rng
    cls_ix = [np.flatnonzero(ysh == c) for c in lab.CLASSES]
    for _ in range(FT_STEPS):
        S, Q = [], []
        for ix in cls_ix:
            pr = rng.permutation(ix)
            ns = 3 if len(ix) <= 5 else len(ix) // 2
            S.append(pr[:ns]); Q.append(pr[ns:])
        ys = torch.from_numpy(np.concatenate([np.full(len(s), j) for j, s in enumerate(S)]))
        yq = torch.from_numpy(np.concatenate([np.full(len(q), j) for j, q in enumerate(Q)]))
        S, Q = np.concatenate(S), np.concatenate(Q)
        loss = proto_loss(Hs[S] @ Wt.T + bt, ys, Hs[Q] @ Wt.T + bt, yq)
        mS, mys, mQ, myq = smp.episode()
        loss = loss + proto_loss(Hm_t[mS] @ Wt.T + bt, torch.from_numpy(mys),
                                 Hm_t[mQ] @ Wt.T + bt, torch.from_numpy(myq))
        loss = loss + FT_L2SP * (((Wt - W0t) ** 2).sum() + ((bt - b0t) ** 2).sum())
        opt.zero_grad()
        loss.backward()
        opt.step()
    W, b = Wt.detach().numpy(), bt.detach().numpy()
    Em, Ea = cl2n(Hm @ W.T + b), cl2n(Ha @ W.T + b)
    return head_pn(Ea, sh, ysh, class_means(Em, ym))


# --------------------------------------------------------------------------- #
# Main loop
# --------------------------------------------------------------------------- #
def meta_for(enc, head, sel):
    return dict(family="metric", encoder=enc, head=head, smoke=SMOKE,
                seeds=SEEDS[:FT_SEEDS] if head == "ftproj" else SEEDS,
                input="lab.robustz per city (clip 10)",
                encoder_arch="393-256-BN-ReLU-Do0.2-256-BN-ReLU-Do0.2-64, L2-normalised",
                optim=f"AdamW lr {LR} wd {WD}, constant LR",
                encoder_cfg=dict(batch=BATCH, supcon_T=TEMP, proto_scale=PROTO_SCALE, proto_k=PROTO_K,
                                 proto_q=PROTO_Q, cross_block="2 km blocks random halves per episode",
                                 coral_w=CORAL_W, da_batch=DA_BATCH, cos_scale=COS_SCALE,
                                 cos_balance="class-balanced batches (128 per class), plain CE"),
                length_selection=sel,
                madrid_choices=f"training steps chosen on lab.madrid_split pseudo-target: {sel['chosen']}",
                head_cfg=dict(CL2N="per-city mean, re-normalised", k_mix=K_MIX, s=S_PN, shrink=SHRINK,
                              ridge=RIDGE, tpn_iters=TPN_IT, tpn_cap=TPN_CAP, knn_k=KNN_K,
                              knn_weight="1/(1-cos+1e-3)", knn_n5="pn fallback",
                              ft_steps=FT_STEPS, ft_lr=FT_LR, ft_l2sp=FT_L2SP,
                              ft_episode="shots half/half (n=5: 3/2) + one Madrid cross-block episode"),
                hyperparams_from_amsterdam="none")


MAD_SMP = None


def main():
    global MAD_SMP
    MAD_SMP = Sampler(ym, pm, np.random.default_rng(0))
    for enc, heads in WANT.items():
        todo = [(n, r) for n, r in DRAWS if not all(lab.have_pred(mname(enc, h), n, r) for h in heads)]
        if not todo:
            log(f"{enc}: all predictions exist, skipping")
            continue
        sel = select_length(enc)
        L = sel["chosen"]
        caches = [encoder_cache(enc, s, L) for s in SEEDS]
        prep = []
        for c in caches:
            Zm, Za = c["Hm"] @ c["W"].T + c["b"], c["Ha"] @ c["W"].T + c["b"]
            Em, Ea = cl2n(Zm), cl2n(Za)
            prep.append(dict(Ea=Ea, mu=class_means(Em, ym), Sm=within_cov(Em, ym)))
            del Em
        metas = {h: meta_for(enc, h, sel) for h in heads}
        for n, r in todo:
            t0 = time.time()
            sh = lab.shots_for(d, n, r)
            ysh = lab.shot_labels(d, sh)
            for h in heads:
                if lab.have_pred(mname(enc, h), n, r):
                    continue
                Ps = []
                for si, (c, pp) in enumerate(zip(caches, prep)):
                    if h == "ftproj" and si >= FT_SEEDS:
                        break
                    Ea, mu = pp["Ea"], pp["mu"]
                    if h == "pn":
                        P = head_pn(Ea, sh, ysh, mu)
                    elif h == "maha":
                        P = head_maha(Ea, sh, ysh, mu, pp["Sm"])
                    elif h == "tpn":
                        P = head_tpn(Ea, sh, ysh, mu)
                    elif h == "knn":
                        P = head_knn(Ea, sh, ysh, mu, n)
                    else:
                        P = head_ftproj(enc, c["Hm"], c["Ha"], c["W"], c["b"], sh, ysh, n, SEEDS[si] + 10 * r)
                    Ps.append(P)
                lab.save_pred(mname(enc, h), n, r, np.mean(Ps, 0), metas[h])
            log(f"{enc} n={n} rep={r}: saved {len(heads)} heads ({time.time() - t0:.1f}s)")
        del caches, prep
    log("done")


if __name__ == "__main__":
    main()