#include <stdio.h>
#include <stdlib.h>

#define UNDEF -1
#define INFINITY 1073741824
#define MAX_SEG_NO 16

/* We store the RGB components in three two-dimensional arrays. The row
   number i grows vertically, and the column number j horozontally. We
   identify a pixel by these indices (i,j). This is in contrast with the
   common convention of representing a pixel by the coordinates (x,y)
   with x in the horizontal direction and y in the vertical direction,
   and with the top left pixel the origin. We denote (i,j) as the
   coordinates of a pixel indicating y = i and x = j. We will always
   work with (i,j) and never with (x,y). */

/* Data type for image */
typedef struct {
   int wd, ht;                                            /* Width and height */
   unsigned char **R, **G, **B;    /* Three 2-d arrays for the RGB components */
} image;

/* Each pixel has eight neighbors. We plan to store each edge in the
   graph only once. For each vertex/pixel (i,j), we store only four of
   its neighbors (i-1,j+1), (i,j+1), (i+1,j+1) and (i+1,j). In that order,
   the pixel distances are stored in the cost array. The edges corresponding
   to the other four neighbors are stored against other pixels. For example,
   the neighbor (i,j-1) stores the distance between (i,j-1) and (i,j) at
   index 1 of the cost array of (i,j-1). Pixels at the image boundary do
   not have all the neighbors. The corresponding cost entries are set
   to UNDEF. 

   We use the same data type also to store trees and forests. Every
   edge present will let the corresponding cost entry store the pixel
   distance standing for the edge. A non-existent tree or forest edge
   will correspond to a storage of UNDEF in the appropriate cost entry. */

typedef int cost[4];

/* Data type graph (also to represent trees and forests) */
typedef struct {
   int wd, ht;  /* Width and height of image -- needed for indexing nbr array */
   int e;                  /* Number of edges -- needed for allocating memory */
   cost **nbr;          /* Adjacency list: 2-d array of 4-element cost arrays */
   int **flag; /* Needed for a variety of purposes -- see the functions below */
} graph;

/* A heap needs to be constructed on the vertex set of the graph. Each entry
   stores the pixel coordinates (for a vertex in Q) and the distance of Q to
   its nearest neighbor in P. */
typedef struct {
   int i, j;                                               /* Own coordinates */
   int ni, nj;                              /* Nearest neighbor's coordinates */
   int d;                                 /* Distance to the nearest neighbor */
} heapNode;

/* Data type heap - a dynamic array of heap nodes */
typedef heapNode *heap;

/* Linked list representation of a queue */
typedef struct _qnode {
   int i, j;                                             /* Pixel coordinates */
   struct _qnode *next;                               /* Pointer to next node */
} qnode;

typedef struct {
   qnode *front, *back;   /* Pointers to the first and last nodes in the list */
} queue;

/* Read a ppm image from a file */
void readImage ( image *I, char infname[] )
{
   FILE *fp;
   char dummy = '\0';
   int i, j, r, g, b, maxval;

   fp = (FILE *)fopen(infname,"r");
   if (fp == NULL) {
      fprintf(stderr, "!!! ERROR: Unable to open input file\n");
      exit(2);
   }
   while (dummy != '\n') fscanf(fp,"%c",&dummy);
   fscanf(fp, "%d%d", &(I -> wd), &(I -> ht));
   fscanf(fp, "%d", &maxval);
   I -> R = (unsigned char **)malloc((I -> ht) * sizeof(unsigned char *));
   I -> G = (unsigned char **)malloc((I -> ht) * sizeof(unsigned char *));
   I -> B = (unsigned char **)malloc((I -> ht) * sizeof(unsigned char *));
   for (i = 0; i < I -> ht; ++i) {
      (I -> R)[i] = (unsigned char *)malloc((I -> wd) * sizeof(unsigned char));
      (I -> G)[i] = (unsigned char *)malloc((I -> wd) * sizeof(unsigned char));
      (I -> B)[i] = (unsigned char *)malloc((I -> wd) * sizeof(unsigned char));
      for (j = 0; j < I -> wd; ++j) {
         fscanf(fp, "%d%d%d", &r, &g, &b);
         (I -> R)[i][j] = (unsigned char)r;
         (I -> G)[i][j] = (unsigned char)g;
         (I -> B)[i][j] = (unsigned char)b;
      }
   }
   fclose(fp);
}

/* Write image to a ppm file */
void writeImage ( image I, char outfname[] )
{
   FILE *fp;
   int i, j;

   fp = fopen(outfname,"w");
   if (fp == NULL) {
      fprintf(stderr, "!!! ERROR: Unable to open output file\n");
      exit(3);
   }
   fprintf(fp, "P3\n%d %d\n255\n", I.wd, I.ht);
   for (i = 0; i < I.ht; ++i)
      for (j = 0; j < I.wd; ++j)
         fprintf(fp, "%d %d %d\n", I.R[i][j], I.G[i][j], I.B[i][j]);
   fclose(fp);
}

/* Distance between two pixels (r1,g1,b1) and (r2,g2,b2) is
   |r1-r2| + |g1-g2| + |b1-b2|. If (i2,j2) is not a valid pixel, the
   distance is returned as UNDEF. */
int pixdist ( image I, int i1, int j1, int i2, int j2 )
{
   int dr, dg, db;

   if ((i2 >= 0) && (i2 < I.ht) && (j2 >= 0) && (j2 < I.wd)) {
      dr = (int)I.R[i1][j1] - (int)I.R[i2][j2]; if (dr < 0) dr = -dr;
      dg = (int)I.G[i1][j1] - (int)I.G[i2][j2]; if (dg < 0) dg = -dg;
      db = (int)I.B[i1][j1] - (int)I.B[i2][j2]; if (db < 0) db = -db;
      return dr + dg + db;
   }
   return UNDEF;
}

/* Generate the pixel graph from the image I */
graph genGraph ( image I )
{
   graph G;
   int i, j, k;

   /* Initialize dimensions */
   G.wd = I.wd; G.ht = I.ht;

   /* Allocate memory */
   G.nbr = (cost **)malloc((G.ht) * sizeof(cost *));
   G.flag = (int **)malloc((G.ht) * sizeof(int *));
   for (i = 0; i < I.ht; ++i) {
      G.nbr[i] = (cost *)malloc((G.wd) * sizeof(cost));
      G.flag[i] = (int *)malloc((G.wd) * sizeof(int));
   }

   /* Store distances of neighboring pixels */
   for (i=0; i<G.ht; ++i) {
      for (j=0; j<G.wd; ++j) {
         G.nbr[i][j][0] = pixdist(I, i, j, i-1, j+1);
         G.nbr[i][j][1] = pixdist(I, i, j,   i, j+1);
         G.nbr[i][j][2] = pixdist(I, i, j, i+1, j+1);
         G.nbr[i][j][3] = pixdist(I, i, j, i+1,   j);
      }
   }

   /* Count the number of edges */
   G.e = 0;
   for (i=0; i<G.ht; ++i) for (j=0; j<G.wd; ++j) for (k=0; k<4; ++k)
      if (G.nbr[i][j][k] != UNDEF) ++G.e;

   return G;
}

/* Standard heapify with respect to the distance field (d) */
void heapify ( graph G, heap H, int n, int i )
{
   int l, r, s;
   int i1, j1, i2, j2;
   heapNode t;

   while (i < n) {
      l = 2*i + 1; r = 2*i + 2;
      if (l >= n) return;
      if (r >= n) s = l;
      else s = (H[l].d <= H[r].d) ? l : r;
      if (H[i].d <= H[s].d) return;

      /* Swapping of H[i] with H[s] is necessary here. First, identify the
         pixel coordinates. */
      i1 = H[i].i; j1 = H[i].j;
      i2 = H[s].i; j2 = H[s].j;

      /* Update the flags to reflect the swap */
      G.flag[i1][j1] = s; G.flag[i2][j2] = i;

      /* Make the swap */
      t = H[i]; H[i] = H[s]; H[s] = t;

      /* Prepare for the next iteration */
      i = s;
   }
}

/* Initialize the heap to contain all pixels except (0,0) */
heap makeHeap ( graph G )
{
   heap H;
   int i, j, cnt;

   H = (heapNode *)malloc((G.wd * G.ht - 1) * sizeof(heapNode));
   cnt = 0;

   /* Handle the first row. Pixel (0,0) is not placed in the heap. */
   for (j=1; j<G.wd; ++j) {
     H[cnt].i = 0; H[cnt].j = j;                           /* Own coordinates */
     H[cnt].ni = 0; H[cnt].nj = 0; /* Initially (0,0) is the nearest neighbor */
     H[cnt].d = INFINITY;                    /* Distance to (0,0) is infinity */
     G.flag[0][j] = cnt;  /* flag[i][j] stores the index of (i,j) in the heap */
     ++cnt;
   }

   /* Handle the remaining rows */
   for (i=1; i<G.ht; ++i) for (j=0; j<G.wd; ++j) {
     H[cnt].i = i; H[cnt].j = j;                           /* Own coordinates */
     H[cnt].ni = 0; H[cnt].nj = 0; /* Initially (0,0) is the nearest neighbor */
     H[cnt].d = INFINITY;                    /* Distance to (0,0) is infinity */
     G.flag[i][j] = cnt;  /* flag[i][j] stores the index of (i,j) in the heap */
     ++cnt;
   }

   /* (0,0) has three adjacent pixels. Replace INFINITY by the actual pixel
      distances for only these three neighbors. */
   H[0].d = G.nbr[0][0][1];
   H[G.wd-1].d = G.nbr[0][0][3];
   H[G.wd].d = G.nbr[0][0][2];

   /* (0,0) does not exist in the heap */
   G.flag[0][0] = UNDEF;

   for (i = (G.ht * G.wd - 1) / 2 - 1; i >= 0; --i)
      heapify(G, H, G.ht * G.wd - 1, i);

   return H;
}

void deleteMin ( graph G, heap H, int n )
{
   int i, j;

   /* Update flag for the pixel residing at H[n-1] */
   i = H[n-1].i; j = H[n-1].j; G.flag[i][j] = 0;

   H[0] = H[n-1];
   heapify(G,H,n-1,0);
}

/* Standard decrease-priority function */
void decPriority ( graph G, heap H, int n, int i )
{
   int i1, j1, i2, j2, p;
   heapNode t;

   while (1) {
      if (i == 0) return;
      p = (i - 1) / 2;
      if (H[p].d <= H[i].d) return;

      /* Swapping H[i] with H[p] is necessary here. Retrieve the coordinates
         of the pixels. */
      i1 = H[i].i; j1 = H[i].j;
      i2 = H[p].i; j2 = H[p].j;

      /* Update flag to remember the swap */
      G.flag[i1][j1] = p;
      G.flag[i2][j2] = i;

      /* Make the swap */
      t = H[i]; H[i] = H[p]; H[p] = t;

      /* Prepare for the next iteration */
      i = p;
   }
}

/* Get the correct index in the G.nbr[i][j] array */
int getnbrno ( int i, int j, int I, int J )
{
   if (J == j+1) return I - i + 1;
   else return 3;
}

/* Check whether (i1,j1) experiences a decrease in priority caused by the
   inclusion of (i,j) in P */
void checkPriority ( graph G, heap H, int i, int j, int i1, int j1, int n )
{
   int k, d;

   /* If (i1,j1) is not a valid pixel, do nothing. */
   if ((i1 < 0) || (i1 >= G.ht) || (j1 < 0) || (j1 >= G.wd)) return;

   /* If (i1,j1) is already shifted to P, do nothing. */
   if (G.flag[i1][j1] == UNDEF) return;

   /* Locate the distance between (i,j) and (i1,j1) */
   if ((j1 > j) || ((j1 == j) && (i1 == i + 1))) {
      k = getnbrno(i,j,i1,j1);
      d = G.nbr[i][j][k];
   } else {
      k = getnbrno(i1,j1,i,j);
      d = G.nbr[i1][j1][k];
   }

   /* Locate (i1,j1) in the heap */
   k = G.flag[i1][j1];

   /* If (i,j) is closer to (i1,j1) than its previous closest neighbor in P */
   if (d < H[k].d) {
      H[k].ni = i; H[k].nj = j;     /* (i,j) is the new closest neighbor in P */
      H[k].d = d;                              /* Record the shorter distance */
      decPriority(G,H,n,k);                          /* Restore heap ordering */
   }
}

/* Prim's algorithm to compute the MST of the graph G */
graph PrimMST ( graph G )
{
   graph T;
   int i, j, k, ni, nj, n;
   heap H;

   /* Initially, P = {(0,0)} and Q is the set of all remaining pixels */
   printf("    Initializing heap\n");
   H = makeHeap(G);

   printf("    Initializing MST\n");
   T.wd = G.wd; T.ht = G.ht;

   /* Allocate memory to the arrays of T */
   T.nbr = (cost **)malloc((T.ht) * sizeof(cost *));
   T.flag = (int **)malloc((T.ht) * sizeof(int *));
   for (i = 0; i < T.ht; ++i) {
      T.nbr[i] = (cost *)malloc((T.wd) * sizeof(cost));
      T.flag[i] = (int *)malloc((T.wd) * sizeof(int));
   }

   /* Initially T contains no edges */
   for (i=0; i<T.ht; ++i) for (j=0; j<T.wd; ++j)
      T.nbr[i][j][0] = T.nbr[i][j][1] = T.nbr[i][j][2] = T.nbr[i][j][3] = UNDEF;

   T.e = 0; n = G.ht * G.wd - 1;

   printf("    Entering Prim loop\n");
   while (n > 0) {
      i = H[0].i; j = H[0].j;      /* Pixel in Q with minimum distance from P */
      ni = H[0].ni; nj = H[0].nj; /* (ni,nj) is the closest nbr of (i,j) in P */

      /* Add to T the edge (i,j)--(ni,nj) */
      if ((nj > j) || ((nj == j) && (ni == i + 1))) {
         k = getnbrno(i, j, ni, nj);
         T.nbr[i][j][k] = G.nbr[i][j][k];
      } else {
         k = getnbrno(ni, nj, i, j);
         T.nbr[ni][nj][k] = G.nbr[ni][nj][k];
      }
      ++T.e;

      /* Delete (i,j) from the heap (that is, move it from Q to P) */
      deleteMin(G,H,n);
      G.flag[i][j] = UNDEF;
      --n;

      /* (i,j) has eight neighbors. Check one by one whether the inclusion
         of (i,j) in P changes the nearest neighbors of these eight pixels. */
      checkPriority(G, H, i, j, i-1, j+1, n);
      checkPriority(G, H, i, j,   i, j+1, n);
      checkPriority(G, H, i, j, i+1, j+1, n);
      checkPriority(G, H, i, j, i+1,   j, n);
      checkPriority(G, H, i, j, i-1, j-1, n);
      checkPriority(G, H, i, j,   i, j-1, n);
      checkPriority(G, H, i, j, i+1, j-1, n);
      checkPriority(G, H, i, j, i-1,   j, n);
   }

   return T;
}

/* Diagnostic function */
void printEdges ( graph T )
{
   int i, j, k, cnt = 0;

   for (i = 0; i < T.ht; ++i) for (j = 0; j < T.wd; ++j) for (k=0; k<4; ++k) {
      if (T.nbr[i][j][k] != UNDEF) {
         ++cnt;
      }
   }
   printf("e   = %d\n", cnt);
   printf("T.e = %d\n", T.e);
}

/* Delete nseg - 1 costliest tree edges */
void delEdges ( graph T, int nseg )
{
   int i, j, k, l, t, d;
   int ei[MAX_SEG_NO], ej[MAX_SEG_NO], ek[MAX_SEG_NO], ecost[MAX_SEG_NO];

   /* Initialize the array of nseg - 1 largest costs */
   for (l=0; l<nseg; ++l) ecost[l] = UNDEF;

   /* Look at all the edges of T */
   for (i = 0; i < T.ht; ++i) {
      for (j = 0; j < T.wd; ++j) {
         for (k = 0; k < 4; ++k) {
            d = T.nbr[i][j][k];
            if (d != UNDEF) {
               /* Locate the position of i,j,k in the largest-cost array */
               l = 0;
               while (d <= ecost[l]) ++l;

               /* If i,j,k belongs to one of the first nseg - 1 entries */
               if (l <= nseg-2) {
                  /* Move the following entries by one cell */
                  for (t = nseg-2; t > l; --t) {
                     ei[t] = ei[t-1]; ej[t] = ej[t-1]; ek[t] = ek[t-1];
                     ecost[t] = ecost[t-1];
                  }

                  /* Store i,j,k */
                  ei[l] = i; ej[l] = j; ek[l] = k; ecost[l] = d;
               }
            }
         }
      }
   }

   /* Delete from T the nseg - 1 costliest edges */
   for (l = 0; l < nseg - 1; ++l) {
      /* printf("%d %d %d %d\n", ei[l], ej[l], ek[l], ecost[l]); */
      T.nbr[ei[l]][ej[l]][ek[l]] = UNDEF;
      --T.e;
   }
   /* printf("T.e = %d\n", T.e); */
}

/* Linked list implementation of the queue ADT */
queue initQ ()
{
   queue Q;

   Q.front = Q.back = NULL;
   return Q;
}

int empty (queue Q)
{
   return (Q.front == NULL);
}

queue enqueue ( queue Q, int i, int j )
{
   qnode *q;

   q = (qnode *)malloc(sizeof(qnode));
   q -> i = i; q -> j = j; q -> next = NULL;
   if (Q.front == NULL) {
      Q.front = Q.back = q;
   } else {
      Q.back -> next = q; Q.back = q;
   }
   return Q;
}

queue dequeue ( queue Q )
{
   qnode *q;

   if (empty(Q)) {
      fprintf(stderr, "!!! Error: Dequeue from empty queue\n");
      exit(4);
   }

   q = Q.front;
   Q.front = q -> next;
   if (Q.front == NULL) Q.back = NULL;
   free(q);
   return Q;
}

/* Try to enqueue the neighbor (i1,j1) of (i,j) */
queue tryEnqueue ( graph T, queue Q, int i, int j, int i1, int j1, int color )
{
   int k;

   /* If (i1,j1) is not a valid pixel, do nothing */
   if ((i1 < 0) || (i1 >= T.ht) || (j1 < 0) || (j1 >= T.wd)) return Q;

   /* (i1,j1) has already received a color in this or an earlier invocation
      of the BFS. Do nothing */
   if (T.flag[i1][j1] != 0) return Q;

   /* If (i,j)--(i1,j1) is not an edge in T, do nothing */
   if ((j1 == j + 1) || ((j1 == j) && (i1 == i + 1))) {
      k = getnbrno(i,j,i1,j1);
      if (T.nbr[i][j][k] == UNDEF) return Q;
   } else {
      k = getnbrno(i1,j1,i,j);
      if (T.nbr[i1][j1][k] == UNDEF) return Q;
   }

   /* Mark (i1,j1) visited by assigning it a non-zero color */
   T.flag[i1][j1] = color;

   /* Enqueue (i1,j1) */
   return enqueue(Q,i1,j1);
}

/* T contains nseg components. By multiple invocations of BFS, we locate these
   components. The BFS starts from pixel (starti,startj). */
void BFS ( graph T, int starti, int startj, int color )
{
   queue Q;
   int i, j;

   Q = initQ();
   Q = enqueue(Q,starti,startj); T.flag[starti][startj] = color;
   while (!empty(Q)) {
      /* Get the pixel (i,j) from the front of the queue, and delete it from
         the queue. */
      i = Q.front -> i; j = Q.front -> j;
      Q = dequeue(Q);

      /* Check whether each of the eight neighbors of (i,j) can be enqueued */
      Q = tryEnqueue(T, Q, i, j, i-1, j+1, color);
      Q = tryEnqueue(T, Q, i, j,   i, j+1, color);
      Q = tryEnqueue(T, Q, i, j, i+1, j+1, color);
      Q = tryEnqueue(T, Q, i, j, i+1,   j, color);
      Q = tryEnqueue(T, Q, i, j, i-1, j-1, color);
      Q = tryEnqueue(T, Q, i, j,   i, j-1, color);
      Q = tryEnqueue(T, Q, i, j, i+1, j-1, color);
      Q = tryEnqueue(T, Q, i, j, i-1,   j, color);
   }
}

/* This function segments the image from the forest T */
void segment ( graph T, int nseg )
{
   int i, j, color;

   /* Initially all vertices in T are unvisited. So we set the flag for
      each vertex to 0. This also stands for Color 0. */
   for (i = 0; i < T.ht; ++i) for (j = 0; j < T.wd; ++j) T.flag[i][j] = 0;

   /* Make nseg - 1 calls of BFS */
   for (color = 1; color < nseg; ++color) {
      /* First identify an unvisited pixel (i,j) */
      for (i = 0; i < T.ht; ++i) {
         for (j = 0; j < T.wd; ++j) {
            if (T.flag[i][j] == 0) break;
         }
         if (j < T.wd) break;
      }
      /* printf("    i = %d, j = %d, color = %d\n", i, j, color); */

      /* Run BFS from (i,j). Set the visited flag to the current color
         for all vertices visited in this invocation of the BFS. */
      BFS(T,i,j,color);
   }
}

/* Change I to the segmented version of I */
void updateImage ( image I, graph T )
{
   /* Sixteen default colors */
   int COLOR[16][3] = {
      {255, 255, 160},
      {0, 255, 255},
      {255, 0, 255},
      {255, 255, 0},
      {255, 0, 0},
      {0, 255, 0},
      {0, 0, 255},
      {0, 127, 127},
      {127, 0, 127},
      {127, 127, 0},
      {127, 0, 0},
      {0, 127, 0},
      {0, 0, 127},
      {255, 160, 160},
      {160, 255, 160},
      {160, 160, 255}
   };
   int i, j, k;

   /* Assign colors to the components using the flag array of T */
   for (i=0; i<I.ht; ++i) {
      for (j=0; j<I.wd; ++j) {
         k = T.flag[i][j];
         I.R[i][j] = COLOR[k][0];
         I.G[i][j] = COLOR[k][1];
         I.B[i][j] = COLOR[k][2];
      }
   }
}

int main ( int argc, char *argv[] )
{
   char infname[1024], outfname[1024];
   image I;
   graph G, T;
   int nseg;

   if (argc == 1) {
      fprintf(stderr, "!!! ERROR: Basename of input file must be supplied as an argument\n");
      exit(1);
   }

   nseg = 0;
   if (argc >= 3) nseg = atoi(argv[2]);
   if ((nseg <= 0) || (nseg > MAX_SEG_NO)) nseg = MAX_SEG_NO;

   sprintf(infname, "%s.ppm", argv[1]);
   sprintf(outfname, "%s-seg.ppm", argv[1]);

   printf("+++ Reading input image\n");
   readImage(&I,infname);

   printf("+++ Generating graph from input image\n");
   G = genGraph(I);

   printf("+++ Generating MST of the graph\n");
   T = PrimMST(G);
   /* printEdges(T); */

   printf("+++ Deleting largest cost edges\n");
   delEdges(T,nseg);

   printf("+++ Run BFS on the edge-deleted tree\n");
   segment(T,nseg);

   printf("+++ Updating image based upon segmentation stored in the tree\n");
   updateImage(I,T);

   writeImage(I,outfname);

   exit(0);
}