"""Canonical transformer trace sites used by the generic tracer."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

SiteName = Literal[
    "block_input",
    "attn_o_proj_pre",
    "attn_output",
    "post_attn_residual",
    "mlp_input",
    "mlp_output",
    "block_output",
    "final_norm",
    "lm_head_input",
    "logits",
]


@dataclass(frozen=True)
class SiteSpec:
    """Metadata describing where a site sits in the transformer stack."""

    name: SiteName
    layer_scoped: bool
    synthetic: bool
    description: str


SITE_SPECS: dict[SiteName, SiteSpec] = {
    "block_input": SiteSpec(
        name="block_input",
        layer_scoped=True,
        synthetic=False,
        description="Residual stream entering a transformer block.",
    ),
    "attn_o_proj_pre": SiteSpec(
        name="attn_o_proj_pre",
        layer_scoped=True,
        synthetic=False,
        description="Input to the attention output projection before heads are mixed.",
    ),
    "attn_output": SiteSpec(
        name="attn_output",
        layer_scoped=True,
        synthetic=False,
        description="Output of the attention module before residual addition.",
    ),
    "post_attn_residual": SiteSpec(
        name="post_attn_residual",
        layer_scoped=True,
        synthetic=True,
        description="Residual stream immediately after the attention write.",
    ),
    "mlp_input": SiteSpec(
        name="mlp_input",
        layer_scoped=True,
        synthetic=False,
        description="Input to the MLP/feed-forward submodule.",
    ),
    "mlp_output": SiteSpec(
        name="mlp_output",
        layer_scoped=True,
        synthetic=False,
        description="Output of the MLP/feed-forward submodule before residual addition.",
    ),
    "block_output": SiteSpec(
        name="block_output",
        layer_scoped=True,
        synthetic=False,
        description="Residual stream leaving a transformer block.",
    ),
    "final_norm": SiteSpec(
        name="final_norm",
        layer_scoped=False,
        synthetic=False,
        description="Output of the model's final normalization before the LM head.",
    ),
    "lm_head_input": SiteSpec(
        name="lm_head_input",
        layer_scoped=False,
        synthetic=False,
        description="Input to the LM head projection.",
    ),
    "logits": SiteSpec(
        name="logits",
        layer_scoped=True,
        synthetic=True,
        description="Per-layer logits projected from intermediate hidden states.",
    ),
}

ALL_SITES = tuple(SITE_SPECS)
LAYER_SITES = tuple(name for name, spec in SITE_SPECS.items() if spec.layer_scoped)
GLOBAL_SITES = tuple(name for name, spec in SITE_SPECS.items() if not spec.layer_scoped)

