"""MLX-LM implementation for IFM/K2-Horizon-MoVA-36B-A4B.

Designed to be loaded through mlx-lm's custom `model_file` mechanism.
Reference architecture: vLLM's K2 Horizon implementation and the HF config.
"""

from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union
import math
import re

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

from mlx_lm.models.activations import swiglu
from mlx_lm.models.base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
from mlx_lm.models.switch_layers import SwitchGLU, SwitchLinear


@dataclass
class ModelArgs(BaseModelArgs):
    model_type: str
    vocab_size: int
    hidden_size: int
    intermediate_size: int
    moe_intermediate_size: int
    num_hidden_layers: int
    num_attention_heads: int
    num_key_value_heads: int
    head_dim: int
    max_position_embeddings: int
    rms_norm_eps: float
    layernorm_num_groups: int
    num_experts: int
    num_experts_per_tok: int
    num_shared_experts: int
    mova_num_experts: int
    mova_num_experts_per_tok: int
    decoder_sparse_step: int
    mlp_only_layers: List[int]
    norm_topk_prob: bool
    router_scaling_factor: float
    router_score_func: str
    moe_gate_bias: bool
    query_key_norm: bool
    rope_head_dim: int
    attention_gate_func: Optional[str]
    attention_bias: bool
    tie_word_embeddings: bool
    hidden_act: str = "silu"
    rope_parameters: Optional[Dict[str, Any]] = None
    rope_theta: Optional[float] = None
    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
    sliding_window: Optional[int] = None
    use_sliding_window: bool = False

    def __post_init__(self):
        if self.hidden_act != "silu":
            raise ValueError(f"K2 Horizon expects silu, got {self.hidden_act!r}")
        if self.router_score_func not in ("sigmoid", "softmax"):
            raise ValueError(f"Unsupported router score function: {self.router_score_func}")
        if self.rope_head_dim != self.head_dim:
            raise NotImplementedError(
                "This MLX port currently requires rope_head_dim == head_dim. "
                "The released K2-Horizon-MoVA-36B-A4B checkpoint uses 128 == 128."
            )
        if self.rope_theta is None:
            params = self.rope_parameters or {}
            self.rope_theta = float(params.get("rope_theta", 10000.0))


def _topk_indices(scores: mx.array, k: int) -> mx.array:
    # Match torch.topk selection semantics; ordering among selected experts is
    # irrelevant because we sum their weighted outputs.
    return mx.argpartition(scores, kth=-k, axis=-1)[..., -k:]


def _route(
    router_logits: mx.array,
    correction_bias: Optional[mx.array],
    score_func: str,
    top_k: int,
    scaling_factor: Optional[float],
    renormalize: bool,
):
    # vLLM computes routing probabilities in fp32, then applies the bias only
    # to expert *selection*. The unbiased probabilities become the weights.
    logits32 = router_logits.astype(mx.float32)
    if score_func == "softmax":
        routing_scores = mx.softmax(logits32, axis=-1, precise=True)
    elif score_func == "sigmoid":
        routing_scores = mx.sigmoid(logits32)
    else:
        raise ValueError(f"Unsupported router score function: {score_func}")

    selection_scores = routing_scores
    if correction_bias is not None:
        selection_scores = selection_scores + correction_bias.astype(mx.float32)

    inds = _topk_indices(selection_scores, top_k)
    inds = mx.stop_gradient(inds)
    weights = mx.take_along_axis(routing_scores, inds, axis=-1)
    if renormalize:
        weights = weights / mx.sum(weights, axis=-1, keepdims=True)
    if scaling_factor is not None:
        weights = weights * scaling_factor
    return weights.astype(router_logits.dtype), inds


def _softplus_beta(x: mx.array, beta: float) -> mx.array:
    # Stable softplus(beta*x) / beta.
    z = beta * x
    return (mx.maximum(z, 0) + mx.log1p(mx.exp(-mx.abs(z)))) / beta


class GroupedRMSNorm(nn.Module):
    """K2 Horizon grouped RMSNorm.

    Each hidden vector is split into `num_groups` contiguous groups and each
    group is independently RMS-normalized. A single full-width learned weight
    is applied afterwards.
    """

    def __init__(self, hidden_size: int, num_groups: int, eps: float):
        super().__init__()
        if hidden_size % num_groups != 0:
            raise ValueError("hidden_size must be divisible by num_groups")
        self.hidden_size = hidden_size
        self.num_groups = num_groups
        self.group_size = hidden_size // num_groups
        self.eps = eps
        self.weight = mx.ones((hidden_size,))

    def __call__(self, x: mx.array) -> mx.array:
        original_shape = x.shape
        grouped = x.reshape(*original_shape[:-1], self.num_groups, self.group_size)
        # Compute variance in fp32, matching the usual RMSNorm numerical path.
        variance = mx.mean(grouped.astype(mx.float32) ** 2, axis=-1, keepdims=True)
        grouped = grouped * mx.rsqrt(variance + self.eps).astype(grouped.dtype)
        return grouped.reshape(original_shape) * self.weight


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

    def __call__(self, x: mx.array) -> mx.array:
        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))


class SparseMoE(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.num_experts = args.num_experts
        self.top_k = args.num_experts_per_tok
        self.norm_topk_prob = args.norm_topk_prob
        self.score_func = args.router_score_func
        self.scaling_factor = args.router_scaling_factor

        # Important: checkpoint gate.bias is a correction bias for expert
        # selection. It must NOT be included in the linear router logits.
        self.gate = nn.Linear(args.hidden_size, args.num_experts, bias=False)
        self.gate_bias = (
            mx.zeros((args.num_experts,), dtype=mx.float32)
            if args.moe_gate_bias
            else None
        )
        self.experts = SwitchGLU(
            args.hidden_size,
            args.moe_intermediate_size,
            args.num_experts,
            bias=False,
        )
        self.shared_experts = (
            MLP(args, args.moe_intermediate_size * args.num_shared_experts)
            if args.num_shared_experts > 0
            else None
        )

    def __call__(self, x: mx.array) -> mx.array:
        logits = self.gate(x)
        scores, inds = _route(
            logits,
            self.gate_bias,
            self.score_func,
            self.top_k,
            self.scaling_factor,
            self.norm_topk_prob,
        )
        y = self.experts(x, inds)
        y = (y * scores[..., None]).sum(axis=-2).astype(x.dtype)
        if self.shared_experts is not None:
            y = y + self.shared_experts(x)
        return y


class Attention(nn.Module):
    def __init__(self, args: ModelArgs, use_mova: bool):
        super().__init__()
        self.n_heads = args.num_attention_heads
        self.n_kv_heads = args.num_key_value_heads
        self.head_dim = args.head_dim
        self.scale = self.head_dim ** -0.5
        self.use_mova = use_mova
        self.query_key_norm = args.query_key_norm
        self.gate_func = args.attention_gate_func

        self.q_proj = nn.Linear(
            args.hidden_size,
            self.n_heads * self.head_dim,
            bias=args.attention_bias,
        )
        self.k_proj = nn.Linear(
            args.hidden_size,
            self.n_kv_heads * self.head_dim,
            bias=args.attention_bias,
        )

        if use_mova:
            self.v_router = nn.Linear(
                args.hidden_size, args.mova_num_experts, bias=False
            )
            self.v_router_bias = (
                mx.zeros((args.mova_num_experts,), dtype=mx.float32)
                if args.moe_gate_bias
                else None
            )
            self.v_experts = SwitchLinear(
                args.hidden_size,
                self.n_kv_heads * self.head_dim,
                args.mova_num_experts,
                bias=False,
            )
            self.v_top_k = args.mova_num_experts_per_tok
            self.v_score_func = args.router_score_func
            self.v_scaling_factor = args.router_scaling_factor
        else:
            self.v_proj = nn.Linear(
                args.hidden_size,
                self.n_kv_heads * self.head_dim,
                bias=args.attention_bias,
            )

        self.o_proj = nn.Linear(
            self.n_heads * self.head_dim,
            args.hidden_size,
            bias=args.attention_bias,
        )

        if self.query_key_norm:
            # The checkpoint currently disables this, but keep support for
            # architecture completeness. K2 normalizes each head separately.
            self.q_norm = GroupedRMSNorm(
                self.n_heads * self.head_dim, self.n_heads, args.rms_norm_eps
            )
            self.k_norm = GroupedRMSNorm(
                self.n_kv_heads * self.head_dim, self.n_kv_heads, args.rms_norm_eps
            )

        self.rope = nn.RoPE(
            args.rope_head_dim,
            traditional=False,
            base=args.rope_theta,
        )

        if self.gate_func is not None:
            self.gate_proj = nn.Linear(
                args.hidden_size,
                self.n_heads * self.head_dim,
                bias=False,
            )

    def _mova_value(self, x: mx.array) -> mx.array:
        logits = self.v_router(x)
        weights, inds = _route(
            logits,
            self.v_router_bias,
            self.v_score_func,
            self.v_top_k,
            self.v_scaling_factor,
            True,  # vLLM default renormalize for top_k > 1
        )
        # gather_mm needs x as (..., 1, 1, in_dim); result is (..., k, 1, out_dim).
        values = self.v_experts(x[..., None, None, :], inds).squeeze(-2)
        # The reference applies silu to each value expert output before mixing.
        return (nn.silu(values) * weights[..., None]).sum(axis=-2).astype(x.dtype)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
        cache: Optional[Any] = None,
    ) -> mx.array:
        B, L, _ = x.shape

        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self._mova_value(x) if self.use_mova else self.v_proj(x)

        if self.query_key_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)

        q = q.reshape(B, L, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
        k = k.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
        v = v.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)

        if cache is not None:
            q = self.rope(q, offset=cache.offset)
            k = self.rope(k, offset=cache.offset)
            k, v = cache.update_and_fetch(k, v)
        else:
            q = self.rope(q)
            k = self.rope(k)

        out = scaled_dot_product_attention(
            q, k, v, cache=cache, scale=self.scale, mask=mask
        )
        out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)

        if self.gate_func is not None:
            gate = self.gate_proj(x)
            if self.gate_func == "silu":
                gate = nn.silu(gate)
            elif self.gate_func == "softplus":
                gate = _softplus_beta(gate, math.log(2.0))
            else:
                raise ValueError(f"Unsupported attention gate: {self.gate_func}")
            out = out * gate

        return self.o_proj(out)


class DecoderLayer(nn.Module):
    def __init__(self, args: ModelArgs, layer_idx: int):
        super().__init__()
        is_sparse = (
            layer_idx not in args.mlp_only_layers
            and args.num_experts > 0
            and (layer_idx + 1) % args.decoder_sparse_step == 0
        )
        self.self_attn = Attention(args, use_mova=is_sparse and args.mova_num_experts > 0)
        self.mlp = SparseMoE(args) if is_sparse else MLP(args, args.intermediate_size)
        self.input_layernorm = GroupedRMSNorm(
            args.hidden_size, args.layernorm_num_groups, args.rms_norm_eps
        )
        self.post_attention_layernorm = GroupedRMSNorm(
            args.hidden_size, args.layernorm_num_groups, args.rms_norm_eps
        )

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
        cache: Optional[Any] = None,
    ) -> mx.array:
        h = x + self.self_attn(self.input_layernorm(x), mask, cache)
        return h + self.mlp(self.post_attention_layernorm(h))


class K2HorizonModel(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.args = args
        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
        self.layers = [DecoderLayer(args, i) for i in range(args.num_hidden_layers)]
        self.norm = GroupedRMSNorm(
            args.hidden_size, args.layernorm_num_groups, args.rms_norm_eps
        )

    def __call__(
        self,
        inputs: mx.array,
        cache=None,
        input_embeddings: Optional[mx.array] = None,
    ) -> mx.array:
        h = self.embed_tokens(inputs) if input_embeddings is None else input_embeddings
        if cache is None:
            cache = [None] * len(self.layers)
        mask = create_attention_mask(h, cache[0])
        for layer, c in zip(self.layers, cache):
            h = layer(h, mask, c)
        return self.norm(h)


class Model(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.args = args
        self.model_type = args.model_type
        self.model = K2HorizonModel(args)
        if not args.tie_word_embeddings:
            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)

    def __call__(
        self,
        inputs: mx.array,
        cache=None,
        input_embeddings: Optional[mx.array] = None,
    ) -> mx.array:
        out = self.model(inputs, cache, input_embeddings)
        if self.args.tie_word_embeddings:
            return self.model.embed_tokens.as_linear(out)
        return self.lm_head(out)

    @property
    def layers(self):
        # mlx_lm generation (prompt cache construction) expects Model.layers.
        return self.model.layers

    def sanitize(self, weights):
        """Map the HF checkpoint's per-expert tensors to MLX Switch layers."""
        if self.args.tie_word_embeddings:
            weights.pop("lm_head.weight", None)

        # Router biases are correction biases, not Linear.bias parameters.
        for layer_idx in range(self.args.num_hidden_layers):
            base = f"model.layers.{layer_idx}"
            sparse = (
                layer_idx not in self.args.mlp_only_layers
                and self.args.num_experts > 0
                and (layer_idx + 1) % self.args.decoder_sparse_step == 0
            )
            if not sparse:
                continue

            gate_bias = f"{base}.mlp.gate.bias"
            if gate_bias in weights:
                weights[f"{base}.mlp.gate_bias"] = weights.pop(gate_bias)

            v_router_bias = f"{base}.self_attn.v_router.bias"
            if v_router_bias in weights:
                weights[f"{base}.self_attn.v_router_bias"] = weights.pop(v_router_bias)

            # Stack routed FFN experts: [E, out_dim, in_dim].
            for proj in ("gate_proj", "up_proj", "down_proj"):
                first = f"{base}.mlp.experts.0.{proj}.weight"
                if first in weights:
                    tensors = [
                        weights.pop(f"{base}.mlp.experts.{e}.{proj}.weight")
                        for e in range(self.args.num_experts)
                    ]
                    weights[f"{base}.mlp.experts.{proj}.weight"] = mx.stack(tensors)

            # Stack MoVA value experts: [E, kv_dim, hidden_dim].
            first_v = f"{base}.self_attn.v_experts.0.weight"
            if first_v in weights:
                tensors = [
                    weights.pop(f"{base}.self_attn.v_experts.{e}.weight")
                    for e in range(self.args.mova_num_experts)
                ]
                weights[f"{base}.self_attn.v_experts.weight"] = mx.stack(tensors)

        return weights

    @property
    def quant_predicate(self):
        def predicate(path, module):
            # Router matrices are tiny and routing is sensitive; keep them at
            # higher precision. Quantize the large expert/projection weights.
            if path.endswith("mlp.gate") or path.endswith("self_attn.v_router"):
                return False
            return True

        return predicate