/* Iterated Conditional Modes (ICM) binary image restoration
 *
 * Jerod Weinman
 * 9 June 2008
 */

#include "pbmio.h"
#include <stdio.h>


#define DEFAULT_ALPHA 2
#define DEFAULT_BETA 1
#define MAX_ITER 20


/* icmupdate - Calculate one ICM update on an img
 *
 * Preconditions:
 *  img is the original binary image
 *  restImg is the current restored image result
 *  workImg is a buffer for storing the update
 *  img, restImg, and workImg are all the same size
 *
 * Postconditions:
 *  update is stored in workImg
 *  Return value indicates when update yields no change from restImg
 */


int icmupdate(const pbm_t *img, const pbm_t *restImg, pbm_t *workImg, 
              double alpha, double beta)
{
     int converged = 1; /* Flag indicating whether algorithm has converged */
     double cost[2]; /* cost array for both binary states at a pixel */
     int i,j;
     
     for (i=0 ; i < img->rows ; i++)
          for (j=0 ; j < img->cols ; j++)
          {
               cost[0] = 0;  /* Initialize costs to zero */
               cost[1] = 0;
                    
               /* Local cost: for flipping pixel (i,j). This
                * assigns the local cost alpha to the opposite
                * state of that at (i,j) */
                    
               cost[1-img->bits[i][j]] = alpha;
                    
               /* Neighborhood cost: Adds beta to the cost for
                * the opposite value of each neighboring
                * state. Note that the neighboring state values
                * used are from the most recent iteration of the
                * restored image. */

               /* Top-Left */
               if (j > 0 && i>0)
                    cost[1-restImg->bits[i-1][j-1]] += beta;
                    
               /* Left */
               if (j > 0)
                    cost[1-restImg->bits[i][j-1]] += beta;
                    
               /* Bottom-Left */
               if (i < restImg->rows-1 && j > 0)
                    cost[1-restImg->bits[i+1][j-1]] += beta;
                    
               /* Top-Right */
               if (i>0 && j < restImg->cols-1)
                    cost[1-restImg->bits[i-1][j+1]] += beta;
                    
               /* Right */
               if (j < restImg->cols-1)
                    cost[1-restImg->bits[i][j+1]] += beta;
                    
               /* Bottom Right */
               if (i < restImg->rows-1 && j < restImg->cols-1)
                    cost[1-restImg->bits[i+1][j+1]] += beta;
                    
               /* Top */
               if (i>0)
                    cost[1-restImg->bits[i-1][j]] += beta;
                    
               /* Bottom */
               if (i < restImg->rows-1)
                    cost[1-restImg->bits[i+1][j]] += beta;
                    
               /* Assign whichever state has lower cost to the
                * intermediate "working" restored image */
               workImg->bits[i][j] = (cost[0] > cost[1]);
                    
               /* If we still think we are converging, check
                * whether the new value (in workImg) differs from
                * our previous restored value. If they differ, we
                * have not converged */
               if (converged && workImg->bits[i][j]!=restImg->bits[i][j])
                    converged = 0;
          }
     return converged;
}

/* runicm - Run the ICM algorthm with parameters alpha and beta on an image
 *
 * Preconditions:
 *  img is the original binary image
 *  restImg is an already allocated buffer for the restored image result
 *  img and restImg are the same size
 *
 * Postconditions:
 *  ICM is run on img until convergence or an iteration limit is reached
 *  Result is stored in restImg buffer
 */
void runicm(const pbm_t *img, pbm_t *restImg, double alpha, double beta)
{
     pbm_t *workImg; /* An image buffer for storing intermediate results */

     int converged;
     int iter;

     workImg = (pbm_t*) malloc (sizeof(pbm_t));

     if (workImg==NULL)
     {
       fprintf(stderr,"Unable to allocate work image");
       exit(1);
     }
     
     pbmalloc(workImg, img->rows, img->cols); /* Allocate work image */

     pbmcopy(restImg, img); /* Copy original image into the restored image */

     
     for (iter=0 ; iter<MAX_ITER ; iter++)
     {
          converged = icmupdate (img, restImg, workImg, alpha, beta);

          if (converged)
               /* Nothing changed, so we are done and can exit the loop */
               break;
          else          
               /* Now that all pixels have been updated, we can copy the
                * working buffer over into the restored image buffer. */
               pbmcopy(restImg,workImg);
     }
     
     pbmfree(workImg); /* Free our temporary work image buffer */
					free(workImg);
}

int main(int argc, char* argv[])
{

     char *origFile, *cleanFile; /* Input and output file names */
     double alpha,beta;          /* ICM algorithm parameters */

     pbm_t origImg, cleanImg;  /* Input and output images */


     /* Process command line arguments */
     if (argc<3)
     {
          fprintf(stderr,"Usage: %s input output [alpha] [beta]", argv[0]);
          exit (1);
     }
     
     origFile = argv[1];
     cleanFile = argv[2];

     /* Process optional alpha and beta arguments */
     if (argc>3)
          alpha = strtod(argv[3], NULL);
     else
          alpha = DEFAULT_ALPHA;

     if (argc>4)
          beta = strtod(argv[4], NULL);
     else
          beta = DEFAULT_BETA;

     /* Read input image */
     if (pbmread(origFile, &origImg)<0)
          exit (1);

     /* Allocate space for the restored image */
     if (pbmalloc(&cleanImg, origImg.rows, origImg.cols)<0)
          exit (1);

     /* Run the restoring ICM algorithm */
     runicm(&origImg, &cleanImg, alpha, beta);

     /* Write the restored image */
     if (pbmwrite(cleanFile, &cleanImg)<0)
          exit (1);

     /* Free our allocated image */
     pbmfree(&cleanImg);

     /* Exit cleanly */
     exit (0); 
}