"""Hidden state extraction from transformer models via forward hooks."""

from __future__ import annotations

import torch
import numpy as np

__all__ = ["get_layer_count", "get_candidate_layers", "extract_hidden_states"]


def get_layer_count(model: object) -> int:
    """Return the number of transformer layers in *model*.

    Supports LLaMA/Mistral (``model.model.layers``) and GPT-style
    (``model.transformer.h``) architectures.

    Raises:
        ValueError: If the layer structure is not recognised.
    """
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        return len(model.model.layers)
    if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
        return len(model.transformer.h)
    raise ValueError(
        f"Cannot determine layer count for {type(model).__name__}. "
        "Please specify layer= manually."
    )


def get_candidate_layers(n_layers: int) -> list[int]:
    """Return five evenly spaced layer indices for auto-selection."""
    if n_layers <= 5:
        return list(range(n_layers))
    return [
        0,
        n_layers // 4,
        n_layers // 2,
        3 * n_layers // 4,
        n_layers - 1,
    ]


def _get_transformer_layers(model: object):
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        return model.model.layers
    if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
        return model.transformer.h
    raise ValueError(f"Unsupported architecture: {type(model).__name__}")


def extract_hidden_states(
    model: object,
    tokenizer: object,
    texts: list[str],
    layer_indices: list[int],
    batch_size: int = 4,
) -> dict[int, list[np.ndarray]]:
    """Extract mean-pooled hidden states at the requested layers.

    Registers temporary forward hooks, runs batched inference, and
    returns one numpy vector per (layer, sample) pair.
    """
    device = next(model.parameters()).device
    transformer_layers = _get_transformer_layers(model)
    all_hidden: dict[int, list[np.ndarray]] = {idx: [] for idx in layer_indices}
    captured: dict[int, torch.Tensor] = {}
    hooks = []

    def make_hook(layer_idx: int):
        def hook_fn(module, input, output):
            if isinstance(output, tuple):
                captured[layer_idx] = output[0].detach()
            else:
                captured[layer_idx] = output.detach()
        return hook_fn

    for idx in layer_indices:
        h = transformer_layers[idx].register_forward_hook(make_hook(idx))
        hooks.append(h)

    try:
        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            inputs = tokenizer(
                batch,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=256,
            ).to(device)
            with torch.no_grad():
                model(**inputs)
            for idx in layer_indices:
                h = captured[idx]
                for j in range(len(batch)):
                    mask = inputs["attention_mask"][j].bool()
                    pooled = h[j][mask].float().mean(dim=0).cpu().numpy()
                    all_hidden[idx].append(pooled)
            captured.clear()
    finally:
        for h in hooks:
            h.remove()

    return all_hidden
