#include "epiceye.h"
#include "lodepng.h"

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

static bool saveAsPNG(int height, int width, int bitdepth, int channels, std::string filename, void* data) {
    lodepng::State state;
    if (channels == 3) {
        state.info_raw.colortype = LCT_RGB;
        state.info_png.color.colortype = LCT_RGB;
    } else if (channels == 1) {
        state.info_raw.colortype = LCT_GREY;
        state.info_png.color.colortype = LCT_GREY;
    }
    state.info_raw.bitdepth = bitdepth;
    state.info_png.color.bitdepth = bitdepth;

    std::vector<unsigned char> buffer;

    unsigned error = lodepng::encode(buffer, (unsigned char*)data, width, height, state);

    if (error) {
        std::cout << "encoder error " << error << ": " << lodepng_error_text(error) << std::endl;
        return false;
    }
    lodepng::save_file(buffer, filename);
    return true;
}

static bool saveDepthGrey(int height, int width, std::string filename, float* data) {
    float maxDepth = 0;
    float minDepth = 3.4e+38f; // MAXFLOAT
    for (int i = 0; i < width; i++) {
        for (int j = 0; j < height; j++) {
            int index = i * height + j;
            if (maxDepth < data[index]) {
                maxDepth = data[index];
                continue;
            }
            if (minDepth > data[index]) {
                minDepth = data[index];
            }
        }
    }
    float depthRange = maxDepth - minDepth;
    uint8_t* depthGreyBuffer = (uint8_t*)malloc(width * height);
    for (int i = 0; i < width; i++) {
        for (int j = 0; j < height; j++) {
            int index = i * height + j;
            depthGreyBuffer[index] = (uint8_t)(round((data[index] - minDepth) / depthRange * 255));
        }
    }
    saveAsPNG(height, width, 8, 1, filename, depthGreyBuffer);
    free(depthGreyBuffer);
    return true;
}

void savePointCloudAsPly(float* cloudMapBuffer, const std::string &savePath, const int width, const int height) {
    std::ofstream outfile(savePath);
    outfile << "ply\n"
            << "format ascii 1.0\n"
            << "obj_info EpicEye PLY PointCloud \n"
            << "obj_info num_cols " << width << "\n"
            << "obj_info num_rows " << height << "\n"
            << "element vertex " << width * height << "\n"
            << "property float x\n"
            << "property float y\n"
            << "property float z\n"
            << "end_header\n";
    for (int i = 0; i < width; i++) {
        for (int j = 0; j < height; j++) {
            int startIndex = (i * height + j) * 3;
            outfile << cloudMapBuffer[startIndex] << " ";
            outfile << cloudMapBuffer[startIndex + 1] << " ";
            outfile << cloudMapBuffer[startIndex + 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;
    }
    // 16 bit BGR image buffer -> 8 bit RGB image buffer
    uint8_t* imageBufferPtr8bit = (uint8_t*)malloc(info.width * info.height * 3);
    for (unsigned int i = 0; i < info.width; i++) {
        for (unsigned int j = 0; j < info.height; j++) {
            unsigned int currentIndex = (i * info.height + j) * 3;
            // BGR to RGB, color channel convert
            uint16_t tempValue = imageBufferPtr[currentIndex + 2];
            imageBufferPtr[currentIndex + 2] = imageBufferPtr[currentIndex + 0];
            imageBufferPtr[currentIndex + 0] = tempValue;
            // BGR -> RGB
            imageBufferPtr8bit[currentIndex] = (uint8_t)(imageBufferPtr[currentIndex] / 4 & 0xFF);
            imageBufferPtr8bit[currentIndex + 1] = (uint8_t)(imageBufferPtr[currentIndex + 1] / 4 & 0xFF);
            imageBufferPtr8bit[currentIndex + 2] = (uint8_t)(imageBufferPtr[currentIndex + 2] / 4 & 0xFF);
        }
    }
    saveAsPNG(info.height, info.width, 8, 3, "image8bit.png", imageBufferPtr8bit);
    free(imageBufferPtr8bit);
    free(imageBufferPtr);

    // 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 << "getImage failed!" << std::endl;
        free(depthBufferPtr);
        return 0;
    }

    saveDepthGrey(info.height, info.width, "depthGrey.png", depthBufferPtr);
    free(depthBufferPtr);

    // 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;
    }
    savePointCloudAsPly(pointCloudBufferPtr, "pointcloud.ply", info.width, info.height);

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