#include <metal_stdlib>
using namespace metal;

struct Uniforms {
    float2 resolution;
    float  time;
    float  timeDelta;
    float2 mouse;
    uint   frame;
    uint   pad;
};

struct Params {
    float gravity;
    float tilt;
    float spin;
    float doppler;
    float disk;
    float heat;
    float volume;
};

struct RKVertexOut {
    float4 position [[position]];
    float2 uv;
};

vertex RKVertexOut rk_vertex(uint vid [[vertex_id]]) {
    float2 p = float2((vid << 1) & 2, vid & 2);
    RKVertexOut o;
    o.position = float4(p * 2.0 - 1.0, 0.0, 1.0);
    o.uv = p;
    return o;
}

#line 1 "shader.metal"
// Ray Kiln · Black Hole
//
// @episode Black Hole
// @length 60
// @variation 30
// @variations 1,2,3,7,8,9,11,12
// @still 22.0
// @teaches Rays bent by gravity, an accretion disk, Doppler colour
// @category Fragment
// @tags raymarching, black holes, gravitational lensing, doppler, accretion disk, value noise, stars, procedural sound
// @sweeps gravity, tilt, doppler
// @final Black hole | A disk of fire wraps over and under the shadow it feeds, the near side rushing towards you white-hot, the far side falling away in ember, the stars behind it bent into a ring.
//
// A black hole with a disk of hot gas around it, seen from close by. Every
// pixel's ray is walked outwards step by step, and at every step gravity
// bends it towards the hole: a ray that comes near enough circles it, and one
// that comes nearer falls in, which is the black shadow in the middle. The
// disk is in one plane; where the ray crosses that plane it picks up the
// disk's light, and because the rays behind the hole are bent up and over it,
// the far side of the disk shows above and below the shadow. The gas orbits
// at close to the speed of light, so the side coming towards you is beamed
// bright and white-hot and the side going away is dim and ember: the Doppler
// effect, and the one asymmetry in an otherwise symmetric picture.
// Each @param line below becomes a slider in the host app and a row in the page's parameters table:
//   // @param name default min max description

// @param gravity 1 0 1.6 How strongly the hole bends light: 0 is straight rays, 1 is Einstein's
// @param tilt 0.14 0 0.5 How far the camera looks down on the disk: 0 is edge on, 0.5 from above
// @param spin 1 0 3 How fast the gas goes round: at 1 the inner edge orbits once in 10 s
// @param doppler 1 0 1 How much the gas's speed beams the near side and dims the far side
// @param disk 1 0 1.5 How bright the disk burns
// @param heat 0.5 0 1 Slides the palette from deep ember to white-hot
// @param volume 0.8 0 1 Loudness of the sound

constant float TAU = 6.2831853;
// Distances are in Schwarzschild radii: the horizon is at r = 1 and the
// disk runs from r = 3 out to r = 9.
constant float HORIZON = 1.0;
constant float DISK_IN = 3.0;
constant float DISK_OUT = 9.0;
// How far the camera is from the hole.
constant float DISTANCE = 19.0;
// The most steps a ray is walked, and how far out counts as escaped.
constant int STEPS = 200;
constant float FAR = 28.0;
// At spin 1 the disk's inner edge goes round once every ten seconds.
constant float ORBIT = 3.27;

// Cracked Clay's hash: a whole number in, a random 0 to 1 out.
float hashU(uint n) {
    n = (n << 13u) ^ n;
    n = n * (n * n * 15731u + 789221u) + 1376312589u;
    return float(n & 0x7fffffffu) / float(0x7fffffff);
}

// Flowing Color Field's palette: near-black, ember, orange, gold, white-hot.
float3 kiln(float t) {
    float3 c0 = float3(0.020, 0.012, 0.020);
    float3 c1 = float3(0.280, 0.040, 0.030);
    float3 c2 = float3(0.880, 0.260, 0.050);
    float3 c3 = float3(1.000, 0.680, 0.200);
    float3 c4 = float3(1.000, 0.970, 0.840);
    t = clamp(t, 0.0, 1.0);
    float3 c = mix(c0, c1, smoothstep(0.00, 0.22, t));
    c = mix(c, c2, smoothstep(0.22, 0.50, t));
    c = mix(c, c3, smoothstep(0.50, 0.76, t));
    return mix(c, c4, smoothstep(0.76, 1.00, t));
}

// A random 0 to 1 for each whole-numbered point in space.
float hash3(float3 c) {
    uint x = uint(int(c.x) + 4096);
    uint y = uint(int(c.y) + 4096);
    uint z = uint(int(c.z) + 4096);
    return hashU(x * 73856093u ^ y * 19349663u ^ z * 83492791u);
}

// Value noise in three dimensions: random values at the whole-numbered
// points, blended across the cube between eight of them.
float valueNoise3(float3 x) {
    float3 i = floor(x);
    float3 f = x - i;
    f = f * f * (3.0 - 2.0 * f);
    float a = mix(hash3(i), hash3(i + float3(1, 0, 0)), f.x);
    float b = mix(hash3(i + float3(0, 1, 0)), hash3(i + float3(1, 1, 0)), f.x);
    float c = mix(hash3(i + float3(0, 0, 1)), hash3(i + float3(1, 0, 1)), f.x);
    float d = mix(hash3(i + float3(0, 1, 1)), hash3(i + float3(1, 1, 1)), f.x);
    return mix(mix(a, b, f.y), mix(c, d, f.y), f.z);
}

// Fractal noise: `octaves` layers of value noise, each twice as fine and half
// as strong as the one before, shifted so their grids don't line up. 0 to 1.
float fbm3(float3 x, int octaves) {
    float sum = 0.0, amp = 0.5, total = 0.0;
    for (int k = 0; k < octaves; k++) {
        sum += amp * valueNoise3(x);
        total += amp;
        x = x * 2.03 + float3(1.7, 9.2, 4.1);
        amp *= 0.5;
    }
    return sum / total;
}

// The stars behind everything, from the direction a ray finally leaves in:
// a grid of cells over the sky, a star in some, a few of them bright.
float3 stars(float3 dir) {
    float2 sph = float2(atan2(dir.z, dir.x), asin(clamp(dir.y, -1.0, 1.0)));
    sph *= 90.0;
    float2 cell = floor(sph);
    uint k = uint(int(cell.x) + 1000) * 7919u;
    k += uint(int(cell.y) + 1000) * 104729u;
    float2 at = float2(hashU(k), hashU(k + 1u)) * 0.8 + 0.1;
    float star = step(0.84, hashU(k + 2u));
    star *= smoothstep(0.11, 0.0, length(fract(sph) - at));
    float bright = pow(hashU(k + 3u), 3.0);
    float3 tint = mix(float3(0.7, 0.8, 1.0), kiln(0.85), hashU(k + 4u));
    // A faint band of dust across the sky, so the lensing has something
    // to bend.
    float band = exp(-dir.y * dir.y * 18.0) * 0.15;
    band *= 0.4 + 0.6 * fbm3(dir * 6.0, 3);
    return kiln(0.45) * band + tint * star * (0.3 + 2.0 * bright);
}

// How fast the gas goes round at radius r, in turns a second at spin 1:
// faster nearer in, as the planets are.
float orbitRate(float r) {
    return ORBIT / (r * sqrt(r));
}

// The disk's glow where a ray crosses its plane at q. Streaks of noise are
// carried round at the orbit's rate for that radius, so the inner edge
// shears past the outer gas, and it burns brightest just outside the edge.
float diskGlow(float3 q, float t, constant Params& p) {
    float r = length(q.xz);
    if (r < DISK_IN || r > DISK_OUT) return 0.0;
    float phi = atan2(q.z, q.x) - orbitRate(r) * p.spin * t;
    // Noise on a cylinder, so there is no seam where phi wraps round.
    float3 at = float3(cos(phi) * 1.6, sin(phi) * 1.6, r * 2.2);
    float n = fbm3(at, 4);
    float streak = 0.35 + 0.65 * smoothstep(0.3, 0.8, n);
    // Brightest just outside the inner edge, fading as the square of r.
    float profile = 1.0 - sqrt(DISK_IN / r) * 0.9;
    profile *= (DISK_IN * DISK_IN) / (r * r);
    float edge = smoothstep(DISK_IN, DISK_IN + 0.4, r);
    edge *= smoothstep(DISK_OUT, DISK_OUT - 2.5, r);
    return streak * profile * edge * 24.0;
}

// The Doppler factor for gas at q seen along the ray direction d: above 1
// where the gas comes towards you, below 1 where it goes away. The gas
// orbits at a fraction beta of the speed of light; the light travels
// back along -d.
float dopplerAt(float3 q, float3 d, constant Params& p) {
    float r = length(q.xz);
    float beta = sqrt(0.5 / r);
    float3 along = normalize(float3(-q.z, 0.0, q.x));
    float gamma = 1.0 / sqrt(1.0 - beta * beta);
    float shift = 1.0 / (gamma * (1.0 + beta * dot(along, d)));
    // Climbing out of the hole's gravity costs the light some energy too.
    float well = sqrt(1.0 - HORIZON / r);
    return mix(1.0, shift, p.doppler) * well;
}

// Walks a ray from the camera, bent by the hole, and gathers what it sees:
// the disk where it crosses the disk's plane, black if it falls into the
// hole, the stars if it escapes.
float3 trace(float3 ro, float3 rd, float t, constant Params& p) {
    float3 pos = ro;
    float3 vel = rd;
    // Angular momentum about the hole, which the bending conserves.
    float3 L = cross(pos, vel);
    float h2 = dot(L, L) * p.gravity;
    float3 light = float3(0.0);
    float carry = 1.0;
    bool escaped = false;
    for (int i = 0; i < STEPS; i++) {
        float r2 = dot(pos, pos);
        if (r2 < HORIZON * HORIZON) break;
        if (r2 > FAR * FAR && dot(pos, vel) > 0.0) { escaped = true; break; }
        // Smaller steps near the hole, where the bending is strong.
        float dt = clamp(0.045 * sqrt(r2), 0.03, 0.5);
        float3 next = pos + vel * dt;
        // Crossing the disk's plane: find where, and take its light.
        if (pos.y * next.y < 0.0) {
            float f = pos.y / (pos.y - next.y);
            float3 q = mix(pos, next, f);
            float glow = diskGlow(q, t, p);
            if (glow > 0.0) {
                float dop = dopplerAt(q, vel, p);
                float seen = glow * dop * dop * dop * p.disk;
                light += carry * seen;
                // The gas is not quite opaque: a little of what is behind
                // shows through.
                carry *= 0.35;
            }
        }
        // Gravity: the pull on a ray of light in Schwarzschild's geometry,
        // felt as a bend towards the hole that grows fast as r shrinks.
        vel += -1.5 * h2 * pos / (r2 * r2 * sqrt(r2)) * dt;
        vel = normalize(vel);
        pos = next;
    }
    if (escaped) light += carry * stars(vel);
    return light;
}

// The camera: on a slow drift around the hole, looking at it, tilted up out
// of the disk's plane as far as the slider says.
float3 cameraPos(float t, constant Params& p) {
    float a = 0.35 + 0.04 * t;
    float up = sin(p.tilt * TAU * 0.25);
    float flat = cos(p.tilt * TAU * 0.25);
    return DISTANCE * float3(flat * sin(a), up, flat * cos(a));
}

float4 shade(float2 uv, constant Uniforms& u, constant Params& p) {
    float t = u.time;
    float side = min(u.resolution.x, u.resolution.y);
    float2 q = (uv - 0.5) * u.resolution / side * 2.0;
    float3 ro = cameraPos(t, p);
    float3 fw = normalize(-ro);
    float3 rt = normalize(cross(fw, float3(0.0, 1.0, 0.0)));
    float3 up = cross(rt, fw);
    float3 rd = normalize(q.x * rt + q.y * up + 1.7 * fw);
    float3 seen = trace(ro, rd, t, p);
    // The light through the palette: its amount picks the heat, and what
    // colour it had (a blue star, the white of beamed gas) tints it.
    float lum = dot(seen, float3(0.3, 0.5, 0.2));
    float heat = 1.0 - exp(-lum * (0.7 + 0.8 * p.heat));
    float3 col = kiln(heat);
    col *= mix(float3(1.0), seen / max(lum, 1e-3), 0.25);
    return float4(clamp(col, 0.0, 1.0), 1.0);
}

// ---- Sound -----------------------------------------------------------------------------------------------------
// The sound is part of the shader: a function of time, run on the GPU 48,000 times a second and put in the video.
// Two voices orbit the hole as the gas does, a low one at the inner edge of
// the disk and a higher one farther out, each going round at its own
// Kepler rate. Each is Doppler shifted as it comes and goes, up in pitch and
// louder as it swings towards you, down and quieter as it swings away, by the
// same factor that beams the picture. Under them the hole hums, a deep slow
// swell, and now and then a star glints.

// A value that drifts smoothly: random values `rate` times a second, blended.
float noise1(float t, float rate, uint salt) {
    float x = t * rate;
    float i = floor(x);
    float f = x - i;
    f = f * f * (3.0 - 2.0 * f);
    uint n = uint(int(i) + 1000000);
    return mix(hashU(n * 31u + salt), hashU((n + 1u) * 31u + salt), f);
}

// One orbiting voice: a tone at `f0` whose pitch rides the Doppler shift of
// gas at radius r. The shift is 1 + beta cos(angle) to a first order, and
// the phase is the integral of the frequency, so the pitch glides smoothly
// and the sound stays a pure function of time.
float2 orbitVoice(float t, float f0, float r, float phase0,
                  constant Params& p) {
    float beta = sqrt(0.5 / r) * p.doppler;
    float w = TAU * orbitRate(r) * p.spin / TAU;
    float ang = w * t + phase0;
    float ph = TAU * f0 * (t + (w > 1e-4 ? beta * sin(ang) / w : 0.0));
    // Louder coming towards you, as the picture is, by the cube of the shift.
    float shift = 1.0 + beta * cos(ang);
    float loud = shift * shift * shift;
    float tone = sin(ph) + 0.5 * sin(2.0 * ph + 0.3) + 0.2 * sin(3.0 * ph);
    // Which ear it is nearer to, as it goes round.
    float pan = 0.5 + 0.4 * sin(ang);
    return float2(1.0 - pan, pan) * tone * loud;
}

float2 sound(float t, constant Params& p) {
    float2 v = float2(0.0);
    v += orbitVoice(t, 82.41, DISK_IN, 0.0, p) * 0.22;
    v += orbitVoice(t, 123.47, 5.0, 2.1, p) * 0.14;
    v += orbitVoice(t, 164.81, 7.5, 4.0, p) * 0.09;
    // The hum of the hole: a low E and its octave, swelling slowly, a little
    // apart in each ear.
    float swell = 0.5 + 0.5 * sin(TAU * t / 13.0 + 2.0 * noise1(t, 0.1, 7u));
    float2 hum = float2(sin(TAU * 41.2 * t), sin(TAU * 41.3 * t + 1.0));
    hum += 0.5 * float2(sin(TAU * 82.4 * t + 0.5), sin(TAU * 82.6 * t));
    v += hum * 0.13 * (0.4 + 0.6 * swell) * p.disk;
    // Star glints: a short high ping every so often, somewhere in the stereo.
    float slot = floor(t / 1.7);
    for (int k = -1; k <= 0; k++) {
        float j = slot + float(k);
        if (j < 0.0) continue;
        uint id = uint(j) * 5u + 11u;
        if (hashU(id) > 0.45) continue;
        float s = t - j * 1.7 - 1.2 * hashU(id + 1u);
        if (s < 0.0) continue;
        float f = 1760.0 * exp2(floor(hashU(id + 2u) * 5.0) * 2.0 / 12.0);
        float ping = sin(TAU * f * s) * exp(-s * 9.0) * (1.0 - exp(-s * 400.0));
        float pan = hashU(id + 3u);
        v += float2(1.0 - pan, pan) * ping * 0.05;
    }
    float fadeIn = smoothstep(0.0, 2.0, t);
    return tanh(v * 1.3) / 1.5 * fadeIn * p.volume;
}

#line 1 "rk_host"
fragment float4 rk_fragment(RKVertexOut in [[stage_in]],
                                constant Uniforms& u [[buffer(0)]],
                                constant Params& p [[buffer(1)]]) {
    return shade(in.uv, u, p);
}

struct SoundUniforms { float start; float rate; uint count; uint pad; };
kernel void rk_sound(device float2* out [[buffer(0)]],
                         constant SoundUniforms& su [[buffer(1)]],
                         constant Params& p [[buffer(2)]],
                         uint i [[thread_position_in_grid]]) {
    if (i >= su.count) return;
    out[i] = clamp(sound(su.start + float(i) / su.rate, p), -1.0, 1.0);
}