"""Byte-level decoder language models.

`variant="plain"` is the P1 small Llama-style causal decoder over raw UTF-8
bytes. `variant="spacebyte"` is the P2 word-boundary hierarchy: cheap local
byte blocks run everywhere, expensive global blocks run only on patch-boundary
positions, then the latest patch context is scattered back to every byte before
optional local output blocks.

forward() returns (logits, loss, parts) where `parts` breaks the loss into the
cross-entropy term and the z-loss term for logging.
"""
from __future__ import annotations

import math
from typing import Dict, Optional, Tuple

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

from .config import BOS_ID, EOS_ID, VOCAB_SIZE, ByteLMConfig
from .layers import Block, RMSNorm, RotaryEmbedding, build_doc_attn_mask


# ASCII whitespace bytes that delimit Kashmiri words in the normalized corpus.
# UTF-8 non-ASCII punctuation is deliberately not treated as a patch boundary.
_SPACEBYTE_WHITESPACE = (9, 10, 11, 12, 13, 32)  # \t \n \v \f \r space


def build_spacebyte_boundary_mask(
    x: torch.Tensor,
    seg_ids: Optional[torch.Tensor] = None,
    spacelike_bytes: Tuple[int, ...] = _SPACEBYTE_WHITESPACE,
) -> torch.Tensor:
    """Return [B,T] positions promoted to SpaceByte global patches.

    A text patch starts at the first byte of each document/window and at a
    spacelike byte that is not itself preceded by another spacelike byte. BOS/EOS
    ids are also promoted because packed byte streams use them as hard document
    boundaries. Marking only the first byte in a run of spaces prevents wasting
    global compute on repeated whitespace.
    """
    if x.dim() != 2:
        raise ValueError(f"x must be [B,T], got shape {tuple(x.shape)}")
    B, T = x.shape
    is_space = torch.zeros_like(x, dtype=torch.bool)
    for byte in spacelike_bytes:
        is_space |= x == byte

    prev_space = torch.zeros_like(is_space)
    prev_space[:, 1:] = is_space[:, :-1]
    boundary = is_space & ~prev_space
    boundary |= (x == BOS_ID) | (x == EOS_ID)

    # Every sampled window needs at least one patch. When segment ids are known,
    # also seed each new document segment so document-aware global attention has
    # a patch before the first visible space.
    boundary[:, 0] = True
    if seg_ids is not None:
        new_doc = torch.zeros_like(boundary)
        new_doc[:, 1:] = seg_ids[:, 1:] != seg_ids[:, :-1]
        boundary |= new_doc
    return boundary


class ByteDecoder(nn.Module):
    def __init__(self, cfg: ByteLMConfig):
        super().__init__()
        cfg.validate()
        self.cfg = cfg
        self.embed = nn.Embedding(VOCAB_SIZE, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.norm_f = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, VOCAB_SIZE, bias=False)
        self.rope = RotaryEmbedding(cfg.head_dim, cfg.rope_theta)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.embed.weight

        self.apply(self._init_weights)
        # Scale residual-path projections by 1/sqrt(2*n_layers) for stable depth.
        scale = 1.0 / math.sqrt(2 * cfg.n_layers)
        for name, p in self.named_parameters():
            if name.endswith("wo.weight") or name.endswith("w_down.weight"):
                with torch.no_grad():
                    p.mul_(scale)

    @staticmethod
    def _init_weights(m: nn.Module):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def num_params(self, non_embedding: bool = False) -> int:
        n = sum(p.numel() for p in self.parameters())
        if non_embedding and not self.cfg.tie_embeddings:
            n -= self.embed.weight.numel()
        return n

    def forward(
        self,
        x: torch.Tensor,
        y: Optional[torch.Tensor] = None,
        seg_ids: Optional[torch.Tensor] = None,
        pos_ids: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Dict[str, float]]:
        B, T = x.shape
        if pos_ids is None:
            pos_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T)
        cos, sin = self.rope(pos_ids)
        attn_mask = build_doc_attn_mask(seg_ids) if seg_ids is not None else None

        h = self.drop(self.embed(x))
        for blk in self.blocks:
            h = blk(h, cos, sin, attn_mask)
        h = self.norm_f(h)
        logits = self.lm_head(h)

        loss, parts = _loss_parts(logits, y, self.cfg)
        return logits, loss, parts

    @torch.no_grad()
    def generate(self, idx: torch.Tensor, max_new_tokens: int,
                 temperature: float = 1.0, top_k: Optional[int] = None,
                 eos_id: Optional[int] = None) -> torch.Tensor:
        """Autoregressive sampling (no doc-mask; single growing sequence)."""
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.cfg.ctx_len:]
            logits, _, _ = self(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-6)
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, nxt), dim=1)
            if eos_id is not None and (nxt == eos_id).all():
                break
        return idx


class SpaceByteDecoder(nn.Module):
    """SpaceByte-style hierarchy over raw bytes.

    Local-in blocks process every byte cheaply. Global blocks process only patch
    boundary states (word/document boundaries), with causal/document masks over
    the patch sequence. The latest available global patch state is then scattered
    back to each byte and added as context before local-out blocks and the shared
    LM head. Setting n_local_in=n_local_out=0 and n_global=n_layers makes this
    reduce exactly to `ByteDecoder` when every position is a boundary.
    """

    def __init__(self, cfg: ByteLMConfig):
        super().__init__()
        cfg.validate()
        self.cfg = cfg
        self.embed = nn.Embedding(VOCAB_SIZE, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.local_in_blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_local_in)])
        self.global_blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_global)])
        self.local_out_blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_local_out)])
        self.norm_f = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, VOCAB_SIZE, bias=False)
        self.rope = RotaryEmbedding(cfg.head_dim, cfg.rope_theta)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.embed.weight

        self.apply(ByteDecoder._init_weights)
        total_layers = max(1, cfg.n_local_in + cfg.n_global + cfg.n_local_out)
        scale = 1.0 / math.sqrt(2 * total_layers)
        for name, p in self.named_parameters():
            if name.endswith("wo.weight") or name.endswith("w_down.weight"):
                with torch.no_grad():
                    p.mul_(scale)

    def num_params(self, non_embedding: bool = False) -> int:
        n = sum(p.numel() for p in self.parameters())
        if non_embedding and not self.cfg.tie_embeddings:
            n -= self.embed.weight.numel()
        return n

    @staticmethod
    def gather_patches(
        h: torch.Tensor,
        boundary_mask: torch.Tensor,
        pos_ids: torch.Tensor,
        seg_ids: Optional[torch.Tensor],
        max_patches: int,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Gather boundary states into a padded [B,P,D] patch sequence.

        Returns `(patch_h, patch_valid, patch_pos, patch_seg, patch_ids)`, where
        `patch_ids[b,t]` is the index of the latest patch at or before byte `t`.
        If a sequence has more than `max_patches` boundaries, later bytes reuse
        the last retained patch context; this keeps memory bounded and explicit.
        """
        if h.dim() != 3 or boundary_mask.dim() != 2:
            raise ValueError("h must be [B,T,D] and boundary_mask must be [B,T]")
        B, T, D = h.shape
        counts = boundary_mask.sum(dim=1).clamp_min(1)
        P = int(min(max_patches, counts.max().item()))
        device = h.device
        gather_idx = torch.zeros((B, P), dtype=torch.long, device=device)
        patch_valid = torch.zeros((B, P), dtype=torch.bool, device=device)
        patch_ids = torch.empty((B, T), dtype=torch.long, device=device)

        for b in range(B):
            idx = torch.nonzero(boundary_mask[b], as_tuple=False).flatten()
            if idx.numel() == 0:
                idx = torch.zeros(1, dtype=torch.long, device=device)
            keep = idx[:P]
            n = keep.numel()
            gather_idx[b, :n] = keep
            patch_valid[b, :n] = True
            ordinal = torch.cumsum(boundary_mask[b].to(torch.long), dim=0) - 1
            patch_ids[b] = ordinal.clamp(min=0, max=max(n - 1, 0))

        patch_h = h.gather(1, gather_idx.unsqueeze(-1).expand(B, P, D))
        patch_pos = pos_ids.gather(1, gather_idx)
        if seg_ids is None:
            patch_seg = torch.zeros((B, P), dtype=torch.long, device=device)
        else:
            patch_seg = seg_ids.gather(1, gather_idx)
        return patch_h, patch_valid, patch_pos, patch_seg, patch_ids

    @staticmethod
    def scatter_patches(patch_h: torch.Tensor, patch_ids: torch.Tensor) -> torch.Tensor:
        """Scatter latest patch states from [B,P,D] back to each byte [B,T,D]."""
        B, T = patch_ids.shape
        D = patch_h.size(-1)
        return patch_h.gather(1, patch_ids.unsqueeze(-1).expand(B, T, D))

    @staticmethod
    def _patch_attn_mask(patch_seg: torch.Tensor, patch_valid: torch.Tensor) -> torch.Tensor:
        """Causal/document patch mask that keeps padded query rows finite."""
        mask = build_doc_attn_mask(patch_seg)
        valid = patch_valid.unsqueeze(1)
        mask = mask & valid.unsqueeze(-1) & valid.unsqueeze(-2)
        eye = torch.eye(patch_seg.size(1), dtype=torch.bool, device=patch_seg.device)
        # Padded rows are ignored later, but SDPA still needs at least one True.
        mask = mask | (eye.unsqueeze(0).unsqueeze(0) & ~valid.unsqueeze(-1))
        return mask

    def forward(
        self,
        x: torch.Tensor,
        y: Optional[torch.Tensor] = None,
        seg_ids: Optional[torch.Tensor] = None,
        pos_ids: Optional[torch.Tensor] = None,
        boundary_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Dict[str, float]]:
        B, T = x.shape
        if pos_ids is None:
            pos_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T)
        byte_cos, byte_sin = self.rope(pos_ids)
        byte_attn_mask = build_doc_attn_mask(seg_ids) if seg_ids is not None else None

        h = self.drop(self.embed(x))
        for blk in self.local_in_blocks:
            h = blk(h, byte_cos, byte_sin, byte_attn_mask)

        if boundary_mask is None:
            boundary_mask = build_spacebyte_boundary_mask(x, seg_ids=seg_ids)
        patch_h, patch_valid, patch_pos, patch_seg, patch_ids = self.gather_patches(
            h, boundary_mask, pos_ids, seg_ids, self.cfg.max_patches
        )
        patch_cos, patch_sin = self.rope(patch_pos)
        patch_attn_mask = self._patch_attn_mask(patch_seg, patch_valid)
        for blk in self.global_blocks:
            patch_h = blk(patch_h, patch_cos, patch_sin, patch_attn_mask)
        patch_context = self.scatter_patches(patch_h, patch_ids)
        if self.cfg.n_local_in == 0 and self.cfg.n_local_out == 0:
            # Degenerate hierarchy: the global stream is the whole residual
            # stream, so all-boundary inputs exactly match the plain decoder.
            h = patch_context
        else:
            h = h + patch_context

        for blk in self.local_out_blocks:
            h = blk(h, byte_cos, byte_sin, byte_attn_mask)
        h = self.norm_f(h)
        logits = self.lm_head(h)

        loss, parts = _loss_parts(logits, y, self.cfg)
        return logits, loss, parts

    @torch.no_grad()
    def generate(self, idx: torch.Tensor, max_new_tokens: int,
                 temperature: float = 1.0, top_k: Optional[int] = None,
                 eos_id: Optional[int] = None) -> torch.Tensor:
        """Autoregressive sampling with SpaceByte boundary detection per window."""
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.cfg.ctx_len:]
            logits, _, _ = self(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-6)
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, nxt), dim=1)
            if eos_id is not None and (nxt == eos_id).all():
                break
        return idx


def _loss_parts(
    logits: torch.Tensor,
    y: Optional[torch.Tensor],
    cfg: ByteLMConfig,
) -> Tuple[Optional[torch.Tensor], Dict[str, float]]:
    """Shared CE + z-loss computation for both decoder variants."""
    if y is None:
        return None, {}
    ce = F.cross_entropy(
        logits.view(-1, VOCAB_SIZE), y.reshape(-1),
        label_smoothing=cfg.label_smoothing,
    )
    z = logits.logsumexp(dim=-1).pow(2).mean()
    loss = ce + cfg.z_loss_weight * z
    return loss, {"ce": ce.item(), "z": z.item(), "loss": loss.item()}


def build_model(cfg: ByteLMConfig) -> nn.Module:
    if cfg.variant == "plain":
        return ByteDecoder(cfg)
    if cfg.variant == "spacebyte":
        return SpaceByteDecoder(cfg)
    raise ValueError(f"unknown variant {cfg.variant!r}")