"""LP-FT with the number of fine-tuning steps chosen on Madrid.

tools/nnlab/lpft.py fixed STEPS = 60 in advance. Its Madrid check showed the gain of
whole-network fine-tuning growing with n, so this driver asks, on the Madrid
pseudo-target only (lab.madrid_split, the same cached source-part base network),
which of STEPS in {60, 150, 300, 600} is best, and then runs that one setting on
the Amsterdam protocol as  tab_dann_rz_lpft_s{STEPS}.
Everything else (learning rates, replay, L2-SP, lam, epochs) is unchanged.
Amsterdam labels: only lab.shot_labels(lab.shots_for(...)) of the current draw.
"""
import json, sys, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
import lab
import numpy as np
import lpft as L
import tab_nn
W = lab.W
GRID = (60, 150, 300, 600)
cache = lab.CACHE / "lpft"
sel = json.loads((lab.CACHE / "tab_nn" / "selection.json").read_text())
epochs, lam = int(sel["epochs"][L.KIND]), float(sel["lam"])
d = lab.load(official60=False)
t00 = time.time()
plain = {"lpft": L.METHODS["lpft"]}
sel_path = cache / "steps_selection.json"
if sel_path.exists():
    S = json.loads(sel_path.read_text())
else:
    src, tgt = lab.madrid_split(d["pm"])
    Xs, Xt = tab_nn.TRANSFORM[L.INP](d["Xm"][src]), tab_nn.TRANSFORM[L.INP](d["Xm"][tgt])
    ys, yt = d["ym"][src], d["ym"][tgt]
    prior = np.array([(ys == c).mean() for c in lab.CLASSES])
    state, _ = L.get_base(cache / f"{L.KIND}_{L.INP}_src_s{L.SEED}.pt", Xs, ys, Xt, epochs)
    rows = []
    for steps in GRID:
        L.STEPS = steps
        for n in L.CHECK_NS:
            for rep in range(L.CHECK_REPS):
                rng = np.random.default_rng(9000 + 100 * n + rep)       # same draws as lpft.madrid_check
                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 = L.lpft(state, Xt[shot], yt[shot], Xt[keep], Xs, ys, Xt, lam, 10_000 + 1000 * n + rep, plain)
                f = W.macro_f1(yt[keep], W.decide(P["lpft"], prior))
                rows.append(dict(steps=steps, n=n, rep=rep, f1=f))
                print(f"[{time.time()-t00:6.0f}s] madrid steps={steps} n={n} rep={rep} f1={f:.4f}", flush=True)
    import pandas as pd
    df = pd.DataFrame(rows)
    by = df.groupby(["steps", "n"]).f1.mean().unstack()
    mean = df.groupby("steps").f1.mean()
    best = int(mean.idxmax())
    S = dict(grid=list(GRID), rows=rows, mean_by_steps={int(k): float(v) for k, v in mean.items()},
             by_n={int(s): {int(n): float(v) for n, v in by.loc[s].items()} for s in by.index}, chosen=best,
             rule="highest mean macro F1 over n in (5, 50, 200) x 3 draws on the Madrid pseudo-target")
    sel_path.write_text(json.dumps(S, indent=1))
print("madrid selection:", S["mean_by_steps"], "by n:", S["by_n"], "-> chosen", S["chosen"], flush=True)
steps = int(S["chosen"])
if steps == 60:
    print("60 steps stays best on Madrid: nothing new to run")
    sys.exit(0)
L.STEPS = steps
name = f"tab_{L.KIND}_{L.INP}_lpft_s{steps}"
Xm, Xa = tab_nn.city_inputs(d, L.INP)
state, _ = L.get_base(cache / f"{L.KIND}_{L.INP}_s{L.SEED}.pt", Xm, d["ym"], Xa, epochs)
for n, rep in lab.draws():
    if lab.have_pred(name, n, rep):
        continue
    t0 = time.time()
    shot = lab.shots_for(d, n, rep)
    P = L.lpft(state, Xa[shot], lab.shot_labels(d, shot), Xa, Xm, d["ym"], Xa, lam, 1000 * n + rep, plain)
    lab.save_pred(name, n, rep, P["lpft"], dict(script="tools/nnlab/lpft_steps.py", steps=steps, steps_selection=S,
                                                lam=lam, epochs=epochs, base="lpft.py base, seed 0"))
    print(f"[{time.time()-t00:6.0f}s] {name} n={n} rep={rep} {time.time()-t0:.1f}s", flush=True)
print("done")