"""
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
import os
from typing import Callable

from importlib_resources import files as importlib_resources_files
import pytest

import spio.generators as gen
import spio.compiler


TRUTHS = ["true", "1", "yes", "y", "t"]

ENABLE_CPP_TESTS = os.environ.get("SPIO_ENABLE_CPP_TESTS", "false").lower() in TRUTHS

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 = []

CPP_TESTS_FILTER = None

# Select specific C++ unit tests with a filter like this:
# CPP_TESTS_FILTER = "Dim.*"


@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("block_i", "i", 64),
        gen.Fold("block_j", "j", 64),
        # Test a fold of a fold. Produces using BLOCK_C8 = Fold<Fold<C, 8>, 4>;
        gen.Fold("block_c8", "c8", 4),
    ]
    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_index():
    """Return the C++ source code that tests a custom index class."""

    specs = [gen.CompoundIndex("MyIndex", gen.Dims(n=4, h=32, w=64, c=128))]
    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_checkerboard_index():
    specs = [gen.Checkerboard("Checkers", "r", "c8", "offset", 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(
        "ContiguousTensor", gen.dtype.float, gen.Dims(n=n, h=h, w=w, c=c)
    )

    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("ContiguousTensor", gen.dtype.float, gen.Dims(n=n), 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(
            "StrideTensor",
            gen.dtype.float,
            gen.Dims(n=n, h=h, w=w, c=c),
            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_fragment_index():

    specs = [
        gen.FragmentIndex("A", gen.FragmentType.M16_K16_F16_A, "r", "s"),
        gen.FragmentIndex("B", gen.FragmentType.N16_K16_F16_B, "s", "t"),
        gen.FragmentIndex("C", gen.FragmentType.M16_N16_F32_C, "r", "s"),
        gen.FragmentIndex("B2", gen.FragmentType.N8_K8_F16_B, "s", "t"),
    ]
    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("Input", gen.FragmentType.M16_K16_F16_A, "x", "c"),
        gen.FragmentLoadIndex("Weights", gen.FragmentType.N8_K16_F16_B, "c", "k"),
    ]
    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


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
    )
