How to use the wgpu.utils.compute_with_buffers function in wgpu

with raises(ValueError):  # output_arrays format invalid
        compute_with_buffers({0: in1}, {0: "10xfoo"}, compute_shader)
    with raises(ValueError):  # output_arrays shape invalid
        compute_with_buffers({0: in1}, {0: ("i",)}, compute_shader)
    with raises(ValueError):  # output_arrays shape invalid
        compute_with_buffers({0: in1}, {0: (0, "i",)}, compute_shader)
    with raises(ValueError):  # output_arrays shape invalid
        compute_with_buffers({0: in1}, {0: (-1, "i",)}, compute_shader)

    with raises(TypeError):  # invalid n
        compute_with_buffers({0: in1}, {0: c_int32 * 100}, compute_shader, n="100")
    with raises(ValueError):  # invalid n
        compute_with_buffers({0: in1}, {0: c_int32 * 100}, compute_shader, n=-1)

    with raises(TypeError):  # invalid shader
        compute_with_buffers({0: in1}, {0: c_int32 * 100}, "not a shader")

def test_compute_0_1_int():
    @python2shader
    def compute_shader(
        index: ("input", "GlobalInvocationId", i32), out: ("buffer", 0, Array(i32)),
    ):
        out[index] = index

    out = compute_with_buffers({}, {0: 400}, compute_shader)
    assert isinstance(out, dict) and len(out) == 1
    assert isinstance(out[0], memoryview)
    assert out[0].cast("i").tolist() == list(range(100))

with raises(TypeError):  # input_arrays key not int
        compute_with_buffers({"0": in1}, {0: c_int32 * 100}, compute_shader)
    with raises(TypeError):  # input_arrays value not ctypes array
        compute_with_buffers({0: list(in1)}, {0: c_int32 * 100}, compute_shader)

    with raises(TypeError):  # output_arrays is not a dict
        compute_with_buffers({0: in1}, [c_int32 * 100], compute_shader)
    with raises(TypeError):  # output_arrays key not int
        compute_with_buffers({0: in1}, {"0": c_int32 * 100}, compute_shader)
    with raises(TypeError):  # output_arrays value not a ctypes Array type
        compute_with_buffers({0: in1}, {0: "foobar"}, compute_shader)

    with raises(ValueError):  # output_arrays format invalid
        compute_with_buffers({0: in1}, {0: "10xfoo"}, compute_shader)
    with raises(ValueError):  # output_arrays shape invalid
        compute_with_buffers({0: in1}, {0: ("i",)}, compute_shader)
    with raises(ValueError):  # output_arrays shape invalid
        compute_with_buffers({0: in1}, {0: (0, "i",)}, compute_shader)
    with raises(ValueError):  # output_arrays shape invalid
        compute_with_buffers({0: in1}, {0: (-1, "i",)}, compute_shader)

    with raises(TypeError):  # invalid n
        compute_with_buffers({0: in1}, {0: c_int32 * 100}, compute_shader, n="100")
    with raises(ValueError):  # invalid n
        compute_with_buffers({0: in1}, {0: c_int32 * 100}, compute_shader, n=-1)

    with raises(TypeError):  # invalid shader
        compute_with_buffers({0: in1}, {0: c_int32 * 100}, "not a shader")

def test_compute_0_1_str():
    @python2shader
    def compute_shader(
        index: ("input", "GlobalInvocationId", i32), out: ("buffer", 0, Array(i32)),
    ):
        out[index] = index

    out = compute_with_buffers({}, {0: "100xi"}, compute_shader)
    assert isinstance(out, dict) and len(out) == 1
    assert isinstance(out[0], memoryview)
    assert out[0].tolist() == list(range(100))

def test_compute_0_1_ctype():
    @python2shader
    def compute_shader(
        index: ("input", "GlobalInvocationId", i32), out: ("buffer", 0, Array(i32)),
    ):
        out[index] = index

    # Create some ints!
    out = compute_with_buffers({}, {0: c_int32 * 100}, compute_shader)
    assert isinstance(out, dict) and len(out) == 1
    assert isinstance(out[0], ctypes.Array)
    assert iters_equal(out[0], range(100))

    # Same, but specify in bytes
    out = compute_with_buffers({}, {0: c_ubyte * 80}, compute_shader, n=20)
    assert isinstance(out, dict) and len(out) == 1
    assert isinstance(out[0], ctypes.Array)
    out0 = (c_int32 * 20).from_buffer(out[0])  # cast (a view in np)
    assert iters_equal(out0, range(20))

with raises(TypeError):  # output_arrays key not int
        compute_with_buffers({0: in1}, {"0": c_int32 * 100}, compute_shader)
    with raises(TypeError):  # output_arrays value not a ctypes Array type
        compute_with_buffers({0: in1}, {0: "foobar"}, compute_shader)

    with raises(ValueError):  # output_arrays format invalid
        compute_with_buffers({0: in1}, {0: "10xfoo"}, compute_shader)
    with raises(ValueError):  # output_arrays shape invalid
        compute_with_buffers({0: in1}, {0: ("i",)}, compute_shader)
    with raises(ValueError):  # output_arrays shape invalid
        compute_with_buffers({0: in1}, {0: (0, "i",)}, compute_shader)
    with raises(ValueError):  # output_arrays shape invalid
        compute_with_buffers({0: in1}, {0: (-1, "i",)}, compute_shader)

    with raises(TypeError):  # invalid n
        compute_with_buffers({0: in1}, {0: c_int32 * 100}, compute_shader, n="100")
    with raises(ValueError):  # invalid n
        compute_with_buffers({0: in1}, {0: c_int32 * 100}, compute_shader, n=-1)

    with raises(TypeError):  # invalid shader
        compute_with_buffers({0: in1}, {0: c_int32 * 100}, "not a shader")

    @python2shader
    def compute_shader(
        index: ("input", "GlobalInvocationId", i32),
        in1: ("buffer", 0, Array(i32)),
        out1: ("buffer", 1, Array(i32)),
        out2: ("buffer", 2, Array(i32)),
    ):
        out1[index] = in1[index]
        out2[index] = index

    # Create an array of 100 random int32
    in1 = [int(random.uniform(0, 100)) for i in range(100)]
    in1 = (c_int32 * 100)(*in1)

    outspecs = {0: 100 * c_int32, 1: 100 * c_int32, 2: 100 * c_int32}
    out = compute_with_buffers({0: in1}, outspecs, compute_shader)
    assert isinstance(out, dict) and len(out) == 3
    assert isinstance(out[0], ctypes.Array)
    assert isinstance(out[1], ctypes.Array)
    assert isinstance(out[2], ctypes.Array)
    assert iters_equal(out[0], in1)  # because it's the same buffer
    assert iters_equal(out[1], in1)  # because the shader copied the data
    assert iters_equal(out[2], range(100))  # because this is the index

def test_compute_0_1_ctype():
    @python2shader
    def compute_shader(
        index: ("input", "GlobalInvocationId", i32), out: ("buffer", 0, Array(i32)),
    ):
        out[index] = index

    # Create some ints!
    out = compute_with_buffers({}, {0: c_int32 * 100}, compute_shader)
    assert isinstance(out, dict) and len(out) == 1
    assert isinstance(out[0], ctypes.Array)
    assert iters_equal(out[0], range(100))

    # Same, but specify in bytes
    out = compute_with_buffers({}, {0: c_ubyte * 80}, compute_shader, n=20)
    assert isinstance(out, dict) and len(out) == 1
    assert isinstance(out[0], ctypes.Array)
    out0 = (c_int32 * 20).from_buffer(out[0])  # cast (a view in np)
    assert iters_equal(out0, range(20))

def test_compute_fails():
    @python2shader
    def compute_shader(
        index: ("input", "GlobalInvocationId", i32),
        in1: ("buffer", 0, Array(i32)),
        out1: ("buffer", 1, Array(i32)),
    ):
        out1[index] = in1[index]

    in1 = [int(random.uniform(0, 100)) for i in range(100)]
    in1 = (c_int32 * 100)(*in1)

    # Baseline; this works
    out = compute_with_buffers(
        {0: in1}, {0: c_int32 * 100}, compute_shader, n=(100, 1, 1)
    )
    assert iters_equal(out[0], in1)

    with raises(TypeError):  # input_arrays is not a dict
        compute_with_buffers([in1], {0: c_int32 * 100}, compute_shader)
    with raises(TypeError):  # input_arrays key not int
        compute_with_buffers({"0": in1}, {0: c_int32 * 100}, compute_shader)
    with raises(TypeError):  # input_arrays value not ctypes array
        compute_with_buffers({0: list(in1)}, {0: c_int32 * 100}, compute_shader)

    with raises(TypeError):  # output_arrays is not a dict
        compute_with_buffers({0: in1}, [c_int32 * 100], compute_shader)
    with raises(TypeError):  # output_arrays key not int
        compute_with_buffers({0: in1}, {"0": c_int32 * 100}, compute_shader)
    with raises(TypeError):  # output_arrays value not a ctypes Array type

):
    data2[index] = data1[index]


# Create input data as a memoryview
n = 20
data = memoryview(bytearray(n * 4)).cast("i")
for i in range(n):
    data[i] = i


# %% The short version, using memoryview

# The first arg is the input data, per binding
# The second arg are the ouput types, per binding
out = compute_with_buffers({0: data}, {1: (n, "i")}, compute_shader, n=n)

# The result is a dict matching the output types
# Select data from buffer at binding 1
result = out[1]
print(result.tolist())


# %% The short version, using numpy

# import numpy as np
#
# numpy_data = np.frombuffer(data, np.int32)
# out = compute_with_buffers({0: numpy_data}, {1: numpy_data.nbytes}, compute_shader, n=n)
# result = np.frombuffer(out[1], dtype=np.int32)
# print(result)