#include <stdio.h>
#include <math.h>
#include <vector>
#include <png.h>
#include "Tri.h"

using namespace std; 

#ifndef _WritePng_
#define _WritePng_

//---------------------------------------------------------------------------
struct Level
{
	float Value;
	char R;
	char G;
	char B;
	char A;	// channel alpha (transparence : 0=transparent, 255=opaque)
};

bool   writePNG(String filename, int width, int height, float **buffer, vector<Level> echelle_couleurs);
bool   writePNG(String filename, int width, int height, float **buffer);

//---------------------------------------------------------------------------
bool writePNG(String filename, int width, int height, float **buffer)
{
	int x, y, l, p;
	vector<Level> echelle_couleurs;
	Level level;
    vector<int> tri, unique;
	
	FILE *fp = NULL;
	png_structp png_ptr = NULL;
	png_infop info_ptr = NULL;
	png_bytep row = NULL;
	
	if(FileExists(filename))
	{
		DeleteFile(filename);
		if(FileExists(filename))
			return false;
	}
	
	// Open file for writing (binary mode)
	fp = fopen(filename.c_str(), "wb");
	if (fp == NULL) {
		fprintf(stderr, "Could not open file %s for writing\n", filename.c_str());
		return false;
	}

	// Initialize write structure
	png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	if (png_ptr == NULL) {
		fprintf(stderr, "Could not allocate write struct\n");
		if (fp != NULL) fclose(fp);
		if (info_ptr != NULL) png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
		if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
		if (row != NULL) free(row);
		return false;
	}

	// Initialize info structure
	info_ptr = png_create_info_struct(png_ptr);
	if (info_ptr == NULL) {
		fprintf(stderr, "Could not allocate info struct\n");
		if (fp != NULL) fclose(fp);
		if (info_ptr != NULL) png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
		if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
		if (row != NULL) free(row);
		return false;
	}

	// Setup Exception handling
	if (setjmp(png_jmpbuf(png_ptr))) {
		fprintf(stderr, "Error during png creation\n");
		if (fp != NULL) fclose(fp);
		if (info_ptr != NULL) png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
		if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
		if (row != NULL) free(row);
		return false;
	}

	png_init_io(png_ptr, fp);

	// Write header (8 bit colour depth)
	png_set_IHDR(png_ptr, info_ptr, width, height,
			8, PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE,
			PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

	png_write_info(png_ptr, info_ptr);

	// Allocate memory for one row (4 bytes per pixel - RGBA)
	row = (png_bytep) malloc(4 * width * sizeof(png_byte));

    // Calcul de l'échelle de couleurs automatique
	level.Value = 0;       level.R = 255;  level.G = 255;   level.B = 255;	level.A = 0;  	echelle_couleurs.push_back(level);
	level.Value = 1;       level.R = 85;   level.G = 255;   level.B = 0;    level.A = 255;	echelle_couleurs.push_back(level);
	level.Value = 2;       level.R = 145;  level.G = 255;   level.B = 0;    level.A = 255;	echelle_couleurs.push_back(level);
	level.Value = 3;       level.R = 200;  level.G = 255;   level.B = 0;    level.A = 255;	echelle_couleurs.push_back(level);
	level.Value = 4;       level.R = 255;  level.G = 255;   level.B = 0;    level.A = 255;	echelle_couleurs.push_back(level);
	level.Value = 5;       level.R = 255;  level.G = 225;   level.B = 0;    level.A = 255;	echelle_couleurs.push_back(level);
	level.Value = 6;       level.R = 255;  level.G = 175;   level.B = 0;    level.A = 255;	echelle_couleurs.push_back(level);
	level.Value = 7;       level.R = 255;  level.G = 100;   level.B = 0;    level.A = 255;	echelle_couleurs.push_back(level);
	level.Value = 8;       level.R = 255;  level.G = 50;    level.B = 0;    level.A = 255;	echelle_couleurs.push_back(level);
	level.Value = 9;       level.R = 225;  level.G = 0;     level.B = 0;    level.A = 255;	echelle_couleurs.push_back(level);
	for (x=0;x<width;x++)
		for (y=0;y<height;y++)
             tri.push_back(int(buffer[x][y]));
	TriVonNeumann(&tri);
	for(l=0;l<tri.size();++l)
	{
		if(l==0)
			unique.push_back(tri[l]);
		else if(tri[l]!=tri[l-1])
			unique.push_back(tri[l]);
	}
	if(unique.size()<echelle_couleurs.size())
	{
		p = echelle_couleurs.size()-1;
		for(l=unique.size()-1;l>=0;l--)
			if(unique[l]>0)
			{
				echelle_couleurs[p].Value = unique[l];
				p--;
			}
		while(p>1)
		{
			echelle_couleurs[p].Value = echelle_couleurs[p+1].Value;
			p--;
		}
	}
	else
	{
		float l0 = (float)tri[tri.size()-1]/(float)(echelle_couleurs.size()-1);
		for(p=1;p<echelle_couleurs.size();++p)
			echelle_couleurs[p].Value = l0*p;
	}

	/*for(p=0;p<echelle_couleurs.size();++p)
		cout << ToString(echelle_couleurs[p].Value).c_str() << endl;*/

	// Write image data
	for (y=height-1;y>=0;y--)
	{
		for (x=0;x<width;x++)
		{			
			if(buffer[x][y]<echelle_couleurs[0].Value)
			{
				(&row[x*4])[0] = 255;
				(&row[x*4])[1] = 255;
				(&row[x*4])[2] = 255;
				(&row[x*4])[3] = 0;
			}
			else if(buffer[x][y]>=echelle_couleurs[echelle_couleurs.size()-1].Value)
			{
				(&row[x*4])[0] = echelle_couleurs[echelle_couleurs.size()-1].R;
				(&row[x*4])[1] = echelle_couleurs[echelle_couleurs.size()-1].G;
				(&row[x*4])[2] = echelle_couleurs[echelle_couleurs.size()-1].B;
				(&row[x*4])[3] = echelle_couleurs[echelle_couleurs.size()-1].A;
			}	
            else
            {
                for(l=0;l<echelle_couleurs.size()-1;++l)
                {
                    if(buffer[x][y]>=echelle_couleurs[l].Value && buffer[x][y]<echelle_couleurs[l+1].Value)
                    {				
						// Sans interpolation des couleurs entre les seuils :
						(&row[x*4])[0] = echelle_couleurs[l].R;
                        (&row[x*4])[1] = echelle_couleurs[l].G;
                        (&row[x*4])[2] = echelle_couleurs[l].B;
						(&row[x*4])[3] = echelle_couleurs[l].A; 
                        break;
                    }
				}
			}
		}
		png_write_row(png_ptr, row);
	}

	// End write
	png_write_end(png_ptr, NULL);

	if (fp != NULL) fclose(fp);
	if (info_ptr != NULL) png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
	if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
	if (row != NULL) free(row);

	return true;
}

//---------------------------------------------------------------------------
bool writePNG(String filename, int width, int height, float **buffer, vector<Level> echelle_couleurs)
{
	int x, y, l, p;
	FILE *fp = NULL;
	png_structp png_ptr = NULL;
	png_infop info_ptr = NULL;
	png_bytep row = NULL;
	
	if(FileExists(filename))
	{
		DeleteFile(filename);
		if(FileExists(filename))
			return false;
	}
	
	// Open file for writing (binary mode)
	fp = fopen(filename.c_str(), "wb");
	if (fp == NULL) {
		fprintf(stderr, "Could not open file %s for writing\n", filename.c_str());
		return false;
	}

	// Initialize write structure
	png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	if (png_ptr == NULL) {
		fprintf(stderr, "Could not allocate write struct\n");
		if (fp != NULL) fclose(fp);
		if (info_ptr != NULL) png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
		if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
		if (row != NULL) free(row);
		return false;
	}

	// Initialize info structure
	info_ptr = png_create_info_struct(png_ptr);
	if (info_ptr == NULL) {
		fprintf(stderr, "Could not allocate info struct\n");
		if (fp != NULL) fclose(fp);
		if (info_ptr != NULL) png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
		if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
		if (row != NULL) free(row);
		return false;
	}

	// Setup Exception handling
	if (setjmp(png_jmpbuf(png_ptr))) {
		fprintf(stderr, "Error during png creation\n");
		if (fp != NULL) fclose(fp);
		if (info_ptr != NULL) png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
		if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
		if (row != NULL) free(row);
		return false;
	}

	png_init_io(png_ptr, fp);

	// Write header (8 bit colour depth)
	png_set_IHDR(png_ptr, info_ptr, width, height,
			8, PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE,
			PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

	png_write_info(png_ptr, info_ptr);

	// Allocate memory for one row (4 bytes per pixel - RGBA)
	row = (png_bytep) malloc(4 * width * sizeof(png_byte));

   	// Write image data
	for (y=height-1;y>=0;y--)
	{
		for (x=0;x<width;x++)
		{			
			if(buffer[x][y]<echelle_couleurs[0].Value)
			{
				(&row[x*4])[0] = 255;
				(&row[x*4])[1] = 255;
				(&row[x*4])[2] = 255;
				(&row[x*4])[3] = 0;
			}
			else if(buffer[x][y]>=echelle_couleurs[echelle_couleurs.size()-1].Value)
			{
				(&row[x*4])[0] = echelle_couleurs[echelle_couleurs.size()-1].R;
				(&row[x*4])[1] = echelle_couleurs[echelle_couleurs.size()-1].G;
				(&row[x*4])[2] = echelle_couleurs[echelle_couleurs.size()-1].B;
				(&row[x*4])[3] = echelle_couleurs[echelle_couleurs.size()-1].A;
			}	
            else
            {
                for(l=0;l<echelle_couleurs.size()-1;++l)
                {
                    if(buffer[x][y]>=echelle_couleurs[l].Value && buffer[x][y]<echelle_couleurs[l+1].Value)
                    {
						// Interpolation des couleurs entre les seuils :
						float difflvl  = echelle_couleurs[l+1].Value - echelle_couleurs[l].Value;
                        float difflvl1 = buffer[x][y] - echelle_couleurs[l].Value;
                        difflvl1 = difflvl1 / difflvl;
                        float difflvl2 = 1-difflvl1;
                        (&row[x*4])[0] = (echelle_couleurs[l].R * difflvl2)+ (echelle_couleurs[l+1].R * difflvl1);
                        (&row[x*4])[1] = (echelle_couleurs[l].G * difflvl2)+ (echelle_couleurs[l+1].G * difflvl1);
                        (&row[x*4])[2] = (echelle_couleurs[l].B * difflvl2)+ (echelle_couleurs[l+1].B * difflvl1);
						(&row[x*4])[3] = (echelle_couleurs[l].A * difflvl2)+ (echelle_couleurs[l+1].A * difflvl1);
                        break;
                    }
				}
			}
		}
		png_write_row(png_ptr, row);
	}

	// End write
	png_write_end(png_ptr, NULL);

	if (fp != NULL) fclose(fp);
	if (info_ptr != NULL) png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
	if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
	if (row != NULL) free(row);

	return true;
}
#endif