//Import neccessary libs
#include "argsParser.h"
#include "buffers.h"
#include "common.h"
#include "logger.h"
#include "parserOnnxConfig.h"

#include "NvInfer.h"
#include "MAC_PythonInterface.h"
#include <cuda_runtime_api.h>
#include <opencv2/core/matx.hpp>
# include <opencv2/opencv.hpp>

#include <cstdlib>
#include <fstream>
#include <iostream>
#include <sstream>
#include <chrono>
#include <bits/stdc++.h>
#include <dlfcn.h>

# include <stdio.h>

using namespace cv;
using namespace std;
using namespace std::chrono;

const std::string gSampleName = "TensorRT.cpp_midas_hq_inference";

//Normalisation function
struct Normalizer
{
    float mean_r, mean_g, mean_b;
    float std_r, std_g, std_b;
    cv::Mat out;
    Normalizer()
    {
        float scale = 255.0f;

        mean_r = 0.485f*scale;
        mean_g = 0.456f*scale;
        mean_b = 0.406f*scale;
        
        std_r = 0.229f*scale;
        std_g = 0.224f*scale;
        std_b = 0.225f*scale;
    }
    
    cv::Mat Process_MIDASHQ(cv::Mat input)
    {
        if(out.cols != input.cols || out.rows != input.rows)
        {
            out = cv::Mat::zeros(input.rows, input.cols, CV_32FC3);
        }
        
        int jump = out.step*out.rows/3;
        unsigned char* src = input.data;
        float* dst1 = (float*)out.data;
        float* dst2 = (float*)(out.data + jump);
        float* dst3 = (float*)(out.data + jump*2);
        
        for(int i = 0; i < input.rows; i++)
        {
            unsigned char* src = input.data + i*input.step;
            float* dst = (float*)(out.data + i*out.step);
            for(int j = 0; j < input.cols; j++)
            {
                float b = *src++;
                float g = *src++;
                float r = *src++;

                *dst1++ = (r - mean_r)/std_r;
                *dst2++ = (g - mean_g)/std_g;
                *dst3++ = (b - mean_b)/std_b;
            }
        }
        return out;
    }
    
};

//Class declaration
class SampleInferenceMIDASHQ
{
    template <typename T>
    using SampleUniquePtr = std::unique_ptr<T, samplesCommon::InferDeleter>;

public:
    SampleInferenceMIDASHQ(const samplesCommon::OnnxSampleParams& params)
        : mParams(params)
        , mEngine_midas_hq(nullptr)
    {
    }
    
    bool build();

    vector<cv::Mat> infer(vector<cv::Mat> &inputs_fin);
    
    void destEngine();

private:
    samplesCommon::OnnxSampleParams mParams;

    std::shared_ptr<nvinfer1::IExecutionContext> context_midas_hq;   

    std::shared_ptr<nvinfer1::ICudaEngine> mEngine_midas_hq; 
    
    bool processInput(const samplesCommon::BufferManager& buffers, vector<cv::Mat>& input);

    vector<cv::Mat> processOutput(const samplesCommon::BufferManager& buffers);
};

//Function that loads and deserialises the engine
bool SampleInferenceMIDASHQ::build()
{
    std::vector<char> trtModelStream_;
    size_t size{ 0 };

    std::ifstream file("midas_384.engine", std::ios::binary);

    if (file.good())
    {
       
        file.seekg(0, file.end);
        size = file.tellg();
        file.seekg(0, file.beg);
        trtModelStream_.resize(size);
        file.read(trtModelStream_.data(), size);
        file.close();
    }

    IRuntime* runtime = createInferRuntime(sample::gLogger);
    
    mEngine_midas_hq = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(trtModelStream_.data(), size, nullptr), samplesCommon::InferDeleter());
    
    if (!mEngine_midas_hq)
    {
        return false;
    }

    //printf("Loaded MIDAS HQ deserialized engine\n");
    
    context_midas_hq = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine_midas_hq->createExecutionContext());
    nvinfer1::Dims4 input_dimensions(4,3,384,1120);
    context_midas_hq->setBindingDimensions(0,input_dimensions);
   
    return true;
}

//Function that performs inferenc
vector<cv::Mat> SampleInferenceMIDASHQ::infer(vector<cv::Mat> &inputs_fin)
{
    cout << "in infer" << endl;

    auto t1 = high_resolution_clock::now();
    samplesCommon::BufferManager buffers(mEngine_midas_hq);
    cout << "created buffers" << endl;
    auto t2 = high_resolution_clock::now();

    bool status_processInput = processInput(buffers,inputs_fin);
    cout << "processInput complete" << endl;
    auto t3 = high_resolution_clock::now();
    
    buffers.copyInputToDevice();
    cout << "copy to buffers" << endl;
    auto t4 = high_resolution_clock::now();

    
    bool status_inference = context_midas_hq->executeV2(buffers.getDeviceBindings().data());
    cout << "inference done" << endl;
    auto t5 = high_resolution_clock::now();
    
    buffers.copyOutputToHost();
    cout << "copy from buffers" << endl;
    auto t6 = high_resolution_clock::now();

    vector<cv::Mat> output_fin = processOutput(buffers);
    cout << "processOutput" << endl;
    auto t7 = high_resolution_clock::now();

    auto d1 = duration_cast<milliseconds>(t2 - t1);
    auto d2 = duration_cast<milliseconds>(t3 - t2);
    auto d3 = duration_cast<milliseconds>(t4 - t3);
    auto d4 = duration_cast<milliseconds>(t5 - t4);
    auto d5 = duration_cast<milliseconds>(t6 - t5);
    auto d6 = duration_cast<milliseconds>(t7 - t6);
    auto d7 = duration_cast<milliseconds>(t7 - t1);

    printf("buffer_declared: %f\n",(float)d1.count());
    printf("preprocessing  : %f\n",(float)d2.count());
    printf("copy_to_buffer : %f\n",(float)d3.count());
    printf("inference_done : %f\n",(float)d4.count());
    printf("copy_from_bufer: %f\n",(float)d5.count());
    printf("processOutput  : %f\n",(float)d6.count());
    printf("complete_infer : %f\n",(float)d7.count());

    return output_fin;

}

int h_orig_hq, w_orig_hq;

//Function that pre-processes the input
bool SampleInferenceMIDASHQ::processInput(const samplesCommon::BufferManager& buffers, vector<cv::Mat>& input)
{
    int batch = 4;
    float* hostDataBuffer = static_cast<float*>(buffers.getHostBuffer("INPUTS"));
    for (int batch_i = 0; batch_i < batch; batch_i++)
    {
        h_orig_hq = input[batch_i].rows;
        w_orig_hq = input[batch_i].cols;    
        
        cv::Mat input1;
        cv::resize(input[batch_i], input1, cv::Size(1120, 384), 0, 0, cv::INTER_CUBIC);
        
        Normalizer normalizer;
        cv::Mat refined = normalizer.Process_MIDASHQ(input1);

        cv::Mat linear_refined = refined.reshape(1,refined.total()*refined.channels());
        for (int i = (int)(linear_refined.rows)*batch_i; i < (batch_i + 1)*(int)(linear_refined.rows); i++)
        {
            hostDataBuffer[i] = (float)linear_refined.at<float>(cv::Point(i-(int)(linear_refined.rows)*batch_i,0));
        }
    }
    return true;
}

//Function that post-processes the output
vector<cv::Mat> SampleInferenceMIDASHQ::processOutput(const samplesCommon::BufferManager& buffers)
{
    int batch = 4;
    vector<cv::Mat> out;
    float* output = static_cast<float*>(buffers.getHostBuffer("OUTPUTS"));
    for(int batch_i=0; batch_i < batch; batch_i++)
    {
        cv::Mat outputs = cv::Mat(384, 1120, CV_32FC1, output+batch_i*(384*1120));   
        out.push_back(6000.0f*(1.0f/outputs));
    }    
    return out;
}

samplesCommon::OnnxSampleParams initializeSampleParams(const samplesCommon::Args& args)
{
    samplesCommon::OnnxSampleParams params;
    return params;
}

samplesCommon::Args args;
SampleInferenceMIDASHQ sample_midas_hq(initializeSampleParams(args));

bool MAC_MIDAS_HQ::Init()
{


    return sample_midas_hq.build();
}

int MAC_MIDAS_HQ::Process(vector<cv::Mat>* inputs, int inSize, vector<cv::Mat>* outs)
{
    vector<cv::Mat> input = inputs[0];
    cout << input[0].rows << input[0].cols << endl;
    vector<cv::Mat> outputs = sample_midas_hq.infer(input);
    *outs = outputs;
    return 1;
}

//Function that destroys engine and context
void  SampleInferenceMIDASHQ::destEngine()
{
    context_midas_hq = nullptr;
    mEngine_midas_hq = nullptr;
        
}

void MAC_MIDAS_HQ::Destroy()
{
    sample_midas_hq.destEngine();
}

vector<cv::Mat> MAC_MIDAS_HQ::ComputeDepthSDE(vector<cv::Mat> &input)
{
    vector<cv::Mat> output;
    int count = Process(&input, 1, &output);
    return output;
}

static MAC_MIDAS_HQ g_Midas_hq;

bool InitMidas_HQ()
{
    g_Midas_hq.Init();
    return true;
}

vector<cv::Mat> ProcessMidas_HQ(vector<cv::Mat> input)
{


    return g_Midas_hq.ComputeDepthSDE(input);
}

void CloseMidas_HQ()
{
    g_Midas_hq.Destroy();
}

int main()
{
    //MAC_MIDAS_HQ midas;
    bool ack = InitMidas_HQ();
    int key = ' ';
    int batch = 4;
    char filename[300];
    for(int i = 351; i < 708-batch+1; i+=batch){
        auto t1 = high_resolution_clock::now();
        vector<cv::Mat> images;
        for (int j = 0; j < batch; j++)
        {
            sprintf(filename,"/media/31A079936F39FBF9/romil/Imgs/img_%d.png",i+j);
            cout << "reading: " << filename << endl;
            images.push_back(imread(filename));
        }
        vector<cv::Mat> output;
        output = ProcessMidas_HQ(images);
        auto t2 = high_resolution_clock::now();
        auto d1 = duration_cast<milliseconds>(t2 - t1);
        printf("with_imread    : %f\n",(float)d1.count());
    }
    CloseMidas_HQ();
    return 0;
}