import taichi as ti
import taichi.math as tm

ti.init(arch=ti.gpu)
W = 640
H = 320

xShift = 0.0
dx = 0.2
pixels = ti.Vector.field(3, dtype=float, shape=(W, H))
img = ti.Vector.field(3, dtype=float, shape=(W, H))

@ti.kernel
def init():
    for i, j in img:
        H = 360 * i / W
        S = j / H
        V = 1

        Hp = H / 60
        C = V * S
        X = C * (1 - abs(Hp % 2 - 1))
        m = V - C
        if Hp >= 0 and Hp < 1:
            img[i, j] = tm.vec3(C, X, 0) + m
        elif Hp >= 1 and Hp < 2:
            img[i, j] = tm.vec3(X, C, 0) + m
        elif Hp >= 2 and Hp < 3:
            img[i, j] = tm.vec3(0, C, X) + m
        elif Hp >= 3 and Hp < 4:
            img[i, j] = tm.vec3(0, X, C) + m
        elif Hp >= 4 and Hp < 5:
            img[i, j] = tm.vec3(X, 0, C) + m
        else:
            img[i, j] = tm.vec3(C, 0, X) + m


@ti.kernel
def paint(xShift: float):
    for i, j in pixels:
        # Parallized over all pixels
        pixels[i, j] = img[int(ti.round(i + xShift)) % W, j]

init()
gui = ti.GUI("Julia Set", res=(W, H))
for i in range(1000000):
    paint(xShift)
    gui.set_image(pixels)
    gui.show()

    xShift = (dx + xShift) % W