# python/taichi/examples/simulation/fractal.py

import taichi as ti

ti.init(arch=ti.gpu)

n = 320
pixels = ti.field(dtype=float, shape=(n * 2, n))

@ti.func
def complex_sqr(z):
    return ti.Vector([z[0]**2 - z[1]**2, z[1] * z[0] * 2])

@ti.kernel
def paint(t: float):
    for i, j in pixels:  # Parallelized over all pixels
        c = ti.Vector([-0.8, ti.cos(t) * 0.2])
        z = ti.Vector([i / n - 1, j / n - 0.5]) * 2
        iterations = 0
        while z.norm() < 20 and iterations < 50:
            z = complex_sqr(z) + c
            iterations += 1
        pixels[i, j] = 1 - iterations * 0.02

gui = ti.GUI("Julia Set", res=(n * 2, n))

for i in range(1000000):
    paint(i * 0.03)
    gui.set_image(pixels)
    gui.show()

#JULIASET IN JULIA
# Julia Set GPU Kernel using CUDA.jl for better performance
function julia_kernel!(img, w, h, max_iter, c1, c2, zoom, move_x, move_y)
    @cuda threads=256 blocks=ceil(Int, w * h / 256) -> kernel_func(img, w, h, max_iter, c1, c2, zoom, move_x, move_y)
end

@cuda kernel function kernel_func(img, w, h, max_iter, c1, c2, zoom, move_x, move_y)
    tid = threadIdx().x + (blockIdx().x - 1) * blockDim().x
    x = (tid % w) + 1
    y = (tid ÷ w) + 1

    if x <= w && y <= h
        zx, zy = (x / w - 1) * zoom + move_x, (y / h - 0.5) * zoom + move_y
        iter = 0
        while zx * zx + zy * zy < 20.0f0 && iter < max_iter
            zx_new = zx * zx - zy * zy + c1
            zy = 2.0f0 * zx * zy + c2
            zx, iter = zx_new, iter + 1
        end
        img[x, y] = 1.0f0 - iter * 0.02f0
    end
end

# Function to generate and update Julia set dynamically
function generate_julia_set(; w=640, h=480, max_iter=50, c=-0.8 + 0.156im, zoom=2.0, move_x=0.0, move_y=0.0)
    img = CUDA.zeros(Float32, w, h)
    julia_kernel!(img, w, h, max_iter, real(c), imag(c), zoom, move_x, move_y)
    return Array(img)
end

using Oceananigans
grid = RectilinearGrid(CPU(), size=(128, 128), x=(0, 2π), y=(0, 2π), topology=(Periodic, Periodic, Flat))
model = NonhydrostaticModel(; grid, advection=WENO())
ϵ(x, y) = 2rand() - 1
set!(model, u=ϵ, v=ϵ)
simulation = Simulation(model; Δt=0.01, stop_time=4)
run!(simulation)

using DitherPunk
using Images
d = dither(img)                   # apply default algorithm: FloydSteinberg()
d = dither(img, Bayer())          # apply algorithm of choice

dither!(img)                      # or in-place modify image
dither!(img, Bayer())             # with the algorithm of your choice

using Metal

function vadd(a, b, c)
    i = thread_position_in_grid_1d()
    c[i] = a[i] + b[i]
    return
end

dims = (3,4)
a = round.(rand(Float32, dims) * 100)
b = round.(rand(Float32, dims) * 100)
c = similar(a)

d_a = MtlArray(a)
d_b = MtlArray(b)
d_c = MtlArray(c)

len = prod(dims)
@metal threads=len vadd(d_a, d_b, d_c)
c = Array(d_c)