#include "img_utils.h"


ImageUtils::ImageUtils() {}

QPixmap ImageUtils::mat2QPixmap(const cv::Mat& mat)
{
    if(mat.empty())
        return QPixmap(); // 返回空的 QPixmap

    QImage img;

    // 根据 Mat 的通道数选择不同的转换方式
    if(mat.channels() == 1){
        // 灰度图像
        img = QImage(mat.data, mat.cols, mat.rows, static_cast<int>(mat.step), QImage::Format_Grayscale8).copy();
    }
    else if(mat.channels() == 3)
    {
        // 彩色图像 (OpenCV 默认是 BGR，需要转换为 RGB)
        cv::Mat rgb;
        cv::cvtColor(mat, rgb, cv::COLOR_BGR2RGB);
        img = QImage(rgb.data, rgb.cols, rgb.rows, static_cast<int>(rgb.step), QImage::Format_RGB888).copy();
    }
    else if(mat.channels() == 4)
    {
        // 如果需要处理带有透明通道的图像 (BGRA 转 RGBA)
        cv::Mat rgba;
        cv::cvtColor(mat, rgba, cv::COLOR_BGRA2RGBA);
        img = QImage(rgba.data, rgba.cols, rgba.rows, static_cast<int>(rgba.step), QImage::Format_RGBA8888).copy();
    }
    else
    {
        // 不支持的图像格式
        std::cout << "Unsupported Mat format with channels:" << mat.channels();
        return QPixmap();
    }

    return QPixmap::fromImage(img);
}

cv::Mat ImageUtils::mil2Mat(const MIL_ID mil_img)
{
    using namespace cv;

        // 获取 MIL 图像的宽度、高度和通道数
        MIL_INT width, height, channels, bitDepth;

        MbufInquire(mil_img, M_SIZE_X, &width);
        MbufInquire(mil_img, M_SIZE_Y, &height);
        MbufInquire(mil_img, M_SIZE_BAND, &channels);
        MbufInquire(mil_img, M_SIZE_BIT, &bitDepth);

        if (channels == 1) {
            // 单通道图像，直接读取整个缓冲区
            Mat grayImage(height, width, CV_8UC1);
            if (bitDepth == 1) {
                MIL_ID temp_img;
                convert_to_uint8(mil_img, temp_img);
                MbufGet(temp_img, grayImage.data);
                MbufFree(temp_img);
            } else {
                MbufGet(mil_img, grayImage.data);
            }
            return grayImage;
        }
        if (channels == 3) {
            // 多通道图像，分通道读取
            MIL_ID redChannel, greenChannel, blueChannel;
            MbufAlloc2d(MilSystem, width, height, 8 + M_UNSIGNED, M_IMAGE + M_PROC, &redChannel);
            MbufAlloc2d(MilSystem, width, height, 8 + M_UNSIGNED, M_IMAGE + M_PROC, &greenChannel);
            MbufAlloc2d(MilSystem, width, height, 8 + M_UNSIGNED, M_IMAGE + M_PROC, &blueChannel);

            // 将 MIL 图像的各个通道复制到单通道缓冲区
            MbufCopyColor(mil_img, redChannel, M_RED);
            MbufCopyColor(mil_img, greenChannel, M_GREEN);
            MbufCopyColor(mil_img, blueChannel, M_BLUE);

            // 分别读取每个通道的数据
            Mat redMat(height, width, CV_8UC1);
            Mat greenMat(height, width, CV_8UC1);
            Mat blueMat(height, width, CV_8UC1);
            MbufGet(redChannel, redMat.data);
            MbufGet(greenChannel, greenMat.data);
            MbufGet(blueChannel, blueMat.data);
            // 释放通道缓冲区
            MbufFree(redChannel);
            MbufFree(greenChannel);
            MbufFree(blueChannel);

            // 合并通道
            std::vector<Mat> bgrChannels = {blueMat, greenMat, redMat};
            Mat colorImage;
            cv::merge(bgrChannels, colorImage);

            return colorImage;
        }
        // 不支持的通道数
        std::cerr << "[Error] Unsupported number of channels: " << channels << std::endl;
        return Mat();
}

void ImageUtils::convert_to_uint8(const MIL_ID& input_img, MIL_ID& output_img) {
    MIL_INT size_x = MbufInquire(input_img, M_SIZE_X, M_NULL);
    MIL_INT size_y = MbufInquire(input_img, M_SIZE_Y, M_NULL);
    MIL_INT channel_num = MbufInquire(input_img,  M_SIZE_BAND, M_NULL);

    MbufAlloc2d(MilSystem, size_x, size_y, 8 + M_UNSIGNED, M_IMAGE + M_PROC + M_DISP, &output_img);
    if(channel_num == 1) {
        MimArith(output_img, input_img, output_img, M_ADD);
        MimArith(output_img, 255.0, output_img, M_MULT_CONST);
    } else if(channel_num == 3) {
        MimConvert(input_img, output_img, M_RGB_TO_L);
        MimArith(output_img, M_NULL, output_img, M_NOT);
    } else {
        std::cout << "Unsupported channel number!" << std::endl;
    }
}

cv::Mat ImageUtils::overlayResultOnInput(const cv::Mat &cv_input, const cv::Mat &total_result, double alpha, int colormap)
{
    // 1. 确保 total_result 的尺寸与 cv_input 一致
    cv::Mat resized_total_result;
    if (total_result.size() != cv_input.size()) {
        cv::resize(total_result, resized_total_result, cv_input.size(), 0, 0, cv::INTER_LINEAR);
    } else {
        resized_total_result = total_result;
    }

    // 2. 将 resized_total_result 转换为伪彩色图像
    cv::Mat total_result_color;
    // 确保 total_result 是单通道图像以应用色图
    if (resized_total_result.channels() == 1) {
        cv::applyColorMap(resized_total_result, total_result_color, colormap);
    } else {
        // 如果已经是多通道图像，可以选择直接使用或根据需要处理
        total_result_color = resized_total_result.clone();
    }

    // 3. 确保 cv_input 是三通道图像（如果是灰度图像，则转换为 BGR）
    cv::Mat cv_input_rgb;
    if (cv_input.channels() == 1) {
        cv::cvtColor(cv_input, cv_input_rgb, cv::COLOR_GRAY2BGR);
    } else {
        cv_input_rgb = cv_input.clone(); // 保证不修改原始图像
    }

    // 4. 设置叠加透明度（alpha: 0.0-1.0）
    double beta = 1.0 - alpha;

    // 5. 使用加权和将 total_result_color 叠加到 cv_input_rgb 上
    cv::Mat overlay;
    cv::addWeighted(cv_input_rgb, alpha, total_result_color, beta, 0.0, overlay);

    // 6. 返回叠加后的图像
    return overlay;
}

std::vector<std::vector<uint8_t> > ImageUtils::mergeMasks(const std::vector<std::vector<uint8_t> > &mask1, const std::vector<std::vector<uint8_t> > &mask2, int offset_y)
{
    if (mask1.empty() || mask1[0].empty() || mask2.empty() || mask2[0].empty()) {
        throw std::invalid_argument("输入的掩码不能为空");
    }

    // 假设两个掩码的宽度相同
    size_t width = mask1[0].size();
    for (const auto& row : mask1)
    {
        if (row.size() != width)
        {
            throw std::invalid_argument("mask1 的所有行必须具有相同的宽度");
        }
    }
    for (const auto& row : mask2)
    {
        if (row.size() != width)
        {
            throw std::invalid_argument("mask2 的所有行必须具有相同的宽度，并且与 mask1 的宽度一致");
        }
    }

    size_t height1 = mask1.size();
    size_t height2 = mask2.size();

    // 计算合并后的高度
    int merged_top = std::min(0, offset_y);
    int merged_bottom = std::max(static_cast<int>(height1), static_cast<int>(height2) + offset_y);
    size_t merged_height = merged_bottom - merged_top;

    // 初始化合并后的掩码为全 0
    std::vector<std::vector<uint8_t>> merged_mask(merged_height, std::vector<uint8_t>(width, 0));

    // 将 mask1 放入 merged_mask
    for (size_t y = 0; y < height1; ++y)
    {
        int merged_y = y - merged_top;
        for (size_t x = 0; x < width; ++x)
        {
            merged_mask[merged_y][x] = mask1[y][x] ? 1 : 0;
        }
    }

    // 将 mask2 放入 merged_mask，考虑 offset_y
    for (size_t y = 0; y < height2; ++y)
    {
        int merged_y = y + offset_y - merged_top;
        if (merged_y < 0 || merged_y >= static_cast<int>(merged_height))
        {
            // 超出合并掩码的范围，跳过
            continue;
        }
        for (size_t x = 0; x < width; ++x)
        {
            merged_mask[merged_y][x] = (merged_mask[merged_y][x] || mask2[y][x]) ? 1 : 0;
        }
    }

    return merged_mask;
}

std::pair<std::vector<std::vector<uint8_t> >, std::vector<std::vector<uint8_t> > > ImageUtils::extractROI(const std::vector<std::vector<uint8_t> > &merged_mask, int roi_x, int roi_y, int roi_width, int roi_height)
{
    if (merged_mask.empty() || merged_mask[0].empty()) {
        throw std::invalid_argument("merged_mask 不能为空");
    }

    size_t mask_height = merged_mask.size();
    size_t mask_width = merged_mask[0].size();

    // 校正 ROI 的边界
    int x_start = std::max(0, roi_x);
    int y_start = std::max(0, roi_y);
    int x_end = std::min(static_cast<int>(mask_width), roi_x + roi_width);
    int y_end = std::min(static_cast<int>(mask_height), roi_y + roi_height);

    // 初始化 ROI 和剩余部分的掩码
    std::vector<std::vector<uint8_t>> roi_mask(roi_height, std::vector<uint8_t>(roi_width, 0));
    std::vector<std::vector<uint8_t>> remaining_mask(mask_height, std::vector<uint8_t>(mask_width, 0));

    for (size_t y = 0; y < mask_height; ++y) {
        for (size_t x = 0; x < mask_width; ++x) {
            if (x >= static_cast<size_t>(roi_x) && x < static_cast<size_t>(roi_x + roi_width) &&
                y >= static_cast<size_t>(roi_y) && y < static_cast<size_t>(roi_y + roi_height)) {
                // 属于 ROI
                if (x >= static_cast<size_t>(x_start) && y >= static_cast<size_t>(y_start) &&
                    x < static_cast<size_t>(x_end) && y < static_cast<size_t>(y_end)) {
                    roi_mask[y - roi_y][x - roi_x] = merged_mask[y][x];
                }
            } else {
                // 属于剩余部分
                remaining_mask[y][x] = merged_mask[y][x];
            }
        }
    }

    return { roi_mask, remaining_mask };
}

// cv::Mat overlayResultOnInput(const cv::Mat& cv_input, const cv::Mat& total_result, double alpha = 0.5, int colormap = cv::COLORMAP_JET) {

// }


// 合并两个二值掩码，通过按位或操作，并考虑垂直偏移 offset_y
// std::vector<std::vector<uint8_t>> mergeMasks(
//     const std::vector<std::vector<uint8_t>>& mask1,
//     const std::vector<std::vector<uint8_t>>& mask2,
//     int offset_y = 0
// ) {

// }

// 提取 ROI 并返回剩余部分
// std::pair<std::vector<std::vector<uint8_t>>, std::vector<std::vector<uint8_t>>>
// extractROI(
//     const std::vector<std::vector<uint8_t>>& merged_mask,
//     int roi_x,
//     int roi_y,
//     int roi_width,
//     int roi_height
// ) {

// }