model.py

"""Unlimited-OCR MLX Model Implementation.

High-precision OCR model fully implemented in MLX for Apple Silicon acceleration.
Architecture: Vision Encoder (SAM-ViT-B + CLIP-L) → DeepSeek-V2 MoE Language Model.
"""

import math
from typing import Optional, Tuple, List, Dict
from dataclasses import dataclass

import mlx.core as mx
import mlx.nn as nn

from .config import UnlimitedOCRConfig, VisionConfig, LanguageConfig, ProjectorConfig


# =============================================================================
# Utility Functions
# =============================================================================

def _compute_default_rope_freqs(
    dim: int, max_position_embeddings: int = 32768, base: float = 10000.0
) -> mx.array:
    """Compute RoPE frequencies. Returns (max_pos, dim/2) for rotation."""
    theta = 1.0 / (base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim))
    t = mx.arange(max_position_embeddings, dtype=mx.float32)
    freqs = mx.outer(t, theta)
    return freqs


def _apply_rotary_pos_emb(q, k, cos, sin, position_ids=None):
    """Apply rotary position embeddings to query and key tensors.

    Args:
        q, k: [B, heads, seq_len, head_dim]
        cos, sin: [seq_len, half_dim] already sliced/indexed by caller
    """
    B, H, L, D = q.shape
    half_D = D // 2

    # cos/sin are already properly shaped by RotaryEmbedding
    # They should be [L, half_D] or [1, L, half_D]
    if cos.ndim == 3:
        cos = cos.reshape(-1, cos.shape[-1])
        sin = sin.reshape(-1, sin.shape[-1])

    # Ensure correct length
    cos = cos[:L]
    sin = sin[:L]

    # Reshape for broadcasting: [1, 1, L, half_D]
    cos = cos.reshape(1, 1, L, half_D)
    sin = sin.reshape(1, 1, L, half_D)

    def _rotate_half(x):
        x1 = x[..., :half_D]
        x2 = x[..., half_D:]
        return mx.concatenate([-x2, x1], axis=-1)

    # Duplicate cos/sin to full head_dim for element-wise multiply
    cos2 = mx.concatenate([cos, cos], axis=-1)
    sin2 = mx.concatenate([sin, sin], axis=-1)

    q_rot = q * cos2 + _rotate_half(q) * sin2
    k_rot = k * cos2 + _rotate_half(k) * sin2

    return q_rot, k_rot


def silu(x):
    """SiLU activation function."""
    return x * mx.sigmoid(x)


# =============================================================================
# RMSNorm
# =============================================================================

class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization."""

    def __init__(self, dims: int, eps: float = 1e-6):
        super().__init__()
        self.weight = mx.ones((dims,))
        self.eps = eps

    def __call__(self, x):
        return mx.fast.rms_norm(x, self.weight, self.eps)


# =============================================================================
# RoPE
# =============================================================================

class RotaryEmbedding:
    """Rotary Position Embedding."""

    def __init__(self, dim: int, max_position_embeddings: int = 32768, base: float = 10000.0):
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        self._freqs_cos_sin = None

    def _ensure_freqs(self):
        if self._freqs_cos_sin is None:
            freqs = _compute_default_rope_freqs(self.dim, self.max_position_embeddings, self.base)
            self._freqs_cos_sin = (mx.cos(freqs), mx.sin(freqs))

    @property
    def cos_cached(self):
        self._ensure_freqs()
        return self._freqs_cos_sin[0]

    @property
    def sin_cached(self):
        self._ensure_freqs()
        return self._freqs_cos_sin[1]

    def __call__(self, x, position_ids=None, seq_len=None):
        self._ensure_freqs()
        cos, sin = self.cos_cached, self.sin_cached
        # position_ids carry ABSOLUTE positions. Index the full table directly;
        # slicing by seq_len first would break decode (abs pos >= seq_len) and
        # is required for correctness under R-SWA KV eviction.
        if position_ids is not None:
            return cos[position_ids], sin[position_ids]
        if seq_len is not None:
            cos, sin = cos[:seq_len], sin[:seq_len]
        return cos, sin


# =============================================================================
# Standard Multi-Head Attention
# =============================================================================

class MultiHeadAttention(nn.Module):
    """Standard Multi-Head Attention with RoPE."""

    def __init__(self, config: LanguageConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.layer_idx = layer_idx

        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)

        self.rotary_emb = RotaryEmbedding(
            self.head_dim,
            max_position_embeddings=config.max_position_embeddings,
            base=config.rope_theta,
        )
        self.scale = self.head_dim ** -0.5

    def __call__(
        self,
        hidden_states: mx.array,
        attention_mask: Optional[mx.array] = None,
        position_ids: Optional[mx.array] = None,
        past_key_value: Optional[Tuple[mx.array, mx.array]] = None,
        use_cache: bool = False,
    ) -> Tuple[mx.array, Optional[Tuple[mx.array, mx.array]]]:
        B, L, _ = hidden_states.shape

        q = self.q_proj(hidden_states).reshape(B, L, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        k = self.k_proj(hidden_states).reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
        v = self.v_proj(hidden_states).reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)

        cos, sin = self.rotary_emb(q, position_ids=position_ids, seq_len=L)
        q, k = _apply_rotary_pos_emb(q, k, cos, sin, position_ids)

        if past_key_value is not None:
            pk, pv = past_key_value
            k = mx.concatenate([pk, k], axis=2)
            v = mx.concatenate([pv, v], axis=2)

        past_kv = (k, v) if use_cache else None

        # GQA: repeat k/v heads
        n_rep = self.num_heads // self.num_kv_heads
        if n_rep > 1:
            k = mx.repeat(k, n_rep, axis=1)
            v = mx.repeat(v, n_rep, axis=1)

        # Scaled dot-product attention
        scores = (q @ k.transpose(0, 1, 3, 2)) * self.scale
        if attention_mask is not None:
            scores = scores + attention_mask

        attn_weights = mx.softmax(scores.astype(mx.float32), axis=-1).astype(q.dtype)
        attn_output = attn_weights @ v

        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)
        output = self.o_proj(attn_output)
        return output, past_kv


# =============================================================================
# MLP (SwiGLU)
# =============================================================================

class SwiGLUMLP(nn.Module):
    """SwiGLU MLP used in dense layers and experts."""

    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def __call__(self, x):
        return self.down_proj(silu(self.gate_proj(x)) * self.up_proj(x))


# =============================================================================
# MoE (Mixture of Experts)
# =============================================================================

class MoEGate(nn.Module):
    """Top-k gating for MoE."""

    def __init__(self, config: LanguageConfig):
        super().__init__()
        self.top_k = config.num_experts_per_tok
        self.n_routed_experts = config.n_routed_experts
        self.scoring_func = config.scoring_func
        self.topk_method = config.topk_method
        self.norm_topk_prob = config.norm_topk_prob

        # Gate weight: [n_experts, hidden_size]
        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))

    def __call__(self, hidden_states: mx.array) -> Tuple[mx.array, mx.array]:
        # hidden_states: [B*L, hidden_size]
        logits = hidden_states.astype(mx.float32) @ self.weight.astype(mx.float32).T

        if self.scoring_func == "softmax":
            scores = mx.softmax(logits, axis=-1)
        else:
            scores = mx.sigmoid(logits)

        # Top-k selection (MLX topk returns indices, then we gather weights)
        topk_indices = mx.argpartition(-scores, kth=self.top_k - 1, axis=-1)[:, :self.top_k]
        # Gather the actual scores for these indices
        topk_weights = mx.take_along_axis(scores, topk_indices, axis=-1)

        if self.norm_topk_prob:
            denom = topk_weights.sum(axis=-1, keepdims=True) + 1e-20
            topk_weights = topk_weights / denom

        return topk_indices, topk_weights


class DeepSeekMoE(nn.Module):
    """DeepSeek-V2 MoE block with shared experts."""

    def __init__(self, config: LanguageConfig):
        super().__init__()
        self.num_experts_per_tok = config.num_experts_per_tok
        self.n_routed_experts = config.n_routed_experts
        self.moe_intermediate_size = config.moe_intermediate_size

        # Create routed experts
        self.experts = [
            SwiGLUMLP(config.hidden_size, self.moe_intermediate_size)
            for _ in range(self.n_routed_experts)
        ]

        self.gate = MoEGate(config)

        # Shared experts (2 experts with combined intermediate size)
        if config.n_shared_experts is not None:
            shared_dim = self.moe_intermediate_size * config.n_shared_experts
            self.shared_experts = SwiGLUMLP(config.hidden_size, shared_dim)

    def _stack_experts(self):
        """Lazily stack expert weights for gather-based inference (cached)."""
        if getattr(self, "_GW", None) is None:
            self._GW = mx.stack([e.gate_proj.weight for e in self.experts])  # [E, M, D]
            self._UW = mx.stack([e.up_proj.weight for e in self.experts])    # [E, M, D]
            self._DW = mx.stack([e.down_proj.weight for e in self.experts])  # [E, D, M]
        return self._GW, self._UW, self._DW

    def _moe_infer(self, x: mx.array, topk_ids: mx.array, topk_weights: mx.array) -> mx.array:
        """Inference-time MoE computation."""
        B, L, D = x.shape
        T = B * L
        K = self.num_experts_per_tok
        x_flat = x.reshape(T, D)
        ids = topk_ids.reshape(T, K)
        tw = topk_weights.reshape(T, K)

        # 解码步 (T 小): 纯 gather+einsum, 无 CPU 同步、无 Python 循环
        if T <= 32:
            GW, UW, DW = self._stack_experts()
            idf = ids.reshape(-1)                        # [T*K]
            xe = mx.repeat(x_flat, K, axis=0)            # [T*K, D]
            gw = mx.take(GW, idf, axis=0)                # [T*K, M, D]
            uw = mx.take(UW, idf, axis=0)                # [T*K, M, D]
            dw = mx.take(DW, idf, axis=0)                # [T*K, D, M]
            h = silu(mx.einsum("nmd,nd->nm", gw, xe)) * mx.einsum("nmd,nd->nm", uw, xe)
            o = mx.einsum("ndm,nm->nd", dw, h)           # [T*K, D]
            o = o.reshape(T, K, D) * tw[:, :, None]
            return o.sum(axis=1).reshape(B, L, D)

        # 预填充 (T 大): 按专家分组, gather 全权重会爆显存, 故保留循环
        tk_flat = ids.reshape(-1)
        tw_flat = tw.reshape(-1)
        import numpy as np
        token_counts = np.bincount(np.array(tk_flat, dtype=np.int32),
                                    minlength=self.n_routed_experts)
        sort_indices = mx.argsort(tk_flat)
        repeated_x = mx.repeat(x_flat, K, axis=0)
        sorted_tokens = repeated_x[sort_indices]
        sorted_weights = tw_flat[sort_indices]

        outputs = []
        start = 0
        for i in range(self.n_routed_experts):
            count = int(token_counts[i])
            if count == 0:
                continue
            end = start + count
            expert_out = self.experts[i](sorted_tokens[start:end])
            expert_out = expert_out * sorted_weights[start:end][:, None]
            outputs.append((sort_indices[start:end], expert_out))
            start = end

        if not outputs:
            return mx.zeros_like(x)

        all_indices = mx.concatenate([o[0] for o in outputs], axis=0)
        all_outputs = mx.concatenate([o[1] for o in outputs], axis=0)
        restore = mx.argsort(all_indices)
        final = all_outputs[restore]
        final = final.reshape(T, K, D).sum(axis=1)
        return final.reshape(B, L, D)

    def __call__(self, hidden_states: mx.array) -> mx.array:
        identity = hidden_states
        B, L, D = hidden_states.shape
        x_flat = hidden_states.reshape(-1, D)

        topk_idx, topk_weight = self.gate(x_flat)

        # Reshape routing back
        topk_idx = topk_idx.reshape(B * L, self.num_experts_per_tok)
        topk_weight = topk_weight.reshape(B * L, self.num_experts_per_tok)

        moe_out = self._moe_infer(hidden_states, topk_idx.reshape(B, L, -1), topk_weight.reshape(B, L, -1))

        if hasattr(self, 'shared_experts'):
            moe_out = moe_out + self.shared_experts(identity)

        return moe_out


# =============================================================================
# DeepSeek-V2 Decoder Layer
# =============================================================================

class DeepSeekDecoderLayer(nn.Module):
    """Single decoder layer with attention + MLP/MoE."""

    def __init__(self, config: LanguageConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size

        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.self_attn = MultiHeadAttention(config, layer_idx)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # Layer 0 is dense MLP, rest are MoE
        is_dense = layer_idx < config.first_k_dense_replace
        if is_dense:
            self.mlp = SwiGLUMLP(config.hidden_size, config.intermediate_size)
            self.is_moe = False
        else:
            self.mlp = DeepSeekMoE(config)
            self.is_moe = True

    def __call__(
        self,
        hidden_states: mx.array,
        attention_mask: Optional[mx.array] = None,
        position_ids: Optional[mx.array] = None,
        past_key_value: Optional[Tuple[mx.array, mx.array]] = None,
        use_cache: bool = False,
    ) -> Tuple[mx.array, Optional[Tuple[mx.array, mx.array]]]:
        # Self-attention
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states, present_kv = self.self_attn(
            hidden_states, attention_mask, position_ids, past_key_value, use_cache
        )
        hidden_states = residual + hidden_states

        # MLP / MoE
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states, present_kv


# =============================================================================
# DeepSeek-V2 Language Model
# =============================================================================

class DeepSeekModel(nn.Module):
    """DeepSeek-V2 Language Model (12 layers, MoE)."""

    def __init__(self, config: LanguageConfig):
        super().__init__()
        self.config = config

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = [
            DeepSeekDecoderLayer(config, i)
            for i in range(config.num_hidden_layers)
        ]
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def __call__(
        self,
        input_ids: Optional[mx.array] = None,
        inputs_embeds: Optional[mx.array] = None,
        attention_mask: Optional[mx.array] = None,
        position_ids: Optional[mx.array] = None,
        past_key_values: Optional[List[Tuple[mx.array, mx.array]]] = None,
        use_cache: bool = False,
    ) -> Tuple[mx.array, Optional[List[Tuple[mx.array, mx.array]]]]:

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        B, L, _ = inputs_embeds.shape

        # Create causal mask
        if attention_mask is None:
            attention_mask = mx.tril(mx.ones((L, L), dtype=mx.bool_))
            attention_mask = mx.where(attention_mask, 0.0, float('-inf'))
            attention_mask = attention_mask[None, None, :, :]  # [1, 1, L, L]

        # Create position IDs
        if position_ids is None:
            if past_key_values is not None and past_key_values[0] is not None:
                cache_len = past_key_values[0][0].shape[2]
                position_ids = mx.arange(cache_len, cache_len + L, dtype=mx.int32)[None, :]
            else:
                position_ids = mx.arange(0, L, dtype=mx.int32)[None, :]

        hidden_states = inputs_embeds
        new_kv_cache = [] if use_cache else None

        for i, layer in enumerate(self.layers):
            pkv = past_key_values[i] if past_key_values else None
            hidden_states, nkv = layer(
                hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_value=pkv,
                use_cache=use_cache,
            )
            if use_cache:
                new_kv_cache.append(nkv)

        hidden_states = self.norm(hidden_states)
        return hidden_states, new_kv_cache


# =============================================================================
# SAM-ViT-B Vision Encoder
# =============================================================================

def _interp_linear(rel_pos: mx.array, out_size: int) -> mx.array:
    """1D linear interpolation of a rel-pos table (align_corners=False)."""
    L, C = rel_pos.shape
    if L == out_size:
        return rel_pos
    scale = L / out_size
    dst = mx.arange(out_size).astype(mx.float32)
    src = mx.clip((dst + 0.5) * scale - 0.5, 0.0, L - 1.0)
    lo = mx.floor(src).astype(mx.int32)
    hi = mx.minimum(lo + 1, L - 1)
    w = (src - lo.astype(mx.float32))[:, None]
    return rel_pos[lo] * (1.0 - w) + rel_pos[hi] * w


def _get_rel_pos(q_size: int, k_size: int, rel_pos: mx.array) -> mx.array:
    """Standard SAM get_rel_pos: interpolate + gather by relative coords."""
    max_rel_dist = 2 * max(q_size, k_size) - 1
    rel_pos_resized = _interp_linear(rel_pos, max_rel_dist)
    q_coords = mx.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
    k_coords = mx.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
    relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
    return rel_pos_resized[relative_coords.astype(mx.int32)]


def _add_decomposed_rel_pos(attn, q, rel_pos_h, rel_pos_w, q_size, k_size):
    """Add decomposed relative position bias to attention (standard SAM)."""
    q_h, q_w = q_size
    k_h, k_w = k_size
    Rh = _get_rel_pos(q_h, k_h, rel_pos_h)  # [q_h, k_h, C]
    Rw = _get_rel_pos(q_w, k_w, rel_pos_w)  # [q_w, k_w, C]
    B, _, dim = q.shape
    r_q = q.reshape(B, q_h, q_w, dim)
    rel_h = mx.einsum("bhwc,hkc->bhwk", r_q, Rh)
    rel_w = mx.einsum("bhwc,wkc->bhwk", r_q, Rw)
    attn = attn.reshape(B, q_h, q_w, k_h, k_w)
    attn = attn + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]
    return attn.reshape(B, q_h * q_w, k_h * k_w)


def _window_partition(x: mx.array, window_size: int):
    """Partition [B,H,W,C] into windows [B*nw, ws, ws, C] with padding."""
    B, H, W, C = x.shape
    pad_h = (window_size - H % window_size) % window_size
    pad_w = (window_size - W % window_size) % window_size
    if pad_h > 0 or pad_w > 0:
        x = mx.pad(x, [(0, 0), (0, pad_h), (0, pad_w), (0, 0)])
    Hp, Wp = H + pad_h, W + pad_w
    x = x.reshape(B, Hp // window_size, window_size, Wp // window_size, window_size, C)
    windows = x.transpose(0, 1, 3, 2, 4, 5).reshape(-1, window_size, window_size, C)
    return windows, (Hp, Wp)


def _window_unpartition(windows, window_size, pad_hw, hw):
    """Inverse of _window_partition; crop back to original [B,H,W,C]."""
    Hp, Wp = pad_hw
    H, W = hw
    B = windows.shape[0] // (Hp * Wp // window_size // window_size)
    x = windows.reshape(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
    x = x.transpose(0, 1, 3, 2, 4, 5).reshape(B, Hp, Wp, -1)
    if Hp > H or Wp > W:
        x = x[:, :H, :W, :]
    return x


class SAMAttention(nn.Module):
    """SAM attention block with relative position bias."""

    def __init__(
        self,
        dim: int,
        num_heads: int,
        window_size: int = 0,
        use_rel_pos: bool = True,
        input_size: Tuple[int, int] = (64, 64),
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.window_size = window_size
        self.scale = self.head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=True)
        self.proj = nn.Linear(dim, dim, bias=True)

        self.use_rel_pos = use_rel_pos
        if use_rel_pos:
            self.rel_pos_h = mx.zeros((2 * input_size[0] - 1, self.head_dim))
            self.rel_pos_w = mx.zeros((2 * input_size[1] - 1, self.head_dim))

    def __call__(self, x: mx.array) -> mx.array:
        # x: [B, H, W, C] spatial format (standard SAM)
        B, H, W, C = x.shape
        qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, self.head_dim)
        qkv = qkv.transpose(2, 0, 3, 1, 4).reshape(3, B * self.num_heads, H * W, self.head_dim)
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q * self.scale) @ k.transpose(0, 2, 1)
        if self.use_rel_pos:
            attn = _add_decomposed_rel_pos(
                attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W)
            )
        attn = mx.softmax(attn.astype(mx.float32), axis=-1).astype(q.dtype)
        out = attn @ v  # [B*heads, H*W, head_dim]
        out = out.reshape(B, self.num_heads, H, W, self.head_dim)
        out = out.transpose(0, 2, 3, 1, 4).reshape(B, H, W, C)
        return self.proj(out)



class SAMMLP(nn.Module):
    """SAM MLP block."""

    def __init__(self, dim: int, mlp_dim: int):
        super().__init__()
        self.lin1 = nn.Linear(dim, mlp_dim)
        self.lin2 = nn.Linear(mlp_dim, dim)

    def __call__(self, x):
        return self.lin2(nn.gelu(self.lin1(x)))


class SAMBlock(nn.Module):
    """SAM ViT block."""

    def __init__(
        self,
        dim: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        window_size: int = 0,
        use_rel_pos: bool = True,
        input_size: Tuple[int, int] = (64, 64),
    ):
        super().__init__()
        self.window_size = window_size
        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
        self.attn = SAMAttention(
            dim, num_heads,
            window_size=window_size,
            use_rel_pos=use_rel_pos,
            input_size=input_size,
        )
        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
        self.mlp = SAMMLP(dim, int(dim * mlp_ratio))

    def __call__(self, x):
        # x: [B, H, W, C] spatial format
        shortcut = x
        x = self.norm1(x)
        # Window partition
        if self.window_size > 0:
            H, W = x.shape[1], x.shape[2]
            x, pad_hw = _window_partition(x, self.window_size)
        x = self.attn(x)
        # Reverse windows
        if self.window_size > 0:
            x = _window_unpartition(x, self.window_size, pad_hw, (H, W))
        x = shortcut + x
        x = x + self.mlp(self.norm2(x))
        return x


class PatchEmbed(nn.Module):
    """Patch embedding for SAM. Uses NHWC format for MLX."""

    def __init__(self, kernel_size=16, stride=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size, stride=stride, bias=True)

    def __call__(self, x):
        # x: [B, H, W, C] (NHWC)
        return self.proj(x)


class SAMVisionEncoder(nn.Module):
    """SAM-ViT-B vision encoder."""

    def __init__(self, config: VisionConfig):
        super().__init__()
        self.img_size = config.sam_img_size
        self.patch_size = config.sam_patch_size
        grid_size = self.img_size // self.patch_size  # 64

        self.patch_embed = PatchEmbed(
            kernel_size=config.sam_patch_size,
            stride=config.sam_patch_size,
            in_chans=3,
            embed_dim=config.sam_embed_dim,
        )
        self.pos_embed = mx.zeros((1, grid_size, grid_size, config.sam_embed_dim))

        input_size = (grid_size, grid_size)
        self.window_size = config.sam_window_size
        self.blocks = []
        for i in range(config.sam_depth):
            use_global = i in config.sam_global_attn_indexes
            window_size = 0 if use_global else config.sam_window_size
            # rel_pos table size follows the attention span:
            # global blocks see the full grid; window blocks see one window.
            blk_input_size = input_size if use_global else (window_size, window_size)
            self.blocks.append(SAMBlock(
                dim=config.sam_embed_dim,
                num_heads=config.sam_num_heads,
                mlp_ratio=config.sam_mlp_ratio,
                window_size=window_size,
                input_size=blk_input_size,
            ))

        # Neck
        self.neck = nn.Sequential(
            nn.Conv2d(config.sam_embed_dim, config.sam_out_chans, 1, bias=False),
            nn.LayerNorm(config.sam_out_chans, eps=1e-6),
            nn.Conv2d(config.sam_out_chans, config.sam_out_chans, 3, padding=1, bias=False),
            nn.LayerNorm(config.sam_out_chans, eps=1e-6),
        )

        # Downsampling convolutions
        self.net_2 = nn.Conv2d(256, 512, 3, stride=2, padding=1, bias=False)
        self.net_3 = nn.Conv2d(512, 1024, 3, stride=2, padding=1, bias=False)

    def __call__(self, x: mx.array) -> mx.array:
        # x: [B, H, W, C] (NHWC format for MLX)
        B, H_in, W_in, C_in = x.shape

        x = self.patch_embed(x)  # [B, H_p, W_p, 768]
        H_p, W_p = x.shape[1], x.shape[2]

        # Add positional embedding (spatial 2D, bilinear interpolated to grid)
        if self.pos_embed.shape[1] != H_p:
            pos = _interpolate_pos_embed(self.pos_embed, H_p)
        else:
            pos = self.pos_embed
        x = x + pos  # [B, H_p, W_p, 768]

        # Transformer blocks operate in spatial format
        for blk in self.blocks:
            x = blk(x)

        # Neck (Conv2d with NHWC)
        x = self.neck(x)  # [B, 64, 64, 256]

        # Downsampling (NHWC)
        x = self.net_2(x)  # [B, 32, 32, 512]
        x = self.net_3(x)  # [B, 16, 16, 1024]

        return x


def _interpolate_pos_embed(pos_embed, target_size):
    """Bilinear interpolation of [1, src, src, dim] pos embed to target grid."""
    src = pos_embed.shape[1]
    if src == target_size:
        return pos_embed
    # interpolate rows then cols using the shared 1D linear helper
    dim = pos_embed.shape[-1]
    x = pos_embed[0]  # [src, src, dim]
    # rows
    x = _interp_linear(x.reshape(src, src * dim), target_size).reshape(target_size, src, dim)
    # cols
    x = x.transpose(1, 0, 2).reshape(src, target_size * dim)
    x = _interp_linear(x, target_size).reshape(target_size, target_size, dim).transpose(1, 0, 2)
    return x[None]


# =============================================================================
# CLIP-L Vision Encoder
# =============================================================================

class CLIPAttention(nn.Module):
    """CLIP multi-head self-attention."""

    def __init__(self, hidden_size: int, num_heads: int):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.qkv_proj = nn.Linear(hidden_size, hidden_size * 3, bias=True)
        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=True)
        self.scale = self.head_dim ** -0.5

    def __call__(self, x):
        B, N, C = x.shape
        qkv = self.qkv_proj(x).reshape(B, N, 3, self.num_heads, self.head_dim)
        q, k, v = qkv[:, :, 0].transpose(0, 2, 1, 3), qkv[:, :, 1].transpose(0, 2, 1, 3), qkv[:, :, 2].transpose(0, 2, 1, 3)

        attn = (q @ k.transpose(0, 1, 3, 2)) * self.scale
        attn = mx.softmax(attn.astype(mx.float32), axis=-1).astype(q.dtype)
        out = attn @ v
        out = out.transpose(0, 2, 1, 3).reshape(B, N, C)
        return self.out_proj(out)


class CLIPMLP(nn.Module):
    """CLIP MLP with QuickGELU."""

    def __init__(self, hidden_size: int, ffn_hidden_size: int):
        super().__init__()
        self.fc1 = nn.Linear(hidden_size, ffn_hidden_size, bias=True)
        self.fc2 = nn.Linear(ffn_hidden_size, hidden_size, bias=True)

    def __call__(self, x):
        # QuickGELU: fc1 → QuickGELU → fc2
        h = self.fc1(x)
        h = h * mx.sigmoid(1.702 * h)
        return self.fc2(h)


class CLIPTransformerLayer(nn.Module):
    """CLIP transformer layer."""

    def __init__(self, hidden_size: int, num_heads: int, ffn_hidden_size: int, eps: float = 1e-5):
        super().__init__()
        self.layer_norm1 = nn.LayerNorm(hidden_size, eps=eps)
        self.self_attn = CLIPAttention(hidden_size, num_heads)
        self.layer_norm2 = nn.LayerNorm(hidden_size, eps=eps)
        self.mlp = CLIPMLP(hidden_size, ffn_hidden_size)

    def __call__(self, x):
        x = x + self.self_attn(self.layer_norm1(x))
        x = x + self.mlp(self.layer_norm2(x))
        return x


class CLIPVisionEmbeddings(nn.Module):
    """CLIP vision embeddings that takes SAM features as input."""

    def __init__(self, hidden_size: int = 1024, image_size: int = 224, patch_size: int = 14):
        super().__init__()
        self.embed_dim = hidden_size
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_patches = (image_size // patch_size) ** 2
        self.num_positions = self.num_patches + 1

        self.class_embedding = mx.zeros((hidden_size,))

        # Patch embedding (projects SAM features) - NHWC conv
        self.patch_embedding = nn.Conv2d(3, hidden_size, patch_size, stride=patch_size, bias=False)

        # Position embedding
        self.position_embedding = nn.Embedding(self.num_positions, hidden_size)
        self.position_ids = mx.arange(self.num_positions)[None, :]

    def __call__(self, pixel_values, patch_embeds=None):
        batch_size = pixel_values.shape[0]

        if patch_embeds is not None:
            # Use pre-computed SAM features
            # patch_embeds: [B, H, W, C] (NHWC from SAM)
            B, H, W, C = patch_embeds.shape
            patch_embeds = patch_embeds.reshape(B, H * W, C)
        else:
            # Use raw conv on NHWC input
            patch_embeds = self.patch_embedding(pixel_values)
            B, H, W, C = patch_embeds.shape
            patch_embeds = patch_embeds.reshape(B, H * W, C)

        class_embeds = mx.tile(self.class_embedding.reshape(1, 1, -1), (batch_size, 1, 1))
        embeddings = mx.concatenate([class_embeds, patch_embeds], axis=1)

        # Add position embeddings with interpolation
        pos_ids = self.position_ids[:, :embeddings.shape[1]]
        pos_embeds = self.position_embedding(pos_ids)
        embeddings = embeddings + pos_embeds

        return embeddings


class CLIPVisionTransformer(nn.Module):
    """CLIP-L vision transformer."""

    def __init__(self, config: VisionConfig):
        super().__init__()
        self.embeddings = CLIPVisionEmbeddings(
            hidden_size=config.clip_hidden_size,
            image_size=config.clip_image_size,
            patch_size=config.clip_patch_size,
        )
        self.pre_layrnorm = nn.LayerNorm(config.clip_hidden_size, eps=config.clip_layernorm_epsilon)
        self.transformer = nn.Sequential(*[
            CLIPTransformerLayer(
                config.clip_hidden_size,
                config.clip_num_heads,
                config.clip_ffn_hidden_size,
                eps=config.clip_layernorm_epsilon,
            )
            for _ in range(config.clip_num_layers)
        ])

    def __call__(self, pixel_values, patch_embeds=None):
        x = self.embeddings(pixel_values, patch_embeds)
        x = self.pre_layrnorm(x)
        x = self.transformer(x)
        return x


# =============================================================================
# Projector
# =============================================================================

class MlpProjector(nn.Module):
    """Linear projector from vision to language space."""

    def __init__(self, config: ProjectorConfig):
        super().__init__()
        self.layers = nn.Linear(config.input_dim, config.n_embed, bias=True)

    def __call__(self, x):
        return self.layers(x)


# =============================================================================
# Unlimited OCR Model
# =============================================================================

@dataclass
class ModelOutput:
    logits: mx.array
    past_key_values: Optional[List[Tuple[mx.array, mx.array]]] = None


class UnlimitedOCRModel(nn.Module):
    """Complete Unlimited-OCR model with vision + language.

    Architecture:
    Image → SAM-ViT-B → CLIP-L → Projector → DeepSeek-V2 MoE → Text
    """

    def __init__(self, config: UnlimitedOCRConfig):
        super().__init__()
        self.config = config

        # Vision
        self.sam_model = SAMVisionEncoder(config.vision)
        self.vision_model = CLIPVisionTransformer(config.vision)

        # Projector: 2048 → 1280
        self.projector = MlpProjector(config.projector)

        # Language
        self.language_model = DeepSeekModel(config.language)
        self.lm_head = nn.Linear(config.language.hidden_size, config.language.vocab_size, bias=False)

        # Image special tokens
        embed_std = 1.0 / math.sqrt(config.language.hidden_size)
        self.image_newline = mx.random.normal((config.language.hidden_size,)) * embed_std
        self.view_seperator = mx.random.normal((config.language.hidden_size,)) * embed_std

    def encode_images(self, images: mx.array, images_spatial_crop=None) -> List[mx.array]:
        """Encode images through vision encoder.

        Args:
            images: List of [patches, original] image tensors (in NCHW from preprocessing)
            images_spatial_crop: List of (width_crops, height_crops) tuples

        Returns:
            List of image feature tensors [N, hidden_size]
        """
        all_features = []

        for idx, image_pair in enumerate(images):
            patches = image_pair[0]  # [N, 3, 640, 640] NCHW
            image_ori = image_pair[1]  # [1, 3, 1024, 1024] NCHW

            has_patches = patches is not None and patches.shape[0] > 0

            # Convert to NHWC for MLX conv
            def to_nhwc(t):
                if t is None:
                    return None
                ndim = len(t.shape)
                if ndim == 4:
                    return t.transpose(0, 2, 3, 1)  # NCHW → NHWC
                return t

            patches_nhwc = to_nhwc(patches)
            image_ori_nhwc = to_nhwc(image_ori)

            if has_patches and images_spatial_crop is not None:
                crop_shape = images_spatial_crop[idx]
                width_crop_num, height_crop_num = crop_shape

                # Process patches (local features)
                sam_local = self.sam_model(patches_nhwc)  # [P, 16, 16, 1024]
                clip_local = self.vision_model(patches_nhwc, sam_local)  # [P, 257, 1024]

                # Combine: CLIP[:, 1:] + SAM flatten
                # SAM: [P, 16, 16, 1024] → [P, 256, 1024]
                sam_flat = sam_local.reshape(patches.shape[0], -1, 1024)
                local_feats = mx.concatenate([
                    clip_local[:, 1:, :],  # [P, 256, 1024]
                    sam_flat,  # [P, 256, 1024]
                ], axis=-1)  # [P, 256, 2048]
                local_feats = self.projector(local_feats)  # [P, 256, 1280]

                # Process original (global features)
                sam_global = self.sam_model(image_ori_nhwc)  # [1, 16, 16, 1024]
                clip_global = self.vision_model(image_ori_nhwc, sam_global)  # [1, 257, 1024]

                sam_gflat = sam_global.reshape(1, -1, 1024)
                global_feats = mx.concatenate([
                    clip_global[:, 1:, :],  # [1, 256, 1024]
                    sam_gflat,  # [1, 256, 1024]
                ], axis=-1)  # [1, 256, 2048]
                global_feats = self.projector(global_feats)  # [1, 256, 1280]

                # Reshape and organize
                _, hw_g, nd = global_feats.shape
                h_g = w_g = int(hw_g ** 0.5)

                _, hw_l, nd2 = local_feats.shape
                h_l = w_l = int(hw_l ** 0.5)

                # Global: reshape to 2D and add newlines
                gf = global_feats.reshape(h_g, w_g, nd)
                gf = mx.concatenate([gf, mx.tile(self.image_newline[None, None, :], (h_g, 1, 1))], axis=1)
                gf = gf.reshape(-1, nd)

                # Local: reshape grid
                lf = local_feats.reshape(height_crop_num, width_crop_num, h_l, w_l, nd2)
                lf = lf.transpose(0, 2, 1, 3, 4).reshape(height_crop_num * h_l, width_crop_num * w_l, nd2)
                lf = mx.concatenate([lf, mx.tile(self.image_newline[None, None, :], (height_crop_num * h_l, 1, 1))], axis=1)
                lf = lf.reshape(-1, nd2)

                # Concat: local + global + separator
                full_feats = mx.concatenate([lf, gf, self.view_seperator[None, :]], axis=0)
                all_features.append(full_feats)

            else:
                # Multiple images or single image without crop
                if len(image_ori_nhwc.shape) == 3:
                    image_ori_nhwc = image_ori_nhwc[None, :, :, :]

                num_imgs = image_ori_nhwc.shape[0]
                for i in range(num_imgs):
                    img = image_ori_nhwc[i:i+1]
                    sam_out = self.sam_model(img)
                    clip_out = self.vision_model(img, sam_out)

                    sam_flat = sam_out.reshape(1, -1, 1024)
                    gf = mx.concatenate([
                        clip_out[:, 1:, :],
                        sam_flat,
                    ], axis=-1)
                    gf = self.projector(gf)

                    _, hw, nd = gf.shape
                    h = w = int(hw ** 0.5)

                    gf_2d = gf.reshape(h, w, nd)
                    gf_2d = mx.concatenate([gf_2d, mx.tile(self.image_newline[None, None, :], (h, 1, 1))], axis=1)
                    gf_2d = gf_2d.reshape(-1, nd)

                    full_feats = mx.concatenate([gf_2d, self.view_seperator[None, :]], axis=0)
                    all_features.append(full_feats)

        return all_features

    def __call__(
        self,
        input_ids: Optional[mx.array] = None,
        attention_mask: Optional[mx.array] = None,
        position_ids: Optional[mx.array] = None,
        past_key_values: Optional[List[Tuple[mx.array, mx.array]]] = None,
        inputs_embeds: Optional[mx.array] = None,
        images: Optional[List[mx.array]] = None,
        images_seq_mask: Optional[mx.array] = None,
        images_spatial_crop: Optional[List[Tuple[int, int]]] = None,
        use_cache: bool = False,
    ) -> ModelOutput:
        B = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]

        if inputs_embeds is None:
            inputs_embeds = self.language_model.embed_tokens(input_ids)

        # Inject image features into embeddings
        if images is not None and images_seq_mask is not None:
            image_features = self.encode_images(images, images_spatial_crop)

            rows = []
            for idx, img_feats in enumerate(image_features):
                emb = inputs_embeds[idx]  # [L, D]
                if img_feats is not None and img_feats.shape[0] > 0:
                    mask = images_seq_mask[idx]  # [L] bool
                    # first True position (image feature block is contiguous)
                    start = int(mx.argmax(mask.astype(mx.int32)).item())
                    n = img_feats.shape[0]
                    # replace the [start:start+n] rows with image features
                    emb = mx.concatenate([emb[:start], img_feats, emb[start + n:]], axis=0)
                rows.append(emb[None])
            inputs_embeds = mx.concatenate(rows, axis=0)

        hidden_states, new_kv = self.language_model(
            input_ids=None,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
        )

        logits = self.lm_head(hidden_states)
        return ModelOutput(logits=logits, past_key_values=new_kv)

    def generate(
        self,
        input_ids: mx.array,
        images: Optional[List] = None,
        images_seq_mask: Optional[mx.array] = None,
        images_spatial_crop: Optional[List] = None,
        max_length: int = 32768,
        temperature: float = 0.0,
        eos_token_id: int = 1,
    ) -> mx.array:
        """Autoregressive text generation."""
        generated = [input_ids]
        past_kv = None
        use_images = (images is not None)

        def step_fn(ids, pkv, first):
            if first:
                out = self(
                    input_ids=ids,
                    images=images if use_images else None,
                    images_seq_mask=images_seq_mask if use_images else None,
                    images_spatial_crop=images_spatial_crop if use_images else None,
                    use_cache=True,
                )
            else:
                out = self(input_ids=ids, past_key_values=pkv, use_cache=True)
            logits = out.logits[:, -1, :]
            if temperature > 0:
                probs = mx.softmax((logits / temperature).astype(mx.float32), axis=-1)
                nt = mx.random.categorical(probs, axis=-1).reshape(1, 1)
            else:
                nt = mx.argmax(logits, axis=-1, keepdims=True)
            return nt, out.past_key_values

        # 流水线解码：GPU 算下一步时不阻塞等 CPU 检查 EOS（async_eval）
        next_token, past_kv = step_fn(input_ids, None, True)
        mx.async_eval(next_token, *[kv for pair in past_kv for kv in pair])
        for step in range(max_length):
            cur = next_token
            if step + 1 < max_length:
                next_token, past_kv = step_fn(cur, past_kv, False)
                mx.async_eval(next_token, *[kv for pair in past_kv for kv in pair])
            tok_id = cur.item()  # 此时 cur 已在上一轮 async_eval 中求值完毕
            if tok_id == eos_token_id:
                break
            generated.append(cur)

        return mx.concatenate(generated, axis=1)

    def generate_batch(
        self,
        input_ids: mx.array,
        images: Optional[List] = None,
        images_seq_mask: Optional[mx.array] = None,
        images_spatial_crop: Optional[List] = None,
        max_length: int = 8192,
        eos_token_id: int = 1,
    ) -> mx.array:
        """批量贪心解码。要求 batch 内各行序列等长（同 crop_shape 分桶即可）。

        返回 [B, T] 的生成 token（不含 prompt）；各行遇 EOS 后填充 EOS，调用方按行截断。
        """
        B = input_ids.shape[0]
        eos = mx.array(eos_token_id, dtype=input_ids.dtype)

        # Prefill
        out = self(
            input_ids=input_ids,
            images=images,
            images_seq_mask=images_seq_mask,
            images_spatial_crop=images_spatial_crop,
            use_cache=True,
        )
        past_kv = out.past_key_values
        next_token = mx.argmax(out.logits[:, -1, :], axis=-1).astype(input_ids.dtype)  # [B]
        mx.async_eval(next_token, *[kv for pair in past_kv for kv in pair])

        tokens = [next_token]
        finished = (next_token == eos)
        for step in range(max_length - 1):
            if bool(mx.all(finished).item()):
                break
            out = self(input_ids=next_token[:, None], past_key_values=past_kv, use_cache=True)
            past_kv = out.past_key_values
            nt = mx.argmax(out.logits[:, -1, :], axis=-1).astype(input_ids.dtype)
            next_token = mx.where(finished, eos, nt)  # 已结束行锁定为 EOS
            tokens.append(next_token)
            finished = finished | (next_token == eos)
            mx.async_eval(next_token, *[kv for pair in past_kv for kv in pair])

        return mx.stack(tokens, axis=1)  # [B, T]

    # =====================================================================
    # Long-horizon (R-SWA) decoding
    # =====================================================================
    def embed_multipage(self, prefix_ids, page_features, suffix_ids):
        """拼接多页 prefill embedding：[prefix_ids] + 各页视觉特征 + [suffix_ids]。

        prefix_ids / suffix_ids: list[int]（如 [bos] 与 encode("\n"+prompt)）
        page_features: list of [N_i, D] —— 每页 encode_images 的输出
        返回 inputs_embeds [1, L_m, D]，图像特征已注入，可直接送入 generate_long。
        """
        emb = self.language_model.embed_tokens
        parts = []
        if prefix_ids:
            parts.append(emb(mx.array([prefix_ids], dtype=mx.int32))[0])
        parts.extend(page_features)  # 每个 [N_i, D]
        if suffix_ids:
            parts.append(emb(mx.array([suffix_ids], dtype=mx.int32))[0])
        return mx.concatenate(parts, axis=0)[None]  # [1, L_m, D]

    @staticmethod
    def _evict_kv(past_kv, prefix_len, window):
        """R-SWA 缓存淘汰：保留全部 prefix + 最近 window 个 decode token 的 KV。

        keys/values 中已包含绝对位置的 RoPE，因此裁剪任意 decode 槽都不影响正确性。
        """
        new = []
        for k, v in past_kv:
            seq = k.shape[2]
            if seq - prefix_len > window:
                keep = seq - window  # 起始裁剪点，保证保留最近 window 个
                k = mx.concatenate([k[:, :, :prefix_len, :], k[:, :, keep:, :]], axis=2)
                v = mx.concatenate([v[:, :, :prefix_len, :], v[:, :, keep:, :]], axis=2)
            new.append((k, v))
        return new

    def generate_long(
        self,
        inputs_embeds: mx.array,
        window: int = 128,
        max_length: int = 32768,
        temperature: float = 0.0,
        eos_token_id: int = 1,
    ) -> mx.array:
        """长程 R-SWA 贪心解码（batch=1）。

        - Prefill：对全部 inputs_embeds 做完整因果注意力（视觉+prompt 即受保护前缀 P）。
        - Decode：每步用绝对位置 position_ids，并在追加 KV 后淘汰窗口外的 decode KV，
          使 KV 缓存恒定为 L_m + window，与论文 R-SWA 一致。

        返回 [1, T] 生成的 token（不含 prefill）。
        """
        _, prefix_len, _ = inputs_embeds.shape
        pos = mx.arange(0, prefix_len, dtype=mx.int32)[None]

        hs, past_kv = self.language_model(
            inputs_embeds=inputs_embeds, position_ids=pos, use_cache=True,
        )
        logits = self.lm_head(hs[:, -1, :])

        def pick(lg):
            if temperature > 0:
                probs = mx.softmax((lg / temperature).astype(mx.float32), axis=-1)
                return mx.random.categorical(probs, axis=-1).reshape(1, 1)
            return mx.argmax(lg, axis=-1, keepdims=True)

        generated = []
        abs_pos = prefix_len
        for _ in range(max_length):
            nt = pick(logits)
            mx.eval(nt)
            if nt.item() == eos_token_id:
                break
            generated.append(nt)

            emb = self.language_model.embed_tokens(nt)  # [1, 1, D]
            step_pos = mx.array([[abs_pos]], dtype=mx.int32)
            hs, past_kv = self.language_model(
                inputs_embeds=emb, position_ids=step_pos,
                past_key_values=past_kv, use_cache=True,
            )
            past_kv = self._evict_kv(past_kv, prefix_len, window)
            logits = self.lm_head(hs[:, -1, :])
            abs_pos += 1

        if not generated:
            return mx.zeros((1, 0), dtype=mx.int32)
        return mx.concatenate(generated, axis=1)