# Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Mk1 wrappers around the borrowed DiffusionGemma trunk layers.

The trunk implementation is imported as an internal dependency. The public
model type remains ``modilify_mk1``. Two Mk1-specific forwards are installed
on the constructed trunk:

* Router: softmax over all experts, then top-k and renormalize
* Canvas merge: RMS-capped latent residual, not previous-logit soft embeds
"""

from __future__ import annotations

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

from mlx_vlm.models.diffusion_gemma.language import (
    DiffusionGemma4Backbone,
    Router,
    geglu,
)

LATENT_RESIDUAL_RMS_RATIO_CAP = 0.5


class Mk1Router(Router):
    """Official router weights with Mk1 log-softmax / top-k routing."""

    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]:
        x = mx.fast.rms_norm(x, None, self.eps)
        x = x * self.scale * self._root_size
        scores = self.proj(x)
        probabilities = mx.softmax(scores, axis=-1, precise=True)
        top_k = self.config.top_k_experts
        indices = mx.argpartition(probabilities, kth=-top_k, axis=-1)[..., -top_k:]
        weights = mx.take_along_axis(probabilities, indices, axis=-1)
        weights = weights / mx.sum(weights, axis=-1, keepdims=True)
        weights = weights * self.per_expert_scale[indices]
        return indices, weights


def merge_latent_context(
    mapper: nn.Module,
    token_embeddings: mx.array,
    latent_context: mx.array | None,
    *,
    rms_ratio_cap: float = LATENT_RESIDUAL_RMS_RATIO_CAP,
) -> mx.array:
    """Apply the native self-conditioning bridge with Mk1 RMS capping."""

    context = (
        mx.zeros_like(token_embeddings)
        if latent_context is None
        else latent_context.astype(token_embeddings.dtype)
    )
    if context.shape != token_embeddings.shape:
        raise ValueError("Latent context must match the canvas embedding shape.")
    normalized = mapper.pre_norm(context)
    mapped = mapper.down_proj(
        geglu(mapper.gate_proj(normalized), mapper.up_proj(normalized))
    )
    mapped_rms = mx.sqrt(
        mx.mean(mx.square(mapped.astype(mx.float32)), axis=-1, keepdims=True)
    )
    token_rms = mx.sqrt(
        mx.mean(
            mx.square(token_embeddings.astype(mx.float32)), axis=-1, keepdims=True
        )
    )
    cap = rms_ratio_cap * token_rms
    scale = cap / mx.sqrt(mx.square(mapped_rms) + mx.square(cap) + 1.0e-12)
    mapped = mapped * scale.astype(mapped.dtype)
    return mapper.post_norm(token_embeddings + mapped)


def _install_mk1_embed_canvas(decoder: nn.Module) -> None:
    def _embed_canvas(
        canvas_ids,
        self_conditioning_logits=None,
        self_conditioning_embeddings=None,
    ):
        if self_conditioning_logits is not None:
            raise ValueError(
                "Modilify Mk1 uses latent embeddings, not logits self-conditioning."
            )
        token_embeddings = decoder.embed_tokens(canvas_ids) * decoder.embed_scale
        return merge_latent_context(
            decoder.self_conditioning,
            token_embeddings,
            self_conditioning_embeddings,
        )

    decoder._embed_canvas = _embed_canvas


def _install_mk1_routers(backbone: DiffusionGemma4Backbone) -> None:
    for layer in backbone.decoder.layers:
        replacement = Mk1Router(layer.router.config)
        replacement.update(layer.router.parameters())
        layer.router = replacement


def build_mk1_backbone(trunk_config) -> DiffusionGemma4Backbone:
    """Construct the borrowed trunk and install Mk1 forwards."""

    backbone = DiffusionGemma4Backbone(trunk_config)
    _install_mk1_routers(backbone)
    _install_mk1_embed_canvas(backbone.decoder)
    return backbone