import numpy as np
import polars as pl
from numba import float64, guvectorize, int64, uint64  # using numba to create some example ufuncs

# Random positive floats
rs = np.random.RandomState(420)
N = 100
x = pl.Series(rs.uniform(100, 200, N))
y = pl.Series(rs.uniform(100, 200, N))


# 1. Inputs and outputs are all uint64; works as expected
@guvectorize([(uint64[:], uint64[:], uint64[:])], "(n),(n)->(n)")
def a(arr, arr2, result):
    for i in range(len(arr)):
        result[i] = np.uint(arr[i] + arr2[i])


a(x.cast(pl.UInt64), y.cast(pl.UInt64))
# => Works as expected, returns a Series of uint64


# 2. Inputs are int64, output is uint64; breaks if dtype is not explicitly specified
@guvectorize([(int64[:], int64[:], uint64[:])], "(n),(n)->(n)")
def b(arr, arr2, result):
    for i in range(len(arr)):
        result[i] = np.uint(arr[i] + arr2[i])


b(x.cast(pl.Int64), y.cast(pl.Int64))
# => TypeError: No loop matching the specified signature and casting was found for ufunc b

b(x.cast(pl.Int64), y.cast(pl.Int64), dtype="L")
# => Works as expected, returns a Series of uint64


# 3. Inputs are float64, so no casting necessary, but output is uint64; breaks like the previous case
@guvectorize([(float64[:], float64[:], uint64[:])], "(n),(n)->(n)")
def c(arr, arr2, result):
    for i in range(len(arr)):
        result[i] = np.uint(arr[i] + arr2[i])


c(x, y)
# => TypeError: No loop matching the specified signature and casting was found for ufunc b

c(x, y, dtype="L")
# => Works as expected, returns a Series of uint64


# 4. Maybe it works if at least _one_ input is uint64?
@guvectorize([(uint64[:], float64[:], uint64[:])], "(n),(n)->(n)")
def d(arr, arr2, result):
    for i in range(len(arr)):
        result[i] = np.uint(arr[i] + arr2[i])


d(x.cast(pl.UInt64), y)
# => Still no luck ...

d(x.cast(pl.UInt64), y, dtype="L")
# => Works as expected, returns a Series of uint64

TypeError: No loop matching the specified signature and casting was found for ufunc {function_name}