"""Repro: ternary on wp.bool in a Warp struct causes CUDA error 700 in Warp 1.11.0.

    uv run python warp_bool_ternary_bug.py broken   # CUDA error 700
    uv run python warp_bool_ternary_bug.py fixed    # silent (OK)
"""
import sys
import warp as wp

wp.init()


@wp.struct
class S:
    f: wp.float32
    flag: wp.bool


mode = sys.argv[1] if len(sys.argv) > 1 else "broken"

if mode == "broken":
    @wp.kernel
    def kernel(arr: wp.array(ndim=2, dtype=S)) -> None:
        i, j = wp.tid()
        s = arr[i, j]
        s.f = s.f if s.flag else wp.float32(0.0)  # BUG
        arr[i, j] = s
else:
    @wp.kernel
    def kernel(arr: wp.array(ndim=2, dtype=S)) -> None:
        i, j = wp.tid()
        s = arr[i, j]
        value = wp.float32(0.0)
        if s.flag:
            value = s.f
        s.f = value  # FIX
        arr[i, j] = s


arr = wp.zeros((128, 2), dtype=S, device="cuda:0")
wp.launch(kernel, dim=arr.shape, inputs=[arr], device="cuda:0")
wp.synchronize_device("cuda:0")