#!/usr/bin/env python
"""Lab 1: find the edit by comparing two weight files. No GPU needed.

    . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
    python labs/lab1_diff.py                          # under a minute on CPU
    python labs/lab1_diff.py --edited qwen35-all2     # diff any other model in /nvme/wsp-lab/models instead;
                                                      # writes to results/lab1/try-<name>, main results untouched

1. Compare every tensor of the base model and the released model. Which ones changed?
2. In each changed matrix: how many entries and which columns (neurons) changed?
3. Take the change apart with an SVD: is it rank one? If so, the left vector is the direction u, and a
   least-squares fit gives the strength alpha in  W' = W - alpha * u (u^T W[:, cols]).
"""
import argparse
import json
import os
import sys

import numpy as np
import torch
from safetensors import safe_open

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from wsp import hf  # noqa: E402
from wsp.edits import recover, weight_path  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--edited", default="qwen35-p50c", help="model folder (in /nvme/wsp-lab/models) to compare with base")
a = ap.parse_args()
TRY = None if a.edited == "qwen35-p50c" else "try-" + os.path.basename(a.edited.rstrip("/"))
OUT = os.path.join(ROOT, "results", "lab1", *([TRY] if TRY else []))
FIG = OUT if TRY else os.path.join(ROOT, "results", "figures")
os.makedirs(OUT, exist_ok=True)
os.makedirs(FIG, exist_ok=True)
base_dir, edit_dir = hf.model_dir("base"), hf.model_dir(a.edited)

# 1. every tensor, bit for bit
changed, n_total, n_params = [], 0, 0
with safe_open(weight_path(base_dir), "pt") as fb, safe_open(weight_path(edit_dir), "pt") as fe:
    for name in fb.keys():
        tb, te = fb.get_tensor(name), fe.get_tensor(name)
        n_total += 1
        n_params += tb.numel()
        if not torch.equal(tb, te):
            changed.append((name, tb, te))
print(f"{len(changed)} of {n_total} tensors differ ({n_params / 1e6:.0f}M numbers in all):")

res = {"tensors_total": n_total, "params_total": n_params, "changed": {}}
col_norms, svals = {}, {}
for name, W0, W1 in changed:
    L = int(name.split("layers.")[1].split(".")[0])
    D = W1.float() - W0.float()
    entries = int((W1 != W0).sum())
    cols = torch.nonzero(D.abs().sum(0) > 0).flatten().tolist()
    r = recover(W0, W1)
    s = r["singular_values"]
    col_norms[L] = D.norm(dim=0).numpy()
    svals[L] = s.numpy()
    res["changed"][name] = {"shape": list(W0.shape), "entries_changed": entries, "n_cols": len(cols),
                            "cols": cols, "rank1_share": float(s[0] ** 2 / (s ** 2).sum()), "alpha": r["alpha"],
                            "relative_change": float(D.norm() / W0.float().norm())}
    print(f"   {name}  shape {list(W0.shape)}")
    print(f"      entries changed: {entries:,} of {W0.numel():,}, all inside {len(cols)} columns (neurons)")
    print(f"      singular values of the change: {[round(float(x), 4) for x in s[:4]]} ...  "
          f"rank-one share {res['changed'][name]['rank1_share']:.6f}")
    print(f"      fitted alpha = {r['alpha']:.4f}; the matrix moved by {res['changed'][name]['relative_change']:.1%}")
    res["changed"][name]["u"] = r["u"].tolist()
names = list(res["changed"])
if len(names) == 2:
    u1, u2 = (torch.tensor(res["changed"][n]["u"]) for n in names)
    res["abs_cos_between_the_two_u"] = abs(float(u1 @ u2))
    print(f"|cos| between the two directions: {res['abs_cos_between_the_two_u']:.3f} "
          f"(1 = same direction, 0 = unrelated; the SVD sign is arbitrary, Lab 3 fixes it)")
json.dump(res, open(os.path.join(OUT, "diff.json"), "w"), indent=1)

import matplotlib  # noqa: E402
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

layers = sorted(col_norms)
fig, ax = plt.subplots(2, len(layers), figsize=(6.5 * len(layers), 6.5), squeeze=False)
for k, L in enumerate(layers):
    ax[0, k].plot(col_norms[L], color="#cf222e", lw=0.8)
    ax[0, k].set_title(f"layer {L}: size of the change in each of 3,584 columns")
    ax[0, k].set_xlabel("neuron (column of down_proj)")
    ax[1, k].bar(np.arange(1, 11), svals[L][:10], color="#0969da")
    ax[1, k].set_yscale("log")
    ax[1, k].set_title(f"layer {L}: singular values of the change (log scale)")
    ax[1, k].set_xlabel("rank")
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab1_diff.png"), dpi=130)
print("wrote", OUT)