#include "Arith_FeaStitch.h"
#include "FileCache.h"
#include "Arith_FeaMatch.h"
#include "Arith_Utils.h"
#include <opencv2/opencv.hpp>
#include <omp.h>

API_FeaStitch* API_FeaStitch::Create(std::string cachedir)
{
	return new FeaStitch();
}


void API_FeaStitch::Destroy(API_FeaStitch* obj)
{
	delete obj;
}



FeaStitch::FeaStitch()
{
	_FeaMatcher = new FeatureMatcher(DetectorType::ORB, MatcherType::FLANN);
}

FeaStitch::~FeaStitch()
{
	if (_FeaMatcher)
	{
		delete _FeaMatcher;
		_FeaMatcher = nullptr;
	}
}

GenPanInfo FeaStitch::Init(GD_VIDEO_FRAME_S img, FrameInfo info)
{
	printf("FeaStitch::Init start\n");
	// 创建全景图
	_panPara.m_pan_width = info.nWidth * 2;
	_panPara.m_pan_height = info.nHeight * 2;

	// 获取Mat
	cv::Mat src = getRGBAMatFromGDFrame(img, img.u64VirAddr[0]);

	printf("FeaStitch::getRGBAMatFromGDFrame\n");

	_panImage = cv::Mat::zeros(_panPara.m_pan_height, _panPara.m_pan_width, CV_8UC3);

	printf("FeaStitch::Init _panImage\n");

	// 使用SetBase初始化，默认从中心开始
	SetBase(src, info, EdgeDirection::NONE);

	printf("FeaStitch::SetBase\n");

	return _panPara;
}





EdgeDirection FeaStitch::isNearEdge(cv::Mat& H,cv::Size size)
{
	// 获取变换后的四个角点
	std::vector<cv::Point2f> corners(4);
	corners[0] = cv::Point2f(0, 0);
	corners[1] = cv::Point2f(size.width, 0);
	corners[2] = cv::Point2f(size.width, size.height);
	corners[3] = cv::Point2f(0, size.height);

	std::vector<cv::Point2f> transformed_corners;
	cv::perspectiveTransform(corners, transformed_corners, H.clone());

	// 检查是否有角点接近全景图边界
	const float margin = 200.0f; // 增加边界容差
	bool nearLeft = false, nearRight = false, nearTop = false, nearBottom = false;
	int nearCount = 0; // 记录接近边界的角点数量

	for (const auto& corner : transformed_corners)
	{
		if (corner.x < margin) 
		{
			nearLeft = true;
			nearCount++;
		}
		if (corner.x > _panPara.m_pan_width - margin) 
		{
			nearRight = true;
			nearCount++;
		}
		if (corner.y < margin) 
		{
			nearTop = true;
			nearCount++;
		}
		if (corner.y > _panPara.m_pan_height - margin) 
		{
			nearBottom = true;
			nearCount++;
		}
	}

	// 只有当至少两个角点接近边界时才触发重新定位
	if (nearCount < 2)
	{
		return EdgeDirection::NONE;
	}

	// 判断出界方向
	if (nearLeft && nearTop) return EdgeDirection::TOP_LEFT;
	if (nearRight && nearTop) return EdgeDirection::TOP_RIGHT;
	if (nearLeft && nearBottom) return EdgeDirection::BOTTOM_LEFT;
	if (nearRight && nearBottom) return EdgeDirection::BOTTOM_RIGHT;
	if (nearLeft) return EdgeDirection::LEFT;
	if (nearRight) return EdgeDirection::RIGHT;
	if (nearTop) return EdgeDirection::TOP;
	if (nearBottom) return EdgeDirection::BOTTOM;

	return EdgeDirection::NONE;
}

void FeaStitch::calculateInitialOffset(EdgeDirection direction, int& offsetX, int& offsetY)
{
	// 默认偏移量（中心位置）
	offsetX = _panPara.m_pan_width / 2;
	offsetY = _panPara.m_pan_height / 2;

	// 根据图像尺寸动态计算margin
	int margin = std::min(_panPara.m_pan_width, _panPara.m_pan_height) / 4;  // 使用图像尺寸的1/4作为margin
	margin = std::max(margin, 300);  // 确保margin至少300像素

	// 根据出界方向调整偏移量
	switch (direction)
	{
		case EdgeDirection::LEFT:
			offsetX = _panPara.m_pan_width - margin;
			break;
		case EdgeDirection::RIGHT:
			offsetX = margin;
			break;
		case EdgeDirection::TOP:
			offsetY = _panPara.m_pan_height - margin;
			break;
		case EdgeDirection::BOTTOM:
			offsetY = margin;
			break;
		case EdgeDirection::TOP_LEFT:
			offsetX = _panPara.m_pan_width - margin;
			offsetY = _panPara.m_pan_height - margin;
			break;
		case EdgeDirection::TOP_RIGHT:
			offsetX = margin;
			offsetY = _panPara.m_pan_height - margin;
			break;
		case EdgeDirection::BOTTOM_LEFT:
			offsetX = _panPara.m_pan_width - margin;
			offsetY = margin;
			break;
		case EdgeDirection::BOTTOM_RIGHT:
			offsetX = margin;
			offsetY = margin;
			break;
		default:
			break;
	}
}

void FeaStitch::SetBase(cv::Mat& img, FrameInfo& info, EdgeDirection edgeDirection)
{
	// 清空所有缓存
	_paraVec.clear();
	_keypointsVec.clear();
	_descriptorsVec.clear();
	_currMatrix.clear();
	_totalFrameCnt = 0;
	_failedFrameCnt = 0;
	// 清空全景图
	_panImage.setTo(0);

	// 计算初始偏移量
	int offsetX = _panPara.m_pan_width / 2;
	int offsetY = _panPara.m_pan_height / 2;
	
	// 根据出界方向计算偏移量
	calculateInitialOffset(edgeDirection, offsetX, offsetY);

	// 初始化基准H矩阵
	float s = 0.3;
	cv::Mat H0 = (cv::Mat_<double>(3, 3) <<
		s, 0, offsetX,
		0, s, offsetY,
		0, 0, 1
	);

	// 保存第一帧的H矩阵
	_currMatrix.push_back(H0);

	// 处理第一帧
	_paraVec.push_back(info);


	// 提取特征点
	std::vector<cv::KeyPoint> keypoints;
	cv::Mat descriptors;
	_FeaMatcher->extractFeatures(img, keypoints, descriptors);

	// 保存特征点和描述子
	_keypointsVec.push_back(keypoints);
	_descriptorsVec.push_back(descriptors);
	_totalFrameCnt = 1;

	// 将第一帧变换到全景图坐标系
	cv::Mat imagetmp(_panImage.size(), CV_8UC3, cv::Scalar(0, 0, 0));
	cv::warpPerspective(img, imagetmp, H0, imagetmp.size(), cv::INTER_LINEAR, cv::BORDER_TRANSPARENT);

	cv::Mat mask = cv::Mat::ones(img.size(), CV_8UC1);
	cv::Mat warped_mask;
	cv::warpPerspective(mask, warped_mask, H0, imagetmp.size(), cv::INTER_LINEAR);
	imagetmp.copyTo(_panImage, warped_mask);
}

float FeaStitch::calculateIOU(const cv::Mat& H1, const cv::Mat& H2, const cv::Size& imgSize)
{
	// 获取第一帧的四个角点
	std::vector<cv::Point2f> corners1(4);
	corners1[0] = cv::Point2f(0, 0);
	corners1[1] = cv::Point2f(imgSize.width, 0);
	corners1[2] = cv::Point2f(imgSize.width, imgSize.height);
	corners1[3] = cv::Point2f(0, imgSize.height);

	// 获取第二帧的四个角点
	std::vector<cv::Point2f> corners2 = corners1;

	// 变换角点
	std::vector<cv::Point2f> transformed_corners1, transformed_corners2;
	cv::perspectiveTransform(corners1, transformed_corners1, H1);
	cv::perspectiveTransform(corners2, transformed_corners2, H2);

	// 计算两个多边形的面积
	float area1 = cv::contourArea(transformed_corners1);
	float area2 = cv::contourArea(transformed_corners2);

	// 计算交集面积
	std::vector<cv::Point2f> intersection;
	cv::intersectConvexConvex(transformed_corners1, transformed_corners2, intersection);
	float intersectionArea = cv::contourArea(intersection);

	// 计算并集面积
	float unionArea = area1 + area2 - intersectionArea;

	// 返回IOU
	return intersectionArea / unionArea;
}

bool FeaStitch::checkHContinuity(const cv::Mat& H, const cv::Size& imgSize)
{
    // 1. 检查平移量
    double tx = H.at<double>(0, 2);
    double ty = H.at<double>(1, 2);
    double max_translation_ratio = 0.1;  // 最大允许平移比例
    if (std::abs(tx) > max_translation_ratio * imgSize.width 
		|| std::abs(ty) > max_translation_ratio * imgSize.height)
    {
		printf("平移量 tx:%.1f,ty:%.1f\n", tx, ty);
        return false;
    }
    
    // 2. 检查旋转角度
    double angle = std::atan2(H.at<double>(1, 0), H.at<double>(0, 0));
    double max_angle = 0.1;  // 最大允许旋转角度（弧度）
    if (std::abs(angle) > max_angle)
    {
		printf("旋转角 angle:%.1f\n", angle);
        return false;
    }
    
    // 3. 检查缩放因子
    double scale_x = std::sqrt(H.at<double>(0, 0) * H.at<double>(0, 0) + 
                             H.at<double>(0, 1) * H.at<double>(0, 1));
    double scale_y = std::sqrt(H.at<double>(1, 0) * H.at<double>(1, 0) + 
                             H.at<double>(1, 1) * H.at<double>(1, 1));
    double max_scale_change = 0.3;  // 最大允许缩放变化
    if (std::abs(scale_x - 1.0) > max_scale_change || 
        std::abs(scale_y - 1.0) > max_scale_change)
    {
		printf("scale_x:%.1f，scale_y:%.1f\n", scale_x,scale_y);
        return false;
    }
    
    // 4. 检查透视变换强度
    double perspective_x = std::abs(H.at<double>(2, 0));
    double perspective_y = std::abs(H.at<double>(2, 1));
    double max_perspective = 0.001;  // 最大允许透视变换强度
    if (perspective_x > max_perspective || perspective_y > max_perspective)
    {
		printf("perspective_x:%.1f，perspective_y:%.1f\n", perspective_x, perspective_y);
        return false;
    }
    
    return true;
}
bool FeaStitch::checkFinalHValidity(const cv::Mat& H, const cv::Size& imgSize)
{
    // 1. 检查图像四个角点变换后的位置
    std::vector<cv::Point2f> corners(4);
    corners[0] = cv::Point2f(0, 0);
    corners[1] = cv::Point2f(imgSize.width, 0);
    corners[2] = cv::Point2f(imgSize.width, imgSize.height);
    corners[3] = cv::Point2f(0, imgSize.height);
    
    std::vector<cv::Point2f> transformed_corners;
    cv::perspectiveTransform(corners, transformed_corners, H);
    
    // 2. 检查变换后的四边形是否合理
    // 2.1 检查面积变化
    double original_area = imgSize.width * imgSize.height;
    double transformed_area = cv::contourArea(transformed_corners);
    double area_ratio = transformed_area / original_area;
    
    const double MIN_AREA_RATIO = 0.05;  // 最小面积比
    const double MAX_AREA_RATIO = 2; // 最大面积比
    if (area_ratio < MIN_AREA_RATIO || area_ratio > MAX_AREA_RATIO)
    {
        return false;
    }
    
    // 2.2 检查四边形是否严重变形
    // 计算相邻边之间的角度
    std::vector<double> angles;
    for (int i = 0; i < 4; i++)
    {
        cv::Point2f v1 = transformed_corners[(i + 1) % 4] - transformed_corners[i];
        cv::Point2f v2 = transformed_corners[(i + 2) % 4] - transformed_corners[(i + 1) % 4];
        double angle = std::abs(std::atan2(v1.x * v2.y - v1.y * v2.x, v1.x * v2.x + v1.y * v2.y));
        angles.push_back(angle);
    }
    
    // 检查角度是否接近90度
    const double ANGLE_TOLERANCE = 0.5;  // 弧度
    for (double angle : angles)
    {
        if (std::abs(angle - CV_PI/2) > ANGLE_TOLERANCE)
        {
            return false;
        }
    }

    return true;
}
SINT32 FeaStitch::Run(GD_VIDEO_FRAME_S img, FrameInfo para)
{
	cv::Mat src = getRGBAMatFromGDFrame(img, img.u64VirAddr[0]);

	// 连续5帧未成功拼接，重置基准图像
	if (_failedFrameCnt > 3)
	{
		SetBase(src, para, EdgeDirection::NONE);
	}

	// 提取特征点
	std::vector<cv::KeyPoint> currKeypoints;
	cv::Mat currDescriptors;
	_FeaMatcher->extractFeatures(src, currKeypoints, currDescriptors);



	// 获取当前帧和前一帧的数据
	std::vector<cv::KeyPoint>& prevKeypoints = _keypointsVec[_totalFrameCnt - 1];
	cv::Mat& prevDescriptors = _descriptorsVec[_totalFrameCnt - 1];

	// 特征匹配
	std::vector<cv::DMatch> matches;
	_FeaMatcher->matchFeatures(prevDescriptors, currDescriptors, matches);

	// 提取匹配点对
	std::vector<cv::Point2f> prevPoints, currPoints;
	for (auto& match : matches)
	{
		prevPoints.push_back(prevKeypoints[match.queryIdx].pt);
		currPoints.push_back(currKeypoints[match.trainIdx].pt);
	}

	// 计算单应性矩阵
	cv::Mat H;
	if (prevPoints.size() >= 4)
	{
		H = cv::findHomography(currPoints, prevPoints, cv::RANSAC, 3.0);
	}

	if (H.empty())
	{
		_failedFrameCnt++;
		return _totalFrameCnt;
	}

	// 更新当前H矩阵
	cv::Mat newH;
	if (_currMatrix.empty())
	{
		newH = H;
	}
	else
	{
		newH = _currMatrix.back() * H;
	}

	// 从第2帧开始
	if (_currMatrix.size() > 1)
	{
		// 检查H矩阵的连续性
		bool bHContinuity =  checkHContinuity(H,src.size());
		// 检查最终H矩阵的有效性
		bool bFinalHValidity = checkFinalHValidity(newH, src.size());

		if (!bHContinuity)
		{
			printf("H矩阵的连续性不足!\n");
		}
		if (!bFinalHValidity)
		{
			printf("最终H矩阵的有效性不足!\n");
		}

		if ( !bHContinuity || !bFinalHValidity)
		{
			_failedFrameCnt++;
			return _totalFrameCnt;
		}
	}



	// 检查是否接近边缘
	EdgeDirection edgeDirection = isNearEdge(newH,cv::Size(img.u32Width,img.u32Height));
	if (edgeDirection != EdgeDirection::NONE && _totalFrameCnt > 3)
	{
		// 重新初始化全景图
		SetBase(src, para, edgeDirection);
		return _totalFrameCnt;
	}

	// 检查IOU，如果与上一帧重叠度太高则跳过
	if (!_currMatrix.empty())
	{
		float iou = calculateIOU(_currMatrix.back(), newH, src.size());
		if (iou > 0.8f)  // 如果重叠度超过80%，认为是无效帧
		{
			printf("IOU低，跳过!\n");
			return _totalFrameCnt;
		}
	}

	// 重置失败计数
	_failedFrameCnt = 0;

	// 保存新的H矩阵
	_currMatrix.push_back(newH);

	// 保存帧参数和数据
	_paraVec.push_back(para);

	// 保存特征点和描述子
	_keypointsVec.push_back(currKeypoints);
	_descriptorsVec.push_back(currDescriptors);
	_totalFrameCnt++;


	// 将当前帧变换到全景图坐标系
	cv::Mat imagetmp(_panImage.size(), CV_8UC3, cv::Scalar(0, 0, 0));
	cv::warpPerspective(src, imagetmp, newH, imagetmp.size(), cv::INTER_LINEAR, cv::BORDER_TRANSPARENT);

	cv::Mat mask = cv::Mat::ones(src.size(), CV_8UC1);
	cv::Mat warped_mask;
	cv::warpPerspective(mask, warped_mask, newH, imagetmp.size(), cv::INTER_LINEAR);
	imagetmp.copyTo(_panImage, warped_mask);

	return _totalFrameCnt;
}

GD_VIDEO_FRAME_S FeaStitch::ExportPanAddr()
{
	GD_VIDEO_FRAME_S pan_out;
	pan_out.enPixelFormat = GD_PIXEL_FORMAT_RGB_PACKED;
	pan_out.u32Width = _panPara.m_pan_width;
	pan_out.u32Height = _panPara.m_pan_height;
	pan_out.u64VirAddr[0] = _panImage.data;
	return pan_out;
}

cv::Mat FeaStitch::ExportPanMat()
{
	return _panImage;
}

