#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 storm;
    float height;
    float green;
    float stars;
    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 · Episode 24 · Aurora
//
// @episode 24 Aurora
// @length 60
// @variation 30
// @variations 1,2,3,5,6,8,9,11
// @still 14.0
// @teaches Thin noise curtains marched through, stars
// @category Fragment
// @tags volume rendering, fractal noise, stars, aurora, reflection, procedural sound
// @final Aurora | The curtains over a ridge and a still lake: they fold and drift, brighten in surges, and the lake gives them back.
//
// The northern lights over a mountain ridge, mirrored in a lake. An aurora is light hanging in the air: tall, thin
// curtains whose lower edge is sharp and whose top fades away. Seen from the ground, each curtain is a line on a high
// layer of the sky, folded by noise, and the same line on every layer above it. So the ray from the eye is marched up
// through the layers, and at each one it asks how near it is to a curtain's line: on the same line at every height,
// the light stacks into a wall. Stars sit behind it on a grid of directions, the ridge is fractal noise against the
// sky, and the lake is the same sky seen upside down, rippled. The pointer turns the view: in the playground, move it,
// steer with the arrow keys, or with a gamepad's stick.
// 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 curtains fold, drift and surge
// @param storm 0.6 0.0 1.0 How active the aurora is: how many curtains are lit, and how brightly
// @param height 0.5 0.0 1.0 How far up the curtains reach before they fade
// @param green 0.0 0.0 1.0 The curtains' colours, from the kiln's gold and ember to oxygen's green and red
// @param stars 0.6 0.0 1.0 How many stars, and how bright
// @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);
}

// 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 than
// the last.
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;
}

float2 rot(float2 p, float a) {
    float c = cos(a), s = sin(a);
    return float2(c * p.x - s * p.y, s * p.x + c * p.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 curtains -----------------------------------------------------------

// The aurora's surges: a slow swell from 0 to 1, its own noise over time, that
// lights more curtains and brightens them. The sound swells with it.
float surge(float time) {
    return smoothstep(0.2, 0.8, vnoise(float2(time * 0.09, 3.7)));
}

// How much curtain there is at a point xz on one layer of the sky. The
// curtains are lines running across the view, one behind another, 6.4 apart
// on layers whose lowest is 1 up; each is folded toward and away from us by
// noise that drifts along it, and the six nearest are drawn. Near its line a
// curtain is bright, and thin: its light falls away fast across it. Its rays
// are fine streaks along the line, the same on every layer, so they stand
// upright. `edge` gives its lower edge here, as a share of the way up.
float curtain(float2 xz, float time, float storm, thread float& edge) {
    float2 p = xz * 0.25;
    float fold = 3.2 * vnoise(float2(p.x * 0.3 + time * 0.05, 1.0))
               + 1.2 * vnoise(float2(p.x * 0.9 - time * 0.09, 5.0))
               + 0.5 * vnoise(float2(p.x * 2.3 + time * 0.2, 9.0));
    float z = (p.y + fold) / 1.6;
    float k = floor(z);
    float across = (fract(z) - 0.5) * 1.6;
    // Only the six curtains from 0 to 5 bands away; the rest is empty sky.
    if (k < 0.0 || k > 5.0) { edge = 1.0; return 0.0; }
    // Each curtain lit along part of its length, by a hash of its number and
    // slow noise along it: more of it the higher storm, and more in a surge.
    float lit = storm * (0.55 + 0.6 * surge(time + k * 3.0));
    float along = 0.4 * hash21(float2(k, 4.2))
                + 0.6 * vnoise(float2(p.x * 0.3 + k * 13.0, time * 0.04));
    float on = smoothstep(1.0 - lit, 1.2 - lit, along);
    // A far curtain is drawn a little wider, so the few layers that meet it
    // still find it.
    float width = 0.04 + 0.04 * hash21(float2(k, 8.9));
    width *= 1.0 + 0.04 * length(xz);
    float sheet = exp(-across * across / (width * width));
    float rays = 0.45 + 0.9 * vnoise(float2(p.x * 18.0 + time * 0.7, k * 7.0));
    rays *= 0.6 + 0.6 * vnoise(float2(p.x * 3.0 - time * 0.3, k * 3.0));
    edge = 0.12 * vnoise(float2(p.x * 1.7 + time * 0.1, k));
    return sheet * on * rays;
}

// The colour a share f of the way up: the kiln's, white-hot to ember, or
// oxygen's, pink at the edge, green, then deep red, blended by `green`.
float3 curtainColour(float f, float green) {
    float3 kilnCol = kiln(0.95 - 0.75 * f);
    float3 oxygen = mix(float3(0.15, 1.0, 0.42), float3(0.85, 0.08, 0.22),
                        smoothstep(0.25, 0.8, f));
    oxygen = mix(float3(0.95, 0.35, 0.6), oxygen, smoothstep(0.0, 0.06, f));
    return mix(kilnCol, oxygen, green);
}

// The aurora along a ray going up: marched through 48 layers of sky from
// height 1 to the curtains' top, each layer adding the light of the curtains
// it meets. `jitter` moves this pixel's layers up by a share of one step, so
// the steps blur rather than show as bands.
float3 aurora(float3 rd, float time, constant Params& p, float jitter) {
    if (rd.y <= 0.0) return float3(0.0);
    const int layers = 48;
    float top = 0.8 + 2.4 * p.height;
    float3 col = float3(0.0);
    for (int i = 0; i < layers; i++) {
        float f = (float(i) + jitter) / float(layers);
        float h = 1.0 + f * top;
        float2 xz = rd.xz * (h / rd.y);
        float edge;
        float c = curtain(xz, time, p.storm, edge);
        // Sharp at the lower edge, fading with height, fainter far away.
        float rise = smoothstep(edge, edge + 0.04, f);
        float fade = exp(-f * 3.0) * exp(-length(xz) * 0.06);
        col += curtainColour(f, p.green) * c * rise * fade;
    }
    return col * (30.0 / float(layers)) * smoothstep(0.02, 0.2, rd.y);
}

// ---- The sky, the stars, the ridge ------------------------------------------

// The stars, as in the first step: a few cells of 1/90 radian hold a soft,
// twinkling dot about a pixel wide (`pixel` is a pixel's angle).
float3 starField(float3 rd, float time, float amount, float pixel) {
    float az = atan2(rd.x, rd.z), el = asin(clamp(rd.y, -1.0, 1.0));
    float2 g = float2(az, el) * 90.0;
    float2 cell = floor(g);
    float h = hash21(cell);
    if (h < 1.0 - 0.07 * amount) return float3(0.0);
    float2 at = 0.2 + 0.6 * float2(hash21(cell + 3.1), hash21(cell + 7.7));
    float2 d = (fract(g) - at) * float2(cos(el), 1.0) / 90.0;
    float spot = exp(-dot(d, d) / (pixel * pixel * 0.8));
    float rate = 3.0 + 5.0 * hash21(cell + 1.3);
    float twinkle = 0.65 + 0.35 * sin(time * rate + h * 40.0);
    float bright = amount * (0.3 + 1.6 * pow(hash21(cell + 5.5), 4.0));
    float3 tint = mix(float3(1.0, 0.82, 0.62), float3(0.8, 0.88, 1.0),
                      hash21(cell + 9.2));
    return tint * spot * bright * twinkle * smoothstep(0.0, 0.25, rd.y);
}

// The ridge's height above the horizon, as an angle, for each direction round:
// fractal noise, higher where a slow wave lifts it.
float ridge(float az) {
    float lift = smoothstep(-0.4, 1.4, sin(az * 1.3 + 0.5) + 1.0);
    return 0.02 + 0.12 * fbm(float2(az * 2.2, 1.7)) * lift;
}

// Everything a ray sees above the horizon: the night sky, a little lighter
// low down, the stars, the curtains, and the ridge in front of them, a black
// shape with a little of the aurora's light on its snow.
float3 skyView(float3 rd, float time, constant Params& p, float jitter,
               float pixel) {
    float3 col = mix(float3(0.035, 0.03, 0.05), float3(0.008, 0.007, 0.014),
                     smoothstep(0.0, 0.5, rd.y));
    col += starField(rd, time, p.stars, pixel);
    col += aurora(rd, time, p, jitter);
    float az = atan2(rd.x, rd.z);
    float r = ridge(az);
    if (rd.y < r) {
        float snow = fbm(float2(az * 40.0, rd.y * 60.0));
        float3 glint = curtainColour(0.1, p.green) * 0.03 * p.storm * snow;
        float3 rock = float3(0.012, 0.011, 0.016) + glint;
        col = mix(col, rock, smoothstep(r, r - pixel * 1.5, rd.y));
    }
    return col;
}

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 / (1.1 * u.resolution.y);

    // The view: looking a little up, turned and tipped by the pointer's offset
    // from the middle of the picture (in an export it is the middle).
    float2 steer = u.mouse / u.resolution - 0.5;
    float3 rd = normalize(float3(q, 1.1));
    rd.yz = rot(rd.yz, -(0.18 + steer.y * 0.9));
    rd.xz = rot(rd.xz, steer.x * 2.4);

    // A different share of a layer for each pixel and frame (aurora above).
    float jitter = hash21(uv * u.resolution + fract(u.time * 7.0) * 61.0);
    float3 col;
    if (rd.y >= 0.0) {
        col = skyView(rd, time, p, jitter, pixel);
    } else {
        // The lake: the sky mirrored, the ray tipped up as much as it pointed
        // down, rippled by noise on the water that drifts with the clock, and
        // darker, since water gives back only part of the light.
        float2 water = rd.xz / -rd.y;
        float2 at = water * float2(1.0, 4.0) + float2(0.0, time * 0.4);
        float ripple = vnoise(at) - 0.5;
        float3 up = normalize(float3(rd.x + ripple * 0.006,
                                     -rd.y + ripple * 0.004, rd.z));
        col = skyView(up, time, p, jitter, pixel) * float3(0.5, 0.55, 0.6);
    }

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

// ---- Sound ------------------------------------------------------------------
// A low drone in A, a fifth and an octave, swelling and brightening with the
// aurora's surges; over it three glassy high tones whose loudness in each ear
// runs like the curtains' rays, so the shimmer moves between the ears; and in
// a surge, faint crackles.

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

// A soft tone with some of its overtones: `bright` from 0 (a pure sine) to 1
// (a reedy, fuller sound).
float voice(float ph, float bright) {
    const float tau = 6.2831853;
    float over = 0.45 * sin(2.0 * tau * ph) + 0.25 * sin(3.0 * tau * ph)
               + 0.12 * sin(4.0 * tau * ph);
    return sin(tau * ph) + bright * over;
}

float2 sound(float t, constant Params& p) {
    const float tau = 6.2831853;
    float time = t * p.speed;
    float swell = p.storm * (0.35 + 0.65 * surge(time));

    // The drone: A1, E2 and A2, each a few cents apart in the two ears.
    float bright = 0.15 + 0.6 * swell;
    float2 a1 = float2(voice(55.0 * t, bright), voice(55.08 * t, bright));
    float2 e2 = float2(voice(82.45 * t, bright), voice(82.35 * t, bright));
    float2 a2 = float2(voice(110.1 * t, bright), voice(109.95 * t, bright));
    float2 drone = a1 * 0.5 + e2 * 0.3 + a2 * 0.2;
    drone *= 0.2 + 0.15 * swell;

    // The shimmer: A4, E5 and A5, each as loud in each ear as noise that runs
    // at the speed of the picture's rays, one sample to the left, one to the
    // right.
    float2 shimmer = float2(0.0);
    for (int i = 0; i < 3; i++) {
        float side = float(i) - 1.0;
        float run = side * 0.7 + time * 0.7;
        float raysL = vnoise(float2(run - 3.0, 14.0 + float(i)));
        float raysR = vnoise(float2(run + 3.0, 21.0 + float(i)));
        float note = i == 0 ? 440.0 : (i == 1 ? 659.3 : 880.0);
        float tone = sin(tau * note * t) / float(i + 2);
        shimmer += tone * float2(raysL, raysR) * swell;
    }
    shimmer *= 0.085;

    // The crackles: in each 1/30 of a second, a chance that grows with the
    // swell of a short burst of noise, dying away in about 10 ms, panned at
    // random.
    float window = t * 30.0;
    uint w = uint(window);
    float age = fract(window) / 30.0;
    float hit = step(hashU(w * 13u + 7u), 0.08 * swell * swell);
    float pan = hashU(w * 29u + 3u);
    float noise = hashU(uint(t * 44100.0)) * 2.0 - 1.0;
    float crack = hit * exp(-age * 100.0) * noise;
    float2 crackle = crack * float2(1.0 - pan, pan) * 0.25;

    // A wind under it all, very low.
    float2 wind = float2(noise1(t, 160.0, 3u), noise1(t, 170.0, 5u)) * 0.04;

    float2 s = drone + shimmer + crackle + wind;
    return s * 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);
}