import numpy as np
import cupy as cp
import awkward as ak
import timeit


# Create random jagged data
np.random.seed(42)
num_lists = 100_000
mean_len = 10
max_len = 40

# Generate random lengths and build offsets
lengths = np.random.poisson(lam=mean_len, size=num_lists).astype(np.int64)
lengths = np.clip(lengths, 0, max_len)
offsets = np.empty(num_lists + 1, dtype=np.int64)
offsets[0] = 0
np.cumsum(lengths, out=offsets[1:])

total_values = int(offsets[-1])
values = np.arange(total_values, dtype=np.int64)

# Build awkward array
layout = ak.contents.ListOffsetArray(
    ak.index.Index64(offsets),
    ak.contents.NumpyArray(values)
)
arr = ak.Array(layout)

# Estimate total pairs
total_pairs = sum(n * (n - 1) // 2 for n in lengths)
print(f"lists={num_lists:,} total_values={total_values:,} mean_len≈{mean_len}")
print(f"estimated total_pairs={total_pairs:,}")

# Move to CUDA
gpu_arr = ak.to_backend(arr, "cuda")
cp.cuda.get_current_stream().synchronize()

# Benchmark using timeit
# warmup:
ak.combinations(gpu_arr, n=2)
result = timeit.timeit(lambda: ak.combinations(gpu_arr, n=2), number=10)
print(f"Time taken for ak.combinations: {result / 10:.4f} seconds")