# Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Weight remapping from the official Mk1 safetensors export to MLX."""

from __future__ import annotations

from collections.abc import Iterable

import mlx.core as mx

# PyTorch → MLX
#   *.experts.down_proj              *.experts.down_proj.weight
#   *.experts.gate_up_proj           *.experts.gate_up_proj.weight
#   *.in_proj_weight [3D, D]         query/key/value_proj.weight
#   *.in_proj_bias [3D]              query/key/value_proj.bias
#   token_ff.0 / token_ff.2          token_ff.layers.0 / token_ff.layers.2
#   memory_ff.0 / memory_ff.2        memory_ff.layers.0 / memory_ff.layers.2

_ATTENTION_MODULES = (
    "local_attention",
    "token_memory_attention",
    "memory_token_attention",
)

_SKIP_SUBSTRINGS = (
    "rotary_emb",
    "lm_head.weight",
)

_CLIP_MARKERS = ("input_max", "input_min", "output_max", "output_min")


def should_keep_source_key(key: str) -> bool:
    if any(marker in key for marker in _SKIP_SUBSTRINGS):
        return False
    if key.startswith("model.encoder.language_model.") and not key.endswith(
        ".layer_scalar"
    ):
        return False
    if key.startswith("model.encoder.vision_tower.") or key.startswith(
        "model.encoder.embed_vision."
    ):
        if any(marker in key for marker in _CLIP_MARKERS):
            return False
    return True


def _split_qkv(prefix: str, value: mx.array) -> list[tuple[str, mx.array]]:
    if value.ndim == 1:
        width = value.shape[0]
        if width % 3:
            raise ValueError(f"Cannot split QKV bias for {prefix}: shape {value.shape}")
        head = width // 3
        pieces = (value[:head], value[head : 2 * head], value[2 * head :])
        names = ("query_proj.bias", "key_proj.bias", "value_proj.bias")
    elif value.ndim == 2:
        width = value.shape[0]
        if width % 3:
            raise ValueError(
                f"Cannot split QKV weight for {prefix}: shape {value.shape}"
            )
        head = width // 3
        pieces = (value[:head], value[head : 2 * head], value[2 * head :])
        names = ("query_proj.weight", "key_proj.weight", "value_proj.weight")
        if pieces[0].shape[0] != pieces[0].shape[1]:
            raise ValueError(
                f"Split QKV weight for {prefix} is not square: {pieces[0].shape}"
            )
    else:
        raise ValueError(f"Unexpected QKV tensor rank for {prefix}: {value.shape}")
    return [(f"{prefix}.{name}", piece) for name, piece in zip(names, pieces)]


def remap_weight(key: str, value: mx.array) -> list[tuple[str, mx.array]]:
    """Map one official Mk1 tensor onto one or more MLX parameter names."""

    if not should_keep_source_key(key):
        return []

    if key.endswith(".experts.down_proj"):
        return [(key + ".weight", value)]
    if key.endswith(".experts.gate_up_proj"):
        return [(key + ".weight", value)]

    for module in _ATTENTION_MODULES:
        in_proj_weight = f".{module}.in_proj_weight"
        in_proj_bias = f".{module}.in_proj_bias"
        if key.endswith(in_proj_weight):
            prefix = key[: -len(".in_proj_weight")]
            return _split_qkv(prefix, value)
        if key.endswith(in_proj_bias):
            prefix = key[: -len(".in_proj_bias")]
            return _split_qkv(prefix, value)

    if ".token_ff.0." in key or key.endswith(".token_ff.0.weight") or key.endswith(
        ".token_ff.0.bias"
    ):
        return [(key.replace(".token_ff.0.", ".token_ff.layers.0."), value)]
    if ".token_ff.2." in key or key.endswith(".token_ff.2.weight") or key.endswith(
        ".token_ff.2.bias"
    ):
        return [(key.replace(".token_ff.2.", ".token_ff.layers.2."), value)]
    if ".memory_ff.0." in key or key.endswith(".memory_ff.0.weight") or key.endswith(
        ".memory_ff.0.bias"
    ):
        return [(key.replace(".memory_ff.0.", ".memory_ff.layers.0."), value)]
    if ".memory_ff.2." in key or key.endswith(".memory_ff.2.weight") or key.endswith(
        ".memory_ff.2.bias"
    ):
        return [(key.replace(".memory_ff.2.", ".memory_ff.layers.2."), value)]

    return [(key, value)]


def remap_state_dict(
    source: Iterable[tuple[str, mx.array]],
) -> dict[str, mx.array]:
    remapped: dict[str, mx.array] = {}
    for key, value in source:
        for new_key, new_value in remap_weight(key, value):
            if new_key in remapped:
                raise ValueError(f"Duplicate remapped key: {new_key}")
            remapped[new_key] = new_value
    return remapped