"""
A_lin: the Linear Accessibility Profile metric from the LAP paper.

A_lin(l) = fraction of prompts where argmax(W_U · LayerNorm(h_l)) == target token.

This is the logit lens applied to intermediate hidden states — no training required.
It measures how "output-aligned" the concept is at each layer.
"""

import numpy as np
import torch
from typing import Dict, List, Tuple

from steering import get_final_norm, get_lm_head, get_layers, resolve_target_token


def compute_alin(
    model,
    tokenizer,
    prompts: List[str],
    target_token: str,
    device: str = "cpu",
) -> List[float]:
    """
    Compute logit-lens accuracy at every layer for `prompts` with `target_token`.

    For each prompt and each layer l:
      - Take hidden state h_l at the last token position
      - Apply final LayerNorm + LM head  (the model's own unembedding projection)
      - Check if argmax == target_token_id
    A_lin(l) = fraction of prompts that pass the check.

    This matches the paper's Eq. (1) exactly. No training required.
    """
    target_id = resolve_target_token(tokenizer, target_token)
    print(f"  Target '{target_token}' → token id {target_id} "
          f"('{tokenizer.decode([target_id])}')")

    final_norm = get_final_norm(model)
    lm_head = get_lm_head(model)
    layers = get_layers(model)
    n_layers = len(layers)

    hits = [0] * n_layers
    model.eval()

    with torch.no_grad():
        for prompt in prompts:
            buf: Dict[int, torch.Tensor] = {}

            def make_hook(idx):
                def hook(module, inp, out):
                    h = out[0] if isinstance(out, tuple) else out
                    buf[idx] = h[0, -1, :].detach()
                return hook

            handles = [layers[i].register_forward_hook(make_hook(i)) for i in range(n_layers)]
            try:
                inputs = tokenizer(
                    prompt, return_tensors="pt", truncation=True, max_length=128
                ).to(device)
                model(**inputs)
            finally:
                for h in handles:
                    h.remove()

            for i in range(n_layers):
                h = buf[i].to(device)
                h_normed = final_norm(h.unsqueeze(0)).squeeze(0)
                logits = lm_head(h_normed)
                if logits.argmax().item() == target_id:
                    hits[i] += 1

    return [h / len(prompts) for h in hits]


# ---------------------------------------------------------------------------
# Layer selection
# ---------------------------------------------------------------------------

def select_optimal_layer(alin: List[float]) -> int:
    """
    Layer with highest A_lin, excluding the final layer.

    The last layer sits immediately before the LM head; A_lin there is
    trivially high (it's what the model actually predicts) but steering
    at that point disrupts generation rather than redirecting it.
    """
    n = len(alin)
    candidates = list(range(n - 1))  # exclude last layer
    return candidates[int(np.argmax([alin[i] for i in candidates]))]


def select_comparison_layer(n_layers: int) -> int:
    """
    Return the middle layer — the standard practitioner heuristic.

    The paper explicitly contrasts LAP-recommended vs middle-layer steering,
    showing that the middle layer fails for entity redirect (A_lin = 0 there).
    """
    return n_layers // 2