#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
from pathlib import Path

import torch
import torch.distributed as dist


def force_megatron_cpu_mode() -> None:
    # Megatron-Core allocates several modules with device=torch.cuda.current_device().
    # Returning "cpu" lets those constructors allocate CPU tensors for this probe.
    torch.cuda.current_device = lambda: "cpu"  # type: ignore[assignment]
    torch.cuda.set_device = lambda *_args, **_kwargs: None  # type: ignore[assignment]


def init_single_process() -> None:
    force_megatron_cpu_mode()
    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
    os.environ.setdefault("MASTER_PORT", "29621")
    os.environ.setdefault("RANK", "0")
    os.environ.setdefault("WORLD_SIZE", "1")
    if not dist.is_initialized():
        dist.init_process_group("gloo", rank=0, world_size=1)

    from megatron.core import parallel_state

    if not parallel_state.is_initialized():
        parallel_state.initialize_model_parallel(
            tensor_model_parallel_size=1,
            pipeline_model_parallel_size=1,
            virtual_pipeline_model_parallel_size=None,
            context_parallel_size=1,
            expert_model_parallel_size=1,
        )


def build_model(seq_length: int):
    import torch.nn.functional as F
    from megatron.bridge.models.gpt_provider import GPTModelProvider
    from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
    from megatron.core.transformer.torch_norm import WrappedTorchNorm
    from megatron.core.utils import init_method_normal, scaled_init_method_normal

    def cpu_rmsnorm_block_spec(config, vp_stage=None):
        spec = get_gpt_decoder_block_spec(
            config,
            use_transformer_engine=False,
            normalization="RMSNorm",
            vp_stage=vp_stage,
        )
        spec.layer_norm = WrappedTorchNorm
        return spec

    from megatron.core.process_groups_config import ProcessGroupCollection

    provider = GPTModelProvider(
        num_layers=12,
        hidden_size=1536,
        num_attention_heads=24,
        num_query_groups=4,
        kv_channels=64,
        ffn_hidden_size=4608,
        seq_length=seq_length,
        vocab_size=154856,
        should_pad_vocab=True,
        share_embeddings_and_output_weights=True,
        position_embedding_type="rope",
        normalization="RMSNorm",
        gated_linear_unit=True,
        activation_func=F.silu,
        add_bias_linear=False,
        num_moe_experts=12,
        moe_layer_freq=[1 if idx in {2, 4, 6, 8, 10} else 0 for idx in range(12)],
        moe_ffn_hidden_size=6144,
        moe_router_topk=4,
        moe_router_load_balancing_type="aux_loss",
        moe_aux_loss_coeff=0.02,
        moe_z_loss_coeff=0.001,
        moe_token_dispatcher_type="allgather",
        moe_expert_capacity_factor=1.25,
        moe_router_enable_expert_bias=False,
        moe_router_bias_update_rate=0.02,
        moe_grouped_gemm=True,
        tensor_model_parallel_size=1,
        pipeline_model_parallel_size=1,
        expert_model_parallel_size=1,
        context_parallel_size=1,
        sequence_parallel=False,
        transformer_impl="local",
        attention_backend="unfused",
        transformer_layer_spec=cpu_rmsnorm_block_spec,
        init_method=init_method_normal(0.02),
        output_layer_init_method=scaled_init_method_normal(0.02, 12),
        init_method_std=0.02,
    )
    provider.perform_initialization = False
    provider._pg_collection = ProcessGroupCollection.use_mpu_process_groups()
    return provider.provide(pre_process=True, post_process=True)


def inspect_model(args: argparse.Namespace) -> None:
    init_single_process()
    model = build_model(args.seq_length)
    print("model_built", type(model))
    print("param_count", sum(p.numel() for p in model.parameters()))
    if args.mode == "sharded-keys":
        state = model.sharded_state_dict(prefix="")
        print("sharded_keys_sample")
    else:
        state = model.state_dict()
        print("state_keys_sample")
    for key in list(state.keys())[:120]:
        print(key)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--seq-length", type=int, default=8192)
    parser.add_argument("--mode", choices=["state-keys", "sharded-keys"], default="state-keys")
    args = parser.parse_args()
    inspect_model(args)


if __name__ == "__main__":
    main()
