#!/usr/bin/env python
"""Lab 7: what else does the edit change?

    . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
    CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py          # a few minutes on one GPU

For the base model and every edited model built in models/:
  1. Ordinary text. Perplexity on the first 64 x 1,024 tokens of WikiText-2 (test), and KL(base || edited):
     how far the edited model's next-token prediction moves from the base model's, averaged over tokens.
  2. The agent's 8 texts. KL on its productive texts vs its looping texts. A good edit moves predictions a lot
     where the model loops and barely anywhere else.
Every edited model differs from the base only in the down_proj weights of the layers its edit_note.json lists
(15 and 19 for most of them), so the script loads the base once and swaps those matrices in and out.

Experiments. Any flag other than the defaults writes to results/lab7/try-NAME/ (figure included), so the main
results stay as they are:
    python labs/lab7_side_effects.py --chunks 8                   # quick check on 8 x 1,024 WikiText tokens
    python labs/lab7_side_effects.py --last 36                    # KL on only the last 36 tokens of the 8 texts
    python labs/lab7_side_effects.py --models p50c capA-a6        # only some models (after Lab 8: your own edit)
"""
import argparse
import json
import os
import sys
import time

import numpy as np
import torch

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 NAME, load  # noqa: E402

ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--models", nargs="+", metavar="NAME",
                help="only these edited models, named like their folders without 'qwen35-' (default: all of them)")
ap.add_argument("--chunks", type=int, default=64, help="how many 1,024-token WikiText-2 chunks to score (default 64)")
ap.add_argument("--last", type=int, default=0,
                help="measure KL on the 8 texts over only their last N tokens (default 0: every token)")
ap.add_argument("--try-name", help="folder name for an experiment; results go to results/lab7/try-NAME/")
a = ap.parse_args()

t_start = time.time()
CHUNK = 1024
parts = []
if a.chunks != 64:
    parts.append(f"chunks-{a.chunks}")
if a.last:
    parts.append(f"last-{a.last}")
if a.models:
    parts.append("models-" + ("-".join(a.models) if len(a.models) <= 4 else str(len(a.models))))
try_name = a.try_name or "-".join(parts)
OUT = os.path.join(ROOT, "results", "lab7", f"try-{try_name}" if try_name else "")
FIG = OUT if try_name else os.path.join(ROOT, "results", "figures")
os.makedirs(OUT, exist_ok=True)
os.makedirs(FIG, exist_ok=True)

MODELS = os.path.join(hf.LAB, "models")
found = sorted(d[len("qwen35-"):] for d in os.listdir(MODELS)
               if d.startswith("qwen35-") and os.path.exists(os.path.join(MODELS, d, "edit_note.json")))
missing = [n for n in a.models or [] if n not in found]
if missing:
    sys.exit(f"no edited model folder for {missing}; found {found}")


def edited_layers(n):
    note = json.load(open(os.path.join(MODELS, "qwen35-" + n, "edit_note.json")))
    return {int(L) for L in note["spec"]["layers"]} if "spec" in note else {15, 19}  # p50c: the released edit


# swap every layer that any chosen model edits; a model that leaves one of them alone just carries the base weights
LAYERS = sorted(set().union(*(edited_layers(n) for n in a.models or found)))
m, tok = hf.load("base")
mods = {L: hf.down_proj(m, L) for L in LAYERS}
base_w = {L: mods[L].weight.detach().clone() for L in LAYERS}
variants = {}
for n in a.models or found:
    t = load(os.path.join(MODELS, "qwen35-" + n), [NAME.format(L) for L in LAYERS])
    variants[n] = {L: t[NAME.format(L)].to("cuda") for L in LAYERS}
print("models:", ["base"] + list(variants), "| down_proj swapped at layers", LAYERS)


def set_weights(w):
    for L in LAYERS:
        mods[L].weight.data.copy_(w[L])


from datasets import load_dataset  # noqa: E402

wiki = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="test")
ids = tok("".join(wiki["text"]), add_special_tokens=False).input_ids
if a.chunks > len(ids) // CHUNK:
    sys.exit(f"WikiText-2 test has only {len(ids) // CHUNK} chunks of {CHUNK} tokens")
chunks = torch.tensor(ids[:CHUNK * a.chunks]).view(a.chunks, CHUNK)
pairs = json.load(open(os.path.join(ROOT, "labs", "data", "astra_pairs.json")))
pair_ids = [hf.encode(tok, t) for t in pairs["texts"]]

names = ["base"] + list(variants)
nll = {n: [] for n in names}
kl = {n: [] for n in names}
kl_pairs = {n: {"productive": [], "looping": []} for n in names}


def logprobs(x):
    return m(x).logits[0, :-1].float().log_softmax(-1)


def kl_text(lb, lv):
    k = (lb.exp() * (lb - lv)).sum(-1)
    return (k[-a.last:] if a.last else k).mean().item()


with torch.no_grad():
    for c in range(a.chunks):
        x = chunks[c:c + 1].cuda()
        set_weights(base_w)
        lb = logprobs(x)
        tgt = x[0, 1:, None]
        nll["base"].append(-lb.gather(1, tgt).mean().item())
        kl["base"].append(0.0)
        for n, w in variants.items():
            set_weights(w)
            lv = logprobs(x)
            nll[n].append(-lv.gather(1, tgt).mean().item())
            kl[n].append((lb.exp() * (lb - lv)).sum(-1).mean().item())
    for i, x in enumerate(pair_ids):
        kind = "looping" if i % 2 else "productive"
        set_weights(base_w)
        lb = logprobs(x)
        kl_pairs["base"][kind].append(0.0)
        for n, w in variants.items():
            set_weights(w)
            lv = logprobs(x)
            kl_pairs[n][kind].append(kl_text(lb, lv))
set_weights(base_w)

res = {}
for n in names:
    res[n] = {"wikitext_ppl": float(np.exp(np.mean(nll[n]))),
              "wikitext_kl": float(np.mean(kl[n])),
              "kl_productive_texts": float(np.mean(kl_pairs[n]["productive"])),
              "kl_looping_texts": float(np.mean(kl_pairs[n]["looping"]))}
for n in names:
    res[n]["ppl_change_pct"] = 100 * (res[n]["wikitext_ppl"] / res["base"]["wikitext_ppl"] - 1)
json.dump({"tokens": CHUNK * a.chunks, "last": a.last, "models": res},
          open(os.path.join(OUT, "side_effects.json"), "w"), indent=1)
print(f"{'model':10s} {'ppl':>8s} {'ppl +%':>7s} {'KL wiki':>9s} {'KL prod':>9s} {'KL loop':>9s}")
for n in names:
    r = res[n]
    print(f"{n:10s} {r['wikitext_ppl']:8.3f} {r['ppl_change_pct']:7.2f} {r['wikitext_kl']:9.5f} "
          f"{r['kl_productive_texts']:9.5f} {r['kl_looping_texts']:9.5f}")

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

order = [n for n in names if n != "base"]
fig, ax = plt.subplots(1, 2, figsize=(14, 4.5))
ax[0].bar(order, [res[n]["ppl_change_pct"] for n in order], color="#0969da")
ax[0].set_ylabel("WikiText-2 perplexity change vs base (%)")
ax[0].set_title("Damage to ordinary text")
xi = np.arange(len(order))
ax[1].bar(xi - 0.2, [res[n]["kl_productive_texts"] for n in order], 0.4, label="productive texts", color="#1a7f37")
ax[1].bar(xi + 0.2, [res[n]["kl_looping_texts"] for n in order], 0.4, label="looping texts", color="#cf222e")
ax[1].set_xticks(xi, order)
ax[1].set_yscale("log")
ax[1].set_ylabel("KL(base || edited), mean per token")
ax[1].set_title("Where the edit moves predictions")
ax[1].legend()
for a_ in ax:
    a_.tick_params(axis="x", rotation=60)
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab7_side_effects.png"), dpi=130)
print(f"wrote {OUT} in {time.time() - t_start:.0f} s")