/*
 * Copyright 2024 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "client.h"

using namespace std;
namespace tc = triton::client;

DdcpTritonClient::DdcpTritonClient(const string& url) {
    // Create the Triton client once during object construction
    FAIL_IF_ERR(tc::InferenceServerHttpClient::Create(&client_, url, false), "unable to create http client");
}

ImgBase64Result DdcpTritonClient::DdcpRequestByPath(string ImagePath) {
    ImgBase64Result res;

    // 使用 ifstream 检查文件是否存在
    std::ifstream file(ImagePath, std::ios::binary);
    if (!file.is_open()) {
        // 如果文件无法打开，返回错误信息
        res.errorCode = 1;
        res.errorMessage = "Can not open file";
        return res;
    }

    unsigned char head[12] = {0};
    file.read((char*)head, 8);
    std::string imgFormat = getFormatFromHeader(head);

    cv::Mat src_img = cv::imread(ImagePath);
    res = DdcpRequestByImage(src_img, imgFormat);

    return res;
}

ImgBase64Result DdcpTritonClient::DdcpRequestByBase64(const char* ImageData) {
    ImgBase64Result res;
    if (strlen(ImageData) < 1) {
        res.errorCode = 1;
        res.errorMessage = "base64_string size error";
        return res;
    }
    std::string decoded_data = base64_decode(ImageData);
    unsigned char head[12] = {0};
    if (decoded_data.size() < 8) {
        res.errorCode = 1;
        res.errorMessage = "decoded_data size error";
        return res;
    }
    memcpy(head, decoded_data.data(), 8);
    std::string imgFormat = getFormatFromHeader(head);
    std::vector<uchar> data(decoded_data.begin(), decoded_data.end());
    cv::Mat Image = cv::imdecode(data, cv::IMREAD_COLOR);
    res = DdcpRequestByImage(Image, imgFormat);

    return res;
}

std::string DdcpTritonClient::GetConfigValue(const char* key, const char* engine) {
    tc::Error err = client_->LoadModel("config_bank");
    std::string model_name = "config_bank";
    tc::InferOptions options(model_name);
    options.model_version_ = "";
    options.priority_ = 1;

    std::vector<std::string> input_key_data{std::string(key)};
    std::vector<std::string> input_engine_data{std::string(engine)};

    tc::InferInput *input0, *input1;
    std::vector<int64_t> shape{1};
    tc::InferInput::Create(&input0, "key", shape, "BYTES");
    tc::InferInput::Create(&input1, "engine", shape, "BYTES");
    std::shared_ptr<tc::InferInput> input0_ptr, input1_ptr;
    input0_ptr.reset(input0);
    input1_ptr.reset(input1);
    input0_ptr->AppendFromString(input_key_data);
    input1_ptr->AppendFromString(input_engine_data);
    std::vector<tc::InferInput*> inputs = {input0_ptr.get(), input1_ptr.get()};

    tc::InferRequestedOutput* output0;
    tc::Error err_out0 = tc::InferRequestedOutput::Create(&output0, "value");
    std::shared_ptr<tc::InferRequestedOutput> output0_ptr;
    output0_ptr.reset(output0);
    std::vector<const tc::InferRequestedOutput*> outputs = {output0_ptr.get()};

    tc::InferResult* results;
    tc::Headers http_headers;
    tc::Error err_infer = client_->Infer(&results, options, inputs, outputs, http_headers);
    if (!err_infer.IsOk()) {
        std::cout << __LINE__ << " err_infer.Message(): " << err_infer.Message() << std::endl;
        return "";
    }

    std::shared_ptr<tc::InferResult> results_ptr;
    results_ptr.reset(results);

    char* output0_data_ptr;
    size_t recv_output0_byte_size;

    std::vector<std::string> text_output_data;
    tc::Error err_paser = results_ptr->StringData("value", &text_output_data);
    if (!err_paser.IsOk()) {
        std::cout << __LINE__ << " err_paser.Message(): " << err_infer.Message() << std::endl;
        return "";
    }
    return text_output_data[0];
}

ImgBase64Result DdcpTritonClient::DdcpRequestByImage(cv::Mat Image, std::string imgFormat) {
    ImgBase64Result res;
    res = UNet_.detect(Image, client_, model_name_, imgFormat);

    return res;
}

tc::Error DdcpTritonClient::LoadModel() {
    model_name_ = GetConfigValue("ddcp", "ddcp");
    if (model_name_.size() < 1) {
        printf("line:%d model_name.size():%ld\n", __LINE__, model_name_.size());
        model_name_ = "document_dewarping";
    }

    tc::Error err = client_->LoadModel(model_name_);

    return err;
}

tc::Error DdcpTritonClient::RemoveModel() {
    tc::Error err;

    return err;
}

cv::Mat Document_Dewarp::linearInterpolation(const cv::Mat& tshape, const cv::Mat& sshape, const cv::Mat& grid_x,
                                             const cv::Mat& grid_y) {
    cv::Mat grid_z(grid_x.size(), CV_32FC2);

    int src_h = tshape.rows;
    int src_w = tshape.cols;
    int dst_h = grid_x.rows;
    int dst_w = grid_x.cols;

    float x_slice = float(dst_w) / float(src_w) * 1.0f;
    float y_slice = float(dst_h) / float(src_h) * 1.0f;

    for (int i = 0; i < grid_x.rows; ++i) {
        for (int j = 0; j < grid_x.cols; ++j) {
            float x = grid_x.at<float>(i, j);
            float y = grid_y.at<float>(i, j);
            float x_1 = x / x_slice;
            float y_1 = y / y_slice;

            int x1 = static_cast<int>(std::floor(x / x_slice));
            int y1 = static_cast<int>(std::floor(y / y_slice));
            int x2 = x1 + 1;
            if (x2 >= src_w) {
                x2 = x1 - 1;
            }
            int y2 = y1 + 1;
            if (y2 >= src_h) {
                y2 = y1 - 1;
            }

            float wx = x_1 - (int)x_1;
            float wy = y_1 - (int)y_1;

            cv::Vec2f p11, p12, p21, p22;
            if (x1 >= 0 && x1 < tshape.rows && y1 >= 0 && y1 < tshape.cols) {
                p11 = sshape.at<cv::Vec2f>(x1, y1);
            } else {
                p11 = cv::Vec2f(0, 0);
            }
            if (x2 >= 0 && x2 < tshape.rows && y1 >= 0 && y1 < tshape.cols) {
                p12 = sshape.at<cv::Vec2f>(x2, y1);
            } else {
                p12 = cv::Vec2f(0, 0);
            }
            if (x1 >= 0 && x1 < tshape.rows && y2 >= 0 && y2 < tshape.cols) {
                p21 = sshape.at<cv::Vec2f>(x1, y2);
            } else {
                p21 = cv::Vec2f(0, 0);
            }
            if (x2 >= 0 && x2 < tshape.rows && y2 >= 0 && y2 < tshape.cols) {
                p22 = sshape.at<cv::Vec2f>(x2, y2);
            } else {
                p22 = cv::Vec2f(0, 0);
            }
            cv::Vec2f p1 = p11 * (1 - wx) + p12 * wx;
            cv::Vec2f p2 = p21 * (1 - wx) + p22 * wx;
            cv::Vec2f p = p1 * (1 - wy) + p2 * wy;

            grid_z.at<cv::Vec2f>(i, j) = p;
        }
    }
    return grid_z;
}

void Document_Dewarp::mesgrid(int xstart, int xend, int xpointnum, int ystart, int yend, int ypointnum, cv::Mat& x,
                              cv::Mat& y) {
    int xsquaresize = (xend - xstart) / xpointnum;
    int ysquaresize = (yend - xstart) / ypointnum;

    x.create(ysquaresize, xsquaresize, CV_32FC1);
    y.create(ysquaresize, xsquaresize, CV_32FC1);
    // 设置边界

    x.col(0).setTo(xstart * 1.0f);
    x.col(xsquaresize - 1).setTo((xend - 1) * 1.0f);
    y.row(0).setTo(ystart * 1.0f);
    y.row(ysquaresize - 1).setTo((yend - 1) * 1.0f);

    float xdelta = float(xend - xstart) / (float(xsquaresize) - 1.0f);  // 两个元素的间隔
    float ydelta = float(yend - ystart) / (float(ysquaresize) - 1.0f);  // 两个元素的间隔

    // 计算其他位置的值
    for (int i = 1; i < xsquaresize - 1; ++i) {
        x.col(i).setTo(float(i * xdelta));
    }

    for (int i = 1; i < ysquaresize - 1; ++i) {
        y.row(i).setTo(float(i * ydelta));
    }
}
cv::Mat Document_Dewarp::flatByRegressWithClassiy_multiProcess(vector<int>& vec_segment,
                                                               std::array<int, 2>& segment_shape,
                                                               vector<float>& vec_regress,
                                                               std::array<int, 4>& regress_shape, cv::Mat& srcImg) {
    cv::Mat srcImg_resize;
    // 图像缩放到（960*1024）;
    cv::resize(srcImg, srcImg_resize, cv::Size(960, 1024));
    // 计算flat_shap
    int flat_shape[2];
    flat_shape[0] = vec_segment[0] * this->fiducial_point_gaps[this->col_gap] * this->fiducial_point_num[this->col_gap];
    flat_shape[1] = vec_segment[1] * this->fiducial_point_gaps[this->row_gap] * this->fiducial_point_num[this->row_gap];

    cv::Mat fiducial_points(regress_shape[1], regress_shape[2], CV_32FC2, vec_regress.data());
    // 定义缩放因子
    cv::Mat scale_factor = (cv::Mat_<float>(1, 2) << 960.0f / 992.0f, 1024.0f / 992.0f);
    // 对点坐标进行缩放
    for (int i = 0; i < fiducial_points.cols; ++i) {
        for (int j = 0; j < fiducial_points.rows; j++) {
            cv::Vec2f point;
            point = fiducial_points.at<cv::Vec2f>(i, j).mul(scale_factor);
            fiducial_points.at<cv::Vec2f>(i, j) = point;
        }
    }

    cv::Mat im_x, im_y;
    mesgrid(0, fiducial_points.cols, this->fiducial_point_gaps[this->col_gap], 0, fiducial_points.rows,
            this->fiducial_point_gaps[this->row_gap], im_x, im_y);
    for (int y = 0; y < im_x.rows; ++y) {
        for (int x = 0; x < im_x.cols; ++x) {
            im_x.at<float>(y, x) = int(im_x.at<float>(y, x)) * vec_segment[1] * 1.0f;
        }
    }
    for (int y = 0; y < im_y.rows; ++y) {
        for (int x = 0; x < im_y.cols; ++x) {
            im_y.at<float>(y, x) = int(im_y.at<float>(y, x)) * vec_segment[0] * 1.0f;
        }
    }
    std::vector<cv::Mat> channels = {im_x, im_y};
    cv::Mat tshape;
    cv::merge(channels, tshape);

    cv::Mat sshape = fiducial_points.clone();
    cv::Mat grid_x, grid_y;
    int output_shape[2];
    output_shape[0] = vec_segment[0] * (fiducial_points.rows - 1);
    output_shape[1] = vec_segment[1] * (fiducial_points.cols - 1);

    mesgrid(0, output_shape[1], 1, 0, output_shape[0], 1, grid_x, grid_y);

    cv::Mat grid_z = linearInterpolation(tshape, sshape, grid_x, grid_y);

    for (int i = 0; i < grid_z.rows; ++i) {
        for (int j = 0; j < grid_z.cols; ++j) {
            cv::Vec2f point = grid_z.at<cv::Vec2f>(i, j);
        }
    }
    std::vector<cv::Mat> channels_1;
    split(grid_z, channels_1);
    cv::Mat output;
    cv::remap(srcImg_resize, output, channels_1[0], channels_1[1], cv::INTER_CUBIC);
#if 0
        cv::imshow("output", output);
        cv::waitKey(0);
#endif
    return output;
}

cv::Mat Document_Dewarp::preprocess(cv::Mat& m_srcImg) {
    cv::Mat dstimg;

    cv::resize(m_srcImg, dstimg, cv::Size(inputWidth, inputHeight), cv::INTER_LINEAR);
    this->input_image_.resize(inputHeight * inputWidth * dstimg.channels());
    int row = dstimg.rows;
    int col = dstimg.cols;
    for (int c = 0; c < 3; c++) {
        for (int i = 0; i < row; i++) {
            for (int j = 0; j < col; j++) {
                float pix = dstimg.ptr<uchar>(i)[j * 3 + c];
                this->input_image_[c * row * col + i * col + j] = pix * 1.0;
            }
        }
    }
    dstimg = cv::dnn::blobFromImage(dstimg);
    return dstimg;
}

ImgBase64Result Document_Dewarp::detect(cv::Mat& m_srcImg,
                                        const std::unique_ptr<triton::client::InferenceServerHttpClient>& client,
                                        const std::string& model_name_, std::string imgFormat) {
    ImgBase64Result res;
    cv::Mat result;
    cv::Mat image;

    if (m_srcImg.empty()) {
        res.errorCode = 2;
        res.errorMessage = "Image is empty";
        // 可以执行后续的图像处理操作
        std::cout << __LINE__ << " res.errorMessage: " << res.errorMessage << std::endl;
        return res;
    }
    // 图像预处理
    image = preprocess(m_srcImg);

#if 0
    // 准备输入数据
    array<int64_t, 4> input_shape_{ 1, 3, inputHeight, inputWidth};
	auto allocator_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
	Value input_tensor_ = Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());
	const std::array<const char*, 1> inputNames = { input_names[0].c_str() };
	const std::array<const char*, 2> outNames = { output_names[0].c_str(), output_names[1].c_str() };
 	std::vector<Ort::Value> ort_outputs;
    try {
        ort_outputs = net->Run(Ort::RunOptions{ nullptr }, inputNames.data(), &input_tensor_, 1, outNames.data(), outNames.size());
    }
    catch (std::exception e) {
        std::cout << e.what() << std::endl;
    }
	const float* f32Outputs = ort_outputs[0].GetTensorMutableData<float>();
    const float* f32Outputs_segment = ort_outputs[1].GetTensorMutableData<float>();

    int iOutputs_b = ort_outputs.at(0).GetTensorTypeAndShapeInfo().GetShape().at(0);
    int iOutputs_c = ort_outputs.at(0).GetTensorTypeAndShapeInfo().GetShape().at(1);
    int iOutputs_w = ort_outputs.at(0).GetTensorTypeAndShapeInfo().GetShape().at(2);
    int iOutputs_h = ort_outputs.at(0).GetTensorTypeAndShapeInfo().GetShape().at(3);

    int iSegment_w = ort_outputs.at(1).GetTensorTypeAndShapeInfo().GetShape().at(0);
    int iSegment_h = ort_outputs.at(1).GetTensorTypeAndShapeInfo().GetShape().at(1);
#else
    std::vector<int64_t> input_shape = {1, 3, inputHeight, inputWidth};
    std::vector<uint8_t> input_data(image.total() * image.channels() * 4);
    std::memcpy(input_data.data(), image.data, input_data.size());
    tc::InferInput* input0;

    tc::Error err_ino = tc::InferInput::Create(&input0, "input", input_shape, "FP32");
    if (!err_ino.IsOk()) {
        res.errorCode = 3;
        // res.errorMessage = "model error:unable to get IN0";
        res.errorMessage = err_ino.Message();
        // model error "unable to get product_reviews"
        std::cout << __LINE__ << " res.errorMessage: " << res.errorMessage << std::endl;
        return res;
    }

    std::shared_ptr<tc::InferInput> input0_ptr;
    input0_ptr.reset(input0);
    input0_ptr->SetBinaryData(true);
    input0_ptr->AppendRaw(input_data.data(), input_data.size());

    tc::InferRequestedOutput* output0;

    tc::Error err_out0 = tc::InferRequestedOutput::Create(&output0, "output");
    if (!err_out0.IsOk()) {
        res.errorCode = 4;
        res.errorMessage = err_out0.Message();
        std::cout << __LINE__ << " res.errorMessage: " << res.errorMessage << std::endl;
        return res;
    }
    tc::InferRequestedOutput* output542;
    tc::Error err_out542 = tc::InferRequestedOutput::Create(&output542, "542");
    if (!err_out542.IsOk()) {
        res.errorCode = 4;
        res.errorMessage = err_out542.Message();
        std::cout << __LINE__ << " err_out542.errorMessage: " << res.errorMessage << std::endl;
        return res;
    }

    std::shared_ptr<tc::InferRequestedOutput> output0_ptr;
    std::shared_ptr<tc::InferRequestedOutput> output542_ptr;
    output0_ptr.reset(output0);
    output542_ptr.reset(output542);

    std::string model_name = model_name_;
    std::string model_version = "";
    // The inference settings. Will be using default for now.
    tc::InferOptions options(model_name);
    options.model_version_ = model_version;
    options.priority_ = 1;

    std::vector<tc::InferInput*> inputs = {input0_ptr.get()};
    std::vector<const tc::InferRequestedOutput*> outputs = {output0_ptr.get(), output542_ptr.get()};

    tc::InferResult* results;
    tc::Headers http_headers;
    tc::Error err_infer = client->Infer(&results, options, inputs, outputs, http_headers);
    if (!err_infer.IsOk()) {
        res.errorCode = 5;
        res.errorMessage = err_infer.Message();
        // Infer error "unable to run model"
        std::cout << __LINE__ << " res.errorMessage: " << res.errorMessage << std::endl;
        return res;
    }
    std::shared_ptr<tc::InferResult> results_ptr;
    results_ptr.reset(results);

    float* f32Outputs;
    float* f32Outputs_segment;
    size_t recv_output0_byte_size;
    size_t recv_output542_byte_size;
    std::vector<int64_t> shape_output;
    std::vector<int64_t> shape_542;
    results_ptr->RawData("output", (const uint8_t**)&f32Outputs, &recv_output0_byte_size);
    tc::Error err_res542 = results_ptr->RawData("542", (const uint8_t**)&f32Outputs_segment, &recv_output542_byte_size);
    if (!err_res542.IsOk()) {
        res.errorCode = 5;
        // res.errorMessage = "Infer error:unable to Infer";
        res.errorMessage = err_res542.Message();
        // Infer error "unable to run model"
        std::cout << __LINE__ << " err_res542.errorMessage: " << res.errorMessage << std::endl;
    }

    results_ptr->Shape("output", &shape_output);
    tc::Error err_Shape542 = results_ptr->Shape("542", &shape_542);
    if (!err_Shape542.IsOk()) {
        res.errorCode = 5;
        res.errorMessage = err_Shape542.Message();
        std::cout << __LINE__ << " err_Shape542.errorMessage: " << res.errorMessage << std::endl;
    }

    int iOutputs_b = shape_output.at(0);
    int iOutputs_c = shape_output.at(1);
    int iOutputs_w = shape_output.at(2);
    int iOutputs_h = shape_output.at(3);

    int iSegment_w = shape_542.at(0);
    int iSegment_h = shape_542.at(1);
#endif

    vector<int> vec_segment;
    // 获取模型输出segment结果(向上取整)
    for (int i = 0; i < iSegment_w; i++) {
        for (int j = 0; j < iSegment_h; j++) {
            float data_ = f32Outputs_segment[i * iSegment_h + j];
            vec_segment.push_back(std::round(data_));
        }
    }
    array<int, 2> Segment_shape{iSegment_w, iSegment_h};

    vector<float> vec_regress;
    // 获取模型输出f32Outputs，并转换为WHC格式
    for (int i = 0; i < iOutputs_w; i++) {
        for (int j = 0; j < iOutputs_h; j++) {
            for (int c = 0; c < iOutputs_c; c++) {
                float data_ = f32Outputs[c * iOutputs_h * iOutputs_w + i * iOutputs_w + j];
                vec_regress.push_back(data_);
            }
        }
    }
    array<int, 4> regress_shape{1, iOutputs_w, iOutputs_h, iOutputs_c};

    result = flatByRegressWithClassiy_multiProcess(vec_segment, Segment_shape, vec_regress, regress_shape, m_srcImg);

    std::string res_base64 = MatToBase64(result);
    res.base64Result = res_base64;
    res.format = imgFormat;  // 图片format
    res.errorCode = -1;      // 假设成功时 errorCode 为 -1
    res.errorMessage.clear();

    return res;
}
