# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
import random
import numpy as np
import pytest

from slangpy.backend import DeviceType, float3, int3, uint3
from slangpy.core.module import Module
from slangpy.experimental.gridarg import grid
from slangpy.tests import helpers
from slangpy.types.buffer import NDBuffer


def load_test_module(device_type: DeviceType):
    device = helpers.get_device(device_type)
    return


def grid_test(device_type: DeviceType, dims: int = 2, datatype: str = 'array', stride: int = 1, offset: int = 0, fixed_shape: bool = True):
    # Generate random shape and the arguments for numpy transpose
    random.seed(42)
    shape = tuple([random.randint(5, 15) for _ in range(dims)])
    transpose = tuple([i+1 for i in range(dims)]) + (0,)
    gen_args = ""

    if datatype == 'vector':
        buffertypename = f"int{dims}"
        slangtypename = buffertypename
    elif datatype == 'array':
        buffertypename = f"int[{dims}]"
        slangtypename = buffertypename
    elif datatype == 'genvector':
        buffertypename = f"int{dims}"
        slangtypename = f"vector<int, N>"
        gen_args = "<let N: int>"
    else:
        raise ValueError(f"Unknown datatype: {datatype}")

    # Create function that just dumps input to output for correct sized int
    device = helpers.get_device(device_type)
    module = helpers.create_module(device, f"""
{slangtypename} get{gen_args}({slangtypename} input) {{
    return input;
}}
"""
    )

    # Buffer for vector results
    res = NDBuffer(device, shape=shape, dtype=module.layout.find_type_by_name(buffertypename))

    # Offset per dimension
    offsets = tuple([offset for s in shape])

    # Call function with grid as input argument
    if fixed_shape:
        if stride == 1:
            module.get(grid(shape, offset=offsets), _result=res)
        else:
            strides = tuple([stride for s in shape])
            module.get(grid(shape, stride=strides, offset=offsets), _result=res)
    else:
        if stride == 1:
            module.get(grid(len(shape), offset=offsets), _result=res)
        else:
            strides = tuple([stride for s in shape])
            module.get(grid(len(shape), stride=strides, offset=offsets), _result=res)

    # Should get random numbers
    resdata = res.to_numpy().view(np.int32).reshape(shape + (dims,))
    expected = np.indices(shape).transpose(*transpose) * stride + offset

    if datatype == 'vector':
        expected = np.flip(expected, axis=-1)

    assert np.all(resdata == expected)


@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
@pytest.mark.parametrize("dims", [1, 2, 3, 4])
@pytest.mark.parametrize("stride", [1, 3])
def test_grid_vectors(device_type: DeviceType, dims: int, stride: int):
    if dims > 4:
        pytest.skip("Vector types only supported up to 4 dimensions")
    grid_test(device_type, dims=dims, datatype='vector', stride=stride)


@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
def test_grid_generic_vectors(device_type: DeviceType):
    pytest.skip("Doesn't currently work due to handling of generic arguments in specialize")
    grid_test(device_type, dims=3, datatype='genvector')


@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
@pytest.mark.parametrize("dims", [1, 3, 6])
def test_grid_arrays(device_type: DeviceType, dims: int):
    grid_test(device_type, dims=dims, datatype='array')


@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
@pytest.mark.parametrize("offsets", [-100, 50])
def test_grid_offsets(device_type: DeviceType, offsets: int):
    grid_test(device_type, offset=offsets)


@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
@pytest.mark.parametrize("stride", [1, 3])
def test_grid_implicit_shape(device_type: DeviceType, stride: int):
    grid_test(device_type, fixed_shape=False, stride=stride)


@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
@pytest.mark.parametrize("stride", [1, 3])
def test_grid_error_pass_to_bad_vector_size(device_type: DeviceType, stride: int):
    device = helpers.get_device(device_type)
    module = helpers.create_module(device, f"""
int3 get(int3 input) {{
    return input;
}}
"""
                                   )
    with pytest.raises(ValueError, match="Could not find suitable"):
        module.get(grid(shape=(2, 2)), _result='numpy')


@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
@pytest.mark.parametrize("stride", [1, 3])
def test_grid_error_pass_to_bad_array_size(device_type: DeviceType, stride: int):
    device = helpers.get_device(device_type)
    module = helpers.create_module(device, f"""
int3 get(int[3] input) {{
    return 0;
}}
"""
                                   )
    with pytest.raises(ValueError, match="Could not find suitable"):
        module.get(grid(shape=(2, 2)), _result='numpy')


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