#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 

#include 
#include 
#include 

#include 
#include 
#include 

#include 

#include "logger.h"
#include "trace.hpp"

using input_info_t = std::vector;

Logger gLogger;

inline int64_t volume(const nvinfer1::Dims &d)
{
    int64_t v = 1;
    for (int i = 0; i < d.nbDims; i++) { v *= d.d[i]; }
    return v;
}

inline size_t elementSize(nvinfer1::DataType t)
{
    switch (t) {
    // TODO: check nvinfer1 version
    // case nvinfer1::DataType::kINT32:
    //     return 4;
    case nvinfer1::DataType::kFLOAT:
        return 4;
    case nvinfer1::DataType::kHALF:
        return 2;
    case nvinfer1::DataType::kINT8:
        return 1;
    }
    assert(0);
    return 0;
}

std::string to_string(const nvinfer1::Dims &d)
{
    std::string s{"("};
    if (d.nbDims != 0) {
        for (int64_t i = 0; i < d.nbDims; i++)
            (s += std::to_string(d.d[i])) += ", ";
        s.pop_back();
        s.pop_back();
    }
    return s + ")";
}

std::string to_string(const nvinfer1::DataType dtype)
{
    return std::to_string(int(dtype));
}

template  struct destroy_deleter {
    void operator()(T *ptr) { ptr->destroy(); }
};

template 
using destroy_ptr = std::unique_ptr;

nvinfer1::ICudaEngine *loadModelAndCreateEngine(const char *uffFile,
                                                int max_batch_size,
                                                nvuffparser::IUffParser *parser,
                                                bool use_f16)
{
    destroy_ptr builder(
        nvinfer1::createInferBuilder(gLogger));
    destroy_ptr network(builder->createNetwork());

    if (use_f16) {
        if (!parser->parse(uffFile, *network, nvinfer1::DataType::kHALF)) {
            return nullptr;
        }
        // builder->setFp16Mode(true);
        builder->setHalf2Mode(true);  // For older version of tensorRT
    } else {
        if (!parser->parse(uffFile, *network, nvinfer1::DataType::kFLOAT)) {
            return nullptr;
        }
    }
    builder->setMaxBatchSize(max_batch_size);
    return builder->buildCudaEngine(*network);
}

nvinfer1::ICudaEngine *
create_engine(const std::string &model_file, const input_info_t &input_info,
              const std::vector &output_names, int max_batch_size,
              bool use_f16)
{
    TRACE_SCOPE(__func__);
    destroy_ptr parser(nvuffparser::createUffParser());
    for (const auto &info : input_info) {
        const auto dims = info.second;
        parser->registerInput(
            info.first.c_str(),
            // Always provide your dimensions in CHW even if your
            // network input was in HWC in yout original framework.
            nvinfer1::DimsCHW(dims[0], dims[1], dims[2]),
            nvuffparser::UffInputOrder::kNCHW  //
        );
    }
    for (auto &name : output_names) { parser->registerOutput(name.c_str()); }
    auto engine = loadModelAndCreateEngine(model_file.c_str(), max_batch_size,
                                           parser.get(), use_f16);
    if (!engine) {
        gLogger.log(nvinfer1::ILogger::Severity::kERROR,
                    "failed to created engine");
        exit(1);
    }
    return engine;
}

class uff_runner_impl : public pose_detection_runner
{
  public:
    uff_runner_impl(const std::string &model_file,
                    const input_info_t &input_info,
                    const std::vector &output_names,
                    int max_batch_size, bool use_f16);
    ~uff_runner_impl() override;

    void operator()(const std::vector &inputs,
                    const std::vector &outputs,
                    int batch_size) override;

  private:
    const int max_batch_size;

    destroy_ptr engine_;

    using cuda_buffer_t = ttl::cuda_tensor;  // [batch_size, data_size]
    std::vector buffers_;

    void createBuffers_(int batch_size);
};

uff_runner_impl::uff_runner_impl(const std::string &model_file,
                                 const input_info_t &input_info,
                                 const std::vector &output_names,
                                 int max_batch_size, bool use_f16)
    : max_batch_size(max_batch_size),
      engine_(create_engine(model_file, input_info, output_names,
                            max_batch_size, use_f16))
{
    createBuffers_(max_batch_size);
}

uff_runner_impl::~uff_runner_impl() { nvuffparser::shutdownProtobufLibrary(); }

void uff_runner_impl::createBuffers_(int batch_size)
{
    TRACE_SCOPE(__func__);
    for (auto i : ttl::range(engine_->getNbBindings())) {
        const nvinfer1::Dims dims = engine_->getBindingDimensions(i);
        const nvinfer1::DataType dtype = engine_->getBindingDataType(i);
        const std::string name(engine_->getBindingName(i));
        std::cout