#include "stereo_rectification_frames.hpp"

#include <opencv2/calib3d.hpp>
#include <opencv2/imgproc.hpp>

const char *StereoRectificationFrames::initDescription() {
	return ("Perform stereo rectification on frames.");
}

void StereoRectificationFrames::initInputs(dv::InputDefinitionList &in) {
	in.addFrameInput("camera0");
	in.addFrameInput("camera1");
}

void StereoRectificationFrames::initOutputs(dv::OutputDefinitionList &out) {
	out.addFrameOutput("rectifiedCamera0");
	out.addFrameOutput("rectifiedCamera1");
}

StereoRectificationFrames::StereoRectificationFrames() {
	for (size_t i = 0; i < NUM_CAMERAS; i++) {
		// Populate frame output info node, keep same as input info node.
		outputs.getFrameOutput("rectifiedCamera" + std::to_string(i))
			.setup(inputs.getFrameInput("camera" + std::to_string(i)));
	}

	// Load calibration from file.
	if (!loadStereoRectificationMatrices(config.getString("calibrationFile"))) {
		throw std::runtime_error("Could not load calibration data.");
	}
}

void StereoRectificationFrames::run() {
	for (size_t i = 0; i < NUM_CAMERAS; i++) {
		const auto inputFrame = inputs.getFrameInput("camera" + std::to_string(i)).frame();

		if (inputFrame) {
			auto outputFrame = outputs.getFrameOutput("rectifiedCamera" + std::to_string(i)).frame();
			stereoRectifyFrame(inputFrame, outputFrame, i);
		}
	}
}

void StereoRectificationFrames::initStereoRectification(
	const CameraInfo loaded[NUM_CAMERAS], const cv::Mat R, const cv::Mat T, const double alpha) {
	cv::Mat RN[NUM_CAMERAS], CN[NUM_CAMERAS], Q;

	cv::stereoRectify(loaded[0].cameraMatrix, loaded[0].distCoeffs, loaded[1].cameraMatrix, loaded[1].distCoeffs,
		size[0], R, T, RN[0], RN[1], CN[0], CN[1], Q, cv::CALIB_ZERO_DISPARITY, alpha, size[0], nullptr, nullptr);

	for (size_t i = 0; i < NUM_CAMERAS; i++) {
		cv::initUndistortRectifyMap(
			loaded[i].cameraMatrix, loaded[i].distCoeffs, RN[i], CN[i], size[i], CV_16SC2, remap1[i], remap2[i]);
	}
}

void StereoRectificationFrames::stereoRectifyFrame(
	const dv::InputDataWrapper<dv::Frame> &in, dv::OutputDataWrapper<dv::Frame> &out, size_t cam) {
	// Setup output frame. Same size.
	out.setPosition(in.position());
	out.setTimestamp(in.timestamp());
	out.setExposure(in.exposure());
	out.setSource(in.source());

	// Get input OpenCV Mat. Lifetime is properly managed.
	auto input = in.getMatPointer();

	// Get output OpenCV Mat. Memory must have been allocated already.
	cv::Mat output;

	cv::remap(*input, output, remap1[cam], remap2[cam], cv::INTER_CUBIC, cv::BORDER_CONSTANT);

	out.setMat(std::move(output));
	out.commit();
}

registerModuleClass(StereoRectificationFrames)
