#include "stackercpp.h"

Stacker::Stacker(bool verbose, struct Imagedata imagedata)
{
	Stacker::verbose = verbose;
	Stacker::export_width = 0;
	Stacker::export_height = 0;
	Stacker::layers = 0;
	Stacker::trimratio = 0;

	Stacker::imagedata = imagedata;

	Stacker::ldb = lf_db_new();
	lf_db_load(Stacker::ldb);

	cv::setNumThreads(0);
	Stacker::stopwatch = clock();

	/*
	Stacker::imagedata.dist_a = 0.025898;
	Stacker::imagedata.dist_b = -0.092305;
	Stacker::imagedata.dist_c = 0.105513;

	Stacker::imagedata.vignette_k1 = -2.4268;
	Stacker::imagedata.vignette_k2 = 3.0906;
	Stacker::imagedata.vignette_k3 = -1.4546;
	 */
}

void
Stacker::add_frame(unsigned char *data, int width, int height)
{
	stopwatch_start();
	Mat warp_matrix = Mat::eye(3, 3, CV_32F);
	Mat mat = Mat(height, width, CV_8UC3, data);
	Mat grayscale = Mat(height, width, CV_8UC1);
	cv::cvtColor(mat, grayscale, cv::COLOR_BGR2GRAY);
	stopwatch_mark("grayscale input");

	stopwatch_start();
	Scalar mean, stddev;
	meanStdDev(grayscale, mean, stddev);
	printf("mean: %f, dev: %f\n", mean[0], stddev[0]);
	if (mean[0] < 10) {
		return;
	}
	stopwatch_mark("filter");

	int number_of_iterations = 5;
	double termination_eps = 1e-10;
	TermCriteria criteria(TermCriteria::COUNT + TermCriteria::EPS, number_of_iterations, termination_eps);
	if (layers == 0) {
		// First image in the stack is used as the reference to align the next frames to
		Stacker::reference = grayscale;

		// Create black image to accumulate the average
		Stacker::stacked = Mat(height, width, CV_64FC3);
		Stacker::stacked.setTo(Scalar(0, 0, 0, 0));

		// Add first frame to the stack
		Stacker::stacked += mat;
		Stacker::layers += 1;
	} else {
		// All frames after the initial one are stacked
		Mat warped = Mat(Stacker::stacked.rows, Stacker::stacked.cols, CV_32FC1);

		stopwatch_start();
		cv::findTransformECC(grayscale, Stacker::reference, warp_matrix, MOTION_HOMOGRAPHY, criteria);
		stopwatch_mark("find alignment");

		stopwatch_start();
		warpPerspective(mat, warped, warp_matrix, warped.size(), INTER_LINEAR);
		stopwatch_mark("warp image");

		// Check how much the image should be cropped to hide the warped edges
		float current_trimratio = cv::videostab::estimateOptimalTrimRatio(warp_matrix, mat.size());
		Stacker::trimratio = std::max(Stacker::trimratio, current_trimratio);

		// Add the warped image to the stack
		Stacker::stacked += warped;
		Stacker::layers += 1;
	}
}

Mat
Stacker::postprocess_mat(Mat
input)
{
	stopwatch_start();

	int h_crop = (int) ((float) input.cols * Stacker::trimratio);
	int v_crop = (int) ((float) input.rows * Stacker::trimratio);
	Mat cropped;
	input(Rect(h_crop, v_crop, input.cols - h_crop - h_crop, input.rows - v_crop - v_crop)
	).
		copyTo(input);
	stopwatch_mark("trim");

	stopwatch_start();

	Mat blur;
	GaussianBlur(input, blur, Size(0, 0),
		1.8);
	std::vector<cv::Mat> rgb_planes(3);
	cv::split(blur, rgb_planes
	);
	double min_r, max_r, min_g, max_g, min_b, max_b;
	minMaxIdx(rgb_planes[0], &min_b, &max_b
	);
	minMaxIdx(rgb_planes[1], &min_g, &max_g
	);
	minMaxIdx(rgb_planes[2], &min_r, &max_r
	);
	input = input - Scalar(min_b + 5, min_g + 5, min_r + 5);
	double scale_r, scale_g, scale_b;
	scale_r = 255 / (max_r - min_r + 5);
	scale_g = (255 / (max_g - min_g + 5));
	scale_b = 255 / (max_b - min_b + 5);
	multiply(input, Scalar(scale_b, scale_g, scale_r), input
	);
	stopwatch_mark("levels");

	stopwatch_start();

	Mat sharpened;
	GaussianBlur(input, sharpened, Size(0, 0),
		1.7);
	addWeighted(input,
		2.5, sharpened, -1.5, 0, sharpened);
	stopwatch_mark("sharpen");

/* Disabled CLAHE local contrast, it's a bit to overpronounced and doesn't
 * seem really necessary at this point

stopwatch_start();
Mat lab;
cvtColor(sharpened, lab, COLOR_BGR2Lab);
std::vector<cv::Mat> lab_planes(3);
cv::split(lab, lab_planes);
stopwatch_mark("to Lab");

stopwatch_start();
cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE();
clahe->setClipLimit(1);
clahe->setTilesGridSize(Size(8, 8));
cv::Mat dst;
clahe->apply(lab_planes[0], dst);
dst.copyTo(lab_planes[0]);
stopwatch_mark("clahe");

stopwatch_start();
Mat result;
cv::merge(lab_planes, lab);
cvtColor(lab, result, COLOR_Lab2BGR);
stopwatch_mark("to RGB");
*/

	return
		sharpened;
}

char *
Stacker::get_result()
{
	// Complete the averaging and go back to an 8-bit image
	stopwatch_start();
	Stacker::stacked.convertTo(Stacker::stacked, CV_8U, 1. / Stacker::layers);
	stopwatch_mark("average");

	// Run the final-frame postprocessing
	Mat result = postprocess_mat(Stacker::stacked);
	Stacker::export_width = result.cols;
	Stacker::export_height = result.rows;

	// Convert mat to bytes
	size_t size = result.total() * result.elemSize();
	char *data = (char *) malloc(size);
	std::memcpy(data, result.data, size * sizeof(char));
	return data;
}

char *
Stacker::postprocess(unsigned char *data, int width, int height)
{
	// Run lensfun
	lfCamera *cam = nullptr;
	const lfCamera **cameras = ldb->FindCamerasExt(nullptr, imagedata.model);
	if (cameras) {
		cam = (lfCamera *) cameras[0];
		fprintf(stderr, "Using camera %s\n", cam->Model);
	} else {
		fprintf(stderr, "No camera found in LensFun database\n");
	}
	lfLens *lens = nullptr;
	const lfLens **lenses = ldb->FindLenses(cam, nullptr, imagedata.model);
	if (lenses) {
		lens = (lfLens *) lenses[0];
		fprintf(stderr, "Using lens %s\n", lens->Model);
	} else {
		fprintf(stderr, "No lens found in LensFun database\n");
	}
	float cropfactor = (float) imagedata.focal_length_35mm / imagedata.focal_length;

	if (!lens && !cam) {
		cam = new lfCamera;
		cam->SetModel(imagedata.model);
		lens = new lfLens;
		lens->Type = LF_RECTILINEAR;
		lens->GuessParameters();
		lens->AddMount("fixed");
		lens->SetModel(imagedata.model);
		lens->MinFocal = imagedata.focal_length;
		lens->MaxFocal = imagedata.focal_length;
		lens->MinAperture = imagedata.fnumber;
		lens->MaxAperture = imagedata.fnumber;
		lens->AspectRatio = 1.5;
		lens->CropFactor = cropfactor;

		if (imagedata.dist_a != 0.0f) {
			auto *distortion = new lfLensCalibDistortion();
			distortion->Model = lfDistortionModel::LF_DIST_MODEL_PTLENS;
			distortion->Focal = imagedata.focal_length;
			distortion->Terms[0] = imagedata.dist_a;
			distortion->Terms[1] = imagedata.dist_b;
			distortion->Terms[2] = imagedata.dist_c;
			lens->AddCalibDistortion(distortion);
		}

		if (imagedata.vignette_k1 != 0.0f) {
			auto *vignette = new lfLensCalibVignetting();
			vignette->Model = LF_VIGNETTING_MODEL_PA;
			vignette->Focal = imagedata.focal_length;
			vignette->Aperture = imagedata.fnumber;
			vignette->Distance = 10.0f;
			vignette->Terms[0] = imagedata.vignette_k1;
			vignette->Terms[1] = imagedata.vignette_k2;
			vignette->Terms[2] = imagedata.vignette_k3;
			lens->AddCalibVignetting(vignette);
		}

		if (!lens->Check()) {
			fprintf(stderr, "Lens check failed\n");
		}
	}

	auto *mod = new lfModifier(lens, cropfactor, width, height);
	mod->Initialize(lens, LF_PF_U8, imagedata.focal_length, imagedata.fnumber, 10.0f, 1.0f, lfLensType::LF_RECTILINEAR,
		LF_MODIFY_ALL, false);

	auto *pos = new float[width * height * 2];
	bool do_remap = mod->ApplyGeometryDistortion(0.0f, 0.0f, width, height, pos);

	// Convert bytes to mat
	Mat mat = Mat(height, width, CV_8UC3, data);
	stopwatch_start();
	mod->ApplyColorModification(mat.data, 0, 0, width, height, LF_CR_3(RED, GREEN, BLUE), width * 3);
	stopwatch_mark("vignette");
	if (do_remap) {
		stopwatch_start();
		Mat dst = Mat(mat.size(), mat.type());
		Mat map_x(mat.size(), CV_32FC1);
		Mat map_y(mat.size(), CV_32FC1);
		for (int x = 0; x < mat.rows; x++) {
			for (int y = 0; y < mat.cols; y++) {
				size_t offset = ((x * mat.cols) + y) * 2;
				map_x.at<float>(x, y) = pos[offset];
				map_y.at<float>(x, y) = pos[offset + 1];
			}
		}
		remap(mat, dst, map_x, map_y, INTER_LANCZOS4, BORDER_CONSTANT, Scalar(0, 0, 0));
		mat = dst;
		stopwatch_mark("distortion");
	}

	// Run the final-frame postprocessing
	Mat result = postprocess_mat(mat);
	Stacker::export_width = result.cols;
	Stacker::export_height = result.rows;

	// Convert mat to bytes
	size_t size = result.total() * result.elemSize();
	char *outdata = (char *) malloc(size);
	std::memcpy(outdata, result.data, size * sizeof(char));
	return outdata;
}

void
Stacker::stopwatch_start()
{
	Stacker::stopwatch = clock();
}

void
Stacker::stopwatch_mark(const char *name)
{
	if (Stacker::verbose) {
		printf("[%.1fms] %s\n", float(clock() - Stacker::stopwatch) / CLOCKS_PER_SEC * 1000, name);
	}
}

int
Stacker::get_width()
{
	return Stacker::export_width;
}

int
Stacker::get_height()
{
	return Stacker::export_height;
}