# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU coverage for MXFP4 checkpoint loading and the W4A8 dtype contract."""

import sys
from types import ModuleType

import pytest
import torch
from vllm.config.load import LoadConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader

from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.forward_context import (
    ForwardContext,
    override_forward_context,
    set_forward_context_denoise_step_idx,
)
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.quantization.mxfp4_config import (
    DiffusionMXFP4Config,
    DiffusionMXFP4DualScaleMixedConfig,
    NPUMxfp4DualScaleLinearMethod,
    NPUMxfp4LinearMethod,
    NPUMxfp4OnlineLinearMethod,
)

pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]


@pytest.fixture(autouse=True)
def _patch_tp_state(monkeypatch):
    monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
    monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)


class _CheckpointLayer(torch.nn.Module):
    def __init__(self, method, dtype=torch.bfloat16):
        super().__init__()
        method.create_weights(
            self,
            input_size_per_partition=512,
            output_partition_sizes=[2],
            input_size=512,
            output_size=2,
            params_dtype=dtype,
            weight_loader=default_weight_loader,
        )

    def load_weights(self, weights):
        params = dict(self.named_parameters())
        loaded = set()
        for name, value in weights:
            if name in params:
                params[name].weight_loader(params[name], value)
                loaded.add(name)
        return loaded


def _load_checkpoint(layer, config, weights, monkeypatch):
    od_config = OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    loader.counter_before_loading_weights = 0.0
    monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
    loader.load_weights(layer)


@pytest.mark.parametrize("serialized", [False, True])
def test_single_scale_requires_checkpoint_scale_only_when_serialized(serialized, monkeypatch):
    config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized)
    # Exercise the shared implementation directly as well as the online MRO
    # below: checkpoint requirements must follow the configured storage mode.
    layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
    weights = {"weight": torch.ones((2, 512), dtype=torch.bfloat16)}
    if serialized:
        with pytest.raises(ValueError, match="Required weights.*weight_scale"):
            _load_checkpoint(layer, config, weights, monkeypatch)
    else:
        _load_checkpoint(layer, config, weights, monkeypatch)


@pytest.mark.parametrize("scale_alg", [0, 2])
def test_online_mxfp4_loads_bf16_checkpoint_without_scale(scale_alg, monkeypatch):
    config = DiffusionMXFP4Config(mxfp4_scale_alg=scale_alg)
    method = NPUMxfp4OnlineLinearMethod(config)
    layer = _CheckpointLayer(method)
    assert set(dict(layer.named_parameters())) == {"weight"}
    assert layer.weight.is_meta

    weight = torch.arange(1024, dtype=torch.bfloat16).reshape(2, 512)
    packed_weight = torch.zeros((2, 256), dtype=torch.uint8)
    generated_scale = torch.full((2, 16), 127, dtype=torch.uint8)
    quantized_inputs = []

    def quantize(value, **kwargs):
        quantized_inputs.append(value.clone())
        expected = dict(dst_type="fp4", axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
        if scale_alg == 2:
            expected["dst_type_max"] = 7.25
        assert kwargs == expected
        return packed_weight, generated_scale

    # Keep real creation, lazy loading, processing, and strict loader checks;
    # substitute only the device quantization operator for this CPU test.
    npu = ModuleType("torch_npu")
    npu.__dict__.update(float4_e2m1fn_x2="fp4", npu_dynamic_mx_quant=quantize)
    monkeypatch.setitem(sys.modules, "torch_npu", npu)
    _load_checkpoint(layer, config, {"weight": weight}, monkeypatch)

    assert len(quantized_inputs) == 1
    torch.testing.assert_close(quantized_inputs[0], weight, rtol=0, atol=0)
    torch.testing.assert_close(layer.weight, packed_weight, rtol=0, atol=0)
    torch.testing.assert_close(layer.weight_scale, generated_scale.reshape(2, 8, 2), rtol=0, atol=0)
    assert not getattr(layer.weight_scale, "is_checkpoint_required", False)
    method.process_weights_after_loading(layer)
    assert len(quantized_inputs) == 1


@pytest.mark.parametrize("require_smooth", [False, True])
@pytest.mark.parametrize("has_smooth", [False, True])
def test_single_scale_checkpoint_smooth_default_and_required(require_smooth, has_smooth, monkeypatch):
    config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=require_smooth)
    layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
    weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
    scale = torch.linspace(0.5, 3.0, 512)
    if has_smooth:
        weights["mul_scale"] = scale
    else:
        weights.pop("mul_scale")

    if require_smooth and not has_smooth:
        with pytest.raises(ValueError, match="mul_scale"):
            _load_checkpoint(layer, config, weights, monkeypatch)
    else:
        _load_checkpoint(layer, config, weights, monkeypatch)
        torch.testing.assert_close(layer.mul_scale, scale if has_smooth else torch.ones(512), rtol=0, atol=0)


@pytest.mark.parametrize("bad_scale", ["shape", "name"])
def test_declared_single_scale_smooth_rejects_malformed_checkpoint(bad_scale, monkeypatch):
    config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
    layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
    weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
    if bad_scale == "shape":
        weights["mul_scale"] = torch.ones(511)
        with pytest.raises(AssertionError):
            _load_checkpoint(layer, config, weights, monkeypatch)
    else:
        weights["multiply_scale"] = weights.pop("mul_scale")
        with pytest.raises(ValueError, match="mul_scale"):
            _load_checkpoint(layer, config, weights, monkeypatch)


@pytest.mark.parametrize("required", [False, True])
@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), 0.0, -1.0])
def test_single_scale_rejects_invalid_smooth_values(required, bad_value, monkeypatch):
    config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=required)
    layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
    weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
    weights["mul_scale"][7] = bad_value
    with pytest.raises(ValueError, match="finite, strictly positive"):
        _load_checkpoint(layer, config, weights, monkeypatch)


@pytest.mark.parametrize("bad_scale", [torch.ones(512, dtype=torch.int32), torch.ones(1, 512)])
def test_single_scale_rejects_invalid_smooth_type_or_rank(bad_scale, monkeypatch):
    config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
    layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
    weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
    weights["mul_scale"] = bad_scale
    with pytest.raises(ValueError, match="one-dimensional floating-point"):
        _load_checkpoint(layer, config, weights, monkeypatch)


@pytest.mark.parametrize("bad_value", [1e10, 1e-10])
def test_single_scale_rejects_smooth_overflow_or_underflow_after_cast(bad_value, monkeypatch):
    config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
    method = NPUMxfp4LinearMethod(config)
    layer = _CheckpointLayer(method, torch.float16)
    layer.mul_scale.data.fill_(bad_value)
    # Validation must fail before any device packing or GEMM.
    monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu"))
    with pytest.raises(ValueError, match="finite, strictly positive"):
        method.process_weights_after_loading(layer)


@pytest.mark.parametrize(
    ("mode", "missing"),
    [
        ("single", "weight"),
        ("single", "weight_scale"),
        ("dual", "weight"),
        ("dual", "weight_scale"),
        ("dual", "weight_dual_scale"),
        ("dual", "mul_scale"),
    ],
)
def test_offline_mxfp4_required_tensors_do_not_use_legacy_scale_tolerance(mode, missing, monkeypatch):
    if mode == "single":
        config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
        method = NPUMxfp4LinearMethod(config)
    else:
        config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
        method = NPUMxfp4DualScaleLinearMethod(config)
    layer = _CheckpointLayer(method)
    weights = {name: torch.ones_like(param) for name, param in layer.named_parameters() if name != missing}
    with pytest.raises(ValueError, match=missing):
        _load_checkpoint(layer, config, weights, monkeypatch)


@pytest.mark.parametrize("rank", [0, 1])
def test_single_scale_smooth_uses_real_row_parallel_loader(rank):
    from vllm.model_executor.layers.linear import RowParallelLinear

    layer = _CheckpointLayer(NPUMxfp4LinearMethod(DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)))
    layer.tp_rank = rank
    checkpoint_scale = torch.arange(1024, dtype=torch.float32)
    RowParallelLinear.weight_loader(layer, layer.mul_scale, checkpoint_scale)
    torch.testing.assert_close(layer.mul_scale, checkpoint_scale[rank * 512 : (rank + 1) * 512], rtol=0, atol=0)


@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("fallback", [False, True])
@pytest.mark.parametrize("with_bias", [False, True])
@pytest.mark.parametrize("scale_alg", [0, 2])
def test_single_scale_smooth_precedes_quantization_and_preserves_output(
    dtype, fallback, with_bias, scale_alg, monkeypatch
):
    config = DiffusionMXFP4Config(
        is_checkpoint_mxfp4_serialized=True, w4a8_fallback_steps=[0], mxfp4_scale_alg=scale_alg
    )
    method = NPUMxfp4LinearMethod(config)
    layer = _CheckpointLayer(method, dtype)
    weights = {
        "weight": torch.linspace(-0.5, 1.5, 1024, dtype=dtype).reshape(2, 512),
        "weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
        "mul_scale": torch.linspace(0.5, 3.0, 512),
    }
    _load_checkpoint(layer, config, weights, monkeypatch)
    x = torch.linspace(-1.0, 2.0, 3072, dtype=dtype).reshape(2, 3, 512)
    bias = torch.tensor([0.5, -1.0], dtype=dtype) if with_bias else None
    quantized_inputs = []

    # Only the two device operators are substituted. Exercise real apply(),
    # step selection, Smooth multiplication, bias handling, and reshape logic.
    npu = ModuleType("torch_npu")

    def quantize(value, **kwargs):
        quantized_inputs.append(value.clone())
        expected = dict(
            dst_type=torch.float8_e4m3fn if fallback else "fp4",
            axis=-1,
            block_size=32,
            round_mode="rint",
            scale_alg=0 if fallback else scale_alg,
        )
        if not fallback and scale_alg == 2:
            expected["dst_type_max"] = 7.25
        assert kwargs == expected
        return value, torch.ones((6, 16), dtype=torch.uint8)

    def matmul(x_q, weight, weight_scale, *, bias, output_dtype, **kwargs):
        if fallback:
            assert output_dtype == torch.bfloat16
            if bias is not None:
                assert bias.dtype == output_dtype
                assert bias.shape == (1, 2)
        result = x_q.float() @ weight.float()
        if bias is not None:
            result += bias.float()
        return result.to(output_dtype)

    npu.__dict__.update(
        float8_e4m3fn=torch.float8_e4m3fn,
        float8_e8m0fnu=torch.uint8,
        float4_e2m1fn_x2="fp4",
        npu_dynamic_mx_quant=quantize,
        npu_quant_matmul=matmul,
    )
    monkeypatch.setitem(sys.modules, "torch_npu", npu)
    with override_forward_context(ForwardContext()):
        set_forward_context_denoise_step_idx(0 if fallback else 1)
        result = method.apply(layer, x, bias)

    expected_input = x.reshape(-1, 512) * weights["mul_scale"].to(dtype)
    assert len(quantized_inputs) == 1
    torch.testing.assert_close(quantized_inputs[0], expected_input, rtol=0, atol=0)
    expected = expected_input.float() @ weights["weight"].float().T
    if bias is not None:
        expected += bias.float()
    if fallback:
        expected = expected.bfloat16()
    torch.testing.assert_close(result, expected.to(dtype).reshape(2, 3, 2), rtol=0, atol=0)
    assert result.dtype == dtype