// main.cpp : Defines the entry point for the console application.
//

#include "config.h"
#include <commdlg.h>
#include <conio.h>
#include <iostream>

#include "cximage/ximage.h"
#include "xbrowseforfolder.h"
// function prototypes
bool CalculateDescriptors(const char *basedir);
bool CategorizeDescriptors(const char *basedir);
bool DetermineWinterSportSelect();
bool DetermineWinterSportBatch(const char *basedir);
bool LoadAverages(const char *basedir);

#define TRAINING_DIR_DEBUG "C:\\Documents and Settings\\rvdzwet\\Desktop\\liacs\\MIR2010\\trainingSet"
#define TESTSET_DIR_DEBUG "C:\\Documents and Settings\\rvdzwet\\Desktop\\liacs\\MIR2010\\testSet"
#define DEBUG 0
// number of bins to use for each color
// Note: number of bins must be a whole fraction of 256. if we would
//       use 256 bins for each color then a single histogram would be
//       64MB, so we should choose a more sensible number
#define BIN_COUNT	32

// names of categories
#define CATEGORY_SIZE 8
#define CATEGORY_1	"cat1.crowd"
#define CATEGORY_2	"cat2.skijump"
#define CATEGORY_3	"cat3.snowboarding"
#define CATEGORY_4  "opt1.bobsleigh"
#define CATEGORY_5  "opt2.icehockey"
#define CATEGORY_6  "opt3.speedskating"
#define CATEGORY_7  "opt4.downhillski"
#define CATEGORY_8  "opt5.curling"

const char *categories[CATEGORY_SIZE] = { CATEGORY_1, CATEGORY_2, CATEGORY_3, 
	CATEGORY_4, CATEGORY_5, CATEGORY_6,  CATEGORY_7, CATEGORY_8};  
float average[CATEGORY_SIZE][BIN_COUNT*BIN_COUNT*BIN_COUNT];

// Some prototyping foo on functions
#define _return(err_msg, retval) printf("DEBUG: %s\n",err_msg); cin.get(); return(retval);

bool verbose = false;

bool file_exists(const char * filename)
{
	if (FILE * file = fopen(filename, "r"))
	{
		fclose(file);
		return true;
	}
	return false;
}


int main(int argc, char **argv)
{
	char dbdir[MAX_PATH];
	char testdir[MAX_PATH];

#if defined(TRAINING_DIR_DEBUG)	
	if (strcpy_s(dbdir, TRAINING_DIR_DEBUG) != 0) {
		_return("Cannot copy TRAINING_DIR_DEBUG",1);
	}
	if (strcpy_s(testdir, TESTSET_DIR_DEBUG) != 0) {
		_return("Cannot copy TESTSET_DIR_DEBUG",2);
	}
#else
	// ask the user for the image database
	// in the image directory, the images are stored in categorized
	// folders. store the descriptors we calculate in the same
	// folder as the image
	if (XBrowseForFolder(NULL, "Please select image database folder", NULL, dbdir, sizeof(dbdir)) == FALSE)
		return 0;
	if (strlen(dbdir) == 0)
		return 0;
	if (XBrowseForFolder(NULL, "Please select image testset folder", NULL, testdir, sizeof(testdir)) == FALSE)
		return 0;
	if (strlen(testdir) == 0)
		return 0;

#endif

#if DEBUG
	if (!LoadAverages(dbdir)){
		_return("Unable to load averages",1);
	}

	if (!DetermineWinterSportBatch(testdir)) {
		_return("could not run winter sport batch",1);
	}
	return(0);
#endif

	// ask the user which option to use
	while (true)
	{
		//system("cls");
		cout << "Using database directory: " << dbdir << endl;
		cout << "Using test directory: " << testdir << endl;
		cout << "Using categories: " << endl;
		for (int i = 0; i < CATEGORY_SIZE; i++) 
		{
			cout << "  - " << categories[i] << endl;
		}
		cout << "***************************" << endl;
		cout << "*  Winter Olympic Imagery *" << endl;
		cout << "***************************" << endl;
		cout << endl;
		cout << "1. calculate descriptors" << endl;
		cout << "2. categorize descriptors (aka averages histograms)" << endl;
		cout << "3. determine winter sport" << endl;
		cout << "4. batch test winter sport" << endl;
		cout << "v. Turn verbose mode ";
		if (verbose)
			cout << "off" << endl;
		else
			cout << "on" << endl;
		cout << endl;
		cout << "Please select option, or type 'q' to quit: ";
		char c = _getch();
		cout << c << endl;
		fflush(stdin);
		// start the chosen option
		switch (c)
		{
		case 'q':
			return 0;
		case '1':
			if (!CalculateDescriptors(dbdir)) {
				_return("could not calculate descriptors",1);
			}
			break;
		case '2':
			if (!CategorizeDescriptors(dbdir)){
				_return("could not categorize descriptors",1);
			}
			break;
		case '3':
			if (!LoadAverages(dbdir)){
				_return("Unable to load averages",1);
			}

			if (!DetermineWinterSportSelect()){
				_return("could not determine winter sport",1);
			}
			break;
		case '4':
			if (!LoadAverages(dbdir)){
				_return("Unable to load averages",1);
			}

			if (!DetermineWinterSportBatch(testdir)) {
				_return("could not run winter sport batch",1);
			}
			break;
		case 'v':
			verbose = (!verbose);
		default:
			continue;
		}
	}
	return 0;
}

// histogram should be a preallocated array of size BIN_COUNT*BIN_COUNT*BIN_COUNT elements and will
// be filled with the color histogram of the image where path points at
bool CalculateDescriptor(const char *path, float *histogram)
{
	// load the image
	CxImage image(path, CXIMAGE_FORMAT_JPG);
	if (!image.IsValid())
		return false;
	// clear histogram
	memset(histogram, 0, BIN_COUNT*BIN_COUNT*BIN_COUNT * sizeof(float));
	// walk through the pixels to fill the histogram
	int width = (int)image.GetWidth();
	int height = (int)image.GetHeight();
	int bin_r, bin_g, bin_b;
	for (int y = 0; y < height; y++)
	{
		for (int x = 0; x < width; x++)
		{
			// Note: CxImage library starts counting at lower-left corner of the image,
			//       which is seen as the top of the image. however, usually images
			//       start counting from the top-left corner of the image. thus if you
			//       want to get pixel(2, 2) from the top-left you would have to ask
			//       for pixel (2, height - 2 - 1) from CxImage. although in this
			//       situation we don't care which pixel is where, we only care about
			//       its color.
			RGBQUAD rgb = image.BlindGetPixelColor(x, y, false);
			// determine the bin this color falls in
			bin_r = rgb.rgbRed   / (256 / BIN_COUNT);
			bin_g = rgb.rgbGreen / (256 / BIN_COUNT);
			bin_b = rgb.rgbBlue  / (256 / BIN_COUNT);
			histogram[bin_r*BIN_COUNT*BIN_COUNT + bin_g*BIN_COUNT + bin_b]++;
		}
	}
	// normalize the histogram so that all together the values will add up
	// to one. since there are width * height pixels, we divide each value
	// by this amount
	for (int i = 0; i < BIN_COUNT*BIN_COUNT*BIN_COUNT; i++)
		histogram[i] /= width * height;
	return true;
}

bool CalculateDescriptors(const char *basedir)
{
	// the histogram that we reuse for each image
	float *histogram = new float[BIN_COUNT*BIN_COUNT*BIN_COUNT];
	// walk through all images
	// Note: each of the three categories has 50 images
	char path[MAX_PATH];
	char catdir[MAX_PATH];
	const char *catname;
	FILE *file = NULL;
	for (int c = 0; c < CATEGORY_SIZE; c++)
	{
		catname = categories[c];
		sprintf(catdir, "%s\\%s\\", basedir, catname);
		cout << "[" << catname << "] Using directory " << catdir << endl;

		// process the images in the directory
		for (int i = 1; i <= 50; i++)
		{
			SAFE_SPRINTF(path, sizeof(path), "%s%i.jpg", catdir, i);
			if (!file_exists(path)) {
				continue;
			}
			cout << "[" << catname << "] processing image " << i << endl;
			// calculate the histogram descriptor
			if (!CalculateDescriptor(path, histogram))
				goto failure;
			// save the descriptor to disk
			SAFE_SPRINTF(path, sizeof(path), "%s%i.dat", catdir, i);
			if ((file = fopen(path, "wb")) == NULL)
				goto failure;
			if (fwrite(histogram, sizeof(float), BIN_COUNT*BIN_COUNT*BIN_COUNT, file) != BIN_COUNT*BIN_COUNT*BIN_COUNT)
				goto failure;
			SAFE_CLOSEFILE(file);
		}
	}
	// release resources
	SAFE_DELETE_ARRAY(histogram);
	return true;

failure:
	SAFE_CLOSEFILE(file);
	SAFE_DELETE_ARRAY(histogram);
	return false;
}

bool CategorizeDescriptors(const char *basedir)
{
	// analyze the descriptors per category to determine the
	// characteristics of that category
	float *histogram = new float[BIN_COUNT*BIN_COUNT*BIN_COUNT];
	float *average = new float[BIN_COUNT*BIN_COUNT*BIN_COUNT];
	// walk through all descriptors
	char path[MAX_PATH];
	char catdir[MAX_PATH];
	const char *catname;
	FILE *file = NULL;
	int c_size = 0;
	for (int c = 0; c < CATEGORY_SIZE; c++)
	{
		c_size = 0;
		catname = categories[c];
		sprintf(catdir,"%s\\%s\\", basedir, catname);

		// average all descriptors
		memset(average, 0, BIN_COUNT*BIN_COUNT*BIN_COUNT * sizeof(float));
		for (int i = 1; i <= 50; i++)
		{
			SAFE_SPRINTF(path, sizeof(path), "%s%i.dat", catdir, i);
			if (!file_exists(path)) {
				continue;
			}
			cout << "[" << catname << "] processing image " << i << endl;
			// load the histogram descriptor
			if ((file = fopen(path, "rb")) == NULL)
				goto failure;
			if (fread(histogram, sizeof(float), BIN_COUNT*BIN_COUNT*BIN_COUNT, file) != BIN_COUNT*BIN_COUNT*BIN_COUNT)
				goto failure;
			SAFE_CLOSEFILE(file);
			// add the value of each bin to the average
			for (int b = 0; b < BIN_COUNT*BIN_COUNT*BIN_COUNT; b++)
				average[b] += histogram[b];

			c_size++;
		}

		for (int b = 0; b < BIN_COUNT*BIN_COUNT*BIN_COUNT; b++)
			average[b] /= c_size;
		// save the average to disk
		SAFE_SPRINTF(path, sizeof(path), "%s%s.dat", catdir, "average");
		if ((file = fopen(path, "wb")) == NULL)
			goto failure;
		if (fwrite(average, sizeof(float), BIN_COUNT*BIN_COUNT*BIN_COUNT, file) != BIN_COUNT*BIN_COUNT*BIN_COUNT)
			goto failure;
		SAFE_CLOSEFILE(file);
	}
	// release resources
	SAFE_DELETE_ARRAY(histogram);
	SAFE_DELETE_ARRAY(average);
	return true;

failure:
	SAFE_CLOSEFILE(file);
	SAFE_DELETE_ARRAY(histogram);
	SAFE_DELETE_ARRAY(average);
	return false;
}

bool LoadAverages(const char *basedir) {
	/* determine the distance to each category */
	const char *catname;
	char catdir[MAX_PATH];
	char path[MAX_PATH];
	FILE *file = NULL;

	for (int c = 0; c < CATEGORY_SIZE; c++)
	{
		catname = categories[c];
		sprintf(catdir, "%s\\%s\\", basedir, catname);

		// load the average from disk
		SAFE_SPRINTF(path, sizeof(path), "%s%s.dat", catdir, "average");
		if ((file = fopen(path, "rb")) == NULL) {
			cout << "Cannot open " << path << endl;
			return false;
		}
		if (fread(average[c], sizeof(float), BIN_COUNT*BIN_COUNT*BIN_COUNT, file) != BIN_COUNT*BIN_COUNT*BIN_COUNT)
			return false;
		SAFE_CLOSEFILE(file);
	}
	return true;
}

int DetermineCategory(const char *path, const int guess=-1) {
	float *histogram = new float[BIN_COUNT*BIN_COUNT*BIN_COUNT];
	float cat2dist[CATEGORY_SIZE];

	/* First category default best canidate */
	int retval = 0;

	/* calculate the histogram of the image */
	if (!CalculateDescriptor(path, histogram))
		return -1;

	/* determine the distance to each category */
	for (int c = 0; c < CATEGORY_SIZE; c++)
	{
		// determine distance
		cat2dist[c] = 0.0f;
		for (int b = 0; b < BIN_COUNT*BIN_COUNT*BIN_COUNT; b++)
			cat2dist[c] += fabs(histogram[b] - average[c][b]);
	}

	/* determine the winning category */
	for (int i = 1; i < CATEGORY_SIZE; i++)
		if (cat2dist[i] < cat2dist[retval])
			retval = i;

	if (verbose) {
		/* Dirty hack to show some more details in case of failure */
		if (guess != -1 && guess != retval) {
			for (int i = 0; i < CATEGORY_SIZE; i++) {		
				printf("%s Distance to %-20s: %f %s\n", (retval == i) ? "*" : " ", 
					categories[i], cat2dist[i],(retval == i) ? "*" : "");
			}
		}
	}

	/* return result */
	return retval;
}

/*  ask for an input image and determine the most likely
* category it belongs to
*/
bool DetermineWinterSportSelect()
{
	float *histogram = new float[BIN_COUNT*BIN_COUNT*BIN_COUNT];
	char path[MAX_PATH] = {0};
	char catdir[MAX_PATH] = {0};
	const char *catname = NULL;
	FILE *file = NULL;
	int c = NULL;
	float cat2dist[CATEGORY_SIZE] = {0};

	OPENFILENAME ofn = {0};
	//ofn.lpstrFilter = "Image files\0*.jpg;*.png;*.bmp\0\0";
	ofn.lpstrFilter = "Image files\0*.jpg\0\0";
	ofn.lpstrFile = path;
	ofn.nMaxFile = MAX_PATH;
	ofn.lpstrTitle = "Choose image file";
	ofn.Flags = OFN_PATHMUSTEXIST | OFN_FILEMUSTEXIST;
	ofn.lStructSize = sizeof(OPENFILENAME);

	if (GetOpenFileName(&ofn) == FALSE)
		goto failure;

	if ((c = DetermineCategory(path,-2)) == -1)
		return false;

	cout << "The category this image belongs is category: " << categories[c] << " [" << c << "]" << endl;
	cout << "Press any key to continue... ";
	_getch();
	fflush(stdin);
	// release resources
	SAFE_DELETE_ARRAY(histogram);
	return true;

failure:
	SAFE_CLOSEFILE(file);
	SAFE_DELETE_ARRAY(histogram);
	return false;
}

bool DetermineWinterSportBatch(const char *basedir)
{
	const char *catname;
	char catdir[MAX_PATH];
	char path[MAX_PATH];

	int all_total = 0;
	int all_succes = 0;
	int c_total[CATEGORY_SIZE] = {0};
	int c_succes[CATEGORY_SIZE] = {0};

	/* determine the distance to each category */
	for (int c = 0; c < CATEGORY_SIZE; c++)
	{
		catname = categories[c];
		sprintf(catdir, "%s\\%s\\", basedir, catname);
		
		/* process the images in the directory */
		for (int i = 1; i <= 50; i++)
		{
			SAFE_SPRINTF(path, sizeof(path), "%s%i.jpg", catdir, i);
			if (!file_exists(path)) {
				continue;
			}

			c_total[c]++;
			/* Check if file matches category */
			if (DetermineCategory(path,c) == c) {
				cout << "[" << catname << "] testing image " << i << " : OK" << endl;
				c_succes[c]++;
			} else {
				cout << "[" << catname << "] testing image " << i << " : FAIL" << endl;			}
		}
		cout << "[" << catname << "] results " << c_succes[c] << "/" << c_total[c] << endl;
	}
	
	/* Display grand total */
	cout << "=== Totals ===" << endl;
	for (int c = 0; c < CATEGORY_SIZE; c++)
	{
		catname = categories[c];
		printf ("[%-20s] %i/%i\n",catname,c_succes[c],c_total[c]);
		all_total += c_total[c];
		all_succes += c_succes[c];
	}
	printf ("[%-20s] %i/%i\n","total",all_succes,all_succes);
	cout << "Press any key to continue..."; cin.get();


	return true;
}