from __future__ import annotations

import importlib.util
import json
import math
import sys
from pathlib import Path

import torch
from torch.nn import functional as F


TRAINING_SCRIPT = Path(__file__).with_name("train-tiny-transformer.py")
spec = importlib.util.spec_from_file_location("tiny_training", TRAINING_SCRIPT)
training = importlib.util.module_from_spec(spec)
assert spec.loader is not None
sys.modules[spec.name] = training
spec.loader.exec_module(training)

TOKENS = ["<bos>", "A", "B", "C", "D", "X", "Y", "<eos>"]


def split_heads(records: torch.Tensor, heads: int) -> torch.Tensor:
    batch, length, width = records.shape
    return records.reshape(batch, length, heads, width // heads).transpose(1, 2)


def round_list(values: torch.Tensor, digits: int = 6) -> list:
    return values.detach().cpu().round(decimals=digits).tolist()


@torch.inference_mode()
def traced_forward(
    model,
    token_ids: list[int],
    *,
    ablate_head: tuple[int, int] | None = None,
    ablate_attention_layer: int | None = None,
    ablate_mlp_layer: int | None = None,
    patch: tuple[str, int, int, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, dict]:
    ids = torch.tensor([token_ids], dtype=torch.long)
    length = ids.shape[1]
    hidden = model.token_embedding(ids) + model.position_embedding[:length]
    trace: dict = {"embedding": hidden.clone(), "layers": []}
    causal = torch.triu(
        torch.full((length, length), float("-inf")), diagonal=1
    )

    for layer_index, block in enumerate(model.blocks):
        residual_before = hidden
        normalized_attention = block.norm1(residual_before)
        qkv = F.linear(normalized_attention, block.attention.in_proj_weight)
        query, key, value = (
            split_heads(part, model.config.n_heads) for part in qkv.chunk(3, dim=-1)
        )
        scores = query @ key.transpose(-2, -1) / math.sqrt(
            model.config.d_model // model.config.n_heads
        )
        weights = (scores + causal).softmax(dim=-1)
        mixed_values = weights @ value

        head_writes = []
        head_width = model.config.d_model // model.config.n_heads
        for head_index in range(model.config.n_heads):
            start = head_index * head_width
            end = start + head_width
            write = F.linear(
                mixed_values[:, head_index],
                block.attention.out_proj.weight[:, start:end],
            )
            if ablate_head == (layer_index, head_index):
                write = torch.zeros_like(write)
            head_writes.append(write)
        stacked_head_writes = torch.stack(head_writes, dim=1)
        attention_write = stacked_head_writes.sum(dim=1)
        if ablate_attention_layer == layer_index:
            attention_write = torch.zeros_like(attention_write)
        if patch is not None and patch[0] == "attention" and patch[1] == layer_index:
            attention_write[:, patch[2]] = patch[3]
        residual_after_attention = residual_before + attention_write

        normalized_mlp = block.norm2(residual_after_attention)
        mlp_preactivation = block.mlp[0](normalized_mlp)
        mlp_activation = block.mlp[1](mlp_preactivation)
        mlp_write = block.mlp[2](mlp_activation)
        if ablate_mlp_layer == layer_index:
            mlp_write = torch.zeros_like(mlp_write)
        hidden = residual_after_attention + mlp_write
        if patch is not None and patch[0] == "residual" and patch[1] == layer_index:
            hidden[:, patch[2]] = patch[3]

        trace["layers"].append(
            {
                "residual_before": residual_before.clone(),
                "attention_weights": weights.clone(),
                "head_writes": stacked_head_writes.clone(),
                "attention_write": attention_write.clone(),
                "residual_after_attention": residual_after_attention.clone(),
                "mlp_preactivation": mlp_preactivation.clone(),
                "mlp_activation": mlp_activation.clone(),
                "mlp_write": mlp_write.clone(),
                "residual_after_layer": hidden.clone(),
            }
        )

    normalized = model.final_norm(hidden)
    logits = model.readout(normalized)
    trace["final_residual"] = hidden.clone()
    trace["final_normalized"] = normalized.clone()
    trace["logits"] = logits.clone()
    return logits, trace


def top_token(model, residual: torch.Tensor) -> tuple[str, float]:
    logits = model.readout(model.final_norm(residual))
    token_id = int(logits.argmax())
    return TOKENS[token_id], float(logits[token_id])


def target_margin(logits: torch.Tensor, target: int, alternative: int) -> float:
    return float(logits[0, -1, target] - logits[0, -1, alternative])


def main() -> None:
    config = training.Config()
    training_result, state = training.train(config, return_state=True)
    model = training.TinyDecoder(config)
    model.load_state_dict(state)
    model.eval()
    torch.set_grad_enabled(False)

    clean_ids = [0, 5, 6]       # <bos> X Y -> X
    corrupted_ids = [0, 6, 5]   # <bos> Y X -> Y
    target_id = 5
    alternative_id = 6
    clean_logits, clean = traced_forward(model, clean_ids)
    corrupted_logits, corrupted = traced_forward(model, corrupted_ids)
    standard_clean = model(torch.tensor([clean_ids], dtype=torch.long))
    trace_logit_error = float((clean_logits - standard_clean).abs().max())
    assert trace_logit_error < 1e-5
    for layer in clean["layers"]:
        assert torch.allclose(
            layer["attention_weights"].sum(dim=-1),
            torch.ones_like(layer["attention_weights"].sum(dim=-1)),
            atol=1e-6,
        )
    clean_margin = target_margin(clean_logits, target_id, alternative_id)
    corrupted_margin = target_margin(corrupted_logits, target_id, alternative_id)

    attention_rows = []
    component_norms = []
    vocabulary_projections = [
        {
            "stage": "embedding_plus_position",
            "top_token": top_token(model, clean["embedding"][0, -1])[0],
            "target_minus_alternative": target_margin(
                model.readout(model.final_norm(clean["embedding"])),
                target_id,
                alternative_id,
            ),
        }
    ]
    for layer_index, layer in enumerate(clean["layers"]):
        for head_index in range(config.n_heads):
            attention_rows.append(
                {
                    "layer": layer_index + 1,
                    "head": head_index + 1,
                    "final_query_weights": round_list(
                        layer["attention_weights"][0, head_index, -1]
                    ),
                }
            )
        component_norms.append(
            {
                "layer": layer_index + 1,
                "residual_before": float(layer["residual_before"][0, -1].norm()),
                "attention_write": float(layer["attention_write"][0, -1].norm()),
                "mlp_write": float(layer["mlp_write"][0, -1].norm()),
                "residual_after_layer": float(
                    layer["residual_after_layer"][0, -1].norm()
                ),
            }
        )
        for stage_name in ("residual_after_attention", "residual_after_layer"):
            stage_logits = model.readout(model.final_norm(layer[stage_name]))
            token_id = int(stage_logits[0, -1].argmax())
            vocabulary_projections.append(
                {
                    "stage": f"layer_{layer_index + 1}_{stage_name}",
                    "top_token": TOKENS[token_id],
                    "target_minus_alternative": target_margin(
                        stage_logits, target_id, alternative_id
                    ),
                }
            )

    ablations = []
    for layer_index in range(config.n_layers):
        for head_index in range(config.n_heads):
            logits, _ = traced_forward(
                model, clean_ids, ablate_head=(layer_index, head_index)
            )
            margin = target_margin(logits, target_id, alternative_id)
            ablations.append(
                {
                    "component": f"layer_{layer_index + 1}_head_{head_index + 1}",
                    "target_minus_alternative": margin,
                    "change_from_clean": margin - clean_margin,
                }
            )
        for kind in ("attention", "mlp"):
            logits, _ = traced_forward(
                model,
                clean_ids,
                ablate_attention_layer=layer_index if kind == "attention" else None,
                ablate_mlp_layer=layer_index if kind == "mlp" else None,
            )
            margin = target_margin(logits, target_id, alternative_id)
            ablations.append(
                {
                    "component": f"layer_{layer_index + 1}_{kind}",
                    "target_minus_alternative": margin,
                    "change_from_clean": margin - clean_margin,
                }
            )

    patches = []
    denominator = clean_margin - corrupted_margin
    for layer_index in range(config.n_layers):
        for position in range(len(clean_ids)):
            for kind, trace_key in (
                ("attention", "attention_write"),
                ("residual", "residual_after_layer"),
            ):
                clean_value = clean["layers"][layer_index][trace_key][0, position]
                logits, _ = traced_forward(
                    model,
                    corrupted_ids,
                    patch=(kind, layer_index, position, clean_value),
                )
                margin = target_margin(logits, target_id, alternative_id)
                patches.append(
                    {
                        "site": f"layer_{layer_index + 1}_{kind}_position_{position}",
                        "target_minus_alternative": margin,
                        "recovery_fraction": (margin - corrupted_margin) / denominator,
                    }
                )

    active_neurons = []
    for layer_index, layer in enumerate(clean["layers"]):
        activations = layer["mlp_activation"][0, -1]
        for neuron_index in torch.argsort(activations, descending=True)[:3]:
            active_neurons.append(
                {
                    "layer": layer_index + 1,
                    "neuron": int(neuron_index),
                    "activation": float(activations[neuron_index]),
                }
            )

    result = {
        "torch_version": torch.__version__,
        "training_final_validation_loss": training_result["final_validation_loss"],
        "trace_logit_error": trace_logit_error,
        "clean_prompt": [TOKENS[index] for index in clean_ids],
        "corrupted_prompt": [TOKENS[index] for index in corrupted_ids],
        "metric": "logit(X) - logit(Y) at the final position",
        "clean_margin": clean_margin,
        "corrupted_margin": corrupted_margin,
        "clean_probabilities": round_list(clean_logits[0, -1].softmax(dim=-1)),
        "corrupted_probabilities": round_list(
            corrupted_logits[0, -1].softmax(dim=-1)
        ),
        "attention_rows": attention_rows,
        "component_norms": component_norms,
        "logit_lens_style_projections": vocabulary_projections,
        "ablations": ablations,
        "activation_patches": patches,
        "largest_absolute_ablation": max(
            ablations, key=lambda row: abs(row["change_from_clean"])
        ),
        "largest_recovery_patch": max(
            patches, key=lambda row: row["recovery_fraction"]
        ),
        "largest_final_position_mlp_activations": active_neurons,
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()