#include <algorithm>
#include <glm/common.hpp>
#include <glm/ext/vector_uint2.hpp>
#include <glm/ext/vector_uint3.hpp>
#include <glm/glm.hpp>
#include <raylib.h>
#include <vector>

#define MAX_PTS_PER_NODE 1 

struct kdtree
{
    kdtree(std::vector<glm::uvec2> &pts)
    {
        calc_aabb(pts);

        if (pts.size() > MAX_PTS_PER_NODE)
        {
            int split_index = aabb_size.x > aabb_size.y ? 0 : 1;
            
            std::sort(pts.begin(), pts.end(),
                      [split_index](const glm::uvec2 &a, const glm::uvec2 &b) 
                      {
                        return a[split_index] < b[split_index]; 
                      });

            {
                std::vector<glm::uvec2> lpts(pts.begin(), pts.begin() + pts.size() / 2);
                left = new kdtree(lpts);
            }

            {
                std::vector<glm::uvec2> rpts(pts.begin() + pts.size() / 2, pts.end());
                right = new kdtree(rpts);
            }
        
        }
    }

    void draw()
    {
        DrawRectangleLines(aabb_min.x, aabb_min.y, aabb_size.x, aabb_size.y, WHITE);
        if(left && right)
        {
            left->draw();
            right->draw();
        }
    }
    
    void calc_aabb(const std::vector<glm::uvec2> &pts)
    {
        for (const auto &p : pts)
        {
            aabb_min = glm::min(aabb_min, p);
            aabb_max = glm::max(aabb_max, p);
        }

        aabb_size = aabb_max - aabb_min;
    }

    kdtree *left = nullptr;
    kdtree *right = nullptr;
    glm::uvec2 aabb_min{999999};
    glm::uvec2 aabb_max{000000};
    glm::uvec2 aabb_size;
};

int main (int argc, char *argv[])
{
    auto img = LoadImage(argv[1]);

    SetConfigFlags(FLAG_WINDOW_HIDDEN);
    InitWindow(img.width, img.height, "");

    std::vector<glm::uvec2> pts;


    for (int i = 0; i < img.width * img.height; ++i)
    {
        int x = i % img.width;
        int y = i / img.width;

        auto val = ((glm::u8vec3 *)img.data + i)->r < 10;

        bool found = false;
        for(int dy = -1; dy <= 1; ++dy)
        {
            for(int dx = -1; dx <= 1; ++dx)
            {
                auto x2 = x + dx;
                auto y2 = y + dy;
                auto index = y2 * img.width + x2;

                if(index >= 0 && index < img.width * img.height)
                {
                    auto v2 = ((glm::u8vec3 *)img.data + index)->r < 10; 
                    if(val != v2)
                    {
                        found = true;
                        break;
                    }
                }
            }
            
            if(found)
            {
                break;
            }
        }
        if(found)
        {
            pts.emplace_back(x, y);
        }
    }

    kdtree tree(pts);

    auto target = LoadRenderTexture(img.width, img.height);

    BeginTextureMode(target);
    ClearBackground(BLACK);
    tree.draw();
    EndTextureMode();
  
    auto t = LoadImageFromTexture(target.texture);
    ImageFlipVertical(&t);
    ExportImage(t, argv[2]);

    CloseWindow();

    return 0;
}
