"""Cognica-PoE modeling (HF transformers `PreTrainedModel` wrapper).

The underlying GPT architecture is a direct port of the nanochat GPT
(https://github.com/karpathy/nanochat), preserved here verbatim so that
checkpoints trained with nanochat load without reshaping. A thin
`CognicaPoEForCausalLM` wraps it to satisfy the HF API contract required by
`AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)`.

Key design points:
- Flash Attention 3 is dropped; everything routes through PyTorch SDPA so the
  model runs on any recent GPU (Ampere/Hopper/Blackwell), MPS, or CPU.
- KV cache is implemented with the nanochat `(B, T, H, D)` layout, pre-allocated
  at `max_seq_len`. `forward()` accepts `past_key_values` / `use_cache` so HF
  `generate()` runs in O(T) per decode step.
- **PoE stage-level inference** is exposed directly on `CognicaPoEForCausalLM`:
  `forward_stage`, `generate_stage`, `generate_wand`, `generate_speculative`.
  Each intermediate stage boundary (layers 5, 11, 17, 23 at `poe_every=6`) is a
  complete predictor through the shared `lm_head`, so prefix pruning / WAND /
  speculative decoding run out-of-the-box on the released weights.
"""

from dataclasses import dataclass
from types import SimpleNamespace
from typing import Optional, Sequence, Union, Tuple, List, Dict

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

from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.generation.utils import GenerationMixin

from .configuration_cognica_poe import CognicaPoEConfig


# -----------------------------------------------------------------------------
# Minimal compute-dtype policy (nanochat defaults to bf16, fp32 weights)
# -----------------------------------------------------------------------------
COMPUTE_DTYPE = torch.bfloat16


# -----------------------------------------------------------------------------
# KV cache (ported from nanochat.engine.KVCache, SDPA-compatible)
# -----------------------------------------------------------------------------
class CognicaKVCache:
    """KV cache with the nanochat layout: (n_layers, B, T, H, D).

    Per-layer `k_cache` / `v_cache` are pre-allocated to `max_seq_len`; each
    forward pass writes new keys/values in place at `cache_seqlens` and reads
    the prefix used for attention. Smear needs the previous token's pre-smear
    embedding; we keep that on the cache object so the next decode step can use
    it.
    """

    def __init__(self, batch_size, num_heads, max_seq_len, head_dim, num_layers, device, dtype):
        self.batch_size = batch_size
        self.max_seq_len = max_seq_len
        self.n_layers = num_layers
        self.n_heads = num_heads
        self.head_dim = head_dim
        self.k_cache = torch.zeros(num_layers, batch_size, max_seq_len, num_heads, head_dim, device=device, dtype=dtype)
        self.v_cache = torch.zeros(num_layers, batch_size, max_seq_len, num_heads, head_dim, device=device, dtype=dtype)
        self.cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device)
        self.prev_embedding = None

    def reset(self):
        self.cache_seqlens.zero_()
        self.prev_embedding = None

    def get_pos(self):
        return int(self.cache_seqlens[0].item())

    def get_layer_cache(self, layer_idx):
        return self.k_cache[layer_idx], self.v_cache[layer_idx]

    def advance(self, num_tokens):
        self.cache_seqlens += num_tokens

    def reorder(self, beam_idx):
        device = self.k_cache.device
        beam_idx = beam_idx.to(device)
        self.k_cache = self.k_cache.index_select(1, beam_idx)
        self.v_cache = self.v_cache.index_select(1, beam_idx)
        self.cache_seqlens = self.cache_seqlens.index_select(0, beam_idx.to(self.cache_seqlens.device))
        if self.prev_embedding is not None:
            self.prev_embedding = self.prev_embedding.index_select(0, beam_idx.to(self.prev_embedding.device))


# -----------------------------------------------------------------------------
# SDPA-only attention, matching the FA3 call sites used by nanochat
# -----------------------------------------------------------------------------
def _sdpa_attention(q, k, v, window_size, enable_gqa):
    Tq = q.size(2)
    Tk = k.size(2)
    window = window_size[0]

    if (window < 0 or window >= Tq) and Tq == Tk:
        return F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=enable_gqa)

    if Tq == 1:
        if 0 <= window < Tk:
            start = max(0, Tk - (window + 1))
            k = k[:, :, start:, :]
            v = v[:, :, start:, :]
        return F.scaled_dot_product_attention(q, k, v, is_causal=False, enable_gqa=enable_gqa)

    device = q.device
    row_idx = (Tk - Tq) + torch.arange(Tq, device=device).unsqueeze(1)
    col_idx = torch.arange(Tk, device=device).unsqueeze(0)
    mask = col_idx <= row_idx
    if 0 <= window < Tk:
        mask = mask & ((row_idx - col_idx) <= window)
    return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, enable_gqa=enable_gqa)


def _flash_attn_func(q, k, v, causal=False, window_size=(-1, -1)):
    q = q.transpose(1, 2)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)
    enable_gqa = q.size(1) != k.size(1)
    y = _sdpa_attention(q, k, v, window_size, enable_gqa)
    return y.transpose(1, 2)


flash_attn = SimpleNamespace(flash_attn_func=_flash_attn_func)


# -----------------------------------------------------------------------------
# Core nanochat layers (RMSNorm, Linear with dtype cast, rotary, attention, MLP)
# -----------------------------------------------------------------------------
def _norm(x):
    return F.rms_norm(x, (x.size(-1),))


class _Linear(nn.Linear):
    """nn.Linear whose weight is cast to the input dtype at forward time."""

    def forward(self, x):
        return F.linear(x, self.weight.to(dtype=x.dtype))


def _has_ve(layer_idx: int, n_layer: int) -> bool:
    return layer_idx % 2 == (n_layer - 1) % 2


def _apply_rotary_emb(x, cos, sin):
    assert x.ndim == 4
    d = x.shape[3] // 2
    x1, x2 = x[..., :d], x[..., d:]
    y1 = x1 * cos + x2 * sin
    y2 = x1 * (-sin) + x2 * cos
    return torch.cat([y1, y2], 3)


class _CausalSelfAttention(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.layer_idx = layer_idx
        self.n_head = config.n_head
        self.n_kv_head = config.n_kv_head
        self.n_embd = config.n_embd
        self.head_dim = self.n_embd // self.n_head
        assert self.n_embd % self.n_head == 0
        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
        self.c_q = _Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
        self.c_k = _Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
        self.c_v = _Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
        self.c_proj = _Linear(self.n_embd, self.n_embd, bias=False)
        self.ve_gate_channels = 12
        self.ve_gate = _Linear(self.ve_gate_channels, self.n_kv_head, bias=False) \
            if _has_ve(layer_idx, config.n_layer) else None

    def forward(self, x, ve, cos_sin, window_size, kv_cache):
        B, T, C = x.size()
        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)

        if ve is not None:
            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
            v = v + gate.unsqueeze(-1) * ve

        cos, sin = cos_sin
        q, k = _apply_rotary_emb(q, cos, sin), _apply_rotary_emb(k, cos, sin)
        q, k = _norm(q), _norm(k)
        q = q * 1.2
        k = k * 1.2

        if kv_cache is None:
            y = flash_attn.flash_attn_func(q, k, v, causal=True, window_size=window_size)
        else:
            k_cache, v_cache = kv_cache.get_layer_cache(self.layer_idx)
            pos = kv_cache.get_pos()
            k_cache[:, pos:pos + T, :, :] = k.to(k_cache.dtype)
            v_cache[:, pos:pos + T, :, :] = v.to(v_cache.dtype)
            # Cast the cache slice back to q's dtype so SDPA sees a consistent
            # dtype (cache may be stored in a wider precision than activations).
            k_used = k_cache[:, :pos + T, :, :].to(q.dtype)
            v_used = v_cache[:, :pos + T, :, :].to(q.dtype)
            y = flash_attn.flash_attn_func(q, k_used, v_used, causal=True, window_size=window_size)
            if self.layer_idx == kv_cache.n_layers - 1:
                kv_cache.advance(T)

        y = y.contiguous().view(B, T, -1)
        return self.c_proj(y)


class _MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = _Linear(config.n_embd, 4 * config.n_embd, bias=False)
        self.c_proj = _Linear(4 * config.n_embd, config.n_embd, bias=False)

    def forward(self, x):
        x = self.c_fc(x)
        x = F.relu(x).square()
        return self.c_proj(x)


class _Block(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.attn = _CausalSelfAttention(config, layer_idx)
        self.mlp = _MLP(config)

    def forward(self, x, ve, cos_sin, window_size, kv_cache):
        x = x + self.attn(_norm(x), ve, cos_sin, window_size, kv_cache)
        x = x + self.mlp(_norm(x))
        return x


@dataclass
class _GPTConfigShim:
    sequence_len: int
    vocab_size: int
    n_layer: int
    n_head: int
    n_kv_head: int
    n_embd: int
    window_pattern: str
    dual_head: bool = False
    frozen_layers: int = 0


class _GPT(nn.Module):
    """Trimmed copy of nanochat's GPT (inference path, supports KV cache)."""

    def __init__(self, config, pad_vocab_size_to: int = 64):
        super().__init__()
        self.config = config
        self.window_sizes = self._compute_window_sizes(config)
        padded_vocab_size = ((config.vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to) * pad_vocab_size_to
        self.padded_vocab_size = padded_vocab_size

        self.transformer = nn.ModuleDict({
            "wte": nn.Embedding(padded_vocab_size, config.n_embd),
            "h": nn.ModuleList([_Block(config, i) for i in range(config.n_layer)]),
        })
        self.lm_head = _Linear(config.n_embd, padded_vocab_size, bias=False)
        # Paper Section 6.5: additive specialist head, zero-init at training start,
        # summed with the frozen base `lm_head` at the final projection. Present
        # only on stage checkpoints where the parent's lm_head remains frozen.
        self.dual_head = bool(getattr(config, "dual_head", False))
        if self.dual_head:
            self.lm_head_stage = _Linear(config.n_embd, padded_vocab_size, bias=False)
        else:
            self.lm_head_stage = None

        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
        self.smear_gate = _Linear(24, 1, bias=False)
        self.smear_lambda = nn.Parameter(torch.zeros(1))
        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))

        head_dim = config.n_embd // config.n_head
        kv_dim = config.n_kv_head * head_dim
        # VE pattern preservation: on stage checkpoints we freeze the parent's
        # layers and add `new_layers` fresh ones. The parent was trained with a
        # VE pattern computed against its own depth (`frozen_layers`), and the
        # paper Section 6.5 training convention keeps those exact VE indices —
        # no new VEs are attached to the stage's new layers. Use
        # `frozen_layers` when present so the VE ModuleDict matches what the
        # delta actually shipped; fall back to `n_layer` for leaf (base) models.
        ve_depth = int(getattr(config, "frozen_layers", 0)) or config.n_layer
        self.value_embeds = nn.ModuleDict({
            str(i): nn.Embedding(padded_vocab_size, kv_dim)
            for i in range(ve_depth) if _has_ve(i, ve_depth)
        })

        # Rotary cache over-provisioned 10x; matches nanochat.
        self.rotary_seq_len = config.sequence_len * 10
        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

    @staticmethod
    def _precompute_rotary_embeddings(seq_len, head_dim, base=100000, device=None, dtype=COMPUTE_DTYPE):
        if device is None:
            device = torch.device("cpu")
        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
        inv_freq = 1.0 / (base ** (channel_range / head_dim))
        t = torch.arange(seq_len, dtype=torch.float32, device=device)
        freqs = torch.outer(t, inv_freq)
        cos, sin = freqs.cos().to(dtype), freqs.sin().to(dtype)
        return cos[None, :, None, :], sin[None, :, None, :]

    def _compute_window_sizes(self, config):
        pattern = config.window_pattern.upper()
        assert all(c in "SL" for c in pattern)
        long_window = config.sequence_len
        short_window = -(-long_window // 4 // 128) * 128
        char_to_window = {"L": (long_window, 0), "S": (short_window, 0)}
        window_sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
        window_sizes[-1] = (long_window, 0)
        return window_sizes

    def make_kv_cache(self, batch_size, max_seq_len, device, dtype=None):
        """Allocate a fresh KV cache matching this model's dims."""
        if dtype is None:
            dtype = COMPUTE_DTYPE if device.type == "cuda" else torch.float32
        head_dim = self.config.n_embd // self.config.n_head
        return CognicaKVCache(
            batch_size=batch_size,
            num_heads=self.config.n_kv_head,
            max_seq_len=max_seq_len,
            head_dim=head_dim,
            num_layers=self.config.n_layer,
            device=device,
            dtype=dtype,
        )

    def _apply_lm_head(self, x: torch.Tensor, use_stage_head: bool = True) -> torch.Tensor:
        """Project a normed hidden state through the shared lm_head with softcap.

        When the model carries a specialist head (`self.lm_head_stage`) and
        `use_stage_head=True`, its output is added to the base head before
        softcapping, per paper Section 6.5 composition `logits = base + stage`.
        Intermediate PoE stage logits pass `use_stage_head=False` so they stay
        on the base path and remain comparable to pre-SFT stage predictions.
        """
        softcap = 15
        logits = self.lm_head(x)
        if use_stage_head and self.lm_head_stage is not None:
            logits = logits + self.lm_head_stage(x)
        logits = logits[..., :self.config.vocab_size].float()
        return softcap * torch.tanh(logits / softcap)

    def forward(self, idx: torch.Tensor, kv_cache: Optional[CognicaKVCache] = None) -> torch.Tensor:
        B, T = idx.size()
        assert T <= self.cos.size(1), f"Sequence length {T} exceeds rotary cache {self.cos.size(1)}"

        T0 = 0 if kv_cache is None else kv_cache.get_pos()
        cos_sin = self.cos[:, T0:T0 + T], self.sin[:, T0:T0 + T]

        x = self.transformer.wte(idx).to(COMPUTE_DTYPE)
        x = _norm(x)

        # Smear: bigram-style mix of previous token's pre-smear embedding.
        if kv_cache is None:
            if T > 1:
                gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
                x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
        else:
            x_pre_smear = kv_cache.prev_embedding
            kv_cache.prev_embedding = x[:, -1:, :].detach().clone()
            if T > 1:
                gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
                x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
            elif x_pre_smear is not None:
                gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, :, :24]))
                x = x + gate * x_pre_smear

        x0 = x
        n_layer = self.config.n_layer
        backout_layer = n_layer // 2
        x_backout = None
        for i, block in enumerate(self.transformer.h):
            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
            ve = self.value_embeds[str(i)](idx).to(x.dtype) if str(i) in self.value_embeds else None
            x = block(x, ve, cos_sin, self.window_sizes[i], kv_cache)
            if i == backout_layer:
                x_backout = x

        if x_backout is not None:
            x = x - self.backout_lambda.to(x.dtype) * x_backout
        x = _norm(x)
        return self._apply_lm_head(x)

    def forward_all_stages(
        self,
        idx: torch.Tensor,
        stage_boundaries: List[int],
    ) -> List[torch.Tensor]:
        """Single forward pass emitting stage-level logits at each boundary layer.

        `stage_boundaries` are layer indices (0-indexed, inclusive) at which to
        snapshot the residual stream. For d24 / poe_every=6 the natural choice is
        `[5, 11, 17, 23]` (the last layer of each of the 4 PoE stages).

        Each stage logit uses `lm_head(norm(x))` — matching the PoE per-stage loss
        path (no backout residual subtraction). The final PoE stage (k=n-1) is
        therefore slightly different from the backout-corrected `forward()` path,
        but the difference is tiny when `backout_lambda` ~ 0 (typical for
        PoE-trained checkpoints). The per-stage top-1 agreement measurements in
        the paper use this same convention.

        Returns a list of `(B, T, vocab_size)` float32 logit tensors, one per
        requested boundary, in the order given.
        """
        B, T = idx.size()
        assert T <= self.cos.size(1), f"Sequence length {T} exceeds rotary cache {self.cos.size(1)}"
        boundaries = sorted(set(stage_boundaries))
        assert all(0 <= b < self.config.n_layer for b in boundaries)

        cos_sin = self.cos[:, :T], self.sin[:, :T]

        x = self.transformer.wte(idx).to(COMPUTE_DTYPE)
        x = _norm(x)
        if T > 1:
            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)

        x0 = x
        stage_logits: Dict[int, torch.Tensor] = {}
        for i, block in enumerate(self.transformer.h):
            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
            ve = self.value_embeds[str(i)](idx).to(x.dtype) if str(i) in self.value_embeds else None
            x = block(x, ve, cos_sin, self.window_sizes[i], None)
            if i in boundaries:
                stage_logits[i] = self._apply_lm_head(_norm(x))

        return [stage_logits[b] for b in stage_boundaries]


# -----------------------------------------------------------------------------
# HF-facing wrappers
# -----------------------------------------------------------------------------
class CognicaPoEPreTrainedModel(PreTrainedModel):
    config_class = CognicaPoEConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False

    def _init_weights(self, module):
        # Weights come from the trained checkpoint - do not reinitialize here.
        return


class CognicaPoEForCausalLM(CognicaPoEPreTrainedModel, GenerationMixin):
    """HuggingFace-compatible causal LM wrapping the nanochat-port GPT."""

    def __init__(self, config: CognicaPoEConfig):
        super().__init__(config)
        gpt_config = _GPTConfigShim(
            sequence_len=config.max_position_embeddings,
            vocab_size=config.vocab_size,
            n_layer=config.num_hidden_layers,
            n_head=config.num_attention_heads,
            n_kv_head=config.num_key_value_heads,
            n_embd=config.hidden_size,
            window_pattern=config.window_pattern,
            dual_head=bool(getattr(config, "dual_head", False)),
            frozen_layers=int(getattr(config, "frozen_layers", 0)),
        )
        self.gpt = _GPT(gpt_config, pad_vocab_size_to=64)

    # -------------------------------------------------------------------------
    # Cascade loader for stage checkpoints (Paper Section 8.8 Elastic Depth).
    #
    # Stage repos carry only the delta tensors the stage itself trained (new
    # layers, `lm_head_stage`, warm-init `wte`, the stage's residual lambdas,
    # and any new value-embed rows), plus a `config.base_model_name_or_path`
    # pointing at the parent stage (or the base). `from_pretrained` here
    # recursively loads the parent, instantiates an architecture with the
    # extended `num_hidden_layers`, copies the parent's weights into the
    # shared prefix, then overlays the stage delta.
    # -------------------------------------------------------------------------
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        from transformers import AutoConfig

        trust_remote_code = kwargs.get("trust_remote_code", True)
        config = kwargs.pop("config", None)
        if config is None:
            config = AutoConfig.from_pretrained(
                pretrained_model_name_or_path,
                trust_remote_code=trust_remote_code,
            )

        base_ref = getattr(config, "base_model_name_or_path", None)
        if not base_ref:
            # This IS the base (leaf). Standard HF load path.
            return super().from_pretrained(
                pretrained_model_name_or_path, *model_args, config=config, **kwargs
            )

        # This is a stage. Recursively load the parent first (may itself be a stage).
        parent_model = cls.from_pretrained(base_ref, trust_remote_code=trust_remote_code)

        # Instantiate the extended architecture. `config.num_hidden_layers`
        # already includes both the parent's layers and this stage's `new_layers`.
        model = cls(config)

        # Copy parent weights into the shared prefix of the extended model.
        # 1-D tensors of different length (resid_lambdas, x0_lambdas) are
        # prefix-copied; the tail is left at the extended model's init values
        # and then overwritten by the stage delta below.
        parent_state = dict(parent_model.state_dict())
        # Multi-stage composition: if the parent is itself a trained stage, fold
        # its lm_head_stage into lm_head_base so each ancestor's specialist head
        # accumulates additively into the final projection.
        parent_stage_head_key = "gpt.lm_head_stage.weight"
        parent_base_head_key = "gpt.lm_head.weight"
        if parent_stage_head_key in parent_state and parent_base_head_key in parent_state:
            parent_stage_head = parent_state[parent_stage_head_key]
            if parent_stage_head.abs().sum() > 0:
                parent_state[parent_base_head_key] = parent_state[parent_base_head_key] + parent_stage_head
                parent_state[parent_stage_head_key] = torch.zeros_like(parent_stage_head)
        own_state = model.state_dict()
        for key, parent_val in parent_state.items():
            if key not in own_state:
                continue
            target = own_state[key]
            if target.shape == parent_val.shape:
                target.copy_(parent_val)
            elif target.dim() == 1 and parent_val.dim() == 1 and parent_val.shape[0] < target.shape[0]:
                target[: parent_val.shape[0]].copy_(parent_val)
            # Any other shape mismatch is left for the stage delta to fill.

        # Resolve and load the stage delta (delta.safetensors or model.safetensors).
        delta_state = cls._load_stage_delta(pretrained_model_name_or_path, **kwargs)
        missing, unexpected = model.load_state_dict(delta_state, strict=False)
        if unexpected:
            raise RuntimeError(
                f"Unexpected keys in stage delta for {pretrained_model_name_or_path}: "
                f"{unexpected[:5]}{'...' if len(unexpected) > 5 else ''}"
            )

        # Move to requested dtype/device if asked. HF's standard from_pretrained
        # handles this via `torch_dtype` / `device_map`; we mirror the minimum.
        torch_dtype = kwargs.get("torch_dtype", None)
        if torch_dtype is not None:
            model.to(dtype=torch_dtype)
        return model

    @staticmethod
    def _load_stage_delta(pretrained_model_name_or_path, **kwargs):
        """Locate `delta.safetensors` (preferred) or `model.safetensors` on disk
        or on the Hub, and return its state dict."""
        import os
        from safetensors.torch import load_file

        candidates = ("delta.safetensors", "model.safetensors")
        if os.path.isdir(pretrained_model_name_or_path):
            for name in candidates:
                path = os.path.join(pretrained_model_name_or_path, name)
                if os.path.isfile(path):
                    return load_file(path)
            raise FileNotFoundError(
                f"No stage delta file found in {pretrained_model_name_or_path}. "
                f"Expected one of: {candidates}"
            )
        # Hub path
        from huggingface_hub import hf_hub_download
        from huggingface_hub.errors import EntryNotFoundError
        revision = kwargs.get("revision", None)
        token = kwargs.get("token", kwargs.get("use_auth_token", None))
        for name in candidates:
            try:
                path = hf_hub_download(
                    repo_id=pretrained_model_name_or_path,
                    filename=name,
                    revision=revision,
                    token=token,
                )
                return load_file(path)
            except EntryNotFoundError:
                continue
        raise FileNotFoundError(
            f"No stage delta file found in repo {pretrained_model_name_or_path}. "
            f"Expected one of: {candidates}"
        )

    def get_input_embeddings(self):
        return self.gpt.transformer.wte

    def set_input_embeddings(self, value):
        self.gpt.transformer.wte = value

    def get_output_embeddings(self):
        return self.gpt.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.gpt.lm_head = new_embeddings

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[CognicaKVCache] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        return_dict: Optional[bool] = True,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        kv_cache = past_key_values
        if use_cache and kv_cache is None:
            dtype = self.gpt.transformer.wte.weight.dtype
            if dtype not in (torch.bfloat16, torch.float16, torch.float32):
                dtype = COMPUTE_DTYPE if input_ids.device.type == "cuda" else torch.float32
            kv_cache = self.gpt.make_kv_cache(
                batch_size=input_ids.size(0),
                max_seq_len=self.config.max_position_embeddings,
                device=input_ids.device,
                dtype=dtype,
            )

        logits = self.gpt(input_ids, kv_cache=kv_cache)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        if not return_dict:
            output = (logits,)
            if use_cache:
                output = output + (kv_cache,)
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=kv_cache if use_cache else None,
            hidden_states=None,
            attentions=None,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values: Optional[CognicaKVCache] = None,
        attention_mask: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = True,
        **kwargs,
    ):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
            "use_cache": use_cache,
        }

    def _reorder_cache(self, past_key_values, beam_idx):
        if past_key_values is None:
            return None
        past_key_values.reorder(beam_idx)
        return past_key_values

    # -------------------------------------------------------------------------
    # PoE stage-level inference
    #
    # During training each stage (poe_every layers) was trained to produce
    # valid next-token logits via the shared lm_head. The methods below expose
    # that directly: prefix pruning, WAND adaptive depth, speculative decoding.
    # No retraining or added parameters — the released weights already carry
    # the per-stage prediction capability.
    #
    # Note: these methods re-forward the full prefix each decode step (no KV
    # cache) for simplicity. Wall-clock speedups come from the reduced layer
    # count per forward; production decode loops can add stage-aware KV caching
    # on top of these primitives.
    # -------------------------------------------------------------------------

    # Paper §5.3 calibrated p99 |ΔL|_∞ bounds for WAND, d24 PoE r=10:
    # transitions stage 0→1, 1→2, 2→3.
    POE_WAND_P99_BOUNDS: Sequence[float] = (7.09, 3.03, 2.15)

    @property
    def poe_n_stages(self) -> int:
        """Number of PoE stages = num_hidden_layers / poe_every."""
        return self.config.num_hidden_layers // self.config.poe_every

    @property
    def poe_stage_boundary_layers(self) -> List[int]:
        """Last layer index (0-indexed, inclusive) of each stage.
        For d24 / poe_every=6: [5, 11, 17, 23].
        """
        pe = self.config.poe_every
        return [pe * (k + 1) - 1 for k in range(self.poe_n_stages)]

    def forward_stage(
        self,
        input_ids: torch.LongTensor,
        stage: int,
    ) -> torch.Tensor:
        """Logits at the end of PoE stage `stage` (0..n_stages-1).

        Stage 0 runs `poe_every` layers (~25 % compute at `poe_every=6, n_layer=24`)
        and achieves ~87.5 % factual accuracy of the full model on the paper's
        8-prompt probe. Stage `n_stages-1` forwards all layers.

        Returns: `(B, T, vocab_size)` float32 logits.
        """
        n = self.poe_n_stages
        if not 0 <= stage < n:
            raise ValueError(f"stage must be in [0, {n}); got {stage}")
        boundary = self.poe_stage_boundary_layers[stage]
        [logits] = self.gpt.forward_all_stages(input_ids, [boundary])
        return logits

    @torch.no_grad()
    def generate_stage(
        self,
        input_ids: torch.LongTensor,
        stage: int,
        max_new_tokens: int = 32,
        do_sample: bool = False,
        temperature: float = 1.0,
        top_p: Optional[float] = None,
    ) -> torch.LongTensor:
        """Generate using only PoE stage `stage` (stage-prefix pruning).

        Example (25 % compute using Stage 0):
            out = model.generate_stage(ids, stage=0, max_new_tokens=64)

        Returns the full sequence (prompt + generated tokens).
        """
        ids = input_ids
        for _ in range(max_new_tokens):
            logits = self.forward_stage(ids, stage=stage)
            next_ids = _sample_next(logits[:, -1, :], do_sample, temperature, top_p)
            ids = torch.cat([ids, next_ids], dim=1)
        return ids

    @torch.no_grad()
    def generate_wand(
        self,
        input_ids: torch.LongTensor,
        max_new_tokens: int = 32,
        p99_bounds: Optional[Sequence[float]] = None,
        safety: float = 1.0,
        do_sample: bool = False,
        temperature: float = 1.0,
        top_p: Optional[float] = None,
        return_stages_used: bool = False,
    ):
        """WAND-style adaptive stage pruning (Jeong 2026 §5.3).

        At each PoE stage boundary, if the current top-1 margin exceeds
        `safety × Σ p99_bounds[stage:]`, no later stage can change the top-1
        token — the prediction is safe to emit immediately. Otherwise the
        next stage is consulted. Paper reports 1.82× wall-clock @ 100 %
        top-1 agreement over 18 probes with `safety=1.0` on d24 r=10.

        Args:
            input_ids: prompt tokens.
            max_new_tokens: generation budget.
            p99_bounds: length `n_stages - 1`; the p99 `|ΔL|_∞` from stage k
                to stage k+1. Defaults to paper calibration (7.09, 3.03, 2.15).
            safety: multiplier on the p99 budget; 1.0 gives 100 % top-1 match
                on the paper's probes. Higher `safety` is more conservative.
            return_stages_used: if True, also return a list of the stage index
                used for each emitted token.

        Returns:
            ids: `(B, prompt_len + max_new_tokens)` token tensor.
            (optionally) stages_used: list of int.
        """
        n_stages = self.poe_n_stages
        bounds = list(p99_bounds if p99_bounds is not None else self.POE_WAND_P99_BOUNDS)
        if len(bounds) != n_stages - 1:
            raise ValueError(
                f"p99_bounds must have length n_stages-1 = {n_stages - 1}; got {len(bounds)}"
            )

        boundaries = self.poe_stage_boundary_layers
        ids = input_ids
        stages_used: List[int] = []

        for _ in range(max_new_tokens):
            stage_logits = self.gpt.forward_all_stages(ids, boundaries)
            emitted = False
            for k, sl in enumerate(stage_logits):
                last = sl[:, -1, :]
                top2 = torch.topk(last, 2, dim=-1)
                margin = top2.values[:, 0] - top2.values[:, 1]
                remaining = safety * sum(bounds[k:])
                if bool((margin > remaining).all().item()):
                    next_ids = _sample_next(last, do_sample, temperature, top_p)
                    ids = torch.cat([ids, next_ids], dim=1)
                    stages_used.append(k)
                    emitted = True
                    break
            if not emitted:
                last = stage_logits[-1][:, -1, :]
                next_ids = _sample_next(last, do_sample, temperature, top_p)
                ids = torch.cat([ids, next_ids], dim=1)
                stages_used.append(n_stages - 1)

        return (ids, stages_used) if return_stages_used else ids

    @torch.no_grad()
    def generate_parallel_composition(
        self,
        input_ids: torch.LongTensor,
        stages: Optional[Sequence[int]] = None,
        stage_weights: Optional[Sequence[float]] = None,
        max_new_tokens: int = 32,
        do_sample: bool = False,
        temperature: float = 1.0,
        top_p: Optional[float] = None,
    ) -> torch.LongTensor:
        """PoE log-space parallel stage composition (Jeong 2026 §6.5.5).

        Each decode step runs one forward pass and combines the logits from
        multiple stage boundaries additively (Log-OP / Product of Experts
        algebra). Paper reports combining `{[1..4], [1..5]}` strengthens
        factual-retrieval margins by +2.4 logit units over the strongest
        single-branch — a quality-positive inference mode with zero retraining.

        On a single GPU this is one forward pass with 4 logit projections; on
        multiple devices the stage branches can be dispatched in parallel
        (`T ≈ T_S1 + max(T_S2..T_Sn)` per paper §4.3) but this reference
        implementation runs on a single device.

        Args:
            stages: stage indices to compose (0..n_stages-1). Default: all stages.
            stage_weights: per-stage multipliers in log-space. Default: uniform 1.0.
                Paper §6.5.6 shows branch weighting as an inference-time tuning axis
                between absolute confidence and margin robustness.
        """
        n_stages = self.poe_n_stages
        if stages is None:
            stages = list(range(n_stages))
        stages = list(stages)
        if stage_weights is None:
            stage_weights = [1.0] * len(stages)
        if len(stage_weights) != len(stages):
            raise ValueError("stage_weights must match len(stages)")
        if not all(0 <= s < n_stages for s in stages):
            raise ValueError(f"all stages must be in [0, {n_stages})")

        boundaries = [self.poe_stage_boundary_layers[s] for s in stages]
        ids = input_ids
        for _ in range(max_new_tokens):
            stage_logits = self.gpt.forward_all_stages(ids, boundaries)
            # Log-OP: weighted sum in logit space. Each stage's softcapped logits
            # are already in a calibrated scale, so simple weighted summation
            # implements the PoE product of distributions.
            combined = sum(w * sl[:, -1, :] for w, sl in zip(stage_weights, stage_logits))
            next_ids = _sample_next(combined, do_sample, temperature, top_p)
            ids = torch.cat([ids, next_ids], dim=1)
        return ids

    @torch.no_grad()
    def generate_speculative(
        self,
        input_ids: torch.LongTensor,
        max_new_tokens: int = 32,
        draft_stage: int = 0,
        k_draft: int = 3,
        return_acceptance: bool = False,
    ):
        """Speculative decoding with PoE stage `draft_stage` as natural drafter.

        (greedy only; matches Jeong 2026 §5.4 protocol.) The drafter is free —
        Stage 0 is already a valid predictor at 25 % compute after PoE training.
        Paper reports 1.87× speedup @ 88 % acceptance with K=3 on d24 r=10.

        Procedure: draft `k_draft` tokens with `draft_stage`, verify in parallel
        with the full forward, accept left-to-right up to the first mismatch,
        then replace the mismatch position with the full-path token and continue.

        Args:
            draft_stage: which stage drafts (must be < n_stages - 1).
            k_draft: tokens drafted per cycle.
            return_acceptance: if True, also return the mean per-cycle acceptance rate.
        """
        n_stages = self.poe_n_stages
        if not 0 <= draft_stage < n_stages - 1:
            raise ValueError(f"draft_stage must be in [0, {n_stages - 1}); got {draft_stage}")

        ids = input_ids
        prompt_len = input_ids.size(1)
        accepts_total = 0
        drafts_total = 0

        while ids.size(1) - prompt_len < max_new_tokens:
            drafted = ids
            for _ in range(k_draft):
                draft_logits = self.forward_stage(drafted, stage=draft_stage)
                nxt = draft_logits[:, -1, :].argmax(dim=-1, keepdim=True)
                drafted = torch.cat([drafted, nxt], dim=1)
            draft_tokens = drafted[:, ids.size(1):]

            full_logits = self.forward_stage(drafted, stage=n_stages - 1)
            verify_pred = full_logits[:, ids.size(1) - 1 : drafted.size(1), :].argmax(dim=-1)

            accepted = 0
            for k in range(k_draft):
                if bool((verify_pred[:, k] == draft_tokens[:, k]).all().item()):
                    accepted += 1
                else:
                    break
            drafts_total += k_draft
            accepts_total += accepted

            ids = torch.cat(
                [
                    ids,
                    draft_tokens[:, :accepted],
                    verify_pred[:, accepted : accepted + 1],
                ],
                dim=1,
            )

        ids = ids[:, : prompt_len + max_new_tokens]
        if return_acceptance:
            rate = accepts_total / drafts_total if drafts_total else 0.0
            return ids, rate
        return ids


# -----------------------------------------------------------------------------
# Sampling helper (used by stage/WAND/speculative generation)
# -----------------------------------------------------------------------------
def _sample_next(
    logits: torch.Tensor,
    do_sample: bool,
    temperature: float,
    top_p: Optional[float],
) -> torch.LongTensor:
    """Greedy (do_sample=False) or nucleus sampling from `(B, V)` logits.
    Returns `(B, 1)` token ids."""
    if not do_sample:
        return logits.argmax(dim=-1, keepdim=True)
    logits = logits / max(temperature, 1e-6)
    if top_p is not None and 0.0 < top_p < 1.0:
        sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True)
        probs = F.softmax(sorted_logits, dim=-1)
        cum = probs.cumsum(dim=-1)
        mask = cum > top_p
        mask[..., 0] = False
        sorted_logits = sorted_logits.masked_fill(mask, float("-inf"))
        logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_idx, sorted_logits)
    probs = F.softmax(logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)


__all__ = [
    "CognicaPoEConfig",
    "CognicaPoEForCausalLM",
    "CognicaPoEPreTrainedModel",
    "CognicaKVCache",
]