/**
 *  gcc -Wall -g -o antidust antidust.c  -lppm -lpgm -lpbm
 *
 */

#include <stdio.h>
#include <errno.h>
#include <string.h>

#include <ppm.h>

//#define HISTOGRAM 1
#define IR_LEVEL 200
#define NSTEPS 30
#define IR_BORDER 1

struct image_info {
  pixval maxval;
  int n_cols;
  int n_rows;
};

static void do_image_infill(pixel **rgb_image, pixel **ir_image, pixel **out_rgb, pixel **out_ir, struct image_info info) {
  int i, j;
  int step;
  int count = 0;
  int b_count = 0;
  int k, l;
  
  pixel pix_rgb;
  pixel pix_ir;
  pixval p_r;
  pixval p_g;
  pixval p_b;
  pixval p_ir;

#if defined(HISTOGRAM) 
  int histo_r[256];
  int histo_g[256];
  int histo_b[256];
  int histo_ir[256];

  for (i = 0; i < 255; i ++) {
    histo_r[i] = 0;
    histo_g[i] = 0;
    histo_b[i] = 0;
    histo_ir[i] = 0;
  }
#endif

  // black out where ir level is below some given value (the black spots);
  for (i = 0; i < info.n_rows; i++) {
    for (j = 0; j < info.n_cols; j++) {
       PPM_ASSIGN(out_ir[i][j], 255, 255, 255);
    }
  }
  for (i = 0; i < info.n_rows; i++) {
    for (j = 0; j < info.n_cols; j++) {
      pix_rgb = rgb_image[i][j];
      pix_ir  = ir_image[i][j];

      p_r = PPM_GETR(pix_rgb);
      p_g = PPM_GETG(pix_rgb);
      p_b = PPM_GETB(pix_rgb);
      p_ir = PPM_GETR(pix_ir);

#if defined(HISTOGRAM)       
      histo_r[p_r]++;
      histo_g[p_g]++;
      histo_b[p_b]++;
      histo_ir[p_ir]++;
#endif

      if (p_ir >= IR_LEVEL) {
	p_ir = 255;
      } else {
	p_ir = 0;
	p_r = p_g = p_b = 255;
      }
      
      if (PPM_GETR(out_ir[i][j]) == 255) {
	PPM_ASSIGN(out_ir[i][j], p_ir, p_ir, p_ir);
	PPM_ASSIGN(out_rgb[i][j], p_r, p_g, p_b);
      }

      // black out 1 pixel border
      if (p_ir == 0) {
	for(k = - IR_BORDER; k <= IR_BORDER; k++) {
	  for (l = - IR_BORDER; l <= IR_BORDER; l++) {
	    if ((i+k)>0 && (i+k)<info.n_rows && (j+l)>0 && (j+l)<info.n_cols) {
	      PPM_ASSIGN(out_ir[i+k][j+l], p_ir, p_ir, p_ir);
	      PPM_ASSIGN(out_rgb[i+k][j+l],p_r, p_g, p_b);
	    }
	  }
	}
      }
    }
  }

#if defined(HISTOGRAM)
  for (i = 0; i < 255; i ++) {
    printf("%3d  ir: %7d  r: %7d  g: %7d  b: %7d\n",
	   i, histo_ir[i], histo_r[i], histo_g[i], histo_b[i]);
  }
#endif

  // image inpainting
  for (step = 0; step < NSTEPS; step++) {
    for (i = 0; i < info.n_rows; i++) {
      for (j = 0; j < info.n_cols; j++) {
	ir_image[i][j] = out_ir[i][j];
      }
    }
    for (i = 0; i < info.n_rows; i++) {
      for (j = 0; j < info.n_cols; j++) {
	pix_rgb = rgb_image[i][j];
	pix_ir  = ir_image[i][j];
	
	p_ir = PPM_GETR(pix_ir);
	       
	if (p_ir < IR_LEVEL) {
	  count = 0;
	  b_count = 0;
	  p_r = 0; p_g = 0; p_b = 0;
	  for(k=-3; k<=3; k++) {
	    for (l=-3; l<=3; l++) {
	      if ((i+k)>0 && (i+k)<info.n_rows &&
		  (j+l)>0 && (j+l)<info.n_cols) {
		if (PPM_GETR(ir_image[i+k][j+l]) != 0) {
		  // ((k*k)+(l*l)) != 0, because at point (i,j) is ir == 0
		  int w = 6 / (abs(k) + abs(l));
		  //int w = 36 / ((k*k) + (l*l));
		  //int w = 1;
		  p_r += PPM_GETR(out_rgb[i+k][j+l]) * w;
		  p_g += PPM_GETG(out_rgb[i+k][j+l]) * w;
		  p_b += PPM_GETB(out_rgb[i+k][j+l]) * w;
		  count += w;
		  if ( k >= -1 && k <= 1 && l >= -1 && l <= 1) {
		    b_count++;
		  }
		}
	      }
	    }
	  }
	  
	  if (count > 0 && b_count >= 3) {
	    p_r = p_r / count;
	    p_g = p_g / count;
	    p_b = p_b / count;
	    PPM_ASSIGN(out_ir[i][j], 255, 255, 255);
	    PPM_ASSIGN(out_rgb[i][j], p_r, p_g, p_b);
	  } 
	}
      }
    }
  }
}

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

  FILE *rgb_file;
  FILE *ir_file;
  FILE *out_rgb;
  FILE *out_ir;

  pixel **rgb_image = NULL;
  pixel **ir_image  = NULL;
  pixel **out_rgb_image = NULL;
  pixel **out_ir_image  = NULL;

  struct image_info rgb_info;
  struct image_info ir_info;
  struct image_info out_info;
  
  if (argc != 5) {
    fprintf(stderr, " Usage image_infill <rgb.pnm> <ir.pnm> <out_rgb.pnm> <out_ir.pnm>\n");
    exit(1);
  }

  if ((rgb_file = fopen(argv[1], "r")) == NULL) {
    fprintf(stderr, "ERROR: open of file %s failed (%s),\n", argv[1], strerror(errno));
    exit(1);
  }

  if ((ir_file = fopen(argv[2], "r")) == NULL) {
    fprintf(stderr, "ERROR: open of file %s failed (%s),\n", argv[2], strerror(errno));
    exit(1);
  }

  if ((out_rgb = fopen(argv[3], "w")) == NULL) {
    fprintf(stderr, "ERROR: open of file %s failed (%s),\n", argv[3], strerror(errno));
    exit(1);
  }
  
  if ((out_ir = fopen(argv[4], "w")) == NULL) {
    fprintf(stderr, "ERROR: open of file %s failed (%s),\n", argv[4], strerror(errno));
    exit(1);
  }

  rgb_image = ppm_readppm(rgb_file, &rgb_info.n_cols, &rgb_info.n_rows, &rgb_info.maxval);
  ir_image = ppm_readppm(ir_file, &ir_info.n_cols, &ir_info.n_rows, &ir_info.maxval);
  
  printf("rgb: n_cols %d n_rows %d maxval %d\n", rgb_info.n_cols, rgb_info.n_rows, rgb_info.maxval);
  printf("ir:  n_cols %d n_rows %d maxval %d\n", ir_info.n_cols, ir_info.n_rows, ir_info.maxval);

  out_info = rgb_info;
  out_rgb_image = ppm_allocarray(out_info.n_cols, out_info.n_rows);
  out_ir_image  = ppm_allocarray(out_info.n_cols, out_info.n_rows);
  
  do_image_infill(rgb_image, ir_image, out_rgb_image, out_ir_image, rgb_info);

  ppm_writeppm(out_rgb, out_rgb_image, out_info.n_cols, out_info.n_rows, out_info.maxval, 0);
  ppm_writeppm(out_ir, out_ir_image, out_info.n_cols, out_info.n_rows, out_info.maxval, 0);
  
  ppm_freearray(rgb_image, out_info.n_rows);
  ppm_freearray(ir_image, out_info.n_rows);
  ppm_freearray(out_rgb_image, out_info.n_rows);
  ppm_freearray(out_ir_image, out_info.n_rows);

  fclose(rgb_file);
  fclose(ir_file);
  fclose(out_rgb);
  fclose (out_ir);

  exit(0);
}