"""Small, HF-compatible causal bootstrap model used by the M0 gate.

The diffusion objective and sampler are deliberately separate modules.  This model
provides the shared transformer backbone and a causal forward path so that the
project can validate shape correctness, parameter accounting, and reproducibility
before any expensive data work begins.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import torch
from torch import Tensor, nn
from torch.nn import functional as F

from .configuration_microloop import MicroLoopConfig

try:
    from transformers import PreTrainedModel
    from transformers.generation import GenerationMixin
    from transformers.utils import ModelOutput
except ImportError:  # pragma: no cover - only used in a minimal environment.

    class GenerationMixin:  # type: ignore[no-redef]
        pass

    class ModelOutput:  # type: ignore[no-redef]
        pass

    class PreTrainedModel(nn.Module):  # type: ignore[no-redef]
        config_class = MicroLoopConfig
        base_model_prefix = "microloop"

        def __init__(self, config: MicroLoopConfig) -> None:
            super().__init__()
            self.config = config


@dataclass
class MicroLoopCausalLMOutput(ModelOutput):
    """Minimal output object with both attribute and mapping-style access."""

    logits: Tensor
    loss: Optional[Tensor] = None
    hidden_states: Optional[Tensor] = None
    loop_applications: Optional[int] = None

    def __getitem__(self, key: str):
        return getattr(self, key)


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

    def forward(self, hidden_states: Tensor) -> Tensor:
        # Explicit computation rather than the fused ``F.rms_norm`` kernel:
        # the fused kernel selects implementations based on process-level state
        # and produces context-dependent numerics inside the training process
        # (the 2026-08-05 provenance incident).  This explicit path is
        # deterministic everywhere; the small speed cost is acceptable here.
        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
        return (hidden_states / torch.sqrt(variance + self.eps)) * self.weight


def _rotate_half(x: Tensor) -> Tensor:
    x_even = x[..., ::2]
    x_odd = x[..., 1::2]
    return torch.stack((-x_odd, x_even), dim=-1).flatten(-2)


def _rope_tables(max_position: int, head_dim: int, theta: float) -> tuple[Tensor, Tensor]:
    """Build interleaved rotary tables once, in float32 for stable reuse."""

    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
    positions = torch.arange(max_position, device=inv_freq.device, dtype=torch.float32)
    angles = positions.unsqueeze(-1) * inv_freq
    angles = torch.stack((angles, angles), dim=-1).flatten(-2)
    return angles.cos(), angles.sin()


def _apply_rope(q: Tensor, k: Tensor, cos: Tensor, sin: Tensor) -> tuple[Tensor, Tensor]:
    """Apply cached interleaved rotary embeddings to query and key tensors."""

    return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin


class GroupedQueryAttention(nn.Module):
    """Grouped-query self-attention with explicit Q/K/V projections.

    Attention uses an explicit scaled-dot-product implementation (matmul +
    softmax) rather than ``F.scaled_dot_product_attention`` because the fused
    kernels select implementations based on process-level state and produce
    deterministic but context-dependent results: evaluations inside the
    training process then disagree with evaluations of the same saved
    checkpoint in a fresh process (the 2026-08-05 provenance incident).  The
    explicit math path is deterministic everywhere at the cost of a small
    amount of speed, which is acceptable at this model size.
    """

    def __init__(self, config: MicroLoopConfig) -> None:
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.head_dim = config.head_dimension
        self.num_groups = self.num_heads // self.num_key_value_heads
        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(
            config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
        )
        self.v_proj = nn.Linear(
            config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
        )
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.output_gate = (
            nn.Linear(config.hidden_size, self.num_heads, bias=False)
            if config.attention_output_gate
            else None
        )
        self.rope_theta = config.rope_theta
        self.attention_implementation = config.attention_implementation
        if config.qk_norm == "per_head":
            # One affine scale is shared by all Q heads and one by all K heads;
            # normalization itself is applied independently over each head.
            self.q_norm: RMSNorm | None = RMSNorm(self.head_dim, config.rms_norm_eps)
            self.k_norm: RMSNorm | None = RMSNorm(self.head_dim, config.rms_norm_eps)
        else:
            self.q_norm = None
            self.k_norm = None

    def forward(
        self,
        hidden_states: Tensor,
        attention_mask: Tensor | None = None,
        position_ids: Tensor | None = None,
        rope_embeddings: tuple[Tensor, Tensor] | None = None,
    ) -> Tensor:
        batch, sequence, _ = hidden_states.shape
        q = self.q_proj(hidden_states).view(batch, sequence, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(
            batch, sequence, self.num_key_value_heads, self.head_dim
        )
        v = self.v_proj(hidden_states).view(
            batch, sequence, self.num_key_value_heads, self.head_dim
        )
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        if position_ids is None:
            position_ids = torch.arange(sequence, device=hidden_states.device).expand(batch, -1)
        if position_ids.shape != (batch, sequence):
            raise ValueError(
                f"position_ids must have shape [batch, sequence], got {tuple(position_ids.shape)}"
            )
        if rope_embeddings is None:
            table_cos, table_sin = _rope_tables(sequence, self.head_dim, self.rope_theta)
            cos = table_cos[position_ids].unsqueeze(1).to(q.dtype)
            sin = table_sin[position_ids].unsqueeze(1).to(q.dtype)
        else:
            cos, sin = rope_embeddings
        q, k = _apply_rope(q, k, cos, sin)
        if self.q_norm is not None:
            q = self.q_norm(q)
            assert self.k_norm is not None
            k = self.k_norm(k)

        if self.num_heads % self.num_key_value_heads != 0:
            raise ValueError(
                f"num_heads {self.num_heads} must divide num_key_value_heads "
                f"{self.num_key_value_heads}"
            )
        if attention_mask is None:
            visible: Tensor | None = None
            is_causal = True
        else:
            # A 2-D mask uses the conventional HF meaning: one means visible.
            if attention_mask.shape == (batch, sequence):
                causal = torch.tril(
                    torch.ones(sequence, sequence, device=q.device, dtype=torch.bool)
                )
                visible = causal.unsqueeze(0).unsqueeze(0) & attention_mask.bool().unsqueeze(
                    1
                ).unsqueeze(2)
                is_causal = False
            elif attention_mask.shape == (batch, sequence, sequence):
                visible = attention_mask.bool().unsqueeze(1)
                is_causal = False
            else:
                raise ValueError(
                    "attention_mask must have shape [batch, sequence] or "
                    "[batch, sequence, sequence], "
                    f"got {tuple(attention_mask.shape)}"
                )
        # GQA: repeat the KV heads so every query head has its own K/V.
        repeat = self.num_heads // self.num_key_value_heads
        if repeat > 1:
            k = k.repeat_interleave(repeat, dim=1)
            v = v.repeat_interleave(repeat, dim=1)
        if self.attention_implementation == "sdpa":
            attended = F.scaled_dot_product_attention(
                q, k, v, attn_mask=visible, dropout_p=0.0, is_causal=is_causal
            )
        else:
            scores = torch.matmul(q, k.transpose(-2, -1)) / float(self.head_dim) ** 0.5
            if is_causal:
                seq_ids = torch.arange(sequence, device=scores.device)
                visible = seq_ids.unsqueeze(0) <= seq_ids.unsqueeze(1)
                visible = visible.expand(batch, self.num_heads, sequence, sequence)
            if visible is not None:
                scores = scores.masked_fill(~visible, float("-inf"))
            probs = torch.softmax(scores, dim=-1)
            attended = torch.matmul(probs, v)
        if self.output_gate is not None:
            gate = F.silu(self.output_gate(hidden_states)).transpose(1, 2).unsqueeze(-1)
            attended = attended * gate
        attended = attended.transpose(1, 2).contiguous().view(batch, sequence, -1)
        return self.o_proj(attended)


class MicroLoopBlock(nn.Module):
    """Pre-norm transformer block with SwiGLU feed-forward network."""

    def __init__(self, config: MicroLoopConfig) -> None:
        super().__init__()
        self.attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attn = GroupedQueryAttention(config)
        self.ffn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.ffn_gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.ffn_up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.ffn_down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.use_attn_residuals = config.attn_res_block_size is not None
        if self.use_attn_residuals:
            self.attn_res_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
            self.ffn_res_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
            self.attn_res_proj = nn.Linear(config.hidden_size, 1, bias=False)
            self.ffn_res_proj = nn.Linear(config.hidden_size, 1, bias=False)
        clamp = dict(config.swiglu_clamp)
        self.swiglu_clamp_enabled = bool(clamp.get("enabled", False))
        self.swiglu_linear_min = float(clamp.get("linear_min", -10.0))
        self.swiglu_linear_max = float(clamp.get("linear_max", 10.0))
        self.swiglu_gate_max = float(clamp.get("gate_max", 10.0))

    def forward(
        self,
        hidden_states: Tensor,
        attention_mask: Tensor | None = None,
        position_ids: Tensor | None = None,
        rope_embeddings: tuple[Tensor, Tensor] | None = None,
        block_residuals: list[Tensor] | None = None,
    ) -> Tensor:
        if self.use_attn_residuals and block_residuals:
            residual_stack = torch.stack(block_residuals, dim=-2)
            scores = torch.cat(
                [self.attn_res_proj(self.attn_res_norm(state)) for state in block_residuals], dim=-1
            )
            hidden_states = hidden_states + (
                torch.softmax(scores, dim=-1).unsqueeze(-1) * residual_stack
            ).sum(dim=-2)
        hidden_states = hidden_states + self.attn(
            self.attn_norm(hidden_states),
            attention_mask,
            position_ids,
            rope_embeddings,
        )
        if self.use_attn_residuals and block_residuals:
            residual_stack = torch.stack(block_residuals, dim=-2)
            scores = torch.cat(
                [self.ffn_res_proj(self.ffn_res_norm(state)) for state in block_residuals], dim=-1
            )
            hidden_states = hidden_states + (
                torch.softmax(scores, dim=-1).unsqueeze(-1) * residual_stack
            ).sum(dim=-2)
        ffn_input = self.ffn_norm(hidden_states)
        gate_linear = self.ffn_gate(ffn_input)
        up_linear = self.ffn_up(ffn_input)
        if self.swiglu_clamp_enabled:
            gate_linear = gate_linear.clamp(self.swiglu_linear_min, self.swiglu_linear_max)
            up_linear = up_linear.clamp(self.swiglu_linear_min, self.swiglu_linear_max)
        gate = F.silu(gate_linear)
        if self.swiglu_clamp_enabled:
            gate = gate.clamp(max=self.swiglu_gate_max)
        ffn_output = self.ffn_down(gate * up_linear)
        return hidden_states + ffn_output


class MicroLoopPreTrainedModel(PreTrainedModel):
    config_class = MicroLoopConfig
    base_model_prefix = "microloop"


class MicroLoopForDiffusionLM(MicroLoopPreTrainedModel, GenerationMixin):
    """Backbone plus tied output head for causal bootstrap and diffusion training."""

    _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}

    def __init__(self, config: MicroLoopConfig) -> None:
        config.validate()
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            [MicroLoopBlock(config) for _ in range(config.num_hidden_layers)]
        )
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.mtp_proj = (
            nn.Linear(config.hidden_size, config.hidden_size, bias=False)
            if config.mtp_enabled
            else None
        )
        rope_cos, rope_sin = _rope_tables(
            config.max_position_embeddings, config.head_dimension, config.rope_theta
        )
        # PERSISTENT buffers on purpose: transformers 5.x ``from_pretrained``
        # re-initializes non-persistent buffers that are missing from the
        # checkpoint (``_initialize_missing_keys`` -> ``initialize_weights``),
        # overwriting the rotary tables with garbage.  That made every fresh-
        # process evaluation of saved checkpoints compute with corrupted rope
        # tables (the historical "external collapse" at ~0.03 accuracy was this
        # artifact).  Persisting the tables makes the saved checkpoint carry
        # exactly the tables used in training.
        self.register_buffer("rope_cos", rope_cos, persistent=True)
        self.register_buffer("rope_sin", rope_sin, persistent=True)
        if config.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight
        # Canonical HF pattern: post_init() installs all_tied_weights_keys and
        # dispatches _initialize_weights per module.
        self.post_init()

    def _reset_rope_buffers(self) -> None:
        """Recompute the rotary tables from the config (they are deterministic)."""

        rope_cos, rope_sin = _rope_tables(
            self.config.max_position_embeddings,
            self.config.head_dimension,
            self.config.rope_theta,
        )
        self.rope_cos.copy_(rope_cos)
        self.rope_sin.copy_(rope_sin)

    def _initialize_weights(self, module: nn.Module, is_custom_code: bool = False) -> None:
        # The caller controls the RNG through seed_everything; this method performs
        # no hidden reseeding and is therefore reproducible by construction.
        if getattr(module, "_is_hf_initialized", False):
            return
        if module is self:
            # The main module owns the rotary tables.  transformers 5.x
            # re-initializes buffers that are missing from a loaded checkpoint
            # ("_initialize_missing_keys"), which zeroes/garbles the tables;
            # restore the deterministic canonical tables here instead.
            self._reset_rope_buffers()
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if getattr(module, "bias", None) is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, RMSNorm):
            nn.init.ones_(module.weight)

    def get_input_embeddings(self) -> nn.Embedding:
        return self.embed_tokens

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.embed_tokens = value
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight

    def forward(
        self,
        input_ids: Tensor | None = None,
        inputs_embeds: Tensor | None = None,
        attention_mask: Tensor | None = None,
        document_ids: Tensor | None = None,
        position_ids: Tensor | None = None,
        labels: Tensor | None = None,
        logit_mask: Tensor | None = None,
        mtp_loss_weight: float = 0.0,
        loop_count: int = 1,
        output_hidden_states: bool = False,
        **_: object,
    ) -> MicroLoopCausalLMOutput:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("exactly one of input_ids or inputs_embeds must be provided")
        if loop_count < 1:
            raise ValueError("loop_count must be at least one")
        if input_ids is not None:
            if input_ids.dim() != 2:
                raise ValueError(
                    f"input_ids must have shape [batch, sequence], got {input_ids.dim()}-D"
                )
            batch, sequence = input_ids.shape
        else:
            if inputs_embeds is None or inputs_embeds.dim() != 3:
                raise ValueError(
                    "inputs_embeds must have shape [batch, sequence, hidden], got "
                    f"{None if inputs_embeds is None else inputs_embeds.dim()}-D"
                )
            batch, sequence, _ = inputs_embeds.shape
        device = (input_ids if input_ids is not None else inputs_embeds).device
        if document_ids is not None:
            if document_ids.shape != (batch, sequence):
                raise ValueError("document_ids must have the same shape as the input")
            if attention_mask is not None and attention_mask.dim() != 2:
                raise ValueError(
                    "document_ids cannot be combined with a precomputed attention mask"
                )
            causal = torch.tril(torch.ones(sequence, sequence, dtype=torch.bool, device=device))
            valid = (
                attention_mask.bool()
                if attention_mask is not None
                else torch.ones((batch, sequence), dtype=torch.bool, device=device)
            )
            attention_mask = (
                causal.unsqueeze(0)
                & document_ids.unsqueeze(2).eq(document_ids.unsqueeze(1))
                & valid.unsqueeze(1)
                & valid.unsqueeze(2)
            )
        if position_ids is None:
            position_ids = torch.arange(sequence, device=device).expand(batch, -1)
        rope_dtype = (
            torch.get_autocast_dtype("cuda")
            if device.type == "cuda" and torch.is_autocast_enabled("cuda")
            else self.embed_tokens.weight.dtype
        )
        rope_embeddings = (
            self.rope_cos[position_ids].unsqueeze(1).to(rope_dtype),
            self.rope_sin[position_ids].unsqueeze(1).to(rope_dtype),
        )
        hidden_states = self.embed_tokens(input_ids) if input_ids is not None else inputs_embeds
        loop_layers = set(self.config.looping.get("layers", [4, 5, 6]))
        block_residuals: list[Tensor] = []
        block_size = self.config.attn_res_block_size
        total_applications = 0
        for layer_number, layer in enumerate(self.layers, start=1):
            if block_size is not None and (layer_number - 1) % block_size == 0:
                block_residuals.append(hidden_states)
            prior_block_residuals = block_residuals[:-1] if block_size is not None else None
            repetitions = loop_count if layer_number in loop_layers else 1
            for _ in range(repetitions):
                hidden_states = layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    rope_embeddings=rope_embeddings,
                    block_residuals=prior_block_residuals,
                )
                total_applications += 1
        hidden_states = self.norm(hidden_states)
        if logit_mask is not None:
            # Diffusion training: only the masked positions carry loss, so the
            # output head runs on those rows instead of the full sequence.
            if labels is not None:
                raise ValueError("logit_mask cannot be combined with labels")
            if logit_mask.shape != (batch, sequence):
                raise ValueError("logit_mask must have the same shape as the input")
            flat = hidden_states.reshape(-1, self.config.hidden_size)
            logits = self.lm_head(flat[logit_mask.reshape(-1)])
        else:
            logits = self.lm_head(hidden_states)
        loss = None
        if labels is not None:
            if labels.shape != (batch, sequence):
                raise ValueError("labels must have the same shape as the input")
            if labels.size(1) < 2:
                raise ValueError("causal training requires sequences with at least two tokens")
            # Causal next-token prediction: position t predicts the label at t + 1.
            loss = F.cross_entropy(
                logits[:, :-1, :].reshape(-1, logits.size(-1)),
                labels[:, 1:].reshape(-1),
                ignore_index=-100,
            )
            if mtp_loss_weight:
                if self.mtp_proj is None:
                    raise ValueError("mtp_loss_weight requires mtp_enabled=true")
                if mtp_loss_weight < 0:
                    raise ValueError("mtp_loss_weight must be non-negative")
                mtp_logits = self.lm_head(self.mtp_proj(hidden_states[:, :-2, :]))
                mtp_loss = F.cross_entropy(
                    mtp_logits.reshape(-1, mtp_logits.size(-1)),
                    labels[:, 2:].reshape(-1),
                    ignore_index=-100,
                )
                loss = loss + float(mtp_loss_weight) * mtp_loss
        return MicroLoopCausalLMOutput(
            logits=logits,
            loss=loss,
            hidden_states=hidden_states if output_hidden_states else None,
            loop_applications=total_applications,
        )

    @torch.no_grad()
    def generate_greedy(
        self, input_ids: Tensor, max_new_tokens: int, eos_token_id: int | None = None
    ) -> Tensor:
        """Small causal smoke decoder; diffusion sampling belongs in ``sampler.py``."""

        return self.generate_causal(
            input_ids, max_new_tokens=max_new_tokens, eos_token_id=eos_token_id, do_sample=False
        )

    @torch.no_grad()
    def generate_causal(
        self,
        input_ids: Tensor,
        *,
        max_new_tokens: int,
        eos_token_id: int | None = None,
        do_sample: bool = False,
        temperature: float = 1.0,
        top_k: int | None = None,
    ) -> Tensor:
        """Generate a batched causal continuation for M2 validation and serving.

        This deliberately recomputes the context on every step. KV caching is a
        later optimization; keeping this reference path simple makes M2 output
        semantics straightforward to test.
        """

        if max_new_tokens < 0:
            raise ValueError("max_new_tokens must be non-negative")
        if temperature <= 0:
            raise ValueError("temperature must be positive")
        if top_k is not None and top_k <= 0:
            raise ValueError("top_k must be positive when supplied")
        generated = input_ids
        finished = torch.zeros(input_ids.size(0), dtype=torch.bool, device=input_ids.device)
        for _ in range(max_new_tokens):
            next_logits = self(generated).logits[:, -1, :]
            if do_sample:
                next_logits = next_logits / temperature
                if top_k is not None and top_k < next_logits.size(-1):
                    threshold = torch.topk(next_logits, top_k, dim=-1).values[:, -1:]
                    next_logits = next_logits.masked_fill(next_logits < threshold, float("-inf"))
                next_token = torch.multinomial(torch.softmax(next_logits, dim=-1), 1)
            else:
                next_token = next_logits.argmax(dim=-1, keepdim=True)
            if eos_token_id is not None:
                next_token = torch.where(
                    finished.unsqueeze(1),
                    torch.full_like(next_token, eos_token_id),
                    next_token,
                )
                finished |= next_token.squeeze(1).eq(eos_token_id)
            generated = torch.cat((generated, next_token), dim=1)
            if eos_token_id is not None and bool(finished.all()):
                break
        return generated