"""Native MLX implementation of LiquidAI/LFM2.5-Encoder-350M.

This is an encoder/masked-language model: attention and short convolutions are
bidirectional.  It intentionally does not use mlx-lm's causal generation API.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

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


@dataclass
class ModelArgs:
    vocab_size: int
    hidden_size: int
    num_hidden_layers: int
    num_attention_heads: int
    num_key_value_heads: int
    norm_eps: float
    conv_bias: bool
    conv_L_cache: int
    block_dim: int
    intermediate_size: int
    block_multiple_of: int
    block_ffn_dim_multiplier: float
    block_auto_adjust_ff_dim: bool
    layer_types: list[str]
    rope_theta: float = 1_000_000.0

    @classmethod
    def from_dict(cls, config: dict) -> "ModelArgs":
        rope = config.get("rope_parameters") or {}
        fields = {
            "vocab_size",
            "hidden_size",
            "num_hidden_layers",
            "num_attention_heads",
            "num_key_value_heads",
            "norm_eps",
            "conv_bias",
            "conv_L_cache",
            "block_dim",
            "intermediate_size",
            "block_multiple_of",
            "block_ffn_dim_multiplier",
            "block_auto_adjust_ff_dim",
            "layer_types",
        }
        values = {key: config[key] for key in fields}
        values["rope_theta"] = rope.get("rope_theta", config.get("rope_theta", 1_000_000.0))
        return cls(**values)


class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        dim = args.hidden_size
        self.n_heads = args.num_attention_heads
        self.n_kv_heads = args.num_key_value_heads
        self.head_dim = dim // self.n_heads
        self.scale = self.head_dim**-0.5

        self.q_layernorm = nn.RMSNorm(self.head_dim, eps=args.norm_eps)
        self.k_layernorm = nn.RMSNorm(self.head_dim, eps=args.norm_eps)
        self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
        self.out_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False)
        self.rope = nn.RoPE(self.head_dim, base=args.rope_theta, traditional=False)

    def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx.array:
        batch, length, _ = x.shape
        queries = self.q_proj(x).reshape(batch, length, self.n_heads, -1)
        keys = self.k_proj(x).reshape(batch, length, self.n_kv_heads, -1)
        values = self.v_proj(x).reshape(batch, length, self.n_kv_heads, -1)

        queries = self.q_layernorm(queries).transpose(0, 2, 1, 3)
        keys = self.k_layernorm(keys).transpose(0, 2, 1, 3)
        values = values.transpose(0, 2, 1, 3)
        queries = self.rope(queries)
        keys = self.rope(keys)

        output = mx.fast.scaled_dot_product_attention(
            queries, keys, values, scale=self.scale, mask=mask
        )
        output = output.transpose(0, 2, 1, 3).reshape(batch, length, -1)
        return self.out_proj(output)


class ShortConv(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.kernel_size = args.conv_L_cache
        self.conv = nn.Conv1d(
            in_channels=args.hidden_size,
            out_channels=args.hidden_size,
            kernel_size=self.kernel_size,
            groups=args.hidden_size,
            bias=args.conv_bias,
        )
        self.in_proj = nn.Linear(args.hidden_size, 3 * args.hidden_size, bias=args.conv_bias)
        self.out_proj = nn.Linear(args.hidden_size, args.hidden_size, bias=args.conv_bias)

    def __call__(self, x: mx.array, padding_mask: Optional[mx.array] = None) -> mx.array:
        if padding_mask is not None:
            x = mx.where(padding_mask[..., None], x, 0)
        gate_b, gate_c, value = mx.split(self.in_proj(x), 3, axis=-1)
        gated = gate_b * value
        pad = self.kernel_size // 2
        gated = mx.pad(gated, [(0, 0), (pad, pad), (0, 0)])
        conv_out = self.conv(gated)
        if conv_out.shape[1] > x.shape[1]:
            conv_out = conv_out[:, : x.shape[1], :]
        elif conv_out.shape[1] < x.shape[1]:
            conv_out = mx.pad(conv_out, [(0, 0), (0, x.shape[1] - conv_out.shape[1]), (0, 0)])
        return self.out_proj(gate_c * conv_out)


class MLP(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        ff_dim = args.intermediate_size
        if args.block_auto_adjust_ff_dim:
            ff_dim = int(2 * ff_dim / 3)
            ff_dim = int(args.block_ffn_dim_multiplier * ff_dim)
            multiple = args.block_multiple_of
            ff_dim = multiple * ((ff_dim + multiple - 1) // multiple)
        self.w1 = nn.Linear(args.block_dim, ff_dim, bias=False)
        self.w3 = nn.Linear(args.block_dim, ff_dim, bias=False)
        self.w2 = nn.Linear(ff_dim, args.block_dim, bias=False)

    def __call__(self, x: mx.array) -> mx.array:
        return self.w2(nn.silu(self.w1(x)) * self.w3(x))


class EncoderLayer(nn.Module):
    def __init__(self, args: ModelArgs, layer_type: str):
        super().__init__()
        self.is_attention_layer = layer_type == "full_attention"
        if self.is_attention_layer:
            self.self_attn = Attention(args)
        else:
            self.conv = ShortConv(args)
        self.feed_forward = MLP(args)
        self.operator_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
        self.ffn_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)

    def __call__(
        self,
        x: mx.array,
        attention_mask: Optional[mx.array],
        padding_mask: Optional[mx.array],
    ) -> mx.array:
        normalized = self.operator_norm(x)
        if self.is_attention_layer:
            residual = self.self_attn(normalized, mask=attention_mask)
        else:
            residual = self.conv(normalized, padding_mask=padding_mask)
        hidden = x + residual
        return hidden + self.feed_forward(self.ffn_norm(hidden))


class Lfm2BidirectionalModel(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
        self.layers = [EncoderLayer(args, layer_type) for layer_type in args.layer_types]
        self.embedding_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)

    def __call__(
        self, input_ids: mx.array, attention_mask: Optional[mx.array] = None
    ) -> mx.array:
        hidden = self.embed_tokens(input_ids)
        padding_mask = None
        additive_mask = None
        if attention_mask is not None:
            padding_mask = attention_mask.astype(mx.bool_)
            additive_mask = mx.where(
                padding_mask[:, None, None, :],
                mx.array(0.0, dtype=hidden.dtype),
                mx.array(-1e9, dtype=hidden.dtype),
            )
        for layer in self.layers:
            hidden = layer(hidden, additive_mask, padding_mask)
        return self.embedding_norm(hidden)


class Model(nn.Module):
    """Bidirectional encoder with a tied masked-language-model head."""

    def __init__(self, args: ModelArgs):
        super().__init__()
        self.args = args
        self.lfm2 = Lfm2BidirectionalModel(args)

    def __call__(
        self,
        input_ids: mx.array,
        attention_mask: Optional[mx.array] = None,
        *,
        return_hidden: bool = False,
    ) -> mx.array | tuple[mx.array, mx.array]:
        hidden = self.lfm2(input_ids, attention_mask=attention_mask)
        logits = self.lfm2.embed_tokens.as_linear(hidden)
        return (logits, hidden) if return_hidden else logits


def load_model(path: str | Path) -> Model:
    path = Path(path)
    with (path / "config.json").open() as handle:
        args = ModelArgs.from_dict(json.load(handle))
    model = Model(args)
    weights = mx.load(str(path / "model.safetensors"))
    model.load_weights(list(weights.items()), strict=True)
    mx.eval(model.parameters())
    return model