#include <metal_stdlib>
using namespace metal;

/// Jump Flooding Algorithm (JFA) for Euclidean Distance Transform.
/// Pass 1 (seed): Initialize seed coordinates from binary input.
/// Pass 2..N (jump): Propagate nearest seed with decreasing step sizes.
/// Final: Compute Euclidean distance from seed coordinates.

struct EdtParams {
    uint width;
    uint height;
    int  step_size;  // Current JFA step size; 0 for seed pass, -1 for final distance pass
    uint _pad;
};

// Seed coordinate encoding: pack (x, y) into a uint2.
// Use (0xFFFF, 0xFFFF) as sentinel for "no seed".
constant uint NO_SEED = 0xFFFFu;

kernel void edt_seed(
    device const float*   input     [[buffer(0)]],
    device uint2*         seeds     [[buffer(1)]],
    constant EdtParams&   params    [[buffer(2)]],
    uint2 gid [[thread_position_in_grid]]
) {
    uint w = params.width;
    uint h = params.height;
    if (gid.x >= w || gid.y >= h) return;

    uint idx = gid.y * w + gid.x;
    // Foreground pixels (> 0.5) become seeds pointing to themselves
    if (input[idx] > 0.5f) {
        seeds[idx] = uint2(gid.x, gid.y);
    } else {
        seeds[idx] = uint2(NO_SEED, NO_SEED);
    }
}

kernel void edt_jump(
    device uint2*         seeds     [[buffer(0)]],
    device uint2*         seeds_out [[buffer(1)]],
    constant EdtParams&   params    [[buffer(2)]],
    uint2 gid [[thread_position_in_grid]]
) {
    uint w = params.width;
    uint h = params.height;
    int step = params.step_size;
    if (gid.x >= w || gid.y >= h) return;

    uint idx = gid.y * w + gid.x;
    uint2 best_seed = seeds[idx];
    float best_dist = 1e30f;

    if (best_seed.x != NO_SEED) {
        float dx = (float)gid.x - (float)best_seed.x;
        float dy = (float)gid.y - (float)best_seed.y;
        best_dist = dx * dx + dy * dy;
    }

    // Check 8 neighbors at current step distance
    for (int dy = -1; dy <= 1; dy++) {
        for (int dx = -1; dx <= 1; dx++) {
            if (dx == 0 && dy == 0) continue;

            int nx = (int)gid.x + dx * step;
            int ny = (int)gid.y + dy * step;

            if (nx < 0 || nx >= (int)w || ny < 0 || ny >= (int)h) continue;

            uint2 neighbor_seed = seeds[(uint)ny * w + (uint)nx];
            if (neighbor_seed.x == NO_SEED) continue;

            float ndx = (float)gid.x - (float)neighbor_seed.x;
            float ndy = (float)gid.y - (float)neighbor_seed.y;
            float dist = ndx * ndx + ndy * ndy;

            if (dist < best_dist) {
                best_dist = dist;
                best_seed = neighbor_seed;
            }
        }
    }

    seeds_out[idx] = best_seed;
}

kernel void edt_distance(
    device const uint2*   seeds     [[buffer(0)]],
    device float*         output    [[buffer(1)]],
    constant EdtParams&   params    [[buffer(2)]],
    uint2 gid [[thread_position_in_grid]]
) {
    uint w = params.width;
    uint h = params.height;
    if (gid.x >= w || gid.y >= h) return;

    uint idx = gid.y * w + gid.x;
    uint2 seed = seeds[idx];

    if (seed.x == NO_SEED) {
        output[idx] = sqrt((float)(w * w + h * h)); // max possible distance
    } else {
        float dx = (float)gid.x - (float)seed.x;
        float dy = (float)gid.y - (float)seed.y;
        output[idx] = sqrt(dx * dx + dy * dy);
    }
}