import math
from typing import List, Optional, Tuple

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

from .config import LMConfig


class RMSNorm(nn.Module):
    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)


class MiniCPMLongRoPE(nn.Module):
    def __init__(self, config: LMConfig):
        super().__init__()
        self.config = config
        # Use kv_channels for head dim if specified
        self.dim = (
            config.kv_channels
            if config.kv_channels
            else config.hidden_size // config.num_attention_heads
        )

        self.base = config.rope_theta
        self.max_position_embeddings = config.max_position_embeddings
        self.original_max_position_embeddings = config.original_max_position_embeddings

        self.short_factor = mx.array(config.rope_short_factor)
        self.long_factor = mx.array(config.rope_long_factor)

        scale = self.max_position_embeddings / self.original_max_position_embeddings
        self.scaling_factor = math.sqrt(
            1
            + math.log(max(scale, 1.0))
            / math.log(self.original_max_position_embeddings)
        )

        half_dim = self.dim // 2
        exponents = mx.arange(0, half_dim, dtype=mx.float32) / half_dim
        inv_freq = 1.0 / (self.base**exponents)
        self.inv_freq = inv_freq

    def __call__(self, position_ids):
        seq_len = position_ids.max().item() + 1

        factors = (
            self.long_factor
            if seq_len > self.original_max_position_embeddings
            else self.short_factor
        )

        t = mx.arange(seq_len, dtype=mx.float32)

        freqs = (t[:, None] * (1.0 / factors[None, :])) * self.inv_freq[None, :]
        emb = mx.concatenate([freqs, freqs], axis=-1)

        cos = mx.cos(emb) * self.scaling_factor
        sin = mx.sin(emb) * self.scaling_factor

        return cos[position_ids], sin[position_ids]


def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return mx.concatenate([-x2, x1], axis=-1)


def apply_rotary_pos_emb(q, k, cos, sin):
    cos = cos[:, :, None, :]
    sin = sin[:, :, None, :]

    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class Attention(nn.Module):
    def __init__(self, config: LMConfig):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        # Use kv_channels for head dim if specified
        self.head_dim = (
            config.kv_channels
            if config.kv_channels
            else config.hidden_size // self.num_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_kv_heads * self.head_dim, bias=False
        )
        self.v_proj = nn.Linear(
            config.hidden_size, self.num_kv_heads * self.head_dim, bias=False
        )
        self.o_proj = nn.Linear(
            self.num_heads * self.head_dim, config.hidden_size, bias=False
        )

    def __call__(self, x, cos, sin, mask=None, cache=None):
        B, L, _ = x.shape

        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)

        q = q.reshape(B, L, self.num_heads, self.head_dim)
        k = k.reshape(B, L, self.num_kv_heads, self.head_dim)
        v = v.reshape(B, L, self.num_kv_heads, self.head_dim)

        # Apply RoPE only if cos/sin are provided (no_rope=False)
        if cos is not None and sin is not None:
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        if cache is not None:
            k_cache, v_cache = cache
            k = mx.concatenate([k_cache, k], axis=1)
            v = mx.concatenate([v_cache, v], axis=1)

        new_cache = (k, v)

        q = q.transpose(0, 2, 1, 3)
        k = k.transpose(0, 2, 1, 3)
        v = v.transpose(0, 2, 1, 3)

        out = mx.fast.scaled_dot_product_attention(
            q, k, v, scale=1 / math.sqrt(self.head_dim), mask=mask
        )

        out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)
        return self.o_proj(out), new_cache


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

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


class MiniCPMDecoderLayer(nn.Module):
    def __init__(self, config: LMConfig):
        super().__init__()
        self.self_attn = Attention(config)
        self.mlp = MLP(config)
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

        self.scale_depth = config.scale_depth
        self.num_hidden_layers = config.num_hidden_layers
        self.use_mup = config.use_mup

    def __call__(self, x, cos, sin, mask=None, cache=None):
        r = x
        x = self.input_layernorm(x)
        h, new_cache = self.self_attn(x, cos, sin, mask, cache)

        if self.use_mup:
            x = r + h * (self.scale_depth / math.sqrt(self.num_hidden_layers))
        else:
            x = r + h

        r = x
        x = self.post_attention_layernorm(x)
        h = self.mlp(x)

        if self.use_mup:
            x = r + h * (self.scale_depth / math.sqrt(self.num_hidden_layers))
        else:
            x = r + h

        return x, new_cache


class MiniCPMModel(nn.Module):
    def __init__(self, config: LMConfig):
        super().__init__()
        self.config = config

        if config.vocab_size > 0:
            self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)

        self.layers = [
            MiniCPMDecoderLayer(config) for _ in range(config.num_hidden_layers)
        ]
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # Skip RoPE when no_rope=True (used for residual LM)
        if not config.no_rope:
            self.rope = MiniCPMLongRoPE(config)
        else:
            self.rope = None

    def __call__(
        self,
        inputs_embeds=None,
        input_ids=None,
        mask=None,
        cache=None,
        is_causal=True,
    ):
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        B, L, D = inputs_embeds.shape

        offset = 0
        if cache is not None:
            offset = cache[0][0].shape[1]

        # Compute RoPE if available
        if self.rope is not None:
            position_ids = mx.arange(offset, offset + L).astype(mx.int32)
            cos, sin = self.rope(position_ids)
            cos = cos[None, :, :]
            sin = sin[None, :, :]
        else:
            cos, sin = None, None

        if mask is None and is_causal and L > 1:
            causal_mask = mx.triu(mx.full((L, L), float("-inf")), k=1)
            mask = causal_mask[None, None, :, :]

        h = inputs_embeds
        new_caches = []

        for i, layer in enumerate(self.layers):
            layer_cache = cache[i] if cache is not None else None
            h, c = layer(h, cos, sin, mask=mask, cache=layer_cache)
            new_caches.append(c)

        h = self.norm(h)
        return h, new_caches