# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# 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 atexit
import logging as _logging
import os
import re
import shutil
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Dict, Literal, Optional, Union

import lightning.pytorch as pl
import torch
import torch.distributed as dist
from lightning.fabric.plugins import CheckpointIO
from lightning.fabric.utilities.rank_zero import rank_zero_info
from lightning.fabric.utilities.seed import reset_seed
from lightning.pytorch.strategies.model_parallel import ModelParallelStrategy as PLModelParallelStrategy
from lightning.pytorch.trainer.states import TrainerFn
from lightning.pytorch.utilities.types import STEP_OUTPUT
from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel, SequenceParallel
from typing_extensions import override

from nemo.lightning import io
from nemo.lightning.pytorch.strategies.utils import (
    _destroy_dist_connection,
    ckpt_to_dir,
    create_checkpoint_io,
    fsdp2_strategy_parallelize,
)
from nemo.utils import logging
from nemo.utils.import_utils import safe_import_from

try:
    from torch.distributed.tensor._api import distribute_tensor
    from torch.distributed.tensor.placement_types import Replicate, Shard
except ImportError:
    from torch.distributed._tensor.api import distribute_tensor
    from torch.distributed._tensor.placement_types import Replicate, Shard

MixedPrecisionPolicy, HAS_MIXED_PRECISION_POLICY = safe_import_from(
    "torch.distributed.fsdp", "MixedPrecisionPolicy", fallback_module="torch.distributed._composable.fsdp"
)
CPUOffloadPolicy, HAS_CPU_OFFLOAD_POLICY = safe_import_from(
    "torch.distributed.fsdp", "CPUOffloadPolicy", fallback_module="torch.distributed._composable.fsdp"
)

_logger = _logging.getLogger(__name__)


class FSDP2Strategy(PLModelParallelStrategy, io.IOMixin):
    """FSDP2Strategy implementing FSDP via FSDP 2."""

    def __init__(
        self,
        data_parallel_size: Union[Literal["auto"], int] = "auto",
        tensor_parallel_size: Union[Literal["auto"], int] = "auto",
        context_parallel_size: Optional[int] = 1,
        sequence_parallel: bool = False,
        offload_policy: 'CPUOffloadPolicy' = None,
        data_sampler=None,
        checkpoint_io=None,
        mp_policy=None,
        parallelize_fn=fsdp2_strategy_parallelize,
        use_hf_tp_plan: bool = True,
        custom_tp_plan: Optional[Dict[str, Union[RowwiseParallel, ColwiseParallel, SequenceParallel]]] = None,
        **kwargs,
    ):
        """Initializes the FSDP2Strategy with specified parallelization settings.

        Args:
            data_parallel_size (Union[Literal["auto"], int]): Size of data parallel. Defaults to "auto".
            tensor_parallel_size (Union[Literal["auto"], int]): Size of tensor parallel. Defaults to "auto".
            context_parallel_size (optional): Number of context-parallel groups. Defaults to 1.
            sequence_parallel (bool): Whether to enable sequence parallelism when use_hf_tp_plan is False and
                custom_tp_plan is not provided. Defaults to False. Only effective when tensor_parallel_size > 1.
            data_sampler (optional): Custom data sampler to process dataloaders.
            mp_policy (optional): Mixed precision policy for parameter and operation casting.
                Defaults to:
                ```python
                MixedPrecisionPolicy(
                    param_dtype=torch.bfloat16,
                    reduce_dtype=torch.float32,
                    output_dtype=None,
                    cast_forward_inputs=True,
                )
                ```
            parallelize_fn (callable, optional): Function for parallelizing the model. Defaults to None.
            use_hf_tp_plan (bool, optional): Whether to use the huggingface TP plan. This will be used if
                custom_tp_plan is not provided. Also, sequence_parallel option will be ignored if use_hf_tp_plan
                is set to True. Defaults to True.
            custom_tp_plan (Optional[Dict[str, Any]], optional): Custom tensor parallel plan for the model.
                tensor_parallel_size need to be > 1 to use this option. If provided, it overrides the
                default tensor parallel plan. sequence_parallel option will be ignored if custom_tp_plan
                is provided.
            **kwargs: Additional arguments for base class initialization.
        """
        super().__init__(data_parallel_size=data_parallel_size, tensor_parallel_size=tensor_parallel_size, **kwargs)
        self._checkpoint_io = checkpoint_io
        self.context_parallel_size = context_parallel_size
        self.data_sampler = data_sampler
        self.checkpoint = None
        self.mp_policy = mp_policy
        if self.mp_policy is None:
            assert HAS_MIXED_PRECISION_POLICY is not None, "Expected to have MixedPrecisionPolicy"
            self.mp_policy = MixedPrecisionPolicy(
                param_dtype=torch.bfloat16,
                reduce_dtype=torch.bfloat16,
                output_dtype=torch.bfloat16,
                cast_forward_inputs=True,
            )
        self.store: Optional[torch.distributed.Store] = None
        self.parallelize_fn = parallelize_fn
        self.offload_policy = offload_policy
        self.sequence_parallel = sequence_parallel
        self.use_hf_tp_plan = use_hf_tp_plan

        self.tp_shard_plan = None
        if custom_tp_plan is not None:
            self.tp_shard_plan = custom_tp_plan
            logging.info(
                "You are using a custom TP plan. Make sure it is compatible with the model. Parallelization would ",
                "not raise errors if the custom TP plan is not compatible. SP option will also be ignored.",
            )
        elif self.use_hf_tp_plan:
            logging.info(
                "You are using a huggingface TP plan. Make sure your model is a huggingface model. Certain ",
                "parallelizations might not be supported. SP option will also be ignored.",
            )
        else:
            # Parallelize the first embedding and the last linear out projection
            base_model_tp_plan = {
                "model.embed_tokens": RowwiseParallel(input_layouts=Replicate()),
                "model.layers.*.self_attn.q_proj": ColwiseParallel(),
                "model.layers.*.self_attn.k_proj": ColwiseParallel(),
                "model.layers.*.self_attn.v_proj": ColwiseParallel(),
                "model.layers.*.self_attn.o_proj": RowwiseParallel(),
                "model.layers.*.mlp.up_proj": ColwiseParallel(),
                "model.layers.*.mlp.gate_proj": ColwiseParallel(),
                "model.layers.*.mlp.down_proj": RowwiseParallel(),
                "lm_head": ColwiseParallel(output_layouts=Replicate()),
            }

            base_model_sp_plan = {
                "model.embed_tokens": RowwiseParallel(input_layouts=Replicate(), output_layouts=Shard(1)),
                "model.norm": SequenceParallel(),
                "model.layers.*.input_layernorm": SequenceParallel(),
                "model.layers.*.self_attn.o_proj": RowwiseParallel(output_layouts=Shard(1)),
                "model.layers.*.post_attention_layernorm": SequenceParallel(),
                "model.layers.*.mlp.down_proj": RowwiseParallel(output_layouts=Shard(1)),
                "lm_head": ColwiseParallel(input_layouts=Shard(1), output_layouts=Replicate()),
            }

            if self.sequence_parallel:
                # Enable sequence parallelism only if TP size > 1
                base_model_tp_plan.update(base_model_sp_plan)

            self.tp_shard_plan = base_model_tp_plan
            logging.info(
                "Using default TP plan for parallelization. It is compatible with huggingface llama3-style models."
            )

    @property
    @override
    def lightning_restore_optimizer(self) -> bool:
        """Optim state restoration is enabled"""
        return True

    def load_optimizer_state_dict(self, checkpoint) -> None:
        """Stores a reference to the optimizer state-dict for later restoration.

        Instead of immediately restoring the optimizer's state, this method saves the checkpoint
        reference and defers the restoration until the first training step. This is necessary
        because, in NeMo 2.0, PeFT adapters are added dynamically just before the first training
        step. Attempting to restore the optimizer state-dict before the adapters are initialized
        would result in an error.

        Args:
            checkpoint (dict): A dictionary containing the trainer's checkpoint,
                            including the optimizer state-dict.
        """
        # TODO(@akoumparouli): refactor.
        self.checkpoint = checkpoint

    def _load_optimizer_state_dict(self) -> None:
        """Restores the optimizer state-dict from the stored checkpoint.

        This method applies the optimizer state stored in the checkpoint to the corresponding
        optimizers. It ensures that the optimizer states are correctly restored after the
        PeFT adapters have been added in the first training step.

        If no checkpoint is stored, the method exits without performing any restoration.

        Note: This operation runs only once, as the checkpoint reference is cleared after execution.
        """
        from torch.distributed.checkpoint.state_dict import set_optimizer_state_dict

        if self.checkpoint is None:
            for optimizer, opt_state in zip(self.optimizers, self.checkpoint["optimizer_states"]):
                set_optimizer_state_dict(
                    self.lightning_module,
                    optimizer,
                    optim_state_dict=opt_state,
                    options={},
                )
            # run this only once
            self.checkpoint = None

    @override
    def setup_environment(self) -> None:
        """setup distributed environment and device mesh"""
        from torch.distributed.device_mesh import init_device_mesh

        self.accelerator.setup_device(self.root_device)

        self._setup_distributed()
        if self._data_parallel_size == "auto":
            self._data_parallel_size = self.num_nodes

        if self._tensor_parallel_size == "auto":
            self._tensor_parallel_size = self.num_processes

        mesh_shape = []
        mesh_dim_names = []
        # TP needs to be the last dimension as innermost dimension, DP-CP-TP
        for dim, name in zip(
            [self._data_parallel_size, self.context_parallel_size, self._tensor_parallel_size],
            ["data_parallel", "context_parallel", "tensor_parallel"],
        ):
            mesh_shape.append(int(dim))
            mesh_dim_names.append(name)

        self._device_mesh = init_device_mesh(
            device_type=self.root_device.type,
            mesh_shape=tuple(mesh_shape),
            mesh_dim_names=mesh_dim_names,
        )

        # Construct sharding and reduction meshes for specific configurations.
        # Replace existing mesh strategies if a custom mesh design is provided.
        # WARNING: ADDING MESHES INCREASES MEMORY USAGE. Use device meshes for
        # multiple dimensions of parallelism if possible.
        if self._device_mesh["context_parallel"].size() > 1:
            # Context parallelism loss reduction mesh. Remember to divide out CP size in tokens per second throughput.
            self._device_mesh[("data_parallel", "context_parallel")]._flatten(mesh_dim_name="dp_cp")

        self.lightning_module._device_mesh = self._device_mesh

    @override
    def setup(self, trainer: pl.Trainer) -> None:
        """Configures the strategy within the PyTorch Lightning trainer.

        Args:
            trainer (pl.Trainer): The PyTorch Lightning trainer instance.
        """
        self.trainer = trainer
        # connect trainer to accelerator.
        self.accelerator.setup(trainer)
        # Parallelize model
        if getattr(self, '_init_model_parallel', True):
            self.parallelize()
        # Corner case, as FSDP2 expected to be used multi-device.
        if self._data_parallel_size == 1:
            self._lightning_module = self._lightning_module.to(self.root_device)
        # setup optim
        if getattr(self, '_setup_optimizers', True) and trainer.state.fn == TrainerFn.FITTING:
            super().setup_optimizers(trainer)

    def parallelize(self):
        """Applies fully_shard on model"""
        if self.parallelize_fn is not None:
            # TODO(@akoumparouli): self.lightning_module is an nn.Module child, use it directly?
            # Apply FSDP2 and TP to the model
            self.parallelize_fn(
                self.lightning_module.model,
                device_mesh=self._device_mesh,
                mp_policy=self.mp_policy,
                use_hf_tp_plan=self.use_hf_tp_plan,
                tp_shard_plan=self.tp_shard_plan,
                offload_policy=self.offload_policy,
            )
            # Apply this only once
            self.parallelize_fn = None
        else:
            logging.warning("Called parallelize more than once.")

    @override
    def _setup_distributed(self) -> None:
        """Initializes process group for communications."""

        # Implementation from superclass copied below in order to pass the store to the process group init
        reset_seed()
        self.set_world_ranks()
        self._process_group_backend = self._get_process_group_backend()
        # See https://github.com/pytorch/pytorch/issues/148532 for details.
        if self.offload_policy is not None:
            self._process_group_backend = "cuda:nccl,cpu:gloo"

        assert self.cluster_environment is not None
        if not torch.distributed.is_available():
            raise RuntimeError("torch.distributed is not available. Cannot initialize distributed process group")
        if torch.distributed.is_initialized():
            _logger.debug("torch.distributed is already initialized. Exiting early")
            return

        global_rank = self.cluster_environment.global_rank()
        world_size = self.cluster_environment.world_size()
        os.environ["MASTER_ADDR"] = self.cluster_environment.main_address
        os.environ["MASTER_PORT"] = str(self.cluster_environment.main_port)
        _logger.info(f"Initializing distributed: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")
        torch.distributed.init_process_group(
            self._process_group_backend, rank=global_rank, world_size=world_size, store=self.store
        )

        if self._process_group_backend == "nccl":
            atexit.register(_destroy_dist_connection)

        # On rank=0 let everyone know training is starting
        rank_zero_info(
            f"{'-' * 100}\n"
            f"distributed_backend={self._process_group_backend}\n"
            f"All distributed processes registered. Starting with {world_size} processes\n"
            f"{'-' * 100}\n"
        )

    def _get_loss_reduction(self, step_type: str):
        """Retrieves the loss reduction method for a given step type.

        Args:
            step_type (str): The type of step (e.g., "training", "validation").

        Returns:
            Callable: The loss reduction function, if defined; otherwise, None.
        """
        for fn_name in [f"{step_type}_loss_reduction", "loss_reduction"]:
            if hasattr(self.lightning_module, fn_name):
                return getattr(self.lightning_module, fn_name)
        return None

    def _step_proxy(self, step_type, batch, batch_idx=None):
        """Executes a training, validation, or test step and applies loss reduction if available.

        Args:
            step_type (str): The step type ("training", "validation", "test", "predict").
            batch: The input batch.
            batch_idx (optional): Index of the batch.

        Returns:
            Tuple: The computed loss and a dictionary with reduced loss metrics.
        """
        method_name = f"{step_type}_step"
        if self.model != self.lightning_module:
            loss = self._forward_redirection(self.model, self.lightning_module, method_name, batch, batch_idx)
        else:
            loss = getattr(self.lightning_module, method_name)(batch, batch_idx)

        _loss_reduction = self._get_loss_reduction(step_type)
        if _loss_reduction:
            return _loss_reduction.forward(batch, loss)
        return loss, {'avg': loss}

    @override
    def optimizer_state(self, optimizer):
        """returns the sharded optim state"""
        return optimizer.state_dict()

    @override
    def training_step(self, batch, batch_idx=None) -> STEP_OUTPUT:
        """Defines the training step, logging relevant metrics.

        Args:
            batch: The input batch.
            batch_idx (optional): The index of the batch.

        Returns:
            STEP_OUTPUT: The loss for backpropagation.
        """

        # See load_optimizer_state_dict to understand why we call this here.
        if self.checkpoint is not None:
            self._load_optimizer_state_dict()

        assert self.lightning_module is not None
        assert self.model is not None

        if self.context_parallel_size > 1:
            # Only pass context_parallel=True if AutoModel supports and has non-trivial CP.
            loss = self.lightning_module.training_step(batch, batch_idx, context_parallel=True)
        else:
            loss = self.lightning_module.training_step(batch, batch_idx)

        self.lightning_module.log(
            'global_step',
            self.trainer.global_step,
            prog_bar=True,
            rank_zero_only=True,
            batch_size=1,
        )

        self.lightning_module.log(
            'step',
            self.trainer.global_step,
        )

        return loss

    @override
    def validation_step(self, batch, batch_idx=None) -> Any:
        """Defines the validation step, logging validation loss.

        Args:
            batch: The input batch.
            batch_idx (optional): The index of the batch.

        Returns:
            Any: The validation loss.
        """
        assert self.lightning_module is not None
        assert self.model is not None
        with self.precision_plugin.val_step_context():
            loss, reduced = self._step_proxy("validation", batch, batch_idx)
            if reduced["avg"]:
                self.lightning_module.log('val_loss', reduced['avg'], rank_zero_only=True, batch_size=1)
            return loss

    @override
    def test_step(self, batch, batch_idx=None) -> STEP_OUTPUT:
        """Defines the test step, logging test loss.

        Args:
            batch: The input batch.
            batch_idx (optional): The index of the batch.

        Returns:
            Any: The test loss.
        """
        assert self.lightning_module is not None
        assert self.model is not None
        with self.precision_plugin.test_step_context():
            loss, reduced = self._step_proxy("test", batch, batch_idx)
            self.lightning_module.log('test_loss', reduced['avg'], rank_zero_only=True, batch_size=1)

            return loss

    @override
    def predict_step(self, batch, batch_idx=None) -> STEP_OUTPUT:
        """Runs one predict step.

        Args:
            batch (dict): the batch to use for pred.
            batch_idx (int, optional): the batch index. Defaults to None.

        Returns:
            STEP_OUTPUT: the reduced loss.
        """
        assert self.lightning_module is not None
        assert self.model is not None
        with self.precision_plugin.predict_step_context():
            loss, reduced = self._step_proxy("predict", batch, batch_idx)
            return reduced

    @contextmanager
    @override
    def tensor_init_context(self, empty_init: Optional[bool] = None):
        """Context manager used for initialization"""
        # Materializaton happens in `setup()`
        # @akoumparouli: using Parent's tensor_init_context causes mcore
        # parameters to be initialized on GPU instead of (assumed) CPU.
        yield

    @property
    @override
    def checkpoint_io(self) -> CheckpointIO:
        """CheckpointIO getter

        Returns:
            CheckpointIO: _description_
        """
        if self._checkpoint_io is None:
            self._checkpoint_io = create_checkpoint_io()

        return self._checkpoint_io

    @checkpoint_io.setter
    def checkpoint_io(self, io: CheckpointIO) -> None:
        """CheckpointIO setter

        Args:
            io (CheckpointIO): the checkpointio to use.
        """
        self._checkpoint_io = io

    @property
    def current_epoch_step(self) -> int:
        """Gets the current step within an epoch.

        Returns:
            int: The step index within the epoch.
        """
        return max(
            self.trainer.fit_loop.epoch_loop.automatic_optimization.optim_progress.optimizer.step.current.completed,
            self.trainer.fit_loop.epoch_loop.manual_optimization.optim_step_progress.current.completed,
        )

    @override
    def remove_checkpoint(self, filepath: Union[str, Path]) -> None:
        """Removes a checkpoint from the filesystem.

        Args:
            filepath (Union[str, Path]): Path to the checkpoint to be removed.
        """
        ckpt = ckpt_to_dir(filepath)
        if self.is_global_zero:
            if os.path.islink(ckpt):
                os.unlink(ckpt)
            elif os.path.exists(ckpt):
                shutil.rmtree(ckpt)

    @override
    def lightning_module_state_dict(self) -> Dict[str, Any]:
        """Collects the state dict of the model."""
        from nemo.lightning.pytorch.strategies.utils import to_cpu

        assert self.lightning_module is not None
        tmp_sd = self.lightning_module.state_dict()
        # Get adapter_only value from checkpoint io. There's two cases:
        # - nemo.lightning.pytorch.callbacks.peft.WrappedAdapterIO
        #   In this case, the self._checkpoint_io object is a wrapper and holds a `checkpoint_io`
        #   attribute which we query for the `adapter_only` attribute
        # - otherwise, it's the base case which has the adapter_only attribute directly accesible.
        is_adapter_only = getattr(
            self._checkpoint_io,
            'adapter_only',
            getattr(getattr(self._checkpoint_io, 'checkpoint_io', {}), 'adapter_only', False),
        )

        if is_adapter_only:
            # if any key has "lora" in FQN, then it will only move lora keys to cpu, since only
            # the adapter weights are saved.
            name_has_lora = lambda x: 'lora' in x.lower()
            module_names = list(filter(name_has_lora, tmp_sd.keys()))
        else:
            module_names = list(tmp_sd.keys())

        state_dict = {}
        for name in module_names:
            param = tmp_sd.pop(name)
            state_dict[name] = to_cpu(param)

        dist.barrier()
        return state_dict

    @override
    def save_checkpoint(
        self, checkpoint: Dict[str, Any], filepath: Union[str, Path], storage_options: Optional[Any] = None
    ) -> None:
        """
        Unshards FSDP2 checkpoint and passes it to checkpoint_io for saving to a file.
        """

        if self.is_global_zero:
            self.checkpoint_io.save_checkpoint(checkpoint, filepath, storage_options=storage_options)

    @override
    def load_checkpoint(self, checkpoint_path: str | Path) -> Dict[str, Any]:
        """Loads checkpoint with checkpoint_io"""
        return self.checkpoint_io.load_checkpoint(checkpoint_path)

    @override
    @torch.no_grad
    def load_model_state_dict(
        self,
        ckpt,
        strict=False,
    ):
        """Shards a full state dict"""
        if self._tensor_parallel_size > 1 and self._device_mesh["context_parallel"].size() == 1:
            # Gather TP/SP strategy keys
            colwise_keys = [k for k in self.tp_shard_plan if isinstance(self.tp_shard_plan[k], ColwiseParallel)]
            rowwise_keys = [k for k in self.tp_shard_plan if isinstance(self.tp_shard_plan[k], RowwiseParallel)]
            seq_parallel_keys = [k for k in self.tp_shard_plan if isinstance(self.tp_shard_plan[k], SequenceParallel)]

            sharded_state = {k: v for k, v in ckpt['state_dict'].items()}

            # placement is (dp, tp)
            for k, v in sharded_state.items():
                if any(re.match(x, k) for x in seq_parallel_keys):
                    sharded_state[k] = distribute_tensor(v, self.device_mesh, placements=(Shard(dim=0), Replicate()))
                elif any(re.match(x, k) for x in colwise_keys):
                    sharded_state[k] = distribute_tensor(v, self.device_mesh, placements=(Shard(dim=0), Shard(dim=0)))
                elif any(re.match(x, k) for x in rowwise_keys):
                    sharded_state[k] = distribute_tensor(v, self.device_mesh, placements=(Shard(dim=0), Shard(dim=1)))
                else:
                    # This is for layers not sharded by TP/SP
                    sharded_state[k] = distribute_tensor(
                        v, self.device_mesh["data_parallel"], placements=(Shard(dim=0),)
                    )
        # TODO(@akoumparouli): update `placements` value once TP is enabled.
        elif self._tensor_parallel_size == 1 and self._device_mesh["context_parallel"].size() > 1:
            # Shard across the CP device mesh, associated with the fully_shard() call
            # in utils.fsdp2_strategy_parallelize().
            sharded_state = {
                k: distribute_tensor(
                    v, self._device_mesh[("data_parallel", "context_parallel")], placements=(Replicate(), Shard(dim=0))
                )
                for k, v in ckpt['state_dict'].items()
            }
        else:
            # Default shard across DP for FSDP2.
            sharded_state = {
                k: distribute_tensor(v, self._device_mesh["data_parallel"], placements=(Shard(dim=0),))
                for k, v in ckpt['state_dict'].items()
            }

        self.lightning_module.load_state_dict(sharded_state, strict=strict)
