"""Early-exit heads: full-vocabulary heads and pruned heads with a reject class.

A pruned head k predicts over its vocabulary V_k plus one extra output, the
reject class (⊥). Its training target is the full-model distribution collapsed
onto V_k ∪ {⊥}: p̄(v) = p(v) for v ∈ V_k and p̄(⊥) = p(V \\ V_k).
"""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * x.to(dtype)


class Adapter(nn.Module):
    """Residual low-rank MLP before an early head (Kangaroo/EAGLE-style, no attention,
    hence no extra KV cache): h + up(silu(down(norm(h)))), identity at initialisation."""

    def __init__(self, hidden_size: int, rank: int, eps: float):
        super().__init__()
        self.norm = RMSNorm(hidden_size, eps)
        self.down = nn.Linear(hidden_size, rank, bias=False)
        self.up = nn.Linear(rank, hidden_size, bias=False)
        nn.init.zeros_(self.up.weight)

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        return h + self.up(F.silu(self.down(self.norm(h))))


class ExitHead(nn.Module):
    """Norm + linear projection onto `vocab_ids` (+ optional reject class).

    vocab_ids=None means the full vocabulary (output index == token id).
    """

    def __init__(self, hidden_size: int, vocab_size: int, eps: float,
                 vocab_ids: torch.Tensor | None = None, reject: bool = False, adapter_rank: int = 0):
        super().__init__()
        self.adapter = Adapter(hidden_size, adapter_rank, eps) if adapter_rank > 0 else None
        self.norm = RMSNorm(hidden_size, eps)
        self.vocab_size = vocab_size
        n_out = vocab_size if vocab_ids is None else int(vocab_ids.numel())
        self.reject = reject
        self.proj = nn.Linear(hidden_size, n_out + int(reject), bias=False)
        # The ⊥ logit gets its own scalar bias (initialised to the log of the
        # mean out-of-vocabulary mass when pruning).
        self.reject_bias = nn.Parameter(torch.zeros(1)) if reject else None
        if vocab_ids is None:
            self.register_buffer("vocab_ids", None)
        else:
            self.register_buffer("vocab_ids", vocab_ids.long().clone())

    @property
    def n_vocab(self) -> int:
        return self.proj.out_features - int(self.reject)

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        """Logits over V_k (∪ ⊥ as the last column), float32."""
        if self.adapter is not None:
            h = self.adapter(h)
        logits = self.proj(self.norm(h)).float()
        if self.reject:
            logits = torch.cat([logits[..., :-1], logits[..., -1:] + self.reject_bias.float()], dim=-1)
        return logits

    def to_token_ids(self, idx: torch.Tensor) -> torch.Tensor:
        # gather (not advanced indexing) so that it stays capturable in CUDA/HIP graphs
        if self.vocab_ids is None:
            return idx
        return self.vocab_ids.gather(0, idx.reshape(-1)).reshape(idx.shape)

    def scatter_full(self, probs_k: torch.Tensor) -> torch.Tensor:
        """Map probabilities over V_k (no ⊥ column) to the full vocabulary."""
        if self.vocab_ids is None:
            return probs_k
        out = probs_k.new_zeros(*probs_k.shape[:-1], self.vocab_size)
        out[..., self.vocab_ids] = probs_k
        return out

    @classmethod
    def from_full(cls, full: "ExitHead", vocab_ids: torch.Tensor, reject: bool = True,
                  reject_bias: float = 0.0, adapter_rank: int = 0) -> "ExitHead":
        """Prune a full head to vocab_ids, appending a zero-weight ⊥ row (optional adapter)."""
        assert full.vocab_ids is None
        d = full.proj.in_features
        head = cls(d, full.vocab_size, full.norm.eps, vocab_ids, reject, adapter_rank)
        with torch.no_grad():
            head.norm.weight.copy_(full.norm.weight)
            head.proj.weight[: vocab_ids.numel()].copy_(full.proj.weight[vocab_ids])
            if reject:
                head.proj.weight[-1].zero_()
                head.reject_bias.fill_(reject_bias)
        return head


def collapse_target(p_full: torch.Tensor, vocab_ids: torch.Tensor) -> torch.Tensor:
    """p̄ over V_k ∪ {⊥} from full-vocabulary probabilities (…, V)."""
    kept = p_full[..., vocab_ids]
    rest = (1.0 - kept.sum(-1, keepdim=True)).clamp_min(0.0)
    return torch.cat([kept, rest], dim=-1)


def head_draft_distribution(logits: torch.Tensor, reject: bool, temperature: float):
    """From head logits return (q over V_k renormalised, reject prob, confidence).

    The confidence is the max probability over V_k ∪ {⊥} (⊥ included in the
    softmax), so a head that puts its mass on ⊥ is not confident. At T=0 the
    returned q is one-hot on the argmax over V_k.
    """
    probs = F.softmax(logits, dim=-1)
    if reject:
        p_rej = probs[..., -1]
        conf, arg = probs.max(-1)
        is_rej = arg == probs.shape[-1] - 1
        conf = torch.where(is_rej, torch.zeros_like(conf), conf)
        body_logits = logits[..., :-1]
    else:
        p_rej = torch.zeros(probs.shape[:-1], device=probs.device)
        conf = probs.max(-1).values
        body_logits = logits
    if temperature == 0.0:
        q = F.one_hot(body_logits.argmax(-1), body_logits.shape[-1]).float()
    else:
        q = F.softmax(body_logits / temperature, dim=-1)
    return q, p_rej, conf


def entropy(logits: torch.Tensor) -> torch.Tensor:
    logp = F.log_softmax(logits, dim=-1)
    return -(logp.exp() * logp).sum(-1)