#include "BoF.h"

const int BoF::DEFAULT_DESCRIPTOR_DIMENSION = 128;
const int BoF::DEFAULT_GRID = 20;


BoF::BoF()
{
	descriptorDimension = DEFAULT_DESCRIPTOR_DIMENSION;
	grid = DEFAULT_GRID;

	gridScale.push_back(30);
	gridScale.push_back(23);
	gridScale.push_back(18);
	gridScale.push_back(15);
	gridScale.push_back(12);
	gridScale.push_back(0);
}


BoF::~BoF() {}



// 特徴検出
// pathToImage : 画像のあるフォルダ (フォルダの中身はimage(i).jpgとなっている必要有)
// start : 開始の番号
// end : 終了の番号
// featureType : 特徴量の種類
void BoF::detectFeature(const std::string pathToFolder, const int start, const int end, const int featureType)
{
	cv::Mat img;
	// 各画像ごとに特徴点の検出
	for (int i = start; i <= end; ++i) {
		// 画像の読み込み
		std::stringstream str;
		str << pathToFolder << "image(" << i << ").jpg";
		img = cv::imread(str.str());
		images.push_back(img);
		
		if(!img.empty()) {
			std::cout << "Load Image : " << str.str() << "\n";
		}
		// grid sampring
		if (featureType == SIFT || featureType == SURF) {
			std::vector<cv::KeyPoint> intoKeyPoints; // 現在の画像における特徴点の集合

			int startX = (int)((img.cols & grid) + grid) / 2 - 1;
			int startY = (int)((img.rows & grid) + grid) / 2 - 1;
			
			cv::KeyPoint intoIntoKeyPoint;
						
			for (int s = 0; s < gridScale.size(); ++s) {
				for (int y = startY; y < img.rows; y += grid) {
					for (int x = startX; x < img.cols; x += grid) {			
						intoIntoKeyPoint.pt.x = (float)x;
						intoIntoKeyPoint.pt.y = (float)y;
						intoIntoKeyPoint.size = (float)gridScale[s];

						intoKeyPoints.push_back(intoIntoKeyPoint);
					}
				}
			}

			keyPoints.push_back(intoKeyPoints);

		}
	}

	return;
}


// 特徴量記述子の計算
// featureType : 特徴量の種類
void BoF::computeDescriptors(const int featureType)
{
	if (featureType == SIFT) {
		for (int i = 0; i < images.size(); ++i) {
			cv::Mat img = images[i];

			cv::SiftDescriptorExtractor extractor;
			descriptorDimension = extractor.descriptorSize();

			cv::Mat intoDescriptors;
			extractor.compute(img, keyPoints[i], intoDescriptors);
			descriptors.push_back(intoDescriptors);

			std::cout << "Compute Descriptors [" << i << "]\n";
		}
	}
	else if (featureType == SURF) {
		for (int i = 0; i < images.size(); ++i) {
			cv::Mat img = images[i];

			cv::SurfDescriptorExtractor extractor(500, 4, 2, true);
			descriptorDimension = extractor.descriptorSize();

			cv::Mat intoDescriptors;
			extractor.compute(img, keyPoints[i], intoDescriptors);
			descriptors.push_back(intoDescriptors);

			std::cout << "Compute Descriptors [" << i << "]\n";
		}
	}

	return;
}

// KMeansクラスタリングを用いたコードブックの生成
// size : コードブックのサイズ
// tc : 終了条件
// retries : 試行回数
// flags : K-meansクラスタリングにおけるフラグ 
void BoF::createCodebook(const int size, const cv::TermCriteria tc, const int retries, const int flags)
{
	codebookSize = size;

	BoWTrainer = new cv::BOWKMeansTrainer(size, tc, retries, flags);

	for (int i = 0; i < descriptors.size(); ++i) {
		BoWTrainer->add(descriptors[i]);
	}

	std::cout << "start K-means clustering\n";
	codebook = BoWTrainer->cluster();
	std::cout << "end K-means clustering\n";

	delete BoWTrainer;

	return;
}


// 画像記述子の生成
// featureType : 特徴量の種類
void BoF::createImageDescriptorsForImages(const int featureType)
{
	if (featureType == SIFT) {
		extractor = new cv::SiftDescriptorExtractor;
		matcher = cv::DescriptorMatcher::create("FlannBased");

		bowDE = new cv::BOWImgDescriptorExtractor(extractor, matcher);
		bowDE->setVocabulary(codebook);

		for (int i = 0; i < images.size(); ++i) {
			cv::Mat intoImageDescriptors;
			bowDE->compute(images[i], keyPoints[i], intoImageDescriptors);
			imageDescriptors.push_back(intoImageDescriptors);

			std::cout << "create ImageDescriptor [" << i << "]\n";
		}
	}
	else if (featureType == SURF) {
		extractor = new cv::SurfDescriptorExtractor(500, 4, 2, true);
		matcher = cv::DescriptorMatcher::create("FlannBased");

		bowDE = new cv::BOWImgDescriptorExtractor(extractor, matcher);
		bowDE->setVocabulary(codebook);

		for (int i = 0; i < images.size(); ++i) {
			cv::Mat intoImageDescriptors;
			bowDE->compute(images[i], keyPoints[i], intoImageDescriptors);
			imageDescriptors.push_back(intoImageDescriptors);

			std::cout << "create ImageDescriptor [" << i << "]\n";
		}
	}

	return;
}


// コードブックの保存
// pathToXML : 保存するファイル名
void BoF::saveCodebook(const std::string pathToXML)
{
	cv::FileStorage cvfs(pathToXML, CV_STORAGE_WRITE);
	cv::write(cvfs, "codebook", codebook);

	return;
}


// コードブックの読み込み
// pathToXML : 読み込むファイル名
void BoF::loadCodebook(const std::string pathToXML)
{
	cv::FileStorage cvfs(pathToXML, CV_STORAGE_READ);
	cv::FileNode node(cvfs.fs, NULL);
	cv::read(node["codebook"], codebook);

	codebookSize = codebook.rows;

	std::cout << "codebook : " << codebook.size() << "\n";

	return;
}



// 類似度が最も高い画像を求める
// img : 認識対象の画像
// similarImg : 類似度が最も高い学習画像
// mostSimilarImageIndex : 類似度が最も高い学習画像の番号
// dist : ヒストグラムの距離
// featureType : 特徴量の種類
void BoF::computeMostSimilarImage(const cv::Mat img, cv::Mat& similarImg, int& mostSimilarImageIndex, double& dist, const int featureType)
{
	cv::Mat imgDescriptor;
	createImageDescriptor(img, imgDescriptor, featureType);

	computeSimilarityToImageDescriptors(imgDescriptor, mostSimilarImageIndex, dist);

	similarImg = images[mostSimilarImageIndex].clone();

	return;
}


// 特徴検出(画像1枚分)
// img : 入力画像
// imgKeyPoints : 検出された特徴点の集合
// featureType : 特徴量の種類
void BoF::detectFeatureOfImage(const cv::Mat img, std::vector<cv::KeyPoint>& imgKeyPoints, const int featureType)
{
	if (featureType == SIFT || featureType == SURF) {
		//Grid Sampling
		int startX = (int)((img.cols % grid) + grid ) / 2 - 1;
		int startY = (int)((img.rows % grid) + grid ) / 2 - 1;

		cv::KeyPoint into_keyPoints;
		for (int s = 0; s < gridScale.size(); s++) {
			for (int y = startY; y < img.rows; y += grid) {
				for (int x = startX; x < img.cols; x += grid) {
					into_keyPoints.pt.x = (float)x;
					into_keyPoints.pt.y = (float)y;
					into_keyPoints.size = (float)gridScale[s];

					imgKeyPoints.push_back(into_keyPoints);
				}
			}
		}
	}

	return;
}


// 画像記述子の生成(画像1枚分)
// img : 入力画像
// imgDescriptor : 生成された画像記述子
// featureType : 特徴点の種類
void BoF::createImageDescriptor(const cv::Mat img, cv::Mat& imgDescriptor, const int featureType)
{
	if (featureType == SIFT) {
		extractor = new cv::SiftDescriptorExtractor;
		matcher = cv::DescriptorMatcher::create("FlannBased");

		bowDE = new cv::BOWImgDescriptorExtractor(extractor, matcher);
		bowDE->setVocabulary(codebook);

		std::vector<cv::KeyPoint> imgKeyPoints;
		detectFeatureOfImage(img, imgKeyPoints, featureType);
		bowDE->compute(img, imgKeyPoints, imgDescriptor);
	}
	else if (featureType == SURF) {
		extractor = new cv::SurfDescriptorExtractor(500, 4, 2, true);
		matcher = cv::DescriptorMatcher::create("FlannBased");

		bowDE = new cv::BOWImgDescriptorExtractor(extractor, matcher);
		bowDE->setVocabulary(codebook);

		std::vector<cv::KeyPoint> imgKeyPoints;
		detectFeatureOfImage(img, imgKeyPoints, featureType);
		bowDE->compute(img, imgKeyPoints, imgDescriptor);
	}

	return;
}



// 類似度の計算
// imgDescriptorをimgaeDescriptorsと比較し，もっとも類似度が高い画像の番号と，ヒストグラムの距離を返す
// imgDescriptor : 識別させたい画像の記述子
// mostSimilarImageIndex : 最も類似度が高い学習画像の番号
// dist : ヒストグラムの距離
void BoF::computeSimilarityToImageDescriptors(const cv::Mat imgDescriptor, int& mostSimilarImageIndex, double& dist)
{
	double dMax = -2.0;
	int index = -1;
	cv::Mat tmpDescriptor = imgDescriptor.clone();

	for (int i = 0; i < imageDescriptors.size(); ++i) {
		imageDescriptors[i].convertTo(imageDescriptors[i], CV_32FC1);
		tmpDescriptor.convertTo(tmpDescriptor, CV_32FC1);

		double d = cv::compareHist(imageDescriptors[i], tmpDescriptor, CV_COMP_CORREL);

		if (d > dMax) {
			dMax = d;
			index = i;
		}
	}

	dist = dMax;
	mostSimilarImageIndex = index;

	return;
}

