# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import torch
import torch.nn as nn
from typing import Callable, Optional
from torch.distributed.device_mesh import DeviceMesh
from torch.distributed._composable.replicate import replicate
from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard, MixedPrecisionPolicy
from torch.distributed.tensor import Shard
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ParallelStyle,
)
from cosmos_rl.utils.parallelism import ParallelDims, pre_parallelize_sanity_check
from cosmos_rl.utils.logging import logger
from cosmos_rl.utils.util import str2torch_dtype
from cosmos_rl.policy.config import Config as CosmosConfig
from cosmos_rl.patch import PipelineStage, Schedule1F1B, ScheduleGPipe
from cosmos_rl.utils.ulysses import (
    ulysses_attn_func,
    swizzle_cp_forward,
    ulysses_attn_func_varlen,
)
from cosmos_rl.policy.kernel.moe.moe import MoE, GroupedExpertsDeepEP
from torch.distributed.tensor import distribute_module, distribute_tensor


@pre_parallelize_sanity_check
def parallelize(
    model: nn.Module,
    parallel_dims: ParallelDims,
    config: CosmosConfig,
    pp_loss_fn: Optional[Callable] = None,
):
    """
    Apply tensor parallelism, activation checkpointing, torch.compile, and data
    parallelism to the model.

    NOTE: The passed-in model preferably should be on meta device. Otherwise,
    the model must fit on GPU or CPU memory.
    """
    world_mesh = parallel_dims.mesh

    pipeline_parallelize(model, parallel_dims, config)

    if parallel_dims.tp_enabled:
        apply_tp_ep(
            model,
            world_mesh["tp"],
            enable_float8_tensorwise_tp=config.train.fp8.enable_fp8
            and config.train.fp8.quant_recipe == "tensorwise",
            enable_async_tp=config.train.async_tp_enabled,
            parallel_dims=parallel_dims,
        )

    if parallel_dims.cp_enabled:
        apply_cp(model, parallel_dims)
        logger.info("Applied Context Parallel to the model")

    # turn on per-TransformerBlock compile after AC wrapping and before FSDP
    if config.train.compile:
        """
        Why we need to apply compile after AC wrapping and before FSDP?
        https://github.com/pytorch/torchtitan/issues/472#issuecomment-2242200809
        """
        apply_compile(model)

    reshard_after_forward_policy = config.train.fsdp_reshard_after_forward
    # For visual model, TP mesh should be merged into DP_Shard
    if model.visual is not None and (
        parallel_dims.tp_enabled
        or parallel_dims.dp_shard_enabled
        or parallel_dims.cp_enabled
    ):
        logger.info(
            f"Applying FSDP(TP-merged: {parallel_dims.tp_enabled}) to the visual model"
        )
        mp_policy = MixedPrecisionPolicy(
            param_dtype=str2torch_dtype(config.train.param_dtype),
            reduce_dtype=str2torch_dtype(config.train.fsdp_reduce_dtype),
            cast_forward_inputs=True,
        )
        fsdp_config = {"mesh": world_mesh["dp_cp_tp"], "mp_policy": mp_policy}
        if config.train.fsdp_offload:
            fsdp_config["offload_policy"] = CPUOffloadPolicy()
        for layer_id, transformer_block in model.visual.blocks.items():
            if reshard_after_forward_policy == "always":
                reshard_after_forward = True
            elif reshard_after_forward_policy == "never":
                reshard_after_forward = False
            elif reshard_after_forward_policy == "default":
                if parallel_dims.pp_enabled:
                    # For PP, do not reshard after forward to avoid per-microbatch
                    # all-gathers, which can be expensive and non-overlapped
                    reshard_after_forward = False
                else:
                    # As an optimization, do not reshard after forward for the last
                    # transformer block since FSDP would prefetch it immediately
                    reshard_after_forward = int(layer_id) < len(model.visual.blocks) - 1
            else:
                raise ValueError(
                    f"Invalid reshard_after_forward_policy: {reshard_after_forward_policy}."
                )
            fully_shard(
                transformer_block,
                **fsdp_config,
                reshard_after_forward=reshard_after_forward,
            )
        fully_shard(
            model.visual,
            **fsdp_config,
            reshard_after_forward=not parallel_dims.pp_enabled,
        )

    if parallel_dims.dp_shard_enabled:
        dp_group_names = ["dp_shard"]
        if parallel_dims.cp_enabled:
            dp_group_names.append("cp")
        if parallel_dims.tp_enabled:
            # _flatten() modifies world_mesh in place, adding fsdp_no_moe/fsdp_moe dims
            world_mesh[tuple(dp_group_names + ["tp"])]._flatten(
                mesh_dim_name="fsdp_no_moe"
            )
            fsdp_mesh_no_moe_name = "fsdp_no_moe"
            world_mesh[tuple(dp_group_names)]._flatten(mesh_dim_name="fsdp_moe")
            fsdp_mesh_moe_name = "fsdp_moe"
        else:
            fsdp_mesh_no_moe_name = fsdp_mesh_moe_name = "dp_shard_cp"

        if parallel_dims.dp_replicate_enabled:
            dp_mesh_dim_names_for_no_moe = ("dp_replicate", fsdp_mesh_no_moe_name)
            dp_mesh_dim_names_for_moe = ("dp_replicate", fsdp_mesh_moe_name)
        else:
            dp_mesh_dim_names_for_no_moe = (fsdp_mesh_no_moe_name,)
            dp_mesh_dim_names_for_moe = (fsdp_mesh_moe_name,)

        # apply FSDP or HSDP, potentially with Context Parallel
        apply_fsdp(
            model,
            world_mesh[tuple(dp_mesh_dim_names_for_no_moe)],
            world_mesh[tuple(dp_mesh_dim_names_for_moe)],
            param_dtype=str2torch_dtype(config.train.param_dtype),
            reduce_dtype=str2torch_dtype(config.train.fsdp_reduce_dtype),
            pp_enabled=parallel_dims.pp_enabled,
            cpu_offload=config.train.fsdp_offload,
            reshard_after_forward_policy=reshard_after_forward_policy,
        )

        if parallel_dims.dp_replicate_enabled:
            logger.info("Applied HSDP to the model")
        else:
            logger.info("Applied FSDP to the model")

        if config.train.fsdp_offload:
            logger.info("Applied CPU Offloading to the model")
    elif parallel_dims.dp_replicate_enabled:
        if world_mesh.ndim > 1:
            raise RuntimeError("DDP has not supported > 1D parallelism")
        apply_ddp(
            model,
            world_mesh,
            enable_compile=config.train.compile,
            enable_compiled_autograd=config.train.compile,
        )

    pp_rank, pp_size = parallel_dims.pp_coord
    if pp_size > 1:
        proc_group = parallel_dims.mesh["pp"].get_group()
        # IMPORTANT:
        # Only CUDA device is acceptable for PP Schedule, `meta` device will cause hanging issue
        device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}")

        stage = PipelineStage(
            model,
            pp_rank,
            pp_size,
            device,
            group=proc_group,
        )

        assert (
            config.train.train_batch_per_replica
            % config.policy.parallelism.pp_micro_batch_size
            == 0
        ), "train_batch must be divisible by pp_micro_batch_size"
        assert (
            config.train.train_batch_per_replica
            // config.policy.parallelism.pp_micro_batch_size
        ) % pp_size == 0, (
            "train_batch / pp_micro_batch_size must be divisible by pp_size"
        )
        assert pp_loss_fn is not None, "pp_loss_fn must be provided"
        n_microbatches = (
            config.train.train_batch_per_replica
            // config.policy.parallelism.pp_micro_batch_size
        )
        if config.validation.enable:
            assert (
                config.validation.batch_size
                % config.policy.parallelism.pp_micro_batch_size
                == 0
            ), "validation_batch must be divisible by pp_micro_batch_size"
            assert (
                config.validation.batch_size
                // config.policy.parallelism.pp_micro_batch_size
            ) % pp_size == 0, (
                "validation_batch / pp_micro_batch_size must be divisible by pp_size"
            )
            n_val_microbatches = (
                config.validation.batch_size
                // config.policy.parallelism.pp_micro_batch_size
            )

        if config.train.train_policy.type == "grpo":
            assert (
                config.train.train_batch_per_replica
                % (
                    config.policy.parallelism.dp_shard_size
                    * config.train.train_policy.mini_batch
                )
                == 0
            ), "train_batch must be divisible by dp_shard_size * mini_batch"
            assert (
                config.train.train_policy.mini_batch
                % config.policy.parallelism.pp_micro_batch_size
                == 0
            ), "mini_batch must be divisible by pp_micro_batch_size"
            n_microbatches = (
                config.train.train_policy.mini_batch
                // config.policy.parallelism.pp_micro_batch_size
            )
        else:
            assert (
                config.train.train_batch_per_replica
                % config.train.train_policy.mini_batch
                == 0
            ), "train_batch must be divisible by mini_batch"
            assert (
                config.train.train_policy.mini_batch
                % config.policy.parallelism.pp_micro_batch_size
                == 0
            ), "mini_batch must be divisible by pp_micro_batch_size"
            n_microbatches = (
                config.train.train_policy.mini_batch
                // config.policy.parallelism.pp_micro_batch_size
            )

        schedule = Schedule1F1B(
            stage=stage,
            n_microbatches=n_microbatches,
            loss_fn=pp_loss_fn,
        )
        if config.validation.enable:
            val_schedule = ScheduleGPipe(
                stage=stage,
                n_microbatches=n_val_microbatches,
            )
            return schedule, val_schedule
        else:
            return schedule, None
    else:
        return None, None


class _ExpertParallel(ParallelStyle):
    """
    ExpertParallel class is used to shard the MoE parameters on the EP mesh.
    Dim `0` of each parameter is sharded since that is the expert dimension.
    """

    def _partition_fn(self, name, module, device_mesh):
        # shard on the expert dimension
        assert device_mesh.ndim == 1

        for name, param in module.named_parameters(recurse=False):
            dist_param = nn.Parameter(distribute_tensor(param, device_mesh, [Shard(0)]))
            module.register_parameter(name, dist_param)

        if isinstance(module, GroupedExpertsDeepEP):
            module.init_token_dispatcher(ep_mesh=device_mesh)

    def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
        return distribute_module(
            module,
            device_mesh,
            self._partition_fn,
        )


def apply_cp(model: nn.Module, parallel_dims: ParallelDims):
    """Apply Context Parallel to the model."""
    # check if cp is compatible with model

    cp_mesh = parallel_dims.mesh["cp"]
    # For language
    for _, transformer_block in model.model.layers.items():
        original_attn_func = transformer_block.self_attn.attn_func
        transformer_block.self_attn.attn_func = ulysses_attn_func(
            original_attn_func, cp_mesh
        )
    for _, transformer_block in model.model.layers.items():
        original_attn_func_varlen = transformer_block.self_attn.attn_func_varlen
        transformer_block.self_attn.attn_func_varlen = ulysses_attn_func_varlen(
            original_attn_func_varlen, cp_mesh
        )

    # For visual model since we merge cp into fsdp, no need handling.
    swizzle_cp_forward(model, parallel_dims)


def apply_tp_ep(
    model: nn.Module,
    tp_ep_mesh: DeviceMesh,
    enable_float8_tensorwise_tp: bool,
    enable_async_tp: bool,
    parallel_dims: ParallelDims,
):
    """Apply tensor parallelism."""
    # - Apply expert parallelism to MLP
    for layer_id, transformer_block in model.model.layers.items():
        if isinstance(transformer_block.mlp, MoE):
            parallelize_module(
                module=transformer_block.mlp.experts,
                device_mesh=tp_ep_mesh,
                parallelize_plan=_ExpertParallel(),
            )

    if enable_async_tp:
        from torch.distributed._symmetric_memory import enable_symm_mem_for_group

        torch._inductor.config._micro_pipeline_tp = True
        enable_symm_mem_for_group(tp_ep_mesh.get_group().group_name)

    logger.info(
        f"Applied {'Float8 tensorwise ' if enable_float8_tensorwise_tp else ''}{'Async ' if enable_async_tp else ''}"
        "Tensor Parallelism to the model"
    )


# for selective op activation checkpointing
_save_list = {
    torch.ops.aten.mm.default,
    torch.ops.aten._scaled_dot_product_efficient_attention.default,
    torch.ops.aten._scaled_dot_product_flash_attention.default,
    # for low precision training, it's useful to always save
    # the result of max, since the absolute maximum is
    # used to compute the scaling factor for quantization.
    torch.ops.aten.max.default,
}


def apply_compile(model: nn.Module, fullgraph: bool = True):
    """
    Apply torch.compile to each TransformerBlock, which makes compilation efficient due to
    repeated structure. Alternatively one can compile the whole model (after applying DP).
    """
    for layer_id, transformer_block in model.model.layers.named_children():
        transformer_block = torch.compile(transformer_block, fullgraph=fullgraph)
        model.model.layers.register_module(layer_id, transformer_block)

    # ``model.visual`` could get deleted by pipeline split
    if model.visual is not None:
        for layer_id, transformer_block in model.visual.blocks.named_children():
            transformer_block = torch.compile(
                transformer_block, fullgraph=fullgraph, dynamic=True
            )
            model.visual.blocks.register_module(layer_id, transformer_block)

    logger.info("Each TransformerBlock compiled with torch.compile")


def apply_fsdp(
    model: nn.Module,
    fsdp_mesh_no_moe: DeviceMesh,
    fsdp_mesh_moe: DeviceMesh,
    param_dtype: torch.dtype,
    reduce_dtype: torch.dtype,
    pp_enabled: bool,
    cpu_offload: bool = False,
    reshard_after_forward_policy: str = "default",
):
    """
    Apply data parallelism (via FSDP2) to the model.

    Args:
        model (nn.Module): The model to apply data parallelism to.
        dp_mesh (DeviceMesh): The device mesh to use for data parallelism.
        param_dtype (torch.dtype): The data type to use for model parameters.
        reduce_dtype (torch.dtype): The data type to use for reduction operations.
        pp_enabled (bool): Whether pipeline parallelism is enabled.
        cpu_offload (bool, optional): Whether to offload model parameters to CPU. Defaults to False.
        reshard_after_forward_policy (str, optional): The policy to use for resharding after forward pass. Defaults to "default".
            Other options: "never", "always".
            - "default" applies default resharding behavior, implementing "smart defaults" for known optimal scenarios.
            - "always" will enable `reshard_after_forward` for all forward passes.
            - "never" will disable `reshard_after_forward` for all forward passes.

    """
    mp_policy = MixedPrecisionPolicy(
        param_dtype=param_dtype, reduce_dtype=reduce_dtype, cast_forward_inputs=True
    )
    fsdp_config_no_moe = {"mesh": fsdp_mesh_no_moe, "mp_policy": mp_policy}
    fsdp_config_moe = {"mesh": fsdp_mesh_moe, "mp_policy": mp_policy}
    if cpu_offload:
        fsdp_config_no_moe["offload_policy"] = CPUOffloadPolicy()
        fsdp_config_moe["offload_policy"] = CPUOffloadPolicy()

    for layer_id, transformer_block in model.model.layers.items():
        if reshard_after_forward_policy == "always":
            reshard_after_forward = True
        elif reshard_after_forward_policy == "never":
            reshard_after_forward = False
        elif reshard_after_forward_policy == "default":
            if pp_enabled:
                # For PP, do not reshard after forward to avoid per-microbatch
                # all-gathers, which can be expensive and non-overlapped
                reshard_after_forward = False
            else:
                # As an optimization, do not reshard after forward for the last
                # transformer block since FSDP would prefetch it immediately
                reshard_after_forward = int(layer_id) < len(model.model.layers) - 1
        else:
            raise ValueError(
                f"Invalid reshard_after_forward_policy: {reshard_after_forward_policy}."
            )

        # Bottom-up: first shard the experts in EP group, then shard the transformer block in big-dp group
        if isinstance(transformer_block.mlp, MoE):
            fully_shard(
                transformer_block.mlp.experts,
                **fsdp_config_moe,
                reshard_after_forward=reshard_after_forward,
            )
        fully_shard(
            transformer_block,
            **fsdp_config_no_moe,
            reshard_after_forward=reshard_after_forward,
        )
    if model.model.embed_tokens is not None:
        fully_shard(
            model.model.embed_tokens, **fsdp_config_no_moe, reshard_after_forward=True
        )
    fully_shard(model, **fsdp_config_no_moe, reshard_after_forward=not pp_enabled)


def apply_ddp(
    model: nn.Module,
    dp_mesh: DeviceMesh,
    enable_compile: bool,
    enable_compiled_autograd: bool,
):
    if enable_compile:
        if enable_compiled_autograd:
            torch._dynamo.config.optimize_ddp = (
                "python_reducer_without_compiled_forward"
            )
        else:
            torch._dynamo.config.optimize_ddp = "ddp_optimizer"

    replicate(model, device_mesh=dp_mesh, bucket_cap_mb=100)

    logger.info("Applied DDP to the model")


def pipeline_parallelize(
    model: nn.Module,
    parallel_dims: ParallelDims,
    config: CosmosConfig,
):
    # TODO(huik): support pipeline parallelism for Qwen3VLMoe
    assert not parallel_dims.pp_enabled, (
        "Pipeline parallelism is not supported for Qwen3VLMoe currently."
    )
    if parallel_dims.pp_enabled:
        pp_rank, pp_size = parallel_dims.pp_coord
        model.apply_pipeline_split(pp_rank, pp_size)
    else:
        logger.info("Pipeline is not enabled, skipping")
