#include "runtime/cpp/core/py_runtime.h"
#include "runtime/cpp/core/process_runtime.h"
#include "built_in/io_ops.h"
#include "built_in/numeric_ops.h"
#include "std/math.h"
#include "std/time.h"
#include "utils/gif.h"

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

float64 clamp01(float64 v) {
    if (v < 0.0)
        return 0.0;
    if (v > 1.0)
        return 1.0;
    return v;
}

float64 dot(float64 ax, float64 ay, float64 az, float64 bx, float64 by, float64 bz) {
    return ax * bx + ay * by + az * bz;
}

float64 length(float64 x, float64 y, float64 z) {
    return pytra::std::math::sqrt(x * x + y * y + z * z);
}

::std::tuple<float64, float64, float64> normalize(float64 x, float64 y, float64 z) {
    float64 l = length(x, y, z);
    if (l < 1e-9)
        return ::std::make_tuple(0.0, 0.0, 0.0);
    return ::std::make_tuple(x / l, y / l, z / l);
}

::std::tuple<float64, float64, float64> reflect(float64 ix, float64 iy, float64 iz, float64 nx, float64 ny, float64 nz) {
    float64 d = dot(ix, iy, iz, nx, ny, nz) * 2.0;
    return ::std::make_tuple(ix - d * nx, iy - d * ny, iz - d * nz);
}

::std::tuple<float64, float64, float64> refract(float64 ix, float64 iy, float64 iz, float64 nx, float64 ny, float64 nz, float64 eta) {
    // Simple IOR-based refraction. Return reflection direction on total internal reflection.
    float64 cosi = -(dot(ix, iy, iz, nx, ny, nz));
    float64 sint2 = eta * eta * (1.0 - cosi * cosi);
    if (sint2 > 1.0)
        return reflect(ix, iy, iz, nx, ny, nz);
    float64 cost = pytra::std::math::sqrt(1.0 - sint2);
    float64 k = eta * cosi - cost;
    return ::std::make_tuple(eta * ix + k * nx, eta * iy + k * ny, eta * iz + k * nz);
}

float64 schlick(float64 cos_theta, float64 f0) {
    float64 m = 1.0 - cos_theta;
    return f0 + (1.0 - f0) * m * m * m * m * m;
}

::std::tuple<float64, float64, float64> sky_color(float64 dx, float64 dy, float64 dz, float64 tphase) {
    // Sky gradient + neon band
    float64 t = 0.5 * (dy + 1.0);
    float64 r = 0.06 + 0.20 * t;
    float64 g = 0.10 + 0.25 * t;
    float64 b = 0.16 + 0.45 * t;
    float64 band = 0.5 + 0.5 * pytra::std::math::sin(8.0 * dx + 6.0 * dz + tphase);
    r += 0.08 * band;
    g += 0.05 * band;
    b += 0.12 * band;
    return ::std::make_tuple(clamp01(r), clamp01(g), clamp01(b));
}

float64 sphere_intersect(float64 ox, float64 oy, float64 oz, float64 dx, float64 dy, float64 dz, float64 cx, float64 cy, float64 cz, float64 radius) {
    float64 lx = ox - cx;
    float64 ly = oy - cy;
    float64 lz = oz - cz;
    float64 b = lx * dx + ly * dy + lz * dz;
    float64 c = lx * lx + ly * ly + lz * lz - radius * radius;
    float64 h = b * b - c;
    if (h < 0.0)
        return -(1.0);
    float64 s = pytra::std::math::sqrt(h);
    float64 t0 = -(b) - s;
    if (t0 > 1e-4)
        return t0;
    float64 t1 = -(b) + s;
    if (t1 > 1e-4)
        return t1;
    return -(1.0);
}

bytes palette_332() {
    // 3-3-2 quantized palette. Lightweight quantization that stays fast after transpilation.
    bytearray p = bytearray(256 * 3);
    float64 __hoisted_cast_1 = float64(7);
    float64 __hoisted_cast_2 = float64(3);
    for (int64 i = 0; i < 256; ++i) {
        int64 r = i >> 5 & 7;
        int64 g = i >> 2 & 7;
        int64 b = i & 3;
        p[i * 3] = int64(float64(255 * r) / __hoisted_cast_1);
        p[i * 3 + 1] = int64(float64(255 * g) / __hoisted_cast_1);
        p[i * 3 + 2] = int64(float64(255 * b) / __hoisted_cast_2);
    }
    return p;
}

int64 quantize_332(float64 r, float64 g, float64 b) {
    int64 rr = int64(clamp01(r) * 255.0);
    int64 gg = int64(clamp01(g) * 255.0);
    int64 bb = int64(clamp01(b) * 255.0);
    return (rr >> 5 << 5) + (gg >> 5 << 2) + (bb >> 6);
}

bytes render_frame(int64 width, int64 height, int64 frame_id, int64 frames_n) {
    float64 t = float64(frame_id) / float64(frames_n);
    auto tphase = 2.0 * pytra::std::math::pi * t;
    
    // Camera slowly orbits.
    float64 cam_r = 3.0;
    float64 cam_x = cam_r * pytra::std::math::cos(tphase * 0.9);
    float64 cam_y = 1.1 + 0.25 * pytra::std::math::sin(tphase * 0.6);
    float64 cam_z = cam_r * pytra::std::math::sin(tphase * 0.9);
    float64 look_x = 0.0;
    float64 look_y = 0.35;
    float64 look_z = 0.0;
    
    auto [fwd_x, fwd_y, fwd_z] = normalize(look_x - cam_x, look_y - cam_y, look_z - cam_z);
    auto [right_x, right_y, right_z] = normalize(fwd_z, 0.0, -(fwd_x));
    auto [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.
    float64 s0x = 0.9 * pytra::std::math::cos(1.3 * tphase);
    float64 s0y = 0.15 + 0.35 * pytra::std::math::sin(1.7 * tphase);
    float64 s0z = 0.9 * pytra::std::math::sin(1.3 * tphase);
    float64 s1x = 1.2 * pytra::std::math::cos(1.3 * tphase + 2.094);
    float64 s1y = 0.10 + 0.40 * pytra::std::math::sin(1.1 * tphase + 0.8);
    float64 s1z = 1.2 * pytra::std::math::sin(1.3 * tphase + 2.094);
    float64 s2x = pytra::std::math::cos(1.3 * tphase + 4.188);
    float64 s2y = 0.20 + 0.30 * pytra::std::math::sin(1.5 * tphase + 1.9);
    float64 s2z = pytra::std::math::sin(1.3 * tphase + 4.188);
    float64 lr = 0.35;
    float64 lx = 2.4 * pytra::std::math::cos(tphase * 1.8);
    float64 ly = 1.8 + 0.8 * pytra::std::math::sin(tphase * 1.2);
    float64 lz = 2.4 * pytra::std::math::sin(tphase * 1.8);
    
    bytearray frame = bytearray(width * height);
    float64 aspect = float64(width) / float64(height);
    float64 fov = 1.25;
    float64 __hoisted_cast_3 = float64(height);
    float64 __hoisted_cast_4 = float64(width);
    
    for (int64 py = 0; py < height; ++py) {
        int64 row_base = py * width;
        float64 sy = 1.0 - 2.0 * (float64(py) + 0.5) / __hoisted_cast_3;
        for (int64 px = 0; px < width; ++px) {
            float64 sx = (2.0 * (float64(px) + 0.5) / __hoisted_cast_4 - 1.0) * aspect;
            float64 rx = fwd_x + fov * (sx * right_x + sy * up_x);
            float64 ry = fwd_y + fov * (sx * right_y + sy * up_y);
            float64 rz = fwd_z + fov * (sx * right_z + sy * up_z);
            auto [dx, dy, dz] = normalize(rx, ry, rz);
            
            // Search for the nearest hit.
            float64 best_t = 1e9;
            int64 hit_kind = 0;
            float64 r = 0.0;
            float64 g = 0.0;
            float64 b = 0.0;
            
            // Floor plane y=-1.2
            if (dy < -(1e-6)) {
                float64 tf = (-(1.2) - cam_y) / dy;
                if ((tf > 1e-4) && (tf < best_t)) {
                    best_t = tf;
                    hit_kind = 1;
                }
            }
            float64 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;
            }
            float64 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;
            }
            float64 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;
            }
            if (hit_kind == 0) {
                auto __tuple_5 = sky_color(dx, dy, dz, tphase);
                r = ::std::get<0>(__tuple_5);
                g = ::std::get<1>(__tuple_5);
                b = ::std::get<2>(__tuple_5);
            } else if (hit_kind == 1) {
                float64 hx = cam_x + best_t * dx;
                float64 hz = cam_z + best_t * dz;
                int64 cx = int64(pytra::std::math::floor(hx * 2.0));
                int64 cz = int64(pytra::std::math::floor(hz * 2.0));
                int64 checker = ((cx + cz) % 2 == 0 ? 0 : 1);
                float64 base_r = (checker == 0 ? 0.10 : 0.04);
                float64 base_g = (checker == 0 ? 0.11 : 0.05);
                float64 base_b = (checker == 0 ? 0.13 : 0.08);
                // Emissive sphere contribution.
                float64 lxv = lx - hx;
                float64 lyv = ly - -(1.2);
                float64 lzv = lz - hz;
                auto [ldx, ldy, ldz] = normalize(lxv, lyv, lzv);
                float64 ndotl = ::std::max<float64>(ldy, 0.0);
                float64 ldist2 = lxv * lxv + lyv * lyv + lzv * lzv;
                float64 glow = 8.0 / (1.0 + ldist2);
                r = base_r + 0.8 * glow + 0.20 * ndotl;
                g = base_g + 0.5 * glow + 0.18 * ndotl;
                b = base_b + 1.0 * glow + 0.24 * ndotl;
            } else {
                float64 cx = 0.0;
                float64 cy = 0.0;
                float64 cz = 0.0;
                float64 rad = 1.0;
                if (hit_kind == 2) {
                    cx = s0x;
                    cy = s0y;
                    cz = s0z;
                    rad = 0.65;
                } else if (hit_kind == 3) {
                    cx = s1x;
                    cy = s1y;
                    cz = s1z;
                    rad = 0.72;
                } else {
                    cx = s2x;
                    cy = s2y;
                    cz = s2z;
                    rad = 0.58;
                }
                float64 hx = cam_x + best_t * dx;
                float64 hy = cam_y + best_t * dy;
                float64 hz = cam_z + best_t * dz;
                auto [nx, ny, nz] = normalize((hx - cx) / rad, (hy - cy) / rad, (hz - cz) / rad);
                
                // Simple glass shading (reflection + refraction + light highlights).
                auto [rdx, rdy, rdz] = reflect(dx, dy, dz, nx, ny, nz);
                auto [tdx, tdy, tdz] = refract(dx, dy, dz, nx, ny, nz, 1.0 / 1.45);
                auto [sr, sg, sb] = sky_color(rdx, rdy, rdz, tphase);
                auto [tr, tg, tb] = sky_color(tdx, tdy, tdz, tphase + 0.8);
                float64 cosi = ::std::max<float64>(float64(-(dx * nx + dy * ny + dz * nz)), 0.0);
                float64 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;
                
                float64 lxv = lx - hx;
                float64 lyv = ly - hy;
                float64 lzv = lz - hz;
                auto [ldx, ldy, ldz] = normalize(lxv, lyv, lzv);
                float64 ndotl = ::std::max<float64>(nx * ldx + ny * ldy + nz * ldz, 0.0);
                auto [hvx, hvy, hvz] = normalize(ldx - dx, ldy - dy, ldz - dz);
                float64 ndoth = ::std::max<float64>(nx * hvx + ny * hvy + nz * hvz, 0.0);
                float64 spec = ndoth * ndoth;
                spec = spec * spec;
                spec = spec * spec;
                spec = spec * spec;
                float64 glow = 10.0 / (1.0 + lxv * lxv + lyv * lyv + lzv * lzv);
                r += 0.20 * ndotl + 0.80 * spec + 0.45 * glow;
                g += 0.18 * ndotl + 0.60 * spec + 0.35 * glow;
                b += 0.26 * ndotl + 1.00 * spec + 0.65 * glow;
                
                // Slight tint variation per sphere.
                if (hit_kind == 2) {
                    r *= 0.95;
                    g *= 1.05;
                    b *= 1.10;
                } else if (hit_kind == 3) {
                    r *= 1.08;
                    g *= 0.98;
                    b *= 1.04;
                } else {
                    r *= 1.02;
                    g *= 1.10;
                    b *= 0.95;
                }
            }
            // Slightly stronger tone mapping.
            r = pytra::std::math::sqrt(clamp01(r));
            g = pytra::std::math::sqrt(clamp01(g));
            b = pytra::std::math::sqrt(clamp01(b));
            frame[row_base + px] = quantize_332(r, g, b);
        }
    }
    return frame;
}

void run_16_glass_sculpture_chaos() {
    int64 width = 320;
    int64 height = 240;
    int64 frames_n = 72;
    str out_path = "sample/out/16_glass_sculpture_chaos.gif";
    
    float64 start = pytra::std::time::perf_counter();
    rc<list<bytes>> frames = rc_list_from_value(list<bytes>{});
    rc_list_ref(frames).reserve((frames_n <= 0) ? 0 : frames_n);
    for (int64 i = 0; i < frames_n; ++i)
        rc_list_ref(frames).append(render_frame(width, height, i, frames_n));
    pytra::utils::gif::save_gif(out_path, width, height, rc_list_ref(frames), palette_332(), 6, 0);
    float64 elapsed = pytra::std::time::perf_counter() - start;
    py_print("output:", out_path);
    py_print("frames:", frames_n);
    py_print("elapsed_sec:", elapsed);
}

int main(int argc, char** argv) {
    pytra_configure_from_argv(argc, argv);
    run_16_glass_sculpture_chaos();
    return 0;
}