#include "graph/prim.h"
#include "graph/search.h"
#include <errno.h>
#include <stdlib.h>

/*
 * todo priority queue, currently its broken
 * */

static int
prim_ (struct graph *g, struct graph *pt, struct state *vs, unsigned int ri,
       unsigned int (*w) (unsigned int, unsigned int))
{
  struct squeue sq;
  struct vertex *r, *n;
  struct state *rs, *ns;
  struct edge *e;
  unsigned int ni;
  int err;

  if (!(sq.q = calloc (g->nv, sizeof *sq.q)))
    return -errno;

  sq.f = sq.l = sq.q;
  *sq.f = &g->v[ri];
  vs[ri].c = RED;

  while (sq.f <= sq.l)
    {
      r = *sq.f;
      ri = vertex_id (g, r);
      rs = &vs[ri];
      e = r->adj;

      while (e && (n = e->to))
        {
          ni = vertex_id (g, n);
          ns = &vs[ni];

          switch (ns->c)
            {
            case WHITE:
              ns->c = RED;

              if ((err = add_arc (pt, ri, ni)))
                {
                  free (sq.q);
                  return err;
                }

              ns->w = w (ri, ni);
              ns->from = r;

              *(++sq.l) = n;
              break;

            case RED:
              if (ns->w < w (ri, ni))
                break;

              del_arc (pt, vertex_id (g, ns->from), ni);

              if ((err = add_arc (pt, ri, ni)))
                {
                  free (sq.q);
                  return err;
                }

              ns->w = w (ri, ni);
              ns->from = r;

              break;
            }

          e = e->next;
        }

      rs->c = BLUE;
      ++sq.f;
    }

  free (sq.q);
  return 0;
}

struct graph *
prim (struct graph *g, unsigned int (*w) (unsigned int, unsigned int))
{
  struct graph *pt;
  struct state *vs;
  unsigned int vi;

  if (!(pt = new_graph (g->nv)))
    goto err_new_graph;

  if (!(vs = calloc (g->nv, sizeof *vs)))
    goto err_calloc;

  for (vi = 0; vi < g->nv; ++vi)
    if (vs[vi].c == WHITE)
      if (prim_ (g, pt, vs, vi, w))
        goto err_prim_;

  free (vs);
  return pt;

err_prim_:
  free (vs);
err_calloc:
  free_graph (pt);
err_new_graph:
  return NULL;
}