"""Run: PYTHONPATH=src python examples/quickstart.py (CPU; no downloads)."""
import torch

from vlm_probing.adapters import HookPoint, ModelReadout, TorchModelAdapter
from vlm_probing.attention import AttentionProfile
from vlm_probing.causal import ActivationPatching, AttentionKnockout
from vlm_probing.lenses import EmbedLens, LogitLens
from vlm_probing.metrics import token_logit_margin
from tiny_model import TinyModel


def main():
    torch.manual_seed(7)
    model = TinyModel().eval()
    readout = ModelReadout(model.lm_head, norm=model.norm)
    adapter = TorchModelAdapter(model, {
        "layer0": HookPoint("layers.0"),
        "layer1": HookPoint("layers.1"),
        "attention0": HookPoint("layers.0.attention.probs"),
        "scores0": HookPoint("layers.0.attention.scores"),
    }, model_id="local/tiny-fixture", readout=readout)
    clean = {"input_ids": torch.tensor([[1, 2, 3]]), "image_tokens": torch.randn(1, 2, 6)}
    corrupt = {**clean, "image_tokens": torch.zeros_like(clean["image_tokens"])}
    trace = adapter.run(clean, capture=["layer0", "layer1", "attention0"])

    lens = LogitLens(readout)
    layer_logits = lens.run(trace.activations["layer1"], layer=1).tensors["logits"]
    torch.testing.assert_close(layer_logits, trace.logits)
    print("Final-layer logit lens matches model logits.")

    embedding_lens = EmbedLens(model.embedding.weight)
    neighbors = embedding_lens.run(clean["image_tokens"], top_k=3)
    print("Visual-token embedding neighbors:", neighbors.tensors["token_ids"].tolist())

    visual = torch.tensor([[True, True, False, False, False]])
    profile = AttentionProfile().run(trace.activations["attention0"], groups={"visual": visual})
    print("Last-query visual attention mass:", profile.tensors["group_mass"][0, 0, -1].tolist())

    patcher = ActivationPatching()
    patched = adapter.run(corrupt, interventions={
        "layer0": lambda receiver: patcher.run(receiver, trace.activations["layer0"],
                                               mask=visual[..., None]).tensors["edited"]
    })
    baseline = adapter.run(corrupt)
    margin = lambda output: token_logit_margin(output.logits[:, -1], 4, 5)
    print("Visual-state patch margin change:", (margin(patched) - margin(baseline)).tolist())

    blocked = torch.zeros(1, 1, 5, 5, dtype=torch.bool)
    blocked[..., -1, :2] = True
    knockout = AttentionKnockout()
    intervention = adapter.run(clean, capture=["attention0"], interventions={
        "scores0": lambda logits: knockout.run(logits, blocked=blocked).tensors["edited"]
    })
    assert intervention.activations["attention0"][..., -1, :2].eq(0).all()
    print("Attention knockout blocked real query-to-image edges before value aggregation.")
    print("Knockout margin change:", (margin(intervention) - margin(trace)).tolist())


if __name__ == "__main__":
    main()