#include <metal_stdlib>
using namespace metal;

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

struct Params {
    float tempo;
    float ratio;
    float index;
    float attack;
    float decay;
    float sustain;
    float release;
    float detune;
    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 · Pulse Synth
//
// @episode Pulse Synth
// @length 60
// @variation 30
// @variations 2,3,4,7,8,9,10,11
// @still 9.64
// @teaches FM synthesis and envelopes, a picture drawn from its own waveform
// @category Fragment
// @tags FM synthesis, envelope, oscilloscope, Bessel functions, procedural sound
// @sweeps index, ratio, detune
// @final Pulse Synth | Two channels a few cents apart, drawn as two glowing traces that slide past each other, over the envelope and the sidebands.
//
// A small synthesizer and the oscilloscope watching it. The sound plays an arpeggio, every note one sine wave whose
// phase is bent by another (frequency modulation), its loudness shaped by an envelope. The picture calls the sound's
// own function: the traces across the middle are the samples you hear, left and right, a few milliseconds of them.
// Above them is the note's envelope, and below them the sidebands the bending makes, from the same numbers.
// 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 tempo 100 60 160 Beats a minute; the notes are eighths
// @param ratio 2 0.5 5 The modulator's frequency, as a multiple of the note's
// @param index 3 0 6 How hard the modulator bends the note: 0 is a pure sine, more is brighter
// @param attack 0.01 0.003 0.3 Seconds each note takes to rise
// @param decay 0.25 0.05 1 Seconds, roughly, to fall from the peak to the sustain
// @param sustain 0.4 0 1 The level a note holds while its key is down
// @param release 0.3 0.02 0.6 Seconds, roughly, a note takes to fade once let go
// @param detune 0.25 0 1 How far apart left and right are, up to 8 cents
// @param volume 0.8 0 1 Loudness of the sound

constant float TAU = 6.2831853;

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

uint hashBits(uint x) {
    x ^= x >> 16;
    x *= 0x7feb352du;
    x ^= x >> 15;
    x *= 0x846ca68bu;
    x ^= x >> 16;
    return x;
}

// ---- The sequencer -------------------------------------------------------------------------------------------------
// Eighth notes, eight to a bar; each bar arpeggiates a chord of A minor pentatonic.

float noteLength(constant Params& p) {
    return 30.0 / max(p.tempo, 1.0);
}

// A minor pentatonic from A2: three octaves.
float scaleNote(int i) {
    const float notes[15] = {110.00, 130.81, 146.83, 164.81, 196.00,
                             220.00, 261.63, 293.66, 329.63, 392.00,
                             440.00, 523.25, 587.33, 659.25, 783.99};
    return notes[clamp(i, 0, 14)];
}

// Step n's pitch: the bar's chord, then the arpeggio's shape.
float noteFrequency(int n) {
    const int roots[4] = {5, 8, 7, 6};
    const int shape[8] = {0, 2, 4, 5, 7, 5, 4, 2};
    return scaleNote(roots[(n / 8) % 4] + shape[n % 8]);
}

// The first note of a bar always plays; about one in five others rests.
bool notePlays(int n) {
    if (n < 0) return false;
    if (n % 8 == 0) return true;
    return (hashBits(uint(n % 32) + 777u) & 255u) > 50u;
}

// ---- The envelope --------------------------------------------------------------------------------------------------
// Attack, decay, sustain, release. The key is held for 60% of a step (the gate): the level rises in a straight line,
// falls towards the sustain, and once let go falls from wherever it was towards silence.

float gateLength(constant Params& p) {
    return 0.6 * noteLength(p);
}

float heldLevel(float since, constant Params& p) {
    float a = max(p.attack, 0.003);
    if (since < a) return since / a;
    float fall = exp(-(since - a) / max(p.decay, 0.01));
    return mix(p.sustain, 1.0, fall);
}

float adsr(float since, constant Params& p) {
    if (since < 0.0) return 0.0;
    float g = gateLength(p);
    if (since < g) return heldLevel(since, p);
    return heldLevel(g, p) * exp(-(since - g) / max(p.release, 0.01));
}

// ---- The voice -----------------------------------------------------------------------------------------------------
// Frequency modulation (John Chowning, 1973): a sine whose phase is pushed back and forth by a second sine, the
// modulator, at `ratio` times its frequency, by up to `index` radians: strongest at the strike, then 30%.

float modIndex(float since, constant Params& p) {
    return p.index * (0.3 + 0.7 * exp(-since * 6.0));
}

float fmVoice(float f, float since, constant Params& p) {
    float phase = f * since;
    float bend = modIndex(since, p) * sin(TAU * p.ratio * phase);
    return sin(TAU * phase + bend) * adsr(since, p);
}

// ---- The mix -------------------------------------------------------------------------------------------------
// Every note of the last twelve steps, summed; the oldest let go over three steps, so it never clicks. The right
// side plays a few cents sharp, so the two drift in and out of step: a chorus. synth() leaves the volume out.

float2 synth(float t, constant Params& p) {
    float len = noteLength(p);
    int now = int(floor(t / len));
    // The right side plays each note a few cents sharp.
    float sharp = exp2(p.detune * 8.0 / 1200.0);
    float2 sum = float2(0.0);
    for (int back = 0; back < 12; back++) {
        int n = now - back;
        if (n < 0) break;
        if (!notePlays(n)) continue;
        float since = t - float(n) * len;
        float held = 1.0 - smoothstep(9.0, 12.0, since / len);
        float f = noteFrequency(n);
        float2 pair = float2(fmVoice(f, since, p), fmVoice(f * sharp, since, p));
        sum += held * pair;
    }
    float fadeIn = smoothstep(0.0, 1.5, t);
    return tanh(sum * 0.25) * fadeIn;
}

// ---- The picture ---------------------------------------------------------------------------------------------------

// The Bessel function J_n(x), from its series: the sum over m of (-1)^m (x/2)^(2m+n) / (m! (m+n)!). The sideband
// k modulator steps from the note is |J_k(index)| big.
float besselJ(int n, float x) {
    float h = 0.5 * x;
    float term = 1.0;
    for (int i = 1; i <= n; i++) term *= h / float(i);
    float sum = term;
    for (int m = 1; m < 16; m++) {
        term *= -h * h / (float(m) * float(m + n));
        sum += term;
    }
    return sum;
}

// The latest note to have started: its start time and frequency.
float2 latestNote(float t, constant Params& p) {
    float len = noteLength(p);
    int now = int(floor(t / len));
    for (int back = 0; back < 12; back++) {
        int n = now - back;
        if (n < 0) break;
        if (notePlays(n)) return float2(float(n) * len, noteFrequency(n));
    }
    return float2(0.0, noteFrequency(0));
}

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

// 1 on a line d away, fading out over w; and 1 when v is between a and b.
float lineAt(float d, float w) {
    return 1.0 - smoothstep(0.0, w, d);
}

float within(float v, float a, float b) {
    return step(a, v) * step(v, b);
}

// The scope's graticule: faint lines every tenth of the height, the middle
// one brighter.
float graticule(float2 uv, float2 q, float px) {
    float2 cell = abs(fract(q / 0.1 + 0.5) - 0.5) * 0.1;
    float grid = lineAt(min(cell.x, cell.y), 1.5 * px);
    float heat = 0.05 * grid * within(uv.y, 0.2, 0.8);
    return heat + 0.06 * lineAt(abs(uv.y - 0.5), 1.5 * px);
}

// How far q is from each trace, left and right: 9 ms of synth() from
// `start`, across the width. Seven samples 3 pixels apart, joined by lines.
float2 traceDistance(float2 q, float start, float aspect, float px,
                     constant Params& p) {
    float2 off = float2(1.0);
    float2 last = float2(0.0);
    for (int i = -3; i <= 3; i++) {
        float x = q.x + float(i) * 3.0 * px;
        float share = clamp((x / aspect - 0.04) / 0.92, 0.0, 1.0);
        float2 s = synth(start + share * 0.009, p) * 0.6 + 0.5;
        if (i > -3) {
            float x0 = x - 3.0 * px;
            off.x = min(off.x, segment(q, float2(x0, last.x), float2(x, s.x)));
            off.y = min(off.y, segment(q, float2(x0, last.y), float2(x, s.y)));
        }
        last = s;
    }
    return off;
}

// The envelope, top left: the latest note's shape from strike to silence,
// the floor, a mark where the key is let go, and a dot where it is now.
float envelopePanel(float2 q, float since, float aspect, float px,
                    constant Params& p) {
    float2 box = float2(0.04 * aspect, 0.84);
    float2 size = float2(0.36, 0.11);
    float span = gateLength(p) + 3.0 * p.release;
    float ex = (q.x - box.x) / size.x;
    float on = within(ex, 0.0, 1.0);
    float heat = 0.0;
    if (q.y > box.y - 0.02 && q.y < box.y + size.y + 0.02) {
        float ea = clamp(ex, 0.0, 1.0), eb = ea + px / size.x;
        float2 a = box + float2(ea, adsr(ea * span, p)) * size;
        float2 b = box + float2(eb, adsr(eb * span, p)) * size;
        heat += 0.5 * lineAt(segment(q, a, b), 1.5 * px) * on;
        heat += 0.08 * lineAt(abs(q.y - box.y), px) * on;
        float gx = box.x + gateLength(p) / span * size.x;
        float tall = within(q.y, box.y, box.y + size.y);
        heat += 0.1 * lineAt(abs(q.x - gx), px) * tall;
    }
    float2 now = float2(clamp(since / span, 0.0, 1.0), adsr(since, p));
    float d = length(q - (box + now * size));
    return heat + 0.9 * (1.0 - smoothstep(4.0 * px, 5.5 * px, d));
}

// The sidebands, along the bottom: energy at the note's frequency plus and
// minus whole multiples of the modulator's, each |J_k(index)| tall, as
// bright as the note is loud. 0 to 4 kHz across; below zero folds back.
float sidebands(float2 uv, float2 note, float since, float aspect, float px,
                constant Params& p) {
    float floorY = 0.04, tall = 0.18;
    if (uv.y < floorY - 0.01 || uv.y > floorY + tall + 0.01) return 0.0;
    if (uv.x < 0.03 || uv.x > 0.97) return 0.0;
    float hz = (uv.x - 0.04) / 0.92 * 4000.0;
    float I = modIndex(since, p);
    float level = adsr(since, p);
    float bars = 0.0;
    for (int k = -8; k <= 8; k++) {
        float f = abs(note.y + float(k) * p.ratio * note.y);
        float w = abs(hz - f) / 4000.0 * 0.92 * aspect;
        if (w > 3.0 * px) continue;
        float a = min(abs(besselJ(abs(k), I)), 1.0);
        float bar = 1.0 - smoothstep(px, 2.5 * px, w);
        bar *= within(uv.y, floorY, floorY + tall * a);
        float bright = 0.3 + 0.5 * level + 0.15 * float(k == 0);
        bars = max(bars, bar * bright);
    }
    float axis = lineAt(abs(uv.y - floorY), px) * within(uv.x, 0.04, 0.96);
    return bars + 0.08 * axis;
}

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;
    float2 note = latestNote(t, p);
    float since = t - note.x;
    float heat = 0.03 + graticule(uv, q, px);

    // The traces start where the latest note's wave last began a cycle, so
    // a steady wave stands still, as a scope's trigger holds it.
    float start = note.x + floor(since * note.y) / note.y;
    float2 off = traceDistance(q, start, aspect, px, p);
    float inside = within(uv.x, 0.04, 0.96);
    float2 core = (1.0 - smoothstep(0.0, 2.0 * px, off)) * inside;
    float2 glow = exp(-off / (8.0 * px)) * inside;
    // Right in ember, left over it in gold: where they agree, they burn white.
    heat += 0.45 * core.y + 0.14 * glow.y;
    heat += 0.55 * core.x + 0.18 * glow.x;

    heat += envelopePanel(q, since, aspect, px, p);
    heat += sidebands(uv, note, since, aspect, px, p);

    // The whole picture warms a little with the note.
    float level = adsr(since, p);
    heat += 0.06 * level * exp(-abs(uv.y - 0.5) * 4.0);
    return float4(kiln(heat), 1.0);
}

// ---- Sound -----------------------------------------------------------------------------------------------------
// What you hear is synth(), the same samples the picture draws, at the volume the slider sets.

float2 sound(float t, constant Params& p) {
    return synth(t, p) * 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);
}