"""
Run all C++ unit tests as a single pytest test.

The C++ tests conver generated index and tensor classes. These classes
work in both C++ and CUDA programs.
"""

from subprocess import CalledProcessError
from tempfile import NamedTemporaryFile
from typing import Callable
import os

from importlib_resources import files as importlib_resources_files
import pytest

import spio.generators as gen
import spio.compiler
from spio.util import env_var_is_true


ENABLE_CPP_TESTS = env_var_is_true("SPIO_ENABLE_CPP_TESTS")

UTEST_HEADER = '#include "utest.h"'

TEST_MODULES = [
    gen.compound_index,
    gen.tensor,
    gen.checkerboard,
    gen.fragment_index,
    gen.dim,
]

CPP_SOURCES = [
    "test_main.cpp",
    "test_tensor.cpp",
    "test_index.cpp",
    "test_fragment_index.cpp",
    "test_mathutil.cpp",
    "test_checkerboard_index.cpp",
    "test_strip_loader_params.cpp",
    "test_dim.cpp",
    "test_coordinates.cpp",
    "test_meta.cpp",
]

TEST_SOURCES = []

# Select specific C++ unit tests with a filter like this:
# "Tensor.checkerboardindex_derived_dim_subscript"
CPP_TESTS_FILTER = os.environ.get("SPIO_CPP_TESTS_FILTER", None)


@pytest.mark.skipif(
    not ENABLE_CPP_TESTS, reason="NVCC support not required by default."
)
def test_cpp_tests():
    """Compile and run all C++ unit tests.

    This PyTest test will succeed if all C++ unit tests pass.
    Run tests with "pytest -s test_cpp.py" to see the detailed output.
    """
    test_headers = [getattr(module, "header")() for module in TEST_MODULES]

    code = "\n".join([UTEST_HEADER] + test_headers + TEST_SOURCES)
    test_source_file = NamedTemporaryFile(prefix="spio_", suffix=".cpp")
    with open(test_source_file.name, "w", encoding="utf-8") as f:
        f.write(code)
    try:
        run_args = _format_utest_command_line(CPP_TESTS_FILTER)
        _compile_cpp_tests([test_source_file.name], run_args=run_args)
    except CalledProcessError as e:
        assert False, f"{e.stdout} {e.stderr}"


def _format_utest_command_line(test_filter: str = None) -> str:
    """Format a UTEST command line argument to filter tests."""
    if test_filter is None:
        return None
    return [f"--filter={test_filter}"]


def _cpp_test(func: Callable[[], str]):
    """Register the C++ source code generated by the given python function."""
    TEST_SOURCES.append(func())


@_cpp_test
def _test_generate_dim():
    """Return the C++ source code that tests a custom dimension class."""
    specs = [gen.Dim("i"), gen.Dim("j")]
    generated_code = gen.generate(specs, namespace="I_Dim_GenCode")
    return f"""
{generated_code}

UTEST(I, methods)
{{
    using namespace I_Dim_GenCode;
    EXPECT_EQ(I(7).get(), 7);
    EXPECT_TRUE(I(7) == I(7));
    EXPECT_TRUE(I(7) < I(8));
    EXPECT_TRUE(I(7) <= I(8));
    EXPECT_TRUE(I(8) > I(7));
    EXPECT_TRUE(I(8) >= I(7));
    EXPECT_TRUE(I(7) == I(7));
    EXPECT_TRUE(I(7) <= I(7));
    EXPECT_TRUE(I(7) >= I(7));
    EXPECT_TRUE(I(7) != I(8));
    EXPECT_TRUE(I(7) + I(8) == I(15));
    EXPECT_TRUE(I(8) - I(7) == I(1));
    EXPECT_TRUE(I(8) % I(3) == I(2));
    EXPECT_TRUE((I(32).fold<4>() == spio::Fold<I, 4>(8)));
    EXPECT_TRUE((I(32).cast<J>() == J(32)));
}}

UTEST(I, range)
{{
    using namespace I_Dim_GenCode;
    int j = 0;
    for (auto i : range(I(7))) {{
        EXPECT_EQ(i.get(), j++);
    }}
    EXPECT_EQ(j, 7);
}}
"""


@_cpp_test
def _test_generate_fold():
    specs = [
        gen.Dim("c"),
        gen.Fold("i", 64, fold_name="block_i"),
        gen.Fold("j", 64, fold_name="block_j"),
        # Test a fold of a fold. Produces using BLOCK_C8 = Fold<Fold<C, 8>, 4>;
        gen.Fold("c8", 4, fold_name="block_c8"),
    ]
    generated_code = gen.generate(specs, namespace="I_Fold_GenCode")
    return f"""
{generated_code}

UTEST(I_Fold, methods)
{{
    using namespace I_Fold_GenCode;
    EXPECT_EQ(BLOCK_I(7).get(), 7);
    EXPECT_TRUE(BLOCK_I(7).unfold() == I(7 * 64));
    EXPECT_TRUE((BLOCK_I(7).fold<32>() == spio::Fold<I, 32>(7 * 2)));

    EXPECT_EQ(BLOCK_C8(7).get(), 7);

    // Unfold the outer fold, producing Fold<C, 8>.
    EXPECT_TRUE((BLOCK_C8(7).unfold() == spio::Fold<C, 8>(7 * 4)));

    EXPECT_TRUE(BLOCK_I(7).cast<J>() == BLOCK_J(7));
}}
"""


@_cpp_test
def _test_generate_fold_with_init():
    """Test Fold generation with a BuiltIn initializer."""
    specs = [
        gen.Dim("i"),
        gen.Fold("i", 64, fold_name="block_i", init=gen.BuiltIn.BLOCK_IDX_Y),
    ]
    generated_code = gen.generate(specs, namespace="FoldInit_GenCode")
    return f"""
// Mock blockIdx for host testing
namespace {{
    struct {{ int x = 7, y = 128, z = 9; }} blockIdx;
}}

{generated_code}

UTEST(FoldInit, default_constructor)
{{
    using namespace FoldInit_GenCode;
    BLOCK_I bi;  // Uses default constructor with blockIdx.y = 128
    EXPECT_EQ(bi.get(), 128);
}}

UTEST(FoldInit, unfold)
{{
    using namespace FoldInit_GenCode;
    BLOCK_I bi;
    // blockIdx.y = 128, so unfold() should give I(128 * 64) = I(8192)
    EXPECT_TRUE(bi.unfold() == I(128 * 64));
}}
"""


@_cpp_test
def _test_generate_coordinates():
    """Test Coordinates generation with Dim and Fold.

    This mirrors the usage in test_mma_checkerboard.py:
        g.block_i = Fold("i", block_x, init=BuiltIn.BLOCK_IDX_Y)
        g.block_j = Fold("j", block_x, init=BuiltIn.BLOCK_IDX_X)
        g.BlockIdx = Coordinates(g.block_i, g.block_j)
    """
    block_i = gen.Fold("i", 64, fold_name="block_i", init=gen.BuiltIn.BLOCK_IDX_Y)
    block_j = gen.Fold("j", 64, fold_name="block_j", init=gen.BuiltIn.BLOCK_IDX_X)
    specs = [
        gen.Dim("i"),
        gen.Dim("j"),
        block_i,
        block_j,
        gen.Coordinates(block_i, block_j, coord_name="BlockIdx"),
    ]
    generated_code = gen.generate(specs, namespace="Coordinates_GenCode")
    return f"""
// Mock blockIdx for host testing - inside namespace to avoid conflicts
namespace Coordinates_GenCode {{
    struct {{ int x = 7, y = 128, z = 9; }} blockIdx;
}}

{generated_code}

UTEST(Coordinates, generate_with_folds)
{{
    using namespace Coordinates_GenCode;
    auto coords = BlockIdx();
    // blockIdx.y = 128, blockIdx.x = 7
    EXPECT_EQ(coords.get<BLOCK_I>().get(), 128);
    EXPECT_EQ(coords.get<BLOCK_J>().get(), 7);
}}
"""


@_cpp_test
def _test_generate_coordinates_with_anonymous_folds():
    """Test Coordinates generation with anonymous (unnamed) Folds.

    This mirrors the cleaner usage in test_mma_checkerboard.py where Folds
    are created inline without explicit fold_name - the generate() function
    automatically assigns names like _FOLD_1, _FOLD_2 (uppercase).
    """
    specs = [
        gen.Dim("i"),
        gen.Dim("j"),
        gen.Coordinates(
            gen.Fold("i", 64, init=gen.BuiltIn.BLOCK_IDX_Y),
            gen.Fold("j", 64, init=gen.BuiltIn.BLOCK_IDX_X),
            coord_name="BlockIdx",
        ),
    ]
    generated_code = gen.generate(specs, namespace="AnonCoordinates_GenCode")
    return f"""
// Mock blockIdx for host testing - inside namespace to avoid conflicts
namespace AnonCoordinates_GenCode {{
    struct {{ int x = 3, y = 42, z = 0; }} blockIdx;
}}

{generated_code}

UTEST(AnonCoordinates, anonymous_folds_are_generated)
{{
    using namespace AnonCoordinates_GenCode;
    auto coords = BlockIdx();
    // blockIdx.y = 42, blockIdx.x = 3
    // Anonymous folds get auto-generated names, but order is not guaranteed
    auto f1 = coords.get<_FOLD_1>().get();
    auto f2 = coords.get<_FOLD_2>().get();
    // Check that we got both values (in some order)
    EXPECT_EQ(std::max(f1, f2), 42);
    EXPECT_EQ(std::min(f1, f2), 3);
}}

UTEST(AnonCoordinates, anonymous_folds_unfold_correctly)
{{
    using namespace AnonCoordinates_GenCode;
    auto coords = BlockIdx();
    // Each fold has stride 64, so unfold multiplies by 64
    // Order of _FOLD_1 vs _FOLD_2 is not guaranteed
    auto f1 = coords.get<_FOLD_1>().unfold().get();
    auto f2 = coords.get<_FOLD_2>().unfold().get();
    // Check that we got both unfolded values (in some order)
    EXPECT_EQ(std::max(f1, f2), 42 * 64);
    EXPECT_EQ(std::min(f1, f2), 3 * 64);
}}
"""


@_cpp_test
def _test_wave_based_block_idx():
    """Test wave-based BlockIdx pattern used in test_mma_checkerboard.py.

    This tests the pattern where blocks iterate in waves:
    - First iterate over wave_size blocks in the i-dimension (block_i_local)
    - Then iterate over all blocks in the j-dimension (block_j)
    - Finally iterate over waves in the i-dimension (block_i_wave)
    """
    # Use values similar to test_mma_checkerboard.py test case
    block_m = 256
    block_n = 128
    wave_size = 4
    m = 8192
    n = 1024

    blocks_m = m // block_m  # 32
    blocks_n = n // block_n  # 8
    block_waves = blocks_m // wave_size  # 8
    block_local = wave_size  # 4
    wave_stride = wave_size * block_m  # 1024

    block_i_wave = gen.Fold("i", wave_stride, fold_name="block_i_wave")
    block_i_local = gen.Fold("i", block_m, fold_name="block_i_local")
    block_j = gen.Fold("j", block_n, fold_name="block_j")

    specs = [
        gen.Dim("i"),
        gen.Dim("j"),
        block_i_wave,
        block_i_local,
        block_j,
        gen.CompoundIndex(
            gen.Dims(
                block_i_wave=block_waves, block_j=blocks_n, block_i_local=block_local
            ),
            class_name="BlockIdx",
            init=gen.BuiltIn.BLOCK_IDX_X,
        ),
    ]
    generated_code = gen.generate(specs, namespace="WaveBlockIdx_GenCode")

    total_blocks = block_waves * blocks_n * block_local  # 256

    return f"""
// Mock blockIdx for host testing
namespace WaveBlockIdx_GenCode {{
    struct {{ int x = 0, y = 0, z = 0; }} blockIdx;
}}

{generated_code}

UTEST(WaveBlockIdx, decomposition_and_unfold)
{{
    using namespace WaveBlockIdx_GenCode;

    // Test a representative set of blockIdx.x values
    struct TestCase {{
        int block_idx_x;
        int expected_local, expected_j, expected_wave;
    }};

    TestCase cases[] = {{
        {{0, 0, 0, 0}},      // First block
        {{3, 3, 0, 0}},      // Last in first local group
        {{4, 0, 1, 0}},      // First in second j-block
        {{31, 3, 7, 0}},     // Last in first wave
        {{32, 0, 0, 1}},     // First in second wave
        {{45, 1, 3, 1}},     // Middle of second wave
        {{255, 3, 7, 7}},    // Last block overall
    }};

    for (const auto& tc : cases) {{
        WaveBlockIdx_GenCode::blockIdx.x = tc.block_idx_x;
        BlockIdx idx;

        // Check decomposition
        EXPECT_EQ(idx.get<BLOCK_I_LOCAL>().get(), tc.expected_local);
        EXPECT_EQ(idx.get<BLOCK_J>().get(), tc.expected_j);
        EXPECT_EQ(idx.get<BLOCK_I_WAVE>().get(), tc.expected_wave);

        // Check unfold produces correct I and J offsets
        int expected_i = tc.expected_wave * {wave_stride} + tc.expected_local * {block_m};
        int expected_j = tc.expected_j * {block_n};

        EXPECT_TRUE(idx.get<BLOCK_I_WAVE>() == I(tc.expected_wave * {wave_stride}));
        EXPECT_TRUE(idx.get<BLOCK_I_LOCAL>() == I(tc.expected_local * {block_m}));
        EXPECT_TRUE(idx.get<BLOCK_J>() == J(expected_j));

        // Verify coordinates combine correctly
        EXPECT_TRUE(idx.coordinates() == make_coordinates(I(expected_i), J(expected_j)));
    }}

    // Verify total size
    EXPECT_EQ(BlockIdx::size(), {total_blocks});
}}
"""


@_cpp_test
def _test_generate_index():
    """Return the C++ source code that tests a custom index class."""

    specs = [gen.CompoundIndex(gen.Dims(n=4, h=32, w=64, c=128), class_name="MyIndex")]
    generated_code = gen.generate(specs, namespace="MyIndex_GenCode")
    size = 4 * 32 * 64 * 128
    test_code = f"""
{generated_code}

UTEST(MyIndex, index_from_offset)
{{
    using namespace MyIndex_GenCode;
    int offset = 532523;
    MyIndex idx(offset);
    EXPECT_TRUE(idx.get<N>() == N(offset / (32 * 64 * 128) ) );
    EXPECT_TRUE(idx.get<H>() == H( (offset / (64 * 128)) % 32) );
    EXPECT_TRUE(idx.get<W>() == W( (offset / 128) % 64));
    EXPECT_TRUE(idx.get<C>() == C(offset % 128));
}}

UTEST(MyIndex, size)
{{
    using namespace MyIndex_GenCode;
    EXPECT_EQ(MyIndex::size(), {size});
}}

UTEST(MyIndex, dim_sizes)
{{
    using namespace MyIndex_GenCode;
    EXPECT_TRUE(MyIndex::size<N>() == N(4));
    EXPECT_TRUE(MyIndex::size<H>() == H(32));
    EXPECT_TRUE(MyIndex::size<W>() == W(64));
    EXPECT_TRUE(MyIndex::size<C>() == C(128));
}}
"""
    return test_code


@_cpp_test
def _test_generate_index_with_init():
    """Test CompoundIndex generation with a BuiltIn initializer."""
    specs = [
        gen.CompoundIndex(
            gen.Dims(i=4, j=8),
            class_name="WarpIndex",
            init=gen.BuiltIn.THREAD_IDX_X,
        )
    ]
    generated_code = gen.generate(specs, namespace="IndexInit_GenCode")
    return f"""
// Mock threadIdx for host testing
namespace {{
    struct {{ int x = 17, y = 0, z = 0; }} threadIdx;
}}

{generated_code}

UTEST(IndexInit, default_constructor)
{{
    using namespace IndexInit_GenCode;
    WarpIndex idx;  // Uses default constructor with threadIdx.x = 17
    // i has stride 8, j has stride 1
    // offset 17: i = 17 / 8 = 2, j = 17 % 8 = 1
    EXPECT_TRUE(idx.get<I>() == I(2));
    EXPECT_TRUE(idx.get<J>() == J(1));
}}

UTEST(IndexInit, size)
{{
    using namespace IndexInit_GenCode;
    EXPECT_EQ(WarpIndex::size(), 32);
}}
"""


@_cpp_test
def _test_generate_checkerboard_index():
    specs = [
        gen.Checkerboard("r", "c8", class_name="Checkers", offset_dim="offset", ranks=8)
    ]
    generated_code = gen.generate(specs, namespace="IndexSpec_GenCode")
    code = f"""
{generated_code}

UTEST(IndexSpec, checkerboard_fused_dim)
{{
    using namespace IndexSpec_GenCode;
    using namespace spio;
    for (int offset = 0; offset < 128; ++offset) {{
        Checkers idx(offset);
        int row = offset / 8;
        int pair = offset / 2;
        int color = ((offset & 1) ^ (row & 1));
        EXPECT_EQ((idx.get<R>().get()), pair);
        EXPECT_EQ((idx.get<Fold<C, 8>>().get()),  color);
        EXPECT_EQ((idx.get<OFFSET>().get()), offset);
        EXPECT_EQ(idx.pair(), pair);
        EXPECT_EQ(idx.color(), color);
    }}
}}
"""
    return code


@_cpp_test
def _test_contiguous_tensor():
    """Return the C++ source code that tests a custom tensor class."""
    n = 7
    h = 16
    w = 33
    c = 42

    tensor_spec = gen.Tensor(
        gen.dtype.float, gen.Dims(n=n, h=h, w=w, c=c), class_name="ContiguousTensor"
    )

    size = n * h * w * c
    assert tensor_spec.size == size
    assert tensor_spec.num_bytes == size * gen.dtype.float.value.size

    generated_code = gen.generate([tensor_spec], namespace="DenseTensor_GenCode")
    test_code = f"""

{generated_code}

UTEST(ContiguousTensor, offset_from_tensor)
{{
    using namespace DenseTensor_GenCode;
    constexpr int N_Size = {n};
    constexpr int H_Size = {h};
    constexpr int W_Size = {w};
    constexpr int C_Size = {c};
    constexpr int size = N_Size * H_Size * W_Size * C_Size;
    constexpr size_t num_bytes = sizeof(float) * size;

    EXPECT_EQ(ContiguousTensor::size(), size);
    EXPECT_EQ(ContiguousTensor::num_bytes(), num_bytes);

    EXPECT_EQ(ContiguousTensor::size<N>().get(), N_Size);
    EXPECT_EQ(ContiguousTensor::size<H>().get(), H_Size);
    EXPECT_EQ(ContiguousTensor::size<W>().get(), W_Size);
    EXPECT_EQ(ContiguousTensor::size<C>().get(), C_Size);
    
    // Initialize the tensor with data.
    float data[ContiguousTensor::size()];
    ContiguousTensor tensor(data);
    for (int n = 0; n < N_Size; ++n) {{
        for (int h = 0; h < H_Size; ++h) {{
            for (int w = 0; w < W_Size; ++w) {{
                for (int c = 0; c < C_Size; ++c) {{
                    data[
                        n*(H_Size*W_Size*C_Size) +
                        h*(W_Size*C_Size) +
                        w*C_Size +
                        c
                    ] = n*(H_Size*W_Size*C_Size) + h*(W_Size*C_Size) + w*C_Size + c;
                }}
            }}
        }}
    }}

    // Test the use of named-index accessors and subscript operators with the tensor class.
    // Note all permutations of the subscript operators are equivalent. Because the operator
    // is overloaded for dimension type, the order of the subscript operator does not matter.
    for (auto n : range(tensor.size<N>())) {{
        for (auto h : range(tensor.size<H>())) {{
            for (auto w : range(tensor.size<W>())) {{
                for (auto c : range(tensor.size<C>())) {{
                    float val = n.get()*(H_Size*W_Size*C_Size) + h.get()*(W_Size*C_Size) + w.get()*C_Size + c.get();
                    EXPECT_EQ(*tensor[n][h][w][c], val);
                    EXPECT_EQ(*tensor[h][n][c][w], val);
                    EXPECT_EQ(*tensor[h][w][c][n], val);
                    EXPECT_EQ(*tensor[w][c][n][h], val);
                }}
            }}
        }}
    }}

    // Test the use of an integer offset with the tensor offset method.
    for (int offset = 0; offset < size; ++offset) {{
        EXPECT_EQ(*tensor.offset(offset), data[offset]);
    }}

    // Test the use of a generated index class with the tensor subscript operator.
    for (int offset = 0; offset < size; ++offset) {{
        ContiguousTensor::index_type idx(offset);
        EXPECT_EQ(*tensor[idx], data[offset]);
    }}

    int nn = 0;
    for (auto n : range(tensor.size<N>())) {{
        EXPECT_EQ(n.get(), nn++);
    }}
    EXPECT_EQ(nn, N_Size);
}}
"""
    return test_code


@_cpp_test
def _test_genrate_1d_contiguous_tensor():
    """Return the C++ source code that tests a custom 1D tensor class"""
    n = 7
    specs = [
        gen.Tensor(
            gen.dtype.float, gen.Dims(n=n), class_name="ContiguousTensor", constant=True
        ),
    ]
    generated_code = gen.generate(specs, namespace="ContiguousTensor_1D_GenCode")
    test_code = f"""

{generated_code}

UTEST(DensTensor1D, offset_from_tensor)
{{
    using namespace ContiguousTensor_1D_GenCode;
    constexpr int N_Size = {n};
    constexpr int size = N_Size;
    constexpr int num_bytes = static_cast<int>(sizeof(float) * size);

    float data[N_Size];
    for (int n = 0; n < N_Size; ++n) {{
        data[n] = n;
    }}
    ContiguousTensor tensor(data);
    EXPECT_EQ(tensor.size(), size);
    EXPECT_EQ(tensor.num_bytes(), num_bytes);

    for (auto n : range(ContiguousTensor::size<N>())) {{
        // Tensor subscript accessors.
        EXPECT_EQ(*tensor[n], data[n.get()]);
    }}
}}
"""
    return test_code


@_cpp_test
def _test_generate_tensor_with_strides():
    """Return the C++ source code that tests a custom tensor class."""
    n = 7
    h = 16
    w = 33
    c = 42

    stride_w = c + 2
    stride_h = (w + 1) * stride_w

    specs = [
        gen.Tensor(
            gen.dtype.float,
            gen.Dims(n=n, h=h, w=w, c=c),
            class_name="StrideTensor",
            strides=gen.Strides(h=stride_h, w=stride_w),
            constant=True,
        ),
    ]
    generated_code = gen.generate(specs, namespace="StrideTensor_GenCode")
    test_code = f"""

{generated_code}

UTEST(StrideTensor, offset_from_tensor)
{{
    using namespace StrideTensor_GenCode;
    constexpr int N_Size = {n};
    constexpr int H_Size = {h};
    constexpr int W_Size = {w};
    constexpr int C_Size = {c};

    constexpr int stride_w = {stride_w};
    constexpr int stride_h = {stride_h};
    constexpr int stride_n = H_Size * stride_h;
    constexpr int size = N_Size * H_Size * W_Size * C_Size;
    constexpr int storage_size = (N_Size - 1) * stride_n + (H_Size - 1) * stride_h + (W_Size - 1) * stride_w + (C_Size - 1) + 1;
    constexpr size_t num_bytes = sizeof(float) * storage_size;

    EXPECT_EQ(StrideTensor::storage_size(), storage_size);
    EXPECT_EQ(StrideTensor::size(), size); 
    EXPECT_EQ(StrideTensor::num_bytes(), num_bytes);

    float data[N_Size * stride_n];
    for (int n = 0; n < N_Size; ++n) {{
        for (int h = 0; h < H_Size; ++h) {{
            for (int w = 0; w < W_Size; ++w) {{
                for (int c = 0; c < C_Size; ++c) {{
                    data[n*stride_n + h*stride_h + w*stride_w +c] = n*(H_Size*W_Size*C_Size) + h*(W_Size*C_Size) + w*C_Size + c;
                }}
            }}
        }}
    }}
    StrideTensor tensor(data);
    for (auto n : range(tensor.size<N>())) {{
        for (auto h : range(tensor.size<H>())) {{
            for (auto w : range(tensor.size<W>())) {{
                for (auto c : range(tensor.size<C>())) {{
                auto val = n.get()*(H_Size*W_Size*C_Size) + h.get()*(W_Size*C_Size) + w.get()*C_Size + c.get();
                    EXPECT_EQ(*tensor[n][h][w][c], val);
                    EXPECT_EQ(*tensor[h][n][c][w], val);
                    EXPECT_EQ(*tensor[h][w][c][n], val);
                    EXPECT_EQ(*tensor[w][c][n][h], val);
                }}
            }}
        }}
    }}
}}
"""
    return test_code


@_cpp_test
def _test_generate_fold_auto_size():
    """Return the C++ source code that tests automatic fold size inference."""
    # Test case 1: k8=16, k4=-1, k=-1
    # k4 = 8/4 = 2, k = 4/1 = 4
    specs1 = [
        gen.Tensor(gen.dtype.float, gen.Dims(k8=16, i=32, k4=-1, k=-1), class_name="A")
    ]
    generated_code1 = gen.generate(specs1, namespace="AutoFold_GenCode1")

    # Test case 2: k8=4, k=-1 (skip k4)
    # k = 8/1 = 8
    specs2 = [gen.Tensor(gen.dtype.float, gen.Dims(k8=4, j=16, k=-1), class_name="B")]
    generated_code2 = gen.generate(specs2, namespace="AutoFold_GenCode2")

    # Test case 3: Multiple base dimensions with auto-inference
    # k8=8, k4=-1, k=-1, i16=4, i=-1
    specs3 = [
        gen.Tensor(
            gen.dtype.float, gen.Dims(k8=8, k4=-1, k=-1, i16=4, i=-1), class_name="C"
        )
    ]
    generated_code3 = gen.generate(specs3, namespace="AutoFold_GenCode3")

    return f"""
{generated_code1}

UTEST(AutoFold, tensor_k8_k4_k)
{{
    using namespace AutoFold_GenCode1;

    // k8=16 (specified), k4=2 (8/4), k=4 (4/1)
    // Total K extent = 16 * 2 * 4 = 128
    EXPECT_EQ(A::size<K8>().get(), 16);
    EXPECT_EQ(A::size<K4>().get(), 2);
    EXPECT_EQ(A::size<K>().get(), 4);
    EXPECT_EQ(A::size<I>().get(), 32);

    // Verify total extent
    EXPECT_EQ(A::size(), 16 * 2 * 4 * 32);
}}

{generated_code2}

UTEST(AutoFold, tensor_k8_k)
{{
    using namespace AutoFold_GenCode2;

    // k8=4 (specified), k=8 (8/1)
    // Total K extent = 4 * 8 = 32
    EXPECT_EQ(B::size<K8>().get(), 4);
    EXPECT_EQ(B::size<K>().get(), 8);
    EXPECT_EQ(B::size<J>().get(), 16);

    // Verify total size
    EXPECT_EQ(B::size(), 4 * 8 * 16);
}}

{generated_code3}

UTEST(AutoFold, tensor_multiple_base_dims)
{{
    using namespace AutoFold_GenCode3;

    // k8=8, k4=2 (8/4), k=4 (4/1)
    // i16=4, i=16 (16/1)
    EXPECT_EQ(C::size<K8>().get(), 8);
    EXPECT_EQ(C::size<K4>().get(), 2);
    EXPECT_EQ(C::size<K>().get(), 4);
    EXPECT_EQ(C::size<I16>().get(), 4);
    EXPECT_EQ(C::size<I>().get(), 16);

    // Verify total size
    EXPECT_EQ(C::size(), 8 * 2 * 4 * 4 * 16);
}}
"""


@_cpp_test
def _test_generate_fold_explicit_validation():
    """Return the C++ source code that tests explicit size validation."""
    # Explicit sizes that match computed values should work
    specs = [gen.Tensor(gen.dtype.float, gen.Dims(k8=4, k4=2, k=4), class_name="D")]
    generated_code = gen.generate(specs, namespace="ExplicitFold_GenCode")

    return f"""
{generated_code}

UTEST(ExplicitFold, matching_explicit_sizes)
{{
    using namespace ExplicitFold_GenCode;

    // All sizes explicitly specified and consistent
    // k8=4, k4=2 (8/4=2 ✓), k=4 (4/1=4 ✓)
    EXPECT_EQ(D::size<K8>().get(), 4);
    EXPECT_EQ(D::size<K4>().get(), 2);
    EXPECT_EQ(D::size<K>().get(), 4);
}}
"""


@_cpp_test
def _test_with_vector_width():
    """Test the with_vector_width() method for Tensor.

    This tests that a tensor with dtype.float and dims (a, b, c, d=4) can be
    transformed to dtype.float4 with dims (a, b, c) by eliminating the stride-1
    dimension and adjusting strides.
    """
    # Original tensor: float with dims (warp=4, j4=16, i=32, j=4)
    # Stride for j4 is (32 + 1) * 4 = 132 (with padding for bank conflicts)
    original = gen.Tensor(
        gen.dtype.float,
        gen.Dims(warp=4, j4=16, i=32, j=4),
        strides=gen.Strides(j4=132),
    )

    # with_vector_width(4) should:
    # - Change dtype to float4 (veclen 4 vs 1, ratio = 4)
    # - Eliminate j dimension (size 4 / ratio 4 = 1)
    # - Divide all strides by 4: j4 stride becomes 132/4 = 33
    wide_tensor = original.vector_length(4, constant=True)

    generated_code = gen.generate(
        [
            gen.Tensor(
                original.data_type,
                original.dims,
                class_name="OriginalTensor",
                strides=original.strides,
            ),
            gen.Tensor(
                wide_tensor.data_type,
                wide_tensor.dims,
                class_name="WideTensor",
                strides=wide_tensor.strides,
                constant=True,
            ),
        ],
        namespace="WithVectorWidth_GenCode",
    )

    return f"""
{generated_code}

UTEST(WithVectorWidth, dimensions_and_strides)
{{
    using namespace WithVectorWidth_GenCode;

    // Original tensor has 4 dimensions
    EXPECT_EQ(OriginalTensor::size<WARP>().get(), 4);
    EXPECT_EQ(OriginalTensor::size<J4>().get(), 16);
    EXPECT_EQ(OriginalTensor::size<I>().get(), 32);
    EXPECT_EQ(OriginalTensor::size<J>().get(), 4);

    // Wide tensor has 3 dimensions (j eliminated since 4/4=1)
    EXPECT_EQ(WideTensor::size<WARP>().get(), 4);
    EXPECT_EQ(WideTensor::size<J4>().get(), 16);
    EXPECT_EQ(WideTensor::size<I>().get(), 32);
}}

UTEST(WithVectorWidth, storage_layout)
{{
    using namespace WithVectorWidth_GenCode;

    // Both tensors should access the same memory with different element types
    // Storage size should be the same in terms of bytes
    // Original: uses float (4 bytes), Wide: uses float4 (16 bytes)
    // Wide storage_size should be 1/4 of original (in elements)

    constexpr auto orig_storage = OriginalTensor::storage_size();
    constexpr auto wide_storage = WideTensor::storage_size();

    // storage_size is in elements, so wide should be 1/4 of original
    EXPECT_EQ(orig_storage, wide_storage * 4);
}}

UTEST(WithVectorWidth, memory_access)
{{
    using namespace WithVectorWidth_GenCode;

    // Verify that both tensors can index the same underlying data
    float orig_data[OriginalTensor::storage_size()];
    OriginalTensor orig_tensor(orig_data);

    // Cast to float4 for wide tensor (same memory, different view)
    const float4* wide_data = reinterpret_cast<const float4*>(orig_data);
    WideTensor wide_tensor(wide_data);

    // Access pattern: for each (warp, j4, i) in wide tensor,
    // this should correspond to 4 consecutive j values in original
    for (auto warp : range(WideTensor::size<WARP>())) {{
        for (auto j4 : range(WideTensor::size<J4>())) {{
            for (auto i : range(WideTensor::size<I>())) {{
                // Get wide tensor pointer (use .get() to get pointer from Cursor)
                const float4* wide_ptr = wide_tensor[warp][j4][i].get();

                // Get original tensor pointer at j=0
                float* orig_ptr = orig_tensor[warp][j4][i][J(0)].get();

                // They should point to the same memory
                EXPECT_EQ(reinterpret_cast<const void*>(wide_ptr),
                          reinterpret_cast<const void*>(orig_ptr));
            }}
        }}
    }}
}}
"""


@_cpp_test
def _test_cursor_initializer_with_ancestors():
    """Test CursorInitializer factory functions with ancestor tensors.

    This tests that when a tensor is derived (via derive_dim or vector_length),
    the CursorInitializer generates factory overloads for all ancestor types.
    This is important for patterns like:
        auto a_load_smem = ALoadSmem(a_smem);
    where a_smem is an ASmem instance and ALoadSmem is derived from ASmem.
    """
    # Thread index for implicit dims
    thread_idx = gen.CompoundIndex(
        gen.Dims(warp=4, lane=32),
        init=gen.BuiltIn.THREAD_IDX_X,
        class_name="ThreadIdx",
    )

    # Base shared memory tensor (like ASmem)
    base_smem = gen.Tensor(
        gen.dtype.float,
        gen.Dims(k_chunk=2, i16=8, k8=4),
        class_name="BaseSmem",
    )

    # Fragment load index (simulates a derived dimension from a Fragment)
    fragment_index = gen.CompoundIndex(
        gen.Dims(i16=8, k8=4),
        class_name="FragmentIndex",
    )

    # Derived tensor with implicit dims (like ALoadSmem = ASmem.derive_dim(...).initializer(...))
    derived_smem = base_smem.with_dim(fragment_index)
    load_smem_init = derived_smem.initializer(thread_idx)
    load_smem_init._set_class_name("LoadSmem")

    # We need to set the class name on the derived tensor too
    derived_smem._set_class_name("DerivedSmem")

    specs = [
        thread_idx,
        base_smem,
        fragment_index,
        derived_smem,
        load_smem_init,
    ]

    generated_code = gen.generate(specs, namespace="AncestorInit_GenCode")

    return f"""
{generated_code}

UTEST(CursorInitializer, accepts_ancestor_tensor)
{{
    using namespace AncestorInit_GenCode;

    // Allocate shared memory for the base tensor
    float data[BaseSmem::storage_size()];
    for (int i = 0; i < BaseSmem::storage_size(); ++i)
        data[i] = static_cast<float>(i);

    // Create a BaseSmem cursor
    BaseSmem base_smem(data);

    // Create a LoadSmem cursor from the BaseSmem cursor (ancestor type)
    // This should work because LoadSmem::initializer() generates an overload
    // that accepts BaseSmem.
    auto load_cursor = LoadSmem(base_smem);

    // Verify we can access data through the derived cursor
    // The cursor should point to the same underlying data
    EXPECT_EQ(reinterpret_cast<const void*>(load_cursor.get()),
              reinterpret_cast<const void*>(base_smem.get()));
}}

UTEST(CursorInitializer, accepts_derived_tensor)
{{
    using namespace AncestorInit_GenCode;

    float data[BaseSmem::storage_size()];
    for (int i = 0; i < BaseSmem::storage_size(); ++i)
        data[i] = static_cast<float>(i);

    // Create a DerivedSmem cursor directly
    DerivedSmem derived_smem(data);

    // Create a LoadSmem cursor from the DerivedSmem cursor (direct parent)
    auto load_cursor = LoadSmem(derived_smem);

    EXPECT_EQ(reinterpret_cast<const void*>(load_cursor.get()),
              reinterpret_cast<const void*>(derived_smem.get()));
}}

UTEST(CursorInitializer, accepts_raw_pointer)
{{
    using namespace AncestorInit_GenCode;

    float data[BaseSmem::storage_size()];
    for (int i = 0; i < BaseSmem::storage_size(); ++i)
        data[i] = static_cast<float>(i);

    // Create a LoadSmem cursor from a raw pointer
    auto load_cursor = LoadSmem(data);

    EXPECT_EQ(reinterpret_cast<const void*>(load_cursor.get()),
              reinterpret_cast<const void*>(data));
}}
"""


@_cpp_test
def _test_cursor_initializer_with_vector_length():
    """Test CursorInitializer with vector_length() creating different data types.

    This tests the case where vector_length() changes the data type (e.g., half2 -> uint4),
    requiring reinterpret_cast in the generated factory function.
    This mirrors the CLoadSmem pattern in mma_checkerboard_16c.cu.
    """
    # Simulate the CSmem -> CLoadSmem pattern:
    # CSmem uses half2, CLoadSmem uses half8 (uint4) via vector_length(8)

    # Thread index for implicit dims
    thread_idx = gen.CompoundIndex(
        gen.Dims(warp=4, lane=32),
        init=gen.BuiltIn.THREAD_IDX_X,
        class_name="ThreadIdx",
    )

    # Base shared memory tensor with float (simulating half2)
    # Dims: (warp_i=4, j8=16, i=32, k=4) with padding stride
    # Using k instead of j for the stride-1 dimension to avoid fold size conflict
    base_smem = gen.Tensor(
        gen.dtype.float,
        gen.Dims(warp_i=4, j8=16, i=32, k=4),
        strides=gen.Strides(j8=132),  # (32 + 1) * 4 for padding
        class_name="CSmem",
    )

    # Use vector_length(4) to create wider tensor (simulating half2 -> uint4)
    # This changes dtype from float to float4 and eliminates the k dimension
    wide_smem = base_smem.vector_length(4, constant=True)
    wide_smem._set_class_name("CSmemWide")

    # Create initializer for the wide tensor
    load_smem_init = wide_smem.initializer(thread_idx)
    load_smem_init._set_class_name("CLoadSmem")

    specs = [
        thread_idx,
        base_smem,
        wide_smem,
        load_smem_init,
    ]

    generated_code = gen.generate(specs, namespace="VectorLengthInit_GenCode")

    return f"""
// Define a simple uint4-like struct for testing if not available
#ifndef __CUDA_ARCH__
#ifndef uint4_defined
#define uint4_defined
// float4 is usually available, but let's make sure
#endif
#endif

{generated_code}

UTEST(CursorInitializer, vector_length_accepts_ancestor)
{{
    using namespace VectorLengthInit_GenCode;

    // Allocate data for the base tensor
    float data[CSmem::storage_size()];
    for (int i = 0; i < CSmem::storage_size(); ++i)
        data[i] = static_cast<float>(i);

    // Create the base tensor
    CSmem c_smem(data);

    // Create the wide tensor from the base tensor via the factory
    // This should work because CLoadSmem generates an overload for CSmem
    // that does: return CSmemWide(reinterpret_cast<const float4*>(tensor.get()))[ThreadIdx()];
    auto load_cursor = CLoadSmem(c_smem);

    // Verify the pointer is correctly cast and offset
    // The base data pointer should match (before any subscript offset)
    const float4* expected_base = reinterpret_cast<const float4*>(c_smem.get());
    
    // The actual pointer may differ due to ThreadIdx() subscript, but the base
    // memory region should be the same
    auto load_ptr = reinterpret_cast<const char*>(load_cursor.get());
    auto base_ptr = reinterpret_cast<const char*>(data);
    auto end_ptr = base_ptr + sizeof(data);
    
    // Verify load_cursor points within the data range
    EXPECT_TRUE(load_ptr >= base_ptr && load_ptr < end_ptr);
}}

UTEST(CursorInitializer, vector_length_data_access)
{{
    using namespace VectorLengthInit_GenCode;

    // Allocate and initialize data
    float data[CSmem::storage_size()];
    for (int i = 0; i < CSmem::storage_size(); ++i)
        data[i] = static_cast<float>(i);

    // Create base tensor and subscript to a known position
    CSmem c_smem(data);

    // Access via base tensor at a specific location
    auto base_cursor = c_smem[WARP_I(0)][J8(0)][I(0)][K(0)];
    float* base_ptr = base_cursor.get();

    // Access via wide tensor at same location (k dimension eliminated)
    CSmemWide c_smem_wide(reinterpret_cast<const float4*>(data));
    auto wide_cursor = c_smem_wide[WARP_I(0)][J8(0)][I(0)];
    const float4* wide_ptr = wide_cursor.get();

    // They should point to the same memory
    EXPECT_EQ(reinterpret_cast<const void*>(wide_ptr),
              reinterpret_cast<const void*>(base_ptr));

    // Verify data values through wide pointer
    EXPECT_EQ(wide_ptr->x, data[0]);
    EXPECT_EQ(wide_ptr->y, data[1]);
    EXPECT_EQ(wide_ptr->z, data[2]);
    EXPECT_EQ(wide_ptr->w, data[3]);
}}

UTEST(CursorInitializer, vector_length_via_factory)
{{
    using namespace VectorLengthInit_GenCode;

    // This test verifies the full factory path with reinterpret_cast
    float data[CSmem::storage_size()];
    for (int i = 0; i < CSmem::storage_size(); ++i)
        data[i] = static_cast<float>(i);

    CSmem c_smem(data);

    // Use the factory to create the wide cursor from base tensor
    // The factory should handle the reinterpret_cast internally
    auto load_cursor = CLoadSmem(c_smem);

    // Verify it's a valid float4 pointer within the data range
    const float4* ptr = load_cursor.get();
    EXPECT_TRUE(ptr != nullptr);

    // The pointer should be aligned for float4 access
    EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % alignof(float4), 0);
}}
"""


@_cpp_test
def _test_cursor_inbounds():
    """Test the Cursor::inbounds() method.

    This tests that inbounds() correctly checks if cursor coordinates
    are within the tensor extents.
    """
    specs = [
        gen.Tensor(
            gen.dtype.float,
            gen.Dims(i=8, j=16),
            class_name="Tensor2D",
        ),
    ]

    generated_code = gen.generate(specs, namespace="CursorInbounds_GenCode")

    return f"""
{generated_code}

UTEST(CursorInbounds, basic)
{{
    using namespace CursorInbounds_GenCode;

    float data[Tensor2D::storage_size()];
    Tensor2D tensor(data);

    // Cursor at origin should be inbounds
    auto cursor = tensor[I(0)][J(0)];
    EXPECT_TRUE(cursor.inbounds());

    // Cursor at last valid position should be inbounds
    auto last = tensor[I(7)][J(15)];
    EXPECT_TRUE(last.inbounds());

    // Cursor at i=8 should be out of bounds
    auto oob_i = tensor[I(8)][J(0)];
    EXPECT_FALSE(oob_i.inbounds());

    // Cursor at j=16 should be out of bounds
    auto oob_j = tensor[I(0)][J(16)];
    EXPECT_FALSE(oob_j.inbounds());

    // Cursor at both out of bounds
    auto oob_both = tensor[I(8)][J(16)];
    EXPECT_FALSE(oob_both.inbounds());
}}

UTEST(CursorInbounds, extents)
{{
    using namespace CursorInbounds_GenCode;

    float data[Tensor2D::storage_size()];
    Tensor2D tensor(data);

    // Verify extents() returns the correct values
    auto ext = tensor[I(0)][J(0)].extents();
    auto i_ext = ext.template get<I>();
    auto j_ext = ext.template get<J>();
    EXPECT_EQ(i_ext.get(), 8);
    EXPECT_EQ(j_ext.get(), 16);
}}

UTEST(CursorInbounds, extent_single_dim)
{{
    using namespace CursorInbounds_GenCode;

    // Test extent<DimType>() static method
    using CursorType = decltype(Tensor2D(nullptr)[I(0)]);
    auto i_extent = CursorType::extent<I>();
    auto j_extent = CursorType::extent<J>();
    EXPECT_EQ(i_extent.get(), 8);
    EXPECT_EQ(j_extent.get(), 16);
}}

UTEST(CursorInbounds, partial_subscript)
{{
    using namespace CursorInbounds_GenCode;

    float data[Tensor2D::storage_size()];
    Tensor2D tensor(data);

    // Cursor with only one dimension subscripted
    // Should be inbounds if that dimension is valid
    auto partial = tensor[I(5)];
    EXPECT_TRUE(partial.inbounds());

    // Cursor with one dimension out of bounds
    auto partial_oob = tensor[I(10)];
    EXPECT_FALSE(partial_oob.inbounds());
}}
"""


@_cpp_test
def _test_compound_index_derived_dims():
    """Test the derived dimensions interface for CompoundIndex.

    CompoundIndex should work as a derived dimension type, mapping
    OFFSET to multi-dimensional coordinates.
    """
    specs = [
        gen.CompoundIndex(gen.Dims(i=8, j=16), class_name="Index2D"),
        gen.Tensor(
            gen.dtype.float,
            gen.Dims(i=8, j=16),
            class_name="Tensor2D",
        ),
    ]

    generated_code = gen.generate(specs, namespace="CompoundIndexDerivedDims_GenCode")

    return f"""
{generated_code}

UTEST(CompoundIndexDerivedDims, input_output_dims)
{{
    using namespace CompoundIndexDerivedDims_GenCode;

    // Verify input_dims contains OFFSET with correct size
    using input = Index2D::input_dims;
    static_assert(spio::detail::tuple_size<input>::value == 1, "Should have 1 input dim");

    // Verify output_dims contains both I and J
    using output = Index2D::output_dims;
    static_assert(spio::detail::tuple_size<output>::value == 2, "Should have 2 output dims");
}}

UTEST(CompoundIndexDerivedDims, compute_coordinates)
{{
    using namespace CompoundIndexDerivedDims_GenCode;
    using spio::Coordinates;
    using spio::OFFSET;

    // Test that compute_coordinates returns the same result as coordinates()
    for (int offset = 0; offset < Index2D::size(); offset++) {{
        Index2D idx(offset);

        // Get expected values directly from CompoundIndex
        auto expected_i = idx.template get<I>().get();
        auto expected_j = idx.template get<J>().get();

        // Compute via derived dims interface
        auto computed = Index2D::compute_coordinates(Coordinates<OFFSET>(OFFSET(offset)));

        // Use idx.get() return type to access computed coordinates
        using I_Coord = decltype(idx.template get<I>());
        using J_Coord = decltype(idx.template get<J>());
        auto computed_i = computed.template get<I_Coord>().get();
        auto computed_j = computed.template get<J_Coord>().get();

        EXPECT_EQ(computed_i, expected_i);
        EXPECT_EQ(computed_j, expected_j);
    }}
}}

UTEST(CompoundIndexDerivedDims, tensor_subscript_with_offset)
{{
    using namespace CompoundIndexDerivedDims_GenCode;

    // Create a tensor with the same dimensions as the compound index
    float data[Tensor2D::storage_size()];
    Tensor2D tensor(data);

    // Initialize data with linear values
    for (int i = 0; i < Tensor2D::storage_size(); i++) {{
        data[i] = static_cast<float>(i);
    }}

    // Verify that subscripting with OFFSET uses derived dimensions
    // to expand to multi-dimensional coordinates
    for (int offset = 0; offset < Index2D::size(); offset++) {{
        Index2D idx(offset);

        // Direct subscript with i and j
        float* direct = tensor[idx.template get<I>()][idx.template get<J>()].get();

        // Using coordinates
        float* via_coords = tensor[idx.coordinates()].get();

        EXPECT_EQ(direct, via_coords);
    }}
}}
"""


@_cpp_test
def _test_fragment_index():

    specs = [
        gen.FragmentIndex(gen.FragmentType.M16_K16_F16_A, "r", "s", class_name="A"),
        gen.FragmentIndex(gen.FragmentType.N16_K16_F16_B, "s", "t", class_name="B"),
        gen.FragmentIndex(gen.FragmentType.M16_N16_F32_C, "r", "s", class_name="C"),
        gen.FragmentIndex(gen.FragmentType.N8_K8_F16_B, "s", "t", class_name="B2"),
    ]
    generated_code = gen.generate(specs, namespace="FragmentIndex_GenCode")
    test_code = f"""
{generated_code}

UTEST(FragmentIndex, MMA_M16_K16_F16_A)
{{
    using namespace FragmentIndex_GenCode;
    using namespace spio;
    constexpr int num_fragments = 4;
    for (int lane = 0; lane < 32; ++lane) {{
        auto expect = MMA_A_88_F16_Index<R, S>(lane);
        auto a = A(lane);
        for (int idx = 0; idx < num_fragments; ++idx) {{
            EXPECT_TRUE(a.get<R>(idx).get() == expect.get<R>(idx).get());
            EXPECT_TRUE((a.get<Fold<S, 2>>(idx).get()) == (expect.get<Fold<S, 2>>(idx).get()));
            EXPECT_TRUE((a.get<Fold<S, 8>>(idx).get()) == (expect.get<Fold<S, 8>>(idx).get()));
            EXPECT_TRUE((a.get<Module<S, 4, 2>>(idx).get()) == (expect.get<Module<S, 4, 2>>(idx).get()));
        }}
    }}
}}

UTEST(FragmentIndex, MMA_B_N16_K16_F16_Index)
{{
    using namespace FragmentIndex_GenCode;
    using namespace spio;
    constexpr int num_fragments = 4;
    for (int lane = 0; lane <32; ++lane) {{
        auto expect = MMA_B_88_F16_Index<S, T>(lane);
        auto b = B(lane);
        for (int idx = 0; idx < num_fragments; ++idx) {{
            EXPECT_TRUE((b.get<T>(idx).get()) == (expect.get<T>(idx).get()));
            EXPECT_TRUE((b.get<Fold<S, 2>>(idx).get()) == (expect.get<Fold<S, 2>>(idx).get()));
            EXPECT_TRUE((b.get<Fold<S, 8>>(idx).get()) == (expect.get<Fold<S, 8>>(idx).get()));
            EXPECT_TRUE((b.get<Module<S, 4, 2>>(idx).get()) == (expect.get<Module<S, 4, 2>>(idx).get()));
        }}
    }}
}}

UTEST(FragmentIndex, MMA_N8_K8_F16_B)
{{
    using namespace FragmentIndex_GenCode;
    using namespace spio;
    constexpr int num_fragments = 1;
    for (int lane = 0; lane <32; ++lane) {{
        auto expect = MMA_B_88_F16_Index<S, T>(lane);
        auto b = B2(lane);
        for (int idx = 0; idx < num_fragments; ++idx) {{
            EXPECT_TRUE((b.get<T>(idx).get()) == (expect.get<T>(idx).get()));
            EXPECT_TRUE((b.get<Fold<S, 2>>(idx).get()) == (expect.get<Fold<S, 2>>(idx).get()));
            EXPECT_TRUE((b.get<Fold<S, 8>>(idx).get()) == (expect.get<Fold<S, 8>>(idx).get()));
            EXPECT_TRUE((b.get<Module<S, 4, 2>>(idx).get()) == (expect.get<Module<S, 4, 2>>(idx).get()));
        }}
    }}
}}


UTEST(FragmentIndex, MMA_M16_N16_F32_C)
{{
    using namespace FragmentIndex_GenCode;
    using namespace spio;
    constexpr int num_fragments = 4;
    for (int lane = 0; lane <32; ++lane) {{
        auto expect = MMA_C_88_F32_Index<R, S>(lane);
        auto c = C(lane);
        for (int idx = 0; idx < num_fragments; ++idx) {{
            EXPECT_TRUE((c.get<R>(idx).get()) == (expect.get<R>(idx).get()));
            EXPECT_TRUE((c.get<Fold<S, 2>>(idx).get()) == (expect.get<Fold<S, 2>>(idx).get()));
            EXPECT_TRUE((c.get<Fold<S, 8>>(idx).get()) == (expect.get<Fold<S, 8>>(idx).get()));
            EXPECT_TRUE((c.get<Module<S, 4, 2>>(idx).get()) == (expect.get<Module<S, 4, 2>>(idx).get()));
        }}
    }}
}}
"""
    return test_code


@_cpp_test
def _test_fragment_load_index():
    specs = [
        gen.Dim("lane"),
        gen.FragmentLoadIndex(
            gen.FragmentType.M16_K16_F16_A, "x", "c", class_name="Input"
        ),
        gen.FragmentLoadIndex(
            gen.FragmentType.N8_K16_F16_B, "c", "k", class_name="Weights"
        ),
    ]
    generated_code = gen.generate(specs, namespace="FragmentLoadIndex_GenCode")
    test_code = f"""
{generated_code}

UTEST(FragmentLoadIndex, MMA_M16_K16_F16_A)
{{
    using namespace FragmentLoadIndex_GenCode;
    using namespace spio;
    for (int lane = 0; lane < 32; ++lane) {{
        auto input = Input(lane);
        auto expect = MMA_A_M16_K16_F16_LoadIndex<X, C>(lane);
        EXPECT_TRUE(input.get<X>() == expect.get<X>());
        EXPECT_TRUE((input.get<Fold<C, 8>>() == expect.get<Fold<C, 8>>()));
    }}

    for (int lane = 0; lane < 32; ++lane) {{
        auto expect = MMA_B_N8_K16_F16_LoadIndex<C, K>(lane);
        auto weights = Weights(lane);
        EXPECT_TRUE((weights.get<Fold<C, 8>>() == expect.get<Fold<C, 8>>()));
        EXPECT_TRUE(weights.get<K>() == expect.get<K>());
     }}
}}
"""
    return test_code


@_cpp_test
def _test_tensor_with_checkerboard_derived_dim():
    """Test Tensor generator with a Checkerboard derived dimension."""
    # Create a Checkerboard derived dimension similar to the C++ test
    # CheckerboardIndex<8, I, K8, CHECKERS, 16>
    g = gen.Generators()
    g.Swizzle = gen.Checkerboard(
        pairs_dim="i",
        colors_dim="k8",
        offset_dim="swizzle",
        size=16,
        ranks=8,
    )

    # Create a Tensor that uses the Checkerboard as a derived dimension
    # The dimension name "checkers" maps to the Checkerboard's output dimension
    g.MyTensor = gen.Tensor(
        gen.dtype.float,
        gen.Dims(swizzle=g.Swizzle),
    )

    generated_code = gen.generate(g, namespace="TensorCheckerboard_GenCode")

    test_code = f"""
{generated_code}

UTEST(TensorCheckerboard, derived_dim_subscript)
{{
    using namespace TensorCheckerboard_GenCode;
    using namespace spio;

    constexpr int Size = 16;
    float data[Size];
    for (int i = 0; i < Size; ++i)
        data[i] = static_cast<float>(i);

    MyTensor tensor(data);

    // Test subscript with input dimensions (pairs_dim=I, colors_dim=K8)
    using K8 = Fold<K, 8>;
    for (int idx = 0; idx < Size; ++idx) {{
        auto i_val = idx / 2;
        auto k8_val = idx % 2;
        auto coords = make_coordinates(I(i_val), K8(k8_val));

        // Compute expected offset using CheckerboardIndex logic
        Swizzle expected_idx{{I(i_val), K8(k8_val)}};
        int expected_offset = expected_idx.offset().get();
        EXPECT_EQ(*tensor[coords], data[expected_offset]);
    }}
}}
"""
    return test_code


@_cpp_test
def _test_tensor_with_derived_dim_and_compound_index():
    """Test Tensor with derived dimension subscripted by CompoundIndex."""
    g = gen.Generators()
    g.Swizzle = gen.Checkerboard(
        pairs_dim="i",
        colors_dim="k8",
        offset_dim="swizzle",
        size=16,
        ranks=8,
    )

    # Tensor with derived dimension
    g.MyTensor = gen.Tensor(
        gen.dtype.float,
        gen.Dims(swizzle=g.Swizzle),
    )

    # CompoundIndex with the input dimensions of the derived dimension
    g.MyIndex = gen.CompoundIndex(gen.Dims(i=8, k8=2))

    generated_code = gen.generate(g, namespace="TensorDerivedCompound_GenCode")

    test_code = f"""
{generated_code}

UTEST(TensorDerivedDim, compound_index_subscript)
{{
    using namespace TensorDerivedCompound_GenCode;
    using namespace spio;

    constexpr int Size = 16;
    float data[Size];
    for (int i = 0; i < Size; ++i)
        data[i] = static_cast<float>(i);

    MyTensor tensor(data);

    // Test subscript with CompoundIndex
    for (int offset = 0; offset < Size; ++offset) {{
        MyIndex idx(offset);

        // CompoundIndex has coordinates() method, so tensor[idx] should work
        // The tensor should project the coordinates onto the derived dimension
        using K8 = Fold<K, 8>;
        auto i_val = idx.get<I>().get();
        auto k8_val = idx.get<K8>().get();

        // Compute expected offset using CheckerboardIndex logic
        Swizzle expected_idx{{I(i_val), K8(k8_val)}};
        int expected_offset = expected_idx.offset().get();
        EXPECT_EQ(*tensor[idx], data[expected_offset]);
    }}
}}
"""
    return test_code


@_cpp_test
def _test_tensor_with_derived_dim_and_fragment_load_index():
    """Test Tensor with derived dimension subscripted by fragment load index.

    This tests whether the fragment load index (MMA_A_M16_K16_F16_LoadIndex)
    can be used to subscript a tensor with a derived dimension.
    The fragment load index has get<Dim>() methods but may not have coordinates().
    """
    g = gen.Generators()
    # MMA_A_M16_K16_F16_LoadIndex produces I=0-15 for lanes 0-31, and K8=0-1
    # size=32 accommodates 16 pairs * 2 colors = 32 elements
    # ranks=8 means 8 elements per row (for 128-byte shared memory lines)
    g.Swizzle = gen.Checkerboard(
        pairs_dim="i",
        colors_dim="k8",
        offset_dim="swizzle",
        size=32,
        ranks=8,
    )

    # Tensor with derived dimension
    g.MyTensor = gen.Tensor(
        gen.dtype.float,
        gen.Dims(swizzle=g.Swizzle),
    )

    generated_code = gen.generate(g, namespace="TensorDerivedFragment_GenCode")

    test_code = f"""
{generated_code}

UTEST(TensorDerivedDim, fragment_load_index_subscript)
{{
    using namespace TensorDerivedFragment_GenCode;
    using namespace spio;

    constexpr int Size = 32;
    float data[Size];
    for (int i = 0; i < Size; ++i)
        data[i] = static_cast<float>(i);

    MyTensor tensor(data);

    // Test subscript with fragment load index
    // MMA_A_M16_K16_F16_LoadIndex maps lane -> (I, K8) coordinates
    // For lanes 0-15: I = lane, K8 = 0
    // For lanes 16-31: I = lane & 15, K8 = 1
    using LoadIndex = MMA_A_M16_K16_F16_LoadIndex<I, K>;
    using K8 = Fold<K, 8>;

    // Test for all 32 lanes
    for (int lane = 0; lane < 32; ++lane) {{
        LoadIndex load_idx(lane);

        // Get the I and K8 values from the load index
        auto i_val = load_idx.get<I>().get();
        auto k8_val = load_idx.get<K8>().get();

        // Compute expected offset using CheckerboardIndex logic
        Swizzle expected_idx{{I(i_val), K8(k8_val)}};
        int expected_offset = expected_idx.offset().get();

        // This is the key test: can we subscript with the load index?
        // The tensor needs to either:
        // 1. Use load_idx.coordinates() if it exists, or
        // 2. Use the derived dimension's ability to extract from load_idx via get<>()
        EXPECT_EQ(*tensor[load_idx], data[expected_offset]);
    }}
}}
"""
    return test_code


@_cpp_test
def _test_async_strip_loader_2d_iteration():
    """Test AsyncStripLoader2D iteration pattern.

    This test verifies that the 2D strip loader correctly iterates over
    two dimensions when loading data. We test the pattern for 4 warps (128 threads)
    loading an 8x1 tile (same as 8 i16 elements x 1 k16 element).

    We don't include the actual async_strip_loader.cuh since it requires CUDA
    device code. Instead, we test the iteration logic directly.
    """
    return """

// Test helper: track which positions were visited and compute offsets
// This mirrors the AsyncStripLoader2D iteration pattern without CUDA dependencies
template <int smem_stride_inner, int global_stride_inner, int num_inner,
          int smem_stride_outer, int global_stride_outer, int num_outer>
struct TestLoader2DIteration {
    static constexpr int total_loads = num_outer * num_inner;

    int smem_offsets[total_loads];
    int global_offsets[total_loads];

    // Simulate the iteration pattern from a given base offset
    void simulate(int smem_base, int global_base) {
        int idx = 0;
        for (int j = 0; j < num_outer; ++j) {
            for (int i = 0; i < num_inner; ++i) {
                smem_offsets[idx] = smem_base + j * smem_stride_outer + i * smem_stride_inner;
                global_offsets[idx] = global_base + j * global_stride_outer + i * global_stride_inner;
                idx++;
            }
        }
    }
};

UTEST(AsyncStripLoader2D, four_warps_8x1_tile_iteration)
{
    // 4 warps (128 threads) loading 8 i16 elements x 1 k16 element
    // Each warp loads 2 elements along i16 (8 / 4 = 2)
    // With 4 warps, warp 0 loads positions 0,4; warp 1 loads 1,5; etc.
    //
    // Layout (showing which warp loads which position):
    //     i16:  0   1   2   3   4   5   6   7
    // k16=0:  [w0  w1  w2  w3  w0  w1  w2  w3]
    //          ^---first load--^---second load-^
    //
    // So for warp w:
    //   - First load: position w (offset = w * smem_stride_element)
    //   - Second load: position w+4 (offset = (w+4) * smem_stride_element)
    //
    // Template params for AsyncStripLoader2D:
    //   smem_stride_inner = 4 * element_stride (skip 4 warps)
    //   num_inner = 2 (each warp loads 2 elements)
    //   smem_stride_outer = k16 stride
    //   num_outer = 1

    constexpr int num_warps = 4;
    constexpr int i16_size = 8;
    constexpr int k16_size = 1;

    // Smem layout: k16 x i16 x 2 (checkerboard k8)
    // Stride for i16 = 2 (the k8 checkerboard pair)
    constexpr int smem_i16_stride = 2;
    constexpr int smem_k16_stride = i16_size * smem_i16_stride;  // 8 * 2 = 16

    // Global layout: k16 x i x k8
    // Each element is half8, so stride for i = 1
    constexpr int global_i_stride = 1;
    constexpr int global_k16_stride = 8192;  // Some large value for outer dim

    // Each warp needs to load i16_size / num_warps = 2 elements
    constexpr int num_inner = i16_size / num_warps;  // 2
    constexpr int num_outer = k16_size;  // 1

    // Inner stride: skip num_warps positions in smem/global
    constexpr int smem_stride_inner = num_warps * smem_i16_stride;  // 4 * 2 = 8
    constexpr int global_stride_inner = num_warps * global_i_stride;  // 4 * 1 = 4

    // Verify the parameters
    EXPECT_EQ(num_inner, 2);
    EXPECT_EQ(num_outer, 1);
    EXPECT_EQ(smem_stride_inner, 8);
    EXPECT_EQ(global_stride_inner, 4);

    // Test loader for warp 0
    TestLoader2DIteration<smem_stride_inner, global_stride_inner, num_inner,
                          smem_k16_stride, global_k16_stride, num_outer> loader0;

    // Warp 0 starts at smem offset 0, global offset 0
    int warp0_smem_base = 0 * smem_i16_stride;  // 0
    int warp0_global_base = 0 * global_i_stride;  // 0
    loader0.simulate(warp0_smem_base, warp0_global_base);

    // Warp 0 should load i16 positions 0 and 4
    // Position 0: smem offset = 0
    // Position 4: smem offset = 0 + 8 = 8 (skip 4 warps * stride 2)
    EXPECT_EQ(loader0.smem_offsets[0], 0);      // i16=0
    EXPECT_EQ(loader0.smem_offsets[1], 8);      // i16=4
    EXPECT_EQ(loader0.global_offsets[0], 0);
    EXPECT_EQ(loader0.global_offsets[1], 4);

    // Test warp 1
    TestLoader2DIteration<smem_stride_inner, global_stride_inner, num_inner,
                          smem_k16_stride, global_k16_stride, num_outer> loader1;
    int warp1_smem_base = 1 * smem_i16_stride;  // 2
    int warp1_global_base = 1 * global_i_stride;  // 1
    loader1.simulate(warp1_smem_base, warp1_global_base);

    // Warp 1 should load i16 positions 1 and 5
    EXPECT_EQ(loader1.smem_offsets[0], 2);     // i16=1
    EXPECT_EQ(loader1.smem_offsets[1], 10);    // i16=5 (2 + 8)
    EXPECT_EQ(loader1.global_offsets[0], 1);
    EXPECT_EQ(loader1.global_offsets[1], 5);

    // Verify all 8 positions are covered by all 4 warps
    // Each position should be loaded exactly once
    int smem_positions_loaded[8] = {0};
    for (int warp = 0; warp < 4; ++warp) {
        TestLoader2DIteration<smem_stride_inner, global_stride_inner, num_inner,
                              smem_k16_stride, global_k16_stride, num_outer> loader;
        loader.simulate(warp * smem_i16_stride, warp * global_i_stride);
        for (int load = 0; load < num_inner; ++load) {
            // smem_offset / smem_i16_stride gives the i16 position
            int pos = loader.smem_offsets[load] / smem_i16_stride;
            EXPECT_TRUE(pos >= 0 && pos < 8);
            smem_positions_loaded[pos]++;
        }
    }
    // Each position should be loaded exactly once
    for (int pos = 0; pos < 8; ++pos) {
        EXPECT_EQ(smem_positions_loaded[pos], 1);
    }
}

UTEST(AsyncStripLoader2D, four_warps_8x2_tile_iteration)
{
    // 4 warps (128 threads) loading 8 i16 elements x 2 k16 elements
    // Each warp loads 2 i16 positions and iterates over both k16 values
    //
    // num_inner = 8/4 = 2, num_outer = 2
    // Total: 4 loads per warp

    constexpr int num_warps = 4;
    constexpr int i16_size = 8;
    constexpr int k16_size = 2;

    constexpr int smem_i16_stride = 2;
    constexpr int smem_k16_stride = i16_size * smem_i16_stride;  // 16

    constexpr int global_i_stride = 1;
    constexpr int global_k16_stride = 8192;

    constexpr int num_inner = i16_size / num_warps;  // 2
    constexpr int num_outer = k16_size;  // 2

    constexpr int smem_stride_inner = num_warps * smem_i16_stride;  // 8
    constexpr int global_stride_inner = num_warps * global_i_stride;  // 4

    EXPECT_EQ(num_inner, 2);
    EXPECT_EQ(num_outer, 2);

    TestLoader2DIteration<smem_stride_inner, global_stride_inner, num_inner,
                          smem_k16_stride, global_k16_stride, num_outer> loader;

    // Warp 0 base offsets
    loader.simulate(0, 0);

    // 4 loads total: outer loop j=0,1; inner loop i=0,1
    // j=0, i=0: smem = 0, global = 0
    // j=0, i=1: smem = 0 + 8 = 8, global = 0 + 4 = 4
    // j=1, i=0: smem = 0 + 16 = 16, global = 0 + 8192 = 8192
    // j=1, i=1: smem = 0 + 16 + 8 = 24, global = 0 + 8192 + 4 = 8196
    EXPECT_EQ(loader.smem_offsets[0], 0);       // k16=0, i16=0
    EXPECT_EQ(loader.smem_offsets[1], 8);       // k16=0, i16=4
    EXPECT_EQ(loader.smem_offsets[2], 16);      // k16=1, i16=0
    EXPECT_EQ(loader.smem_offsets[3], 24);      // k16=1, i16=4

    EXPECT_EQ(loader.global_offsets[0], 0);
    EXPECT_EQ(loader.global_offsets[1], 4);
    EXPECT_EQ(loader.global_offsets[2], 8192);
    EXPECT_EQ(loader.global_offsets[3], 8196);
}

UTEST(AsyncStripLoader2D, eight_warps_8x1_tile_iteration)
{
    // 8 warps loading 8 i16 x 1 k16 = perfect 1:1 mapping
    // This is the baseline case that works with AsyncStripLoader
    // num_inner = 1, num_outer = 1

    constexpr int num_warps = 8;
    constexpr int i16_size = 8;
    constexpr int k16_size = 1;

    constexpr int num_inner = i16_size / num_warps;  // 1
    constexpr int num_outer = k16_size;  // 1

    EXPECT_EQ(num_inner, 1);
    EXPECT_EQ(num_outer, 1);
    EXPECT_EQ(num_inner * num_outer, 1);  // Single load per warp

    // With num_inner=1 and num_outer=1, there's only one load per warp
    // The strides don't matter since we never step
    constexpr int smem_stride_inner = 16;  // doesn't matter
    constexpr int global_stride_inner = 8;  // doesn't matter
    constexpr int smem_k16_stride = 32;
    constexpr int global_k16_stride = 8192;

    TestLoader2DIteration<smem_stride_inner, global_stride_inner, num_inner,
                          smem_k16_stride, global_k16_stride, num_outer> loader;

    // Warp 3 starts at its position
    int warp3_smem_base = 3 * 2;  // position 3 * stride 2
    int warp3_global_base = 3;
    loader.simulate(warp3_smem_base, warp3_global_base);

    // Only one load, at the base position
    EXPECT_EQ(loader.smem_offsets[0], 6);
    EXPECT_EQ(loader.global_offsets[0], 3);
}
"""


def _compile_cpp_tests(extra_cpp_test_files=None, run_args=None):
    """Compile C++ tests with NVCC."""
    if extra_cpp_test_files is None:
        extra_cpp_test_files = []
    includes = [
        importlib_resources_files("spio.include"),
        importlib_resources_files("spio.src_tests"),
    ]
    sources = [
        importlib_resources_files("spio.src_tests") / src for src in CPP_SOURCES
    ] + extra_cpp_test_files
    includes = [str(include) for include in includes]
    return spio.compiler.compile_with_nvcc(
        sources=sources, includes=includes, run=True, run_args=run_args
    )
