#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 power;
    float breathe;
    float iterations;
    float trap;
    float glow;
    float heat;
    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 20 · Mandelbulb
//
// @episode 20 Mandelbulb
// @length 60
// @variation 30
// @variations 1,2,3,6,7,8,9,11
// @sweeps power, breathe, iterations
// @still 21.0
// @teaches 3D fractals, distance estimators, orbit-trap colour
// @category Fragment
// @tags mandelbulb, fractals, distance estimators, orbit traps, raymarching, ambient occlusion, procedural sound
// @final Mandelbulb | The bulb glowing in the kiln's dark, its power breathing so the buds swell and fold, coloured by how near each point's orbit came to its equator.
//
// Episode 7 squared complex numbers over and over and coloured each point by whether it escaped. The Mandelbulb does
// the same in three dimensions: a point is turned into spherical coordinates, its distance raised to a power and its
// two angles multiplied by it, and the starting point added back. Points that never escape make the bulb. A ray can't
// step through that set by guessing, so the loop also keeps a running derivative, and from it a distance estimate:
// how far the ray can safely step before it could touch the surface. Colour comes from the orbit itself: how close
// each point's path came to the bulb's equator (its orbit trap).
// 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 0.5 0.0 2.0 How fast the camera circles and the power breathes
// @param power 8.0 3.0 12.0 The power each point is raised to: 3 is a folded lump, 8 the classic bulb, 12 a bristling one
// @param breathe 0.4 0.0 1.0 How far the power swells and shrinks around its setting as time goes on
// @param iterations 8.0 1.0 12.0 How many times each point is raised to the power: 1 is a smooth lump, 2 grows buds, 3 buds on them
// @param trap 0.7 0.0 1.0 How much the orbit trap colours the bulb, against plain light and shade
// @param glow 0.5 0.0 1.0 How brightly rays that pass near the bulb glow
// @param heat 0.5 0.0 1.0 Slides the palette from deep ember to white-hot
// @param volume 0.8 0.0 1.0 Loudness of the sound

constant float TAU = 6.2831853;

// The kiln palette from episode 1.
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 bulb ---------------------------------------------------------------------------------------------------

// The distance estimate and orbit trap from stages 1 and 3, with up to 12
// turns: iterations is now a slider.
float bulb(float3 p, float power, int iterations, thread float& trap) {
    float3 z = p;
    float dr = 1.0;
    float r = max(length(z), 1e-6);
    trap = 1e3;
    for (int i = 0; i < 12; i++) {
        if (i >= iterations || r > 2.0) break;
        float theta = acos(clamp(z.y / r, -1.0, 1.0)) * power;
        float phi = atan2(z.z, z.x) * power;
        dr = pow(r, power - 1.0) * power * dr + 1.0;
        float zr = pow(r, power);
        z = zr * float3(sin(theta) * cos(phi), cos(theta),
                        sin(theta) * sin(phi)) + p;
        r = max(length(z), 1e-6);
        trap = min(trap, abs(z.y));
    }
    return 0.5 * log(r) * r / dr;
}

float bulbDist(float3 p, float power, int iterations) {
    float trap;
    return bulb(p, power, iterations, trap);
}

// The normal from four samples, as before.
float3 bulbNormal(float3 p, float power, int iterations, float e) {
    const float2 k = float2(1.0, -1.0);
    return normalize(k.xyy * bulbDist(p + k.xyy * e, power, iterations)
                   + k.yyx * bulbDist(p + k.yyx * e, power, iterations)
                   + k.yxy * bulbDist(p + k.yxy * e, power, iterations)
                   + k.xxx * bulbDist(p + k.xxx * e, power, iterations));
}

// Ambient occlusion, as in stage 2.
float occlusion(float3 p, float3 n, float power, int iterations) {
    float shut = 0.0;
    float weight = 1.0;
    for (int i = 1; i <= 5; i++) {
        float h = 0.01 + 0.03 * float(i);
        shut += (h - bulbDist(p + n * h, power, iterations)) * weight;
        weight *= 0.75;
    }
    return clamp(1.0 - 3.0 * shut, 0.0, 1.0);
}

// Where the ray enters and leaves a sphere of radius 1.6 around the bulb.
// Nothing outside it is marched.
float2 bound(float3 ro, float3 rd) {
    float b = dot(ro, rd);
    float c = dot(ro, ro) - 1.6 * 1.6;
    float h = b * b - c;
    if (h < 0.0) return float2(1.0, -1.0);
    h = sqrt(h);
    return float2(max(-b - h, 0.0), -b + h);
}

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

float4 shade(float2 uv, constant Uniforms& u, constant Params& p) {
    float t = u.time * p.speed;
    // The power breathes slowly around its setting, so the buds swell and fold.
    float power = max(p.power + p.breathe * 1.5 * sin(t * 0.4), 2.0);
    int iterations = int(p.iterations + 0.5);

    // The camera circles the bulb. The pointer's offset from the middle (zero in
    // every video) turns it further round and tips it up or down.
    float2 steer = u.mouse / u.resolution - 0.5;
    float yaw = t * 0.3 + steer.x * TAU * 0.5;
    float pitch = 0.35 + 0.25 * sin(t * 0.17) + steer.y * 2.0;
    pitch = clamp(pitch, -1.4, 1.4);
    float3 ro = 2.6 * float3(cos(pitch) * sin(yaw), sin(pitch),
                             cos(pitch) * cos(yaw));
    float3 fw = normalize(-ro);
    float3 rt = normalize(cross(fw, float3(0.0, 1.0, 0.0)));
    float3 up = cross(rt, fw);
    float side = min(u.resolution.x, u.resolution.y);
    float2 q = (uv - 0.5) * u.resolution / side * 2.0;
    float3 rd = normalize(q.x * rt + q.y * up + 2.0 * fw);
    // How wide one pixel is a unit away: the march stops this close to the
    // surface, so detail is as fine as the picture is sharp (a zoomed picture
    // has more pixels, and finer detail).
    float pixel = 1.0 / side;

    float shift = (p.heat - 0.5) * 0.4;
    // The kiln's dark behind it, warmer toward the middle of the picture.
    float back = 0.06 + 0.10 * exp(-dot(q, q) * 0.6);

    float2 span = bound(ro, rd);
    if (span.y < span.x) return float4(kiln(back + shift), 1.0);

    float dist = span.x;
    float nearest = 1e3;
    bool hit = false;
    float trap = 1.0;
    for (int i = 0; i < 160; i++) {
        float3 pos = ro + rd * dist;
        float d = bulb(pos, power, iterations, trap);
        nearest = min(nearest, d);
        if (d < pixel * dist) { hit = true; break; }
        dist += d * 0.9;
        if (dist > span.y) break;
    }

    if (!hit) {
        // A ray that missed glows by how near it passed the bulb.
        float halo = exp(-nearest * 18.0) * p.glow * 0.55;
        return float4(kiln(back + halo + shift), 1.0);
    }

    float3 pos = ro + rd * dist;
    float3 n = bulbNormal(pos, power, iterations, pixel * dist * 0.5);
    // The light is above and to the left of the camera, and goes round with it.
    float3 l = normalize(up * 0.8 - rt * 0.5 - fw * 0.4);
    float lit = max(dot(n, l), 0.0);
    float ao = occlusion(pos, n, power, iterations);
    // The rim: edges facing away from the camera catch the kiln's glow.
    float rim = pow(1.0 - max(dot(n, -rd), 0.0), 3.0);
    float shadeLevel = (0.35 + 0.65 * lit) * (0.4 + 0.6 * ao)
                     + 0.25 * rim * p.glow;
    // The orbit trap: points whose orbit passed close to the equator run hot,
    // the ones that kept away from it stay dark. Mixed with plain light and shade.
    float orbit = 1.0 - sqrt(clamp(trap * 1.2, 0.0, 1.0));
    float level = mix(0.15 + 0.75 * shadeLevel,
                      0.1 + 0.6 * shadeLevel + 0.6 * orbit * orbit, p.trap);
    return float4(kiln(level + shift), 1.0);
}

// ---- Sound -------------------------------------------------------------------------------------------------------
// A fractal looks the same however far in you go; this sound does the same in pitch. Six tones an octave apart glide
// up together, each fading in at the bottom and out at the top (a Shepard-Risset glide), so the whole seems to rise
// forever without going anywhere. A second glide a fifth above follows the orbit trap's colour, and the breathing
// power swells both.

float glide(float t, float rate, float base, float pan) {
    float a = 0.0;
    for (int k = 0; k < 6; k++) {
        // x runs 0 to 1 as this tone climbs six octaves, then starts again.
        float x = fract(t * rate + float(k) / 6.0);
        // The phase is the frequency's integral: base * 2^(6x) / (6 * rate * ln 2).
        float phase = base * exp2(6.0 * x) / (6.0 * rate * 0.6931472);
        float env = sin(3.14159265 * x);
        a += sin(TAU * fract(phase) + pan * float(k)) * env * env;
    }
    return a / 3.0;
}

float2 sound(float t, constant Params& p) {
    float rate = 0.004 + 0.01 * p.speed;
    float swell = 0.8 + 0.2 * sin(t * p.speed * 0.4) * p.breathe;
    float2 low = float2(glide(t, rate, 55.0, 0.0),
                        glide(t + 0.37, rate, 55.0, 0.9));
    float2 fifth = float2(glide(t + 0.21, rate, 82.41, 0.4),
                          glide(t, rate, 82.41, 1.3));
    float2 mixed = low + fifth * 0.6 * p.trap;
    float fadeIn = smoothstep(0.0, 2.0, t);
    return tanh(mixed * 0.5 * swell) * 0.8 * fadeIn * 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);
}