#include "DescriptorManager.h"

//Default constructor, don't use this
DescriptorManager::DescriptorManager() {
	img_path = "data"; //default value
}

//Sets image path and defines the image descriptors.
DescriptorManager::DescriptorManager(const fs::path & image_path) {
	img_path = image_path;

	//Manual definitions.. SAD!
	descriptor_types[7] = new DCTMedian();
	descriptor_types[6] = new DCTPearson();
	descriptor_types[5] = new threshold();
	descriptor_types[4] = new pearson();
	descriptor_types[3] = new sobel();
	descriptor_types[2] = new median();
	descriptor_types[1] = new numColors();
	descriptor_types[0] = new averageColor();

	mutators[0] = new Grayscale();
	mutators[1] = new Flip();
	mutators[2] = new Flop();
	mutators[3] = new JPEG();
	mutators[4] = new LowQ();
	mutators[5] = new Hue();
	mutators[6] = new Brighten();
}

//No need to call destructor.. yet?
DescriptorManager::~DescriptorManager() {
	//Do some clean-up?
	//Realistically, when this gets destroyed the application is done anyway.
}

//Loads all base images, calculates their descriptors,
//calculates all image mutations and stores their descriptors
//on disk.
void DescriptorManager::Init() {
	//First load all images and store descriptors in cache and disk
	if (!fs::exists(img_path)) {
		cout << "Error. Data folder not found." << endl;
		return;
	}
	fs::ifstream ifs;
	fs::ofstream ofs;
	int cnt = std::count_if(
		fs::directory_iterator(img_path),
		fs::directory_iterator(),
		static_cast<bool(*)(const fs::path&)>(fs::is_regular_file));
	//foreach image in directory (non-recursive)
	int filecount = 0;
	for (fs::directory_iterator itr(img_path); itr != fs::directory_iterator(); ++itr) {
		cout << "\r" << "Progress: " << filecount++*100/cnt << "% - (" << filecount << "/" << cnt << ")" << flush;
		if (!itr->path().has_extension()) continue; //skip directories and such
		string image_name = itr->path().filename().generic_string();
		Image* image = new Image;
		image->read(itr->path().string());
		//foreach descriptor
		for (int i = 0; i < numDescriptors; i++) {
			//folder to store descriptors; create if not exists
			fs::path descriptorBase("meta" / descriptor_types[i]->ToPath());
			if (!fs::is_directory(descriptorBase))
				fs::create_directory(descriptorBase);
			//cout << "Starting with descriptor " << i << " on location " << descriptorBase.generic_string() << endl;

			fs::path featurePath(descriptorBase / itr->path().stem());
			bitset<64>* b1;
			if (fs::exists(featurePath)) { //feature on disk, load it in
				ifs.open(featurePath, fs::fstream::binary);
				unsigned long long n;
				ifs.read(reinterpret_cast<char*>(&n), sizeof(n));
				ifs.close();
				b1 = new bitset<64>(n);
			}
			else { //not on disk, calculate it
				b1 = new bitset<64>(descriptor_types[i]->feature(image));
				unsigned long long n = b1->to_ullong();
				ofs.open(featurePath, fs::fstream::binary);
				ofs.write(reinterpret_cast<const char*>(&n), sizeof(n));
				ofs.close();
			}
			features[i][image_name] = *b1;
			delete b1;
		}
		for (int i = 0; i < numMutations; i++) {
			//Check if we have somewhere to store the mutation
			fs::path mutationBase(img_path / mutators[i]->ToPath());
			if (!fs::is_directory(mutationBase))
				fs::create_directory(mutationBase);
			fs::path mutPath(mutationBase / itr->path().filename().generic_string());
			Image* mutated;
			if (!fs::exists(mutPath)) { //only mutate if nothing on disk
				mutated = mutators[i]->getMutated(image);
				mutated->write(mutPath.generic_string());
			}
			else {
				mutated = new Image;
				mutated->read(mutPath.generic_string());
			}

			//now calculate all the features for the mutated images as well
			//foreach descriptor
			for (int j = 0; j < numDescriptors; j++) {
				//folder to store descriptors; create if not exists
				fs::path descriptorBase("meta" / descriptor_types[j]->ToPath() / mutators[i]->ToPath());
				if (!fs::is_directory(descriptorBase))
					if (fs::create_directory(descriptorBase)) {}
				fs::path featurePath(descriptorBase / itr->path().stem());
				//cout << featurePath << endl;
				string image_name = itr->path().filename().generic_string();
				bitset<64>* b1 = NULL;
				if (!fs::exists(featurePath)) { //feature not yet calculated
												//descriptor_types[i]->distance(b1, b1);
												//descriptors[i][image_name] = b1;
					b1 = new bitset<64>(descriptor_types[j]->feature(mutated));
					unsigned long long n = b1->to_ullong();
					//cout << "Writing " << *b1 << "(" << n << ")" << " to " << featurePath << endl;
					ofs.open(featurePath, fs::fstream::binary); //in binary mode; only open file after calculations have been done
					ofs.write(reinterpret_cast<const char*>(&n), sizeof(n));
					ofs.close();
				}
				else {
					ifs.open(featurePath, fs::fstream::binary);
					unsigned long long n;
					ifs.read(reinterpret_cast<char*>(&n), sizeof(n));
					ifs.close();
					b1 = new bitset<64>(n);
				}

				mutatedFeatures[j][i][image_name] = *b1;
				if(b1) delete b1;
			}
			delete mutated;
		}
		delete image;
	}
	initialized = true;
}

//Sorts all images we have on distance from the query image
void DescriptorManager::rankImages(string image_name, int method) {
	bitset<64> target;
	if (features[method].count(image_name) && initialized) { //if we have image
		target = features[method][image_name];
	}
	else {
		return;
	}
	distances.clear(); //start fresh
	for (map<string, bitset<64>>::iterator it = features[method].begin(); it != features[method].end(); ++it) {
		unsigned long long d = descriptor_types[method]->distance(target, it->second);
		pair<unsigned long long, string> p(d, it->first);
		distances.insert(p); //pairs are compared by their first element..
	}
}

//inserts the descriptors of the mutated versions of the image.
//Assumes the distances array is filled with the normal distances
void DescriptorManager::rankMutations(string image_name, int method) {
	bitset<64> target;
	if (features[method].count(image_name) && initialized) { //if we have image
		target = features[method][image_name];
	}
	else {
		return;
	}
	//cout << "Finding hits for " << image_name << " which has feature-number " << target.to_ullong() << endl;
	multiset<pair<unsigned long long, string> > local_distance = distances; //get our local copy of the distances.
	for (int i = 0; i < numMutations; i++) {
		bitset<64> b1 = mutatedFeatures[method][i][image_name];
		unsigned long long d = descriptor_types[method]->distance(target, b1);
		pair<unsigned long long, string> p(d, mutators[i]->ToString());
		local_distance.insert(p);
		upper_bound
	}

	//cout << "..." << endl;
	for (int i = 0; i < numMutations; i++) {
		int j = 0;
		for (std::multiset<pair<unsigned long long, string>>::iterator it = local_distance.begin(); it != local_distance.end(); ++it) {
			//cout << "hey " << endl;
			if (it->second == mutators[i]->ToString()) {
				cout << j << "th position: ";
				cout << "(" << it->first << ", " << it->second << ")" << endl;
			}
			j++;
		}
	}
}

void DescriptorManager::runExperiments() {
	for (int i = 0; i < numDescriptors; i++) {
		for (map<string, bitset<64>>::iterator it = features[i].begin(); it != features[i].end(); ++it) {
			cout << "=== " << descriptor_types[i]->ToString() << ": " << it->first << " ===" << endl;
			rankImages(it->first, i);
			rankMutations(it->first, i);
		}
	}
}

//outputs the gathered data to iostream
//can be called directly or is called by its overload
//assumes a < numDescriptors
void DescriptorManager::outputMap(int a) {
	cout << "Outputting map " << descriptor_types[a]->ToString() << endl;
	map<string, bitset<64>> a_map = features[a];
	if (a_map.begin() == a_map.end()) {
		cout << "------------------------" << endl << "No data found" << endl << endl;
	}

	for (map<string, bitset<64>>::iterator it = a_map.begin(); it != a_map.end(); ++it) {
		cout << it->first << " => " << it->second << " (" << it->second.to_ullong() << ")" << endl;
	}
} 

//calls its buddy function to output all data to iostream
void DescriptorManager::outputMap() {
	for (int i = 0; i < numDescriptors; i++) {
		outputMap(i);
	}
}

//Outputs all generated features ordered on distance from the query image.
void DescriptorManager::outputRank(string imagename, int method) {
	for (std::multiset<pair<unsigned long long, string>>::iterator it = distances.begin(); it != distances.end(); ++it) {
			cout << "(" << it->first << ", " << it->second << ")" << endl;
	}
}