# 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.
from functools import partial
from typing import Any, Callable, Dict, List, Set, Union

import torch
from lightning.fabric.utilities.imports import (
    _TORCH_GREATER_EQUAL_2_0,
    _TORCH_GREATER_EQUAL_2_1,
)


if torch.distributed.is_available():
    from torch.distributed.fsdp.wrap import _or_policy, lambda_auto_wrap_policy

    if _TORCH_GREATER_EQUAL_2_1:
        # N.B. while imports using the public API are preferred, avoiding use of `_Policy` here would mean excluding
        # support for user policies that directly subclass `_Policy` (rather than `CustomPolicy` or `ModuleWrapPolicy`)
        from torch.distributed.fsdp.wrap import _Policy, CustomPolicy
    elif _TORCH_GREATER_EQUAL_2_0:
        from torch.distributed.fsdp.wrap import _FSDPPolicy


if _TORCH_GREATER_EQUAL_2_1:

    class NameDrivenCustomPolicy(CustomPolicy):
        """An auto-wrapping policy extension that applies module name-based override directives on top of a given
        base ``auto_wrap_policy``.

        The composition of module name-based wrapping directives with a given ``auto_wrap_policy`` is
        achieved here by:
            1. Subclassing ``CustomPolicy``, initializing it with an object id-based module name mapping lambda to
                create a module name-driven policy with a base policy handle.
            2. Generating a base wrapping configuration dictionary by running the user's provided ``auto_wrap_policy``
            3. Updating the base configuration dictionary with any wrapping configuration generated by the module
                name-driven policy and returning the composed configuration.
        """

        def __init__(self, auto_wrap_policy_handle: _Policy, override_ids: List):
            """Compose the provided ``auto_wrap_policy`` with any provided override directives.

            Args:
                auto_wrap_policy_handle (Union[Callable, _Policy]): The user's base ``auto_wrap_policy``.
                override_ids (List): Object ids of the desired modules to wrap even if the provided ``auto_wrap_policy``
                    otherwise would not dictate so.
            """
            super().__init__(lambda_fn=lambda m: id(m) in override_ids)
            self._base_awp = auto_wrap_policy_handle

        def _run_policy(
            self,
            root_module: torch.nn.Module,
            ignored_modules: Set[torch.nn.Module],
            root_kwargs: Dict[str, Any],
        ) -> Dict[torch.nn.Module, Dict[str, Any]]:
            target_module_to_kwargs = self._base_awp._run_policy(root_module, ignored_modules, root_kwargs)
            nb_policy_target_module_to_kwargs = super()._run_policy(root_module, ignored_modules, root_kwargs)
            target_module_to_kwargs.update(nb_policy_target_module_to_kwargs)
            return target_module_to_kwargs

    NameDrivenPolicy = NameDrivenCustomPolicy

elif _TORCH_GREATER_EQUAL_2_0:

    class NameDrivenFSDPPolicy(_FSDPPolicy):
        """An auto-wrapping policy extension that applies module name-based override directives on top of a given
        base ``auto_wrap_policy``. This implementation uses the deprecated `_FSDPPolicy` API and is only kept for
        backwards compatibility with PyTorch 2.0. For PyTorch >= 2.1, ``NameDrivenCustomPolicy`` should be used.

        The composition of module name-based wrapping directives with a given ``auto_wrap_policy`` is
        achieved here by:
            1. Generating an object id-based module name mapping lambda and passing it to the standard
                ``lambda_auto_wrap_policy``.
            2. Composing the user's provided ``auto_wrap_policy`` with the above name-based policy using the standard
                ``_or_policy``.
        """

        def __init__(self, auto_wrap_policy_handle: Union[Callable, _FSDPPolicy], override_ids: List):
            """Compose the provided ``auto_wrap_policy`` with any provided override directives.

            Args:
                auto_wrap_policy_handle (Union[Callable, _FSDPPolicy]): The user's base ``auto_wrap_policy``.
                override_ids (List): Object ids of the desired modules to wrap even if the provided ``auto_wrap_policy``
                    otherwise would not dictate so.
            """
            if isinstance(auto_wrap_policy_handle, _FSDPPolicy):
                auto_wrap_policy_handle = auto_wrap_policy_handle.policy
            name_driven_policy = partial(lambda_auto_wrap_policy, lambda_fn=lambda m: id(m) in override_ids)
            self._policy: Callable = partial(_or_policy, policies=[auto_wrap_policy_handle, name_driven_policy])

        @property
        def policy(self) -> Callable:
            return self._policy

    NameDrivenPolicy = NameDrivenFSDPPolicy  # type: ignore[assignment,misc]
