/* Voronoi diagram generator.

   This is same as voronoi.c, but also tries to limit output
   to just just a few colors.

   ./voronoi_4colors <width> <height> > output.pgm

   Fewer colors means 4 colors, which is what Four color theorem
   says is the minimum needed.
   https://en.wikipedia.org/wiki/Four_color_theorem

   But I did left COLOR_COUNT to be a compile time constant because
   earlier implementation was buggy, and I wasn't so confident about
   using just 4 colors.  Algorithm for the current version came from
   map.c in Simon Tatham's puzzle collection:
   https://www.chiark.greenend.org.uk/~sgtatham/puzzles/
                                                                     */

#include<assert.h>
#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<time.h>
#ifdef _WIN32
   #include<fcntl.h>
   #include<io.h>
#endif

#define POINT_COUNT     50
#define COLOR_COUNT     4

/* Point locations. */
static int point[POINT_COUNT][2];

/* Adjacency matrix for all point pairs. */
static int adjacent[POINT_COUNT][POINT_COUNT];

/* Color assignment for each point. */
static int color[POINT_COUNT];

/* Initialize seed point locations. */
static void InitializeSeedPoints(int width, int height, unsigned char *pixels)
{
   int x, y, i, j;

   /* Generate random points locations and also mark seed positions. */
   for(i = 0; i < POINT_COUNT; i++)
   {
      /* We would like all seed locations to be unique, but there is some
         chance that we can't pick all unique locations if output area is
         too small, so we give up after finite number of attempts.        */
      for(j = 0; j < 128; j++)
      {
         x = (int)((double)rand() / (RAND_MAX + 1u) * width);
         y = (int)((double)rand() / (RAND_MAX + 1u) * height);
         if( pixels[y * width + x] == 0 )
            break;
      }

      pixels[y * width + x] = i + 1;
      point[i][0] = x;
      point[i][1] = y;
   }
}

/* Compute distance between two coordinate values. */
static int DistanceSquared(int ax, int ay, int bx, int by)
{
   int dx = bx - ax;
   int dy = by - ay;
   return dx * dx + dy * dy;
}

/* Assign colors to each pixel. */
static void Voronoi(int width, int height, unsigned char *pixels)
{
   int step, x, y, px, py, i, j;
   int c_neighbor, c_current, d_neighbor, d_current;

   /* Set initial step size to be largest power of 2 that is
      half of the long edge length.                          */
   for(step = 1; step * 2 < width && step * 2 < height; step *= 2);

   /* Apply jump flooding. */
   for(; step >= 1; step /= 2)
   {
      for(y = 0; y < height; y++)
      {
         for(x = 0; x < width; x++)
         {
            /* Skip current pixel if it's not colored yet. */
            c_current = pixels[y * width + x];
            if( c_current == 0 )
               continue;

            for(i = -step; i <= step; i += step)
            {
               py = y + i;
               if( py < 0 || py >= height )
                  continue;
               for(j = -step; j <= step; j += step)
               {
                  px = x + j;
                  if( px < 0 || px >= width || (i == 0 && j == 0) )
                     continue;
                  c_neighbor = pixels[py * width + px];
                  if( c_neighbor == 0 )
                  {
                     /* Neighbor pixel is not colored yet, propagate
                        color of current pixel there.                */
                     pixels[py * width + px] = c_current;
                  }
                  else
                  {
                     /* Neighbor pixel is already colored.  If neighbor
                        pixel is closer to the seed pixel of current
                        pixel, use current pixel color.                 */
                     d_neighbor = DistanceSquared(
                        px, py,
                        point[c_neighbor - 1][0], point[c_neighbor - 1][1]);
                     d_current = DistanceSquared(
                        px, py,
                        point[c_current - 1][0], point[c_current - 1][1]);
                     if( d_neighbor > d_current )
                        pixels[py * width + px] = c_current;
                  }
               }
            }
         }
      }
   }
}

/* Build adjacency matrix between all point pairs. */
static void BuildAdjacencyMatrix(int width, int height, unsigned char *pixels)
{
   int i, c_current, c_neighbor;

   /* Find all color pairs that are adjacent to each other. */
   memset(adjacent, 0, sizeof(adjacent));
   for(i = 1; i < width * height; i++)
   {
      /* Add edge for pixel to the left. */
      c_current = pixels[i] - 1;
      if( i % width != 0 )
      {
         c_neighbor = pixels[i - 1] - 1;
         adjacent[c_current][c_neighbor] = adjacent[c_neighbor][c_current] = 1;
      }

      /* Add edge for pixel above. */
      if( i >= width )
      {
         c_neighbor = pixels[i - width] - 1;
         adjacent[c_current][c_neighbor] = adjacent[c_neighbor][c_current] = 1;
      }
   }
}

/* Assign colors recursively until all colors are assigned.

   Returns 1 on success.                                    */
static int RecursiveAssignColors()
{
   int neighbor, i;
   int lowest_degree_of_freedom = COLOR_COUNT + 1;
   int degree_of_freedom;
   int selected_node = -1;
   int colors_used[POINT_COUNT];

   /* Find node that does not have color assigned yet, keeping the
      one with fewest degree of freedom.                           */
   memset(colors_used, 0, sizeof(colors_used));
   for(i = 0; i < POINT_COUNT; i++)
   {
      if( color[i] != 0 )
         continue;

      degree_of_freedom = COLOR_COUNT;
      for(neighbor = 0; neighbor < POINT_COUNT; neighbor++)
      {
         if( i == neighbor ||
             adjacent[i][neighbor] == 0 ||
             color[neighbor] == 0 )
         {
            continue;
         }
         if( (colors_used[i] & (1 << color[neighbor])) == 0 )
         {
            colors_used[i] |= (1 << color[neighbor]);
            degree_of_freedom--;

            /* Return failure if we found an uncolored node
               that has no possible coloring.               */
            if( degree_of_freedom == 0 )
               return 0;
         }
      }
      if( lowest_degree_of_freedom > degree_of_freedom )
      {
         lowest_degree_of_freedom = degree_of_freedom;
         selected_node = i;
      }
   }

   /* Return success if there are no uncolored nodes left. */
   if( selected_node < 0 )
      return 1;

   /* Try different colors for the selected node. */
   for(i = 1; i <= COLOR_COUNT; i++)
   {
      if( (colors_used[selected_node] & (1 << i)) != 0 )
         continue;
      color[selected_node] = i;
      if( RecursiveAssignColors() )
         return 1;
      color[selected_node] = 0;
   }

   /* Couldn't find a valid color assignment. */
   return 0;
}

/* Assign colors to each seed point. */
static void AssignColors(void)
{
   int i;

   memset(color, 0, sizeof(color));
   if( !RecursiveAssignColors() )
      fputs("Failed to assign colors\n", stderr);

   /* Scale colors. */
   for(i = 0; i < POINT_COUNT; i++)
      color[i] = (color[i] - 1) * 255 / (COLOR_COUNT - 1);
}


int main(int argc, char **argv)
{
   int width, height, i;
   unsigned char *pixels;

   assert(COLOR_COUNT < 30);
   assert(POINT_COUNT > 0);
   assert(POINT_COUNT <= 254);

   if( argc != 3 ||
       (width = atoi(argv[1])) <= 0 ||
       (height = atoi(argv[2])) <= 0 )
   {
      return printf("%s <width> <height>\n", *argv);
   }
   if( width >= 0x8000 || height >= 0x8000 )
      return !puts("Output size too large.");

   if( (pixels = (unsigned char*)calloc(width * height, 1)) == NULL )
      return !puts("Not enough memory");

   #ifdef _WIN32
      setmode(STDOUT_FILENO, O_BINARY);
   #endif

   srand(time(NULL));

   InitializeSeedPoints(width, height, pixels);
   Voronoi(width, height, pixels);
   BuildAdjacencyMatrix(width, height, pixels);
   AssignColors();

   /* Write pixels. */
   printf("P5\n%d %d\n255\n", width, height);
   for(i = 0; i < width * height; i++)
      fputc(color[pixels[i] - 1], stdout);
   free(pixels);
   return 0;
}