#include <metal_stdlib>
using namespace metal;

struct VertexOut {
    float4 position [[position]];
};

struct DitherUniforms {
    float2 texScale;
    float2 viewSize;
    float spread;
    float contrast;
    float brightness;
    uint paletteCount;
    uint matrixSize;
    uint pixelSize;
};

constant float bayer2[4] = {
    0.0, 2.0,
    3.0, 1.0
};

constant float bayer4[16] = {
     0.0,  8.0,  2.0, 10.0,
    12.0,  4.0, 14.0,  6.0,
     3.0, 11.0,  1.0,  9.0,
    15.0,  7.0, 13.0,  5.0
};

constant float bayer8[64] = {
     0.0, 32.0,  8.0, 40.0,  2.0, 34.0, 10.0, 42.0,
    48.0, 16.0, 56.0, 24.0, 50.0, 18.0, 58.0, 26.0,
    12.0, 44.0,  4.0, 36.0, 14.0, 46.0,  6.0, 38.0,
    60.0, 28.0, 52.0, 20.0, 62.0, 30.0, 54.0, 22.0,
     3.0, 35.0, 11.0, 43.0,  1.0, 33.0,  9.0, 41.0,
    51.0, 19.0, 59.0, 27.0, 49.0, 17.0, 57.0, 25.0,
    15.0, 47.0,  7.0, 39.0, 13.0, 45.0,  5.0, 37.0,
    63.0, 31.0, 55.0, 23.0, 61.0, 29.0, 53.0, 21.0
};

static float thresholdAt(uint2 cell, uint size) {
    if (size <= 2u) {
        uint index = (cell.y % 2u) * 2u + (cell.x % 2u);
        return (bayer2[index] + 0.5) / 4.0;
    }
    if (size <= 4u) {
        uint index = (cell.y % 4u) * 4u + (cell.x % 4u);
        return (bayer4[index] + 0.5) / 16.0;
    }
    uint index = (cell.y % 8u) * 8u + (cell.x % 8u);
    return (bayer8[index] + 0.5) / 64.0;
}

vertex VertexOut ditherVertex(uint vertexID [[vertex_id]]) {
    float2 corners[4] = {
        float2(-1.0, -1.0),
        float2( 1.0, -1.0),
        float2(-1.0,  1.0),
        float2( 1.0,  1.0)
    };
    VertexOut out;
    out.position = float4(corners[vertexID], 0.0, 1.0);
    return out;
}

fragment float4 ditherFragment(VertexOut in [[stage_in]],
                               texture2d<float> source [[texture(0)]],
                               constant DitherUniforms &u [[buffer(0)]],
                               constant float4 *palette [[buffer(1)]])
{
    constexpr sampler linearSampler(filter::linear, mip_filter::none, address::clamp_to_edge);

    float block = float(max(u.pixelSize, 1u));
    uint2 cell = uint2(in.position.xy / block);
    float2 center = (float2(cell) + 0.5) * block;
    float2 base = center / max(u.viewSize, float2(1.0, 1.0));
    float2 uv = (base - 0.5) * u.texScale + 0.5;
    uv = clamp(uv, float2(0.0, 0.0), float2(1.0, 1.0));

    float3 color = source.sample(linearSampler, uv).rgb;
    color = (color - 0.5) * u.contrast + 0.5 + u.brightness;

    float threshold = thresholdAt(cell, u.matrixSize) - 0.5;
    color = clamp(color + threshold * u.spread, 0.0, 1.0);

    uint count = max(u.paletteCount, 1u);
    float3 weights = float3(0.299, 0.587, 0.114);
    float bestDistance = 1e9;
    float3 bestColor = palette[0].rgb;

    for (uint i = 0u; i < count; i++) {
        float3 candidate = palette[i].rgb;
        float3 delta = color - candidate;
        float distance = dot(delta * delta, weights);
        if (distance < bestDistance) {
            bestDistance = distance;
            bestColor = candidate;
        }
    }

    return float4(bestColor, 1.0);
}