import torch
import torch.distributed as dist

import megatron.core.parallel_state as mpu
from megatron.core.distributed import DistributedDataParallelConfig
from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel
from megatron.core.distributed.fsdp.src.megatron_fsdp.megatron_fsdp import MegatronFSDP
from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
from megatron.core.models.gpt.gpt_model import GPTModel
from megatron.core.transformer import TransformerConfig
from megatron.core.transformer.transformer_layer import TransformerLayer
from megatron.core.utils import unwrap_model


def init_distributed():
    """Initialize torch.distributed and Megatron parallel state."""
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())
    mpu.initialize_model_parallel(
        tensor_model_parallel_size=1,
        pipeline_model_parallel_size=1,
    )


def build_fsdp_model():
    """Build a GPTModel wrapped with FullyShardedDataParallel, just like training does."""
    transformer_config = TransformerConfig(
        num_layers=2,
        hidden_size=64,
        num_attention_heads=4,
        use_cpu_initialization=True,
    )

    transformer_layer_spec = get_gpt_layer_local_spec()

    gpt_model = GPTModel(
        config=transformer_config,
        transformer_layer_spec=transformer_layer_spec,
        vocab_size=256,
        max_sequence_length=128,
        pre_process=True,
        post_process=True,
    ).cuda()

    ddp_config = DistributedDataParallelConfig(
        data_parallel_sharding_strategy="optim_grads_params",
        overlap_grad_reduce=True,
        overlap_param_gather=True,
        bucket_size=10000,
        use_megatron_fsdp=True,
    )

    fsdp_model = FullyShardedDataParallel(
        config=transformer_config,
        ddp_config=ddp_config,
        module=gpt_model,
        fsdp_unit_modules=[TransformerLayer],
    )

    return gpt_model, fsdp_model


def main():
    init_distributed()

    gpt_model, fsdp_model = build_fsdp_model()

    # Print the wrapping hierarchy
    print("=" * 60)
    print("Model wrapping hierarchy:")
    print(f"  fsdp_model              : {type(fsdp_model).__name__}")
    print(f"  fsdp_model.module       : {type(fsdp_model.module).__name__}")
    print(f"  fsdp_model.module.module: {type(fsdp_model.module.module).__name__}")
    print()

    # Verify the hierarchy:
    #   FullyShardedDataParallel -> MegatronFSDP -> GPTModel
    assert isinstance(fsdp_model, FullyShardedDataParallel)
    assert isinstance(fsdp_model.module, MegatronFSDP)
    assert isinstance(fsdp_model.module.module, GPTModel)

    # Now test unwrap_model
    unwrapped = unwrap_model(fsdp_model)

    print("unwrap_model result:")
    print(f"  type  : {type(unwrapped).__name__}")
    print(f"  is original GPTModel: {unwrapped is gpt_model}")

    # Cleanup
    mpu.destroy_model_parallel()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()