"""Tests for vsa_utils — standalone VSA metadata utilities.

These tests use CPU index computation except for an optional multi-GPU
regression covering cache isolation between CUDA devices.
"""

import math
import pytest
import torch

from fastvideo_kernel.vsa_utils import (
    VSA_TILE_SIZE,
    _canonicalize_device,
    get_tile_partition_indices,
    get_reverse_tile_partition_indices,
    construct_variable_block_sizes,
    get_non_pad_index,
    build_vsa_metadata,
)


class TestDeviceCanonicalization:

    def test_unindexed_cuda_resolves_current_device(self, monkeypatch):
        monkeypatch.setattr(torch.cuda, "current_device", lambda: 3)
        assert _canonicalize_device("cuda") == torch.device("cuda:3")
        assert _canonicalize_device(torch.device("cuda")) == torch.device("cuda:3")

    @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices")
    def test_metadata_cache_isolated_between_cuda_devices(self):
        shape = (7, 8, 8)
        tensor_keys = (
            "tile_partition_indices",
            "reverse_tile_partition_indices",
            "variable_block_sizes",
            "non_pad_index",
        )

        with torch.cuda.device(0):
            metadata_0 = build_vsa_metadata(shape, device="cuda")
        with torch.cuda.device(1):
            metadata_1 = build_vsa_metadata(shape, device="cuda")

        for key in tensor_keys:
            assert metadata_0[key].device == torch.device("cuda:0")
            assert metadata_1[key].device == torch.device("cuda:1")


class TestGetTilePartitionIndices:

    @pytest.mark.parametrize("dit_seq_shape,tile_size", [
        ((8, 16, 16), (4, 4, 4)),
        ((4, 8, 8), (4, 4, 4)),
        ((9, 10, 7), (4, 4, 4)),
    ])
    def test_is_valid_permutation(self, dit_seq_shape, tile_size):
        """Output must be a permutation of 0..N-1."""
        device = torch.device("cpu")
        idx = get_tile_partition_indices(dit_seq_shape, tile_size, device)
        n = math.prod(dit_seq_shape)
        assert idx.shape == (n, )
        assert idx.dtype == torch.long
        assert set(idx.tolist()) == set(range(n))

    def test_exact_values_small(self):
        """Manually verify a small case: (2,2,2) with tile (2,2,2) = 1 tile."""
        device = torch.device("cpu")
        idx = get_tile_partition_indices((2, 2, 2), (2, 2, 2), device)
        assert idx.tolist() == list(range(8))

    def test_non_divisible_shape(self):
        """When shape doesn't divide evenly by tile_size, all tokens still covered."""
        device = torch.device("cpu")
        shape = (5, 7, 3)
        idx = get_tile_partition_indices(shape, (4, 4, 4), device)
        assert idx.shape == (5 * 7 * 3, )
        assert set(idx.tolist()) == set(range(5 * 7 * 3))


class TestGetReverseTilePartitionIndices:

    @pytest.mark.parametrize("dit_seq_shape", [
        (8, 16, 16),
        (9, 10, 7),
    ])
    def test_inverse_of_forward(self, dit_seq_shape):
        """reverse[forward[i]] == i for all i."""
        device = torch.device("cpu")
        tile_size = (4, 4, 4)
        fwd = get_tile_partition_indices(dit_seq_shape, tile_size, device)
        rev = get_reverse_tile_partition_indices(dit_seq_shape, tile_size, device)
        n = math.prod(dit_seq_shape)
        identity = torch.arange(n, device=device)
        assert torch.equal(rev[fwd], identity)
        assert torch.equal(fwd[rev], identity)


class TestConstructVariableBlockSizes:

    def test_sum_equals_total_tokens(self):
        """Sum of block sizes must equal T*H*W."""
        device = torch.device("cpu")
        shape = (8, 16, 16)
        tile_size = (4, 4, 4)
        num_tiles = tuple(math.ceil(s / t) for s, t in zip(shape, tile_size))
        vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size)
        assert vbs.sum().item() == math.prod(shape)

    def test_max_block_size(self):
        """No block can exceed tile volume."""
        device = torch.device("cpu")
        shape = (9, 10, 7)
        tile_size = (4, 4, 4)
        num_tiles = tuple(math.ceil(s / t) for s, t in zip(shape, tile_size))
        vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size)
        assert vbs.max().item() <= math.prod(tile_size)

    def test_num_blocks(self):
        """Number of blocks = product of num_tiles."""
        device = torch.device("cpu")
        shape = (8, 16, 16)
        tile_size = (4, 4, 4)
        num_tiles = tuple(math.ceil(s / t) for s, t in zip(shape, tile_size))
        vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size)
        assert vbs.shape[0] == math.prod(num_tiles)

    def test_exact_divisible(self):
        """When perfectly divisible, all blocks have the same size."""
        device = torch.device("cpu")
        shape = (8, 8, 8)
        tile_size = (4, 4, 4)
        num_tiles = (2, 2, 2)
        vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size)
        assert (vbs == 64).all()

    def test_non_divisible_last_tile_smaller(self):
        """When not divisible, at least one block is smaller than max."""
        device = torch.device("cpu")
        shape = (9, 8, 8)
        tile_size = (4, 4, 4)
        num_tiles = (3, 2, 2)
        vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size)
        assert vbs.min().item() < math.prod(tile_size)

    def test_custom_tile_size(self):
        """tile_size parameter overrides default VSA_TILE_SIZE."""
        device = torch.device("cpu")
        shape = (6, 8, 16)
        tile_size = (2, 4, 8)
        num_tiles = (3, 2, 2)
        vbs = construct_variable_block_sizes(shape, num_tiles, device, tile_size)
        assert vbs.sum().item() == math.prod(shape)
        assert (vbs == 64).all()


class TestGetNonPadIndex:

    def test_length_equals_sum_block_sizes(self):
        """Output length must equal sum of variable_block_sizes."""
        vbs = torch.tensor([32, 48, 64], dtype=torch.long)
        idx = get_non_pad_index(vbs, 64)
        assert idx.shape[0] == 32 + 48 + 64

    def test_indices_in_valid_range(self):
        """All indices must be in [0, num_blocks * max_block_size)."""
        vbs = torch.tensor([32, 48], dtype=torch.long)
        idx = get_non_pad_index(vbs, 64)
        assert idx.min().item() >= 0
        assert idx.max().item() < 2 * 64

    def test_block_boundary_alignment(self):
        """First token of block i starts at i * max_block_size."""
        vbs = torch.tensor([20, 40], dtype=torch.long)
        idx = get_non_pad_index(vbs, 64)
        assert idx[0].item() == 0
        assert idx[20].item() == 64

    def test_full_blocks(self):
        """When all blocks are full, output is just 0..N-1."""
        vbs = torch.tensor([64, 64], dtype=torch.long)
        idx = get_non_pad_index(vbs, 64)
        assert torch.equal(idx, torch.arange(128))


class TestBuildVsaMetadata:

    def test_all_keys_present(self):
        """build_vsa_metadata returns all expected keys."""
        meta = build_vsa_metadata((8, 16, 16), device="cpu")
        expected_keys = {
            "tile_partition_indices",
            "reverse_tile_partition_indices",
            "variable_block_sizes",
            "non_pad_index",
            "num_tiles",
            "max_block_size",
        }
        assert set(meta.keys()) == expected_keys

    def test_types(self):
        """Return types are correct."""
        meta = build_vsa_metadata((8, 16, 16), device="cpu")
        assert isinstance(meta["tile_partition_indices"], torch.Tensor)
        assert isinstance(meta["reverse_tile_partition_indices"], torch.Tensor)
        assert isinstance(meta["variable_block_sizes"], torch.Tensor)
        assert isinstance(meta["non_pad_index"], torch.Tensor)
        assert isinstance(meta["num_tiles"], tuple)
        assert isinstance(meta["max_block_size"], int)

    def test_num_tiles_correct(self):
        meta = build_vsa_metadata((9, 10, 7), tile_size=(4, 4, 4), device="cpu")
        assert meta["num_tiles"] == (3, 3, 2)
        assert meta["max_block_size"] == 64

    @pytest.mark.parametrize("tile_size,expected_num_tiles,expected_block_size", [
        ((2, 4, 8), (3, 2, 2), 64),
        ((4, 8, 8), (2, 1, 2), 256),
    ])
    def test_supported_custom_tile_size(self, tile_size, expected_num_tiles, expected_block_size):
        meta = build_vsa_metadata((6, 8, 16), tile_size=tile_size, device="cpu")
        assert meta["num_tiles"] == expected_num_tiles
        assert meta["max_block_size"] == expected_block_size

    def test_unsupported_tile_volume(self):
        with pytest.raises(ValueError, match="Unsupported VSA tile volume 27"):
            build_vsa_metadata((6, 6, 6), tile_size=(3, 3, 3), device="cpu")

    def test_consistency(self):
        """All components are internally consistent."""
        shape = (8, 16, 16)
        meta = build_vsa_metadata(shape, device="cpu")
        n = math.prod(shape)
        assert meta["tile_partition_indices"].shape == (n, )
        assert meta["reverse_tile_partition_indices"].shape == (n, )
        assert meta["variable_block_sizes"].sum().item() == n
        assert meta["non_pad_index"].shape[0] == n


class TestConsistencyWithFramework:
    """Verify vsa_utils matches the framework-level functions exactly.

    Only runs if fastvideo is importable (skip otherwise).
    """

    @pytest.fixture(autouse=True)
    def _skip_if_no_framework(self):
        try:
            from fastvideo.attention.backends.video_sparse_attn import (
                get_tile_partition_indices as fw_get_tile, )
        except ImportError:
            pytest.skip("fastvideo framework not installed")

    @pytest.mark.parametrize("shape", [(8, 16, 16), (9, 10, 7)])
    def test_tile_indices_match(self, shape):
        from fastvideo.attention.backends.video_sparse_attn import (
            get_tile_partition_indices as fw_get_tile, )
        device = torch.device("cpu")
        tile_size = (4, 4, 4)
        ours = get_tile_partition_indices(shape, tile_size, device)
        theirs = fw_get_tile(shape, tile_size, device)
        assert torch.equal(ours, theirs)

    @pytest.mark.parametrize("shape", [(8, 16, 16), (9, 10, 7)])
    def test_variable_block_sizes_match(self, shape):
        from fastvideo.attention.backends.video_sparse_attn import (
            construct_variable_block_sizes as fw_construct_vbs, )
        device = torch.device("cpu")
        tile_size = (4, 4, 4)
        num_tiles = tuple(math.ceil(s / t) for s, t in zip(shape, tile_size))
        ours = construct_variable_block_sizes(shape, num_tiles, device, tile_size)
        theirs = fw_construct_vbs(shape, num_tiles, device)
        assert torch.equal(ours, theirs)


if __name__ == "__main__":
    pytest.main([__file__, "-v", "-s"])