#include <cassert>
#include <fstream>
#include <sstream>
#include <iostream>
#include <cmath>
#include <sys/stat.h>
#include <cmath>
#include <time.h>
#include <cuda_runtime_api.h>
#include <cudnn.h>
#include <cublas_v2.h>
#include <memory>
#include <cstring>
#include <algorithm>
#include <opencv2/opencv.hpp>

#include "NvInfer.h"
#include "NvCaffeParser.h"


using namespace nvinfer1;
using namespace nvcaffeparser1;
using namespace std;

#define CHECK(status)                                   \
{                                                       \
    if (status != 0)                                    \
    {                                                   \
        std::cout << "Cuda failure: " << status;        \
        abort();                                        \
    }                                                   \
}

// stuff we know about the network and the caffe input/output blobs
static const int INPUT_C = 3;
static const int INPUT_H = 160;
static const int INPUT_W = 320;


const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "seg-prob-region";

// Logger for GIE info/warning/errors
class Logger : public ILogger
{
    void log(Severity severity, const char* msg) override
    {
        // suppress info-level messages
        if (severity != Severity::kINFO || false)
            std::cout << msg << std::endl;
    }
} gLogger;


void caffeToGIEModel(const std::string& deployFile,                 // name for caffe prototxt
                     const std::string& modelFile,                  // name for model
                     const std::vector<std::string>& outputs,       // network outputs
                     unsigned int maxBatchSize,                     // batch size - NB must be at least as large as the batch we want to run with)
                     nvcaffeparser1::IPluginFactory* pluginFactory, // factory for plugin layers
                     IHostMemory *&gieModelStream)                  // output stream for the GIE model
{
    // create the builder
    IBuilder* builder = createInferBuilder(gLogger);

    // parse the caffe model to populate the network, then set the outputs
    INetworkDefinition* network = builder->createNetwork();
    ICaffeParser* parser = createCaffeParser();
    parser->setPluginFactory(pluginFactory);


    const IBlobNameToTensor* blobNameToTensor = parser->parse(deployFile.c_str(),
                                                              modelFile.c_str(),
                                                              *network,
                                                              DataType::kFLOAT);

    // specify which tensors are outputs
    for (auto& s : outputs)
    {
        if (blobNameToTensor->find(s.c_str()) == nullptr)
        {
            std::cout << "could not find output blob " << s << std::endl;
            return ;
        }
        network->markOutput(*blobNameToTensor->find(s.c_str()));
    }

    // Build the engine
    builder->setMaxBatchSize(maxBatchSize);
    builder->setMaxWorkspaceSize( 1000 * (1 << 20));
    //builder->setHalf2Mode(fp16);

    ICudaEngine* engine = builder->buildCudaEngine(*network);
    assert(engine);

    // we don't need the network any more, and we can destroy the parser
    network->destroy();
    parser->destroy();

    // serialize the engine, then close everything down
    gieModelStream = engine->serialize();

    engine->destroy();
    builder->destroy();
    shutdownProtobufLibrary();

}


class ELU_Plugin: public IPlugin
{
public:


    ELU_Plugin ()
    {

    }

    ELU_Plugin(const Weights *weights, int nbWeights, int nbOutputChannels): mNbOutputChannels(nbOutputChannels)
    {

    }
    // create the plugin at runtime from a byte stream
    ELU_Plugin(const void* data, size_t length)
    {

    }

    ~ELU_Plugin()
    {

    }

    int getNbOutputs() const override
    {
        return 1;
    }

    Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override
    {
        cout << "ELU Shape " << inputs[0].d[0] << ", " << inputs[0].d[1] << ", " << inputs[0].d[2] <<endl;
        return DimsCHW(inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]);
    }

    void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override
    {

    }

    int initialize() override
    {
        CHECK(cudnnCreate(&mCudnn));                            // initialize cudnn and cublas
        CHECK(cublasCreate(&mCublas));
        CHECK(cudnnCreateTensorDescriptor(&mSrcDescriptor));    // create cudnn tensor descriptors we need for bias addition
        CHECK(cudnnCreateTensorDescriptor(&mDstDescriptor));
        CHECK(cudnnCreateActivationDescriptor(&mActDescriptor));
        return 0;
    }

    void terminate() override
    {
        CHECK(cublasDestroy(mCublas));
        CHECK(cudnnDestroy(mCudnn));
    }

    size_t getWorkspaceSize(int maxBatchSize) const override
    {
        return 0;
    }


    int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override
    {

        printf("Enqueue start \n");
        const float kONE = 1.0f, kZERO = 0.0f;
        cublasSetStream(mCublas, stream);
        cudnnSetStream(mCudnn, stream);
        CHECK(cudnnSetActivationDescriptor(mActDescriptor, CUDNN_ACTIVATION_ELU, CUDNN_PROPAGATE_NAN, kONE));
        CHECK(cudnnActivationForward(mCudnn, mActDescriptor, &kONE, mSrcDescriptor, inputs, &kZERO, mDstDescriptor, outputs));

        return 0;
    }

    size_t getSerializationSize() override
    {
        return 0;
    }

    void serialize(void* buffer) override
    {

    }
private:
    int mNbOutputChannels, mNbInputChannels;
    cudnnHandle_t mCudnn;
    cublasHandle_t mCublas;
    cudnnTensorDescriptor_t mSrcDescriptor, mDstDescriptor;
    cudnnActivationDescriptor_t mActDescriptor;
};

// integration for serialization
class PluginFactory : public nvinfer1::IPluginFactory, public nvcaffeparser1::IPluginFactory
{
public:

    // caffe parser plugin implementation
    bool isPlugin(const char* name) override
    {
        string s (name);

        cout << "Layer name: " << s << endl;

        if (s.substr(0, 3) == "elu")
            return 1;
        else
            return 0;
    }

    virtual nvinfer1::IPlugin* createPlugin(const char* layerName, const nvinfer1::Weights* weights, int nbWeights) override
    {
        // there's no way to pass parameters through from the model definition, so we have to define it here explicitly
        mPlugin = std::unique_ptr<ELU_Plugin>(new ELU_Plugin());
        return mPlugin.get();
    }

    // deserialization plugin implementation
    IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) override
    {
        mPlugin = std::unique_ptr<ELU_Plugin>(new ELU_Plugin(serialData, serialLength));
        return mPlugin.get();
    }

    // the application has to destroy the plugin when it knows it's safe to do so
    void destroyPlugin()
    {
        mPlugin.release();
    }

    std::unique_ptr<ELU_Plugin> mPlugin{ nullptr };
};

int main(int argc, char** argv)
{
    // create a GIE model from the caffe model and serialize it to a stream
    PluginFactory pluginFactory;
    IHostMemory *gieModelStream{ nullptr };
    caffeToGIEModel("/home/models/u-net0/deploy_short.prototxt", "/home/models/u-net0/model.caffemodel", std::vector < std::string > { OUTPUT_BLOB_NAME }, 1, &pluginFactory, gieModelStream);
    pluginFactory.destroyPlugin();


    return 0;
}