from typing import List, Optional

import torch
import torch.nn as nn
import torch.distributed as dist

from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging
try:
    from .configuration_mimo_v2_pro import MiMoV2ProConfig
except ImportError:
    from configuration_mimo_v2_pro import MiMoV2ProConfig

logger = logging.get_logger(__name__)


class FakeRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(hidden_size))

    def forward(self, x):
        raise NotImplementedError("This is a fake model and does not support forward pass.")


class FakeMLP(nn.Module):
    def __init__(self, config: MiMoV2ProConfig, is_expert=False):
        super().__init__()
        if is_expert:
            self.intermediate_size = config.moe_intermediate_size
        else:
            self.intermediate_size = config.intermediate_size
        self.hidden_size = config.hidden_size

        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)

    def forward(self, x):
        raise NotImplementedError("This is a fake model and does not support forward pass.")


class FakeMoEGate(nn.Module):
    def __init__(self, config: MiMoV2ProConfig):
        super().__init__()
        self.weight = nn.Parameter(
            torch.empty((config.n_routed_experts, config.hidden_size), dtype=torch.float32)
        )
        if config.topk_method == "noaux_tc":
            self.e_score_correction_bias = nn.Parameter(
                torch.empty((config.n_routed_experts), dtype=torch.float32)
            )

    def forward(self, hidden_states):
        raise NotImplementedError("This is a fake model and does not support forward pass.")


class FakeMoE(nn.Module):
    def __init__(self, config: MiMoV2ProConfig):
        super().__init__()
        self.experts = nn.ModuleList(
            [FakeMLP(config, is_expert=True) for _ in range(config.n_routed_experts)]
        )
        self.gate = FakeMoEGate(config)

    def forward(self, hidden_states):
        raise NotImplementedError("This is a fake model and does not support forward pass.")


class FakeFullAttention(nn.Module):
    """Fake Attention for full attention layers (hybrid_layer_pattern == 0)."""

    def __init__(self, config: MiMoV2ProConfig):
        super().__init__()
        q_size = config.num_attention_heads * config.head_dim
        k_size = config.num_key_value_heads * config.head_dim
        v_size = config.num_key_value_heads * config.v_head_dim
        o_size = config.num_attention_heads * config.v_head_dim

        self.qkv_proj = nn.Linear(
            config.hidden_size, q_size + k_size + v_size, bias=config.attention_bias
        )
        self.o_proj = nn.Linear(o_size, config.hidden_size, bias=False)

        if getattr(config, "add_full_attention_sink_bias", False):
            self.attention_sink_bias = nn.Parameter(
                torch.empty(config.num_attention_heads)
            )

    def forward(self, *args, **kwargs):
        raise NotImplementedError("This is a fake model and does not support forward pass.")


class FakeSWAAttention(nn.Module):
    """Fake Attention for SWA layers (hybrid_layer_pattern == 1)."""

    def __init__(self, config: MiMoV2ProConfig):
        super().__init__()
        q_size = config.swa_num_attention_heads * config.swa_head_dim
        k_size = config.swa_num_key_value_heads * config.swa_head_dim
        v_size = config.swa_num_key_value_heads * config.swa_v_head_dim
        o_size = config.swa_num_attention_heads * config.swa_v_head_dim

        self.qkv_proj = nn.Linear(
            config.hidden_size, q_size + k_size + v_size, bias=config.attention_bias
        )
        self.o_proj = nn.Linear(o_size, config.hidden_size, bias=False)

        if getattr(config, "add_swa_attention_sink_bias", False):
            self.attention_sink_bias = nn.Parameter(
                torch.empty(config.swa_num_attention_heads)
            )

    def forward(self, *args, **kwargs):
        raise NotImplementedError("This is a fake model and does not support forward pass.")


class FakeDecoderLayer(nn.Module):
    def __init__(self, config: MiMoV2ProConfig, layer_idx: int):
        super().__init__()
        self.input_layernorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
        self.post_attention_layernorm = FakeRMSNorm(
            config.hidden_size, eps=config.layernorm_epsilon
        )

        is_moe = (
            getattr(config, "n_routed_experts", None) is not None
            and config.moe_layer_freq[layer_idx]
        )
        self.mlp = FakeMoE(config) if is_moe else FakeMLP(config)

        is_swa = config.hybrid_layer_pattern[layer_idx] == 1
        self.self_attn = FakeSWAAttention(config) if is_swa else FakeFullAttention(config)

    def forward(self, *args, **kwargs):
        raise NotImplementedError("This is a fake model and does not support forward pass.")


class FakeMTPLayer(nn.Module):
    def __init__(self, config: MiMoV2ProConfig):
        super().__init__()
        self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
        self.enorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
        self.hnorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
        self.final_layernorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
        self.input_layernorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
        self.pre_mlp_layernorm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
        self.self_attn = FakeSWAAttention(config)
        self.mlp = FakeMLP(config)

    def forward(self, *args, **kwargs):
        raise NotImplementedError("This is a fake model and does not support forward pass.")


class MiMoV2ProModel(PreTrainedModel):
    config_class = MiMoV2ProConfig

    def __init__(self, config: MiMoV2ProConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            [FakeDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
        )
        self.norm = FakeRMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
        self.mtp = nn.ModuleDict(
            {"layers": nn.ModuleList([FakeMTPLayer(config)])}
        )

    def forward(self, *args, **kwargs):
        raise NotImplementedError("This is a fake model and does not support forward pass.")


class MiMoV2ProForCausalLM(PreTrainedModel):
    config_class = MiMoV2ProConfig
    _keys_to_ignore_on_load_unexpected = [
        r"model.layers\.\d+\.self_attn\.rotary_emb\.inv_freq"
    ]

    def __init__(self, config: MiMoV2ProConfig):
        super().__init__(config)
        self.model = MiMoV2ProModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

    def forward(self, *args, **kwargs):
        raise NotImplementedError("This is a fake model and does not support forward pass.")