# 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 GLM4MoEGPTQ(BaseQModel):
    shared_input_verified_model_types = frozenset({"glm4_moe"})

    # GLM-4.5-Air MoE Model Structure:
    # Layer 0: Standard MLP (no MoE experts) - handled by ["mlp.down_proj"], ["mlp.gate_proj"], ["mlp.up_proj"]
    # Layers 1-46: MoE with shared_experts and individual experts (128 experts total) - handled by MoE components
    # Layer 46: Additional special structure with expanded parameters (embed_tokens, shared_head, eh_proj, etc.)
    #   This is handled dynamically through layer_modules_strict = False
    #
    # allow dynamic expert index for layer_modules so we don't need to write out 128 layers here
    # config.n_routed_experts contains the actual expert count used for index
    dynamic_expert_index = "n_routed_experts"

    pre_lm_head_norm_module = "model.norm"

    # Set to False since GLM-4.5-Air may have dynamic module structures
    layer_modules_strict = False

    out_of_model_tensors = {"files": ["mtp.safetensors"]}

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

    module_tree = [
        "model",
        "layers",
        "#",
        {
            "input_layernorm": ("input_layernorm:!",),
            "self_attn": ("q_proj:0:in=x", "q_norm:0:!", "k_proj:0:in=x", "k_norm:0:!", "v_proj:0:in=x", "o_proj:1"),
            "post_attention_layernorm": ("post_attention_layernorm:!",),
            "mlp:moe": {  # MoE module - can be Glm4MoeMLP (layer 0) or Glm4MoeMoE (layers 1-46)
                "gate": ("gate:!",),
                # Glm4MoeTopKRouter, ~1.6MB float32 per layer.  We really do not quant to quantize this.
                "experts": {
                    "#": ("gate_proj:0:in=x", "up_proj:0:in=x", "down_proj:1"),
                },
                "shared_experts": {
                    "gate_proj": ("gate_proj:0:in=x",),
                    "up_proj": ("up_proj:0:in=x",),
                    "down_proj": ("down_proj:1",),
                },
                # Standard MLP components for layer 0
                "": ("gate_proj:0:in=x", "up_proj:0:in=x", "down_proj:1"),
            },
        }
    ]