#include "epiceye.h"
#include "opencv2/core.hpp"
#include "opencv2/imgproc.hpp"
#include <opencv2/highgui.hpp>

#include <iostream>
#include <fstream>
#include <string>
#include <iomanip>
#include <stdlib.h>


cv::Mat imageToCVMat(void* data, uint32_t width, uint32_t height, bool convertToU8C3 = true, uint32_t channel = 1) {
    if (channel == 1) {
        cv::Mat mat(cv::Size(width, height), CV_16UC1, data);
        if (convertToU8C3) {
            mat.convertTo(mat, CV_8UC1, 0.25);
            cv::cvtColor(mat, mat, cv::COLOR_GRAY2BGR);
        }
        return mat.clone();
    }
    cv::Mat mat(cv::Size(width, height), CV_16UC3, data);
    if (convertToU8C3) {
        mat.convertTo(mat, CV_8UC3, 0.25);
    }
    return mat.clone();
}

cv::Mat depthToCVMat(void* data, uint32_t width, uint32_t height) {
    return cv::Mat(cv::Size(width, height), CV_32FC1, data).clone();
}

cv::Mat pointCloudToCVMat(void* data, uint32_t width, uint32_t height) {
    return cv::Mat(cv::Size(width, height), CV_32FC3, data).clone();
}

void savePointCloudAsPly(cv::Mat cloudMap, const std::string &savePath) {
    std::ofstream outfile(savePath);
    outfile << "ply\n"
            << "format ascii 1.0\n"
            << "obj_info EpicEye PLY PointCloud \n"
            << "obj_info num_cols " << cloudMap.cols << "\n"
            << "obj_info num_rows " << cloudMap.rows << "\n"
            << "element vertex " << cloudMap.rows * cloudMap.cols << "\n"
            << "property float x\n"
            << "property float y\n"
            << "property float z\n"
            << "end_header\n";
    for (int i = 0; i < cloudMap.rows; i++) {
        for (int j = 0; j < cloudMap.cols; j++) {
            cv::Vec3f point = cloudMap.at<cv::Vec3f>(i, j);
            outfile << point[0] << " ";
            outfile << point[1] << " ";
            outfile << point[2] << " ";
            outfile << "\n";
        }
    }
    outfile.close();
}

int main(int argc, char** argv) {
    std::string ip;
    if (argc > 1) {
        ip = argv[1];
    } else {
        std::cout << "---------------search EpicEye camera---------------" << std::endl;
        std::vector<TFTech::EpicEyeInfo> cameraList;
        if (TFTech::EpicEye::searchCamera(cameraList)) {
            std::cout << "Camera found: " << cameraList.size() << std::endl;
            for (int i = 0; i < cameraList.size(); i++) {
                std::cout << std::right << std::setw(3) << i << ": "
                    << std::left << std::setw(20) << std::setfill(' ') << cameraList[i].IP
                    << std::left << std::setw(30) << cameraList[i].SN << std::endl;
            }
            ip = cameraList[0].IP;
        } else {
            std::cout << "Camera not found!" << std::endl;
            return 0;
        }
    }

    std::cout << std::endl;

    std::cout << "---------------get " << ip << " EpicEyeInfo-------------- - " << std::endl;
    TFTech::EpicEyeInfo info;
    if (!TFTech::EpicEye::getInfo(ip, info)) {
        std::cout << "getInfo failed!" << std::endl;
        return 0;
    }
    std::cout << "SN: " << info.SN << std::endl;
    std::cout << "IP: " << info.IP << std::endl;
    std::cout << "model: " << info.model << std::endl;
    std::cout << "alias: " << info.alias << std::endl;
    std::cout << "resolution: " << info.width << "x" << info.height << std::endl;

    // trigger one frame
    std::cout << "---------------TriggerFrame---------------" << std::endl;
    std::string frameID;
    bool requestPointCloud = true;
    if (!TFTech::EpicEye::triggerFrame(ip, frameID, requestPointCloud)) {
        std::cout << "triggerFrame failed!" << std::endl;
        return 0;
    }
    std::cout << "frameID: " << frameID << std::endl;

    // get image, image channel order -> BGR
    std::cout << "---------------get image---------------" << std::endl;
    uint16_t* imageBufferPtr = (uint16_t*)malloc(info.width * info.height * 3 * sizeof(uint16_t));
    int pixelByteSize = 0;
    if (!TFTech::EpicEye::getImage(ip, frameID, imageBufferPtr, pixelByteSize)) {
        std::cout << "getImage failed!" << std::endl;
        free(imageBufferPtr);
        return 0;
    }
    cv::Mat image = imageToCVMat(imageBufferPtr, info.width, info.height, true, pixelByteSize / 2);
    free(imageBufferPtr);
    cv::namedWindow("image", cv::WINDOW_KEEPRATIO);
    cv::resizeWindow("image", 960, 800);
    cv::imshow("image", image);
    cv::waitKey();
    cv::imwrite("image.png", image);

    // get depth
    std::cout << "---------------get depth---------------" << std::endl;
    float* depthBufferPtr = (float*)malloc(info.width * info.height * sizeof(float));
    if (!TFTech::EpicEye::getDepth(ip, frameID, depthBufferPtr)) {
        std::cout << "getDepth failed!" << std::endl;
        free(depthBufferPtr);
        return 0;
    }
    cv::Mat depth = depthToCVMat(depthBufferPtr, info.width, info.height);
    free(depthBufferPtr);

    depth.convertTo(depth, CV_8UC1, 1 / 10.0);
    cv::applyColorMap(depth, depth, cv::COLORMAP_JET);
    cv::namedWindow("depth", cv::WINDOW_KEEPRATIO);
    cv::resizeWindow("depth", 960, 800);
    cv::imshow("depth", depth);
    cv::waitKey();
    cv::imwrite("depth.png", depth);

    // get point cloud
    std::cout << "---------------get point cloud---------------" << std::endl;
    float* pointCloudBufferPtr = (float*)malloc(info.width * info.height * sizeof(float) * 3);
    if (!TFTech::EpicEye::getPointCloud(ip, frameID, pointCloudBufferPtr)) {
        std::cout << "getPointCloud failed!" << std::endl;
        free(pointCloudBufferPtr);
        return 0;
    }
    cv::Mat pointCloud = pointCloudToCVMat(pointCloudBufferPtr, info.width, info.height);
    free(pointCloudBufferPtr);

    savePointCloudAsPly(pointCloud, "pointcloud.ply");

    std::cout << std::endl;
#ifdef _WIN32
    system("pause");
#endif
    return 0;
}