/*********************************************************************
 * This file is part of the PRAPI library.
 *
 * Copyright (C) 2001 Jaakko Viertola
 * Copyright (C) 2002 Topi Mäenpää
 * All rights reserved.
 *
 * This program is free software. You can redistribute and/or modify
 * it under the terms of the free software licence found in the
 * accompanying file "COPYING". The licence terms must always be
 * redistributed with this source file. The above copyright notice
 * must be reproduced in all modified and unmodified copies of this
 * source file.
 *
 * $Revision: 1.5 $
 *********************************************************************/

#ifndef _OUTEXCLASSIFICATIONENGINE_H
#define _OUTEXCLASSIFICATIONENGINE_H

#include <fstream>
#include <string>

#include <Util.h>
#include <List.h>
#include <Matrix.h>
#include <Math.h>
#include <MatrixCodec.h>

#include "../Sample.h"
#include "../ConfusionMatrix.h"
#include "../RasterCodec.h"
#include "../BmpCodec.h"

#include "OutexResult.h"

namespace prapi { namespace extras {

	/**
	 * Thrown when something is wrong with a test suite.
	 **/
	class OutexException : public util::Exception
	{
	public:
		OutexException(std::string message) : util::Exception(message) {}
	};
	
	/**
	 * This class automatically parses and classifies Outex test suites.
	 * It automatically creates sample sets and asks a subclass to
	 * perform the actual classification on each iteration of the test
	 * suite.
	 **/	
	template <class T, class I=std::string, class C=int> class OutexClassificationEngine
	{
	public:
		/**
		 * Construct a classification engine.
		 *
		 * @param precalculate a flag that controls the memory consumption
		 * versus speed trade-off. If set to true (the default), feature
		 * vectors for all images are calculated and stored in the main
		 * memory before classification. If there is not enough memory for
		 * this, the flag can be set to false, in which case a feature
		 * vector is calculated each time an image is accessed. Memory
		 * consumption is kept at minimum, but feature vectors must be
		 * calculated many times for each image.
		 **/
		OutexClassificationEngine(bool precalculate = true) :
			_bPrecalculate(precalculate),
			_iStartIndex(0), _iEndIndex(-1) {}

		/**
		 * Clean up any reserved resources.
		 **/
		virtual ~OutexClassificationEngine(){}

		/**
		 * Set the precalculation flag.
		 *
		 * @see OutexClassificationEngine(bool)
		 **/
		void setPrecalculation(bool precalculate) { _bPrecalculate = precalculate; }

		/**
		 * Get the current value of the precalculation flag.
		 **/
		bool getPrecalculation() { return _pPrecalculate; }

		/**
		 * Set the start and end indices for the classification problems
		 * to be run. (By default, all problems are classified.)
		 *
		 * @param startIndex the first problem to be classified.
		 * @param endIndex the last problem to be classified. Negative
		 * value means 'as many as possible'.
		 **/
		void setRange(int startIndex=0, int endIndex=-1)
		{
			_iStartIndex = startIndex;
			_iEndIndex = endIndex;
		}

		/**
		 * Get the current start index.
		 **/
		int getStartIndex() const { return _iStartIndex; }

		/**
		 * Get the current end index.
		 **/
		int getEndIndex() const { return _iEndIndex; }
		
		/**
		 * Whenever the engine encounters a gray-scale image, this method
		 * is called. The method should extract the needed features and
		 * return them as a feature vector. The default implementation
		 * returns an empty list.
		 *
		 * @param mat a gray-scale image matrix
		 **/
		virtual util::List<T> getFeatureVector(const util::Matrix<int>& mat) { List<T> lst; return lst; }

		/**
		 * Whenever a color image is encountered, this method is consulted
		 * to extract the features. The default implementation returns an
		 * empty list.
		 *
		 * @param mat a color image
		 **/
		virtual util::List<T> getFeatureVector(const util::Matrix<RGBColor<> >& mat){ List<T> lst;return lst; }
		
		/**
		 * Once the engine has constructed the necessary sample sets for
		 * classification, it calls this method. The method can use
		 * whatever classification principle you feel suitable. As a
		 * result, the classification of each sample in the test set
		 * should be performed.
		 *
		 * @param train The train set which is used for training.
		 * @param test The test set which is used for testing.
		 * @param classCount the total number of classes
		 **/
		virtual void classify(util::List<Sample<T,I,C> >& train, util::List<Sample<T,I,C> >& test, int classCount) = 0;

		/**
		 * Once each iteration has been completed, the confusion matrices
		 * from each classification are collected in a list. You can get
		 * the list by calling this method.
		 **/
		util::List<ConfusionMatrix> getConfusionMatrices() { return _lstConfusionMatrices; }
		/**
		 * Get the class names from each classification.
		 **/
		util::List<util::List<std::string> > getClasses() { return _lstClasses; }
		/**
		 * Get the class names at the specified classification round.
		 **/
		util::List<std::string> getClasses(int index) throw (InvalidArgumentException&)
		{
			if (index >= _lstClasses.getLength())
				throw InvalidArgumentException("OutexClassificationEngine::getClasses(int): Invalid round index: " + String::toString(index));
			return _lstClasses[index >= _lstClasses.getLength() ? 0 : index];
		}
		
		/**
		 * The function evaluateSuite classifies all the problems in the test suite
		 * and returns the result in OutexResult format.
		 *
		 * @param path A path to the root of the test suite.
		 *
		 * @return The result of the classification.
		 **/
		OutexResult evaluateSuite(std::string path) throw (OutexException&);

	private:
		/**
		 * function reads the Outex format txt file and returns
		 * the list of strings in file.
		 **/
		static util::List<std::string> readOutexTxtFile(std::string fileName) throw (OutexException&);
		/**
		 * function reads the Outex format classes txt file.
		 * Note that the index of class is the same as the place
		 * in the classes list.
		 **/
		static void readOutexClassesFile(std::string fileName,util::List<std::string>& classes,util::List<int>& cost)
			throw (OutexException&);
		/**
		 * function reads the Outex format train/test txt file.
		 * Note that the index of class is the same as the place
		 * in the classes list.
		 **/
		static void readOutexTrainTestFile(std::string fileName,util::List<std::string>& names,util::List<int>& classIndexes)
			throw (OutexException&);


	protected:
		/**
		 * The pre-calculation flag.
		 **/
		bool _bPrecalculate;
		/**
		 * Start and end indices for the problems to be evaluated.
		 **/
		int _iStartIndex, _iEndIndex;
		/**
		 * A list stroring a confusion matrix for each classification
		 * round.
		 **/
		util::List<ConfusionMatrix> _lstConfusionMatrices;
		/**
		 * Class names for each classification round.
		 **/
		util::List<util::List<std::string> > _lstClasses;
	};

	
	template <class T,class I,class C> OutexResult OutexClassificationEngine<T,I,C>::evaluateSuite(std::string path)
		throw (OutexException&)
	{
		//variables
		int i=0,imageCount=0;
		RasterCodec ras;
		ColorBmpCodec bmp;

		if (!util::String::endsWith(path,"/"))
			path += '/';
		
		// first calculate all features from the given images
		util::List<std::string> images(readOutexTxtFile(path+"images.txt"));
		imageCount=images.getLength();
		util::List<util::List<T> > features(imageCount);

		if(_bPrecalculate) // check if precalculation is wanted
			{
				//decide which codec will be used.
				if(util::String::endsWith(images[0], ".ras"))
					for(i=0;i<imageCount;i++)// read the image and calculate features
						features += getFeatureVector(ras.readFromFile(path+images[i]));
				else // color images
					for(i=0;i<imageCount;i++)// read the image and calculate features
						features += getFeatureVector(bmp.readFromFile(path+images[i]));
			}
				
		// then start processing the problems.
		util::List<std::string> problems(readOutexTxtFile(path+"problems.txt"));
		util::List<double> errors(problems.getLength()); // make the result list
		int end = _iEndIndex >= _iStartIndex ? _iEndIndex : problems.getLength()-1;
		if (end >= problems.getLength() || _iStartIndex < 0)
			throw OutexException("OutexClassificationEngine::evaluateSuite(string): Invalid problem range (" +
													 String::toString(_iStartIndex) + "-" + String::toString(_iEndIndex) +
													 "). This suite has problems 0-" + String::toString(problems.getLength()-1) + ".");

		for(int pIndex=_iStartIndex; pIndex<=end; pIndex++)
			{
				//read the class and cost information
				util::List<std::string> classes;
				util::List<int> cost;
				readOutexClassesFile(path+problems[pIndex]+"/classes.txt",classes,cost); //read the class and cost information

				_lstClasses += classes;
				// make the train set
				util::List<std::string> names; // the amount of image names in train set
				util::List<int> classIndices;  // the indices of image classes
				readOutexTrainTestFile(path+problems[pIndex]+"/train.txt",names,classIndices);
				int trainSetSize=names.getLength();
				util::List<Sample<T,I,C> > train(trainSetSize);

				// start making the train sample set
				for(i=0; i<trainSetSize; i++)
					{ 
						Sample<T,I,C> sample(classIndices[i]); // make new sample
						sample.setIdentifier(I(names[i])); // set the identifier

						if(!_bPrecalculate) // check if precalculation is not used
							{
							
								//decide which codec will be used.
								if(images[0].substr(images[0].length()-3)=="ras")
									sample.setFeatureVector(getFeatureVector(ras.readFromFile(path+"images/"+names[i])));
								else // color images
									sample.setFeatureVector(getFeatureVector(bmp.readFromFile(path+"images/"+names[i])));
							}
						else
							{
								// calculate the index of image in the list
								int imageIndex=images.indexOf("images/"+names[i]);
								if(imageIndex == -1)
									throw OutexException("OutexClassificationEngine<T,I,C>::evaluateSuite(std::string) : Image "+names[i]+" not found.");
						
								sample.setFeatureVector(features[imageIndex]);
							}
						train += sample;
						
						//cerr << "\rMaking train sample: "<<i;
					}
				//cerr << endl;

				
				// make the test set
				readOutexTrainTestFile(path+problems[pIndex]+"/test.txt",names,classIndices);
				int testSetSize=names.getLength();
				util::List<Sample<T,I,C> >  test(testSetSize);

				// start making the test sample set
				for(i=0;i<testSetSize;i++)
					{ 
						Sample<T,I,C> sample(classIndices[i]); // make new sample
						sample.setIdentifier(I(names[i])); // set the identifier

						if(!_bPrecalculate) // check if precalculation is not used
							{
							
								//decide which codec will be used.
								if(images[0].substr(images[0].length()-3)=="ras")
									sample.setFeatureVector(getFeatureVector(ras.readFromFile(path+"images/"+names[i])));
								else // color images
									sample.setFeatureVector(getFeatureVector(bmp.readFromFile(path+"images/"+names[i])));
							}
						else
							{
								// calculate the index of image in the list
								int imageIndex=images.indexOf("images/"+names[i]);
								if(imageIndex == -1)
									throw OutexException("OutexClassificationEngine<T,I,C>::evaluateSuite(std::string) : Image "+names[i]+" not found.");
						
								sample.setFeatureVector(features[imageIndex]);
							}
						test += sample;
						
						//cerr << "\rMaking test sample: "<<i;
					}
				//cerr << endl;
				
				// then classify the problem
				int classCount = classes.getLength();
				classify(train,test,classCount);
				
				//Then make the confusion Matrix.
				ConfusionMatrix cm(test,classCount);
				_lstConfusionMatrices += cm; //add to the list.

				//cm.print(cout,classes);
				
				// then calculate the result (remember to use cost functions).
				double score = 0, maxScore = 0;
				for (int r=0;r<classCount;r++)
					{
						int correct=0, incorrect=0;
						for (int c=0;c<classCount;c++)
							{
								if (r == c)correct += cm(r,c);
								else incorrect += cm(r,c);
							}
						
						score += cost[r] * incorrect;
						maxScore += (incorrect+correct) * cost[r];
					}
				errors += (score/maxScore); // add the result to the errors list.

				//cerr << "problem: "<<pIndex<<" procecced"<<endl;
			}
		
		// then make the results
		OutexResult result(1-Math::min(errors),1-Math::max(errors),1-Math::mean(errors),Math::stdev(errors));
		
		return result;
	}

	template <class T,class I,class C>
	util::List<std::string> OutexClassificationEngine<T,I,C>::readOutexTxtFile(std::string fileName)
		throw (OutexException&)
	{
		// read the first line
		std::string line="";
		std::ifstream file(fileName.c_str());
		if (!file)
			throw OutexException("OutexClassificationEngine::readOutexTxtFile(string): Cannot open " + fileName + ".");
		file >> line;

		// make the List which length is
		int numberOfProblems=Util::parseInt(line);
		util::List<std::string> result(numberOfProblems);

		// and take the 
		for(int i=0;i<numberOfProblems;i++)
			{
				file >> line;
				result += line;
			}
		file.close();
		
		return result;
	}
	
	template <class T,class I,class C>
	void OutexClassificationEngine<T,I,C>::readOutexClassesFile(std::string fileName,util::List<std::string>& classes,util::List<int>& cost)
		throw (OutexException&)
	{
		// read the first line
		std::string line="";
		std::ifstream file(fileName.c_str());
		if (!file)
			throw OutexException("OutexClassificationEngine::readOutexClassesFile(string): Cannot open " + fileName + ".");
		
		file >> line;
		
		// make the List which length is
		int numberOfClasses=Util::parseInt(line);
		classes.setCapacity(numberOfClasses);
		classes.setLength(0);
		cost.setCapacity(numberOfClasses);
		cost.setLength(0);

		// and take the 
		for(int i=0;i<numberOfClasses;i++)
			{
				file >> line; // read the class name
				classes += line; 
				file >> line; // read the class index
				file >> line; // read the cost
				cost += Util::parseInt(line);
			}
		file.close();
	}

	template <class T,class I,class C>
	void OutexClassificationEngine<T,I,C>::readOutexTrainTestFile(std::string fileName,util::List<std::string>& names,util::List<int>& classIndices)
		throw (OutexException&)
	{
		// read the first line
		std::string line="";
		std::ifstream file(fileName.c_str());
		if (!file)
			throw OutexException("OutexClassificationEngine::readOutexTrainTestFile(string): Cannot open " + fileName + ".");
		file >> line;

		int numberOfSamples=Util::parseInt(line);
		names.setCapacity(numberOfSamples);
		names.setLength(0);
		classIndices.setCapacity(numberOfSamples);
		classIndices.setLength(0);

		for(int i=0; i<numberOfSamples; i++)
			{
				file >> line; // read the sample name
				names += line;
				file >> line; // read the class Index
				classIndices += Util::parseInt(line);
			}
		file.close();
	}
	
}}
#endif