from typing import List

from mindspeed.deprecate import AutoExecuteFunction
from mindspeed.features_manager import (
    DisableGlooGroupFeature,
    FusedEmaAdamwFeature,
    FusedMoEPermuteFeature,
    FusedSoftmaxFeature,
    GroupedMatmulFeature,
    MC2Feature,
    MoEAlltoAllOverLapFeature,
    MoEAllGatherOverLapFeature,
    MoEFwdBwdOverlapFeature,
    MoEGmmFeature,
    MoEZeroMemoryFeature,
    OptimizeSendRecvCommFeature,
    SwapOptimizerFeature,
    ReuseFP32Param,
    RiPipeSchedulesAdvanceFeature,
    RiPipeSchedulesBubbleFeature,
    UnalignedLinearFeature,
    UnalignedPipelineFeature,
    VirtualOptimizerFeature,
    HcclBufferAdaptiveFeature,
    HcclBufferSetFeature,
    HcclOpModeSetFeature,
    RecomputeNormFeature,
    RecomputeActivationFeature,
    NPUDeterministicFeature,
    NPUDataDumpFeature,
    EnableRecomputeLayersPerPPRank,
    RecomputeMethodFeature,
    SmartSwapFeature,
    SwapAttentionFeature,
    ContextParallelKvCacheFeature,
    TorchFullyShardedDataParallelFeature,
    ProfilerDefaultFeature,
    OptimizeP2PCommFeature,
    FusionAttentionV2Feature,
    MoEAlltoAllMC2Feature,
)
from mindspeed.features_manager.feature import MindSpeedFeature
from mindspeed.features_manager.features_manager import MindSpeedFeaturesManager

from mindspeed_llm.features_manager.fusions.swiglu_limit_feature import SwigluLimitFeature
from mindspeed_llm.features_manager.low_precision.low_precision_optimizer_feature import LowPrecisionOptimizerFeature
from mindspeed_llm.features_manager.affinity.affinity import AffinityFeature
from mindspeed_llm.features_manager.context_parallel.context_parallel_feature import ContextParallelFeature
from mindspeed_llm.features_manager.context_parallel.ulysses_context_parallel import UlyssesContextParallelFeature
from mindspeed_llm.features_manager.context_parallel.mamba_context_parallel import MambaContextParallelFeature
from mindspeed_llm.features_manager.common.data import DataFeature
from mindspeed_llm.features_manager.models.module import ModuleFeature
from mindspeed_llm.features_manager.common.embedding import LanguageModelEmbeddingFeature
from mindspeed_llm.features_manager.common.rotary import RotaryPositionEmbeddingFeature
from mindspeed_llm.features_manager.common.training import TrainingDefaultFeature
from mindspeed_llm.features_manager.tensor_parallel.coc import CoCFeature
from mindspeed_llm.features_manager.dataset.dataset import DatasetFeature
from mindspeed_llm.features_manager.finetune.finetune import FinetuneFeature
from mindspeed_llm.features_manager.dataset.data_preprocess import DatasetPreprocessFeature
from mindspeed_llm.features_manager.finetune.lora import LoraFeature
from mindspeed_llm.features_manager.finetune.lu_lora import LuLoraFeature
from mindspeed_llm.features_manager.high_availability.high_availability import HighAvailabilityFeature
from mindspeed_llm.features_manager.inference.inference import InferenceFeature
from mindspeed_llm.features_manager.evaluation.evaluation import EvaluationFeature
from mindspeed_llm.features_manager.megatron_basic.megatron_basic import MegatronBasicFeature
from mindspeed_llm.features_manager.megatron_basic.model_basic import ModelBasicFeature
from mindspeed_llm.features_manager.megatron_basic.requirements_basic import RequirementsBasicFeature
from mindspeed_llm.features_manager.megatron_basic.training_basic import TrainingBasicFeature
from mindspeed_llm.features_manager.megatron_basic.transformer_engine_basic import TransformerEngineBasicFeature
from mindspeed_llm.features_manager.models.mamba import MambaModel
from mindspeed_llm.features_manager.transformer.multi_latent_attention.dsa_indexer_feature import DSAIndexerFeature
from mindspeed_llm.features_manager.moe.moe_router import MoERouter
from mindspeed_llm.features_manager.moe.shared_expert import MoESharedExpertsFeature
from mindspeed_llm.features_manager.moe.tp_extend_ep import MoETpExtendEpFeature
from mindspeed_llm.features_manager.dpo.dpo import DPOFeature
from mindspeed_llm.features_manager.pipeline_parallel.dualpipev_feature import DualpipeVFeature
from mindspeed_llm.features_manager.pipeline_parallel.noop_layers import NoopLayersFeature
from mindspeed_llm.features_manager.functional.profiling import ProfilingFeature
from mindspeed_llm.features_manager.tokenizer.build_tokenizer import BuildTokenizerFeature
from mindspeed_llm.features_manager.transformer.flash_attention.fusion_attention_feature import FusionAttentionFeature
from mindspeed_llm.features_manager.transformer.flash_attention.alibi_feature import AlibiFeature
from mindspeed_llm.features_manager.transformer.flash_attention.reset_attention_mask_feature import (
    ResetAttentionMaskFeature,
)
from mindspeed_llm.features_manager.transformer.mtp import MultiTokenPredictionFeature
from mindspeed_llm.features_manager.transformer.multi_latent_attention.mla_feature import MLAFeature
from mindspeed_llm.features_manager.transformer.qwen3_next_attention.qwen3_next_feature import Qwen3NextFeature
from mindspeed_llm.features_manager.transformer.transformer_block import TransformerBlockFeature
from mindspeed_llm.features_manager.pipeline_parallel.num_layer_list import NumLayerListFeature
from mindspeed_llm.features_manager.ai_framework.ms_patch_feature import MindSporePatchFeature
from mindspeed_llm.features_manager.tensor_parallel.tp_2d import TP2dFeature
from mindspeed_llm.features_manager.arguments.deprecated_args import DeprecatedArgsFeature
from mindspeed_llm.features_manager.convert_checkpoint.convert_checkpoint import CheckpointFeature
from mindspeed_llm.features_manager.memory.chunk_loss import ChunkLossFeature
from mindspeed_llm.features_manager.layerwise_disaggregated_training.u_shaped_split_feature import UShapedSplitFeature
from mindspeed_llm.features_manager.layerwise_disaggregated_training.vtp_feature import VTPFeature
from mindspeed_llm.features_manager.qat.qat_quant_engine import QATQuantEngineFeature
from mindspeed_llm.features_manager.layerwise_disaggregated_training.vdp_feature import VirtualDPFeature
from mindspeed_llm.features_manager.optimizer.muon_optimizer_feature import MuonOptimizerFeature
from mindspeed_llm.features_manager.transformer.mhc_feature import MHCFeature
from mindspeed_llm.features_manager.transformer.multi_latent_attention.g2_feature import G2Feature

FEATURES_LIST = [
    # MindSpeed Legacy Features
    # MindSpeed Mcore Features
    UnalignedLinearFeature(),
    # MindSpeed-LLM Mcore Features
    TrainingDefaultFeature(),
    DataFeature(),
    LoraFeature(),
    DisableGlooGroupFeature(),
    RotaryPositionEmbeddingFeature(),
    LanguageModelEmbeddingFeature(),
    MambaModel(),
    MoERouter(),
    CoCFeature(),
    HighAvailabilityFeature(),
    MultiTokenPredictionFeature(),
    # MindSpeed-LLM Legacy Features
]


def add_megatron_basic_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            RequirementsBasicFeature(),
            MegatronBasicFeature(),
            TransformerEngineBasicFeature(),
            Qwen3NextFeature(),
            ChunkLossFeature(),
            MuonOptimizerFeature(),
        ]
    )


def add_llm_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            ModelBasicFeature(),
            TrainingBasicFeature(),
            DatasetFeature(),
            DatasetPreprocessFeature(),
            ModuleFeature(),
            NumLayerListFeature(),
            DPOFeature(),
            InferenceFeature(),
            EvaluationFeature(),
            DeprecatedArgsFeature(),
            DSAIndexerFeature(),
            MambaModel(),
            LanguageModelEmbeddingFeature(),
            CheckpointFeature(),
            MHCFeature(),
        ]
    )


def add_affinity_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            AffinityFeature(),
        ]
    )


def add_context_parallel_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            ContextParallelFeature(),
            UlyssesContextParallelFeature(),
            ContextParallelKvCacheFeature(),
            MambaContextParallelFeature(),
        ]
    )


def add_fusions_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            SwigluLimitFeature(),
            FusedSoftmaxFeature(),
            RotaryPositionEmbeddingFeature(),
            GroupedMatmulFeature(),
            FusedMoEPermuteFeature(),
        ]
    )


def add_tensor_parallel_features(features_list: List[MindSpeedFeature]):
    features_list.extend([CoCFeature(), MC2Feature(), TP2dFeature()])


def add_pipeline_parallel_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            RiPipeSchedulesBubbleFeature(),
            RiPipeSchedulesAdvanceFeature(),
            NoopLayersFeature(),
            OptimizeP2PCommFeature(),
            OptimizeSendRecvCommFeature(),
            UnalignedPipelineFeature(),
            DualpipeVFeature(),
        ]
    )


def add_transformer_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            FusionAttentionFeature(),
            # LLM feature
            MLAFeature(),
            # LLM feature
            MultiTokenPredictionFeature(),
            # LLM feature
            TransformerBlockFeature(),
            # LLM feature
            AlibiFeature(),
            # LLM feature
            ResetAttentionMaskFeature(),
            FusionAttentionV2Feature(),
            G2Feature(),
        ]
    )


def add_tokenizer_features(features_list: List[MindSpeedFeature]):
    features_list.extend([BuildTokenizerFeature()])


def add_distributed_features(features_list: List[MindSpeedFeature]):
    features_list.extend([TorchFullyShardedDataParallelFeature()])


def add_reuse_param_features(features_list: List[MindSpeedFeature]):
    features_list.extend([ReuseFP32Param()])


def add_swap_manage_features(features_list: List[MindSpeedFeature]):
    features_list.extend([SmartSwapFeature(), SwapAttentionFeature()])


def add_moe_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            MoEGmmFeature(),
            MoEAlltoAllMC2Feature(),
            # LLM feature
            MoERouter(),
            MoETpExtendEpFeature(),
            # LLM feature
            MoESharedExpertsFeature(),
            MoEAllGatherOverLapFeature(),
            MoEFwdBwdOverlapFeature(),
            MoEAlltoAllOverLapFeature(),
            MoEZeroMemoryFeature(),
        ]
    )


def add_hccl_buffer_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            HcclBufferSetFeature(),
            HcclBufferAdaptiveFeature(),
            HcclOpModeSetFeature(),
        ]
    )


def add_optimizer_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            FusedEmaAdamwFeature(),
            VirtualOptimizerFeature(),
            LowPrecisionOptimizerFeature(),
        ]
    )


def add_functional_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [ProfilingFeature(), NPUDeterministicFeature(), NPUDataDumpFeature(), ProfilerDefaultFeature()]
    )


def add_recompute_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            RecomputeActivationFeature(),
            RecomputeNormFeature(),
            EnableRecomputeLayersPerPPRank(),
            RecomputeMethodFeature(),
        ]
    )


def add_swap_optimizer_feature(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            SwapOptimizerFeature(),
        ]
    )


def add_disable_gloo_group_feature(features_list: List[MindSpeedFeature]):
    features_list.extend([DisableGlooGroupFeature()])


def add_high_availability_feature(features_list: List[MindSpeedFeature]):
    features_list.extend([HighAvailabilityFeature()])


def add_finetune_feature(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            FinetuneFeature(),
            LoraFeature(),
            LuLoraFeature(),
        ]
    )


def add_ai_framework_feature(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            MindSporePatchFeature(),
        ]
    )


def add_layerwise_diaggregated_training_feature(feature_list: List[MindSpeedFeature]):
    feature_list.extend(
        [
            UShapedSplitFeature(),
            VTPFeature(),
            VirtualDPFeature(),
        ]
    )


def add_qat_features(features_list: List[MindSpeedFeature]):
    features_list.extend(
        [
            QATQuantEngineFeature(),
        ]
    )


def create_features_list():
    features_list = []
    add_megatron_basic_features(features_list)
    add_context_parallel_features(features_list)
    add_llm_features(features_list)
    add_affinity_features(features_list)
    add_fusions_features(features_list)
    add_recompute_features(features_list)
    add_functional_features(features_list)
    add_tensor_parallel_features(features_list)
    add_pipeline_parallel_features(features_list)
    add_transformer_features(features_list)
    add_tokenizer_features(features_list)
    add_distributed_features(features_list)
    add_reuse_param_features(features_list)
    add_swap_manage_features(features_list)
    add_moe_features(features_list)
    add_hccl_buffer_features(features_list)
    add_optimizer_features(features_list)
    add_swap_optimizer_feature(features_list)
    add_disable_gloo_group_feature(features_list)
    add_high_availability_feature(features_list)
    add_finetune_feature(features_list)
    add_ai_framework_feature(features_list)
    add_layerwise_diaggregated_training_feature(features_list)
    add_qat_features(features_list)

    return features_list


@AutoExecuteFunction
def set_default_features_list():
    MindSpeedFeaturesManager.set_features_list(create_features_list())


set_default_features_list()