/******************************************************************************

Image I/O and Image processing code 

Copyright 2000 by the University of Waterloo
All rights reserved.
   
Developed by: Richard Mann, Department of Computer Science

******************************************************************************/
#include <stdio.h>
#include <math.h>
#include <string.h>
#include <stdlib.h>

#define min(A,B) ((A)<(B)?(A):(B))
#define max(A,B) ((A)>(B)?(A):(B))
#define abs(A) ((A)>(0)?(A):(-A))
#define sign(A,B) ((B)>=(0)?(A):(-(A)))
#define sqr(A) ((A)*(A))
#define round(A) (floor(0.5+(A)))
#ifndef PI
  #define PI M_PI
#endif

/* image accessor: image(x,y) */
#define I(imageP, nx, ny, x, y) imageP[x+y*nx]

/* maximum # of image modes (excluding outlier) */
#define MAXMODES 4

void error(char *s1, char *s2)
{
  fprintf(stderr, "ERROR: %s %s\n", s1, s2);
  exit(1);
}

/*
 * 2d interpolation of image pixel data.
 * Assumes point (x,y) within range of pixels array (ie., 0..nx-1,0..ny-1).
 */
float interp2(float *imageP, int nx, int ny, float x, float y)
{
  int xl, xh, yl, yh;
  float dx, dy;
  float p1, p2, p3, p4;

  xl=(int)floor(x); xh=(int)ceil(x);
  yl=(int)floor(y); yh=(int)ceil(y);
  if (xl<0 || yl<0 || xh>=nx || yh>=ny) {
    fprintf(stderr,
            "INTERP2: x:%f, y:%f, nx:%d, ny:%d, xl:%d, xh:%d, yl:%d, yh:%d\n",
            x,y,nx,ny,xl,xh,yl,yh);
    exit(1);
  }
  dx=x-(float)xl; dy=y-(float)yl;
  p1=I(imageP,nx,ny,xl,yl);
  p2=I(imageP,nx,ny,xl,yh);
  p3=I(imageP,nx,ny,xh,yl);
  p4=I(imageP,nx,ny,xh,yh);
  return(dx*dy*p4+(1-dx)*dy*p3+dx*(1-dy)*p2+(1-dx)*(1-dy)*p1);
}

/*
 * 1d filter.  Does work in place
 */
void filter_image_x(float *imageP, int nx, int ny, float *filtP, int nfilt)
{
  float *rowP;
  float s;
  int c, x, y, fi, xi;
  
  c = (int)floor(nfilt/2.0); /* center of filter */
  if ((rowP=(float*)malloc((size_t)nx*sizeof(float)))==NULL)
    error("Cannot allocate memory", "image_filter_x");
  for (y=0; y<ny; y++) {
    for (x=0; x<nx; x++) {
      rowP[x]=I(imageP,nx,ny,x,y);
    }
    for (x=0; x<nx; x++) {
      s=0;
      for (fi=0; fi<nfilt; fi++) {
        xi=x-c+fi;
	if (xi>=0 && xi<nx) s+=rowP[xi]*filtP[fi];
      }
      I(imageP,nx,ny,x,y)=s; /* update original image pixel */
    }
  }
  free(rowP);
}

void filter_image_y(float *imageP, int nx, int ny, float *filtP, int nfilt)
{
  float *colP;
  float s;
  int c, x, y, fi, yi;

  c = (int)floor(nfilt/2); /* center of filter */
  if ((colP=(float*)malloc((size_t)ny*sizeof(float)))==NULL)
    error("Cannot allocate memory", "image_filter_y");
  for (x=0; x<nx; x++) {
    for (y=0; y<ny; y++) {
      colP[y]=I(imageP,nx,ny,x,y);
    }
    for (y=0; y<ny; y++) {
      s=0;
      for (fi=0; fi<nfilt; fi++) {
        yi=y-c+fi;
	if (yi>=0 && yi<ny) s+=colP[yi]*filtP[fi];
      }
      I(imageP,nx,ny,x,y)=s; /* update original image pixel */
    }
  }
  free(colP);
}

/*
 * Derviatives of image in x and y-directions, using five-point difference
 * operator.
 */
void diff_image_x(float *imageP, int nx, int ny)
{
  static float filt[5]={1.0/12.0, -8.0/12.0, 0.0, 8.0/12.0, -1.0/12.0};
  
  filter_image_x(imageP, nx, ny, filt, 5);
}

void diff_image_y(float *imageP, int nx, int ny)
{
  static float filt[5]={1.0/12.0, -8.0/12.0, 0.0, 8.0/12.0, -1.0/12.0};
  
  filter_image_y(imageP, nx, ny, filt, 5);
}

void smooth_image(float *imageP, int nx, int ny, float sigma)
{
  float *filtP;
  int i, n, w;
  float s, t;

  n=(int)ceil(4*sigma); /* number of pixels on either side of center */
  w=2*n+1; /* width of filter */

  /* make gaussian filter, then normalize it so it sums to one */
  if ((filtP=(float*)malloc((size_t)w*sizeof(float))) == NULL)
    error("Cannot allocate filter", "smooth_float_image");
  t=0.5/sqr(sigma);
  for (i=0; i<w; i++) {
    filtP[i]=exp(-t*sqr(i-n));
  }
  s=0;
  for (i=0; i<w; i++) s+=filtP[i];
  for (i=0; i<w; i++) filtP[i]/=s;
  
  /* filter image in X and Y direction */
  filter_image_x(imageP, nx, ny, filtP, w);
  filter_image_y(imageP, nx, ny, filtP, w);

  free(filtP);
}

float *make_image(int nx, int ny)
{
  float *imageP;

  if ((imageP=(float*)malloc((size_t)nx*ny*sizeof(float)))==NULL)
    error("Cannot allocate image", "make_image");
  return(imageP);
}

void *free_image(float *imageP)
{
  free(imageP);
}

float *copy_image(float *imageP, int nx, int ny)
{
  float *image2P;
  int x, y;
 
  image2P=make_image(nx,ny);
  for (y=0; y<ny; y++) {
    for (x=0; x<nx; x++) {
      I(image2P,nx,ny,x,y)=I(imageP,nx,ny,x,y);
    }
  }
  return(image2P);
}

float *read_pgm_image(char *filenameP, int *nxP, int *nyP)
{
  char line[255];
  int i, maxval, x, y, nx, ny, nread;
  unsigned char *charP;
  float *imageP;
  FILE *fp;

  if ((fp=fopen(filenameP,"r")) == NULL)
    error("Unable to open PGM image", filenameP);

  if (fgets(line, 255, fp) == NULL)
    error("Unexpected end of file", filenameP);
  if (strncmp(line, "P5", 2) != 0)
    error("File is not raw pgm image", filenameP);
  while (1) {
    if (fgets(line, 255, fp) == NULL)
      error("Unexpected EOF", filenameP);
    if (line[0] != '#') break;
  }
  if (sscanf(line, "%d %d", nxP, nyP) == EOF)
    error("Can't get image size", filenameP);
  nx=*nxP; ny=*nyP;
  while (1) {
    if (fgets(line, 255, fp) == NULL)
      error("Unexpected EOF", filenameP);
    if (line[0] != '#') break;
  }
  if (sscanf(line, "%d %d", &maxval) == EOF)
    error("Can't get image range", filenameP);
  if (maxval != 255)
    error("PGM maxval not 255", filenameP);

  if ((charP=(unsigned char*)malloc((size_t)nx*ny*sizeof(unsigned char)))==NULL)
    error("Unable to allocate PGM image", "read_pgm_image");
  if (nx*ny != fread((void*)charP, sizeof(unsigned char), nx*ny, fp)) {
    error("unexpected end of file on", filenameP);
  }
    
  imageP=make_image(nx,ny);

  for (y=0; y<ny; y++) {
    for (x=0; x<nx; x++) {
      I(imageP,nx,ny,x,y)=(float)I(charP,nx,ny,x,y);
    }
  }
  free(charP);
  return(imageP);
}
  
int write_pgm_image(char *filenameP, float *imageP, int nx, int ny, char *blurb)
{
  FILE *fp;
  unsigned char *charP;
  int x, y;

  if ((charP=(unsigned char*)malloc((size_t)nx*ny*sizeof(char)))==NULL)
    error("Unable to allocate PGM image", "read_pgm_image");
  for (y=0; y<ny; y++) {
    for (x=0; x<nx; x++) {
      I(charP,nx,ny,x,y)=(unsigned char)max(0,min(255,round(I(imageP,nx,ny,x,y))));
    }
  }
  
  if ((fp=fopen(filenameP,"w")) == NULL)
    error("Unable to output output image", filenameP);

  fprintf(fp, "P5\n# %s\n%d %d\n%d\n", blurb, nx, ny, (int)255);
  fwrite((void*)charP, sizeof(unsigned char), nx*ny, fp);
  fflush(fp);
  free(charP);
  return(0);
}

float image_max(float *imageP, int nx, int ny)
{
  int x,y;
  float s;

  s=I(imageP,nx,ny,0,0);
  for (y=0;y<ny;y++) {
    for (x=0;x<nx;x++) {
      if (I(imageP,nx,ny,x,y)>s) s=I(imageP,nx,ny,x,y);
    }
  }
  return(s);
}

float image_min(float *imageP, int nx, int ny)
{
  int x,y;
  float s;

  s=I(imageP,nx,ny,0,0);
  for (y=0;y<ny;y++) {
    for (x=0;x<nx;x++) {
      if (I(imageP,nx,ny,x,y)<s) s=I(imageP,nx,ny,x,y);
    }
  }
  return(s);
}

void scale_image(float *imageP, int nx, int ny, float scale)
{
  int x,y;

  for (y=0; y<ny; y++) {
    for (x=0; x<nx; x++) {
      I(imageP,nx,ny,x,y)*=scale;
    }
  }
}

float *add_images(float *image1P, float *image2P, int nx, int ny)
{
  int x,y;
  float *image3P;

  image3P=make_image(nx,ny);
  for (y=0; y<ny; y++) {
    for (x=0; x<nx; x++) {
      I(image3P,nx,ny,x,y)=I(image1P,nx,ny,x,y)+I(image2P,nx,ny,x,y);
    }
  }
  return(image3P);
}

float *subtract_images(float *image1P, float *image2P, int nx, int ny)
{
  int x,y;
  float *image3P;

  image3P=make_image(nx,ny);
  for (y=0; y<ny; y++) {
    for (x=0; x<nx; x++) {
      I(image3P,nx,ny,x,y)=I(image1P,nx,ny,x,y)-I(image2P,nx,ny,x,y);
    }
  }
  return(image3P);
}

float *multiply_images(float *image1P, float *image2P, int nx, int ny)
{
  int x,y;
  float *image3P;

  image3P=make_image(nx,ny);
  for (y=0; y<ny; y++) {
    for (x=0; x<nx; x++) {
      I(image3P,nx,ny,x,y)=I(image1P,nx,ny,x,y)*I(image2P,nx,ny,x,y);
    }
  }
  return(image3P);
}

void test_richard()
     /* Test image i/o and filtering */
{
  float *im1P, *im2P;
  int nx, ny;

  printf("starting\n");
  im1P=read_pgm_image("coke-pick02.031", &nx, &ny);
  printf("after read\n");
  write_pgm_image("blah.pgm", im1P, nx, ny, "original image");
  printf("after write\n");
  im2P=copy_image(im1P, nx, ny);
  printf("after copy\n");
  write_pgm_image("blah.copy", im2P, nx, ny, "copy");
  diff_image_x(im1P, nx, ny);
  printf("after diffx\n");
  write_pgm_image("blah.diffx", im1P, nx, ny, "diff_x");
  printf("after write\n");
  smooth_image(im2P, nx, ny, 8.0);
  printf("after smooth\n");
  write_pgm_image("blah.smoothed", im2P, nx, ny, "smooth");
  printf("after write\n");
  diff_image_x(im2P, nx, ny);
  printf("after diffx\n");
  write_pgm_image("blah.smoothed.diffx", im2P, nx, ny, "smoothed diff_x");
  printf("after write\n");
}