#include "bundle_adjustment.h"

// ---------- Constructor ---------- //
BundleAdjustment::BundleAdjustment(std::string ba_cost_func,
	float confidence_threshold,
	std::string ba_refine_mask,
	bool do_wave_correct,
	cv::detail::WaveCorrectKind wave_correct) :
	do_wave_correct(do_wave_correct), wave_correct(wave_correct) {
	if (ba_cost_func == "reproj") {
		adjuster = new cv::detail::BundleAdjusterReproj();
	}
	else if (ba_cost_func == "ray") {
		adjuster = new cv::detail::BundleAdjusterRay();
	}

	adjuster->setConfThresh(confidence_threshold);

	refine_mask = cv::Mat::zeros(3, 3, CV_8U);
	if (ba_refine_mask[0] == 'x') refine_mask(0, 0) = 1;
	if (ba_refine_mask[1] == 'x') refine_mask(0, 1) = 1;
	if (ba_refine_mask[2] == 'x') refine_mask(0, 2) = 1;
	if (ba_refine_mask[3] == 'x') refine_mask(1, 1) = 1;
	if (ba_refine_mask[4] == 'x') refine_mask(1, 2) = 1;
	adjuster->setRefinementMask(refine_mask);
}

int BundleAdjustment::bundle_adjustment(std::vector<cv::detail::ImageFeatures> &features,
	std::vector<cv::detail::MatchesInfo> &pairwise_matches,
	std::vector<cv::detail::CameraParams> &cameraParameters,
	float &warped_image_scale) {
	(*adjuster)(features, pairwise_matches, cameraParameters);

	// Find median focal length
	std::vector<double> focals;
	for (size_t i = 0; i < cameraParameters.size(); ++i) {
		focals.push_back(cameraParameters[i].focal);
	}

	std::sort(focals.begin(), focals.end());

	if (focals.size() % 2 == 1) {
		warped_image_scale = static_cast<float>(focals[focals.size() / 2]);
	}
	else {
		warped_image_scale = static_cast<float>(focals[focals.size() / 2 - 1] + focals[focals.size() / 2]) * 0.5f;
	}

	if (do_wave_correct) {
		std::vector<cv::Mat> rmats;
		for (size_t i = 0; i < cameraParameters.size(); ++i) {
			rmats.push_back(cameraParameters[i].R);
		}
		cv::detail::waveCorrect(rmats, wave_correct);
		for (size_t i = 0; i < cameraParameters.size(); ++i) {
			cameraParameters[i].R = rmats[i];
		}
	}

	return 0;
}
