/***
 * function: 360 surrond view combine c++ demo
 * author: joker.mao
 * date: 2023/07/15
 * copyright: ADAS_EYES all right reserved
*/

#include "common.h"

#define DEBUG
#define AWB_LUN_BANLANCE_ENALE    1

// Compute blend weight that handles empty (black) regions
// Outside overlap: G=1 where imA has content, G=0 where imA is black
// Inside overlap: smooth transition using horizontal gradient
static cv::Mat make_blend_weight(const cv::Mat& imA, const cv::Mat& imB, bool l2r) {
    int h = imA.rows, cw = imA.cols;
    cv::Mat grayA, grayB, maskA, maskB;
    cv::cvtColor(imA, grayA, cv::COLOR_BGR2GRAY);
    cv::cvtColor(imB, grayB, cv::COLOR_BGR2GRAY);

    cv::threshold(grayA, maskA, 5, 255, cv::THRESH_BINARY);
    cv::threshold(grayB, maskB, 5, 255, cv::THRESH_BINARY);
    cv::Mat both = maskA & maskB;

    cv::Mat wm(h, cw, CV_32FC1);
    for (int y = 0; y < h; ++y) {
        for (int x = 0; x < cw; ++x) {
            if (both.at<uchar>(y, x)) {
                float t = (float)x / (float)(cw - 1);
                wm.at<float>(y, x) = l2r ? t : (1.0f - t);
            } else if (maskA.at<uchar>(y, x)) {
                wm.at<float>(y, x) = 1.0f;
            } else {
                wm.at<float>(y, x) = 0.0f;
            }
        }
    }
    return wm;
}

int main(int argc, char** argv)
{
    if (argc != 2) {
        std::cout << "usage:\n\t" << argv[0] << " path\n";
        return -1;
    }
    std::cout  << argv[0] << " app start running..." << std::endl;
    std::string data_path = std::string(argv[1]);
    cv::Mat car_img;
    cv::Mat origin_dir_img[4];
    cv::Mat undist_dir_img[4];
    cv::Mat merge_weights_img[4];
    cv::Mat out_put_img;
    CameraPrms prms[4];

    //1.read image
    car_img = cv::imread(data_path + "/images/car.png");
    cv::resize(car_img, car_img, cv::Size(xr - xl, yb - yt));
    out_put_img = cv::Mat(cv::Size(total_w, total_h), CV_8UC3, cv::Scalar(0, 0, 0));

    //1. read calibration prms
    for (int i = 0; i < 4; ++i) {
        auto& prm = prms[i];
        prm.name = camera_names[i];
        auto ok = read_prms(data_path + "/yaml/" + prm.name + ".yaml", prm);
        if (!ok) {
            return -1;
        }
    }

     //2.lum equalization and awb for four channel image
    std::vector<cv::Mat*> srcs;
    for (int i = 0; i < 4; ++i) {
        auto& prm = prms[i];
        origin_dir_img[i] = cv::imread(data_path + "/images/" + prm.name + ".png");
        srcs.push_back(&origin_dir_img[i]);
    }

#if AWB_LUN_BANLANCE_ENALE
    awb_and_lum_banlance(srcs);
#endif

    //3. undistort image
    for (int i = 0; i < 4; ++i) {
        auto& prm = prms[i];
        cv::Mat& src = origin_dir_img[i];
       
        undist_by_remap(src, src, prm);
        cv::warpPerspective(src, src, prm.project_matrix, project_shapes[prm.name]);
        
        if (camera_flip_mir[i] == "r+") {
            cv::rotate(src, src, cv::ROTATE_90_CLOCKWISE);
        } else if (camera_flip_mir[i] == "r-") {
            cv::rotate(src, src, cv::ROTATE_90_COUNTERCLOCKWISE);
        } else if (camera_flip_mir[i] == "m") {
            cv::rotate(src, src, cv::ROTATE_180);
        }
        //display_mat(src, "project");
        //cv::imwrite(prms.name + "_undist.png", src);
        undist_dir_img[i] = src.clone();
    }

    //4. Compute blend weights from projected images (matches Python get_weight_mask_matrix)
    // mask-based: only blend where both images have content, take fully outside overlap
    merge_weights_img[0] = make_blend_weight(
        undist_dir_img[0](cv::Rect(0, 0, xl, yt)),       // FI = front[:, :xl]
        undist_dir_img[1](cv::Rect(0, 0, xl, yt)),        // LI = left[:yt, :]
        true);                                              // G increases left→right
    merge_weights_img[1] = make_blend_weight(
        undist_dir_img[0](cv::Rect(xr, 0, xl, yt)),       // FII = front[:, xr:]
        undist_dir_img[3](cv::Rect(0, 0, xl, yt)),         // RII = right[:yt, :]
        false);                                             // G decreases left→right
    merge_weights_img[2] = make_blend_weight(
        undist_dir_img[2](cv::Rect(0, 0, xl, yt)),        // BIII = back[:, :xl]
        undist_dir_img[1](cv::Rect(0, yb, xl, yt)),        // LIII = left[yb:, :]
        true);                                              // G increases left→right
    merge_weights_img[3] = make_blend_weight(
        undist_dir_img[2](cv::Rect(xr, 0, xl, yt)),       // BIV = back[:, xr:]
        undist_dir_img[3](cv::Rect(0, yb, xl, yt)),        // RIV = right[yb:, :]
        false);                                             // G decreases left→right

    //5.start combine
    std::cout  << argv[0] << " app start combine" << std::endl;
    car_img.copyTo(out_put_img(cv::Rect(xl, yt, car_img.cols, car_img.rows)));
    //5.1 center copy
    for (int i = 0; i < 4; ++i) {
        cv::Rect roi;
        if (std::string(camera_names[i]) == "front") {
            roi = cv::Rect(xl, 0, xr - xl, yt);
            undist_dir_img[i](roi).copyTo(out_put_img(roi));
        } else if (std::string(camera_names[i]) == "left") {
            roi = cv::Rect(0, yt, xl, yb - yt);
            undist_dir_img[i](roi).copyTo(out_put_img(roi));
        } else if (std::string(camera_names[i]) == "right") {
            roi = cv::Rect(0, yt, xl, yb - yt);
            undist_dir_img[i](roi).copyTo(out_put_img(cv::Rect(xr, yt, total_w - xr, yb - yt)));
        } else if (std::string(camera_names[i]) == "back") {
            roi = cv::Rect(xl, 0, xr - xl, yt);
            undist_dir_img[i](roi).copyTo(out_put_img(cv::Rect(xl, yb, xr - xl, yt)));
        }
    }
    //5.2 four corner merge
    //image order: {front, left, back, right}
    //weight index: 0=left_top, 1=right_top, 2=left_bottom, 3=right_bottom
    cv::Rect roi;
    roi = cv::Rect(0, 0, xl, yt);
    merge_image(undist_dir_img[0](roi), undist_dir_img[1](roi), merge_weights_img[0], out_put_img(roi));
    roi = cv::Rect(xr, 0, xl, yt);
    merge_image(undist_dir_img[0](roi), undist_dir_img[3](cv::Rect(0, 0, xl, yt)), merge_weights_img[1], out_put_img(cv::Rect(xr, 0, xl, yt)));
    roi = cv::Rect(0, yb, xl, yt);
    merge_image(undist_dir_img[2](cv::Rect(0, 0, xl, yt)), undist_dir_img[1](roi), merge_weights_img[2], out_put_img(roi));
    merge_image(undist_dir_img[2](cv::Rect(xr, 0, xl, yt)), undist_dir_img[3](cv::Rect(0, yb, xl, yt)), merge_weights_img[3], out_put_img(cv::Rect(xr, yb, xl, yt)));

    cv::imwrite("ADAS_EYES_360_VIEW.png", out_put_img);
    
#ifdef DEBUG
    cv::resize(out_put_img, out_put_img, cv::Size(out_put_img.size())),
    display_mat(out_put_img, "out_put_img");
#endif

    std::cout  << argv[0] << " app finished" << std::endl;
}