"""
lpft - whole-network adaptation of the DANN MLP with the shots: LP-FT (Kumar et al.
2022) with an L2-SP anchor (Li et al. 2018), optionally plus SHOT-style information
maximisation on unlabelled Amsterdam rows (Liang et al. 2020).

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

Base network: tab_nn.py's "dann" backbone on input "rz" (lab.robustz per city), seed 0,
trained exactly as tab_nn.train does (epochs from CACHE/tab_nn/selection.json), once,
with NNLAB_BASE_THREADS (default 3) torch threads - the count tab_nn's cache was made with;
CPU results depend on it (2 threads gave Za max abs diff 14) - then the run's NNLAB_THREADS;
state_dict cached at CACHE/lpft/dann_rz_s0.pt. Its Amsterdam embedding is compared with
CACHE/tab_nn/emb_dann_rz_s0.npz Za (max abs diff logged and stored in meta).

Per protocol draw (fresh copy of the base each time):
  LP  last Linear <- lab.finetune_last_layer(W0, b0, Z_shot, y_shot, lam), lam from
      selection.json (= tab_dann_rz_ft for seed 0 only; not saved).
  FT  everything unfrozen (domain head unused); BatchNorm in eval mode (running stats
      frozen), dropout off (net.eval() throughout); AdamW lr 1e-4 backbone, 1e-3 last
      layer, weight decay 0; STEPS=60 full-batch steps; loss =
        class-balanced CE(shots)
        + 0.1 * CE(fresh class-balanced Madrid replay batch, 128 per class = 512)
        + 0.01 * sum_params ||theta - theta_LP||^2
      [+ 0.1 * (mean entropy - entropy of mean prediction) on 512 random Amsterdam
       rows, variant "_ent"]. No early stopping.
Methods: tab_dann_rz_lpft, tab_dann_rz_lpft_ent  (suffix _SMOKE under --smoke).

Madrid check (evidence only, run once, CACHE/lpft/madrid_check.json, copied into meta):
lab.madrid_split; base trained on the source part with the unlabelled pseudo-target
as adversary (transforms per part, as tab_nn's selection); for n in (5, 50, 200) and
3 draws (rng(9000 + 100 n + rep)) from the pseudo-target: LP vs LP-FT vs LP-FT-ent,
macro F1 of W.decide (source-part prior) on the remaining pseudo-target rows.
Replay batches come from the source part; the entropy batch from the pseudo-target.
"""
from __future__ import annotations

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

sys.path.insert(0, str(Path(__file__).resolve().parent))
import lab                               # noqa: E402  (sets thread env before torch)
import numpy as np                       # noqa: E402
import tab_nn                            # noqa: E402

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

KIND, INP, SEED = "dann", "rz", 0
STEPS = 60
LR_BACKBONE, LR_HEAD = 1e-4, 1e-3
REPLAY_W, REPLAY_PER_CLASS = 0.1, 128
L2SP_W = 0.01
ENT_W, ENT_BATCH = 0.1, 512
CHECK_NS, CHECK_REPS = (5, 50, 200), 3
BASE_THREADS = int(os.environ.get("NNLAB_BASE_THREADS", "3"))  # thread count tab_nn's cache was made with
METHODS = {"lpft": False, "lpft_ent": True}     # suffix -> use entropy term


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


# --------------------------------------------------------------------------- #
# Base network
# --------------------------------------------------------------------------- #
def make_net(state, d_in):
    net = tab_nn.build_net(KIND, d_in)
    net.load_state_dict(state)
    net.eval()
    return net


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


def get_base(path, Xs, ys, Xt, epochs):
    import torch
    if path.exists():
        return torch.load(path), None
    t0 = time.time()
    # tab_nn trains on lab.torch_setup's thread count; its cache was made with 3 threads and
    # CPU torch results depend on the thread count (2 threads: Za max abs diff 14, argmax
    # agreement 0.87), so the base is trained with BASE_THREADS to reproduce it exactly.
    run_threads, lab.THREADS = lab.THREADS, BASE_THREADS
    try:
        net, _ = tab_nn.train(KIND, Xs, ys, Xt, epochs, SEED)
    finally:
        lab.THREADS = run_threads
        torch.set_num_threads(run_threads)
    state = {k: v.detach().clone() for k, v in net.state_dict().items()}
    torch.save(state, path)
    log(f"trained base {path.name} ({BASE_THREADS} threads): {epochs} epochs on {min(len(Xs), tab_nn.SMOKE_ROWS or len(Xs))} rows  {time.time() - t0:.0f}s")
    return state, net


# --------------------------------------------------------------------------- #
# LP-FT for one draw
# --------------------------------------------------------------------------- #
def lpft(state, X_shot, ys, X_pred, X_rep, y_rep, X_unl, lam, seed, variants):
    """Returns {"lp": P, suffix: P ...} probabilities on X_pred (columns classes 1..4)."""
    import torch
    import torch.nn.functional as F
    lab.torch_setup(seed)
    base = make_net(state, X_shot.shape[1])
    Z_shot = tab_nn.embed(base, X_shot)
    W0, b0 = tab_nn.last_layer(base)
    Wf, bf = lab.finetune_last_layer(W0, b0, Z_shot, ys, lam)
    with torch.no_grad():
        base.head.weight.copy_(torch.tensor(Wf, dtype=torch.float32))
        base.head.bias.copy_(torch.tensor(bf, dtype=torch.float32))
    out = {"lp": lab.softmax(logits_of(base, X_pred))}

    Xs = torch.from_numpy(np.ascontiguousarray(X_shot, np.float32))
    Ys = torch.from_numpy(np.asarray(ys) - 1).long()
    cw = torch.tensor(lab.balanced_class_weights(np.asarray(ys)), dtype=torch.float32)
    Xr = torch.from_numpy(np.ascontiguousarray(X_rep, np.float32))
    by_cls = [torch.from_numpy(np.flatnonzero(np.asarray(y_rep) == c)) for c in lab.CLASSES]
    Xu = torch.from_numpy(np.ascontiguousarray(X_unl, np.float32))

    for sfx, use_ent in variants.items():
        net = copy.deepcopy(base)
        net.eval()                                   # BN running stats frozen, dropout off
        if net.dom is not None:
            for p in net.dom.parameters():
                p.requires_grad_(False)              # domain head is not in the classifier path
        named = [(k, p) for k, p in net.named_parameters() if p.requires_grad]
        anchor = {k: p.detach().clone() for k, p in named}
        head_p = [p for k, p in named if k.startswith("head.")]
        body_p = [p for k, p in named if not k.startswith("head.")]
        opt = torch.optim.AdamW([{"params": body_p, "lr": LR_BACKBONE},
                                 {"params": head_p, "lr": LR_HEAD}], weight_decay=0.0)
        g = torch.Generator().manual_seed(seed)
        # reviewer: separate stream for the entropy batch so both variants get the SAME
        # replay batches (the only difference between them is the entropy term)
        g_ent = torch.Generator().manual_seed(seed + 7_777_777)
        for _ in range(STEPS):
            loss = F.cross_entropy(net(Xs)[0], Ys, weight=cw)
            ri = torch.cat([ix[torch.randint(0, len(ix), (REPLAY_PER_CLASS,), generator=g)] for ix in by_cls])
            yr = torch.repeat_interleave(torch.arange(4), REPLAY_PER_CLASS)
            loss = loss + REPLAY_W * F.cross_entropy(net(Xr[ri])[0], yr)
            loss = loss + L2SP_W * sum(((p - anchor[k]) ** 2).sum() for k, p in named)
            if use_ent:
                j = torch.randint(0, len(Xu), (ENT_BATCH,), generator=g_ent)
                pu = F.softmax(net(Xu[j])[0], 1)
                ment = -(pu * torch.log(pu + 1e-8)).sum(1).mean()
                pm = pu.mean(0)
                hm = -(pm * torch.log(pm + 1e-8)).sum()
                loss = loss + ENT_W * (ment - hm)
            opt.zero_grad()
            loss.backward()
            opt.step()
        net.eval()
        out[sfx] = lab.softmax(logits_of(net, X_pred))
    return out


# --------------------------------------------------------------------------- #
# Madrid check
# --------------------------------------------------------------------------- #
def madrid_check(d, cache, epochs, lam, smoke):
    path = cache / "madrid_check.json"
    if path.exists():
        return json.loads(path.read_text())
    src, tgt = lab.madrid_split(d["pm"])
    Xs, Xt = tab_nn.TRANSFORM[INP](d["Xm"][src]), tab_nn.TRANSFORM[INP](d["Xm"][tgt])
    ys, yt = d["ym"][src], d["ym"][tgt]
    prior = np.array([(ys == c).mean() for c in lab.CLASSES])
    state, net = get_base(cache / f"{KIND}_{INP}_src_s{SEED}.pt", Xs, ys, Xt, epochs)
    res = {"ns": list(CHECK_NS), "reps": CHECK_REPS, "shot_rng": "default_rng(9000 + 100 n + rep)",
           "epochs": epochs, "lam": lam, "draws": []}
    sel_f = lab.CACHE / "tab_nn" / f"sel_{KIND}.npz"
    if net is not None and sel_f.exists() and not smoke:
        with np.load(sel_f) as z:
            diff = float(np.abs(tab_nn.embed(net, Xt) - z["Zt"]).max())
        res["src_base_vs_tab_nn_sel_Zt_maxabs"] = diff
        log(f"madrid-check base vs tab_nn sel_{KIND}.npz Zt: max abs diff {diff:.2e}")
    ns = (5, 200) if smoke else CHECK_NS
    reps = 1 if smoke else CHECK_REPS
    for n in ns:
        for rep in range(reps):
            t0 = time.time()
            rng = np.random.default_rng(9000 + 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)
            P = lpft(state, Xt[shot], yt[shot], Xt[keep], Xs, ys, Xt, lam, 10_000 + 1000 * n + rep, METHODS)
            row = {"n": n, "rep": rep}
            for k, Pk in P.items():
                row[k] = W.macro_f1(yt[keep], W.decide(Pk, prior))
            res["draws"].append(row)
            log(f"madrid-check n={n} rep={rep}: " + "  ".join(f"{k}={v:.4f}" for k, v in row.items()
                                                              if k not in ("n", "rep")) + f"  {time.time() - t0:.1f}s")
    import pandas as pd
    df = pd.DataFrame(res["draws"])
    res["mean"] = {k: {int(n): float(v) for n, v in df.groupby("n")[k].mean().items()}
                   for k in ["lp", *METHODS]}
    path.write_text(json.dumps(res, indent=1))
    log(f"madrid-check means: {res['mean']}")
    return res


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--smoke", action="store_true")
    ap.add_argument("--only", default="")
    ap.add_argument("--base-only", action="store_true",
                    help="train/cache the Amsterdam base net, compare with tab_nn's embedding, stop")
    a = ap.parse_args()
    smoke = a.smoke
    sfx = "_SMOKE" if smoke else ""
    if smoke:
        tab_nn.SMOKE_ROWS = 3000
        global BASE_THREADS                  # reviewer: smoke base is never compared, keep NNLAB_THREADS
        BASE_THREADS = lab.THREADS
    cache = lab.CACHE / ("lpft" + sfx)
    cache.mkdir(parents=True, exist_ok=True)
    sel = json.loads((lab.CACHE / "tab_nn" / "selection.json").read_text())
    epochs = 2 if smoke else int(sel["epochs"][KIND])
    lam = float(sel["lam"])
    draws = [(5, 0), (200, 0)] if smoke else list(lab.draws())
    names = {s: f"tab_{KIND}_{INP}_{s}{sfx}" for s in METHODS}
    variants = {s: e for s, e in METHODS.items() if a.only in names[s]}

    d = lab.load(official60=False)
    Xm, Xa = tab_nn.city_inputs(d, INP)
    log("data loaded")

    # base on all Madrid, adversary = unlabelled Amsterdam
    base_path = cache / f"{KIND}_{INP}_s{SEED}.pt"
    state, net = get_base(base_path, Xm, d["ym"], Xa, epochs)
    chk_path = cache / "base_check.json"
    if net is not None or not chk_path.exists():
        if net is None:
            net = make_net(state, Xm.shape[1])
        emb_f = lab.CACHE / "tab_nn" / f"emb_{KIND}_{INP}_s{SEED}.npz"
        chk = {"compared_with": str(emb_f.relative_to(lab.ROOT)), "base_threads": BASE_THREADS}
        if smoke:
            chk["note"] = "smoke base (3000 rows, 2 epochs): no comparison"
        else:
            with np.load(emb_f) as z:
                Za = tab_nn.embed(net, Xa)
                chk["Za_maxabs"] = float(np.abs(Za - z["Za"]).max())
                chk["W0_maxabs"] = float(np.abs(tab_nn.last_layer(net)[0] - z["W0"]).max())
                chk["match_1e-3"] = bool(chk["Za_maxabs"] <= 1e-3)
                assert int(z["epochs"]) == epochs
            msg = "MATCH" if chk["match_1e-3"] else "MISMATCH (continuing; threading?)"
            log(f"base vs tab_nn {emb_f.name}: Za max abs diff {chk['Za_maxabs']:.2e}, "
                f"W0 {chk['W0_maxabs']:.2e} -> {msg}")
        chk_path.write_text(json.dumps(chk, indent=1))
    base_chk = json.loads(chk_path.read_text())
    del net
    if a.base_only:
        log("base only: done")
        return
    if not variants:
        log("nothing selected by --only")
        return

    check = madrid_check(d, cache, epochs, lam, smoke)

    meta = dict(script="tools/nnlab/lpft.py", backbone=KIND, input=INP, seed=SEED, epochs=epochs,
                lp_lam=lam, steps=STEPS, lr_backbone=LR_BACKBONE, lr_head=LR_HEAD, weight_decay=0.0,
                replay_weight=REPLAY_W, replay_batch=4 * REPLAY_PER_CLASS, l2sp_weight=L2SP_W,
                ent_weight=ENT_W, ent_batch=ENT_BATCH, bn="eval (running stats frozen)", dropout="off",
                draw_seed="torch seed / generator 1000 n + rep",
                base_threads=BASE_THREADS, base_check=base_chk, madrid_check=check, smoke=smoke,
                hyperparameters="fixed in advance by the brief; epochs and lam from tab_nn Madrid selection")
    for n, rep in draws:
        todo = {s: e for s, e in variants.items() if not lab.have_pred(names[s], n, rep)}
        if not todo:
            continue
        t0 = time.time()
        shot = lab.shots_for(d, n, rep)
        ys = lab.shot_labels(d, shot)
        P = lpft(state, Xa[shot], ys, Xa, Xm, d["ym"], Xa, lam, 1000 * n + rep, todo)
        for s in todo:
            lab.save_pred(names[s], n, rep, P[s], {**meta, "method": names[s], "entropy_term": todo[s]})
        log(f"{'+'.join(names[s] for s in todo)} n={n} rep={rep} {time.time() - t0:.1f}s")
    log("done")


if __name__ == "__main__":
    main()