#include "Training.h"

using namespace cv;
using namespace std;


Training::Training()
{

}


Training::~Training()
{
}

void Training::CreateTrainingSet()
{
	Camera cam = Camera(1);
	bool running = true;

	while (running)
	{
		string in;
		cout << "Ready for next category" << endl;
		cout << "Please specifiy name or type exit to stop: ";
		cin >> in;

		if (in == "exit")
		{
			running = false;
			break;
		}

		bool takingPhotos = true;
		int i = 0;

		while (takingPhotos)
		{
			Mat image;
			image = cam.getImage();

			imshow("Live feed", image);
			
			int key = waitKey(100);
			switch (key)
			{
			case 32:
				imwrite("training/" + in + "_" + to_string(i) + ".bmp", image);
				cout << "Image taken #" << i << endl;
				i++;
				break;
			case 27:
				destroyWindow("Live feed");
				takingPhotos = false;
				running = false;
				break;
			case 110:
				destroyWindow("Live feed");
				takingPhotos = false;
			default:
				break;
			}
		}
	}
}

void Training::LoadTrainingSet()
{
	NeuralNetwork bpn;

	Mat ITset = Mat_<double>();
	Mat OTset = Mat_<double>();

	vector<string> files;
	string dir = "training/";
	
	read_directory(dir, files);

	cout << "Found " << files.size()-2 << " files in " << dir << endl;

	random_shuffle(files.begin(), files.end());

	for (string file : files)
	{
		string loc = dir + file;

		Mat image;
		image = imread(loc, CV_LOAD_IMAGE_COLOR);

		if (!image.data)
			continue;

		cout << "Loaded " << file << endl;

		string classname;
		class_name(file, classname);

		Mat gray_image, binaryImage;
		cvtColor(image, gray_image, CV_BGR2GRAY);
		threshold(gray_image, binaryImage, 200, 1, CV_THRESH_BINARY_INV);

		//Extract Features
		vector<Point> contour;
		FeatureExtractor::findContour(gray_image, contour);
		FeatureExtractor ftext = FeatureExtractor(contour);

		Mat descriptors;
		ftext.Extract(descriptors);
		Mat output;
		bpn.getClass(classname, output);

		ITset.push_back<double>(descriptors);
		OTset.push_back<double>(output);
	}

	destroyAllWindows();

	cout << "Finished loading all files" << endl;

	bpn.Train(ITset, OTset);

	cout << "Training complete" << endl;

	cin.ignore();
}



void Training::read_directory(const string& name, vector<string> &v)
{
	string pattern = name;
	pattern.append("\\*");
	WIN32_FIND_DATA data;
	HANDLE hFind;
	if ((hFind = FindFirstFile(pattern.c_str() , &data)) != INVALID_HANDLE_VALUE) {
		do {
			v.push_back(data.cFileName);
		} while (FindNextFile(hFind, &data) != 0);
		FindClose(hFind);
	}
}

void Training::class_name(const string& name, string& classname)
{
	string::size_type const p(name.find_last_of('_'));
	classname = name.substr(0, p);
}