# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import pytest
import deepspeed.comm as dist
import torch

from unit.common import DistributedTest, preferred_dtype
from unit.simple_model import random_dataloader

import deepspeed
from deepspeed.utils import set_z3_leaf_modules, unset_z3_leaf_modules, get_z3_leaf_modules, z3_leaf_module, \
    set_z3_leaf_modules_by_name, set_z3_leaf_modules_by_suffix
from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
from deepspeed.runtime.zero.leaf_module_config import (DEFAULT_LEAF_MODULE_CLASSES, DEFAULT_LEAF_MODULE_NAMES,
                                                       DEFAULT_LEAF_MODULE_NAME_SUFFIXES)
from deepspeed.accelerator import get_accelerator
from torch import nn
import time


class ChooseModuleByCounter(torch.nn.Module):

    def __init__(self, hidden_dim):
        super(ChooseModuleByCounter, self).__init__()
        self.linears = torch.nn.ModuleList(
            [torch.nn.Linear(hidden_dim, hidden_dim, bias=False),
             torch.nn.Linear(hidden_dim, hidden_dim, bias=False)])
        self.act = torch.nn.ReLU()
        self.cel = torch.nn.CrossEntropyLoss()
        self.counter = 0

    def forward(self, x, y):
        # This fails without setting this module as a leaf module.
        # See the comment in `set_z3_leaf_modules()`.
        x = self.linears[self.counter % len(self.linears)](x)
        x = self.act(x)
        loss = self.cel(x, y)
        self.counter += 1
        return x, loss


class ChooseModuleByRankModel(torch.nn.Module):

    def __init__(self, hidden_dim):
        super(ChooseModuleByRankModel, self).__init__()
        self.linears = torch.nn.ModuleList(
            [torch.nn.Linear(hidden_dim, hidden_dim, bias=False),
             torch.nn.Linear(hidden_dim, hidden_dim, bias=False)])
        self.act = torch.nn.ReLU()
        self.cel = torch.nn.CrossEntropyLoss()

    def forward(self, x, y):
        # Each rank runs only one of the linear layers
        x = self.linears[dist.get_rank() % len(self.linears)](x)
        x = self.act(x)
        loss = self.cel(x, y)
        return x, loss


class MultiOutputMoEBlock(nn.Module):
    """A simplified MoE block that returns multiple tensors.

    This model mimics Qwen3 MoE which returns (hidden_states, router_logits).
    When used with ZeRO3 leaf modules and autograd multithreading enabled,
    this pattern previously caused race conditions in fetch_sub_module
    because backward hooks could be triggered concurrently from multiple threads.

    See: https://github.com/deepspeedai/DeepSpeed/issues/7824
    """

    def __init__(self, hidden_dim, num_experts=4):
        super(MultiOutputMoEBlock, self).__init__()
        self.num_experts = num_experts
        self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
        self.experts = nn.ModuleList([nn.Linear(hidden_dim, hidden_dim, bias=False) for _ in range(num_experts)])
        self.act = nn.ReLU()
        self.cel = nn.CrossEntropyLoss()

    def forward(self, x, y):
        # Compute router logits - this tensor will have gradients flowing through it
        router_logits = self.gate(x)

        # Process through experts
        for expert in self.experts:
            x = expert(x)
        x = self.act(x)
        loss = self.cel(x, y)

        # Return multiple tensors - this triggers concurrent backward hooks
        # when autograd multithreading is enabled
        return x, loss, router_logits


class MLPBlock(nn.Module):

    def __init__(self, hidden_dim):
        super(MLPBlock, self).__init__()
        self.gate_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.up_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.act_fn = nn.GELU()

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


class FineGrainedBlock(nn.Module):

    def __init__(self, hidden_dim, num_block):
        super(FineGrainedBlock, self).__init__()
        self.num_block = num_block
        self.mlp_layers = torch.nn.ModuleList([MLPBlock(hidden_dim=hidden_dim) for _ in range(self.num_block)])

    def forward(self, x):
        for i in range(self.num_block):
            x = self.mlp_layers[i](x)
        return x


class BaseLeafModule(nn.Module):

    def __init__(self):
        super(BaseLeafModule, self).__init__()


class SubLeafModule(BaseLeafModule):

    def __init__(self, hidden_dim):
        super(SubLeafModule, self).__init__()
        self.proj = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x):
        return self.proj(x)


class WrapperLeafModule(nn.Module):

    def __init__(self, hidden_dim):
        super(WrapperLeafModule, self).__init__()
        self.child = SubLeafModule(hidden_dim)

    def forward(self, x):
        return self.child(x)


def test_set_leaf_modules_with_fully_qualified_name():
    hidden_dim = 16
    model = WrapperLeafModule(hidden_dim)
    fq_name = f"{SubLeafModule.__module__}.{SubLeafModule.__qualname__}"

    matched = set_z3_leaf_modules(model, [fq_name])

    assert len(matched) == 1
    assert matched[0] is model.child
    assert z3_leaf_module(model.child)
    assert not z3_leaf_module(model)


def test_set_leaf_modules_no_raise_when_missing():
    hidden_dim = 16
    model = WrapperLeafModule(hidden_dim)

    matched = set_z3_leaf_modules(model, ["NonExistentClass"], raise_if_not_found=False)

    assert matched == []
    assert not z3_leaf_module(model.child)


def test_set_leaf_modules_by_name():
    hidden_dim = 16
    model = WrapperLeafModule(hidden_dim)

    matched, missing = set_z3_leaf_modules_by_name(model, ["child"])

    assert matched == [model.child]
    assert missing == []
    assert z3_leaf_module(model.child)


def test_set_leaf_modules_by_name_missing():
    hidden_dim = 16
    model = WrapperLeafModule(hidden_dim)

    matched, missing = set_z3_leaf_modules_by_name(model, ["missing"], raise_if_not_found=False)

    assert matched == []
    assert missing == ["missing"]


def test_set_leaf_modules_by_suffix():
    hidden_dim = 16
    model = WrapperLeafModule(hidden_dim)

    matched, missing = set_z3_leaf_modules_by_suffix(model, ["child"])

    assert missing == []
    assert matched == [model.child]
    assert z3_leaf_module(model.child)


def test_set_leaf_modules_by_suffix_missing():
    hidden_dim = 16
    model = WrapperLeafModule(hidden_dim)

    matched, missing = set_z3_leaf_modules_by_suffix(model, ["missing"], raise_if_not_found=False)

    assert matched == []
    assert missing == ["missing"]


def test_zero_leaf_module_default_config():
    config = DeepSpeedZeroConfig()
    assert config.leaf_module.classes == DEFAULT_LEAF_MODULE_CLASSES
    assert config.leaf_module.names == DEFAULT_LEAF_MODULE_NAMES
    assert config.leaf_module.name_suffixes == DEFAULT_LEAF_MODULE_NAME_SUFFIXES


def test_zero_leaf_module_custom_config():
    payload = {
        "leaf_module": {
            "classes": ["custom.module.CustomClass"],
            "names": ["transformer.layer"],
            "name_suffixes": ["experts"]
        }
    }

    config = DeepSpeedZeroConfig(**payload)

    assert config.leaf_module.classes == ["custom.module.CustomClass"]
    assert config.leaf_module.names == ["transformer.layer"]
    assert config.leaf_module.name_suffixes == ["experts"]


def test_zero_leaf_module_string_coercion():
    payload = {"leaf_module": {"classes": "my.Class", "names": "submodule", "name_suffixes": "tail"}}

    config = DeepSpeedZeroConfig(**payload)

    assert config.leaf_module.classes == ["my.Class"]
    assert config.leaf_module.names == ["submodule"]
    assert config.leaf_module.name_suffixes == ["tail"]


@pytest.mark.skip(reason="Requires Hugging Face transformers; run manually when validating defaults.")
def test_default_leaf_module_classes_exist():
    import importlib

    from deepspeed.runtime.zero.leaf_module_config import DEFAULT_LEAF_MODULE_CLASSES

    for cls_path in DEFAULT_LEAF_MODULE_CLASSES:
        module_name, _, class_name = cls_path.rpartition('.')
        module = importlib.import_module(module_name)
        assert hasattr(module, class_name), f"Expected {class_name} in {module_name}"


class modelWithFineGrainedBlock(nn.Module):

    def __init__(self, hidden_dim, num_block):
        super(modelWithFineGrainedBlock, self).__init__()
        self.coarse_grained_layer1 = nn.Linear(hidden_dim, 8 * hidden_dim)
        self.coarse_grained_layer2 = nn.Linear(8 * hidden_dim, hidden_dim)
        self.fine_grained_layer = FineGrainedBlock(hidden_dim, num_block)
        self.cel = torch.nn.CrossEntropyLoss()

    def forward(self, x, y):
        x = self.coarse_grained_layer1(x)
        x = self.coarse_grained_layer2(x)
        x = self.fine_grained_layer(x)
        loss = self.cel(x, y)
        return x, loss


def run_model(model, config_dict, hidden_dim, dtype, requires_grad):
    model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
    data_loader = random_dataloader(model=model,
                                    total_samples=10,
                                    hidden_dim=hidden_dim,
                                    device=model.device,
                                    dtype=dtype)
    dist.barrier()
    for batch in data_loader:
        batch[0].requires_grad = requires_grad
        loss = model(batch[0], batch[1])
        loss = loss[1]
        model.backward(loss)
        model.step()

    # Needed in ZeRO 3. Not doing so can give memory leak
    model.destroy()


class TestSetZ3LeafModule(DistributedTest):
    # Need multiple gpus to test possible hanging
    world_size = 2
    reuse_dist_env = True

    def _create_zero_config(self, hidden_dim, leaf_module=None):
        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 1e-6
                }
            },
            "zero_optimization": {
                "stage": 3,
                "stage3_prefetch_bucket_size": hidden_dim**2,
                "stage3_param_persistence_threshold": 0,
                "stage3_max_reuse_distance": 0,
            }
        }
        if leaf_module is not None:
            config_dict["zero_optimization"]["leaf_module"] = leaf_module

        if preferred_dtype() is torch.float16:
            config_dict["fp16"] = {"enabled": True}
        elif preferred_dtype() is torch.bfloat16:
            config_dict["bf16"] = {"enabled": True}

        return config_dict

    def _test_set_z3_leaf_modules(self, cls, requires_grad):
        hidden_dim = 128
        config_dict = self._create_zero_config(hidden_dim)

        model = cls(hidden_dim)

        assert not z3_leaf_module(model)
        set_z3_leaf_modules(model, [cls])
        assert z3_leaf_module(model)

        run_model(model, config_dict, hidden_dim, preferred_dtype(), requires_grad)

    def test_choose_module_by_counter(self):
        self._test_set_z3_leaf_modules(ChooseModuleByCounter, True)

    def test_choose_module_by_rank(self):
        self._test_set_z3_leaf_modules(ChooseModuleByRankModel, True)

    def test_multi_output_leaf_module_thread_safety(self):
        """Test that leaf modules returning multiple tensors work correctly with autograd multithreading.

        This tests the fix for https://github.com/deepspeedai/DeepSpeed/issues/7824
        where MoE models (like Qwen3) returning multiple tensors caused race conditions
        in fetch_sub_module when autograd executed backward hooks from multiple threads.
        """
        # Ensure autograd multithreading is enabled (this is the default, but be explicit)
        torch.autograd.set_multithreading_enabled(True)

        hidden_dim = 128
        config_dict = self._create_zero_config(hidden_dim)

        model = MultiOutputMoEBlock(hidden_dim, num_experts=4)

        assert not z3_leaf_module(model)
        set_z3_leaf_modules(model, [MultiOutputMoEBlock])
        assert z3_leaf_module(model)

        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)

        data_loader = random_dataloader(model=model,
                                        total_samples=10,
                                        hidden_dim=hidden_dim,
                                        device=model.device,
                                        dtype=preferred_dtype())
        dist.barrier()

        # Run multiple iterations to increase chance of hitting race conditions
        for batch in data_loader:
            batch[0].requires_grad = True
            # Model returns (output, loss, router_logits)
            output, loss, router_logits = model(batch[0], batch[1])
            # Include router_logits in the loss to ensure multiple backward paths
            total_loss = loss + 0.01 * router_logits.mean()
            model.backward(total_loss)
            model.step()

        model.destroy()

    def test_multi_output_non_leaf_module_thread_safety(self):
        """Ensure non-leaf modules returning multiple tensors remain thread-safe.

        This covers the multi-output autograd multithreading case without marking the
        module as a ZeRO leaf module.
        """
        torch.autograd.set_multithreading_enabled(True)

        hidden_dim = 128
        config_dict = self._create_zero_config(hidden_dim)

        model = MultiOutputMoEBlock(hidden_dim, num_experts=4)
        assert not z3_leaf_module(model)

        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)

        data_loader = random_dataloader(model=model,
                                        total_samples=10,
                                        hidden_dim=hidden_dim,
                                        device=model.device,
                                        dtype=preferred_dtype())
        dist.barrier()

        for batch in data_loader:
            batch[0].requires_grad = True
            output, loss, router_logits = model(batch[0], batch[1])
            total_loss = loss + 0.01 * router_logits.mean()
            model.backward(total_loss)
            model.step()

        model.destroy()

    def test_no_grad_input_error(self):
        try:
            self._test_set_z3_leaf_modules(ChooseModuleByCounter, False)
            raise AssertionError(
                "Expected RuntimeError: inputs with requires_grad=False is not supported for a leaf module")
        except RuntimeError as e:
            pass

    def test_set_unset_leaf_modules(self):
        hidden_dim = 128
        model = ChooseModuleByCounter(hidden_dim)
        assert len(set_z3_leaf_modules(model, [torch.nn.ModuleList])) == 1, \
            "Expected only one module to be set as a leaf module"
        assert len(get_z3_leaf_modules(model)) == 1, "Expected there is only one leaf module"

        assert len(unset_z3_leaf_modules(model, [torch.nn.ModuleList])) == 1, \
            "Expected only one module to be unset as a leaf module"
        assert len(get_z3_leaf_modules(model)) == 0, "Expected there is no leaf module"

    def test_set_leaf_modules_with_subclass(self):
        hidden_dim = 32
        model = WrapperLeafModule(hidden_dim)

        leaf_modules = set_z3_leaf_modules(model, [BaseLeafModule])

        assert len(leaf_modules) == 1, "Expected the subclass instance to be marked as leaf"
        assert leaf_modules[0] is model.child, "Expected the subclass instance to be returned"
        assert z3_leaf_module(model.child), "Expected subclass instance flagged as leaf"
        assert not z3_leaf_module(model), "Expected wrapper module to remain non-leaf"

    def test_set_no_match_class(self):
        hidden_dim = 128
        model = ChooseModuleByCounter(hidden_dim)
        try:
            set_z3_leaf_modules(model, [torch.nn.Conv2d])
            raise AssertionError("Expected error that no module is set as a leaf module")
        except ValueError as e:
            pass

    def test_leaf_module_enabled_via_config(self):
        hidden_dim = 128
        leaf_class_fqn = f"{ChooseModuleByCounter.__module__}.{ChooseModuleByCounter.__qualname__}"
        config_dict = self._create_zero_config(hidden_dim,
                                               leaf_module={
                                                   "classes": [leaf_class_fqn],
                                                   "name_suffixes": ["linears"]
                                               })

        model = ChooseModuleByCounter(hidden_dim)
        assert not z3_leaf_module(model)

        run_model(model, config_dict, hidden_dim, preferred_dtype(), True)

        assert z3_leaf_module(model)
        modules_by_name = dict(model.named_modules())
        assert "linears" in modules_by_name
        assert z3_leaf_module(modules_by_name["linears"])


@pytest.mark.parametrize("module_granularity_threshold", [0, 100, 12100, 10000000])
class TestZ3LeafOptimization(DistributedTest):
    world_size = 2
    reuse_dist_env = True

    def test_finegrained_optimization(self, module_granularity_threshold: int):
        hidden_dim = 128
        num_block = 16
        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 1e-6
                }
            },
            "zero_optimization": {
                "stage": 3,
                "stage3_prefetch_bucket_size": hidden_dim**2,
                "stage3_param_persistence_threshold": 0,
                "stage3_max_reuse_distance": 0,
            }
        }
        if preferred_dtype() is torch.float16:
            config_dict["fp16"] = {"enabled": True}
        elif preferred_dtype() is torch.bfloat16:
            config_dict["bf16"] = {"enabled": True}

        def bench_loss_and_time(config):
            warm_up_step = 10
            model = modelWithFineGrainedBlock(hidden_dim, num_block)
            model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)
            data_loader = random_dataloader(model=model,
                                            total_samples=20,
                                            hidden_dim=hidden_dim,
                                            device=model.device,
                                            dtype=preferred_dtype())
            dist.barrier()
            loss_list = []

            for i, batch in enumerate(data_loader):
                if i == warm_up_step:
                    dist.barrier()
                    get_accelerator().synchronize()
                    start_time = time.time()
                batch[0].requires_grad = True
                loss = model(batch[0], batch[1])
                loss = loss[1]
                loss_list.append(loss)
                model.backward(loss)
                model.step()
            get_accelerator().synchronize()
            end_time = time.time()
            duration = end_time - start_time
            model.destroy()
            return loss_list, duration

        baseline_loss_list, baseline_exec_time = bench_loss_and_time(config_dict)

        config_dict["zero_optimization"]["stage3_module_granularity_threshold"] = module_granularity_threshold
        loss, duration = bench_loss_and_time(config_dict)

        if dist.get_rank() == 0:
            print("baseline exec time:", baseline_exec_time)
            print(
                f"finegrained optimziation exec time: {duration},granularity threshold:{module_granularity_threshold} "
            )
            assert baseline_loss_list == loss, f"incorrect loss value with threshold:{module_granularity_threshold}"