#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 rain;
    float mist;
    float blur;
    float lens;
    float storm;
    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 · Rain on Glass
//
// @episode Rain on Glass
// @length 60
// @variation 30
// @variations 2,3,5,6,8,9,10,11
// @still 12.0
// @teaches Refraction through drops on a window, a blurred world behind
// @category Fragment
// @tags refraction, bokeh, rain, lightning, procedural sound
// @final Rain on Glass | A misted window at night: drops slide and leave clear trails, and each drop holds the street upside down.
//
// A window at night, rain on the glass, a street of lights behind it. The camera is focused on the glass, so the
// street is out of focus: each light is a soft disc, bigger the more it is blurred. A drop on the glass is a small lens:
// it bends the light from a wide patch of the street into its own little circle, upside down and sharp. So for a pixel
// inside a drop, the street is read at a point on the far side of the drop's middle, at a distance that grows with how
// far the pixel is from the middle. Around the drops the glass is misted: the street is blurrier and greyer there,
// except where a sliding drop has wiped a clear trail. Now and then lightning lights the sky. The pointer moves the
// street behind the glass, as if you leaned: the drops stay, and the street in them moves the other way.
// 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 3.0 How fast the drops slide, the traffic moves and the lightning comes
// @param rain 0.6 0.0 1.0 How much rain: how many drops, and how loud it is
// @param mist 0.6 0.0 1.0 How misted the glass is between the drops
// @param blur 0.7 0.0 1.0 How far out of focus the street is behind the glass
// @param lens 0.5 0.0 1.0 How much of the street each drop gathers into itself
// @param storm 0.4 0.0 1.0 How often lightning lights the sky
// @param volume 0.8 0.0 1.0 Loudness of the sound

// ---- Noise ------------------------------------------------------------------

float hash21(float2 p) {
    p = fract(p * float2(123.34, 456.21));
    p += dot(p, p + 45.32);
    return fract(p.x * p.y);
}

float2 hashPair(float2 p) {
    return float2(hash21(p), hash21(p + 19.19));
}

// Value noise: random values on a grid, blended smoothly between the grid
// points.
float vnoise(float2 p) {
    float2 i = floor(p), f = fract(p);
    f = f * f * (3.0 - 2.0 * f);
    float a = hash21(i), b = hash21(i + float2(1, 0));
    float c = hash21(i + float2(0, 1)), d = hash21(i + float2(1, 1));
    return mix(mix(a, b, f.x), mix(c, d, f.x), f.y);
}

// The Ray Kiln 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));
}

// ---- The street behind the glass --------------------------------------------

// Lightning: the clock is cut into windows of 9 seconds, and in each, with a
// chance that grows with storm, a strike at a random moment, a flash and a
// second, weaker one just after. The sound's thunder follows the same strikes.
float strikeAt(float k, float storm) {
    if (hash21(float2(k, 3.3)) >= storm * 0.8) return -1.0;
    return 1.0 + 5.0 * hash21(float2(k, 7.1));
}

float flash(float time, float storm) {
    float k = floor(time / 9.0);
    float at = strikeAt(k, storm);
    if (at < 0.0) return 0.0;
    float x = time - k * 9.0 - at;
    if (x < 0.0) return 0.0;
    return exp(-x * 9.0) + 0.6 * step(0.18, x) * exp(-(x - 0.18) * 7.0);
}

// One light's colour: mostly the kiln's ember, orange and gold, a few cool
// white ones, as a street's lamps and windows are.
float3 lampColour(float h) {
    if (h > 0.85) return float3(0.75, 0.85, 1.0);
    return kiln(0.45 + 0.5 * h);
}

// One layer of lights, `scale` cells to a unit; some cells hold one. Out of
// focus a light is a disc most of a cell wide, sharp a small dot. Its light
// is spread over its area, so a big disc is dimmer, its edge a little
// brighter. `aa` is a pixel's width.
float3 lampLayer(float2 p, float scale, float seed, float soft, float aa,
                 float low) {
    float2 g = p * scale + seed * float2(17.3, 5.1);
    float2 cell = floor(g);
    float3 col = float3(0.0);
    float edge = aa * scale;
    float glow = 0.12 * soft;
    for (int y = -1; y <= 1; y++) {
        for (int x = -1; x <= 1; x++) {
            float2 c = cell + float2(x, y);
            float h = hash21(c + seed * 31.0);
            // Fewer lights higher up the street.
            float height = (c.y + 0.5) / scale - seed * 5.1 / scale;
            float dense = low * smoothstep(0.9, -0.4, height);
            if (h > dense) continue;
            float2 at = c + 0.5 + (hashPair(c + seed) - 0.5) * 0.6;
            float rad = mix(0.05, 0.42 + 0.15 * hash21(c + 2.7), soft);
            float d = length(g - at);
            float inner = rad - edge - glow * rad;
            float disc = 1.0 - smoothstep(inner, rad + edge, d);
            float rim = 1.0 + 0.6 * smoothstep(rad * 0.5, rad, d) * soft;
            float energy = 0.05 / ((rad + edge) * (rad + edge) + 0.05);
            col += lampColour(hash21(c + 5.9)) * disc * rim * energy;
        }
    }
    return col;
}

// The traffic: two rows of lights moving across, red going one way and gold
// coming the other, low down on the street.
float3 traffic(float2 p, float time, float soft, float aa) {
    float3 col = float3(0.0);
    for (int i = 0; i < 2; i++) {
        float dir = i == 0 ? 1.0 : -1.0;
        float y = i == 0 ? -0.30 : -0.35;
        float run = time * 0.8 * dir + float(i) * 3.3;
        float2 g = float2(p.x * 7.0 + run, (p.y - y) * 7.0);
        float cell = floor(g.x);
        float3 tint = i == 0 ? float3(1.0, 0.12, 0.05) : kiln(0.85);
        for (int n = -1; n <= 1; n++) {
            float c = cell + float(n);
            if (hash21(float2(c, float(i))) > 0.45) continue;
            float2 at = float2(c + 0.5, 0.0);
            float rad = mix(0.04, 0.35, soft);
            float edge = aa * 7.0;
            float d = length(g - at);
            float disc = 1.0 - smoothstep(rad - edge, rad + edge, d);
            float energy = 0.05 / ((rad + edge) * (rad + edge) + 0.05);
            col += tint * disc * energy * 0.8;
        }
    }
    return col;
}

// The street: a night sky with the city's glow, three layers of lights and
// the traffic, `soft` from 0 (in focus) to 1. Lightning lights it.
float3 street(float2 p, float soft, float aa, float time, float storm) {
    float3 col = mix(float3(0.045, 0.03, 0.045), float3(0.010, 0.010, 0.022),
                     smoothstep(-0.3, 0.5, p.y));
    col += kiln(0.42) * 0.3 * exp(-(p.y + 0.5) * 2.5);
    float f = flash(time, storm);
    float sky = 0.25 + 0.6 * smoothstep(-0.2, 0.6, p.y);
    col += float3(0.55, 0.6, 0.85) * f * sky;
    col += lampLayer(p, 16.0, 1.0, soft, aa, 0.32);
    col += lampLayer(p, 9.0, 2.0, soft, aa, 0.26);
    col += lampLayer(p, 5.0, 3.0, soft, aa, 0.2) * 1.2;
    col += traffic(p, time, soft, aa);
    return col * (1.0 + 0.6 * f);
}

// ---- The drops --------------------------------------------------------------

// A drop over this pixel: where the street is read, where the pixel is in
// it, how much it covers, how much mist is wiped, how much smaller the street
// is in it.
struct Drop {
    float2 at;
    float2 d;
    float mask;
    float clear;
    float zoom;
};

// The lens: the street is read on the far side of the drop's middle, upside
// down, across a patch `span` wide.
Drop lensDrop(float2 q, float2 c, float2 r, float span, float mask) {
    Drop o;
    o.d = (q - c) / r;
    float len = length(o.d);
    o.mask = mask * (1.0 - smoothstep(0.82, 1.0, len));
    o.at = c - o.d * span;
    o.clear = 0.0;
    o.zoom = span / r.x;
    return o;
}

// Whichever of two drops covers more of the pixel; trails add up.
Drop nearer(Drop a, Drop b) {
    Drop o = a.mask >= b.mask ? a : b;
    o.clear = max(a.clear, b.clear);
    return o;
}

// The still drops: in each cell of a grid, maybe a drop, appearing at once
// and slowly drying away, then another later.
Drop stillDrops(float2 q, float time, float rain, float lens, float scale,
                float seed) {
    float2 g = q * scale + seed;
    float2 cell = floor(g);
    float h = hash21(cell + seed * 7.0);
    float2 n = hashPair(cell + seed * 3.0) - 0.5;
    float life = fract(time * 0.05 + h * 7.0);
    float fade = smoothstep(0.0, 0.02, life) * smoothstep(1.0, 0.6, life);
    float there = h < rain * 0.7 ? 1.0 : 0.0;
    float2 c = (cell + 0.5 + n * 0.5 - seed) / scale;
    float r = (0.12 + 0.12 * hash21(cell + 9.1)) / scale;
    float span = r * 2.0 + 0.04 + 0.3 * lens;
    return lensDrop(q, c, float2(r, r), span, fade * there);
}

// In each column, maybe a drop sliding down. It sticks and slips: the clock
// plus a ripple just small enough that it stops but never climbs. It wobbles
// across, and its trail follows the same path, leaving small drops.
Drop slidingDrops(float2 q, float time, float rain, float lens, float cols,
                  float seed) {
    const float tau = 6.2831853;
    float gx = q.x * cols + seed;
    float col = floor(gx);
    float h = hash21(float2(col, seed));
    float there = h < rain ? 1.0 : 0.0;
    float tt = time * (0.07 + 0.05 * h) + h * 10.0;
    float s = fract(tt);
    float yd = 0.7 - 1.4 * (s + 0.024 * sin(tau * s * 6.0 + h * 6.0));
    float r = (0.18 + 0.1 * hash21(float2(col, seed + 4.0))) / cols;
    float wob = 0.18 / cols;
    float pathX = (col + 0.5 - seed) / cols;
    float xAtDrop = pathX + wob * sin(yd * 14.0 + h * 20.0);
    float xAtPixel = pathX + wob * sin(q.y * 14.0 + h * 20.0);
    float span = r * 2.0 + 0.04 + 0.3 * lens;

    // The drop: a little taller than wide, as a running drop is.
    Drop o = lensDrop(q, float2(xAtDrop, yd), float2(r, r * 1.25), span, there);

    // The trail above it, fading with height, and its small drops.
    float above = q.y - yd;
    float side = abs(q.x - xAtPixel);
    float trail = there * step(0.0, above) * exp(-above * 3.0)
                * (1.0 - smoothstep(r * 0.5, r * 0.9, side));
    o.clear = trail;
    float rowH = 0.025;
    float row = floor(q.y / rowH);
    float hr = hash21(float2(col * 13.0 + row, seed));
    float yc = (row + 0.5) * rowH;
    float xc = pathX + wob * sin(yc * 14.0 + h * 20.0) + (hr - 0.5) * r * 0.6;
    float rr = r * (0.2 + 0.25 * hr);
    float keep = there * step(yd + r, yc) * step(hr, 0.6);
    keep *= exp(-(yc - yd) * 2.0);
    float spanS = rr * 2.0 + 0.02 + 0.15 * lens;
    Drop small = lensDrop(q, float2(xc, yc), float2(rr, rr), spanS, keep);
    return nearer(o, small);
}

// ---- The window -------------------------------------------------------------

float4 shade(float2 uv, constant Uniforms& u, constant Params& p) {
    float aspect = u.resolution.x / u.resolution.y;
    float2 q = (uv - 0.5) * float2(aspect, 1.0);
    float time = u.time * p.speed;
    float pixel = 1.0 / u.resolution.y;

    // Leaning: the pointer moves the street behind the glass.
    float2 lean = (u.mouse / u.resolution - 0.5) * float2(0.25, 0.15);

    // The drops: two sizes that slide, two grids of still ones.
    Drop drop = slidingDrops(q, time, p.rain, p.lens, 9.0, 0.0);
    float later = time * 1.3, less = p.rain * 0.8;
    drop = nearer(drop, slidingDrops(q, later, less, p.lens, 15.0, 0.37));
    drop = nearer(drop, stillDrops(q, time, p.rain * 0.75, p.lens, 22.0, 0.0));
    drop = nearer(drop, stillDrops(q, time, p.rain * 0.35, p.lens, 47.0, 0.53));

    // The misted glass: the street blurrier and greyer, the mist uneven, and
    // wiped where a trail has run.
    float grain = vnoise(q * 9.0) * 0.6 + vnoise(q * 31.0) * 0.4;
    float mist = p.mist * (0.75 + 0.5 * grain) * (1.0 - drop.clear);
    float soft = clamp(p.blur * (0.6 + 0.5 * mist), 0.0, 1.0);
    float3 col = street(q + lean, soft, pixel, time, p.storm);
    float f = flash(time, p.storm);
    col = col * (1.0 - 0.45 * mist)
        + float3(0.05, 0.05, 0.065) * mist * (1.0 + 4.0 * f);

    // Through a drop: the street sharp and upside down, gathered from a wide
    // patch, so a pixel in it is a smaller piece of the street. Its rim is
    // dark, where light bends too far to come through, and a highlight sits
    // near its top, the window's reflection of the sky.
    if (drop.mask > 0.0) {
        float aa = pixel * drop.zoom;
        float3 seen = street(drop.at + lean, 0.25, aa, time, p.storm) * 1.4;
        float len = length(drop.d);
        seen *= 1.0 - 0.55 * smoothstep(0.6, 1.0, len);
        float2 glint = drop.d - float2(-0.3, 0.4);
        seen += float3(1.0, 0.95, 0.88) * 0.5 * exp(-dot(glint, glint) * 24.0);
        col = mix(col, seen, drop.mask);
    }

    // Bright light goes to white rather than clipping its colour.
    col = 1.0 - exp(-col * 1.4);
    return float4(col, 1.0);
}

// ---- Sound ------------------------------------------------------------------
// The rain: a hiss whose loudness drifts in gusts; drops ticking on the glass,
// each a short high ring at a random pitch and place between the ears, more of
// them the more it rains; and after each lightning strike, thunder, a low
// rumble that swells and rolls away.

float hashU(uint n) {
    n = (n << 13u) ^ n;
    n = n * (n * n * 15731u + 789221u) + 1376312589u;
    return float(n & 0x7fffffffu) / float(0x7fffffff);
}

// Noise for sound: 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;
}

// The thunder of the strike in window k, `x` seconds after its flash: it
// starts 1.2 seconds later, as sound is slower than light, swells for half a
// second and rolls away over a few.
float thunder(float t, float x) {
    float y = x - 1.2;
    if (y <= 0.0) return 0.0;
    float env = (1.0 - exp(-y * 4.0)) * exp(-y * 0.9);
    float rumble = noise1(t, 70.0, 11u) * 0.6 + noise1(t, 140.0, 12u) * 0.3
                 + noise1(t, 35.0, 13u) * 0.5;
    return rumble * env;
}

float2 sound(float t, constant Params& p) {
    const float tau = 6.2831853;
    float time = t * p.speed;

    // The hiss: noise with its lowest part taken away, in gusts.
    float gust = 0.6 + 0.4 * noise1(t, 0.4, 3u);
    float hissL = noise1(t, 9000.0, 1u) - noise1(t, 900.0, 2u) * 0.6;
    float hissR = noise1(t, 9000.0, 4u) - noise1(t, 900.0, 5u) * 0.6;
    float2 hiss = float2(hissL, hissR) * 0.3 * gust * (0.3 + p.rain);

    // The ticks: in each 1/50 of a second, a chance of a drop, ringing at its
    // own pitch and dying in a few milliseconds.
    float window = t * 50.0;
    uint w = uint(window);
    float age = fract(window) / 50.0;
    float hit = step(hashU(w * 17u + 5u), 0.35 * p.rain);
    float pitch = 1800.0 + 3000.0 * hashU(w * 23u + 1u);
    float pan = hashU(w * 31u + 9u);
    float tick = hit * sin(tau * pitch * age) * exp(-age * 600.0)
               * smoothstep(0.0, 0.0004, age);
    float2 ticks = tick * float2(1.0 - pan, pan) * 0.35;

    // The thunder after each strike, from this window and the one before, so
    // a roll never stops at a window's end.
    float boom = 0.0;
    float k = floor(time / 9.0);
    for (int i = 0; i < 2; i++) {
        float kk = k - float(i);
        float at = strikeAt(kk, p.storm);
        if (at < 0.0) continue;
        boom += thunder(t, (time - kk * 9.0 - at) / max(p.speed, 0.05));
    }
    float2 rumble = float2(boom * 0.55, boom * 0.5);

    // A fade in over the first 0.3 s, so the hiss doesn't start with a click.
    return (hiss + ticks + rumble) * smoothstep(0.0, 0.3, t) * 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);
}