import numpy as np
from numba import cuda
from PIL import Image
import math

WIDTH, HEIGHT = 800, 600

@cuda.jit
def render(img):
    x, y = cuda.grid(2)
    if x >= WIDTH or y >= HEIGHT:
        return

    u = (x / WIDTH) * 2.0 - 1.0
    v = (y / HEIGHT) * 2.0 - 1.0
    v *= HEIGHT / WIDTH

    # Ray-sphere intersection
    ro = (0.0, 0.0, -3.0)
    rd = (u, v, 1.0)
    mag = math.sqrt(rd[0]**2 + rd[1]**2 + rd[2]**2)
    rd = (rd[0]/mag, rd[1]/mag, rd[2]/mag)

    b = ro[2] * rd[2]
    c = ro[2]**2 - 1.0
    h = b*b - c

    idx = (y * WIDTH + x) * 3
    if h > 0:
        img[idx]   = 255
        img[idx+1] = 180
        img[idx+2] = 120
    else:
        img[idx]   = 20
        img[idx+1] = 20
        img[idx+2] = 40

img = np.zeros(WIDTH * HEIGHT * 3, dtype=np.uint8)
d_img = cuda.to_device(img)

threads = (16, 16)
blocks = ((WIDTH+15)//16, (HEIGHT+15)//16)

render[blocks, threads](d_img)
d_img.copy_to_host(img)

Image.frombuffer("RGB", (WIDTH, HEIGHT), img).save("numba_cuda_render.png")
print("Rendered using CUDA")