#include <raylib.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <omp.h>

static bool is_inside_shape(Color pixel) {
    return pixel.r == 255 && pixel.g == 255 && pixel.b == 255;
}

// 1D squared Euclidean distance transform (felzenszwalb & huttenlocher)
static void edt_1d(float *f, int n, int *v, float *z, float *tmp) {
    int k = 0;
    v[0] = 0;
    z[0] = -1e30f;
    z[1] =  1e30f;

    for (int q = 1; q < n; q += 1) {
        float fq = f[q];
        float s = ((fq + (float)(q * q)) - (f[v[k]] + (float)(v[k] * v[k]))) / (float)(2 * (q - v[k]));
        while (s <= z[k]) {
            k -= 1;
            s = ((fq + (float)(q * q)) - (f[v[k]] + (float)(v[k] * v[k]))) / (float)(2 * (q - v[k]));
        }
        k += 1;
        v[k]     = q;
        z[k]     = s;
        z[k + 1] = 1e30f;
    }

    int j = 0;
    for (int q = 0; q < n; q += 1) {
        while (z[j + 1] < (float)q) j += 1;
        float d = (float)(q - v[j]);
        tmp[q] = d * d + f[v[j]];
    }
    memcpy(f, tmp, n * sizeof(float));
}

static float *compute_edt(bool *seeds, int width, int height) {
    float large = (float)((width + height) * (width + height));
    float *grid = malloc(width * height * sizeof(float));

    for (int i = 0; i < width * height; i += 1) {
        grid[i] = seeds[i] ? 0.0f : large;
    }

    int maxdim = width > height ? width : height;

    #pragma omp parallel
    {
        int   *v   = malloc(maxdim * sizeof(int));
        float *z   = malloc((maxdim + 1) * sizeof(float));
        float *tmp = malloc(maxdim * sizeof(float));
        float *col = malloc(maxdim * sizeof(float));

        // 1D EDT along each column
        #pragma omp for nowait
        for (int x = 0; x < width; x += 1) {
            for (int y = 0; y < height; y += 1) col[y] = grid[y * width + x];
            edt_1d(col, height, v, z, tmp);
            for (int y = 0; y < height; y += 1) grid[y * width + x] = col[y];
        }

        #pragma omp barrier

        // 1D EDT along each row
        #pragma omp for nowait
        for (int y = 0; y < height; y += 1) {
            memcpy(col, &grid[y * width], width * sizeof(float));
            edt_1d(col, width, v, z, tmp);
            memcpy(&grid[y * width], col, width * sizeof(float));
        }

        free(v);
        free(z);
        free(tmp);
        free(col);
    }

    return grid;
}

static Image generate_distance_field(Image input_texture, int offset_y) {
    int width  = input_texture.width;
    int height = input_texture.height;

    Color *input_pixels = LoadImageColors(input_texture);
    bool  *inside_mask  = calloc(width * height, sizeof(bool));

    for (int y = 0; y < height; y += 1) {
        int src_y = y + offset_y;
        if (src_y < 0 || src_y >= height) continue;
        for (int x = 0; x < width; x += 1) {
            inside_mask[y * width + x] = is_inside_shape(input_pixels[src_y * width + x]);
        }
    }
    UnloadImageColors(input_pixels);

    bool *outside_mask = malloc(width * height * sizeof(bool));
    for (int i = 0; i < width * height; i += 1) outside_mask[i] = !inside_mask[i];

    TraceLog(LOG_INFO, "Computing inside EDT...");
    float *dist_to_outside = compute_edt(outside_mask, width, height);

    TraceLog(LOG_INFO, "Computing outside EDT...");
    float *dist_to_inside = compute_edt(inside_mask, width, height);

    free(outside_mask);

    // pixel_value = clamp(sdf, -128, 127) + 128
    Color *out_pixels = malloc(width * height * sizeof(Color));
    for (int i = 0; i < width * height; i += 1) {
        float sdf = inside_mask[i]
            ? -sqrtf(dist_to_outside[i])  // negative = inside
            : +sqrtf(dist_to_inside[i]);  // positive = outside

        float clamped = fmaxf(-128.0f, fminf(127.0f, sdf));
        unsigned char v = (unsigned char)(clamped + 128.0f);
        out_pixels[i] = (Color){v, v, v, 255};
    }

    free(inside_mask);
    free(dist_to_inside);
    free(dist_to_outside);

    Image out = {
        .data    = out_pixels,
        .width   = width,
        .height  = height,
        .format  = PIXELFORMAT_UNCOMPRESSED_R8G8B8A8,
        .mipmaps = 1,
    };
    return out;
}

int main(int argc, char *argv[]) {
    if (argc < 3) {
        TraceLog(LOG_ERROR, "Usage: %s <input_image> <output_image> [ball_radius]", argv[0]);
        return 1;
    }

    const char *input_path  = argv[1];
    const char *output_path = argv[2];
    int ball_radius = (argc >= 4) ? atoi(argv[3]) : 0;
    int offset_y    = (ball_radius + 10) / 2;

    Image input_image = LoadImage(input_path);
    if (input_image.data == NULL) {
        TraceLog(LOG_ERROR, "Failed to load input image '%s'", input_path);
        return 1;
    }

    TraceLog(LOG_INFO, "Generating distance field (%dx%d, offset_y=%d)...", input_image.width, input_image.height, offset_y);

    Image distance_field = generate_distance_field(input_image, offset_y);

    if (!ExportImage(distance_field, output_path)) {
        TraceLog(LOG_ERROR, "Failed to save output image '%s'", output_path);
        UnloadImage(input_image);
        UnloadImage(distance_field);
        return 1;
    }

    TraceLog(LOG_INFO, "Distance field saved to '%s'", output_path);

    UnloadImage(input_image);
    UnloadImage(distance_field);

    return 0;
}
