#include <png.h>

void savePNG(const char *filename, int width, int height, float *red, float *green, float *blue, float *alpha)
{
  png_byte colorType=PNG_COLOR_TYPE_RGBA;
  png_byte bitDepth=8;
  png_structp pngPtr;
  png_infop pngInfo;
  png_bytep *rowPointers;
  
  rowPointers=new png_bytep[height];
  for(int j=0;j<height;j++) {
  	rowPointers[j]=new png_byte[width*4];
  	for(int i=0;i<width;i++) {
		float clampedRed=red[j*width+i];
		if(clampedRed>1) {
		  clampedRed=1.0;
		}
		float clampedGreen=green[j*width+i];
		if(clampedGreen>1) {
		  clampedGreen=1.0;
		}
		float clampedBlue=blue[j*width+i];
		if(clampedBlue>1) {
		  clampedBlue=1.0;
		}
		float clampedAlpha=alpha[j*width+i];
		if(clampedAlpha>1) {
		  clampedAlpha=1.0;
		}
  		rowPointers[j][(i*4)]=(unsigned char)(clampedRed*0xFF);
  		rowPointers[j][(i*4)+1]=(unsigned char)(clampedGreen*0xFF);
  		rowPointers[j][(i*4)+2]=(unsigned char)(clampedBlue*0xFF);
  		rowPointers[j][(i*4)+3]=(unsigned char)(clampedAlpha*0xFF);
  	}
  }
  
  // create file
  FILE *fp;
  if (!(fp = fopen(filename, "wb")))
    return;

  pngPtr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
        
  if(!pngPtr) {
    return;
  }

  pngInfo = png_create_info_struct(pngPtr);
  if(!pngInfo) {
    return;
  }

  if(setjmp(png_jmpbuf(pngPtr))) {
    return;
  }
 
  png_init_io(pngPtr, fp);

  // write header
  if(setjmp(png_jmpbuf(pngPtr))) {
    return;
  }

  png_set_IHDR(pngPtr, pngInfo, width, height,
               bitDepth, colorType, PNG_INTERLACE_NONE,
               PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
               
  png_set_gAMA(pngPtr,pngInfo,1.0);

  png_write_info(pngPtr, pngInfo);


  // write image
  if(setjmp(png_jmpbuf(pngPtr))) {
    return;
  }

  png_write_image(pngPtr, (png_byte**)rowPointers);

  if(setjmp(png_jmpbuf(pngPtr))) {
    return;
  }

  png_write_end(pngPtr, NULL);

  for(int i=0; i<height; i++) {
    delete rowPointers[i];
  }
  delete rowPointers;

  fclose(fp);
}

bool loadPNG(const char *filename, unsigned int &width, unsigned int &height, float **red, float **green, float **blue, float **alpha)
{
  png_byte header[8]; // 8 is the maximum size that can be checked
  int colorType;
  int bitDepth;
  png_structp pngPtr;
  png_infop pngInfo;
  png_bytep *rowPointers;
  double gamma;  

  // open file and test for it being a png
  FILE *fp;
  if(!(fp=fopen(filename,"rb"))) {
  	return false;
  }
  fread(header, 1, 8, fp);
  if(png_sig_cmp(header, 0, 8)) {
  	return false;
  }

  // initialise data structures
  pngPtr = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
        
  if(!pngPtr) {
  	return false;
  }

  pngInfo = png_create_info_struct(pngPtr);
  if(!pngInfo) {
  	return false;
  }
  if(setjmp(png_jmpbuf(pngPtr))) {
    return false;
  }

  png_init_io(pngPtr, fp);
  png_set_sig_bytes(pngPtr, 8);

  png_read_info(pngPtr, pngInfo);

  png_get_IHDR(pngPtr, pngInfo, &width, &height, &bitDepth, &colorType, NULL, NULL, NULL);

  png_read_update_info(pngPtr, pngInfo);

  // read file
  if(setjmp(png_jmpbuf(pngPtr))) {
  	return false;
  }

  rowPointers = new png_bytep[height];
  for(int i=0; i<height; i++) {
    rowPointers[i] = new png_byte[png_get_rowbytes(pngPtr, pngInfo)];
  }

  // strip 16 bit color depth image to 8 bit automatically
  if(bitDepth == 16) {
    png_set_strip_16(pngPtr);
  }
  // apply gama value if present in image (set desired gamma exponent to 1 to get linear results):
  if(png_get_gAMA(pngPtr, pngInfo, &gamma)) {
    png_set_gamma(pngPtr, 1, gamma);
  }

  png_read_image(pngPtr, rowPointers);
  fclose(fp);
  
  int maxValue=0;
  for(int i=0;i<bitDepth;i++) {
  	maxValue=maxValue<<1;
  	maxValue=maxValue+1;
  }
  
  *red=new float[width*height];
  *green=new float[width*height];
  *blue=new float[width*height];
  *alpha=new float[width*height];

  switch(colorType) {
  case PNG_COLOR_TYPE_RGB_ALPHA:
    for(int j=0;j<height;j++) {
      for(int i=0;i<width;i++) {
        (*red)[j*width+i]=rowPointers[j][(i*4)]/(float)maxValue;
  	(*green)[j*width+i]=rowPointers[j][(i*4)+1]/(float)maxValue;
  	(*blue)[j*width+i]=rowPointers[j][(i*4)+2]/(float)maxValue;
  	(*alpha)[j*width+i]=rowPointers[j][(i*4)+3]/(float)maxValue;
      }
    }
    break;
  case PNG_COLOR_TYPE_RGB:
    for(int j=0;j<height;j++) {
      for(int i=0;i<width;i++) {
        (*red)[j*width+i]=rowPointers[j][(i*3)]/(float)maxValue;
  	(*green)[j*width+i]=rowPointers[j][(i*3)+1]/(float)maxValue;
  	(*blue)[j*width+i]=rowPointers[j][(i*3)+2]/(float)maxValue;
  	(*alpha)[j*width+i]=1.0;
      }
    }
    break;
  case PNG_COLOR_TYPE_GRAY_ALPHA:
    for(int j=0;j<height;j++) {
      for(int i=0;i<width;i++) {
        (*red)[j*width+i]=rowPointers[j][(i*2)]/(float)maxValue;
  	(*green)[j*width+i]=rowPointers[j][(i*2)]/(float)maxValue;
  	(*blue)[j*width+i]=rowPointers[j][(i*2)]/(float)maxValue;
  	(*alpha)[j*width+i]=rowPointers[j][(i*2)+1]/(float)maxValue;
      }
    }
    break;
  case PNG_COLOR_TYPE_GRAY:
    for(int j=0;j<height;j++) {
      for(int i=0;i<width;i++) {
        (*red)[j*width+i]=rowPointers[j][i]/(float)maxValue;
  	(*green)[j*width+i]=rowPointers[j][i]/(float)maxValue;
  	(*blue)[j*width+i]=rowPointers[j][i]/(float)maxValue;
  	(*alpha)[j*width+i]=1.0;
      }
    }
    break;
  case PNG_COLOR_TYPE_PALETTE:
    // TODO: need palette information to decode
    break;
  }
  
  for(int i=0;i<height;i++) {
  	delete rowPointers[i];
  }
  delete rowPointers;
  
  return true;
}