#include <metal_stdlib>
using namespace metal;

// Cineshader Lava — edankwan
// MSL port of cineshader_lava.vk.glsl.

struct ShadertoyUBO {
    float iTime;
    float2 iResolution;
};

struct PassthroughVOut {
    float4 pos [[position]];
    float2 uv;
};

static float opSmoothUnion(float d1, float d2, float k) {
    float h = clamp(0.5 + 0.5 * (d2 - d1) / k, 0.0, 1.0);
    return mix(d2, d1, h) - k * h * (1.0 - h);
}

static float sdSphere(float3 p, float s) {
    return length(p) - s;
}

static float map(float3 p, float iTime) {
    float d = 2.0;
    for (int i = 0; i < 16; i++) {
        float fi = float(i);
        float time = iTime * (fract(fi * 412.531 + 0.513) - 0.5) * 2.0;
        d = opSmoothUnion(
            sdSphere(p + sin(time + fi * float3(52.5126, 64.62744, 632.25)) * float3(2.0, 2.0, 0.8),
                     mix(0.5, 1.0, fract(fi * 412.531 + 0.5124))),
            d,
            0.4
        );
    }
    return d;
}

static float3 calcNormal(float3 p, float iTime) {
    const float h = 1e-5;
    const float2 k = float2(1, -1);
    return normalize(k.xyy * map(p + k.xyy * h, iTime) +
                     k.yyx * map(p + k.yyx * h, iTime) +
                     k.yxy * map(p + k.yxy * h, iTime) +
                     k.xxx * map(p + k.xxx * h, iTime));
}

fragment float4 cineshader_lava_fs(PassthroughVOut in [[stage_in]],
                                    constant ShadertoyUBO& U [[buffer(0)]]) {
    float2 fragCoord = in.uv * U.iResolution;
    float2 uv = fragCoord / U.iResolution;

    float3 rayOri = float3((uv - 0.5) * float2(U.iResolution.x / U.iResolution.y, 1.0) * 6.0, 3.0);
    float3 rayDir = float3(0.0, 0.0, -1.0);

    float depth = 0.0;
    float3 p = float3(0.0);

    for (int i = 0; i < 64; i++) {
        p = rayOri + rayDir * depth;
        float dist = map(p, U.iTime);
        depth += dist;
        if (dist < 1e-6) { break; }
    }

    depth = min(6.0, depth);
    float3 n = calcNormal(p, U.iTime);
    float b = max(0.0, dot(n, float3(0.577)));
    float3 col = (0.5 + 0.5 * cos((b + U.iTime * 3.0) + uv.xyx * 2.0 + float3(0, 2, 4))) * (0.85 + b * 0.35);
    col *= exp(-depth * 0.15);

    return float4(col, 1.0 - (depth - 0.5) / 2.0);
}