# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
#
# See LICENSE for license information.

import torch
import pytest

torch.manual_seed(1234)
device = torch.device("npu" if torch.npu.is_available() else "cpu")


class _Sequential(torch.nn.Sequential):
    """Sequential model that forwards keyword arguments to modules."""

    def forward(self, input_: torch.Tensor, **kwargs) -> torch.Tensor:
        x = input_
        for module in self:
            x = module(x, **kwargs)
        return x


class ModelConfig:
    def __init__(
        self,
        hidden_size: int = 128,
        ffn_hidden_size: int = 512,
        layers: int = 1,
    ):
        self._hidden_size = hidden_size
        self._ffn_hidden_size = ffn_hidden_size
        self._layers = layers

    def build(self):
        from transformer_engine.pytorch.module.layernorm_mlp import LayerNormMLP
        from transformer_engine.pytorch.distributed import checkpoint

        ln_list, sln_list = [], []
        for _ in range(self._layers):
            ln = LayerNormMLP(self._hidden_size, self._ffn_hidden_size).to(device)
            sln = LayerNormMLP(self._hidden_size, self._ffn_hidden_size).to(device)
            # Copy parameters from ln to sln to ensure identical weights
            with torch.no_grad():
                sln.ln_weight = torch.nn.Parameter(ln.ln_weight.clone())
                if ln.layer_norm_bias is not None:
                    sln.layer_norm_bias = torch.nn.Parameter(ln.layer_norm_bias.clone())
                sln.fc1_weight = torch.nn.Parameter(ln.fc1_weight.clone())
                sln.fc2_weight = torch.nn.Parameter(ln.fc2_weight.clone())
                sln.fc1_bias = torch.nn.Parameter(ln.fc1_bias.clone())
                sln.fc2_bias = torch.nn.Parameter(ln.fc2_bias.clone())
            ln_list.append(ln)
            sln_list.append(sln)

        ln_model = _Sequential(*ln_list)

        class _CheckpointedSequential(torch.nn.Module):
            """Sequential model that applies checkpoint to each module."""

            def __init__(self, modules):
                super().__init__()
                self.modules_list = torch.nn.ModuleList(modules)

            def forward(self, x, **kwargs):
                for module in self.modules_list:
                    x = checkpoint(module, x, **kwargs)
                return x

        sln_model = _CheckpointedSequential(sln_list)

        return ln_model, sln_model


config = {
    "small": ModelConfig(128, 512, 2),
    "medium": ModelConfig(512, 2048, 2),
}

seq_sizes = [2**7, 2**10]


def _warmup(model, tensor):
    for _ in range(3):
        model(tensor).sum().backward()


def _run_fwd(model, tensor):
    if torch.npu.is_available():
        torch.npu.reset_peak_memory_stats(device)
        torch.npu.synchronize()
        start_mem = torch.npu.memory_allocated(device)
    else:
        start_mem = 0

    out = model(tensor)

    if torch.npu.is_available():
        torch.npu.synchronize()
        peak_mem = torch.npu.max_memory_allocated(device)
        mem = float(peak_mem - start_mem)
    else:
        mem = 0.0

    return out, mem


def _run_bwd(model, out):
    model.zero_grad(set_to_none=False)
    loss = out.sum()

    if torch.npu.is_available():
        torch.npu.reset_peak_memory_stats(device)
        torch.npu.synchronize()
        start_mem = torch.npu.memory_allocated(device)

    loss.backward()

    if torch.npu.is_available():
        torch.npu.synchronize()
        peak_mem = torch.npu.max_memory_allocated(device)
        mem = float(peak_mem - start_mem)
    else:
        mem = 0.0

    param_grads = _collect_param_grads(model)
    return param_grads, mem


def _max_diff(ref, other):
    """Return max absolute difference between two tensors or collections."""
    if ref is None or other is None:
        return 0.0
    if isinstance(ref, (list, tuple)):
        diffs = [_max_diff(r, o) for r, o in zip(ref, other)]
        return max(diffs) if diffs else 0.0
    return torch.max(torch.abs(ref.detach() - other.detach())).item()


def _collect_param_grads(model):
    grads = {}
    for name, param in model.named_parameters():
        if param.grad is None:
            continue
        key = _param_key(name)
        if key is not None:
            grads[key] = param.grad.detach().clone()
    return grads


def _param_key(name):
    return name.split(".")[-1]


@pytest.mark.parametrize("size", config.keys())
@pytest.mark.parametrize("seq_size", seq_sizes)
def test_selective_activation_checkpoint(size, seq_size):

    ln_model, sln_model = config[size].build()
    data = torch.randn((seq_size, 1, config[size]._hidden_size), device=device, requires_grad=True)

    # Run baseline model (no checkpointing)
    _warmup(ln_model, data)
    ln_fwd_out, ln_fwd_mem = _run_fwd(ln_model, data)
    ln_grads, ln_bwd_mem = _run_bwd(ln_model, ln_fwd_out)

    # Run checkpointed model
    _warmup(sln_model, data)
    sln_fwd_out, sln_fwd_mem = _run_fwd(sln_model, data)
    sln_grads, sln_bwd_mem = _run_bwd(sln_model, sln_fwd_out)

    # Verify memory reduction (checkpointing should reduce forward memory)
    if torch.npu.is_available():
        assert ln_fwd_mem > sln_fwd_mem, (
            "selective activation checkpointing does not reduce forward memory, "
            f"ln_fwd_mem={ln_fwd_mem}, sln_fwd_mem={sln_fwd_mem}"
        )

    # Verify outputs are numerically identical
    diff = _max_diff(ln_fwd_out, sln_fwd_out)
    assert diff == 0.0, f"outputs are not equal! maximum difference {diff}"

    # Verify gradients are numerically identical
    for key in [
        "ln_weight",
        "layer_norm_bias",
        "fc1_weight",
        "fc1_bias",
        "fc2_weight",
        "fc2_bias",
    ]:
        diff = _max_diff(ln_grads[key], sln_grads[key])
        assert diff == 0.0, f"gradients for {key} are not equal! maximum difference: {diff}"