import torch
from transformers.activations import ACT2FN

from .configuration_aliceai_t5 import AliceAIT5ModuleConfig
from .configuration_aliceai_t5_moe import AliceAIT5MoEConfig
from .modeling_aliceai_t5 import (
    AliceAIT5Decoder,
    AliceAIT5DecoderLayer,
    AliceAIT5Encoder,
    AliceAIT5EncoderLayer,
    AliceAIT5EncoderModel,
    AliceAIT5ForConditionalGeneration,
    AliceAIT5Model,
    AliceAIT5PreTrainedModel,
)
from .moe_layers import AdditionalArgs, Arguments, dMoE


def build_moe_arguments(model_config: AliceAIT5MoEConfig, is_encoder: bool = True) -> dict[str, object]:
    module_config = model_config.encoder if is_encoder else model_config.decoder
    return {
        "args": Arguments(
            hidden_size=module_config.hidden_size,
            ffn_hidden_size=model_config.expert_intermediate_size,
            activation_fn=ACT2FN[model_config.activation_type],
            moe_num_experts=model_config.n_experts,
            moe_top_k=model_config.experts_gating_top,
            moe_normalize_expert_weights=model_config.moe_normalize_expert_weights,
            routed_scaling_factor=model_config.routed_scaling_factor,
            bias=True,
        ),
        "additional_args": AdditionalArgs(
            moe_n_group=model_config.experts_n_group,
            moe_top_k_group=model_config.experts_top_k_group,
            group_routing=model_config.group_routing
            and model_config.experts_n_group is not None
            and model_config.experts_n_group > 1,
        ),
    }


class AliceAIT5MoEEncoderLayer(AliceAIT5EncoderLayer):
    def __init__(self, config: AliceAIT5MoEConfig, layer_idx: int):
        super().__init__(
            config.encoder,
            layer_idx=layer_idx,
            mlp=dMoE(**build_moe_arguments(model_config=config, is_encoder=True)),
        )


class AliceAIT5MoEDecoderLayer(AliceAIT5DecoderLayer):
    def __init__(self, config: AliceAIT5MoEConfig, layer_idx: int):
        super().__init__(
            config.decoder,
            layer_idx=layer_idx,
            mlp=dMoE(**build_moe_arguments(model_config=config, is_encoder=False)),
        )


class AliceAIT5MoEEncoder(AliceAIT5Encoder):
    _no_split_modules = [AliceAIT5MoEEncoderLayer.__name__]

    def __init__(self, config: AliceAIT5MoEConfig):
        self._moe_config = config
        super().__init__(config.encoder)

    def _build_layers(self, config: AliceAIT5ModuleConfig):
        moe_config = self._moe_config
        self.layers = torch.nn.ModuleList(
            [AliceAIT5MoEEncoderLayer(moe_config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )


class AliceAIT5MoEDecoder(AliceAIT5Decoder):
    _no_split_modules = [AliceAIT5MoEDecoderLayer.__name__]

    def __init__(self, config: AliceAIT5MoEConfig):
        self._moe_config = config
        super().__init__(config.decoder)

    def _build_layers(self, config):
        moe_config = self._moe_config
        self.layers = torch.nn.ModuleList(
            [AliceAIT5MoEDecoderLayer(moe_config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )


class AliceAIT5MoEModel(AliceAIT5Model):
    config_class = AliceAIT5MoEConfig
    _no_split_modules = [AliceAIT5MoEEncoderLayer.__name__, AliceAIT5MoEDecoderLayer.__name__]

    def __init__(self, config: AliceAIT5MoEConfig):
        AliceAIT5PreTrainedModel.__init__(self, config)

        if not config.is_encoder_decoder:
            raise ValueError("AliceAIT5MoEModel only supports encoder-decoder modeling.")

        self.encoder = AliceAIT5MoEEncoder(config)
        self.decoder = AliceAIT5MoEDecoder(config)
        self.post_init()


class AliceAIT5MoEEncoderModel(AliceAIT5EncoderModel):
    config_class = AliceAIT5MoEConfig
    _auto_class = "AutoModel"

    def __init__(self, config: AliceAIT5MoEConfig):
        super().__init__(config)
        self.config.auto_map = {
            "AutoConfig": "configuration_aliceai_t5_moe.AliceAIT5MoEConfig",
            "AutoModel": "modeling_aliceai_t5_moe.AliceAIT5MoEEncoderModel",
        }

    def _build_encoder(self, config):
        return AliceAIT5MoEEncoder(config)


class AliceAIT5MoEForConditionalGeneration(AliceAIT5ForConditionalGeneration):
    config_class = AliceAIT5MoEConfig
    _no_split_modules = [AliceAIT5MoEEncoderLayer.__name__, AliceAIT5MoEDecoderLayer.__name__]

    def _build_model(self, config):
        return AliceAIT5MoEModel(config)


__all__ = [
    "AliceAIT5MoEConfig",
    "AliceAIT5MoEEncoderLayer",
    "AliceAIT5MoEDecoderLayer",
    "AliceAIT5MoEEncoder",
    "AliceAIT5MoEDecoder",
    "AliceAIT5MoEModel",
    "AliceAIT5MoEEncoderModel",
    "AliceAIT5MoEForConditionalGeneration",
]