# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
# SPDX-License-Identifier: Apache-2.0
# Contact: qubitium@modelcloud.ai, x.com/qubitium

from ..base import BaseQModel
from ..moe_lifecycle import GateUpDownMoELifecycleHooks


class DeepSeekV3QModel(BaseQModel):
    shared_input_verified_model_types = frozenset({"deepseek_v3"})

    # deepseek_v3 requires custom model code
    require_trust_remote_code = True

    require_fast_init = False

    # allow dynamic expert index for layer_modules so we don't need to write out 64 layers here
    # config.num_experts contains the actual expert count used for index
    dynamic_expert_index = "n_routed_experts"

    pre_lm_head_norm_module = "model.norm"

    # MoE lifecycle hooks for gate_proj/up_proj/down_proj pattern
    moe_lifecycle_hooks = GateUpDownMoELifecycleHooks()

    # DeepSeek V3 uses dynamic modules based on lora(rank):
    layer_modules_strict = False

    module_tree = [
        "model",
        "layers",
        "#",
        {
            "input_layernorm": ("input_layernorm:!",),
            # DeepSeek-V3:
            #   self_attn uses split Q projections:
            #     - q_a_proj
            #     - q_b_proj
            #
            # Moonlight:
            #   self_attn uses a single Q projection:
            #     - q_proj
            "self_attn": ("q_proj:0", "q_a_proj:0:in=h", "kv_a_proj_with_mqa:0:in=h", "q_b_proj:1:in=q_a", "kv_b_proj:1:in=kv_a", "o_proj:2"),
            "post_attention_layernorm": ("post_attention_layernorm:!",),
            "mlp:moe": {
                "": ("gate_proj:0:in=x", "up_proj:0:in=x", "down_proj:1"),
                "experts": {
                    "#": ("gate_proj:0:in=x", "up_proj:0:in=x", "down_proj:1"),
                },
                "shared_experts": ("gate_proj:0:in=x", "up_proj:0:in=x", "down_proj:1"),
            },
        }
    ]