#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 };
}

void VectorToImg(const std::vector<std::vector<uint8_t> > &array, const std::string &image_path)
{
    int height = array.size();
    int width = array[0].size();

    // 创建一个Mat对象来表示图像，CV_8UC1表示单通道8位无符号整数类型（用于黑白图像）
    cv::Mat image(height, width, CV_8UC1);

    // 遍历二维向量，设置图像像素值
    for (int y = 0; y < height; y++)
    {
        for (int x = 0; x < width; x++)
        {
            if (array[y][x] == 1)
            {
                image.at<uchar>(y, x) = 255;  // 白色像素，对应灰度值255
            } else
            {
                image.at<uchar>(y, x) = 0;    // 黑色像素，对应灰度值0
            }
        }
    }

    // 将图像保存为文件，传入的image_path指定保存路径和文件名
    imwrite(image_path, image);
}


