#!/usr/bin/env python3
"""Minimal Quadrants CUDA kernel smoke run."""

from __future__ import annotations

import argparse

import numpy as np
import quadrants as qd


@qd.kernel
def axpy(n: int, x: qd.types.NDArray[qd.f32, 1], y: qd.types.NDArray[qd.f32, 1]) -> None:
    for i in range(n):
        y[i] = 2.0 * x[i] + 1.0


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--arch", default="cuda", choices=("cuda", "gpu", "cpu", "x64"))
    args = parser.parse_args()

    qd.init(arch=getattr(qd, args.arch))

    x_np = np.arange(8, dtype=np.float32)
    expected = 2.0 * x_np + 1.0

    x = qd.ndarray(qd.f32, shape=x_np.shape)
    y = qd.ndarray(qd.f32, shape=x_np.shape)
    x.from_numpy(x_np)

    axpy(x_np.size, x, y)
    qd.sync()

    actual = y.to_numpy()
    if not np.array_equal(actual, expected):
        raise SystemExit(f"Quadrants hello kernel mismatch: {actual} != {expected}")

    print(f"quadrants hello world ok on arch={args.arch}: {actual}")


if __name__ == "__main__":
    main()