#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 twist;
    float oval;
    float young;
    float dust;
    float core;
    float tilt;
    float trails;
    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;
}

struct Particle;
// Every particle as it was after the last step, for an update that reads the others.
struct Swarm {
    const device Particle* ps;
    uint count;
    Particle at(uint j) const;
};

struct Trails {
    texture2d<float> image;
    float4 at(float2 uv) const {
        constexpr sampler s(filter::linear, address::clamp_to_edge);
        return image.sample(s, float2(uv.x, 1.0 - uv.y));
    }
};

#line 1 "shader.metal"
// Ray Kiln · Spiral Galaxy
//
// @episode Spiral Galaxy
// @length 60
// @variation 30
// @variations 1,2,3,4,7,10,11,12
// @still 24.0
// @teaches Stars on orbits, density waves, dust lanes
// @category Compute
// @tags compute shaders, particles, orbits, density waves, spiral galaxy, value noise, dust, stars, procedural sound
// @final Spiral galaxy | Stars on turning ovals crowd into two arms, young stars light where they pass through, and dust darkens the arms' inner edges.
//
// A spiral galaxy, seen at a slant. Its arms look like solid things, but no star stays in them: every star is on an
// orbit of its own, the inner ones going round faster than the outer. Each orbit is an oval, and each oval is turned
// a little further than the one inside it, so the ovals crowd together along two spiral lines. That crowding is the
// arm: a traffic jam that stars drive into and out of, while the jam itself turns slowly. Young stars and glowing gas
// light up where the jam squeezes them, and a lane of dust runs along each arm's inner edge.
// 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
//
// @particles 160000

// @param speed 1.0 0.0 3.0 How fast the stars go round: the outer ones once a minute at 1
// @param twist 3.2 0.5 6.0 How far each orbit is turned from the one inside it: loose arms or tightly wound
// @param oval 0.3 0.0 0.45 How oval the orbits are: round orbits make no arms
// @param young 0.6 0.0 1.0 How many young stars and glowing clouds light up in the arms
// @param dust 0.7 0.0 1.5 How dark the dust lanes are
// @param core 0.5 0.0 1.0 How bright the old stars of the core are
// @param tilt 58.0 0.0 80.0 How far the galaxy leans away from face-on, in degrees
// @param trails 0.6 0.0 0.9 How much of the last image is left at each step: crisp stars or a soft glow
// @param volume 0.8 0.0 1.0 Loudness of the sound

constant float PI = 3.1415927;
constant float TAU = 6.2831853;
// The core's size in the rotation curve: inside it the speed rises from nothing, outside it stays the same.
constant float CORE = 0.12;
// How fast the arms themselves turn, in radians a second at speed 1: far slower than the stars.
constant float PATTERN = 0.035;
// How much of the picture's half height the galaxy's radius fills.
constant float ZOOM = 1.1;

// What one star is, where it is now: its place in the galaxy's own frame (x and y in the disk, height out of it),
// what kind of star it is (0 an old disk star, 1 a star of the core, 2 a young star, 3 a cloud of glowing gas), and
// how lit it is now.
struct Particle {
    float2 position;
    float height;
    float kind;
    float glow;
};

// Cracked Clay's integer hash: a whole number in, a random number from 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);
}

// Kaleidoscope's rotation: a point turned by angle a about the middle.
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 rotation curve: how many radians a second a star goes round at distance a. Its speed along the orbit rises
// through the core and then stays the same, so an outer star, with further to go, takes longer to go round.
float orbitRate(float a, constant Params& p) {
    float v = 0.11 * a / sqrt(a * a + CORE * CORE);
    return p.speed * v / max(a, 0.01);
}

// How oval the orbit of size a is: round in the core, most oval in the disk, rounder again at the edge.
float ovalness(float a, constant Params& p) {
    float edge = 1.0 - 0.6 * smoothstep(0.6, 1.3, a);
    return p.oval * smoothstep(0.03, 0.25, a) * edge;
}

// How far the oval of size a is turned: twist turns each oval a little further than the one inside it, and the
// whole set turns slowly with time. This one line is where the arms come from.
float orbitTurn(float a, float t, constant Params& p) {
    return p.twist * a + PATTERN * p.speed * t;
}

// The point at angle m round the oval of size a, in the galaxy's frame at time t.
float2 onOrbit(float a, float m, float t, constant Params& p) {
    float2 q = float2(a * cos(m), a * (1.0 - ovalness(a, p)) * sin(m));
    return rot(q, orbitTurn(a, t, p));
}

// How deep in an arm a star at angle m round its oval is, 0 to 1. The ovals crowd where the next one out, turned a
// little further, swings in towards this one: an eighth of a turn back from each end of the long axis. None in the
// core, where the orbits are round, and fading out at the edge of the disk.
float armWeight(float m, float a, constant Params& p) {
    float squeeze = smoothstep(0.2, 1.0, -sin(2.0 * m));
    float disk = smoothstep(0.08, 0.3, a) * (1.0 - smoothstep(0.75, 1.15, a));
    return squeeze * disk * smoothstep(0.0, 0.1, p.oval);
}

// Star number id at time t. Nothing is carried from one step to the next: a star on an orbit is wherever its orbit
// has it at that time, so every step works out each star from its number and the clock.
Particle star(uint id, float t, constant Params& p) {
    float pick = hashU(id * 5u);
    float h1 = hashU(id * 5u + 1u), h2 = hashU(id * 5u + 2u);
    float h3 = hashU(id * 5u + 3u), h4 = hashU(id * 5u + 4u);
    Particle q;
    if (pick < 0.16) {
        // The core: small round orbits, each on a plane of its own, leaned by lean and turned by node, so the core
        // is a ball, not a disk.
        float a = 0.004 + 0.09 * -log(1.0 - h1 * 0.97);
        float m = h2 * TAU + orbitRate(a, p) * t;
        float lean = (h3 * 2.0 - 1.0) * 1.3;
        float2 flat = float2(cos(m), sin(m) * cos(lean)) * a;
        q.position = rot(flat, h4 * TAU);
        q.height = sin(m) * sin(lean) * a * 0.8;
        q.kind = 1.0;
        q.glow = 1.0;
        return q;
    }
    // The disk: most stars near the middle, fewer further out. Each goes round its oval at the rate of its size,
    // less the turning of the ovals themselves.
    float a = min(0.012 + 0.3 * -log(1.0 - h1 * 0.995), 1.3);
    float m = h2 * TAU + (orbitRate(a, p) - PATTERN * p.speed) * t;
    q.position = onOrbit(a, m, t, p);
    float arm = armWeight(m, a, p);
    if (pick < 0.16 + 0.1 * p.young) {
        // Young stars: thin in the disk, and lit only while the arm squeezes them.
        q.height = (h3 - 0.5) * 0.01;
        q.kind = 2.0;
        q.glow = arm * arm;
    } else if (pick < 0.27 && h4 < p.young) {
        // Glowing gas, fewer, lit only in the thick of the squeeze.
        q.height = (h3 - 0.5) * 0.01;
        q.kind = 3.0;
        q.glow = arm * arm * arm;
    } else {
        // Old stars: everywhere in the disk, which is thicker in the middle, and dimmer far out so it has no edge.
        q.height = (h3 - 0.5) * 0.05 * (1.0 + 2.0 * exp(-a / 0.15));
        q.kind = 0.0;
        q.glow = 1.0 - 0.8 * smoothstep(0.6, 1.3, a);
    }
    return q;
}

Particle spawn(uint id, constant Params& p) {
    return star(id, 0.0, p);
}

void update(thread Particle& q, uint id,
            constant Uniforms& u, constant Params& p) {
    q = star(id, u.time, p);
}

// How the galaxy is seen: the pointer across turns it about its axis, the pointer up and down leans it more or less.
// Gives the turn and the lean.
float2 viewAngles(constant Uniforms& u, constant Params& p) {
    float2 m = u.mouse / max(u.resolution, float2(1.0)) - 0.5;
    float lean = clamp(p.tilt * PI / 180.0 + m.y * 1.2, 0.0, 1.45);
    return float2(0.4 - m.x * 2.0, lean);
}

// Where a point of the galaxy is on the picture, one unit the half height.
float2 viewOf(float2 position, float height,
              constant Uniforms& u, constant Params& p) {
    float2 v = viewAngles(u, p);
    float2 r = rot(position, v.x);
    return float2(r.x, r.y * cos(v.y) + height * sin(v.y));
}

// How big the galaxy is drawn: its radius a little more than the half height (it leans, so it is less tall than
// wide), or a little short of the half width in a picture taller than wide.
float zoomOf(constant Uniforms& u) {
    return min(ZOOM, 0.95 * u.resolution.x / u.resolution.y);
}

float3 place(Particle q, constant Uniforms& u, constant Params& p) {
    float aspect = u.resolution.x / u.resolution.y;
    float2 spot = viewOf(q.position, q.height, u, p) * zoomOf(u);
    float size = q.kind == 3.0 ? 0.03 : q.kind == 2.0 ? 0.0034 : 0.0024;
    return float3(spot / float2(aspect, 1.0) * 0.5 + 0.5, size);
}

// A star's light, by kind: the old stars gold and dim, the core warmer, the young stars nearly white and bright,
// the gas ember and soft. Only light: the dots add up, and shade rolls the sum off so it never clips. Longer
// trails keep more of each step, so each dot gives less.
float4 ink(Particle q, float2 at, constant Uniforms& u, constant Params& p) {
    float d = length(at);
    float3 c = float3(1.0, 0.8, 0.55) * 0.75;
    float edge = smoothstep(1.0, 0.0, d);
    if (q.kind == 1.0) c = float3(1.0, 0.62, 0.3) * (0.11 + 0.64 * p.core);
    if (q.kind == 2.0) c = float3(1.0, 0.95, 0.88) * 2.6;
    if (q.kind == 3.0) {
        c = float3(1.0, 0.28, 0.1) * 0.1;
        edge = exp(-d * d * 4.0) * step(d, 1.0);
    }
    return float4(c * edge * q.glow * (1.0 - p.trails) * 0.5, 0.0);
}

float4 fade(float2 uv, Trails previous,
            constant Uniforms& u, constant Params& p) {
    return previous.at(uv) * p.trails;
}

// Value noise from Flowing Color Field: hashed corners eased into each other.
float vhash(float2 c) {
    uint x = uint(int(c.x) + 4096), y = uint(int(c.y) + 4096);
    return hashU(x * 7919u + y * 104729u);
}

float vnoise(float2 x) {
    float2 i = floor(x), f = x - i;
    f = f * f * (3.0 - 2.0 * f);
    return mix(mix(vhash(i), vhash(i + float2(1, 0)), f.x),
               mix(vhash(i + float2(0, 1)), vhash(i + float2(1, 1)), f.x), f.y);
}

// The dust, at a point of the disk: a lane a little inside each arm, where the gas crowds first, torn into wisps by
// two layers of noise that turn with the arms.
float dustAt(float2 x, float t, constant Params& p) {
    float r = length(x);
    float angle = atan2(x.y, x.x) - orbitTurn(r, t, p);
    float lane = pow(max(-sin(2.0 * angle - 1.5), 0.0), 12.0);
    float2 w = rot(x, -PATTERN * p.speed * t);
    float wisps = vnoise(w * 9.0) * 0.65 + vnoise(w * 23.0 + 5.0) * 0.35;
    float reach = smoothstep(0.08, 0.2, r) * (1.0 - smoothstep(0.6, 0.95, r));
    float torn = smoothstep(0.25, 0.75, wisps);
    return lane * torn * reach * smoothstep(0.0, 0.1, p.oval);
}

// The picture: the starlight, dimmed where dust lies in front, a smooth haze for the stars too many to see one by
// one, and the stars of our own sky far behind.
float4 shade(float2 uv, constant Uniforms& u, constant Params& p,
             Trails trails) {
    float aspect = u.resolution.x / u.resolution.y;
    float3 light = trails.at(uv).rgb;
    // Back from the picture to the disk: undo the zoom, the lean and the turn, for the point of the disk here.
    float2 v = viewAngles(u, p);
    float2 s = (uv - 0.5) * 2.0 * float2(aspect, 1.0) / zoomOf(u);
    float2 x = rot(float2(s.x, s.y / max(cos(v.y), 0.15)), -v.x);
    // The haze: the core's glow, round on the picture since the core is a ball, and the disk's, flat with it.
    float bulge = exp(-length(s) / 0.07) * (0.2 + 0.8 * p.core) * 0.9;
    light += float3(1.0, 0.6, 0.28) * bulge;
    light += float3(1.0, 0.7, 0.45) * exp(-length(x) / 0.3) * 0.06;
    light *= exp(-dustAt(x, u.time, p) * p.dust * 2.2);
    float3 col = 1.0 - exp(-light * 1.6);
    // Far stars: one in a few hundred little cells.
    float2 cell = floor(uv * u.resolution / 2.5);
    uint n = uint(cell.x + 8192.0) * 73856093u
           ^ uint(cell.y + 8192.0) * 19349663u;
    col += float3(1.0, 0.92, 0.8) * smoothstep(0.9975, 1.0, hashU(n)) * 0.45;
    col += float3(0.035, 0.022, 0.024);
    return float4(col, 1.0);
}

// A smooth random value from 0 to 1 that changes `rate` times a second.
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);
}

// The sound follows two stars of the disk, one near the middle on the left and one further out on the right. Each
// has a drone (D and A) that swells as its star goes into an arm, and while it is squeezed, young stars light round
// it as chimes: short notes of a D pentatonic, a little above, at random moments. Under them, the core's low hum.
float2 sound(float t, constant Params& p) {
    float radii[2] = { 0.38, 0.72 };
    float drones[2] = { 146.83, 220.0 };
    float pent[5] = { 587.33, 659.25, 739.99, 880.0, 987.77 };
    float2 v = float2(0.0);
    for (int i = 0; i < 2; i++) {
        float a = radii[i];
        float rate = orbitRate(a, p) - PATTERN * p.speed;
        float start = 1.3 + float(i) * 2.1;
        float arm = armWeight(start + rate * t, a, p);
        float pan = i == 0 ? 0.3 : 0.7;
        float x = drones[i] * t;
        float tone = sin(TAU * x) + 0.35 * sin(2.0 * TAU * x + 0.4)
                   + 0.15 * sin(3.0 * TAU * x + 1.1);
        tone *= 0.07 * (0.25 + 0.75 * arm);
        v += float2(1.0 - pan, pan) * tone;
        // Chimes, eight slots a second: a slot rings when its hash is under how squeezed the star was as it began,
        // so a note never starts halfway through.
        for (int k = 0; k < 6; k++) {
            float slot = floor(t * 8.0) - float(k);
            uint n = uint(int(slot) + 100000) * 2u + uint(i);
            float then = armWeight(start + rate * slot / 8.0, a, p);
            float chance = then * then * 0.35 * (0.3 + p.young);
            float ring = step(hashU(n * 13u + 5u), chance);
            float age = t - slot / 8.0;
            float f = pent[uint(hashU(n * 13u + 7u) * 4.99)];
            f *= i == 0 ? 1.0 : 1.5;
            float env = ring * exp(-age * 5.0) * smoothstep(0.0, 0.004, age);
            float wide = (hashU(n * 13u + 9u) - 0.5) * 0.5;
            float2 side = float2(0.5 - wide, 0.5 + wide);
            v += side * sin(TAU * f * age) * env * 0.05;
        }
    }
    float hum = sin(TAU * 73.42 * t) * (0.6 + 0.4 * noise1(t, 0.4, 3u))
              + 0.4 * sin(TAU * 36.71 * t);
    v += hum * 0.06 * (0.3 + 0.7 * p.core);
    float fadeIn = smoothstep(0.0, 2.0, t);
    return tanh(v * 5.0) / 1.3 * fadeIn * p.volume;
}

#line 1 "rk_host"
Particle Swarm::at(uint j) const { return ps[metal::min(j, count - 1)]; }

kernel void rk_particle_size(device uint* out [[buffer(0)]], uint i [[thread_position_in_grid]]) {
    if (i == 0) out[0] = sizeof(Particle);
}

kernel void rk_spawn(device Particle* ps [[buffer(0)]],
                          constant Params& p [[buffer(1)]],
                          constant uint& count [[buffer(2)]],
                          uint i [[thread_position_in_grid]]) {
    if (i >= count) return;
    ps[i] = spawn(i, p);
}

kernel void rk_update(device Particle* ps [[buffer(0)]],
                           constant Uniforms& u [[buffer(1)]],
                           constant Params& p [[buffer(2)]],
                           constant uint& count [[buffer(3)]],
                           const device Particle* last [[buffer(4)]],
                           uint i [[thread_position_in_grid]]) {
    if (i >= count) return;
    // The last step: the other buffer when update reads the swarm, the same one otherwise.
    Particle q = last[i];
    update(q, i, u, p);
    ps[i] = q;
}

fragment float4 rk_fade(RKVertexOut in [[stage_in]],
                            texture2d<float> previous [[texture(0)]],
                            constant Uniforms& u [[buffer(0)]],
                            constant Params& p [[buffer(1)]]) {
    return fade(in.uv, Trails{previous}, u, p);
}

struct RKDotOut {
    float4 position [[position]];
    float size [[point_size]];
    uint id [[flat]];
};

vertex RKDotOut rk_dot_vertex(uint id [[vertex_id]],
                                 const device Particle* ps [[buffer(0)]],
                                 constant Uniforms& u [[buffer(1)]],
                                 constant Params& p [[buffer(2)]]) {
    float3 w = place(ps[id], u, p);
    RKDotOut o;
    o.position = float4(w.xy * 2.0 - 1.0, 0.0, 1.0);
    o.size = clamp(w.z * u.resolution.y, 0.0, 511.0);
    o.id = id;
    return o;
}

fragment float4 rk_dot_fragment(RKDotOut in [[stage_in]],
                                   float2 pc [[point_coord]],
                                   const device Particle* ps [[buffer(0)]],
                                   constant Uniforms& u [[buffer(1)]],
                                   constant Params& p [[buffer(2)]]) {
    return ink(ps[in.id], float2(pc.x * 2.0 - 1.0, 1.0 - pc.y * 2.0), u, p);
}

fragment float4 rk_fragment(RKVertexOut in [[stage_in]],
                                texture2d<float> trails [[texture(0)]],
                                constant Uniforms& u [[buffer(0)]],
                                constant Params& p [[buffer(1)]]) {
    return shade(in.uv, u, p, Trails{trails});
}

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