#include <opencv2/opencv.hpp>

#include "FitMapParam.h"

// 使用最小二乘法进行多项式拟合
FitParam FittingByPoints(const double* x, const double* y, int count) {
	// 构造矩阵方程 AW=B
	cv::Mat A(count, FittingOrder + 1, CV_64FC1); // 设计矩阵
	cv::Mat B(count, 1, CV_64FC1); // 观测值矩阵

	// 填充矩阵A和B
	for (int i = 0; i < count; i++) {
		A.ptr<double>(i)[0] = 1; // 常数项设为1
		for (int n = 1; n <= FittingOrder; n++)
			A.ptr<double>(i)[n] = pow(x[i], n); // x的n次方

		B.ptr<double>(i)[0] = y[i]; // 对应的y值
	}

	// 最小二乘解：W = (A^T*A)^-1 * A^T * B
	cv::Mat W = (A.t() * A).inv() * A.t() * B;

	// 将结果存入结构体
	FitParam result;
	for (int n = 0; n <= FittingOrder; n++)
		result.s[n] = W.ptr<double>(0)[n];

	return result;
}

// 对输入矩阵的每一行(或间隔行)进行多项式拟合，返回拟合参数矩阵
cv::Mat GetFitParamMap(const cv::Mat& map, int row_step) {
	// 检查行数是否能被row_step整除
	if (0 != (map.rows % row_step))
		return cv::Mat(); // 不能整除返回空矩阵

	// 计算结果矩阵的行数
	const int rows = 1 == row_step ? map.rows : (map.rows / row_step) + 1;
	cv::Mat fitMap = cv::Mat::zeros(rows, (FittingOrder + 1), map.type());

	for (int _r = 0; _r < map.rows; _r++) {
		// 计算当前行在结果矩阵中的行号
		int fitRow = _r / row_step;
		if ((1 != row_step) && (_r == (map.rows - 1)))
			fitRow++; // 处理最后一行特殊情况
		else if (0 != (_r % row_step))
			continue; // 跳过不需要处理的行

		// 准备拟合数据
		double* x = new double[map.cols];
		double* y = new double[map.cols];
		for (int _c = 0; _c < map.cols; _c++) {
			x[_c] = _c; // x坐标为列号
			y[_c] = map.ptr<float>(_r)[_c]; // y坐标为像素值
		}

		// 进行多项式拟合
		FitParam param = FittingByPoints(x, y, map.cols);
		delete[] x;
		delete[] y;

		// 存储拟合结果
		for (int _c = 0; _c <= FittingOrder; _c++)
			fitMap.ptr<float>(fitRow)[_c] = param.s[_c];
	}

	return fitMap;
}

// 从拟合参数矩阵重建原始图像
cv::Mat GetMapFromFitMap(const cv::Mat& map, const cv::Size size, int row_step) {
	// 检查输入参数是否匹配
	int rows_needed = size.height / row_step;
	if ((size.height % row_step) > 0)
		rows_needed += 1;
	if (map.rows != rows_needed)
		return cv::Mat();

	// 计算稀疏的Map(只计算关键行)
	cv::Mat fitMap = cv::Mat::zeros(size.height, size.width, map.type());
	for (int _r = 0; _r < map.rows; _r++) {
		// 计算实际行号(处理边界情况)
		const int real_row = _r * row_step >= size.height ?
			size.height - 1 : _r * row_step;

		// 使用多项式参数重建该行像素值
		for (int _c = 0; _c < size.width; _c++) {
			fitMap.ptr<float>(real_row)[_c] = map.ptr<float>(_r)[0]; // 常数项
			for (int n = FittingOrder; n > 0; n--)
				fitMap.ptr<float>(real_row)[_c] += map.ptr<float>(_r)[n] * pow(_c, n); // 加上各次项
		}
	}

	// 对非关键行进行线性插值
	if (row_step > 1) {
		for (int _r = 0; _r < size.height; _r++) {
			if (0 == (_r % row_step)) // 跳过关键行
				continue;

			// 找到上下两个关键行
			const int row_l = int(_r / row_step) * row_step;
			int row_h = int(_r / row_step + 1) * row_step;
			row_h = row_h >= map.rows ? map.rows - 1 : row_h;

			// 计算插值权重
			const double sl = double(row_h - _r) / (row_h - row_l);
			const double sh = double(_r - row_l) / (row_h - row_l);

			// 进行线性插值
			for (int _c = 0; _c < map.cols; _c++) {
				fitMap.ptr<float>(_r)[_c] = sl * fitMap.ptr<float>(row_l)[_c]
					+ sh * fitMap.ptr<float>(row_h)[_c];
			}
		}
	}

	return fitMap;
}