import csv
from pathlib import Path

import torch
from transformers import AutoTokenizer, GPT2LMHeadModel


SOURCE_LAYER = 5
TARGET_LAYER = 6
SOURCE_NEURON = 541
CANDIDATES = [2374, 2712, 1506, 2662, 203, 3043]
MODEL_REVISION = "607a30d783dfa663caf39e06633721c8d4cfcd7e"
BOS_TOKEN_ID = 50256

SCRIPT_DIR = Path(__file__).resolve().parent
INPUT_FILE = SCRIPT_DIR / "541 MAE List - Moby Dick.txt"
OUTPUT_DIR = SCRIPT_DIR / "results"


def load_strings():
    if not INPUT_FILE.exists():
        raise FileNotFoundError(
            f"Missing input file:\n{INPUT_FILE}"
        )

    strings = []

    for raw_line in INPUT_FILE.read_text(
        encoding="utf-8"
    ).splitlines():
        line = raw_line.rstrip()

        if not line or line.lower() == "all 80 entries:":
            continue

        if line.startswith("- "):
            line = line[2:].rstrip()

        if line:
            strings.append(line)

    if len(strings) != 80:
        raise RuntimeError(
            f"Expected exactly 80 input rows, found {len(strings)}. "
            "Experiment stopped without writing results."
        )

    unique_strings = list(dict.fromkeys(strings))

    if len(unique_strings) != 62:
        raise RuntimeError(
            f"Expected exactly 62 unique strings, found "
            f"{len(unique_strings)}. Experiment stopped."
        )

    return strings, unique_strings


def top_two(values):
    result = torch.topk(values, k=2)

    return {
        "winner": int(result.indices[0]),
        "winner_value": float(result.values[0]),
        "runner_up": int(result.indices[1]),
        "runner_up_value": float(result.values[1]),
        "margin": float(result.values[0] - result.values[1]),
    }


class NeuronRunner:
    def __init__(self, model, tokenizer, device):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device

    @torch.inference_mode()
    def run(self, text, ablate_position=None):
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
        ).to(self.device)

        # Match the canonical znou/probe_v9 protocol: every candidate is
        # evaluated as BOS + its ordinary GPT-2 tokenisation.
        bos = torch.full(
            (inputs["input_ids"].shape[0], 1),
            BOS_TOKEN_ID,
            dtype=inputs["input_ids"].dtype,
            device=self.device,
        )
        inputs["input_ids"] = torch.cat(
            (bos, inputs["input_ids"]), dim=1
        )
        inputs["attention_mask"] = torch.cat(
            (torch.ones_like(bos), inputs["attention_mask"]), dim=1
        )

        captured = {}

        def source_hook(_module, _inputs, output):
            edited = output.clone()

            if ablate_position is not None:
                edited[
                    :,
                    ablate_position,
                    SOURCE_NEURON,
                ] = 0.0

            captured["source"] = (
                edited.detach().float().cpu()
            )
            return edited

        def target_hook(_module, _inputs, output):
            captured["target"] = (
                output.detach().float().cpu()
            )

        source_handle = self.model.transformer.h[
            SOURCE_LAYER
        ].mlp.act.register_forward_hook(source_hook)

        target_handle = self.model.transformer.h[
            TARGET_LAYER
        ].mlp.act.register_forward_hook(target_hook)

        try:
            self.model(**inputs, use_cache=False)
        finally:
            source_handle.remove()
            target_handle.remove()

        token_ids = inputs["input_ids"][0].detach().cpu()
        tokens = ["<BOS>"] + [
            self.tokenizer.decode([token_id])
            for token_id in token_ids.tolist()[1:]
        ]

        return {
            "source": captured["source"][0],
            "target": captured["target"][0],
            "tokens": tokens,
        }


def candidate_metrics(
    activations,
    candidate,
    reference_position,
):
    neuron_values = activations[:, candidate]
    maximum, maximum_position = torch.max(
        neuron_values,
        dim=0,
    )

    return {
        "max": float(maximum),
        "max_position": int(maximum_position),
        "at_reference": float(
            neuron_values[reference_position]
        ),
    }


def run_experiment(
    runner,
    strings,
    output_path,
    experiment_name,
):
    print()
    print("=" * 70)
    print(experiment_name)
    print(f"Inputs: {len(strings)}")
    print("=" * 70)

    rows = []

    for row_number, text in enumerate(strings, start=1):
        natural = runner.run(text)

        source_per_neuron = natural["source"].amax(dim=0)
        source_top = top_two(source_per_neuron)
        measured_l5_winner = source_top["winner"]

        source_values = natural["source"][
            :,
            SOURCE_NEURON,
        ]

        source_peak_value, source_peak_position = (
            torch.max(source_values, dim=0)
        )
        source_peak_position = int(source_peak_position)

        ablated = runner.run(
            text,
            ablate_position=source_peak_position,
        )

        natural_l6 = natural["target"].amax(dim=0)
        ablated_l6 = ablated["target"].amax(dim=0)

        natural_top = top_two(natural_l6)
        ablated_top = top_two(ablated_l6)

        row = {
            "row": row_number,
            "text": text,
            "token_count": len(natural["tokens"]),
            "l5_541_peak_position": source_peak_position,
            "l5_541_peak_token": (
                natural["tokens"][source_peak_position]
            ),
            "l5_541_peak_activation": float(
                source_peak_value
            ),
            "natural_l5_winner": measured_l5_winner,
            "natural_l5_winner_activation": (
                source_top["winner_value"]
            ),
            "natural_l5_runner_up": source_top["runner_up"],
            "natural_l5_margin": source_top["margin"],
            "l5_541_retained": (
                measured_l5_winner == SOURCE_NEURON
            ),
            "natural_l6_winner": natural_top["winner"],
            "natural_l6_winner_activation": (
                natural_top["winner_value"]
            ),
            "natural_l6_runner_up": (
                natural_top["runner_up"]
            ),
            "natural_l6_margin": natural_top["margin"],
            "ablated_l6_winner": ablated_top["winner"],
            "ablated_l6_winner_activation": (
                ablated_top["winner_value"]
            ),
            "ablated_l6_runner_up": (
                ablated_top["runner_up"]
            ),
            "ablated_l6_margin": ablated_top["margin"],
            "winner_changed": (
                natural_top["winner"]
                != ablated_top["winner"]
            ),
        }

        for candidate in CANDIDATES:
            natural_candidate = candidate_metrics(
                natural["target"],
                candidate,
                source_peak_position,
            )
            ablated_candidate = candidate_metrics(
                ablated["target"],
                candidate,
                source_peak_position,
            )

            prefix = f"l6_{candidate}"

            row[f"{prefix}_natural_max"] = (
                natural_candidate["max"]
            )
            row[f"{prefix}_ablated_max"] = (
                ablated_candidate["max"]
            )
            row[f"{prefix}_max_delta"] = (
                ablated_candidate["max"]
                - natural_candidate["max"]
            )
            row[f"{prefix}_natural_at_541_peak"] = (
                natural_candidate["at_reference"]
            )
            row[f"{prefix}_ablated_at_541_peak"] = (
                ablated_candidate["at_reference"]
            )
            row[f"{prefix}_peak_delta"] = (
                ablated_candidate["at_reference"]
                - natural_candidate["at_reference"]
            )

        rows.append(row)

        print(
            f"{row_number:>3}/{len(strings)}  "
            f"L5-N{measured_l5_winner:>4}  "
            f"L6-N{natural_top['winner']:>4} "
            f"→ L6-N{ablated_top['winner']:>4}"
        )

    output_path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    with output_path.open(
        "w",
        newline="",
        encoding="utf-8",
    ) as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=rows[0].keys(),
        )
        writer.writeheader()
        writer.writerows(rows)

    changed = sum(
        row["winner_changed"]
        for row in rows
    )
    retained = sum(
        row["l5_541_retained"]
        for row in rows
    )

    print()
    print(f"Saved: {output_path}")
    print(f"L5-N541 retained: {retained}/{len(rows)}")
    print(
        f"L6 winner changed: "
        f"{changed}/{len(rows)}"
    )

    for candidate in CANDIDATES:
        delta_key = f"l6_{candidate}_peak_delta"

        relevant = [
            row[delta_key]
            for row in rows
            if row["natural_l6_winner"] == candidate
        ]

        if relevant:
            mean_delta = sum(relevant) / len(relevant)

            print(
                f"L6-N{candidate}: "
                f"{len(relevant)} natural winners; "
                f"mean peak delta {mean_delta:+.4f}"
            )


def main():
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False

    strings, unique_strings = load_strings()

    device = torch.device(
        "cuda" if torch.cuda.is_available() else "cpu"
    )

    print(f"Device: {device}")
    print("Loading GPT-2 Small in float32...")

    tokenizer = AutoTokenizer.from_pretrained(
        "gpt2",
        revision=MODEL_REVISION,
        add_prefix_space=False,
    )
    model = GPT2LMHeadModel.from_pretrained(
        "gpt2",
        revision=MODEL_REVISION,
        dtype=torch.float32,
    ).to(device)
    model.eval()

    runner = NeuronRunner(
        model=model,
        tokenizer=tokenizer,
        device=device,
    )

    run_experiment(
        runner=runner,
        strings=strings,
        output_path=(
            OUTPUT_DIR
            / "n541_l5_to_l6_ablation_all_80.csv"
        ),
        experiment_name="FULL 80-ROW REPLAY",
    )

    run_experiment(
        runner=runner,
        strings=unique_strings,
        output_path=(
            OUTPUT_DIR
            / "n541_l5_to_l6_ablation_unique_62.csv"
        ),
        experiment_name="UNIQUE 62-STRING REPLAY",
    )

    print()
    print("Both experiments completed successfully.")


if __name__ == "__main__":
    main()