"""Check exact restoration with identical batch composition before and after perturbation."""

import importlib.util
import json
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 = "efc3f02e5775d3d8428b4961f1fd94aedc6ca34f"


def main():
    source = hf_hub_download(REPO, "train_restriction_pilot.py", revision=REVISION)
    spec = importlib.util.spec_from_file_location("restriction_helper", source)
    helper = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(helper)
    report = json.loads(
        Path(hf_hub_download(REPO, "restriction_results.json", revision=REVISION)).read_text()
    )
    rows = report["rows"][:14]
    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=REVISION).eval()

    @torch.no_grad()
    def capture():
        texts = []
        for offset in range(0, 14, 8):
            inputs = tokenizer(
                [helper.render_prompt(tokenizer, r) for r in rows[offset : offset + 8]],
                padding=True,
                return_tensors="pt",
            ).to("cuda")
            outputs = model.generate(
                **inputs, max_new_tokens=80, do_sample=False, pad_token_id=tokenizer.pad_token_id
            )
            texts.extend(
                tokenizer.batch_decode(
                    outputs[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True
                )
            )
        return texts

    before = capture()
    modules = [m for m in model.modules() if hasattr(m, "lora_A") and "default" in m.lora_A]
    originals = [
        (m.lora_A["default"].weight.clone(), m.lora_B["default"].weight.clone()) for m in modules
    ]
    torch.manual_seed(1729)
    with torch.no_grad():
        for module in modules:
            module.lora_A["default"].weight.normal_(std=0.01)
            module.lora_B["default"].weight.normal_(std=0.01)
    perturbed = capture()
    with torch.no_grad():
        for module, (a, b) in zip(modules, originals, strict=True):
            module.lora_A["default"].weight.copy_(a)
            module.lora_B["default"].weight.copy_(b)
    after = capture()
    assert all(
        torch.equal(a, m.lora_A["default"].weight) and torch.equal(b, m.lora_B["default"].weight)
        for m, (a, b) in zip(modules, originals, strict=True)
    )
    result = dict(
        source_revision=REVISION,
        rows=14,
        batch_sizes=[8, 6],
        weight_restoration_exact=True,
        raw_matches=sum(a == b for a, b in zip(before, after, strict=True)),
        before=before,
        after=after,
        perturbed=perturbed,
        original_capture_matches=sum(
            a == b["raw"]
            for a, b in zip(before, report["conditions"]["selective"]["outputs"][:14], strict=True)
        ),
        note=("Matched-batch restoration; original 96-case capture "
              "has a different second batch composition."),
    )
    folder = Path("restriction_parity")
    folder.mkdir(exist_ok=True)
    (folder / "parity.json").write_text(json.dumps(result, indent=2))
    (folder / "train_restriction_parity_pilot.py").write_bytes(Path(__file__).read_bytes())
    HfApi().upload_folder(
        repo_id=REPO,
        folder_path=str(folder),
        path_in_repo="parity",
        commit_message="Check exact matched-batch restoration",
    )
    print(
        json.dumps({k: v for k, v in result.items() if k not in ("before", "after", "perturbed")}),
        flush=True,
    )


if __name__ == "__main__":
    main()