include(joinpath(@__DIR__, "built_in", "py_runtime.jl"))

include(joinpath(@__DIR__, "std", "math.jl"))
math = (ceil=ceil, cos=cos, e=e, exp=exp, fabs=fabs, floor=floor, log=log, log10=log10, pi=pi, pow=pow, sin=sin, sqrt=sqrt, tan=tan)
include(joinpath(@__DIR__, "std", "time.jl"))
include(joinpath(@__DIR__, "utils", "gif.jl"))

# 16: Sample that ray-traces chaotic rotation of glass sculptures and outputs a GIF.

function clamp01(v)
    if (v < 0.0)
        return 0.0
    end
    if (v > 1.0)
        return 1.0
    end
    return v
end

function dot(ax, ay, az, bx, by, bz)
    return (((ax * bx) + (ay * by)) + (az * bz))
end

function length_(x, y, z)
    return math.sqrt((((x * x) + (y * y)) + (z * z)))
end

function normalize(x, y, z)
    l = length_(x, y, z)
    if (l < 1e-09)
        return (0.0, 0.0, 0.0)
    end
    return ((x / l), (y / l), (z / l))
end

function reflect(ix, iy, iz, nx, ny, nz)
    d = (dot(ix, iy, iz, nx, ny, nz) * 2.0)
    return ((ix - (d * nx)), (iy - (d * ny)), (iz - (d * nz)))
end

function refract(ix, iy, iz, nx, ny, nz, eta)
    # Simple IOR-based refraction. Return reflection direction on total internal reflection.
    cosi = (-dot(ix, iy, iz, nx, ny, nz))
    sint2 = ((eta * eta) * (1.0 - (cosi * cosi)))
    if (sint2 > 1.0)
        return reflect(ix, iy, iz, nx, ny, nz)
    end
    cost = math.sqrt((1.0 - sint2))
    k = ((eta * cosi) - cost)
    return (((eta * ix) + (k * nx)), ((eta * iy) + (k * ny)), ((eta * iz) + (k * nz)))
end

function schlick(cos_theta, f0)
    m = (1.0 - cos_theta)
    return (f0 + ((1.0 - f0) * ((((m * m) * m) * m) * m)))
end

function sky_color(dx, dy, dz, tphase)
    # Sky gradient + neon band
    t = (0.5 * (dy + 1.0))
    r = (0.06 + (0.2 * t))
    g = (0.1 + (0.25 * t))
    b = (0.16 + (0.45 * t))
    band = (0.5 + (0.5 * math.sin((((8.0 * dx) + (6.0 * dz)) + tphase))))
    r = r + (0.08 * band)
    g = g + (0.05 * band)
    b = b + (0.12 * band)
    return (clamp01(r), clamp01(g), clamp01(b))
end

function sphere_intersect(ox, oy, oz, dx, dy, dz, cx, cy, cz, radius)
    lx = (ox - cx)
    ly = (oy - cy)
    lz = (oz - cz)
    b = (((lx * dx) + (ly * dy)) + (lz * dz))
    c = ((((lx * lx) + (ly * ly)) + (lz * lz)) - (radius * radius))
    h = ((b * b) - c)
    if (h < 0.0)
        return (-1.0)
    end
    s = math.sqrt(h)
    t0 = ((-b) - s)
    if (t0 > 0.0001)
        return t0
    end
    t1 = ((-b) + s)
    if (t1 > 0.0001)
        return t1
    end
    return (-1.0)
end

function palette_332()
    # 3-3-2 quantized palette. Lightweight quantization that stays fast after transpilation.
    p = __pytra_bytearray((256 * 3))
    for i in 0:256 - 1
        r = ((i >> 5) & 7)
        g = ((i >> 2) & 7)
        b = (i & 3)
        p[__pytra_idx(((i * 3) + 0), length(p))] = __pytra_int(((255 * r) / 7))
        p[__pytra_idx(((i * 3) + 1), length(p))] = __pytra_int(((255 * g) / 7))
        p[__pytra_idx(((i * 3) + 2), length(p))] = __pytra_int(((255 * b) / 3))
    end
    return __pytra_bytes(p)
end

function quantize_332(r, g, b)
    rr = __pytra_int((clamp01(r) * 255.0))
    gg = __pytra_int((clamp01(g) * 255.0))
    bb = __pytra_int((clamp01(b) * 255.0))
    return ((((rr >> 5) << 5) + ((gg >> 5) << 2)) + (bb >> 6))
end

function render_frame(width, height, frame_id, frames_n)
    t = (frame_id / frames_n)
    tphase = ((2.0 * math.pi) * t)
    
    # Camera slowly orbits.
    cam_r = 3.0
    cam_x = (cam_r * math.cos((tphase * 0.9)))
    cam_y = (1.1 + (0.25 * math.sin((tphase * 0.6))))
    cam_z = (cam_r * math.sin((tphase * 0.9)))
    look_x = 0.0
    look_y = 0.35
    look_z = 0.0
    
    (fwd_x, fwd_y, fwd_z) = normalize((look_x - cam_x), (look_y - cam_y), (look_z - cam_z))
    (right_x, right_y, right_z) = normalize(fwd_z, 0.0, (-fwd_x))
    (up_x, up_y, up_z) = normalize(((right_y * fwd_z) - (right_z * fwd_y)), ((right_z * fwd_x) - (right_x * fwd_z)), ((right_x * fwd_y) - (right_y * fwd_x)))
    
    # Moving glass sculpture (3 spheres) and an emissive sphere.
    s0x = (0.9 * math.cos((1.3 * tphase)))
    s0y = (0.15 + (0.35 * math.sin((1.7 * tphase))))
    s0z = (0.9 * math.sin((1.3 * tphase)))
    s1x = (1.2 * math.cos(((1.3 * tphase) + 2.094)))
    s1y = (0.1 + (0.4 * math.sin(((1.1 * tphase) + 0.8))))
    s1z = (1.2 * math.sin(((1.3 * tphase) + 2.094)))
    s2x = (1.0 * math.cos(((1.3 * tphase) + 4.188)))
    s2y = (0.2 + (0.3 * math.sin(((1.5 * tphase) + 1.9))))
    s2z = (1.0 * math.sin(((1.3 * tphase) + 4.188)))
    lr = 0.35
    lx = (2.4 * math.cos((tphase * 1.8)))
    ly = (1.8 + (0.8 * math.sin((tphase * 1.2))))
    lz = (2.4 * math.sin((tphase * 1.8)))
    
    frame = __pytra_bytearray((width * height))
    aspect = (width / height)
    fov = 1.25
    
    for py in 0:height - 1
        row_base = (py * width)
        sy = (1.0 - ((2.0 * (py + 0.5)) / height))
        for px in 0:width - 1
            sx = ((((2.0 * (px + 0.5)) / width) - 1.0) * aspect)
            rx = (fwd_x + (fov * ((sx * right_x) + (sy * up_x))))
            ry = (fwd_y + (fov * ((sx * right_y) + (sy * up_y))))
            rz = (fwd_z + (fov * ((sx * right_z) + (sy * up_z))))
            (dx, dy, dz) = normalize(rx, ry, rz)
            
            # Search for the nearest hit.
            best_t = 1000000000.0
            hit_kind = 0
            r = 0.0
            g = 0.0
            b = 0.0
            
            # Floor plane y=-1.2
            if (dy < (-1e-06))
                tf = (((-1.2) - cam_y) / dy)
                if ((tf > 0.0001) && (tf < best_t))
                    best_t = tf
                    hit_kind = 1
                end
            end
            t0 = sphere_intersect(cam_x, cam_y, cam_z, dx, dy, dz, s0x, s0y, s0z, 0.65)
            if ((t0 > 0.0) && (t0 < best_t))
                best_t = t0
                hit_kind = 2
            end
            t1 = sphere_intersect(cam_x, cam_y, cam_z, dx, dy, dz, s1x, s1y, s1z, 0.72)
            if ((t1 > 0.0) && (t1 < best_t))
                best_t = t1
                hit_kind = 3
            end
            t2 = sphere_intersect(cam_x, cam_y, cam_z, dx, dy, dz, s2x, s2y, s2z, 0.58)
            if ((t2 > 0.0) && (t2 < best_t))
                best_t = t2
                hit_kind = 4
            end
            glow = nothing
            hx = nothing
            hz = nothing
            ldx = nothing
            ldy = nothing
            ldz = nothing
            lxv = nothing
            lyv = nothing
            lzv = nothing
            ndotl = nothing
            if (hit_kind == 0)
                (r, g, b) = sky_color(dx, dy, dz, tphase)
            elseif (hit_kind == 1)
                hx = (cam_x + (best_t * dx))
                hz = (cam_z + (best_t * dz))
                cx_i = __pytra_int(math.floor((hx * 2.0)))
                cz_i = __pytra_int(math.floor((hz * 2.0)))
                checker = (__pytra_truthy((((cx_i + cz_i) % 2) == 0)) ? (0) : (1))
                base_r = (__pytra_truthy((checker == 0)) ? (0.1) : (0.04))
                base_g = (__pytra_truthy((checker == 0)) ? (0.11) : (0.05))
                base_b = (__pytra_truthy((checker == 0)) ? (0.13) : (0.08))
                # Emissive sphere contribution.
                lxv = (lx - hx)
                lyv = (ly - (-1.2))
                lzv = (lz - hz)
                (ldx, ldy, ldz) = normalize(lxv, lyv, lzv)
                ndotl = max(ldy, 0.0)
                ldist2 = (((lxv * lxv) + (lyv * lyv)) + (lzv * lzv))
                glow = (8.0 / (1.0 + ldist2))
                r = ((base_r + (0.8 * glow)) + (0.2 * ndotl))
                g = ((base_g + (0.5 * glow)) + (0.18 * ndotl))
                b = ((base_b + (1.0 * glow)) + (0.24 * ndotl))
            else
                cx = 0.0
                cy = 0.0
                cz = 0.0
                rad = 1.0
                if (hit_kind == 2)
                    cx = s0x
                    cy = s0y
                    cz = s0z
                    rad = 0.65
                elseif (hit_kind == 3)
                    cx = s1x
                    cy = s1y
                    cz = s1z
                    rad = 0.72
                else
                    cx = s2x
                    cy = s2y
                    cz = s2z
                    rad = 0.58
                end
                hx = (cam_x + (best_t * dx))
                hy = (cam_y + (best_t * dy))
                hz = (cam_z + (best_t * dz))
                (nx, ny, nz) = normalize(((hx - cx) / rad), ((hy - cy) / rad), ((hz - cz) / rad))
                
                # Simple glass shading (reflection + refraction + light highlights).
                (rdx, rdy, rdz) = reflect(dx, dy, dz, nx, ny, nz)
                (tdx, tdy, tdz) = refract(dx, dy, dz, nx, ny, nz, (1.0 / 1.45))
                (sr, sg, sb) = sky_color(rdx, rdy, rdz, tphase)
                (tr, tg, tb) = sky_color(tdx, tdy, tdz, (tphase + 0.8))
                cosi = max((-(((dx * nx) + (dy * ny)) + (dz * nz))), 0.0)
                fr = schlick(cosi, 0.04)
                r = ((tr * (1.0 - fr)) + (sr * fr))
                g = ((tg * (1.0 - fr)) + (sg * fr))
                b = ((tb * (1.0 - fr)) + (sb * fr))
                
                lxv = (lx - hx)
                lyv = (ly - hy)
                lzv = (lz - hz)
                (ldx, ldy, ldz) = normalize(lxv, lyv, lzv)
                ndotl = max((((nx * ldx) + (ny * ldy)) + (nz * ldz)), 0.0)
                (hvx, hvy, hvz) = normalize((ldx - dx), (ldy - dy), (ldz - dz))
                ndoth = max((((nx * hvx) + (ny * hvy)) + (nz * hvz)), 0.0)
                spec = (ndoth * ndoth)
                spec = (spec * spec)
                spec = (spec * spec)
                spec = (spec * spec)
                glow = (10.0 / (((1.0 + (lxv * lxv)) + (lyv * lyv)) + (lzv * lzv)))
                r = r + (((0.2 * ndotl) + (0.8 * spec)) + (0.45 * glow))
                g = g + (((0.18 * ndotl) + (0.6 * spec)) + (0.35 * glow))
                b = b + (((0.26 * ndotl) + (1.0 * spec)) + (0.65 * glow))
                
                # Slight tint variation per sphere.
                if (hit_kind == 2)
                    r = r * 0.95
                    g = g * 1.05
                    b = b * 1.1
                elseif (hit_kind == 3)
                    r = r * 1.08
                    g = g * 0.98
                    b = b * 1.04
                else
                    r = r * 1.02
                    g = g * 1.1
                    b = b * 0.95
                end
            end
            # Slightly stronger tone mapping.
            r = math.sqrt(clamp01(r))
            g = math.sqrt(clamp01(g))
            b = math.sqrt(clamp01(b))
            frame[__pytra_idx((row_base + px), length(frame))] = quantize_332(r, g, b)
        end
    end
    return __pytra_bytes(frame)
end

function run_16_glass_sculpture_chaos()
    width = 320
    height = 240
    frames_n = 72
    out_path = "sample/out/16_glass_sculpture_chaos.gif"
    
    start = perf_counter()
    frames = Any[]
    for i in 0:frames_n - 1
        push!(frames, render_frame(width, height, i, frames_n))
    end
    save_gif(out_path, width, height, frames, palette_332(), 6, 0)
    elapsed = (perf_counter() - start)
    __pytra_print("output:", out_path)
    __pytra_print("frames:", frames_n)
    __pytra_print("elapsed_sec:", elapsed)
end


run_16_glass_sculpture_chaos();