"""Logit lens: see what the model "thinks" at each layer.

Projects hidden states from each layer through the final layer norm
and unembedding matrix (lm_head) to reveal how predictions evolve
across layers. Early layers often predict generic tokens, with the
final answer crystallizing in later layers.

Reference: https://www.lesswrong.com/posts/AcKRB8wDds238WkfB

Usage:
    python examples/logit_lens.py \\
        --base-url http://localhost:8000 \\
        --prompt "The Eiffel Tower is in the city of"

Requires a running vLLM server with vllm-lens installed.
"""

from __future__ import annotations

import argparse
from typing import Any

import torch
from vllm_lens import Hook

from vllm_lens.client import VLLMLensClient
from _utils import find_norm, get_num_layers


def run_logit_lens(
    client: VLLMLensClient,
    prompt: str,
    top_k: int = 5,
) -> dict[str, Any]:
    """Run the logit lens on a prompt.

    Uses a hook that projects hidden states through the model's own
    lm_head on the GPU, avoiding large activation transfers.
    """
    print(f"Prompt: {prompt!r}\n")

    n_layers = get_num_layers(client.model)
    all_layers = list(range(n_layers))

    def project_hook(ctx, h):
        """Project hidden states through norm + unembed weight, save top-k."""
        weight = ctx.get_parameter("lm_head.weight")
        norm = find_norm(ctx.model)
        with torch.no_grad():
            normed = norm(h) if norm is not None else h
            logits = normed.float() @ weight.float().T  # (seq_len, vocab_size)
            topk = logits.topk(top_k, dim=-1)
        # Key by layer index so PP merge preserves ordering.
        ctx.saved[f"ids_{ctx.layer_idx}"] = topk.indices.cpu()
        ctx.saved[f"logits_{ctx.layer_idx}"] = topk.values.float().cpu()
        return None

    hook = Hook(fn=project_hook, layer_indices=all_layers)
    # Prefetch lm_head weight so hooks on any PP stage can access it.
    client.prefetch_params(["lm_head.weight"])
    # max_tokens=1 so there's a single prefill pass and token count
    # matches between the echo'd logprobs and the hook's captured data.
    output = client.generate(
        prompt,
        max_tokens=1,
        hooks=[hook],
        logprobs=5,
        echo=True,
    )

    # Exclude the generated token — we only have hook data for prompt tokens.
    tokens = output.logprobs["tokens"][:-1] if output.logprobs else []
    print(f"Generated: {output.text!r}")
    print(f"Tokens ({len(tokens)}): {tokens}\n")

    assert output.hook_results is not None, "No hook results returned"
    saved = output.hook_results["0"]
    # Reconstruct ordered lists from layer-keyed results.
    top_ids = [saved[f"ids_{i}"] for i in range(n_layers)]
    top_logits = [saved[f"logits_{i}"] for i in range(n_layers)]

    return {
        "tokens": tokens,
        "generated": output.text,
        "top_ids": top_ids,
        "top_logits": top_logits,
        "n_layers": n_layers,
    }


def print_logit_lens(
    results: dict[str, Any],
    model_name: str,
    focus_position: int = -1,
) -> None:
    """Print the logit lens for a specific token position."""
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    tokens = results["tokens"]
    top_ids = results["top_ids"]
    top_logits = results["top_logits"]
    n_layers = results["n_layers"]

    if focus_position < 0:
        focus_position = len(tokens) + focus_position

    print(
        f"Logit lens at position {focus_position} (token: {tokens[focus_position]!r})"
    )
    print("  Predicting the token AFTER this position.\n")
    print(f"{'Layer':>6s}  {'Top-1':>8s}  {'Logit':>8s}  {'Top-5'}")
    print("-" * 70)

    for layer_idx in range(n_layers):
        ids = top_ids[layer_idx]
        logits = top_logits[layer_idx]

        pos_ids = ids[focus_position]
        pos_logits = logits[focus_position]

        top1 = tokenizer.decode([pos_ids[0].item()])
        top1_logit = pos_logits[0].item()
        top5 = [tokenizer.decode([pos_ids[k].item()]) for k in range(pos_ids.shape[0])]

        print(f"L{layer_idx:02d}     {top1!r:>8s}  {top1_logit:>8.2f}  {top5}")


def main():
    parser = argparse.ArgumentParser(description="Logit lens via vllm-lens")
    parser.add_argument("--base-url", default="http://localhost:8000")
    parser.add_argument("--prompt", default="The Eiffel Tower is in the city of")
    parser.add_argument(
        "--position",
        type=int,
        default=-1,
        help="Token position to focus on (-1 = last)",
    )
    args = parser.parse_args()

    client = VLLMLensClient(args.base_url)
    results = run_logit_lens(client, args.prompt)
    print_logit_lens(results, client.model, focus_position=args.position)

    torch.save(results, "logit_lens_results.pt")
    print("\nResults saved to logit_lens_results.pt")


if __name__ == "__main__":
    main()