#include <iostream>
#include <fstream>
#include <vector>
#include <limits> // For numeric_limits
#include "visualizer.h"
using namespace std;

/* Type: ImageData
 * ---------------------------------------------
 * A type representing image data tagged with a
 * label.
 */
struct ImageData
{
  vector<double> image; // Raw picture data
  int label;            // Label from 0 - 9
};

/* Constant: kImageDimension
 * Value: The size of one side of an image.
 */
const size_t kImageDimension = 28;

/* Constant: kImageSize
 * Value: The number of pixels in an image.
 */
const size_t kImageSize = kImageDimension * kImageDimension;

/* Constant: kLearningRate
 * Value: The learning rate (alpha) for perceptron learning.
 */
const double kLearningRate = 1.0e-5;

/* Constant: kDesiredAccuracy
 * Value: What accuracy we want before considering a perceptron "trained."
 */
const double kDesiredAccuracy = 5.0e-3;

/* Function: LoadData(string filename)
 * Usage: vector<ImageData> data = LoadData("data/myData")
 * --------------------------------------------------------
 * Loads a set of training image examples from the specified
 * data file.  This function assumes the file exists and the
 * data is formatted correctly and will not behave correctly
 * if these assumptions are invalid.
 */
vector<ImageData> LoadData(string filename)
{
	ifstream input(filename.c_str()); // Read from the specified file
	vector<ImageData> result;  // Holds the result of the operation

	/* Continue reading out of the file until no more data exists. */
	while (true)
	{
		ImageData next;
		
		/* We need space to hold the entire picture, plus one additional
		 * slot for the constant +1 term.  This code will ensure that space
		 * exists, then will set the last element of the vector to 1.
		 */
		next.image.resize(kImageSize + 1);
		next.image.back() = 1.0;

		/* Read in data from the file.  We don't do any error-checking here since
		 * we'll do that later.
		 */
		for (size_t i = 0; i < kImageSize; ++i)
			input >> next.image[i];

		/* Read in the label. */
		input >> next.label;

		/* Now, we need to ensure that this read operation completed successfully.
		 * If it didn't, we've exhausted the data in the file and should return.
		 * Otherwise, we should add our new image to the list of examples.
		 */
		if (input.fail()) return result;

		result.push_back(next);
	}
}

/* Function: StepFunction(double value)
 * Usage: int response = StepFunction(x);
 * -----------------------------------------------------------
 * Returns 1 if the value is positive and zero otherwise.
 */
int StepFunction(double value)
{
	/* In C++, all boolean operators yield a 1 or 0 based on whether the result is
	 * true or false.  We abuse this here to implement the function in a single line
	 * of code.
	 */
	return value > 0;
}

/* Function: DotProduct(vector<double>& one, vector<double>& two)
 * Usage: double result = DotProduct(one, two);
 * -------------------------------------------------------------
 * Computes the dot product of two vectors (the sum of the componentwise product).
 * This function will cause a runtime error if the two vectors are not the same
 * size.
 */
double DotProduct(vector<double>& one, vector<double>& two)
{
	double result = 0.0;
	
	for (size_t i = 0; i < one.size(); ++i)
		result += one[i] * two[i];
	
	return result;
}

/* Function: GoodEnough(vector<int>& errors)
 * Usage: if (GoodEnough(errors)) break;
 * ---------------------------------------------------------------
 * Returns whether a perceptron has trained enough on the test set.
 * This is true if the number of errors is less than the
 * kDesiredAccuracy constant defined at the top of the program.
 */
bool GoodEnough(vector<int>& errors)
{
	size_t numWrong = 0;

	/* Count the number of errors (the number of terms which
	 * were nonzero)
	 */
	for (size_t i = 0; i < errors.size(); ++i)
		if (errors[i] != 0) ++numWrong;

	return numWrong < errors.size() * kDesiredAccuracy;
}

/* Function: TrainPerceptron(vector<ImageData>& data, int digit)
 * Usage: vector<double> weights = TrainPerceptron(data, 0);
 * ---------------------------------------------------------------
 * Trains a perceptron to recognize a particular digit given a
 * training set.  This function will returns weights that have an
 * error no greater than kDesiredAccuracy.
 */
vector<double> TrainPerceptron(vector<ImageData>& data, int digit)
{
	/* The weights vector.  Initially, default these to zero, though
	 * in actuality we could pick any random values we'd like.
	 */
	vector<double> weights(kImageSize + 1);

	/* Continue looping until we have the proper accuracy */
	while (true)
	{
		/* Compute the error vector, which is the difference between our
		 * perceptron's output and the actual output.
		 */
		vector<int> errors(data.size());

		for (size_t i = 0; i < data.size(); ++i)
		{
			/* Compute perceptron's response */
			int response = StepFunction(DotProduct(weights, data[i].image));

			/* Compute actual answer, which is 1 if the digit matches what we're
			 * classifying and zero otherwise.
			 */
			int actual = (data[i].label == digit);

			/* Store the difference. */
			errors[i] = actual - response;
		}

	/* If our error is within limits, stop training. */
	if (GoodEnough(errors)) return weights;

	/* Otherwise,, update the weights according to the update rule. */
	for (size_t coord = 0; coord < weights.size(); ++coord)
		for (size_t i = 0; i < data.size(); ++i)
			weights[coord] += kLearningRate * errors[i] * data[i].image[coord];
	}
}

/* Function: TestPerceptrons(vector< vector<double> >& perceptrons, vector<ImageData>&);
 * Usage: TestPerceptrons(myNetwork, testData)
 * -------------------------------------------------------------
 * Tests the perceptron network on a data set and reports the error.
 */
void TestPerceptrons(vector< vector<double> >& perceptrons, vector<ImageData>& data)
{
	size_t numCorrect = 0;
	
	/* Check performance on each data point. */
	for (size_t i = 0; i < data.size(); ++i)
	{
		/* Find the perceptron with the highest confidence.  Initially, make a nonsensical
		 * guess, then update based on the perceptron.
		 */
		int ourGuess = -1;
		double bestConfidence = -numeric_limits<double>::infinity();

		/* Compute the response and update the best known value if necessary. */
		for (size_t p = 0; p < perceptrons.size(); ++p)
		{
			double response = DotProduct(perceptrons[p], data[i].image);

			if (response > bestConfidence)
			{
				bestConfidence = response;
				ourGuess = p;
			}
		}
	}

	/* If we were correct, increment the total. */
	if (data[i].label == ourGuess)
		++numCorrect;
	}

	cout << "Accuracy: " << (100.0 * numCorrect) / data.size() << '%' << endl;
}

/* Main entry point. */
int main()
{
	/* Load the training data.  PC should use "data\\large-training-data" as
	 * the filename.
	 */
	vector<ImageData> trainingData = LoadData("data/large-training-data");
	cout << "Number of training examples: " << trainingData.size() << endl;

	/* Train and visualize perceptrons. */
	vector< vector<double> > perceptrons;

	for (int i = 0; i < 10; ++i)
	{
		cout << "Training a perceptron to recognize " << i << "..." << endl;
		perceptrons.push_back(TrainPerceptron(trainingData, i));

		SaveVisualization(perceptrons.back(), i);
	}

	/* Load test data and see how we did.  PC should use "data\\testing-data"
	 * instead as the filename.
	 */
	vector<ImageData> testingData = LoadData("data/testing-data");
	TestPerceptrons(perceptrons, testingData);
}