"""Constants for Flower Language Model (LLM) configurations."""

from enum import Enum
from types import MappingProxyType


class Tokenizers(Enum):
    """Tokenizers templates."""

    SMOLLM_SHARED = MappingProxyType({
        "name": "HuggingFaceTB/SmolLM-1.7B",
        "kwargs": {
            "max_seq_len": 2048,
        },
    })

    COSMO = MappingProxyType({
        "name": "/nfs-share/anon/projects/mt_dao/trained_tokenizers/tokenizer_cosmo_",
        "kwargs": {
            "max_seq_len": 2048,
        },
    })

    INFINIMATH = MappingProxyType({
        "name": "/nfs-share/anon/projects/mt_dao/"
        "trained_tokenizers/tokenizer_infiwebmath-4plus_",
        "kwargs": {
            "max_seq_len": 2048,
        },
    })


class ModelConfig(Enum):
    """Model configurations templates."""

    MPT_CAUSAL_LM_125M = MappingProxyType(
        {
            "name": "mpt_causal_lm",
            "init_device": "cpu",
            "d_model": 768,
            "n_heads": 12,
            "n_layers": 12,
            "expansion_ratio": 4,
            "max_seq_len": 2048,
            "vocab_size": 50368,
            "attn_config": {
                "attn_impl": "torch",
            },
            "output_hidden_states": True,
        },
    )

    MPT_CAUSAL_LM_3B = MappingProxyType(
        {
            "name": "mpt_causal_lm",
            "init_device": "cpu",
            "d_model": 2560,
            "n_heads": 20,
            "n_layers": 32,
            "expansion_ratio": 4,
            "max_seq_len": 2048,
            "vocab_size": 50368,
            "attn_config": {
                "attn_impl": "torch",
            },
            "output_hidden_states": True,
        },
    )

    MPT_CAUSAL_LM_350M = MappingProxyType(
        {
            "name": "mpt_causal_lm",
            "init_device": "cpu",
            "d_model": 8192,
            "n_heads": 16,
            "n_layers": 24,
            "expansion_ratio": 4,
            "max_seq_len": 2048,
            "vocab_size": 50368,
            "attn_config": {
                "attn_impl": "torch",
            },
            "output_hidden_states": True,
        },
    )

    SMOLLM_135M = MappingProxyType({
        "name": "mpt_causal_lm",
        "d_model": 576,
        "n_heads": 9,
        "n_layers": 30,
        "ffn_hidden_size": 1536,
        "max_seq_len": 2048,
        "vocab_size": 50368,
        "attn_config": {
            "attn_impl": "flash",
            "rope": True,
            "alibi": False,
            "rope_theta": 10000,
        },
        "ffn_config": {
            "ffn_act_fn": {"name": "silu"},
        },
        "init_config": {
            "init_std": 0.041666666666666664,
        },
    })

    SMOLLM_1B = MappingProxyType({
        "name": "mpt_causal_lm",
        "d_model": 2048,
        "n_heads": 32,
        "n_layers": 24,
        "ffn_hidden_size": 8192,
        "max_seq_len": 2048,
        "vocab_size": 50368,
        "attn_config": {
            "attn_impl": "flash",
            "rope": True,
            "alibi": False,
            "rope_theta": 10000,
        },
        "ffn_config": {
            "ffn_act_fn": {"name": "silu"},
        },
        "init_config": {
            "init_std": 0.02,
        },
    })

    SMOLLM_13B = MappingProxyType({
        "name": "mpt_causal_lm",
        "d_model": 5120,
        "n_heads": 40,
        "n_layers": 40,
        "ffn_hidden_size": 20480,
        "max_seq_len": 2048,
        "vocab_size": 50368,
        "attn_config": {
            "attn_impl": "flash",
            "rope": True,
            "alibi": False,
            "rope_theta": 10000,
        },
        "ffn_config": {
            "ffn_act_fn": {"name": "silu"},
        },
        "init_config": {
            "init_std": 0.014,
        },
    })

    DEPTOE_4E_MAX = MappingProxyType({
        "name": "mpt_mup_causal_lm",
        "init_device": "meta",
        "no_bias": True,
        "use_embedding_norm": True,
        "use_peri_norm": True,
        "d_model": 2048,
        "n_heads": 32,
        "n_layers": 24,
        "ffn_hidden_size": 8192,
        "max_seq_len": 2048,  # Example value, adjust as needed
        "attn_config": {
            "attn_impl": "flash",
            "rope": True,
            "alibi": False,
            "rope_theta": 10000,
        },
        "ffn_config": {
            "ffn_type": "sigma_moe",
            "ff_n_experts": 4,
            "ff_expert_size": 8192,
            "ff_n_active_experts": 4,
            "custom_kernel": False,
            "n_repeats": 1,
            "expert_dropout": 0.0,
            "ffn_act_fn": {"name": "silu"},
        },
        "init_config": {
            "name": "baseline_",
            "init_std": 0.02,
            "emb_init_std": 0.02,
        },
        "mup_config": {
            "mup_enabled": True,
            "mup_disable_attention_scaling": True,
            "mup_disable_hidden_lr_scaling": False,
            "mup_width_multiplier": 1.0,
            "mup_input_alpha": 1.0,
            "mup_output_alpha": 1.0,
            "completep_depth_alpha_enabled": True,
            "completep_depth_multiplier": 1.0,
            "completep_depth_alpha_exp": 1.0,
            "completep_eps_scaling_enabled": True,
            "mup_scale_expert_sel": True,
        },
    })

    DEPTOE_1B_32E_MAX = MappingProxyType({
        "name": "mpt_mup_causal_lm",
        "init_device": "meta",
        "no_bias": True,
        "use_embedding_norm": True,
        "use_peri_norm": True,
        "d_model": 2048,
        "n_heads": 32,
        "n_layers": 24,
        "ffn_hidden_size": 256,
        "max_seq_len": 2048,  # Example value, adjust as needed
        "attn_config": {
            "attn_impl": "flash",
            "rope": True,
            "alibi": False,
            "rope_theta": 10000,
        },
        "ffn_config": {
            "ffn_type": "sigma_moe",
            "ff_n_experts": 32,
            "ff_expert_size": 256,
            "ff_n_active_experts": 4,
            "custom_kernel": False,
            "n_repeats": 1,
            "expert_dropout": 0.0,
            "ffn_act_fn": {"name": "silu"},
        },
        "init_config": {
            "name": "baseline_",
            "init_std": 0.02,
            "emb_init_std": 0.02,
        },
        "mup_config": {
            "mup_enabled": True,
            "mup_disable_attention_scaling": True,
            "mup_disable_hidden_lr_scaling": False,
            "mup_width_multiplier": 8.0,
            "mup_input_alpha": 1.0,
            "mup_output_alpha": 1.0,
            "completep_depth_alpha_enabled": True,
            "completep_depth_multiplier": 6.0,
            "completep_depth_alpha_exp": 1.0,
            "completep_eps_scaling_enabled": True,
            "mup_scale_expert_sel": True,
        },
    })


class FakeLLMConfigs:
    """Fake LLM configurations for testing and examples."""

    SMOLLM_1B_CONFIG = MappingProxyType({
        "name": "smollm-1b",
        "data_local": "./my-copy-c4",
        "data_remote": None,
        "max_seq_len": 2048,
        "global_seed": 17,
        "autoresume": False,
        "run_name": "",
        "eval_gauntlet": None,
        "icl_tasks_config": None,
        "icl_tasks": None,
        "model": {
            "name": "mpt_mup_causal_lm",
            "init_device": "meta",
            "d_model": 2048,
            "n_heads": 32,
            "n_layers": 24,
            "ffn_hidden_size": 8192,
            "max_seq_len": 2048,
            "attn_config": {
                "attn_impl": "flash",
                "rope": True,
                "alibi": False,
                "rope_theta": 10000,
            },
            "ffn_config": {
                "ffn_act_fn": {
                    "name": "silu",
                },
            },
            "init_config": {
                "name": "baseline_",
                "init_std": 0.02,
                "emb_init_std": 0.02,
            },
            "mup_config": {
                "mup_enabled": True,
                "mup_disable_attention_scaling": True,
                "mup_disable_hidden_lr_scaling": False,
                "mup_width_multiplier": 1.0,
                "mup_input_alpha": 1.0,
                "mup_output_alpha": 1.0,
                "completep_depth_alpha_enabled": True,
                "completep_depth_multiplier": 1.0,
                "completep_depth_alpha_exp": 1.0,
                "completep_eps_scaling_enabled": True,
                "mup_scale_expert_sel": False,
            },
        },
        "tokenizer": {
            "name": "HuggingFaceTB/SmolLM-1.7B",
            "kwargs": {
                "model_max_length": 2048,
            },
        },
        "train_loader": {
            "name": "text",
            "dataset": "",
            "drop_last": True,
            "num_workers": 8,
        },
        "eval_loader": {
            "name": "text",
            "dataset": "",
            "drop_last": False,
            "num_workers": 8,
        },
        "scheduler": {
            "schedulers": {
                "lr": {
                    "name": "constant_with_sqrt_cooldown_with_warmup",
                    "t_warmup": "2048ba",
                    "t_max": "20480ba",
                    "t_cooldown": "2048ba",
                },
            },
        },
        "optimizer": {
            "name": "adopt",
            "lr": 0.0005,
            "betas": [0.9, 0.9999],
            "eps": 1.0e-06,
            "weight_decay": 0.0,
        },
        "algorithms": {
            "gradient_clipping": {
                "clipping_type": "norm",
                "clipping_threshold": 1.0,
            },
        },
        "max_duration": "20480ba",
        "eval_interval": "512ba",
        "eval_first": False,
        "eval_subset_num_batches": -1,
        "global_train_batch_size": 1024,
        "seed": 17,
        "device_eval_batch_size": 16,
        "device_train_microbatch_size": 16,
        "precision": "amp_bf16",
        "fsdp_config": {
            "sharding_strategy": "FULL_SHARD",
            "mixed_precision": "PURE",
            "activation_checkpointing": False,
            "activation_checkpointing_reentrant": False,
            "activation_cpu_offload": False,
            "limit_all_gathers": True,
        },
        "progress_bar": False,
        "log_to_console": True,
        "console_log_interval": "100ba",
        "callbacks": {
            "speed_monitor": {
                "window_size": 20,
                "gpu_flops_available": None,
            },
            "lr_monitor": {},
            "memory_monitor": {},
            "runtime_estimator": {},
            "activation_monitor_full_model": {
                "interval": "10ba",
            },
            "optimizer_monitor": {
                "interval": "10ba",
                "only_global": True,
            },
        },
        "loggers": {
            "tensorboard": {"flush_interval": 10},
        },
        "save_interval": "512ba",
        "save_num_checkpoints_to_keep": 1,
        "save_folder": "",
        "save_overwrite": False,
        "load_path": None,
    })
