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

void prim(mazegrid_t const* g, mazegrid_t* result) {
    struct pq* q = pq_new(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) pq_add(q, mazegrid_get_edge(g, 0, 0, EDGE_UP), 0, 1, EDGE_DOWN);
    if (g->width > 1) pq_add(q, mazegrid_get_edge(g, 0, 0, EDGE_RIGHT), 1, 0, EDGE_LEFT);

    while (!pq_empty(q) && (node_set_size(s) < g->width * g->height)) {
        size_t x, y;
        mazeedge_dir_t dir;
        pq_pop(q, &x, &y, &dir);

        if (node_set_contains(s, x, y)) continue;

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

        // add neighbors
        if ((x > 0) && !node_set_contains(s, x - 1, y)) {
            pq_add(q, mazegrid_get_edge(g, x, y, EDGE_LEFT), x - 1, y, EDGE_RIGHT);
        }
        if ((x < g->width - 1) && !node_set_contains(s, x + 1, y)) {
            pq_add(q, mazegrid_get_edge(g, x, y, EDGE_RIGHT), x + 1, y, EDGE_LEFT);
        }
        if ((y > 0) && !node_set_contains(s, x, y - 1)) {
            pq_add(q, mazegrid_get_edge(g, x, y, EDGE_DOWN), x, y - 1, EDGE_UP);
        }
        if ((y < g->height - 1) && !node_set_contains(s, x, y + 1)) {
            pq_add(q, mazegrid_get_edge(g, x, y, EDGE_UP), x, y + 1, EDGE_DOWN);
        }

        node_set_add(s, x, y);
    }

    node_set_free(s);
    pq_free(q);
}