#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<time.h>
#include<png.h>

#define REGION_COUNT 50

/* Adjacency matrix between different color regions. */
static int adjacency_matrix[REGION_COUNT][REGION_COUNT];

/* Color assignments for each region (1-based). */
static int color_map[REGION_COUNT];

/* Assign colors recursively.  Returns 1 on success. */
static int RecursiveAssignColors()
{
   /* Set of colors that are adjacent to each region, one bit for each color. */
   int colors_used[REGION_COUNT];

   /* Lowest number of available colors for most constrained region. */
   int lowest_degree_of_freedom = 5;

   /* Selected region index. */
   int selected_region = -1;

   int i, j, degree_of_freedom;

   /* Find region that don't have a color map entry yet, keeping
      one with the lowest degree of freedom.                     */
   memset(colors_used, 0, sizeof(colors_used));
   for(i = 0; i < REGION_COUNT; i++)
   {
      if( color_map[i] != 0 )
         continue;
      degree_of_freedom = 4;
      for(j = 0; j < REGION_COUNT; j++)
      {
         if( i == j || adjacency_matrix[i][j] == 0 || color_map[j] == 0 )
            continue;
         if( (colors_used[i] & (1 << color_map[j])) == 0 )
         {
            /* Neighbor used a color that is not used by current region,
               mark that color as unavailable for use.                   */
            colors_used[i] |= (1 << color_map[j]);
            degree_of_freedom--;
            if( degree_of_freedom == 0 )
               return 0;  /* No colors available for current region. */
         }
      }
      if( lowest_degree_of_freedom > degree_of_freedom )
      {
         lowest_degree_of_freedom = degree_of_freedom;
         selected_region = i;
      }
   }

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

   /* Try different colors for the selected region. */
   for(i = 1; i <= 4; i++)
   {
      if( (colors_used[selected_region] & (1 << i)) != 0 )
         continue;
      color_map[selected_region] = i;
      if( RecursiveAssignColors() )
         return 1;

      /* Failed to make progress with this selected color, backtrack. */
      color_map[selected_region] = 0;
   }

   /* None of the color assignments worked. */
   return 0;
}

static int DistanceSquared(int ax, int ay, int bx, int by)
{
   int dx = bx - ax;
   int dy = by - ay;
   return dx * dx + dy * dy;
}

int main(int argc, char **argv)
{
   int seed_point[REGION_COUNT][2];
   png_image image;
   png_bytep buffer1, buffer2;
   int s, i, j, x, y, px, py, offset, step, c_current, c_neighbor, *pixel;

   if( argc != 4 )
      return printf("%s <input.png> <output1.png> <output2.png>\n", *argv);

   memset(&image, 0, sizeof(image));
   image.version = PNG_IMAGE_VERSION;
   if( !png_image_begin_read_from_file(&image, argv[1]) )
      return printf("Error reading %s\n", argv[1]);

   image.format = PNG_FORMAT_RGBA;
   s = image.width * image.height;
   buffer1 = (png_bytep)malloc(s * 4);
   buffer2 = (png_bytep)calloc(s, 4);
   if( buffer1 == NULL || buffer2 == NULL )
      return puts("Out of memory");
   if( !png_image_finish_read(&image, NULL, buffer1, 0, NULL) )
      return printf("Error loading %s\n", argv[1]);

   srand(time(NULL));

   /* Initialize seed points. */
   pixel = (int*)buffer2;
   for(i = 0; i < REGION_COUNT; i++)
   {
      /* Prefer points that are distinct, but give up after finite number
         of tries.                                                        */
      for(j = 0; j < 99; j++)
      {
         x = (int)((double)rand() / (RAND_MAX + 1u) * image.width);
         y = (int)((double)rand() / (RAND_MAX + 1u) * image.height);
         if( pixel[offset = y * image.width + x] == 0 )
            break;
      }
      pixel[offset] = i + 1;
      seed_point[i][0] = x;
      seed_point[i][1] = y;
   }

   /* Generate voronoi coloring using jumping flood algorithm.
      https://en.wikipedia.org/wiki/Jump_flooding_algorithm    */
   for(step = 1;
       step * 2 < (int)image.width && step * 2 < (int)image.height; step *= 2);
   for(; step >= 1; step /= 2)
   {
      for(y = offset = 0; y < (int)image.height; y++)
      {
         for(x = 0; x < (int)image.width; x++, offset++)
         {
            /* Skip uncolored pixel. */
            c_current = pixel[offset];
            if( c_current == 0 )
               continue;

            for(i = -step; i <= step; i += step)
            {
               py = y + i;
               if( py < 0 || py >= (int)image.height )
                  continue;
               for(j = -step; j <= step; j += step)
               {
                  px = x + j;
                  if( px < 0 || px >= (int)image.width )
                     continue;

                  c_neighbor = pixel[py * image.width + px];
                  if( c_neighbor == 0 )
                  {
                     /* Neighbor pixel is not colored yet, propagate
                        color of current pixel.                      */
                     pixel[py * image.width + px] = c_current;
                  }
                  else
                  {
                     /* Neighbor pixel is already colored, overwrite
                        with current pixel color if current seed pixel
                        is closer.                                     */
                     if( DistanceSquared(px,
                                         py,
                                         seed_point[c_neighbor - 1][0],
                                         seed_point[c_neighbor - 1][1]) >
                         DistanceSquared(px,
                                         py,
                                         seed_point[c_current - 1][0],
                                         seed_point[c_current - 1][1]) )
                     {
                        pixel[py * image.width + px] = c_current;
                     }
                  }
               }
            }
         }
      }
   }

   /* Build adjacency matrix for all regions. */
   memset(adjacency_matrix, 0, sizeof(adjacency_matrix));
   for(i = 0; i < s; i++)
   {
      c_current = pixel[i] - 1;
      if( i % image.width != 0 )
      {
         /* Add edge for pixel to the left. */
         c_neighbor = pixel[i - 1] - 1;
         adjacency_matrix[c_current][c_neighbor] =
         adjacency_matrix[c_neighbor][c_current] = 1;
      }
      if( i >= (int)image.width )
      {
         /* Add edge for pixel above. */
         c_neighbor = pixel[i - image.width] - 1;
         adjacency_matrix[c_current][c_neighbor] =
         adjacency_matrix[c_neighbor][c_current] = 1;
      }
   }

   /* Assign colors. */
   memset(color_map, 0, sizeof(color_map));
   RecursiveAssignColors();

   /* Remap region colors to reduced color set. */
   for(i = 0; i < s; i++)
      pixel[i] = color_map[pixel[i] - 1] - 1;

   /* Move pixels. */
   for(i = 0; i < s; i++)
   {
      if( (double)rand() / RAND_MAX * 3 < pixel[i] )
      {
         memcpy(buffer2 + (i * 4), buffer1 + (i * 4), 4);
         memset(buffer1 + (i * 4), 0, 4);
      }
   }

   if( !png_image_write_to_file(&image, argv[2], 0, buffer1, 0, NULL) )
      return printf("Error writing %s\n", argv[2]);
   if( !png_image_write_to_file(&image, argv[3], 0, buffer2, 0, NULL) )
      return printf("Error writing %s\n", argv[3]);

   return 0;
}