import taichi as ti
import taichi.math as tm
import numpy as np

ti.init(arch=ti.gpu)

# type aliases
vec2 = tm.vec2
vec3 = tm.vec3


# Perlin noise loosely based on https://www.shadertoy.com/view/XdXGW8

@ti.func
def hash(p : vec2):
    p  = 50.0*tm.fract( p*0.3183099 + vec2(0.71,0.113))
    return -1.0+2.0*tm.fract( p.x*p.y*(p.x+p.y) )

@ti.func
def randunit(p : vec2):
    t = tm.pi * hash(p)
    return vec2(tm.cos(t), tm.sin(t))

@ti.func
def gnoise(p : vec2):
    # the four corners of the integer square where p falls
    i00 = tm.floor(p)
    i10 = i00 + vec2(1,0)
    i01 = i00 + vec2(0,1)
    i11 = i00 + vec2(1,1)
    # the values of the four pseudorandom gradients, evaluated at p
    v00 = (p - i00).dot(randunit(i00))
    v01 = (p - i01).dot(randunit(i01))
    v10 = (p - i10).dot(randunit(i10))
    v11 = (p - i11).dot(randunit(i11))
    # the two blending factors (f.x and f.y) we will use to interpolate
    a = p - i00
    f = 3*a*a - 2*a*a*a
    # bilinear interpolation between the four gradient values
    return (
        (v00 * (1-f[0]) + v10 * f[0]) * (1 - f[1]) +
        (v01 * (1-f[0]) + v11 * f[0]) * f[1]
        )

@ti.func
def scalar_noise(p : vec2):
    return gnoise(20*p)


if __name__ == '__main__':

    width = 800
    height = 800

    # draw scalar noise field as intensity
    pixels = ti.field(dtype=float, shape=(width, height))
    @ti.kernel
    def noise_img():
        for i in ti.grouped(pixels):
            p = i / height
            pixels[i] = (1 + scalar_noise(p)) / 2


    gui = ti.GUI("Perlin", res=(width, height))
    while gui.running:
        noise_img()
        for e in gui.get_events(gui.PRESS):
            if e.key == 's':
                print('writing to foo.png')
                ti.tools.imwrite(pixels.to_numpy(), 'foo.png')
        gui.set_image(pixels)
        gui.show()