#include <iostream>

#include "blender.h"

// ---------- Constructor ---------- //
Blender::Blender(double &work_scale,
	float &warped_image_scale,
	cv::Ptr<cv::WarperCreator> &warper_creator,
	int no_of_videos,
	std::vector<cv::detail::CameraParams> &cameraParameters,
	std::vector<cv::Size> &full_image_sizes,
	std::vector<cv::Point> &corners,
	std::vector<cv::Size> &sizes) {

	compose_scale = 1;

	// Compute relative scales
	compose_work_aspect = compose_scale / work_scale;

	// Update warped image scale
	warped_image_scale *= static_cast<float>(compose_work_aspect);

	// Initialise warper
	warper = warper_creator->create(warped_image_scale);

	// Update corners and sizes
	for (int i = 0; i < no_of_videos; ++i) {
		cameraParameters[i].focal *= compose_work_aspect;
		cameraParameters[i].ppx *= compose_work_aspect;
		cameraParameters[i].ppy *= compose_work_aspect;

		cv::Size size = full_image_sizes[i];
		cv::Mat K;
		cameraParameters[i].K().convertTo(K, CV_32F);
		cv::Rect roi = warper->warpRoi(size, K, cameraParameters[i].R);
		corners[i] = roi.tl();
		sizes[i] = roi.size();
	}
}

int Blender::blend(std::vector<cv::Mat> &video_frames,
	std::vector<cv::detail::CameraParams> &cameraParameters,
	std::vector<cv::Point> &corners,
	std::vector<cv::Size> &sizes,
	cv::Ptr<cv::detail::ExposureCompensator> &compensator,
	std::vector<cv::Mat> &masks_warped,
	int blend_type,
	bool try_gpu,
	float blend_strength,
	cv::Mat &result,
	cv::Mat &result_mask) {

	for (size_t i = 0; i < video_frames.size(); ++i) {
		// Read image
		cv::Mat image = video_frames.at(i);
		cv::Size image_size = image.size();

		cv::Mat K;
		cameraParameters[i].K().convertTo(K, CV_32F);

		// Warp the current image
		warper->warp(image, K, cameraParameters[i].R, cv::INTER_LINEAR, cv::BORDER_REFLECT, image_warped);

		// Warp the current image mask
		mask.create(image_size, CV_8U);
		mask.setTo(cv::Scalar::all(255));
		warper->warp(mask, K, cameraParameters[i].R, cv::INTER_LINEAR, cv::BORDER_CONSTANT, mask_warped);

		// Compensate exposure
		compensator->apply(i, corners[i], image_warped, mask_warped);

		image_warped.convertTo(image_warped_temp, CV_16S);
		image_warped.release();
		image.release();
		mask.release();

		cv::dilate(masks_warped[i], dilated_mask, cv::Mat());
		cv::resize(dilated_mask, seam_mask, mask_warped.size());
		mask_warped = seam_mask & mask_warped;

		// Initialise and prepare blender
		if (blender.empty()) {
			blender = cv::detail::Blender::createDefault(blend_type, try_gpu);
			cv::Size dst_size = cv::detail::resultRoi(corners, sizes).size();
			float blend_width = sqrt(static_cast<float>(dst_size.area())) * blend_strength / 100.f;

			if (blend_width < 1.f) {
				blender = cv::detail::Blender::createDefault(cv::detail::Blender::NO, try_gpu);
			}
			else if (blend_type == cv::detail::Blender::MULTI_BAND) {
				cv::detail::MultiBandBlender* mb = dynamic_cast<cv::detail::MultiBandBlender*>(static_cast<cv::detail::Blender*>(blender));
				mb->setNumBands(static_cast<int>(ceil(log(blend_width) / log(2.)) - 1.));
			}
			else if (blend_type == cv::detail::Blender::FEATHER) {
				cv::detail::FeatherBlender* fb = dynamic_cast<cv::detail::FeatherBlender*>(static_cast<cv::detail::Blender*>(blender));
				fb->setSharpness(1.f / blend_width);
			}

			blender->prepare(corners, sizes);
		}

		// Feed blender all the images
		blender->feed(image_warped_temp, mask_warped, corners[i]);
	}

	// Blend all images into a panoramic image
	blender->blend(result, result_mask);
	result.convertTo(result, CV_8U); // Video needs images with depth of 8

	// Destroy blender
	blender.release();

	return 0;
}
