"""Match donor capture to generation prefill, including the final layer before normalization."""

import importlib.util
import json
import re
from pathlib import Path

import torch
from huggingface_hub import HfApi, hf_hub_download
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "itsalissonsilva/Restriction-Learning-Experiment"
REVISION = "3a3cb6bef28c71ffc822f0a68645e2adc4b6aed2"
ADAPTER_REVISION = "efc3f02e5775d3d8428b4961f1fd94aedc6ca34f"


def main():
    torch.manual_seed(1729)
    source = hf_hub_download(REPO, "train_restriction_pilot.py", revision=ADAPTER_REVISION)
    spec = importlib.util.spec_from_file_location("helper", source)
    helper = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(helper)
    report = json.loads(
        Path(hf_hub_download(REPO, "interpretability/results.json", revision=REVISION)).read_text()
    )
    tokenizer = AutoTokenizer.from_pretrained(helper.BASE, revision=helper.REVISION)
    tokenizer.padding_side = "left"
    base = AutoModelForCausalLM.from_pretrained(
        helper.BASE,
        revision=helper.REVISION,
        torch_dtype=torch.bfloat16,
        attn_implementation="eager",
    ).to("cuda")
    model = PeftModel.from_pretrained(
        base, REPO, subfolder="selective", revision=ADAPTER_REVISION
    ).eval()
    layers = [(n, m) for n, m in model.named_modules() if re.search(r"\.layers\.\d+$", n)]
    private = report["patch_pairs"]["private"]
    public = report["patch_pairs"]["public"]

    @torch.no_grad()
    def generate(rows, max_tokens=80):
        inputs = tokenizer(
            [helper.render_prompt(tokenizer, r) for r in rows], padding=True, return_tensors="pt"
        ).to("cuda")
        generated = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
        )
        texts = tokenizer.batch_decode(
            generated[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True
        )
        return [helper.grading(r, t) for r, t in zip(rows, texts, strict=True)]

    def capture(rows):
        stored = {}
        handles = []
        for i, (_, layer) in enumerate(layers):

            def hook(module, args, output, index=i):
                h = output[0] if isinstance(output, tuple) else output
                if h.shape[1] > 1:
                    stored[index] = h[:, -1].detach().clone()

            handles.append(layer.register_forward_hook(hook))
        try:
            generate(rows, 1)
        finally:
            for handle in handles:
                handle.remove()
        assert len(stored) == 26
        return stored

    hp, hc = capture(private), capture(public)
    baseline = generate(private)
    result = dict(
        source_revision=REVISION,
        source_patch_exclusion_reason="Earlier capture used ordinary forward position IDs and post-normalized final hidden state; superseded by identical generation-prefill hooks at every patch site",
        patching=[],
        patch_pairs=dict(
            private=private,
            public=public,
            baseline_private=baseline,
            baseline_public=generate(public),
        ),
        seed=1729,
    )
    for layer in range(26):
        delta = hc[layer].float() - hp[layer].float()
        noise = torch.randn_like(delta)
        noise *= delta.norm(dim=1, keepdim=True) / noise.norm(dim=1, keepdim=True).clamp_min(1e-20)
        for donor_kind, donor in [
            ("public", hc[layer]),
            ("self_restoration", hp[layer]),
            ("random_norm_matched", hp[layer] + noise.to(hp[layer].dtype)),
        ]:

            def patch(module, args, output):
                h = output[0] if isinstance(output, tuple) else output
                if h.shape[1] <= 1:
                    return output
                replacement = h.clone()
                replacement[:, -1] = donor
                return (replacement, *output[1:]) if isinstance(output, tuple) else replacement

            handle = layers[layer][1].register_forward_hook(patch)
            try:
                outputs = generate(private)
            finally:
                handle.remove()
            matches = sum(o["raw"] == b["raw"] for o, b in zip(outputs, baseline, strict=True))
            if donor_kind == "self_restoration":
                assert matches == 8, (layer, matches)
            result["patching"].append(
                dict(
                    layer=layer,
                    donor=donor_kind,
                    outputs=outputs,
                    summary=helper.summarize(private, outputs),
                    raw_matches_private_baseline=matches,
                )
            )
        print(
            json.dumps(dict(stage="patching", layer=layer, self_restoration_matches=8)), flush=True
        )
    assert generate(private) == baseline
    folder = Path("patchfix")
    folder.mkdir(exist_ok=True)
    (folder / "patching.json").write_text(json.dumps(result, indent=2))
    (folder / "train_interpretability_patchfix_pilot.py").write_bytes(Path(__file__).read_bytes())
    HfApi().upload_folder(
        repo_id=REPO,
        folder_path=str(folder),
        path_in_repo="interpretability_patchfix",
        commit_message="Correct donor capture to exact generation-prefill patch sites; verify all 26 self-restorations",
    )
    print(json.dumps(dict(stage="complete", layers=26, self_restoration_matches=208)), flush=True)


if __name__ == "__main__":
    main()