#include <metal_stdlib>
using namespace metal;

// Just Another Cube — mrange
// MSL port of just_another_cube.vk.glsl. The Shadertoy original uses
// mutable file-scope globals; in MSL we pass them by reference.

struct ShadertoyUBO {
    float iTime;
    float2 iResolution;
};

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

constant float M = 1e-3;

static float D(float3 p,
               thread float2x2& R,
               thread float& d,
               thread float& G) {
    p.xy = p.xy * R;
    p.xz = p.xz * R;

    float3 S = sin(123.0 * p);

    float shell = abs(length(p) - 0.6);
    p *= p * p * p;
    d = pow(dot(p, p), 0.125) - 0.5
        - pow(1.0 + S.x * S.y * S.z, 8.0) / 1e5;

    G = min(G, max(shell, d));

    return d;
}

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

    float2x2 R;
    float d = 1.0, z = 0.0, G = 9.0;

    float3 r = float3(U.iResolution, U.iResolution.y);
    float3 I = normalize(float3(C - 0.5 * r.xy, r.y));
    float3 B = float3(1, 2, 9) * M;
    float3 p = float3(0.0);
    float3 O = float3(0.0);

    {
        float4 cs = cos(0.3 * U.iTime + float4(0, 11, 33, 0));
        R = float2x2(cs.x, cs.y, cs.z, cs.w);
    }

    for (; z < 9.0 && d > M; z += D(p, R, d, G)) {
        p = z * I;
        p.z -= 2.0;
    }

    if (z < 9.0) {
        for (int i = 0; i < 3; ) {
            r = float3(0.0);
            r[i] = M;
            O[i++] = D(p + r, R, d, G) - D(p - r, R, d, G);
        }

        O = normalize(O);
        z = 1.0 + dot(O, I);

        r = reflect(I, O);

        C = (p + r * (5.0 - p.y) / abs(r.y)).xz;

        float3 colorTerm;
        if (r.y > 0.0) {
            d = sqrt(length(C * C)) + 1.0;
            colorTerm = 5e2 * smoothstep(5.0, 4.0, d) * d * B;
        } else {
            colorTerm = exp(-2.0 * length(C)) * (B / M - 1.0);
        }

        O = z * z * colorTerm + pow(1.0 + O.y, 5.0) * B;
    }

    float4 result = sqrt(float4(O + B / G, O.x + B.x / G));
    return result;
}