#include <NvInfer.h>
#include <NvInferPlugin.h>
#include <fstream>
#include <iostream>
#include <map>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

namespace {
struct Logger : nvinfer1::ILogger {
    void log(Severity severity, const char* message) noexcept override {
        if (severity <= Severity::kWARNING) std::cerr << message << '\n';
    }
};
void require(bool valid) { if (!valid) throw std::runtime_error("ENGINE_CONTRACT"); }
bool shape(nvinfer1::Dims dims, const std::vector<int>& expected) {
    if (dims.nbDims != static_cast<int>(expected.size())) return false;
    for (int i=0; i<dims.nbDims; ++i) if (dims.d[i] != expected[i]) return false;
    return true;
}
}

int main(int argc, char** argv) {
    try {
        require(argc==4 && NV_TENSORRT_MAJOR==10 && NV_TENSORRT_MINOR==16);
        const std::string key=argv[2]; const int batch=std::stoi(argv[3]);
        require((key=="medium" || key=="pose" || key=="fire") && batch>=1 && batch<=8);
        const bool pose=key=="pose";
        std::map<std::string,std::vector<int>> expected={
            {"input",{-1,3,576,576}}, {"dets",{-1,pose ? 100 : 300,4}},
            {"labels",{-1,pose ? 100 : 300,pose ? 2 : key=="fire" ? 3 : 91}}};
        if (pose) expected["keypoints"]={-1,100,34,8};
        std::ifstream file(argv[1],std::ios::binary); require(file.good());
        std::vector<char> bytes((std::istreambuf_iterator<char>(file)),{});
        Logger logger; initLibNvInferPlugins(&logger,"");
        std::unique_ptr<nvinfer1::IRuntime> runtime(nvinfer1::createInferRuntime(logger));
        require(bool(runtime));
        std::unique_ptr<nvinfer1::ICudaEngine> engine(runtime->deserializeCudaEngine(bytes.data(),bytes.size()));
        require(engine && engine->getNbIOTensors()==static_cast<int>(expected.size())
                && engine->getNbOptimizationProfiles()==1);
        for (int i=0; i<engine->getNbIOTensors(); ++i) {
            const std::string name=engine->getIOTensorName(i);
            require(expected.count(name)==1);
            require(shape(engine->getTensorShape(name.c_str()),expected.at(name)));
            const bool input=name=="input";
            require(engine->getTensorDataType(name.c_str())==(input ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT));
            require(engine->getTensorIOMode(name.c_str())==(input ? nvinfer1::TensorIOMode::kINPUT : nvinfer1::TensorIOMode::kOUTPUT));
        }
        const nvinfer1::OptProfileSelector selectors[]={nvinfer1::OptProfileSelector::kMIN,
            nvinfer1::OptProfileSelector::kOPT,nvinfer1::OptProfileSelector::kMAX};
        for (int i=0; i<3; ++i)
            require(shape(engine->getProfileShape("input",0,selectors[i]),{i==0 ? 1 : batch,3,576,576}));
        std::cout << "{\"contract\":\"PASS\",\"backend\":\"TensorRT_CXX\",\"tensorrt_version\":\""
                  << NV_TENSORRT_MAJOR << '.' << NV_TENSORRT_MINOR << '.' << NV_TENSORRT_PATCH
                  << "\",\"batch_bounds\":[[1,3,576,576],[" << batch << ",3,576,576],["
                  << batch << ",3,576,576]]}" << std::endl;
        return 0;
    } catch (const std::exception&) {
        std::cerr << "ENGINE_CONTRACT_FAILED\n";
        return 2;
    }
}