import vpsdk.window as window
import vpsdk.process as process
import vpsdk.vpui as vpui
import string
import vipermath as vm
import viperlib
from stdint import *
import memhub
import t, c

GW: t.CDefine = 160
GH: t.CDefine = 90
GW_F: t.CDefine = 160.0
GH_F: t.CDefine = 90.0
MAX_STEPS: t.CDefine = 100
ESCAPE_R: t.CDefine = 60.0
DISK_IN: t.CDefine = 2.6
DISK_OUT: t.CDefine = 11.0
CAM_DIST: t.CDefine = 22.0
CAM_ELEV: t.CDefine = 9.0
FOVF: t.CDefine = 0.55
YIELD_ROWS: t.CDefine = 4
V3_POOL_SIZE: t.CDefine = 8192
V3_POOL_BYTES: t.CDefine = V3_POOL_SIZE * 24 + 512

_v3_mem: t.CArray[t.CUInt8T, V3_POOL_BYTES]
_v3_pool: memhub.MemPool | t.CPtr
_fb_mem: t.CArray[t.CUInt32T, GW * GH]

class V3:
    x: t.CFloat64T
    y: t.CFloat64T
    z: t.CFloat64T

    def __new__() -> 'V3' | t.CPtr:
        return t.CVoid(t.CUInt64T(_v3_pool.alloc(24)), t.CPtr)

    def __init__(self, x: t.CFloat64T, y: t.CFloat64T, z: t.CFloat64T):
        self.x = x
        self.y = y
        self.z = z

    def __add__(self, b: 'V3' | t.CPtr) -> 'V3' | t.CPtr:
        return V3(self.x + b.x, self.y + b.y, self.z + b.z)

    def __sub__(self, b: 'V3' | t.CPtr) -> 'V3' | t.CPtr:
        return V3(self.x - b.x, self.y - b.y, self.z - b.z)

    def __mul__(self, s: t.CFloat64T) -> 'V3' | t.CPtr:
        return V3(self.x * s, self.y * s, self.z * s)

    def __neg__(self) -> 'V3' | t.CPtr:
        return V3(-self.x, -self.y, -self.z)

    def dot(self, b: 'V3' | t.CPtr) -> t.CFloat64T:
        return self.x * b.x + self.y * b.y + self.z * b.z

    def cross(self, b: 'V3' | t.CPtr) -> 'V3' | t.CPtr:
        return V3(self.y * b.z - self.z * b.y, self.z * b.x - self.x * b.z, self.x * b.y - self.y * b.x)

    def nrm(self) -> 'V3' | t.CPtr:
        l: t.CFloat64T = vm.sqrt(self.dot(self))
        if l > 1e-12:
            return self * (1.0 / l)
        return V3(0.0, 0.0, 0.0)

    def lerp(self, b: 'V3' | t.CPtr, t_val: t.CFloat64T) -> 'V3' | t.CPtr:
        return self * (1.0 - t_val) + b * t_val

_g_cam:   t.CArray[t.CFloat64T, 3]
_g_fwd:   t.CArray[t.CFloat64T, 3]
_g_right: t.CArray[t.CFloat64T, 3]
_g_up:    t.CArray[t.CFloat64T, 3]

CAM: V3 | t.CPtr
FWD: V3 | t.CPtr
RIGHT: V3 | t.CPtr
UP: V3 | t.CPtr

def v3_pool_reset():
    _v3_pool.reset()

def clmp(x: t.CFloat64T, a: t.CFloat64T, b: t.CFloat64T) -> t.CFloat64T:
    if x < a: return a
    if x > b: return b
    return x

def smoothstep(a: t.CFloat64T, b: t.CFloat64T, x: t.CFloat64T) -> t.CFloat64T:
    t_val: t.CFloat64T = clmp((x - a) / (b - a), 0.0, 1.0)
    return t_val * t_val * (3.0 - 2.0 * t_val)

def setup_camera():
    global CAM, FWD, RIGHT, UP
    e: t.CFloat64T = CAM_ELEV * 3.14159265358979323846 / 180.0
    CAM = t.CVoid(c.Addr(_g_cam), t.CPtr)
    CAM.x = CAM_DIST * vm.cos(e)
    CAM.y = CAM_DIST * vm.sin(e)
    CAM.z = 0.0
    tmp0: V3 | t.CPtr = V3(0.0, 0.0, 0.0)
    FWD = t.CVoid(c.Addr(_g_fwd), t.CPtr)
    tmp1: V3 | t.CPtr = tmp0 - CAM
    n1: V3 | t.CPtr = tmp1.nrm()
    FWD.x = n1.x
    FWD.y = n1.y
    FWD.z = n1.z
    tmp2: V3 | t.CPtr = V3(0.0, 1.0, 0.0)
    RIGHT = t.CVoid(c.Addr(_g_right), t.CPtr)
    tmp3: V3 | t.CPtr = FWD.cross(tmp2)
    tmp4: V3 | t.CPtr = tmp3.nrm()
    RIGHT.x = tmp4.x
    RIGHT.y = tmp4.y
    RIGHT.z = tmp4.z
    UP = t.CVoid(c.Addr(_g_up), t.CPtr)
    tmp5: V3 | t.CPtr = RIGHT.cross(FWD)
    UP.x = tmp5.x
    UP.y = tmp5.y
    UP.z = tmp5.z

def hash3(x: t.CFloat64T, y: t.CFloat64T, z: t.CFloat64T) -> t.CFloat64T:
    n: t.CFloat64T = vm.sin(x * 127.1 + y * 311.7 + z * 74.7) * 43758.5453
    return n - vm.floor(n)

def stars(d: V3 | t.CPtr) -> V3 | t.CPtr:
    col: V3 | t.CPtr = V3(0.0, 0.0, 0.0)
    i: t.CInt = 0
    while i < 3:
        s: t.CFloat64T = 80.0 + t.CFloat64T(i) * 90.0
        px: t.CFloat64T = d.x * s
        py: t.CFloat64T = d.y * s
        pz: t.CFloat64T = d.z * s
        ix: t.CFloat64T = vm.floor(px)
        iy: t.CFloat64T = vm.floor(py)
        iz: t.CFloat64T = vm.floor(pz)
        h: t.CFloat64T = hash3(ix, iy, iz)
        if h > 0.991:
            fx: t.CFloat64T = px - ix - 0.5
            fy: t.CFloat64T = py - iy - 0.5
            fz: t.CFloat64T = pz - iz - 0.5
            d2: t.CFloat64T = fx * fx + fy * fy + fz * fz
            br: t.CFloat64T = (h - 0.991) / 0.009 * vm.exp(-d2 * 9.0) * 1.4
            tv: t.CFloat64T = hash3(iy, iz, ix)
            tint: V3 | t.CPtr = V3(0.0, 0.0, 0.0)
            if tv < 0.3:
                tint = V3(0.7, 0.8, 1.0)
            elif tv > 0.8:
                tint = V3(1.0, 0.8, 0.6)
            else:
                tint = V3(1.0, 1.0, 1.0)
            col = col + tint * br
        i += 1
    return col

def disk_shade(cp: V3 | t.CPtr, rr: t.CFloat64T, gT: t.CFloat64T) -> V3 | t.CPtr:
    tt: t.CFloat64T = (rr - DISK_IN) / (DISK_OUT - DISK_IN)
    edge: t.CFloat64T = smoothstep(0.0, 0.07, tt) * smoothstep(0.0, 0.18, 1.0 - tt)
    ang: t.CFloat64T = vm.atan2(cp.z, cp.x)
    w: t.CFloat64T = 1.0 / vm.pow(rr, 1.5)
    ph: t.CFloat64T = ang - gT * 0.0011 * w * 9.0
    n: t.CFloat64T = 0.5 + 0.5 * vm.sin(ph * 5.0 - rr * 1.6)
    n2: t.CFloat64T = 0.5 + 0.5 * vm.sin(ph * 13.0 + rr * 0.8 + 1.7)
    dens: t.CFloat64T = (0.35 + 0.65 * n) * (0.55 + 0.45 * n2)
    bright: t.CFloat64T = (0.45 + 1.7 * vm.pow(1.0 - tt, 1.8)) * edge
    vdir: V3 | t.CPtr = V3(-cp.z, 0.0, cp.x).nrm()
    beta: t.CFloat64T = 0.46 / vm.sqrt(rr)
    to_cam: V3 | t.CPtr = (CAM - cp).nrm()
    mu: t.CFloat64T = vdir.dot(to_cam)
    dop: t.CFloat64T = 1.0 / (1.0 - beta * mu)
    beam: t.CFloat64T = vm.pow(dop, 3.0)
    c1: V3 | t.CPtr = V3(1.0, 0.92, 0.68)
    c2: V3 | t.CPtr = V3(1.0, 0.42, 0.13)
    base: V3 | t.CPtr = c1.lerp(c2, tt)
    base.z += (dop - 1.0) * 0.35
    base.x -= (dop - 1.0) * 0.05
    inten: t.CFloat64T = bright * dens * beam * 1.25
    return base * inten

def trace(dir_v: V3 | t.CPtr, gT: t.CFloat64T) -> V3 | t.CPtr:
    v3_pool_reset()
    pos: V3 | t.CPtr = V3(CAM.x, CAM.y, CAM.z)
    vel: V3 | t.CPtr = V3(dir_v.x, dir_v.y, dir_v.z)
    hvec: V3 | t.CPtr = pos.cross(vel)
    h2: t.CFloat64T = hvec.dot(hvec)
    T: t.CFloat64T = 1.0
    col: V3 | t.CPtr = V3(0.0, 0.0, 0.0)

    i: t.CInt = 0
    while i < MAX_STEPS:
        r2: t.CFloat64T = pos.dot(pos)
        r: t.CFloat64T = vm.sqrt(r2)
        if r < 1.0:
            break
        if r > ESCAPE_R:
            col = col + stars(vel.nrm()) * T
            break

        dt: t.CFloat64T = clmp(r * 0.08, 0.03, 2.2)
        r5: t.CFloat64T = vm.pow(r2, 2.5)
        coeff: t.CFloat64T = -1.5 * h2 / r5

        k1v: V3 | t.CPtr = pos * coeff
        k2p_a: V3 | t.CPtr = vel + k1v * (dt * 0.5)
        k2v_a: V3 | t.CPtr = pos + vel * (dt * 0.5)
        r2_a: t.CFloat64T = k2v_a.dot(k2v_a)
        coeff_a: t.CFloat64T = -1.5 * h2 / vm.pow(r2_a, 2.5)
        k2v: V3 | t.CPtr = k2v_a * coeff_a

        k3p_a: V3 | t.CPtr = vel + k2v * (dt * 0.5)
        k3v_a: V3 | t.CPtr = pos + k2p_a * (dt * 0.5)
        r2_b: t.CFloat64T = k3v_a.dot(k3v_a)
        coeff_b: t.CFloat64T = -1.5 * h2 / vm.pow(r2_b, 2.5)
        k3v: V3 | t.CPtr = k3v_a * coeff_b

        k4p_a: V3 | t.CPtr = vel + k3v * dt
        k4v_a: V3 | t.CPtr = pos + k3p_a * dt
        r2_c: t.CFloat64T = k4v_a.dot(k4v_a)
        coeff_c: t.CFloat64T = -1.5 * h2 / vm.pow(r2_c, 2.5)
        k4v: V3 | t.CPtr = k4v_a * coeff_c

        npos: V3 | t.CPtr = pos + (vel + k2p_a * 2.0 + k3p_a * 2.0 + k4p_a) * (dt / 6.0)
        nvel: V3 | t.CPtr = vel + (k1v + k2v * 2.0 + k3v * 2.0 + k4v) * (dt / 6.0)

        if pos.y * npos.y < 0.0:
            f: t.CFloat64T = pos.y / (pos.y - npos.y)
            cp: V3 | t.CPtr = pos.lerp(npos, f)
            rr: t.CFloat64T = vm.sqrt(cp.x * cp.x + cp.z * cp.z)
            if rr > DISK_IN and rr < DISK_OUT:
                emit: V3 | t.CPtr = disk_shade(cp, rr, gT)
                col = col + emit * T
                a: t.CFloat64T = clmp(smoothstep(0.0, 0.07, (rr - DISK_IN) / (DISK_OUT - DISK_IN)) * smoothstep(0.0, 0.18, 1.0 - (rr - DISK_IN) / (DISK_OUT - DISK_IN)) * 0.92, 0.0, 1.0)
                T = T * (1.0 - a)
                if T < 0.01:
                    break

        pos = npos
        vel = nvel
        i += 1
    return col

def shade_pixel(px: t.CInt, py: t.CInt, gT: t.CFloat64T) -> t.CUInt32T:
    aspect: t.CFloat64T = GW_F / GH_F
    sx: t.CFloat64T = (2.0 * (t.CFloat64T(px) + 0.5) / GW_F - 1.0) * aspect * FOVF
    sy: t.CFloat64T = (1.0 - 2.0 * (t.CFloat64T(py) + 0.5) / GH_F) * FOVF
    dir_v: V3 | t.CPtr = (FWD + RIGHT * sx + UP * sy).nrm()
    c: V3 | t.CPtr = trace(dir_v, gT)
    ex: t.CFloat64T = 1.15
    c = c * ex
    c.x = c.x / (1.0 + c.x)
    c.y = c.y / (1.0 + c.y)
    c.z = c.z / (1.0 + c.z)
    c.x = vm.pow(c.x, 1.0 / 2.2)
    c.y = vm.pow(c.y, 1.0 / 2.2)
    c.z = vm.pow(c.z, 1.0 / 2.2)
    R: t.CInt = t.CInt(clmp(c.x, 0.0, 1.0) * 255.0)
    G: t.CInt = t.CInt(clmp(c.y, 0.0, 1.0) * 255.0)
    B: t.CInt = t.CInt(clmp(c.z, 0.0, 1.0) * 255.0)
    return t.CUInt32T((R << 16) | (G << 8) | B)

@c.Attribute(t.attr.section(".text.startup"), t.attr.aligned(16))
def _start() -> t.CInt | t.CExport:
    _v3_pool = memhub.MemPool(c.Addr(_v3_mem), V3_POOL_BYTES)
    setup_camera()

    root: vpui.Tk = vpui.Tk()
    root.title("Gargantua")
    root.set_style(window.STYLE_WIN7)
    root.bg_color = t.CUInt32T(0x000000)
    master: vpui.Tk | t.CPtr = c.Addr(root)

    canvas: vpui.FBCanvas = vpui.FBCanvas(master, GW, GH)
    root.pack(0, vpui.PACK_TOP, vpui.FILL_BOTH, 1, 0, 0)

    fb: UINT32PTR = c.Addr(_fb_mem)
    string.memset(fb, 0, GW * GH * 4)
    canvas.set_fb(fb, GW, GH)

    root.geometry(100, 30, GW, GH + 32)

    gT: t.CFloat64T = 0.0
    rendered: t.CInt = 0

    evt: window.InputEvent
    while window.is_active(root.win_id):
        result: t.CInt = window.poll_event(c.Addr(evt))
        if result != 0:
            if evt.type == window.EVENT_TYPE_PING:
                window.flush(root.win_id)
            elif evt.type == window.EVENT_TYPE_RESIZE:
                sz: t.CInt = window.get_win_size(root.win_id)
                new_w: t.CInt = sz >> 16
                new_h: t.CInt = sz & 0xFFFF
                canvas.set_fb(fb, new_w, new_h)

        if rendered == 0:
            y: t.CInt = 0
            while y < GH:
                x: t.CInt = 0
                while x < GW:
                    pixel: t.CUInt32T = shade_pixel(x, y, gT)
                    fb[y * GW + x] = pixel
                    x += 1
                y += 1
                if y % YIELD_ROWS == 0:
                    root.redraw()
                    process.yield_cpu()
                    result2: t.CInt = window.poll_event(c.Addr(evt))
                    if result2 != 0:
                        if evt.type == window.EVENT_TYPE_PING:
                            window.flush(root.win_id)
            root.redraw()
            rendered = 1
            gT += 50.0

        process.yield_cpu()

    process.exit(0)
    return 0