"""Hugging Face remote-code implementation of THOX Microprocessor.

This file is uploaded verbatim to the model repository, so it must be
self-contained: it may not import the ``thox_microprocessor`` package.  That
duplication is deliberate but dangerous — v0.1 shipped a copy whose submodules
were named ``n1``/``post``/``n2`` while the training package used
``attn_norm``/``post_attn_norm``/``mixer_norm``, so a published checkpoint would
have loaded with three randomly-initialised norms per layer and no error.
``tests/test_hub_parity.py`` now asserts strict state-dict compatibility and
logit equality between the two implementations.

Scope note: this reference path recomputes the full prefix for each decoding
step rather than maintaining a KV cache.  It is portable across transformers
versions, which matters more for remote code than decode throughput; the
``thox_microprocessor`` package implements cached decoding for training and
deployment.
"""

from __future__ import annotations

import math
from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

try:  # transformers >= 4.50 stopped mixing GenerationMixin into PreTrainedModel
    from transformers import GenerationMixin
except ImportError:  # pragma: no cover - older transformers
    from transformers.generation import GenerationMixin

from .configuration_thox_microprocessor import ThoxMicroprocessorConfig


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

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


class QKRMSNorm(RMSNorm):
    """Per-head Q/K RMS normalization applied before RoPE."""


def _rope_frequencies(seq_len: int, head_dim: int, theta: float, device) -> Tuple[torch.Tensor, torch.Tensor]:
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
    positions = torch.arange(seq_len, device=device, dtype=torch.float32)
    angles = torch.outer(positions, inv_freq)
    return angles.cos(), angles.sin()


def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    even = x[..., 0::2]
    odd = x[..., 1::2]
    cos = cos[None, None, :, :].to(dtype=x.dtype)
    sin = sin[None, None, :, :].to(dtype=x.dtype)
    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)


def _fwht_block(x: torch.Tensor) -> torch.Tensor:
    n = x.shape[-1]
    if n <= 0 or n & (n - 1):
        raise ValueError(f"FWHT input width must be a power of two, got {n}")
    y = x
    h = 1
    while h < n:
        shape = y.shape
        y = y.reshape(*shape[:-1], -1, 2, h)
        a = y[..., 0, :]
        b = y[..., 1, :]
        y = torch.cat((a + b, a - b), dim=-1).reshape(*shape)
        h *= 2
    return y * (1.0 / math.sqrt(n))


def _stride_shuffle(x: torch.Tensor, block_size: int) -> torch.Tensor:
    nblocks = x.shape[-1] // block_size
    return x.reshape(*x.shape[:-1], block_size, nblocks).transpose(-1, -2).reshape(x.shape)


def _stride_unshuffle(x: torch.Tensor, block_size: int) -> torch.Tensor:
    nblocks = x.shape[-1] // block_size
    return x.reshape(*x.shape[:-1], nblocks, block_size).transpose(-1, -2).reshape(x.shape)


class BlockHadamardMixer(nn.Module):
    def __init__(self, dim: int, block_size: int = 64, center: bool = True, shuffle: bool = False) -> None:
        super().__init__()
        if dim % block_size != 0:
            raise ValueError(f"dim ({dim}) must be divisible by block_size ({block_size})")
        self.dim = dim
        self.block_size = block_size
        self.num_blocks = dim // block_size
        self.center = center
        self.shuffle = shuffle and self.num_blocks > 1
        self.d1 = nn.Parameter(torch.ones(dim))
        self.d2 = nn.Parameter(torch.ones(dim))
        self.d3 = nn.Parameter(torch.full((dim,), 0.02))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shape = x.shape
        y = x * self.d1
        if self.shuffle:
            y = _stride_shuffle(y, self.block_size)
        y = y.reshape(*shape[:-1], self.num_blocks, self.block_size)
        y = _fwht_block(y)
        y = F.silu(y * self.d2.reshape(self.num_blocks, self.block_size))
        if self.center:
            y = y - y.mean(dim=-1, keepdim=True)
        y = _fwht_block(y).reshape(shape)
        if self.shuffle:
            y = _stride_unshuffle(y, self.block_size)
        return y * self.d3


class GroupedQueryAttention(nn.Module):
    def __init__(self, config: ThoxMicroprocessorConfig) -> None:
        super().__init__()
        self.config = config
        self.head_dim = config.d_model // config.num_heads
        kv_dim = config.num_kv_heads * self.head_dim
        self.q_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.k_proj = nn.Linear(config.d_model, kv_dim, bias=False)
        self.v_proj = nn.Linear(config.d_model, kv_dim, bias=False)
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.head_gate = nn.Linear(config.d_model, config.num_heads, bias=False)
        self.q_norm = QKRMSNorm(self.head_dim, config.rms_eps) if config.use_qk_norm else nn.Identity()
        self.k_norm = QKRMSNorm(self.head_dim, config.rms_eps) if config.use_qk_norm else nn.Identity()

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        attn_bias: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        cfg = self.config
        batch, seq_len, _ = x.shape
        q = self.q_proj(x).view(batch, seq_len, cfg.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch, seq_len, cfg.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch, seq_len, cfg.num_kv_heads, self.head_dim).transpose(1, 2)
        q = _apply_rope(self.q_norm(q), cos, sin)
        k = _apply_rope(self.k_norm(k), cos, sin)

        repeats = cfg.num_heads // cfg.num_kv_heads
        if repeats > 1:
            k = k.repeat_interleave(repeats, dim=1)
            v = v.repeat_interleave(repeats, dim=1)

        scores = torch.matmul(q, k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        if attn_bias is not None:
            scores = scores + attn_bias
        probs = F.softmax(scores.float(), dim=-1).to(dtype=q.dtype)
        probs = F.dropout(probs, p=cfg.attention_dropout, training=self.training)
        context = torch.matmul(probs, v)

        gate = torch.sigmoid(self.head_gate(x)).transpose(1, 2).unsqueeze(-1)
        context = (context * gate).transpose(1, 2).contiguous().view(batch, seq_len, cfg.d_model)
        return self.out_proj(context)


class ThoxMicroBlock(nn.Module):
    def __init__(self, config: ThoxMicroprocessorConfig, layer_index: int = 0) -> None:
        super().__init__()
        self.config = config
        self.attn_norm = RMSNorm(config.d_model, config.rms_eps)
        self.attn = GroupedQueryAttention(config)
        self.post_attn_norm = (
            RMSNorm(config.d_model, config.rms_eps) if config.use_sandwich_norm else nn.Identity()
        )
        if config.use_hadamard_mixer:
            self.mixer_norm = RMSNorm(config.d_model, config.rms_eps)
            self.mixer = BlockHadamardMixer(
                config.d_model,
                config.hadamard_block_size,
                center=config.mixer_center,
                shuffle=config.mixer_block_shuffle and (layer_index % 2 == 1),
            )
        else:
            self.mixer_norm = None
            self.mixer = None
        if config.use_residual_gate:
            self.attn_gate = nn.Parameter(torch.tensor(0.0))
            self.mixer_gate = nn.Parameter(torch.tensor(0.0))
        else:
            self.register_parameter("attn_gate", None)
            self.register_parameter("mixer_gate", None)

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        attn_bias: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        update = self.post_attn_norm(self.attn(self.attn_norm(x), cos, sin, attn_bias))
        if self.attn_gate is not None:
            update = torch.sigmoid(self.attn_gate) * update
        x = x + update
        if self.mixer is not None:
            update = self.mixer(self.mixer_norm(x))
            if self.mixer_gate is not None:
                update = torch.sigmoid(self.mixer_gate) * update
            x = x + update
        return x


class ThoxMicroprocessorPreTrainedModel(PreTrainedModel):
    config_class = ThoxMicroprocessorConfig
    base_model_prefix = "thox_microprocessor"
    supports_gradient_checkpointing = False
    _supports_cache_class = False
    _supports_sdpa = False
    _supports_flash_attn_2 = False
    # With tied embeddings there is no second tensor to tie, so nothing is
    # shared and this stays empty.  An empty dict is iterated identically by the
    # transformers 4.x (list) and 5.x (mapping) code paths.
    _tied_weights_keys: dict = {}

    def _init_weights(self, module: nn.Module) -> None:
        std = self.config.initializer_range
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, RMSNorm):
            nn.init.ones_(module.weight)
        elif isinstance(module, BlockHadamardMixer):
            nn.init.ones_(module.d1)
            nn.init.ones_(module.d2)
            nn.init.constant_(module.d3, 0.02)


class ThoxMicroprocessorForCausalLM(ThoxMicroprocessorPreTrainedModel, GenerationMixin):
    def __init__(self, config: ThoxMicroprocessorConfig) -> None:
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
        self.embed_dropout = nn.Dropout(config.embedding_dropout)
        self.layers = nn.ModuleList(
            [ThoxMicroBlock(config, layer_index=i) for i in range(config.num_layers)]
        )
        self.final_norm = RMSNorm(config.d_model, config.rms_eps)
        self.lm_head = (
            None
            if config.tie_word_embeddings
            else nn.Linear(config.d_model, config.vocab_size, bias=False)
        )
        self.post_init()
        # Depth-scaled residual init, applied after post_init so it is not
        # overwritten.  This must match the training package exactly.
        residual_std = config.initializer_range / math.sqrt(2 * config.num_layers)
        for layer in self.layers:
            nn.init.normal_(layer.attn.out_proj.weight, mean=0.0, std=residual_std)

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

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, value) -> None:
        self.lm_head = value

    def _output_weight(self) -> torch.Tensor:
        return self.embed_tokens.weight if self.lm_head is None else self.lm_head.weight

    def _attention_bias(self, batch, seq_len, attention_mask, dtype, device) -> torch.Tensor:
        neg = torch.finfo(dtype).min
        causal = torch.ones(seq_len, seq_len, dtype=torch.bool, device=device).tril()
        bias = torch.where(
            causal,
            torch.zeros((), dtype=dtype, device=device),
            torch.full((), neg, dtype=dtype, device=device),
        ).expand(batch, 1, seq_len, seq_len)
        if attention_mask is not None:
            if attention_mask.ndim != 2 or attention_mask.shape != (batch, seq_len):
                raise ValueError(
                    f"attention_mask must have shape [{batch}, {seq_len}], "
                    f"got {tuple(attention_mask.shape)}"
                )
            pad = torch.where(
                attention_mask.to(torch.bool),
                torch.zeros((), dtype=dtype, device=device),
                torch.full((), neg, dtype=dtype, device=device),
            )
            bias = bias + pad[:, None, None, :]
        return bias

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        use_cache: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("provide exactly one of input_ids or inputs_embeds")
        return_dict = True if return_dict is None else return_dict

        if inputs_embeds is None:
            if input_ids.ndim != 2:
                raise ValueError(
                    f"input_ids must have shape [batch, seq_len], got {tuple(input_ids.shape)}"
                )
            inputs_embeds = self.embed_tokens(input_ids)
        batch, seq_len, _ = inputs_embeds.shape
        if seq_len > self.config.max_seq_len:
            raise ValueError(
                f"sequence length {seq_len} exceeds max_seq_len={self.config.max_seq_len}"
            )

        x = self.embed_dropout(inputs_embeds * math.sqrt(self.config.d_model))
        device = x.device
        cos, sin = _rope_frequencies(seq_len, self.config.head_dim, self.config.rope_theta, device)
        attn_bias = self._attention_bias(batch, seq_len, attention_mask, x.dtype, device)
        for layer in self.layers:
            x = layer(x, cos, sin, attn_bias)
        x = self.final_norm(x)
        logits = F.linear(x, self._output_weight())

        loss = None
        if labels is not None:
            if labels.shape[:2] != (batch, seq_len):
                raise ValueError("labels must have shape [batch, seq_len]")
            loss = F.cross_entropy(
                logits[:, :-1].contiguous().view(-1, self.config.vocab_size).float(),
                labels[:, 1:].contiguous().view(-1),
                ignore_index=-100,
            )

        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)

    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
        # No KV cache: always re-encode the (truncated) prefix.
        window = self.config.max_seq_len
        input_ids = input_ids[:, -window:]
        if attention_mask is not None:
            attention_mask = attention_mask[:, -window:]
        return {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False}


__all__ = [
    "ThoxMicroprocessorConfig",
    "ThoxMicroprocessorPreTrainedModel",
    "ThoxMicroprocessorForCausalLM",
]