#!/usr/bin/env python
"""Build every model folder the labs use.

1. qwen35-p50c: the paper's released edit, rebuilt with the paper's own apply_edit.py (hash-checked).
2. results/edit_info.json: for each edited layer, the direction u, the strength alpha and the 64 columns,
   read back out of the weights (Lab 1 explains how).
3. A self-check: our formula with the recovered u, alpha and columns must give the released bits.
4. One folder per variant in configs/variants.json (qwen35-<name>).
"""
import argparse
import json
import os
import subprocess
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
import torch  # noqa: E402

from wsp import edits as E  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--base", default="/nvme/wsp-lab/models/Qwen3.5-0.8B")
ap.add_argument("--models", default="/nvme/wsp-lab/models")
ap.add_argument("--exp8", default="/nvme/wsp-lab/code/external/lm_understand-logs/exp8_math")
ap.add_argument("--variants", default=os.path.join(ROOT, "configs", "variants.json"))
ap.add_argument("--info", default=os.path.join(ROOT, "results", "edit_info.json"))
ap.add_argument("--only", nargs="*", default=None)
a = ap.parse_args()

LAYERS = (15, 19)
rel_dir = os.path.join(a.models, "qwen35-p50c")
man = json.load(open(os.path.join(a.exp8, "runs/p50c/edit/manifest.json")))

# 1. the released edit
if not os.path.exists(E.weight_path(rel_dir)):
    os.makedirs(rel_dir, exist_ok=True)
    tmp = E.weight_path(rel_dir) + ".tmp.safetensors"
    subprocess.run([sys.executable, os.path.join(a.exp8, "code/apply_edit.py"), E.weight_path(a.base),
                    os.path.join(a.exp8, "runs/p50c/edit"), tmp], check=True)
    os.replace(tmp, E.weight_path(rel_dir))
    E.link_base_files(a.base, rel_dir)
    json.dump({"name": "p50c", "about": "the paper's released edit, rebuilt with its apply_edit.py",
               "edited_sha256": {t["name"]: t["edited_sha256"] for t in man["tensors"]}},
              open(os.path.join(rel_dir, "edit_note.json"), "w"), indent=1)

# 2. read the edit back out of the weights
info = {"layers": {}}
for L in LAYERS:
    n = E.NAME.format(L)
    W0, W1 = E.load_down(a.base, L), E.load_down(rel_dir, L)
    want = next(t["edited_sha256"] for t in man["tensors"] if t["name"] == n)
    assert E.sha256(W1) == want, f"layer {L}: released tensor hash mismatch"
    r = E.recover(W0, W1)
    u = r["u"]
    if u[u.abs().argmax()] < 0:  # SVD's sign is arbitrary; fix one convention (labs/03 finds the loop sign)
        u = -u
    D = (W1.float() - W0.float())
    info["layers"][str(L)] = {
        "u": [round(float(x), 8) for x in u], "alpha": r["alpha"], "cols": r["cols"],
        "singular_values_top5": [float(x) for x in r["singular_values"][:5]],
        "rank1_share": float(r["singular_values"][0] ** 2 / (r["singular_values"] ** 2).sum()),
        "relative_change": float(D.norm() / W0.float().norm()),
        "changed_entries": int((W1.view(torch.int16) != W0.view(torch.int16)).sum()),
    }
    # 3. self-check: our formula with what we recovered vs the released bits
    for alpha in (r["alpha"], 6.0):
        W2 = E.projection_edit(W0, u, alpha, r["cols"])
        diff = (W2.view(torch.int16) != W1.view(torch.int16))
        info["layers"][str(L)][f"rebuild_alpha_{alpha:.4f}_entries_differing"] = int(diff.sum())
        info["layers"][str(L)][f"rebuild_alpha_{alpha:.4f}_max_abs_diff"] = float((W2.float() - W1.float()).abs().max())
    print(f"layer {L}: alpha={r['alpha']:.4f}, columns={len(r['cols'])}, rank-1 share="
          f"{info['layers'][str(L)]['rank1_share']:.6f}, relative change={info['layers'][str(L)]['relative_change']:.5f}")
u15, u19 = (torch.tensor(info["layers"][str(L)]["u"]) for L in LAYERS)
info["cos_u15_u19"] = float(u15 @ u19 / u15.norm() / u19.norm())
os.makedirs(os.path.dirname(a.info), exist_ok=True)
json.dump(info, open(a.info, "w"), indent=1)
print(json.dumps({k: v for L, d in info["layers"].items() for k, v in d.items() if k.startswith("rebuild")}, indent=1))

# 4. variants
released = E.released_spec(a.info)
variants = json.load(open(a.variants))
for name, spec in variants.items():
    if name.startswith("_") or (a.only and name not in a.only):
        continue
    out = os.path.join(a.models, f"qwen35-{name}")
    note_p = os.path.join(out, "edit_note.json")
    weights_spec = {k: v for k, v in spec.items() if k != "about"}
    old = json.load(open(note_p)).get("spec", {}) if os.path.exists(note_p) else None
    if old is not None and {k: v for k, v in old.items() if k != "about"} == weights_spec:
        print("exists", name)
        continue
    new, desc = E.build_from_spec(a.base, spec, released)
    E.write_variant(a.base, out, new, {"name": name, "spec": spec, "built": desc})
    print("built", name, json.dumps(desc))