"""Measured cross-model directions, steering controls and token-level causal patching."""

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

import numpy as np
import torch
from huggingface_hub import HfApi, hf_hub_download
from peft import PeftModel
from safetensors.torch import save_file
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "itsalissonsilva/Restriction-Learning-Experiment"
REVISION = "efc3f02e5775d3d8428b4961f1fd94aedc6ca34f"


def main():
    started = time.monotonic()
    torch.manual_seed(1729)
    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)
    original = json.loads(
        Path(hf_hub_download(REPO, "restriction_results.json", revision=REVISION)).read_text()
    )
    train = [r for r in helper.make_rows("train", 24) if r["kind"] != "ambiguous"]
    dev = [r for r in helper.make_rows("development", 4) if r["kind"] != "ambiguous"]
    test = [r for r in original["rows"][:84] if r["kind"] != "ambiguous"]
    intervention_rows = original["rows"][:14] + original["rows"][-4:]
    result = dict(
        seed=1729,
        adapter_revision=REVISION,
        training_rows=train,
        development_rows=dev,
        test_rows=test,
        intervention_rows=intervention_rows,
        models={},
        cross_model={},
        patching=[],
        scope="Synthetic workspace restriction intent; one seed; no general unsafe-intent or universal-circuit claim",
    )
    tensors = {}
    for key, base_name, revision in [
        ("gemma", helper.BASE, helper.REVISION),
        ("qwen", "Qwen/Qwen3-0.6B", HfApi().model_info("Qwen/Qwen3-0.6B").sha),
    ]:
        tokenizer = AutoTokenizer.from_pretrained(base_name, revision=revision)
        tokenizer.padding_side = "left"
        model = (
            AutoModelForCausalLM.from_pretrained(
                base_name,
                revision=revision,
                torch_dtype=torch.bfloat16,
                attn_implementation="eager",
            )
            .to("cuda")
            .eval()
        )
        if key == "gemma":
            model = PeftModel.from_pretrained(
                model, REPO, subfolder="selective", revision=REVISION
            ).eval()
        layers = [(n, m) for n, m in model.named_modules() if re.search(r"\.layers\.\d+$", n)]
        assert len(layers) >= 26

        @torch.no_grad()
        def capture(rows):
            chunks = []
            for offset in range(0, len(rows), 8):
                inputs = tokenizer(
                    [helper.render_prompt(tokenizer, r) for r in rows[offset : offset + 8]],
                    padding=True,
                    return_tensors="pt",
                ).to("cuda")
                outputs = model(**inputs, output_hidden_states=True)
                chunks.append(torch.stack([h[:, -1].float().cpu() for h in outputs.hidden_states]))
            return torch.cat(chunks, dim=1)

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

        htrain, hdev, htest = capture(train), capture(dev), capture(test)
        y = np.array([r["expected_action"] == "withhold_private" for r in train])
        yd = np.array([r["expected_action"] == "withhold_private" for r in dev])
        yt = np.array([r["expected_action"] == "withhold_private" for r in test])
        entry = dict(
            base_model=base_name,
            base_revision=revision,
            condition="selective restriction LoRA" if key == "gemma" else "base",
            probes=[],
            steering=[],
            activation_shape=list(htrain.shape),
        )
        for layer in range(htrain.shape[0]):
            x = htrain[layer].numpy()
            d = hdev[layer].numpy()
            t = htest[layer].numpy()
            mean = x.mean(0)
            scale = x.std(0).clip(1e-4)
            xs = (x - mean) / scale
            # Ridge linear discriminator; labels and validation threshold stay outside held-out test.
            targets = y.astype(float) * 2 - 1
            selected = np.argsort(np.abs(xs.T @ targets))[-64:]
            selected_x = xs[:, selected]
            sparse = selected_x.T @ np.linalg.solve(
                selected_x @ selected_x.T + 10 * np.eye(len(x)), targets
            )
            coef = np.zeros(x.shape[1])
            coef[selected] = sparse
            scores = ((d - mean) / scale) @ coef
            candidates = np.unique(np.r_[scores, 0])
            threshold = max(candidates, key=lambda th: np.mean((scores > th) == yd))
            prediction = ((t - mean) / scale) @ coef > threshold
            shuffled = np.random.default_rng(1729).permutation(targets)
            random_coef = xs.T @ np.linalg.solve(xs @ xs.T + 10 * np.eye(len(x)), shuffled)
            random_accuracy = float(np.mean((((t - mean) / scale) @ random_coef > 0) == yt))
            entry["probes"].append(
                dict(
                    hidden_state_index=layer,
                    nonzero_features=64,
                    validation_accuracy=float(np.mean((scores > threshold) == yd)),
                    test_accuracy=float(np.mean(prediction == yt)),
                    shuffled_label_accuracy=random_accuracy,
                    threshold=float(threshold),
                    test_scores=(((t - mean) / scale) @ coef).tolist(),
                )
            )
            tensors[f"{key}.probe.{layer}.mean"] = torch.from_numpy(mean)
            tensors[f"{key}.probe.{layer}.scale"] = torch.from_numpy(scale)
            tensors[f"{key}.probe.{layer}.coef"] = torch.from_numpy(coef.astype(np.float32))
        for label, hidden in [("train", htrain), ("development", hdev), ("test", htest)]:
            tensors[f"{key}.hidden.{label}"] = hidden
        directions = {}
        for layer in [12, 24]:
            direction = htrain[layer + 1, y].mean(0) - htrain[layer + 1, ~y].mean(0)
            random_direction = torch.randn_like(direction)
            random_direction *= direction.norm() / random_direction.norm()
            directions[layer] = direction
            tensors[f"{key}.direction.{layer}"] = direction
            tensors[f"{key}.random_direction.{layer}"] = random_direction
            for kind, vector in [("learned", direction), ("random", random_direction)]:
                for alpha in [-2, -1, 0, 1, 2]:

                    def hook(module, args, output):
                        hidden = output[0] if isinstance(output, tuple) else output
                        if hidden.shape[1] <= 1:
                            return output
                        changed = hidden.clone()
                        changed[:, -1] += alpha * vector.to(
                            device=hidden.device, dtype=hidden.dtype
                        )
                        return (changed, *output[1:]) if isinstance(output, tuple) else changed

                    handle = layers[layer][1].register_forward_hook(hook)
                    try:
                        outputs = generate(intervention_rows)
                    finally:
                        handle.remove()
                    entry["steering"].append(
                        dict(
                            layer=layer,
                            alpha=alpha,
                            direction=kind,
                            outputs=outputs,
                            summary=helper.summarize(intervention_rows, outputs),
                        )
                    )
                    print(
                        json.dumps(
                            dict(
                                stage="steering",
                                model=key,
                                layer=layer,
                                direction=kind,
                                alpha=alpha,
                            )
                        ),
                        flush=True,
                    )
        # Zero strength and hook removal use identical batches for restoration.
        baseline = generate(intervention_rows)
        zeros = [s for s in entry["steering"] if s["alpha"] == 0]
        entry["zero_strength_raw_matches"] = [
            sum(o["raw"] == b["raw"] for o, b in zip(s["outputs"], baseline, strict=True))
            for s in zeros
        ]
        assert all(n == len(intervention_rows) for n in entry["zero_strength_raw_matches"])
        if key == "gemma":
            # Same request and content; only the field's PRIVATE/PUBLIC label changes.
            private = []
            public = []
            for r in original["rows"]:
                if r["kind"] == "direct" and len(private) < 8:
                    a = dict(r, request="What is the recovery code in this record?")
                    b = dict(
                        a,
                        context=a["context"].replace("PRIVATE", "PUBLIC"),
                        expected_action="answer_public",
                        private_canary="",
                        required=[a["private_canary"]],
                    )
                    private.append(a)
                    public.append(b)
            hp, hc = capture(private), capture(public)
            baseline_private = generate(private, 8)
            baseline_public = generate(public, 8)
            result["patch_pairs"] = dict(
                private=private,
                public=public,
                baseline_private=baseline_private,
                baseline_public=baseline_public,
            )
            # Patch donor states at matched final-prompt positions only. Donor and receiver have identical batch layout.
            for layer in range(26):
                for donor_kind, donor in [
                    ("public", hc[layer + 1]),
                    ("self_restoration", hp[layer + 1]),
                    ("random_norm_matched", None),
                ]:

                    def patch(module, args, output):
                        hidden = output[0] if isinstance(output, tuple) else output
                        if hidden.shape[1] <= 1:
                            return output
                        changed = hidden.clone()
                        # One complete 8-row batch is used for patch capture/generation below.
                        if donor is not None:
                            replacement = donor.to(hidden.device, hidden.dtype)
                        else:
                            noise = torch.randn(
                                hidden.shape[0],
                                hidden.shape[-1],
                                device=hidden.device,
                                dtype=torch.float32,
                            )
                            noise *= (
                                (hc[layer + 1] - hp[layer + 1]).norm(dim=1).to(hidden.device)
                                / noise.norm(dim=1).clamp_min(1e-20)
                            )[:, None]
                            replacement = hp[layer + 1].to(hidden.device, hidden.dtype) + noise.to(
                                hidden.dtype
                            )
                        changed[:, -1] = replacement
                        return (changed, *output[1:]) if isinstance(output, tuple) else changed

                    # Explicit batch of all eight avoids donor-index ambiguity.
                    inputs = tokenizer(
                        [helper.render_prompt(tokenizer, r) for r in private],
                        padding=True,
                        return_tensors="pt",
                    ).to("cuda")
                    handle = layers[layer][1].register_forward_hook(patch)
                    try:
                        with torch.no_grad():
                            generated = model.generate(
                                **inputs,
                                max_new_tokens=80,
                                do_sample=False,
                                pad_token_id=tokenizer.pad_token_id,
                            )
                    finally:
                        handle.remove()
                    texts = tokenizer.batch_decode(
                        generated[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True
                    )
                    outputs = [
                        helper.grading(r, text) for r, text in zip(private, texts, strict=True)
                    ]
                    result["patching"].append(
                        dict(
                            layer=layer,
                            donor=donor_kind,
                            outputs=outputs,
                            summary=helper.summarize(private, outputs),
                            raw_matches_private_baseline=sum(
                                o["raw"] == b["raw"]
                                for o, b in zip(outputs, baseline_private, strict=True)
                            ),
                        )
                    )
                print(json.dumps(dict(stage="patching", layer=layer)), flush=True)
        result["models"][key] = entry
        model = None
        layers = []
        del htrain, hdev, htest
        torch.cuda.empty_cache()
    # Paired low-rank mapping, fit on training activations only. Report held-out alignment rather than claiming behavioral transfer.
    for layer in [12, 24]:
        x = tensors["qwen.hidden.train"][layer + 1].numpy()
        y = tensors["gemma.hidden.train"][layer + 1].numpy()
        xm, ym = x.mean(0), y.mean(0)
        u, s, vt = np.linalg.svd(x - xm, full_matrices=False)
        rank = 8
        left = vt[:rank].T
        right = np.linalg.solve(
            (s[:rank] ** 2)[:, None] * np.eye(rank) + 10 * np.eye(rank),
            (u[:, :rank] * s[:rank]).T @ (y - ym),
        )
        xt = tensors["qwen.hidden.test"][layer + 1].numpy()
        yt = tensors["gemma.hidden.test"][layer + 1].numpy()
        pred = (xt - xm) @ left @ right + ym
        result["cross_model"][str(layer)] = dict(
            rank=rank,
            heldout_mse=float(np.mean((pred - yt) ** 2)),
            train_mean_baseline_mse=float(np.mean((ym - yt) ** 2)),
            scope="paired activation reconstruction only; behavioral transfer gate not yet tested",
        )
        tensors[f"mapper.{layer}.left"] = torch.from_numpy(left)
        tensors[f"mapper.{layer}.right"] = torch.from_numpy(right)
    result["elapsed_seconds"] = time.monotonic() - started
    folder = Path("interpretability")
    folder.mkdir(exist_ok=True)
    (folder / "results.json").write_text(json.dumps(result, indent=2))
    save_file(
        {k: v.contiguous() for k, v in tensors.items()}, folder / "measured_activations.safetensors"
    )
    (folder / "train_interpretability_pilot.py").write_bytes(Path(__file__).read_bytes())
    HfApi().upload_folder(
        repo_id=REPO,
        folder_path=str(folder),
        path_in_repo="interpretability",
        commit_message="Measure Gemma/Qwen steering, probes, causal patching and paired activations",
    )
    print(json.dumps(dict(stage="complete", seconds=result["elapsed_seconds"])), flush=True)


if __name__ == "__main__":
    main()