#include <metal_stdlib>
using namespace metal;

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

struct Params {
    float every;
    float distance;
    float jag;
    float branches;
    float restrikes;
    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 · Lightning
//
// @episode Lightning
// @length 60
// @variation 30
// @variations 1,2,3,5,8,9,11,12
// @still 45.9
// @teaches Branching bolts by midpoint displacement, glow, thunder that arrives late
// @category Fragment
// @tags midpoint displacement, fractals, lightning, glow, procedural sound
// @sweeps jag, distance, glow
// @final Lightning | Strikes near and far under a ceiling of cloud, each one's thunder a ring on the ground that reaches you seconds after the flash.
//
// A storm at night, seen from a field. Each bolt is drawn by midpoint
// displacement: a straight line from the cloud to the ground, its middle
// pushed sideways, then the middle of each half, and so on, each time by less.
// Branches grow from points along it the same way. A strike comes down as a
// faint leader, then the return stroke lights it white and lights the clouds;
// the bolt flickers as it strikes again down the same channel. The thunder
// leaves the strike at the speed of sound, drawn as a ring on the ground, and
// you hear it when the ring reaches you: about three seconds for every
// kilometre.
// 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 every 5 2 12 Seconds between strikes, on average
// @param distance 1.2 0.4 3 How far away the storm is, in kilometres
// @param jag 0.58 0.4 0.72 How rough the bolts are: each level of displacement keeps this much of the last
// @param branches 5 0 8 Branches on each bolt
// @param restrikes 2 0 3 Strokes after the first down the same channel, in the bolt alone
// @param glow 1 0 2 How far the light spreads from the bolt
// @param volume 0.8 0 1 Loudness of the sound

// The horizon on screen, and our height in km: ground z km away is drawn
// CAMERA / z below the horizon.
constant float HORIZON = 0.24;
constant float CAMERA = 0.06;
// The cloud base's height, in km.
constant float CEILING = 1.2;
// The speed of sound, in km a second.
constant float SOUND = 0.343;
// How many times the line is split: 2^7 = 128 pieces.
constant int LEVELS = 7;
constant float PIECES = 128.0;

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

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

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

// Fractal noise: several octaves of value noise, each finer and fainter.
float fbm(float2 p) {
    float v = 0.0, a = 0.5;
    for (int i = 0; i < 5; i++) {
        v += a * vnoise(p);
        p = p * 2.03 + float2(17.1, 9.2);
        a *= 0.5;
    }
    return v;
}

// Integer hash, so noise stays exact for minutes of sound.
float hashU(uint n) {
    n = (n << 13u) ^ n;
    n = n * (n * n * 15731u + 789221u) + 1376312589u;
    return float(n & 0x7fffffffu) / float(0x7fffffff);
}

// Distance from point a to the segment from b to c.
float segment(float2 a, float2 b, float2 c) {
    float2 ba = a - b, cb = c - b;
    float h = clamp(dot(ba, cb) / dot(cb, cb), 0.0, 1.0);
    return length(ba - cb * h);
}

// How far sideways the point a fraction s along a bolt is pushed: one tent
// over the whole line, two over its halves, 2^l at level l, each level jag
// times as tall as the last. Rough, and still one line.
float boltOffset(float s, uint seed, float jag) {
    float off = 0.0, amp = 0.2, n = 1.0;
    for (int l = 0; l < LEVELS; l++) {
        float u = clamp(s, 0.0, 1.0) * n;
        float j = min(floor(u), n - 1.0), f = u - j;
        uint k = seed * 2048u + uint(l) * 256u + uint(j);
        off += (hashU(k) * 2.0 - 1.0) * amp * (1.0 - abs(2.0 * f - 1.0));
        amp *= jag;
        n *= 2.0;
    }
    return off;
}

// The point a fraction s along the bolt from a to b, pushed sideways.
float2 boltPoint(float s, float2 a, float2 b, uint seed, float jag) {
    float2 d = b - a;
    return a + d * s + float2(-d.y, d.x) * boltOffset(s, seed, jag);
}

// Distance from p to the bolt from a to b, as far as the fraction `upto`.
// Only the few pieces about p's place along the line can be nearest.
float boltDistance(float2 p, float2 a, float2 b, uint seed, float jag,
                   float upto) {
    float2 d = b - a;
    float len2 = dot(d, d);
    float s = dot(p - a, d) / len2;
    float side = abs(dot(p - a, float2(-d.y, d.x))) / len2;
    if (upto <= 0.0 || s < -0.1 || s > 1.1 || side > 0.2 / (1.0 - jag))
        return 1e3;
    float i0 = floor(clamp(s, 0.0, 1.0) * PIECES);
    float best = 1e3;
    float2 last = boltPoint((i0 - 3.0) / PIECES, a, b, seed, jag);
    for (int i = -2; i <= 3; i++) {
        float s1 = (i0 + float(i)) / PIECES;
        float2 next = boltPoint(s1, a, b, seed, jag);
        if (s1 > 0.0 && s1 - 1.0 / PIECES < upto)
            best = min(best, segment(p, last, next));
        last = next;
    }
    return best;
}

// Branches leave the channel at random points and angles, shorter lower
// down. Each is a bolt of its own, from where it leaves to where it ends.
float branchDistance(float2 q, float2 top, float2 foot, uint seed, int b,
                     float jag, float upto) {
    uint bs = seed * 16u + uint(b);
    float from = 0.1 + 0.65 * hashU(bs * 3u);
    float turn = (0.35 + 0.55 * hashU(bs * 3u + 1u))
               * (hashU(bs * 3u + 2u) < 0.5 ? -1.0 : 1.0);
    float len = (0.15 + 0.25 * hashU(bs + 999u)) * (1.0 - 0.5 * from);
    float2 a = boltPoint(from, top, foot, seed, jag);
    float2 axis = foot - top;
    float c = cos(turn), sn = sin(turn);
    float2 dir = float2(axis.x * c - axis.y * sn, axis.x * sn + axis.y * c);
    float reach = (upto - from) / (len * c);
    return boltDistance(q, a, a + dir * len, bs + 77u, jag, reach);
}

// As in Rain on Glass, the clock is cut into windows, here `every` seconds
// long, with a strike at a random moment in the first half of each, so two
// are never closer than half a window. Each strike's place on the ground,
// x across and z away in kilometres, and its seed come from its window.
struct Strike {
    float at;
    float2 ground;
    uint seed;
};

Strike strikeIn(float k, float aspect, constant Params& p) {
    uint w = uint(k);
    Strike s;
    s.at = k * p.every + 0.5 * p.every * hashU(w * 7u + 1u);
    float z = p.distance * (0.55 + 0.9 * hashU(w * 7u + 2u));
    s.ground = float2((hashU(w * 7u + 3u) - 0.5) * 0.75 * aspect * z, z);
    s.seed = w * 7u + 4u;
    return s;
}

// Where a point on the ground or the cloud base is drawn.
float2 toScreen(float2 ground, float height, float aspect) {
    return float2(0.5 * aspect + ground.x / ground.y,
                  HORIZON + (height - CAMERA) / ground.y);
}

// The bolt's light `x` seconds into the strike. The leader takes 0.2 s to
// come down, faint; then the return stroke, and each restrike 0.1 s apart,
// is a sharp pulse that fades.
constant float LEADER = 0.2;
float strokes(float x, constant Params& p) {
    if (x < LEADER) return 0.25 * step(0.0, x);
    float light = 0.0;
    for (int j = 0; j <= 3; j++) {
        float y = x - LEADER - 0.1 * float(j);
        if (float(j) <= p.restrikes + 0.5 && y >= 0.0)
            light += pow(0.8, float(j)) * exp(-y * 14.0);
    }
    return light;
}

// The light thrown on the clouds and the ground: one rise and one smooth fall
// a strike, however often its bolt flickers. A flash that comes and goes more
// than three times a second can cause seizures; this one never does.
float skyLight(float x) {
    float y = x - LEADER;
    if (y < 0.0) return 0.0;
    return smoothstep(0.0, 0.04, y) * exp(-y * 2.2);
}

// One bolt and its branches, as heat, for the strike `x` seconds old.
float boltHeat(float2 q, Strike s, float x, float px, float aspect,
               constant Params& p) {
    float2 top = toScreen(s.ground, CEILING, aspect);
    top.y = min(top.y, 1.15);
    float2 foot = toScreen(s.ground, 0.0, aspect);
    // How far down the leader has come, and how bright the channel is.
    float down = clamp(x / LEADER, 0.0, 1.0);
    float light = strokes(x, p);
    float sky = skyLight(x);
    // Thinner and dimmer the farther it is.
    float near = clamp(0.8 / s.ground.y, 0.3, 1.6);
    float width = px * (0.6 + 0.9 * near);
    float reach = 0.012 * p.glow * near;

    float d = boltDistance(q, top, foot, s.seed, p.jag, down);
    float heat = light * near * (1.0 - smoothstep(0.0, width, d)) * 1.3;
    heat += (sky + 0.3 * step(x, LEADER) * down) * near * 0.5
          * exp(-d / max(reach, px));

    // The branches, lit as the leader reaches them; they fade sooner.
    for (int b = 0; b < 8; b++) {
        if (float(b) >= p.branches) break;
        float bd = branchDistance(q, top, foot, s.seed, b, p.jag, down);
        float fade = x < LEADER ? 0.7 : 0.45 * exp(-(x - LEADER) * 3.0);
        heat += fade * (light + 0.3) * near
              * (1.0 - smoothstep(0.0, width * 0.8, bd)) * 0.9;
        heat += fade * sky * near * 0.3 * exp(-bd / max(reach * 0.6, px));
    }
    return heat;
}


// The cloud base, seen from below: a plane CEILING high, so its noise is
// drawn smaller towards the horizon. Its density, and the place on it.
float cloudAt(float2 q, float aspect, float t, thread float2& place) {
    float z = (CEILING - CAMERA) / max(q.y - HORIZON, 1e-3);
    place = float2((q.x - 0.5 * aspect) * z, z);
    float c = fbm(place * float2(1.6, 1.2) + float2(t * 0.03, 0.0));
    return smoothstep(0.25, 0.85, c) * exp(-z * 0.08);
}

float4 shade(float2 uv, constant Uniforms& u, constant Params& p) {
    float t = u.time;
    float aspect = u.resolution.x / u.resolution.y;
    float2 q = float2(uv.x * aspect, uv.y);
    float px = 1.0 / u.resolution.y;
    bool sky = q.y > HORIZON;

    float2 place;
    float cloud = sky ? cloudAt(q, aspect, t, place) : 0.0;
    // The ground's place, x across and z away.
    float zg = CAMERA / max(HORIZON - q.y, 1e-3);
    float2 ground = float2((q.x - 0.5 * aspect) * zg, zg);
    // A dim sky, darker overhead, and haze along the horizon.
    float heat = sky ? 0.07 + 0.05 * cloud - 0.03 * (q.y - HORIZON)
                     : 0.03 + 0.02 * fbm(ground * 3.0);
    heat += 0.06 * exp(-abs(q.y - HORIZON) * 25.0);

    float k = floor(t / p.every);
    for (int i = 0; i < 4; i++) {
        float kk = k - float(i);
        if (kk < 0.0) break;
        Strike s = strikeIn(kk, aspect, p);
        float x = t - s.at;
        if (x < 0.0) continue;
        float lit = skyLight(x) * clamp(0.9 / s.ground.y, 0.25, 1.5);
        if (sky) {
            // The clouds glow about the strike, a little everywhere.
            float far = length(place - s.ground);
            heat += lit * (0.2 + cloud) * (0.25 + exp(-far * 1.2));
        } else {
            float far = length(ground - s.ground);
            heat += lit * 0.35 * exp(-far * 1.5);
            // The thunder: a ring on the ground leaving the strike at the
            // speed of sound. It reaches you as you hear it.
            float r = SOUND * x;
            float w = px * max(zg * zg / CAMERA, zg) * 1.5;
            float ring = exp(-abs(far - r) / w) * exp(-r * 0.5)
                       * step(r, length(s.ground) + 0.05);
            heat += 0.22 * ring;
        }
        if (x < 1.5) heat += boltHeat(q, s, x, px, aspect, p);
    }
    return float4(kiln(heat), 1.0);
}

// Rain, and thunder. Each strike's thunder starts when its ring reaches you:
// its distance over the speed of sound. A near strike cracks first, then
// rumbles; a far one only rumbles, lower, longer and quieter.

// Noise made by joining random values `rate` times a second:
// a lower rate sounds darker.
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 `x` seconds after it arrives, from `far` kilometres away.
float thunder(float t, float x, float far, uint seed) {
    if (x <= 0.0) return 0.0;
    float near = smoothstep(1.6, 0.4, far);
    float crack = (noise1(t, 3500.0, seed) * 0.6
                 + noise1(t, 1200.0, seed + 1u) * 0.5)
                * exp(-x * 12.0) * smoothstep(0.0, 0.004, x) * near;
    float env = (1.0 - exp(-x * 5.0)) * exp(-x / (0.9 + 0.6 * far))
              * smoothstep(12.0, 9.0, x);
    float swell = 0.55 + 0.45 * noise1(t, 2.5, seed + 2u);
    float low = 1.0 - 0.15 * far;
    float rumble = noise1(t, 50.0 * low, seed + 3u) * 0.6
                 + noise1(t, 110.0 * low, seed + 4u) * 0.4
                 + noise1(t, 260.0 * low, seed + 5u) * 0.25;
    return (crack * 0.7 + rumble * env * swell) * 5.0 / (0.5 + far);
}

float2 sound(float t, constant Params& p) {
    float aspect = 16.0 / 9.0;
    // The rain: noise with its lowest part taken away.
    float2 rain = float2(noise1(t, 9000.0, 1u) - noise1(t, 900.0, 2u) * 0.6,
                         noise1(t, 9000.0, 4u) - noise1(t, 900.0, 5u) * 0.6);
    float2 sum = rain * 0.09;

    float k = floor(t / p.every);
    for (int i = 0; i < 16; i++) {
        float kk = k - float(i);
        if (kk < 0.0) break;
        Strike s = strikeIn(kk, aspect, p);
        float far = length(s.ground);
        float x = t - s.at - far / SOUND;
        if (x <= 0.0 || x > 12.0) continue;
        // Louder on the side it struck, each ear a little different.
        float pan = 0.5 + 0.4 * s.ground.x / far;
        float a = thunder(t, x, far, s.seed);
        float b = thunder(t, x, far, s.seed + 50u);
        sum += float2((0.7 * a + 0.3 * b) * (1.2 - pan),
                      (0.3 * a + 0.7 * b) * (0.2 + pan));
    }
    // A fade in over the first 0.3 s, so the rain doesn't start with a click.
    return tanh(sum * 1.2) * 0.66 * 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);
}