#include "API_FrontStitch.h"
#include "Arith_FrontStitch.h"
#include "Arith_Utils.h"
#include "Arith_CoordModule.h"
#include "Arith_FeaMatch.h"
#include <opencv2/opencv.hpp>
#include <omp.h>


using namespace std;
using namespace cv;

API_FrontStitch* API_FrontStitch::Create(std::string cachedir)
{
    return new FrontStitch(cachedir);
}


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

FrontStitch::FrontStitch(std::string cachedir)
{
    // 创建缓存对象
    _cache = new FileCache<FrameCache>(50, cachedir);

    _GeoSolver = new GeoSolver();

    _FeaMatcher = new FeatureMatcher(DetectorType::SIFT, MatcherType::FLANN);

    _BATask = new BA_Task(_cache);

    _FBlendTask = new FrontMapBlend(_cache);
}

FrontStitch::~FrontStitch()
{
    #ifdef ENABLE_CUDA
        cudaFree(_cuda_mem.global_cuda_Frame);
        cudaFree(_cuda_mem.global_cuda_Pan);
        cudaFree(_cuda_mem.global_cuda_H_inv_data);
        cudaFree(_cuda_mem.global_cuda_pan_mask);
    #endif
}


FPanInfo FrontStitch::Init(FrameInfo info, ScanRange stAzRange, ScanRange stPtRange)
{
    // 帧角分辨率
    if (info.camInfo.fAglReso == 0)
    {
        info.camInfo.fAglReso = info.camInfo.fPixelSize / info.camInfo.nFocus * 0.06;
    }

    // 视场角计算
    float FOV_W = info.camInfo.fAglReso * info.nWidth;
    float FOV_H = info.camInfo.fAglReso * info.nHeight;

    // 设置当前拼接起始位置
    _GeoSolver->SetOriginPoint(info);


    // 大地系下范围计算

    // west
    info.servoInfo.fServoAz = stAzRange.Agl0;
    cv::Mat H1 = _GeoSolver->findHomography(info);
    Pole Pole_west = getPoleFromImgWithH(H1, cv::Point2f(0, info.nHeight / 2), info.nEvHeight);
    // east
    info.servoInfo.fServoAz = stAzRange.Agl1;
    cv::Mat H2 = _GeoSolver->findHomography(info);
    Pole Pole_east = getPoleFromImgWithH(H2, cv::Point2f(info.nWidth, info.nHeight / 2), info.nEvHeight);
    // up
    info.servoInfo.fServoPt = stPtRange.Agl1;
    cv::Mat H3 = _GeoSolver->findHomography(info);
    Pole Pole_up = getPoleFromImgWithH(H3, cv::Point2f(info.nWidth/2, 0), info.nEvHeight);
    // bottom
    info.servoInfo.fServoPt = stPtRange.Agl0;
    cv::Mat H4 = _GeoSolver->findHomography(info);
    Pole Pole_bottom = getPoleFromImgWithH(H4, cv::Point2f(info.nWidth/2, info.nHeight), info.nEvHeight);



    _panPara.centerGeo.fAz = (Pole_east.beta + Pole_west.beta)/2;
    _panPara.centerGeo.fPt = (Pole_up.alpha + Pole_bottom.alpha)/2;
    
    _panPara.stAzRangeGeo.Agl0 = Pole_west.beta;
    _panPara.stAzRangeGeo.Agl1 = Pole_east.beta;
    _panPara.stPtRangeGeo.Agl0 = Pole_bottom.alpha;
    _panPara.stPtRangeGeo.Agl1 = Pole_up.alpha;


    _panPara.fAglRes = info.camInfo.fAglReso * 2;

    // 全景图规模
    _panPara.m_pan_width = DEGLIM360(Pole_east.beta - Pole_west.beta) / _panPara.fAglRes;
    _panPara.m_pan_height = DEGLIM360(Pole_up.alpha - Pole_bottom.alpha) / _panPara.fAglRes;

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

#ifdef ENABLE_CUDA
    cudaMalloc((void**)&_cuda_mem.global_cuda_Frame, info.nWidth * info.nHeight * 3);//rgb
    cudaMalloc((void**)&_cuda_mem.global_cuda_Pan, _panPara.m_pan_width * _panPara.m_pan_height * 3);//rgb
    cudaMalloc((void**)&_cuda_mem.global_cuda_H_inv_data, 9 * sizeof(double));//
    cudaMalloc((void**)&_cuda_mem.global_cuda_pan_mask, _panPara.m_pan_width * _panPara.m_pan_height * 1);//gray
#endif

    return _panPara;
}

BYTE8 FrontStitch::Run(GD_VIDEO_FRAME_S img, FrameInfo info)
{
    // 极坐标拼接
    PoleStitch(img, info);

    //ReceiveFrame(img, info);

    return 0;
}

SINT32 FrontStitch::OptAndOutCurrPan()
{
    // 优化所有帧 效果不佳，放弃
    //ProcessFrame(_recvFrameKey);

    // 重绘所有帧
    //_FBlendTask->DirectMap(_recvFrameKey, _panPara, _panImage, _panMask, _cuda_mem);

    return SINT32();
}

BYTE8 FrontStitch::PoleStitch(GD_VIDEO_FRAME_S img, FrameInfo info)
{
    cv::Mat H = _GeoSolver->findHomography(info);
    float dep = info.nEvHeight;

    // 获取rgb源图像
    cv::Mat rgb = getRGBMatFromGDFrame(img, img.u64VirAddr[0]);

    //Pole CT_map = getPoleFromImgWithH(H, cv::Point2f(img.u32Width/2, img.u32Height/2), info.nEvHeight);

    // 计算帧的map四至，注意这里不是线性变化，不能完全保证四至约束下的原图能完全覆盖，即反采样后图像可能不完整
    Pole leftTop_map = getPoleFromImgWithH(H, cv::Point2f(0, 0), info.nEvHeight);
    Pole rightTop_map = getPoleFromImgWithH(H, cv::Point2f(img.u32Width, 0), info.nEvHeight);
    Pole rightBottom_map = getPoleFromImgWithH(H, cv::Point2f(img.u32Width, img.u32Height), info.nEvHeight);
    Pole leftBottom_map = getPoleFromImgWithH(H, cv::Point2f(0, img.u32Height), info.nEvHeight);

    //leftTop_map.beta = DEGLIM360(leftTop_map.beta);
    //rightTop_map.beta = DEGLIM360(rightTop_map.beta);
    //rightBottom_map.beta = DEGLIM360(rightBottom_map.beta);
    //leftBottom_map.beta = DEGLIM360(leftBottom_map.beta);

    // 转全景图的像方
    cv::Point2f p_leftTop_map = getFPanFromPole(leftTop_map, _panPara);
    cv::Point2f p_rightTop_map = getFPanFromPole(rightTop_map, _panPara);
    cv::Point2f p_rightBottom_map = getFPanFromPole(rightBottom_map, _panPara);
    cv::Point2f p_leftBottom_map = getFPanFromPole(leftBottom_map, _panPara);

    // 计算全景图上范围
    int right = max(max(max(p_leftTop_map.x, p_rightTop_map.x), p_rightBottom_map.x), p_leftBottom_map.x);
    int left = min(min(min(p_leftTop_map.x, p_rightTop_map.x), p_rightBottom_map.x), p_leftBottom_map.x);
    int top = min(min(min(p_leftTop_map.y, p_rightTop_map.y), p_rightBottom_map.y), p_leftBottom_map.y);
    int bottom = max(max(max(p_leftTop_map.y, p_rightTop_map.y), p_rightBottom_map.y), p_leftBottom_map.y);

    int xRange = right - left;
    int yRnage = bottom - top;

    //反映射到像素坐标
    MINMAXRECT32S RANGES = { 0 };
    RANGES.minY = std::max(0, top);
    RANGES.maxY = std::min(_panImage.rows - 1, bottom);

    RANGES.minX = std::max(0, left);
    RANGES.maxX = std::min(_panImage.cols - 1, right);

    cv::Rect2d roi(RANGES.minX, RANGES.minY, RANGES.maxX - RANGES.minX, RANGES.maxY - RANGES.minY);

    if (roi.width < 0)
    {
        int a = 0;
    }

    cv::Mat H_inv = H.inv();

#ifdef ENABLE_CUDA
    UpdatePan_CUDA(rgb, _panImage, roi, H_inv, _panPara, dep, _cuda_mem);
#else
    UpdatePan_CPU(rgb, _panImage, roi, H_inv, _panPara, dep);
#endif

    return 0;
}


SINT32 FrontStitch::ReceiveFrame(GD_VIDEO_FRAME_S img, FrameInfo para)
{
    // 构造单帧缓存
    auto _t_frame_cache = std::make_shared<FrameCache>();

    _t_frame_cache->_para = para;
    _t_frame_cache->_frame_info = img;

    int imgSize = 0;

    if (img.enPixelFormat == GD_PIXEL_FORMAT_GRAY_Y8)
    {
        imgSize = img.u32Height * img.u32Width;
    }
    if (img.enPixelFormat == GD_PIXEL_FORMAT_RGB_PACKED)
    {
        imgSize = img.u32Height * img.u32Width * 3;
    }
    if (img.enPixelFormat == GD_PIXEL_FORMAT_NV12)
    {
        imgSize = img.u32Height * img.u32Width * 1.5;
    }

    // 深拷贝图像数据
    memcpy(_t_frame_cache->_data, img.u64VirAddr[0], imgSize);

    // 提取特征点
    std::vector<cv::KeyPoint> keypoints;
    cv::Mat descriptors;


    cv::Mat src = getRGBAMatFromGDFrame(img, img.u64VirAddr[0]);
    _FeaMatcher->extractFeatures(src, keypoints, descriptors);

    size_t keyNum = MIN(keypoints.size(), FEA_NUM_MAX);

    _t_frame_cache->ptNum = keyNum;

    // 保存特征点
    memcpy(_t_frame_cache->_pt, keypoints.data(), sizeof(cv::KeyPoint) * keyNum);

    // 保存描述子
    memcpy(_t_frame_cache->_desp, descriptors.data, sizeof(float) * keyNum * FEA_DES_SIZE);

    cv::Mat_<double> H0 = _GeoSolver->findHomography(para);

    // 保存初始H
    memcpy(_t_frame_cache->H, H0.data, sizeof(double) * 9);

    // 预处理结果加入文件缓存
    _cache->set(para.nFrmID, _t_frame_cache);

    _recvFrameKey.push_back(para.nFrmID);

    _totalFrameCnt++;

    return _totalFrameCnt;
}

SINT32 FrontStitch::ProcessFrame(vector<KeyType> keys)
{
    if (keys.size() == 0)
    {
        return -1;
    }

    // 优化所有帧
    _BATask->OptFrame(keys, cv::Mat::zeros(cv::Size(3,3),CV_8UC1));

    return 0;
}

GD_VIDEO_FRAME_S FrontStitch::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;
}

void UpdatePan_CPU(cv::Mat rgbFrame, cv::Mat pan, cv::Rect2d roi, cv::Mat H_inv, FPanInfo _panPara, float dep)
{
    cv::Mat subImg = pan(roi);
    // 遍历整个子图
 #pragma omp parallel for
    for (int i = 0; i < (int)roi.height; i++)
    {
        for (int j = 0; j < (int)roi.width; j++)
        {
            cv::Point2f p_img = getImgPosFromPole(H_inv, getPoleFromFPan(Point2f(j + roi.x, i + roi.y), _panPara), dep);
            if (p_img.x >= 200 && p_img.y >= 200 && p_img.x < rgbFrame.cols - 200 && p_img.y < rgbFrame.rows - 200)
            {
                auto rgbVal = Interpolation_RGB24(rgbFrame.data, rgbFrame.cols, rgbFrame.rows, p_img.x, p_img.y);
                uchar* row = subImg.ptr<uchar>(i); 
                row[3 * j + 0] = rgbVal.R;  
                row[3 * j + 1] = rgbVal.G;  
                row[3 * j + 2] = rgbVal.B;  

            }
        }
    }

}

