# MLX port of bytedance/Lance modeling/lance/modeling_utils.py
# Original: Copyright (c) 2025 ByteDance Ltd. and/or its affiliates. Apache 2.0.

import math

import numpy as np
import mlx.core as mx
import mlx.nn as nn


# ---------------------------------------------------------------------------
# Sin-cos position embedding tables (init-time numpy, frozen at runtime)
# ---------------------------------------------------------------------------

def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
    assert embed_dim % 2 == 0
    omega = np.arange(embed_dim // 2, dtype=np.float64)
    omega /= embed_dim / 2.0
    omega = 1.0 / 10000 ** omega
    pos = pos.reshape(-1)
    out = np.einsum("m,d->md", pos, omega)
    return np.concatenate([np.sin(out), np.cos(out)], axis=1)


def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
    assert embed_dim % 2 == 0
    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
    return np.concatenate([emb_h, emb_w], axis=1)


def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False, extra_tokens: int = 0) -> np.ndarray:
    grid_h = np.arange(grid_size, dtype=np.float32)
    grid_w = np.arange(grid_size, dtype=np.float32)
    grid = np.stack(np.meshgrid(grid_w, grid_h), axis=0)
    grid = grid.reshape([2, 1, grid_size, grid_size])
    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
    if cls_token and extra_tokens > 0:
        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
    return pos_embed


def get_3d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
    assert embed_dim % 2 == 0, "Embedding dimension must be even for 3D embeddings"
    d = embed_dim // 3
    d = d if d % 2 == 0 else d - 1
    dim_t, dim_h = d, d
    dim_w = embed_dim - 2 * d
    assert dim_w % 2 == 0
    emb_t = get_1d_sincos_pos_embed_from_grid(dim_t, grid[0])
    emb_h = get_1d_sincos_pos_embed_from_grid(dim_h, grid[1])
    emb_w = get_1d_sincos_pos_embed_from_grid(dim_w, grid[2])
    return np.concatenate([emb_t, emb_h, emb_w], axis=1)


def get_3d_sincos_pos_embed(embed_dim: int, t: int, h: int, w: int) -> np.ndarray:
    grid_t = np.arange(t, dtype=np.float32)
    grid_h = np.arange(h, dtype=np.float32)
    grid_w = np.arange(w, dtype=np.float32)
    tt, hh, ww = np.meshgrid(grid_t, grid_h, grid_w, indexing="ij")
    grid = np.stack([tt, hh, ww], axis=0)
    return get_3d_sincos_pos_embed_from_grid(embed_dim, grid)


# ---------------------------------------------------------------------------
# Activation lookup (ACT2FN equivalent for the subset Lance uses)
# ---------------------------------------------------------------------------

def _gelu_pytorch_tanh(x: mx.array) -> mx.array:
    # Matches torch.nn.functional.gelu(x, approximate="tanh"), which is the
    # default for "gelu_pytorch_tanh" in transformers' ACT2FN.
    return 0.5 * x * (1.0 + mx.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * x * x * x)))


ACT2FN = {
    "gelu": nn.gelu,
    "gelu_pytorch_tanh": _gelu_pytorch_tanh,
    "gelu_new": _gelu_pytorch_tanh,
    "silu": nn.silu,
    "swish": nn.silu,
    "relu": nn.relu,
}


# ---------------------------------------------------------------------------
# Timestep embedder (DiT-style)
# ---------------------------------------------------------------------------

class TimestepEmbedder(nn.Module):
    """Embeds scalar (possibly fractional) timesteps into hidden_size vectors.

    PT checkpoint uses nn.Sequential, producing param names mlp.0.{weight,bias}
    and mlp.2.{weight,bias}. The convert_weights tool maps mlp.0 -> fc1 and
    mlp.2 -> fc2 when loading the Lance safetensors.
    """

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

    @staticmethod
    def timestep_embedding(t: mx.array, dim: int, max_period: float = 10000.0) -> mx.array:
        half = dim // 2
        freqs = mx.exp(
            -math.log(max_period) * mx.arange(0, half, dtype=mx.float32) / half
        )
        args = t.astype(mx.float32)[:, None] * freqs[None]
        embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
        if dim % 2:
            embedding = mx.concatenate([embedding, mx.zeros_like(embedding[:, :1])], axis=-1)
        return embedding

    def __call__(self, t: mx.array) -> mx.array:
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        h = self.fc1(t_freq)
        h = nn.silu(h)
        h = self.fc2(h)
        return h


# ---------------------------------------------------------------------------
# MLP connector (vision -> LLM hidden space)
# ---------------------------------------------------------------------------

class MLPconnector(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, hidden_act: str):
        super().__init__()
        if hidden_act not in ACT2FN:
            raise ValueError(f"Unsupported activation: {hidden_act!r}")
        self._act_name = hidden_act
        self.fc1 = nn.Linear(in_dim, out_dim)
        self.fc2 = nn.Linear(out_dim, out_dim)

    def __call__(self, hidden_states: mx.array) -> mx.array:
        h = self.fc1(hidden_states)
        h = ACT2FN[self._act_name](h)
        h = self.fc2(h)
        return h


# ---------------------------------------------------------------------------
# Frozen sin-cos position embedding tables (2D + 3D)
# ---------------------------------------------------------------------------

class PositionEmbedding(nn.Module):
    """2D sin-cos lookup table.

    Stored as `pos_embed` to match PT param name. Initialized to sin-cos
    values; checkpoint load overwrites with identical values (PT also stores
    the initialized table as a requires_grad=False Parameter).
    """

    def __init__(self, max_num_patch_per_side: int, hidden_size: int):
        super().__init__()
        self.max_num_patch_per_side = max_num_patch_per_side
        self.hidden_size = hidden_size
        table = get_2d_sincos_pos_embed(hidden_size, max_num_patch_per_side).astype(np.float32)
        self.pos_embed = mx.array(table)

    def __call__(self, position_ids: mx.array) -> mx.array:
        return self.pos_embed[position_ids]


class PositionEmbedding3D(nn.Module):
    """3D sin-cos lookup table over (max_t * max_h * max_w)."""

    def __init__(self, max_latent_num_frames: int, max_latent_size: int, hidden_size: int):
        super().__init__()
        self.max_num_latent_frames = max_latent_num_frames
        self.max_latent_size = max_latent_size
        self.hidden_size = hidden_size
        table = get_3d_sincos_pos_embed(
            hidden_size,
            max_latent_num_frames,
            max_latent_size,
            max_latent_size,
        ).astype(np.float32)
        self.pos_embed = mx.array(table)

    def __call__(self, position_ids: mx.array) -> mx.array:
        return self.pos_embed[position_ids]