#include <metal_stdlib>
using namespace metal;

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

struct Params {
    float speed;
    float wind;
    float thickness;
    float puffs;
    float forward;
    float glow;
    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 · Smoke Stack
//
// @episode Smoke Stack
// @length 60
// @variation 30
// @variations 1,2,4,5,7,8,9,12
// @still 27.0
// @teaches Volumetric clouds: light marching, self-shadowing, scattering
// @category Fragment
// @tags volumetric rendering, light marching, self-shadowing, phase function, transmittance, procedural sound
// @final Smoke stack | The kiln's chimney at night, its smoke lit silver by the moon behind it and orange by the fire below, puffing harder each time the kiln is stoked.
//
// The kiln's chimney at night, and the smoke rising out of it. Smoke has no surface to find: it is a density at
// every point, so the ray walks through it in even steps and adds up what each bit does to the light. Some light is
// taken out on the way (absorbed or sent elsewhere), so what lies behind thick smoke fades. Some is sent toward the eye,
// and how much depends on how much light reaches that bit of smoke, which takes a second, shorter march toward the
// light: the smoke shadows itself. And it depends on the angle: smoke sends most light on in the direction it was going,
// so the edges glow where the moon is behind them.
// 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 speed 1.0 0.0 2.0 How fast the smoke rises and the kiln is stoked
// @param wind 0.5 0.0 1.0 How hard the wind bends the plume
// @param thickness 0.6 0.1 1.0 How dense the smoke is
// @param puffs 0.5 0.0 1.0 From a steady stream to the kiln breathing out in puffs
// @param forward 0.6 0.0 0.9 How strongly the smoke sends light on the way it was going (the silver lining)
// @param glow 0.6 0.0 1.0 How bright the fire under the smoke burns
// @param volume 0.8 0.0 1.0 Loudness of the sound

constant float TAU = 6.2831853;
constant float PI = 3.1415927;
// How fast smoke rises, in units a second, and the chimney's radius.
constant float RISE = 0.45;
constant float STACK = 0.3;

// The kiln palette from Flowing Color Field.
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));
}

// Kiln Fire's 3D hash, value noise and fractal noise.
float hash31(float3 p) {
    uint3 v = uint3(int3(floor(p))) * 1664525u + 1013904223u;
    v.x += v.y * v.z; v.y += v.z * v.x; v.z += v.x * v.y;
    v ^= v >> 16u;
    v.x += v.y * v.z; v.y += v.z * v.x; v.z += v.x * v.y;
    return float(v.x) * (1.0 / 4294967295.0);
}

float vnoise3(float3 p) {
    float3 i = floor(p), f = fract(p);
    f = f * f * (3.0 - 2.0 * f);
    float a = hash31(i),                     b = hash31(i + float3(1, 0, 0));
    float c = hash31(i + float3(0, 1, 0)),   d = hash31(i + float3(1, 1, 0));
    float e = hash31(i + float3(0, 0, 1)),   g = hash31(i + float3(1, 0, 1));
    float h = hash31(i + float3(0, 1, 1)),   k = hash31(i + float3(1, 1, 1));
    float lo = mix(mix(a, b, f.x), mix(c, d, f.x), f.y);
    float hi = mix(mix(e, g, f.x), mix(h, k, f.x), f.y);
    return mix(lo, hi, f.z);
}

// Octaves of vnoise3, as Kiln Fire's fbm3, but with the number of octaves
// given: the light march gets by with fewer.
float fbm3(float3 p, int octaves) {
    float v = 0.0, a = 0.5;
    for (int i = 0; i < octaves; i++) {
        v += a * vnoise3(p);
        p = p * 2.02 + float3(17.1, 9.2, 4.7);
        a *= 0.5;
    }
    return v;
}

// ---- The smoke ---------------------------------------------------------------------------------------------------

// How hard the kiln breathes out at time t: a little all the time, a
// stoking every 20 seconds that sends up a thick column for a while,
// and puffs, every two or three seconds, on top.
float emission(float t, constant Params& p) {
    float since = fract(t / 20.0 + 0.6) * 20.0;
    float stoke = smoothstep(0.0, 1.0, since) * exp(-since * 0.22);
    float puff = 0.5 + 0.5 * sin(t * 2.4 + 1.3 * sin(t * 0.9));
    puff = mix(1.0, puff * puff * 1.6, p.puffs);
    return (0.35 + 1.1 * stoke) * puff;
}

// The wind at time t: it swings slowly, so the plume leans further over and
// back across the minute.
float windAt(float t, constant Params& p) {
    return p.wind * (0.55 + 0.45 * sin(t * 0.11 - 0.6));
}

// Smoke per unit of length at a point: 0 for clear air. The smoke at
// height h left the chimney h / RISE seconds ago, so it carries the kiln's
// breath from then; the wind carries it sideways the higher it goes, and it
// spreads and thins as it rises. Noise drifting up with it frays the edges.
float smokeAt(float3 pos, float t, constant Params& p, int octaves) {
    float h = pos.y;
    if (h < 0.0) return 0.0;
    float left = t - h / RISE;
    float w = windAt(left, p);
    float3 axis = float3(w * h * h * 0.32 + w * h * 0.2, h, 0.0);
    float r = STACK + 0.4 * h;
    float3 q = pos - float3(w * t * 0.5, t * RISE, 0.0);
    float n = fbm3(q * 1.7, octaves);
    float edge = 1.0 - length((pos - axis).xz) / r;
    float d = edge * 1.3 + (n - 0.5) * 2.2 - 0.1;
    d = clamp(d, 0.0, 1.0) * emission(left, p);
    d *= smoothstep(0.0, 0.15, h) * exp(-h * 0.28);
    return d * 10.0 * p.thickness;
}

// ---- Light -------------------------------------------------------------------------------------------------------

// The moon: behind the smoke as the camera sees it, up and a little left.
constant float3 MOON = float3(-0.06, 0.55, 0.833);
constant float3 MOONLIGHT = float3(0.72, 0.80, 1.0);

// The fire in the chimney's mouth, glowing up into the smoke.
constant float3 FIRE = float3(0.0, -0.05, 0.0);

// Henyey-Greenstein's phase function: the share of scattered light sent at
// cosine c from its way. g = 0: every way alike; toward 1, straight on.
float phase(float c, float g) {
    float k = 1.0 + g * g - 2.0 * g * c;
    return (1.0 - g * g) / (4.0 * PI * k * sqrt(k));
}

// The light from a direction that reaches a point: a short march toward
// it adding up the smoke, then exp(-sum), Beer and Lambert's law.
float lightThrough(float3 pos, float3 dir, float t, constant Params& p,
                   float len, int steps) {
    float ds = len / float(steps);
    float sum = 0.0;
    for (int i = 0; i < steps; i++) {
        sum += smokeAt(pos + dir * (ds * (float(i) + 0.5)), t, p, 2);
    }
    return exp(-sum * ds);
}

// The night sky: dark blue overhead, warm haze low down from the kiln yard's
// lights, stars, and the moon with a glow round it.
float3 nightSky(float3 rd, float t) {
    float up = clamp(rd.y, 0.0, 1.0);
    float3 col = mix(float3(0.10, 0.045, 0.040), float3(0.012, 0.016, 0.040),
                     smoothstep(-0.05, 0.45, rd.y));
    float2 a = float2(atan2(rd.x, rd.z), rd.y) * 90.0;
    float star = hash31(float3(floor(a), 1.0));
    float2 at = fract(a) - 0.5;
    col += step(0.985, star) * up * smoothstep(0.12, 0.0, length(at)) * 0.7
           * (0.75 + 0.25 * sin(t * 3.0 + star * 60.0));
    float m = dot(rd, MOON);
    col += MOONLIGHT * smoothstep(0.9993, 0.9995, m) * 1.4;
    m = max(m, 0.0);
    col += MOONLIGHT * (pow(m, 300.0) * 0.25 + pow(m, 12.0) * 0.05);
    return col;
}

// ---- The chimney -------------------------------------------------------------------------------------------------

// Where a ray meets the outside of the chimney, a cylinder round the y axis
// up to y = 0, or -1.
float hitStack(float3 ro, float3 rd) {
    float a = dot(rd.xz, rd.xz);
    float b = dot(ro.xz, rd.xz);
    float c = dot(ro.xz, ro.xz) - STACK * STACK;
    float disc = b * b - a * c;
    if (disc < 0.0) return -1.0;
    float d = (-b - sqrt(disc)) / a;
    return (d > 0.0 && ro.y + rd.y * d < 0.0) ? d : -1.0;
}

// The chimney's bricks at a point on it, lit by the moon on one side and,
// near the top, by the fire spilling over the rim.
float3 brick(float3 pos, float glow) {
    float3 n = normalize(float3(pos.x, 0.0, pos.z));
    float2 b = float2(atan2(pos.x, pos.z) * STACK / 0.11, pos.y / 0.06);
    b.x += 0.5 * fmod(floor(b.y), 2.0);
    float2 f = fract(b);
    float mortar = smoothstep(0.0, 0.08, min(f.x, 1.0 - f.x) * 2.2)
                 * smoothstep(0.0, 0.1, min(f.y, 1.0 - f.y));
    float3 col = mix(float3(0.10, 0.05, 0.04), float3(0.24, 0.10, 0.07),
                     hash31(float3(floor(b), 2.0)));
    col *= 0.45 + 0.55 * mortar;
    float moon = max(dot(n, MOON), 0.0) * 0.6 + 0.08;
    float3 lit = col * MOONLIGHT * moon;
    lit += kiln(0.6) * glow * exp(pos.y * 9.0) * 0.8 * col * 6.0;
    return lit;
}

// The light a bit of smoke sends toward the eye, per unit it takes out.
float3 lightAt(float3 pos, float3 rd, float d, float moonPhase, float glow,
               float t, constant Params& p) {
    // The moon, through the smoke between, by the phase function.
    float3 lit = MOONLIGHT * 5.0 * moonPhase
               * lightThrough(pos, MOON, t, p, 1.6, 4);
    // The fire, from below: through the smoke, fading with distance.
    float3 toFire = FIRE - pos;
    float fd = length(toFire);
    toFire /= fd;
    float fireT = lightThrough(pos, toFire, t, p, min(fd, 1.2), 2);
    float firePhase = phase(dot(rd, toFire), p.forward * 0.5) * 4.0 * PI;
    lit += kiln(0.62) * 4.0 * glow * fireT * firePhase / (1.0 + fd * fd * 1.6);
    // The sky's soft light.
    lit += float3(0.03, 0.04, 0.08) * (1.0 - 0.5 * smoothstep(0.0, 1.5, d));
    // Smoke scatters about half of what it takes out. Fresh smoke glows.
    float heat = exp(-pos.y * 5.0) * glow;
    return lit * 0.55 + kiln(0.5 + 0.3 * heat) * heat * 1.2;
}

float4 shade(float2 uv, constant Uniforms& u, constant Params& p) {
    float t = u.time * p.speed + 4.0;

    // A camera on the ground near the chimney, looking up past it.
    float side = min(u.resolution.x, u.resolution.y);
    float2 s = (uv - 0.5) * u.resolution / side * 2.0;
    float3 ro = float3(-0.7, 0.05, -3.4);
    float3 target = float3(0.5, 1.25, 0.0);
    float3 fw = normalize(target - ro);
    float3 rt = normalize(cross(float3(0.0, 1.0, 0.0), fw));
    float3 up = cross(fw, rt);
    float3 rd = normalize(s.x * rt + s.y * up + 1.9 * fw);

    // Behind everything, the sky; in front of it, maybe the chimney, where
    // the smoke march stops.
    float3 col = nightSky(rd, t);
    float far = 14.0;
    float hit = hitStack(ro, rd);
    float glow = p.glow * (0.8 + 0.2 * sin(t * 7.0) * sin(t * 2.3));
    if (hit > 0.0) {
        col = brick(ro + rd * hit, glow);
        far = hit;
    }
    // March through the smoke's box in even steps, from a random place in
    // the first step, fixed for each pixel.
    float3 lo = float3(-2.2, 0.0, -2.2), hi = float3(4.2, 5.5, 2.2);
    float3 inv = 1.0 / rd;
    float3 t0 = (lo - ro) * inv, t1 = (hi - ro) * inv;
    float3 near = min(t0, t1), away = max(t0, t1);
    float enter = max(max(near.x, near.y), max(near.z, 0.0));
    float leave = min(min(away.x, away.y), min(away.z, far));
    if (leave > enter) {
        const int STEPS = 48;
        float dt = (leave - enter) / float(STEPS);
        float dist = enter + dt * hash31(float3(uv * u.resolution, 5.1));
        float transmit = 1.0;
        float3 light = float3(0.0);
        float moonPhase = phase(dot(rd, MOON), p.forward);
        for (int i = 0; i < STEPS; i++) {
            if (dist > leave || transmit < 0.02) break;
            float3 pos = ro + rd * dist;
            float d = smokeAt(pos, t, p, 4);
            if (d > 0.002) {
                // Light sent to the eye: what the step takes out, times the
                // light arriving. Then it dims what lies behind.
                float stepT = exp(-d * dt);
                float3 lit = lightAt(pos, rd, d, moonPhase, glow, t, p);
                light += transmit * lit * (1.0 - stepT);
                transmit *= stepT;
            }
            dist += dt;
        }
        col = col * transmit + light;
    }

    col = 1.0 - exp(-col * 1.6);
    return float4(col, 1.0);
}

// ---- Sound -------------------------------------------------------------------------------------------------------

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

// Cracked Clay's sound noise: random values joined `rate` times a second.
float noise1(float t, float rate, uint seed) {
    float x = t * rate;
    float i = floor(x), f = fract(x);
    f = f * f * (3.0 - 2.0 * f);
    uint k = uint(i) + seed * 7919u;
    return mix(hashU(k), hashU(k + 1u), f) * 2.0 - 1.0;
}

float2 sound(float t, constant Params& p) {
    float tt = t * p.speed + 4.0;
    // The kiln breathing out, as the picture's smoke leaves the chimney.
    float breath = emission(tt, p);
    float since = fract(tt / 20.0 + 0.6) * 20.0;

    // The draft: air rushing up the stack, a dark rumble of slow noise
    // and a hollow note, both louder with each breath.
    float rumble = noise1(t, 90.0, 1u) * 0.6 + noise1(t, 230.0, 2u) * 0.4;
    float hum = sin(TAU * 49.0 * t + 0.6 * noise1(t, 3.0, 3u))
              + 0.4 * sin(TAU * 98.0 * t);
    float draft = rumble * 0.34 * (0.4 + 0.5 * breath)
                + hum * 0.1 * (0.5 + 0.5 * p.glow);

    // The smoke going: a soft hiss, whooshing with the puffs, panned with
    // the wind that carries it.
    float hiss = noise1(t, 2600.0, 4u) * 0.6 + noise1(t, 900.0, 5u) * 0.4;
    hiss *= 0.07 * breath * p.thickness;
    float pan = 0.5 + 0.35 * windAt(tt, p);

    // The fire: crackles, short bright bursts at random moments, many just
    // after a stoking, a few the rest of the time.
    float slot = floor(t * 14.0);
    float age = fract(t * 14.0) / 14.0;
    float chance = 0.08 + 0.6 * exp(-since * 0.4);
    float crack = hashU(uint(slot) * 3u + 11u) < chance * p.glow ? 1.0 : 0.0;
    float burst = noise1(t, 7000.0, 6u) * exp(-age * 260.0)
                * smoothstep(0.0, 0.0006, age) * crack;
    float side2 = hashU(uint(slot) * 3u + 12u);

    float2 mixdown = float2(draft)
                   + hiss * float2(1.0 - pan, pan) * 1.6
                   + burst * 0.45 * float2(1.0 - side2, side2);
    float fadeIn = smoothstep(0.0, 1.5, t);
    return tanh(mixdown * 1.5) * 0.8 * 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);
}