#include "pq.h"
#include "prim.h"
#include "set.h"

typedef struct {
    size_t x, y;
    mazeedge_dir_t dir;
} _point_t;

DEFINE_PQ(_point_t, PointPQ, unsigned int);

void prim(mazegrid_t const* g, mazegrid_t* result) {
    NEW_PQ_SZ(PointPQ, q, g->width*g->height*4);
    struct node_set* s = node_set_new(g->width, g->height);
    *result = mazegrid_new(g->width, g->height);

    node_set_add(s, 0, 0);
    if (g->height > 1) {
        _point_t e = { 0, 1, EDGE_DOWN };
        PUSH_PQ(q, PointPQ, e, mazegrid_get_edge(g, 0, 0, EDGE_UP));
    }
    if (g->width > 1) {
        _point_t e = { 1, 0, EDGE_LEFT };
        PUSH_PQ(q, PointPQ, e, mazegrid_get_edge(g, 0, 0, EDGE_RIGHT));
    }
    while (!EMPTY_PQ(q) && (node_set_size(s) < g->width * g->height)) {
        _point_t e;
        unsigned int p;
        POP_PQ(q, PointPQ, e, p);
        if (node_set_contains(s, e.x, e.y)) continue;

        // add edge to tree
        mazegrid_set_edge(result, e.x, e.y, e.dir, 1);

        // add neighbors
        if ((e.x > 0) && !node_set_contains(s, e.x - 1, e.y)) {
            _point_t pt = { e.x - 1, e.y, EDGE_RIGHT };
            PUSH_PQ(q, PointPQ, pt, mazegrid_get_edge(g, e.x, e.y, EDGE_LEFT));
        }
        if ((e.x < g->width - 1) && !node_set_contains(s, e.x + 1, e.y)) {
            _point_t pt = { e.x + 1, e.y, EDGE_LEFT };
            PUSH_PQ(q, PointPQ, pt, mazegrid_get_edge(g, e.x, e.y, EDGE_RIGHT));
        }
        if ((e.y > 0) && !node_set_contains(s, e.x, e.y - 1)) {
            _point_t pt = { e.x, e.y - 1, EDGE_UP };
            PUSH_PQ(q, PointPQ, pt, mazegrid_get_edge(g, e.x, e.y, EDGE_DOWN));
        }
        if ((e.y < g->height - 1) && !node_set_contains(s, e.x, e.y + 1)) {
            _point_t pt = { e.x, e.y + 1, EDGE_DOWN };
            PUSH_PQ(q, PointPQ, pt, mazegrid_get_edge(g, e.x, e.y, EDGE_UP));
        }

        node_set_add(s, e.x, e.y);
    }

    node_set_free(s);
    FREE_PQ(q);
}