import pytest
import torch

from model.wan_kv_mmd import _accumulate_fp32_pixel_cotangent


def test_pixel_cotangents_accumulate_immediately_in_fp32():
    accumulated = None
    shape = None
    expected = torch.zeros(2, 3, 4, dtype=torch.float32)
    norms = []
    for index, dtype in enumerate(
        (torch.bfloat16, torch.float16, torch.float32), start=1
    ):
        value = torch.full((2, 3, 4), float(index), dtype=dtype)
        expected.add_(value.float())
        accumulated, shape, norm = _accumulate_fp32_pixel_cotangent(
            accumulated, value, shape, branch_name=f"branch{index}")
        norms.append(norm)
    assert accumulated.dtype == torch.float32
    torch.testing.assert_close(accumulated, expected)
    assert shape == (2, 3, 4)
    assert all(item.ndim == 0 for item in norms)


def test_pixel_cotangent_accumulation_rejects_shape_and_nonfinite():
    first, shape, _ = _accumulate_fp32_pixel_cotangent(
        None, torch.ones(1, 2), None, branch_name="first")
    with pytest.raises(ValueError, match="shapes differ"):
        _accumulate_fp32_pixel_cotangent(
            first, torch.ones(1, 3), shape, branch_name="second")
    with pytest.raises(FloatingPointError, match="non-finite"):
        _accumulate_fp32_pixel_cotangent(
            None, torch.tensor([float("nan")]), None, branch_name="bad")