import numpy as np
import warp as wp


@wp.func
def func_1() -> int:
  return 1


@wp.func
def func_2() -> int:
  return 2


funcs = [func_1, func_2]


def builder(funcids):
  _funcs = [funcs[id] for id in funcids]

  @wp.kernel
  def eval(input: wp.array(dtype=int), output: wp.array(dtype=int)):
    for i in range(wp.static(len(_funcs))):
      if wp.static(funcids[i]) == input[i]:
        output[0] = wp.static(_funcs[i])()

  return eval


output1 = wp.array((1,), dtype=int)
wp.launch(
  builder([0]),
  dim=(1,),
  inputs=[wp.array(np.array([0]), dtype=int)],
  outputs=[output1],
)
print(f"output: {output1.numpy()[0]} (should be 1)")

output2 = wp.array((1,), dtype=int)
wp.launch(
  builder([1]),
  dim=(1,),
  inputs=[wp.array(np.array([1]), dtype=int)],
  outputs=[output2],
)
print(f"output: {output2.numpy()[0]} (should be 2)")