Black Hole — gravitational lensing

use system.collections.array
use system.math
use system.io

// Canvas: 1280×720 pixels (16:9); RGBA paint output.
const CW = 1280
const CH = 720
const PIXELS = CW * CH
const PAINT = PIXELS * 4

// Camera state: [yaw, pitch, distance, idle_seconds] — ping-ponged across frames.
gpu var cam_a = Array<f32, 4>()
gpu var cam_b = Array<f32, 4>()

// Paint output (RGBA: 4 floats per pixel).
gpu var paint = Array<f32, PAINT>()

// 2D value hash → [0,1); the base entropy source for the nebula.
fn hash21(px f32, py f32) f32
    var qx = fract(px * 123.34) as f32
    var qy = fract(py * 345.45) as f32
    let d = qx * (qx + 34.345) + qy * (qy + 34.345)
    qx = qx + d
    qy = qy + d
    return fract(qx * qy) as f32

// 3D value hash → [0,1); seeds the volumetric starfield cells.
fn hash31(px f32, py f32, pz f32) f32
    let ax = fract(px * 0.3183099 + 0.1) as f32
    let ay = fract(py * 0.3183099 + 0.1) as f32
    let az = fract(pz * 0.3183099 + 0.1) as f32
    let bx = ax * 17.0
    let by = ay * 17.0
    let bz = az * 17.0
    return fract(bx * by * bz * (bx + by + bz)) as f32

// Smooth 2D value noise (bilinear interpolation over four hashed corners).
fn vnoise(px f32, py f32) f32
    let ix = floor(px) as f32
    let iy = floor(py) as f32
    var fx = px - ix
    var fy = py - iy
    fx = fx * fx * (3.0 - 2.0 * fx)
    fy = fy * fy * (3.0 - 2.0 * fy)
    let a = hash21(ix, iy)
    let b = hash21(ix + 1.0, iy)
    let c = hash21(ix, iy + 1.0)
    let d = hash21(ix + 1.0, iy + 1.0)
    return mix(mix(a, b, fx) as f32, mix(c, d, fx) as f32, fy) as f32

// Four-octave fractional Brownian motion over the value noise.
fn nebula_fbm(px f32, py f32) f32
    var s = 0.0
    var a = 0.5
    var qx = px
    var qy = py
    var i = 0
    while i < 4
        s = s + a * vnoise(qx, qy)
        qx = qx * 2.02
        qy = qy * 2.02
        a = a * 0.5
        i = i + 1
    return s as f32

// One octave of the volumetric starfield: sparse bright points per grid cell.
fn star_layer(dx f32, dy f32, dz f32, scale f32, thr f32) f32
    let px = dx * scale
    let py = dy * scale
    let pz = dz * scale
    let idx = floor(px) as f32
    let idy = floor(py) as f32
    let idz = floor(pz) as f32
    let fx = (px - idx) - 0.5
    let fy = (py - idy) - 0.5
    let fz = (pz - idz) - 0.5
    let h = hash31(idx, idy, idz)
    let offx = hash31(idx + 1.3, idy + 1.3, idz + 1.3) - 0.5
    let offy = hash31(idx + 2.7, idy + 2.7, idz + 2.7) - 0.5
    let offz = hash31(idx + 4.1, idy + 4.1, idz + 4.1) - 0.5
    let sx = fx - offx * 0.7
    let sy = fy - offy * 0.7
    let sz = fz - offz * 0.7
    let len = sqrt(sx * sx + sy * sy + sz * sz) as f32
    let star = smoothstep(0.5, 0.0, len) as f32
    let on = step(1.0 - thr, h) as f32
    return star * star * on * (0.4 + 0.6 * (fract(h * 91.7) as f32))

// Accretion-disk opacity ramp: fades in past the inner edge, out toward the rim.
fn disk_opacity(r f32) f32
    let a = smoothstep(2.6, 3.1, r) as f32
    let b = smoothstep(6.5, 9.0, r) as f32
    return 0.9 * a * (1.0 - b)

// Seed: starting pose (near edge-on for the iconic warped-disk view), idle past
// the auto-rotate threshold so the camera drifts until the pointer takes over.
// The distance is set so the lensed disk's widest extent — the sheared outer rim
// at r = 9, magnified by the lensing — lands inside the horizontal field of view
// with margin to spare, since the page scales the 16:9 render to cover frames of
// other aspect ratios and crops the overflow.
forall i in 0..4
    cam_a[0] = 1.15
    cam_a[1] = 0.18
    cam_a[2] = 26.0
    cam_a[3] = 99.0

// Clear the canvas. This 2-D pass over the exact display extent also tells the
// web-gpu backend the canvas is 1280×720 (a flat paint buffer reads as square).
forall px, py in 0..CW, 0..CH
    let base = (py * CW + px) * 4
    paint[base] = 0.0
    paint[base + 1] = 0.0
    paint[base + 2] = 0.0
    paint[base + 3] = 1.0

gpu frame
    // Pass 1: integrate orbit angles, zoom distance, and idle timer from input.
    forall i in 0..4
        let yaw = cam_a[0]
        let pitch = cam_a[1]
        let dist = cam_a[2]
        let idle = cam_a[3]
        let dragging = 1.0 if frame.mouse_down else 0.0
        let next_idle = 0.0 if dragging > 0.5 else idle + frame.dt
        // After a brief idle the camera auto-orbits; drag always steers it.
        let auto = frame.dt * 0.08 if next_idle > 0.1 else 0.0
        cam_b[0] = yaw + frame.drag_dx * 0.008 + auto
        // `drag_dy` points up the screen while the pointer's travel points down,
        // so raising the pointer must raise the camera: subtract, matching the
        // reference's `rot.y += (clientY - lastY) * 0.006`.
        cam_b[1] = clamp(pitch - frame.drag_dy * 0.006, -1.45, 1.45) as f32
        cam_b[2] = clamp(dist * (exp(frame.wheel * 0.0012) as f32), 4.5, 34.0) as f32
        cam_b[3] = next_idle

    // Pass 2: trace one photon geodesic per pixel and composite the result.
    forall idx in 0..PIXELS
        let px = idx % CW
        let py = idx / CW
        let tm = frame.time
        // Aspect-corrected screen coords: widen x for the 16:9 frame (matches the
        // reference shader's `uv.x *= res.x / res.y`).
        let uvx = ((px as f32) / (CW as f32) - 0.5) * ((CW as f32) / (CH as f32))
        let uvy = 0.5 - (py as f32) / (CH as f32)

        // Camera orbiting the hole on a sphere of radius `dist`.
        let ya = cam_b[0]
        let pa = clamp(cam_b[1], -1.45, 1.45) as f32
        let dist = cam_b[2]
        let cpa = cos(pa) as f32
        let rox = (cos(ya) as f32) * cpa * dist
        let roy = (sin(pa) as f32) * dist
        let roz = (sin(ya) as f32) * cpa * dist
        // Forward / right / up basis (right = cross(fw, up0), up = cross(rt, fw)).
        let flen = sqrt(rox * rox + roy * roy + roz * roz) as f32
        let fwx = (0.0 - rox) / flen
        let fwy = (0.0 - roy) / flen
        let fwz = (0.0 - roz) / flen
        let rl = sqrt(fwz * fwz + fwx * fwx) as f32
        let rtx = (0.0 - fwz) / rl
        let rtz = fwx / rl
        let upx = 0.0 - rtz * fwy
        let upy = rtz * fwx - rtx * fwz
        let upz = rtx * fwy
        let rdx0 = fwx * 1.7 + rtx * uvx + upx * uvy
        let rdy0 = fwy * 1.7 + upy * uvy
        let rdz0 = fwz * 1.7 + rtz * uvx + upz * uvy
        let rdl = sqrt(rdx0 * rdx0 + rdy0 * rdy0 + rdz0 * rdz0) as f32
        let rdx = rdx0 / rdl
        let rdy = rdy0 / rdl
        let rdz = rdz0 / rdl

        // Photon orbital plane {e1 = radial, e2 = forward-tangent}.
        let r0 = flen
        let e1x = rox / r0
        let e1y = roy / r0
        let e1z = roz / r0
        let nrmx0 = e1y * rdz - e1z * rdy
        let nrmy0 = e1z * rdx - e1x * rdz
        let nrmz0 = e1x * rdy - e1y * rdx
        let nl = sqrt(nrmx0 * nrmx0 + nrmy0 * nrmy0 + nrmz0 * nrmz0) as f32
        let usey = 1.0 if nl < 0.0001 else 0.0
        let nrmx = 0.0 if usey > 0.5 else nrmx0 / nl
        let nrmy = 1.0 if usey > 0.5 else nrmy0 / nl
        let nrmz = 0.0 if usey > 0.5 else nrmz0 / nl
        let e2x0 = nrmy * e1z - nrmz * e1y
        let e2y0 = nrmz * e1x - nrmx * e1z
        let e2z0 = nrmx * e1y - nrmy * e1x
        let e2l = sqrt(e2x0 * e2x0 + e2y0 * e2y0 + e2z0 * e2z0) as f32
        var e2x = e2x0 / e2l
        var e2y = e2y0 / e2l
        var e2z = e2z0 / e2l
        let de2 = rdx * e2x + rdy * e2y + rdz * e2z
        let flip = -1.0 if de2 < 0.0 else 1.0
        e2x = e2x * flip
        e2y = e2y * flip
        e2z = e2z * flip

        var u = 1.0 / r0
        let vr = rdx * e1x + rdy * e1y + rdz * e1z
        let vt = rdx * e2x + rdy * e2y + rdz * e2z
        var du = (0.0 - u) * vr / (max(vt, 0.001) as f32)

        var phi = 0.0
        var r = r0
        var col_r = 0.0
        var col_g = 0.0
        var col_b = 0.0
        var trans = 1.0
        var prev_y = roy
        var prev_r = r0
        var prev_px = rox
        var prev_py = roy
        var prev_pz = roz
        var captured = 0.0
        var pos_x = rox
        var pos_y = roy
        var pos_z = roz

        var i = 0
        while i < 260
            let dphi = mix(0.028, 0.11, clamp((r - 3.0) / 14.0, 0.0, 1.0)) as f32
            let acc = 0.0 - u + 1.5 * u * u
            du = du + acc * dphi
            u = u + du * dphi
            phi = phi + dphi
            r = 1.0 / (max(u, 0.0001) as f32)
            let cph = cos(phi) as f32
            let sph = sin(phi) as f32
            pos_x = r * (cph * e1x + sph * e2x)
            pos_y = r * (cph * e1y + sph * e2y)
            pos_z = r * (cph * e1z + sph * e2z)

            if r < 1.02
                captured = 1.0
                i = 260
            else
                // Equatorial-plane crossing → sample the disk (front-to-back).
                if prev_y * pos_y < 0.0
                    let f = prev_y / (prev_y - pos_y)
                    let cr = mix(prev_r, r, f) as f32
                    if cr > 2.6 and cr < 9.0
                        let x = clamp((cr - 2.6) / 6.4, 0.0, 1.0) as f32
                        let cpx = mix(prev_px, pos_x, f) as f32
                        let cpz = mix(prev_pz, pos_z, f) as f32
                        let ang = atan2(cpz, cpx) as f32
                        let omega = pow(2.6 / cr, 1.5) as f32
                        let rot = tm * omega * 0.9
                        let bnd = nebula_fbm((ang - rot) * 2.5, cr * 1.3)
                        let bands = 0.35 + 1.25 * bnd * bnd
                        // Temperature: hot white-orange inner → cool red outer.
                        let sm1 = smoothstep(0.0, 0.35, x) as f32
                        let sm2 = smoothstep(0.35, 1.0, x) as f32
                        // Temperature blend: hot white-orange → mid orange → cool red.
                        let cr0 = mix(mix(1.0, 1.0, sm1) as f32, 0.85, sm2) as f32
                        let cg0 = mix(mix(0.92, 0.55, sm1) as f32, 0.2, sm2) as f32
                        let cb0 = mix(mix(0.78, 0.18, sm1) as f32, 0.05, sm2) as f32
                        // Relativistic Doppler beaming along the orbit tangent.
                        let travx = pos_x - prev_px
                        let travy = pos_y - prev_py
                        let travz = pos_z - prev_pz
                        let tl = sqrt(travx * travx + travy * travy + travz * travz) as f32
                        let tvx = travx / tl
                        let tvz = travz / tl
                        let odl = sqrt(cpz * cpz + cpx * cpx) as f32
                        let odx = cpz / odl
                        let odz = (0.0 - cpx) / odl
                        let beta = min(sqrt(0.5 / cr) as f32, 0.75) as f32
                        let gg = 1.0 / (sqrt(1.0 - beta * beta) as f32)
                        let ddot = odx * beta * tvx + odz * beta * tvz
                        let delta = 1.0 / (gg * (1.0 - ddot))
                        let boost = pow(clamp(delta, 0.2, 3.0) as f32, 3.0) as f32
                        let tint = clamp((delta - 0.85) * 1.2, 0.0, 1.0) as f32
                        let ccr = cr0 * (mix(1.0, 0.75, tint) as f32)
                        let ccg = cg0 * (mix(0.55, 0.85, tint) as f32)
                        let ccb = cb0 * (mix(0.35, 1.2, tint) as f32)
                        // Gravitational redshift dimming near the horizon.
                        let grav = sqrt(clamp(1.0 - 1.0 / cr, 0.04, 1.0) as f32) as f32
                        let bright = mix(1.7, 0.5, x) as f32
                        let em = bands * bright * boost * grav
                        col_r = col_r + trans * ccr * em
                        col_g = col_g + trans * ccg * em
                        col_b = col_b + trans * ccb * em
                        trans = trans * (1.0 - disk_opacity(cr))
                // Stop once the ray has escaped to infinity or the disk has gone
                // fully opaque, leaving `prev_*` one step behind `pos_*` so the
                // escaped-sky direction below is a real displacement, not zero.
                if r > 60.0 or trans < 0.01
                    i = 260
                else
                    prev_y = pos_y
                    prev_r = r
                    prev_px = pos_x
                    prev_py = pos_y
                    prev_pz = pos_z
            i = i + 1

        // Rays that never crossed the horizon escape to the lensed sky.
        if captured < 0.5
            let dnx = pos_x - prev_px
            let dny = pos_y - prev_py
            let dnz = pos_z - prev_pz
            // Photon-ring rays that neither escaped nor were captured within the
            // step budget end with `pos_* == prev_*`; floor the length so the
            // direction stays finite instead of dividing by zero.
            let dl = max(sqrt(dnx * dnx + dny * dny + dnz * dnz) as f32, 0.0001) as f32
            let dirx = dnx / dl
            let diry = dny / dl
            let dirz = dnz / dl
            let s1 = star_layer(dirx, diry, dirz, 14.0, 0.05)
            let s2 = star_layer(dirx, diry, dirz, 30.0, 0.035)
            let s3 = star_layer(dirx, diry, dirz, 60.0, 0.02)
            let nb = nebula_fbm((dirx * 2.0 + dirz) * 1.4 + 5.0, (diry * 2.0 - dirz) * 1.4 + 5.0)
            let bgr = 0.85 * s1 + 1.0 * s2 * 0.8 + 0.7 * s3 * 0.5 + 0.06 * nb * nb + 0.007
            let bgg = 0.92 * s1 + 0.9 * s2 * 0.8 + 0.8 * s3 * 0.5 + 0.03 * nb * nb + 0.0105
            let bgb = 1.0 * s1 + 0.78 * s2 * 0.8 + 1.0 * s3 * 0.5 + 0.11 * nb * nb + 0.0175
            col_r = col_r + trans * bgr
            col_g = col_g + trans * bgg
            col_b = col_b + trans * bgb

        // Tone map (Reinhard + gamma), then write opaque RGBA.
        col_r = col_r * 1.05
        col_g = col_g * 1.05
        col_b = col_b * 1.05
        col_r = col_r / (col_r + 0.85)
        col_g = col_g / (col_g + 0.85)
        col_b = col_b / (col_b + 0.85)
        let base = idx * 4
        paint[base] = pow(max(col_r, 0.0) as f32, 0.82) as f32
        paint[base + 1] = pow(max(col_g, 0.0) as f32, 0.82) as f32
        paint[base + 2] = pow(max(col_b, 0.0) as f32, 0.82) as f32
        paint[base + 3] = 1.0