/*
 * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "BatchStream.h"
#include "EntropyCalibrator.h"
#include "argsParser.h"
#include "buffers.h"
#include "common.h"
#include "logger.h"
#include "parserOnnxConfig.h"
#include "kernel.h"
#include "NvInfer.h"
#include <cuda_runtime_api.h>
#include <random>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <dirent.h>
using namespace std;
//using namespace cv;
using namespace nvonnxparser;
using namespace nvinfer1;

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

static const int INPUT_H = 24;
static const int INPUT_W = 94;
/*
Developed by Nyan.
Plugin creation requirements:
(1)Try to understand plugin process looking at tensorflow code, for example sparsetodence conversion, 
need to look at TF's sparsetodense code (/home/xavier/Desktop/tensorflow-master/tensorflow/core/util/ctc/ctc_decoder.h)
(2)Then use base methods from IPluginV2IOExt and IPluginCreator as shown in the following sample
(3)Important thing is to understand how to process in CUDA (some may not be necessary to process on CUDA)
(4)Just need to make sure number of inputs and number of outputs should be same as Tensorflow API's requirements
(5)Check the model graph with NETRON. SparseToDense has Four key inputs. They are

         Tensorflow' API arguments                          Inputs to CTCGreedyDecoder node (Can check from PB or ONNX graph using netron)               
const typename CTCDecoder<T>::SequenceLength& seq_len                           sequence_length(Fill) 
const std::vector<typename CTCDecoder<T>::Input>& input                         inputs(Transpose)


                                                            Outputs are
                                                             scores
                                                             std::vector<typename CTCDecoder<T>::Output>* output

Input/Output shapes can be seen from pb graph by printing all layer input and output
output = graph.get_tensor_by_name("import/d_predictions:0")
x = sess.run(output, feed_dict={input: [img]})   
print(x)
print(x.shape) 
as shown in the Input_output_layers_size 

For the graph variation, ONNX-Graphsurgeon can be used.
In the latest graph, SparseToDense node was removed using graphsurgeon.                                  
 */

////////////////////////////////////////////////////////////////////////////////////////////////////
/*
TensorRT accepts dynamic batch size. Dynamic means every iteration runs, batch size can be changed.
sampleDynamicReshape.cpp shows how to reshape batch or image size dynamically for every iteration.
For that, we need to have fixed model size 10,24,94,3 with max batch size.
buildPredictionEngine's optimized profile and calibration profile needs to be same shape as input.
buildPreprocessorEngine's optimization profile needs to be dynamic from min to max.
If there are plugins, plugin's shape needs to be max shape.
*/
///////////////////////////////////////////////////////////////////////////////////////////////////	

class CTCGreedyDecoder : public IPluginV2DynamicExt
{
public:
    //CTCGreedyDecoder has two inputs (import/transpose:0 (88, 1, 43), import/Fill:0 (1,)) 
    //and two outputs (import/CTCGreedyDecoder:0 (7, 2), import/ToInt32:0 (7,))
    CTCGreedyDecoder(const PluginFieldCollection& fc)
    {        
        (void) fc;
    }
    //data represents the class member variables data serialized in 
    //void serialize(void* buffer) const override
    //length is the length of data in serialization
    //Need to deserialize in this Constructor method    
    CTCGreedyDecoder(const void* data, size_t length)
    {
        //cout << "length " << length << endl;
        const char* d = static_cast<const char*>(data);
        const char* const a = d;
        assert(d == a + length);
    }

    // It makes no sense to construct CTCGreedyDecoder without arguments.
    CTCGreedyDecoder() = delete;

    virtual ~CTCGreedyDecoder() {}

public:
    int getNbOutputs() const override
    {
        return 1;//it has one output
    }

    DimsExprs getOutputDimensions(int outputIndex, const DimsExprs* inputs, int nbInputs, IExprBuilder& exprBuilder) override
    {
        //cout << "nbInputs from getOutputDimensions " << nbInputs << endl;
        //cout << "inputs's DimsExprs nbDims from getOutputDimensions " << inputs[0].nbDims << endl;
        //cout << "inputs's DimsExprs dimensions from getOutputDimensions " << inputs[0].d[0]<< " " << inputs[0].d[1] << " " << inputs[0].d[2] << endl;
        nvinfer1::DimsExprs output;
        output.nbDims=2;
        output.d[0] = inputs[0].d[1];
        output.d[1] = exprBuilder.constant(20);
        return output;
    }

    int initialize() override
    {        
        return 0;
    }

    void terminate() override
    {
        //To release memory
    }

    size_t getWorkspaceSize(const PluginTensorDesc* inputs, int nbInputs, const PluginTensorDesc* outputs,int 	nbOutputs) const override
    {
        return 0;
    }
   
    int enqueue(const PluginTensorDesc *inputDesc, const PluginTensorDesc *outputDesc, const void *const *inputs, void *const *outputs, void *workspace,cudaStream_t stream) override
    {   
        //cout << "inputDesc[0].dims from enqueue " << inputDesc[0].dims << endl;
        Dims mInputDims = inputDesc[0].dims;
        //cout << "mInputDims from enqueue " << mInputDims.d[0]<< " " << mInputDims.d[1] << " " << mInputDims.d[2] << endl;
        //cout << "reach to enqueue " << endl;
        int rows = mInputDims.d[0];
        int batch = mInputDims.d[1];   
        int widths = mInputDims.d[2];   
        Dims mOutputDims = outputDesc[0].dims;
        //cout << "mOutputDims from enqueue " << mOutputDims.d[0]<< " " << mOutputDims.d[1] << endl;
        interface(stream, inputs[0], outputs[0], rows, batch, widths, mOutputDims.d[1]);
        return 0;
    }
    //All class member private variables are serialized 
    //in the following two methods one after another
    //serializationSize is the size of all member variables
    size_t getSerializationSize() const override
    {
        size_t serializationSize = 0;  
        return serializationSize;
    }
    //serialize to char pointer
    void serialize(void* buffer) const override
    {
        char* d = static_cast<char*>(buffer);
        const char* const a = d;        
        assert(d == a + getSerializationSize());
    }
    //Plugin configuration for the input/output types, formats and sizes
    //PluginTensorDesc are fields that a plugin might see for an input or output.
    //it has 4 attributes (Dims, DataType, TensorFormat, float 	scale)
    //you can assert all match to the requirements
    //you can check all input/output types meets the expectations
    //For this CTCGreedyDecoder
    //first input is 3D vector 
    //second is 1D vector     
    void configurePlugin(const DynamicPluginTensorDesc* in, int nbInput, const DynamicPluginTensorDesc* out, int nbOutput) override
    {    
        //cout << "nbInput from configurePlugin " << nbInput << endl;
        //cout << "nbOutput from configurePlugin " << nbOutput << endl;    
        assert(in && nbInput == 2);
        assert(out && nbOutput == 1);
        return ;
    }

    //! The combination of kLINEAR + kINT8/kHALF/kFLOAT is supported.
    //inOut[pos].type pos 0 is kFLOAT   inOut[pos].type pos 1 is kINT32   
    //inOut[pos].format pos 0 is kLINEAR   inOut[pos].format pos 1 is kLINEAR     
    bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) override
    {     
        assert(nbInputs == 2 && nbOutputs == 1 && pos < nbInputs + nbOutputs);
        bool condition = inOut[pos].format == TensorFormat::kLINEAR;
        if(pos==0){
           condition &= inOut[pos].type == DataType::kFLOAT;
        }
        else if(pos==1){
           condition &= inOut[pos].type == DataType::kINT32;
        }
        return condition;
    }

    DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const override
    {
        //cout << "inputTypes from configugetOutputDataTyperePlugin " << inputTypes << endl; 
        assert(inputTypes && nbInputs == 2);
        (void) index;
        return inputTypes[0];
    }

    const char* getPluginType() const override
    {
        return "CTCGreedyDecoder";
    }

    const char* getPluginVersion() const override
    {
        return "1";
    }

    void destroy() override
    {
        delete this;
    }

    IPluginV2DynamicExt* clone() const override
    {
        auto* plugin = new CTCGreedyDecoder(*this);
        return plugin;
    }

    void setPluginNamespace(const char* libNamespace) override
    {
        mNamespace = libNamespace;
    }

    const char* getPluginNamespace() const override
    {
        return mNamespace.data();
    }
    

private:
    template <typename T>
    void write(char*& buffer, const T& val) const
    {
        *reinterpret_cast<T*>(buffer) = val;
        buffer += sizeof(T);
    }

    template <typename T>
    T read(const char*& buffer) const
    {
        T val = *reinterpret_cast<const T*>(buffer);
        buffer += sizeof(T);
        return val;
    }
    

private:    
    //This sparsetodense plugin doesn't need private members
    //All inputs/outputs are in the stream from previous outputs
    //so it is not necessary to have private members
    std::string mNamespace;
};

class CTCGreedyDecoderCreator : public IPluginCreator
{
public:
    const char* getPluginName() const override
    {
        return "CTCGreedyDecoder";
    }

    const char* getPluginVersion() const override
    {
        return "1";
    }

    const PluginFieldCollection* getFieldNames() override
    {
        return &mFieldCollection;
    }

    IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc) override
    {
        auto plugin = new CTCGreedyDecoder(*fc);
        mFieldCollection = *fc;
        mPluginName = name;
        return plugin;
    }
    
    IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override
    {
        //serialData is all data serialized in fun void serialize(void* buffer) const override
        //serialLength is length of data seiralized in serialize(void* buffer)
        auto plugin = new CTCGreedyDecoder(serialData, serialLength);
        mPluginName = name;
        return plugin;
    }

    void setPluginNamespace(const char* libNamespace) override
    {
        mNamespace = libNamespace;
    }

    const char* getPluginNamespace() const override
    {
        return mNamespace.c_str();
    }

private:
    std::string mNamespace;
    std::string mPluginName;
    PluginFieldCollection mFieldCollection{0, nullptr};
};

REGISTER_TENSORRT_PLUGIN(CTCGreedyDecoderCreator);


class Recognition
{
    template <typename T>
    using SampleUniquePtr = std::unique_ptr<T, samplesCommon::InferDeleter>;

public:
    Recognition(const samplesCommon::OnnxSampleParams& params)
        : mParams(params)
    {
    }

    //!
    //! \brief Builds both engines.
    //!
    bool build();

    //!
    //! \brief Prepares the model for inference by creating execution contexts and allocating buffers.
    //!
    bool prepare();

    //!
    //! \brief Runs inference using TensorRT on a random image.
    //!
    bool infer(std::vector<string> &files);
    vector<string> splitStrings(string str, char dl);
private:
    bool buildPreprocessorEngine(const SampleUniquePtr<nvinfer1::IBuilder>& builder);
    bool buildPredictionEngine(const SampleUniquePtr<nvinfer1::IBuilder>& builder);

    Dims loadJPEGFile(std::vector<std::string> fileName, int num);
    bool validateOutput(int digit);

    samplesCommon::OnnxSampleParams mParams; //!< The parameters for the sample.

    nvinfer1::Dims mPredictionInputDims;  //!< The dimensions of the input of the MNIST model.
    nvinfer1::Dims mPredictionOutputDims; //!< The dimensions of the output of the MNIST model.

    // Engines used for inference. The first is used for resizing inputs, the second for prediction.
    SampleUniquePtr<nvinfer1::ICudaEngine> mPreprocessorEngine{nullptr}, mPredictionEngine{nullptr};

    SampleUniquePtr<nvinfer1::IExecutionContext> mPreprocessorContext{nullptr}, mPredictionContext{nullptr};

    samplesCommon::ManagedBuffer mInput{};          //!< Host and device buffers for the input.
    samplesCommon::DeviceBuffer mPredictionInput{}; //!< Device buffer for the output of the preprocessor, i.e. the
                                                    //!< input to the prediction model.
    samplesCommon::ManagedBuffer mOutput{};         //!< Host buffer for the ouptut

    template <typename T>
    SampleUniquePtr<T> makeUnique(T* t)
    {
        return SampleUniquePtr<T>{t};
    }
};

//!
//! \brief Builds the two engines required for inference.
//!
//! \details This function creates one TensorRT engine for resizing inputs to the correct sizes,
//!          then creates a TensorRT network by parsing the ONNX model and builds
//!          an engine that will be used to run inference (mPredictionEngine).
//!
//! \return Ruturns false if error in build preprocessor or predict engine.
//!
bool Recognition::build()
{
    auto builder = makeUnique(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
    if (!builder)
    {
        sample::gLogError << "Create inference builder failed." << std::endl;
        return false;
    }
    // This function will also set mPredictionInputDims and mPredictionOutputDims,
    // so it needs to be called before building the preprocessor.
    return buildPredictionEngine(builder) && buildPreprocessorEngine(builder);
}

//!
//! \brief Builds an engine for preprocessing (mPreprocessorEngine).
//!
//! \return Ruturns false if error in build preprocessor engine.
//!
bool Recognition::buildPreprocessorEngine(const SampleUniquePtr<nvinfer1::IBuilder>& builder)
{
    // Create the preprocessor engine using a network that supports full dimensions (createNetworkV2).
    auto preprocessorNetwork = makeUnique(builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));
    if (!preprocessorNetwork)
    {
        sample::gLogError << "Create network failed." << std::endl;
        return false;
    }

    // Reshape a dynamically shaped input to the size expected by the model, (1, 1, 28, 28).
    auto input = preprocessorNetwork->addInput("input:0", nvinfer1::DataType::kFLOAT, Dims4{-1, 24, 94, 3});
    auto resizeLayer = preprocessorNetwork->addResize(*input);
    resizeLayer->setOutputDimensions(mPredictionInputDims);
    preprocessorNetwork->markOutput(*resizeLayer->getOutput(0));
    // Finally, configure and build the preprocessor engine.
    auto preprocessorConfig = makeUnique(builder->createBuilderConfig());
    if (!preprocessorConfig)
    {
        sample::gLogError << "Create builder config failed." << std::endl;
        return false;
    }

    // Create an optimization profile so that we can specify a range of input dimensions.
    auto profile = builder->createOptimizationProfile();
    // This profile will be valid for all images whose size falls in the range of [(1, 1, 1, 1), (1, 1, 56, 56)]
    // but TensorRT will optimize for (1, 1, 28, 28)
    // We do not need to check the return of setDimension and addOptimizationProfile here as all dims are explicitly set
    profile->setDimensions(input->getName(), OptProfileSelector::kMIN, Dims4(1, 24, 94, 3));
    profile->setDimensions(input->getName(), OptProfileSelector::kOPT, Dims4(5, 24, 94, 3));
    profile->setDimensions(input->getName(), OptProfileSelector::kMAX, Dims4(10, 24, 94, 3));
    preprocessorConfig->addOptimizationProfile(profile);

    // Create a calibration profile.
    auto profileCalib = builder->createOptimizationProfile();
    const int calibBatchSize{16};
    // We do not need to check the return of setDimension and setCalibrationProfile here as all dims are explicitly set
    profileCalib->setDimensions(input->getName(), OptProfileSelector::kMIN, Dims4{calibBatchSize, 24, 94, 3});
    profileCalib->setDimensions(input->getName(), OptProfileSelector::kOPT, Dims4{calibBatchSize, 24, 94, 3});
    profileCalib->setDimensions(input->getName(), OptProfileSelector::kMAX, Dims4{calibBatchSize, 24, 94, 3});
    preprocessorConfig->setCalibrationProfile(profileCalib);

    std::unique_ptr<IInt8Calibrator> calibrator;
    if (mParams.int8)
    {
        preprocessorConfig->setFlag(BuilderFlag::kINT8);
        const int nCalibBatches{10};
        MNISTBatchStream calibrationStream(
            calibBatchSize, nCalibBatches, "train-images-idx3-ubyte", "train-labels-idx1-ubyte", mParams.dataDirs);
        calibrator.reset(
            new Int8EntropyCalibrator2<MNISTBatchStream>(calibrationStream, 0, "MNISTPreprocessor", "input"));
        preprocessorConfig->setInt8Calibrator(calibrator.get());
    }

    mPreprocessorEngine = makeUnique(builder->buildEngineWithConfig(*preprocessorNetwork, *preprocessorConfig));
    if (!mPreprocessorEngine)
    {
        sample::gLogError << "Preprocessor engine build failed." << std::endl;
        return false;
    }else{
        sample::gLogError << "Preprocessor engine build succeeded." << std::endl;

    }
    sample::gLogInfo << "Profile dimensions in preprocessor engine:" << std::endl;
    sample::gLogInfo << "    Minimum = " << mPreprocessorEngine->getProfileDimensions(0, 0, OptProfileSelector::kMIN)
                     << std::endl;
    sample::gLogInfo << "    Optimum = " << mPreprocessorEngine->getProfileDimensions(0, 0, OptProfileSelector::kOPT)
                     << std::endl;
    sample::gLogInfo << "    Maximum = " << mPreprocessorEngine->getProfileDimensions(0, 0, OptProfileSelector::kMAX)
                     << std::endl;
    return true;
}

//!
//! \brief Builds an engine for prediction (mPredictionEngine).
//!
//! \details This function builds an engine for the MNIST model, and updates mPredictionInputDims and
//! mPredictionOutputDims according to the dimensions specified by the model. The preprocessor reshapes inputs to
//! mPredictionInputDims.
//!
//! \return Ruturns false if error in build prediction engine.
//!
bool Recognition::buildPredictionEngine(const SampleUniquePtr<nvinfer1::IBuilder>& builder)
{
    // Create a network using the parser.
    const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    auto network = makeUnique(builder->createNetworkV2(explicitBatch));
    if (!network)
    {
        sample::gLogError << "Create network failed." << std::endl;
        return false;
    }
    

    auto parser = samplesCommon::infer_object(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
    bool parsingSuccess = parser->parseFromFile(locateFile(mParams.onnxFileName, mParams.dataDirs).c_str(), static_cast<int>(sample::gLogger.getReportableSeverity()));
    if (!parsingSuccess)
    {
        sample::gLogError << "Failed to parse model." << std::endl;
        return false;
    }    

    // Get information about the inputs/outputs directly from the model.
    mPredictionInputDims = network->getInput(0)->getDimensions();
    cout << "mPredictionInputDims " << mPredictionInputDims << endl;
    mPredictionOutputDims = network->getOutput(0)->getDimensions();
    cout << "mPredictionOutputDims " << mPredictionOutputDims << endl;
    builder->setMaxBatchSize(10);
    // Create a builder config
    auto config = makeUnique(builder->createBuilderConfig());
    if (!config)
    {
        sample::gLogError << "Create builder config failed." << std::endl;
        return false;
    }
    config->setMaxWorkspaceSize(16_MiB);
    if (mParams.fp16)
    {
        config->setFlag(BuilderFlag::kFP16);
    }

    auto profileCalib = builder->createOptimizationProfile();
    const auto inputName = mParams.inputTensorNames[0].c_str();
    const int calibBatchSize{16};
    // We do not need to check the return of setDimension and setCalibrationProfile here as all dims are explicitly set
    profileCalib->setDimensions(inputName, OptProfileSelector::kMIN, Dims4{calibBatchSize, 24, 94, 3});
    profileCalib->setDimensions(inputName, OptProfileSelector::kOPT, Dims4{calibBatchSize, 24, 94, 3});
    profileCalib->setDimensions(inputName, OptProfileSelector::kMAX, Dims4{calibBatchSize, 24, 94, 3});
    config->setCalibrationProfile(profileCalib);

    // Create an optimization profile so that we can specify a range of input dimensions.
    auto profile = builder->createOptimizationProfile();
    profile->setDimensions(inputName, OptProfileSelector::kMIN, Dims4(10, 24, 94, 3));
    profile->setDimensions(inputName, OptProfileSelector::kOPT, Dims4(10, 24, 94, 3));
    profile->setDimensions(inputName, OptProfileSelector::kMAX, Dims4(10, 24, 94, 3));
    config->addOptimizationProfile(profile);
    std::unique_ptr<IInt8Calibrator> calibrator;
    if (mParams.int8)
    {
        config->setFlag(BuilderFlag::kINT8);
        int nCalibBatches{10};
        MNISTBatchStream calibrationStream(
            calibBatchSize, nCalibBatches, "train-images-idx3-ubyte", "train-labels-idx1-ubyte", mParams.dataDirs);
        calibrator.reset(
            new Int8EntropyCalibrator2<MNISTBatchStream>(calibrationStream, 0, "MNISTPrediction", inputName));
        config->setInt8Calibrator(calibrator.get());
    }
    // Build the prediciton engine.
    mPredictionEngine = makeUnique(builder->buildEngineWithConfig(*network, *config));
    if (!mPredictionEngine)
    {
        sample::gLogError << "Prediction engine build failed." << std::endl;
        return false;
    }else{
        sample::gLogError << "Prediction engine build succeeded." << std::endl;

    }
    return true;
}

//!
//! \brief Prepares the model for inference by creating an execution context and allocating buffers.
//!
//! \details This function sets up the sample for inference. This involves allocating buffers for the inputs and
//! outputs, as well as creating TensorRT execution contexts for both engines. This only needs to be called a single
//! time.
//!
//! \return Ruturns false if error in build preprocessor or predict context.
//!
bool Recognition::prepare()
{
    mPreprocessorContext = makeUnique(mPreprocessorEngine->createExecutionContext());
    if (!mPreprocessorContext)
    {
        sample::gLogError << "Preprocessor context build failed." << std::endl;
        return false;
    }

    mPredictionContext = makeUnique(mPredictionEngine->createExecutionContext());
    if (!mPredictionContext)
    {
        sample::gLogError << "Prediction context build failed." << std::endl;
        return false;
    }
    
    // Since input dimensions are not known ahead of time, we only allocate the output buffer and preprocessor output
    // buffer.
    mPredictionInput.resize(mPredictionInputDims);
    mOutput.hostBuffer.resize(mPredictionOutputDims);
    mOutput.deviceBuffer.resize(mPredictionOutputDims);
    return true;
}

vector<string> Recognition::splitStrings(string str, char dl) 
{ 
    string word = ""; 
  
    // adding delimiter character at the end 
    // of 'str' 
    str = str + dl; 
  
    // length of 'str' 
    int l = str.size(); 
  
    // traversing 'str' from left to right 
    vector<string> substr_list; 
    for (int i = 0; i < l; i++) { 
  
        // if str[i] is not equal to the delimiter 
        // character then accumulate it to 'word' 
        if (str[i] != dl) 
            word = word + str[i]; 
  
        else { 
  
            // if 'word' is not an empty string, 
            // then add this 'word' to the array 
            // 'substr_list[]' 
            if ((int)word.size() != 0) 
                substr_list.push_back(word); 
  
            // reset 'word' 
            word = ""; 
        } 
    } 
  
    // return the splitted strings 
    return substr_list; 
} 
//!
//! \brief Runs inference for this sample
//!
//! \details This function is the main execution function of the sample.
//! It runs inference for using a random image from the MNIST dataset as an input.
//!
bool Recognition::infer(std::vector<string> &files)
{
    // Load a random PGM file into a host buffer, then copy to device.
    std::random_device rd{};
    std::default_random_engine generator{rd()};
    std::uniform_int_distribution<int> digitDistribution{0, 9};
    int digit = digitDistribution(generator);
    cv::namedWindow( "Display window", cv::WINDOW_AUTOSIZE );// Create a window for display.
    char decode[]={'0','1','2','3','4','5','6','7','8','9','A','B','C','D','E','F','G','H','J','K','L','M','N','P','Q','R','S','T','U','V','W','X','Y','Z','c','d','e','f','g','h','j','k','p','s','u', 'v','w','_'};
    int number_of_files=10;
    std::vector<std::string> fileName;
    std::vector<std::string> numplate;
    std::string chars="\n";
    ofstream myfile;
    myfile.open ("labels_validated.txt");
    float truecnt=0;
    float filecnt=0;
    for(unsigned int i=10; i < files.size(); i=i+number_of_files){
        //number_of_files = (rand() % 10)+1;
        cout << "i " << i << endl;
        for(unsigned int j=i-10; j < i-10+number_of_files; j++){
           cout <<  files[j] << endl;           
           vector<string> res = splitStrings(files[j], ' ');            
           fileName.push_back(res[0]);  
           for(char c:chars){
             res[1].erase(std::remove(res[1].begin(), res[1].end(), c), res[1].end());
           }
           numplate.push_back(res[1]);
        }       
        high_resolution_clock::time_point t1 = high_resolution_clock::now();
        Dims inputDims = loadJPEGFile(fileName, number_of_files);
        mInput.deviceBuffer.resize(inputDims);
        CHECK(cudaMemcpy(mInput.deviceBuffer.data(), mInput.hostBuffer.data(), mInput.hostBuffer.nbBytes(), cudaMemcpyHostToDevice));

        // Set the input size for the preprocessor
        CHECK_RETURN_W_MSG(mPreprocessorContext->setBindingDimensions(0, inputDims), false, "Invalid binding dimensions.");

        // We can only run inference once all dynamic input shapes have been specified.
        if (!mPreprocessorContext->allInputDimensionsSpecified())
        {
            return false;
        }
        // Run the preprocessor to resize the input to the correct shape
        std::vector<void*> preprocessorBindings = {mInput.deviceBuffer.data(), mPredictionInput.data()};
        // For engines using full dims, we can use executeV2, which does not include a separate batch size parameter.
	bool status = mPreprocessorContext->executeV2(preprocessorBindings.data());
	if (!status)
	{
	    return false;
	}

	// Next, run the model to generate a prediction.
	std::vector<void*> predicitonBindings = {mPredictionInput.data(), mOutput.deviceBuffer.data()};
	status = mPredictionContext->executeV2(predicitonBindings.data());
	if (!status)
	{
           return false;
	}

        // Copy the outputs back to the host and verify the output.
        CHECK(cudaMemcpy(mOutput.hostBuffer.data(), mOutput.deviceBuffer.data(), mOutput.deviceBuffer.nbBytes(), cudaMemcpyDeviceToHost));
        high_resolution_clock::time_point t2 = high_resolution_clock::now();
        duration<double, std::milli> time_span = t2 - t1;
        std::cout << "It took me " << time_span.count() << " milliseconds.\n";          
        float *data = static_cast<float*>(mOutput.hostBuffer.data());
        const int d_w=3;
        int d_h=number_of_files/d_w;
        d_h=d_h+number_of_files%d_w;
        const int ind_h=50;
        cv::Mat background(ind_h*d_h, 94*d_w, CV_8UC3, cv::Scalar(0, 0, 0));
        for(int m=0; m < number_of_files; m++){
           cv::Mat image;
           image = cv::imread(fileName[m], cv::IMREAD_COLOR);
           std::vector<char> str;
           int offset=m*20;
           for(int k=0; k < 20; k++){
              if(*(data+offset+k) < 0)
                 break;
              str.push_back(decode[(int)*(data+offset+k)]);
           }
           std::string s(str.begin(), str.end());
           if (numplate[m].compare(s) == 0){
              myfile << fileName[m] <<" " << numplate[m] <<" "<< "true\n";
              truecnt++;
           }else{
              myfile << fileName[m] <<" " << numplate[m] <<" "<< s << "\n";
           }  
           filecnt++;
           int row=m/d_w;
           int col=m%d_w;
           image.copyTo(background(cv::Rect(col*image.cols, row*ind_h, image.cols, image.rows))); 
           cv::putText(background, s, cv::Point(col*image.cols, row*ind_h+40), cv::FONT_HERSHEY_SIMPLEX, 0.5, CV_RGB(255, 255, 255), 2);  
        } 
        cv::imshow( "Display window", background ); 
        cv::waitKey(1000);  
        background.release();     
        fileName.clear();
        numplate.clear();
    }
    float percent = (truecnt/filecnt)*100;
    myfile << "Percentage " << percent << "\n";
    myfile.close();
    return validateOutput(digit);
}

//!
//! \brief Loads a PGM file into mInput and returns the dimensions of the loaded image.
//!
//! \details This function loads the specified PGM file into the input host buffer.
//!
Dims Recognition::loadJPEGFile(std::vector<std::string> fileName, int num)
{
    Dims4 inputDims{num, 24, 94, 3};
    Dims4 inputDims_1img{1, 24, 94, 3};
    const size_t vol = samplesCommon::volume(inputDims);
    const size_t vol_1img = samplesCommon::volume(inputDims_1img);
    unsigned char *data = new unsigned char[vol];
    for(int f=0; f < num; f++){
       cv::Mat image, im_rgb;
       image = cv::imread(fileName[f], cv::IMREAD_COLOR);
       cv::cvtColor(image, im_rgb, cv::COLOR_BGR2RGB);
       image.release();
       memcpy(data+(f*vol_1img), im_rgb.ptr<unsigned char>(), vol_1img);
       im_rgb.release();
       mInput.hostBuffer.resize(inputDims);
       float* hostDataBuffer = static_cast<float*>(mInput.hostBuffer.data());
       std::transform(data, data+vol, hostDataBuffer, [](uint8_t x) { return 1.0 - static_cast<float>(x / 255.0); });      
    }
    delete[] data;
    return inputDims;    
}

//!
//! \brief Checks whether the model prediction (in mOutput) is correct.
//!
bool Recognition::validateOutput(int digit)
{
    const float* bufRaw = static_cast<const float*>(mOutput.hostBuffer.data());
    std::vector<float> prob(bufRaw, bufRaw + mOutput.hostBuffer.size());

    int curIndex{0};
    for (const auto& elem : prob)
    {
        sample::gLogInfo << " Prob " << curIndex << "  " << std::fixed << std::setw(5) << std::setprecision(4) << elem
                         << " "
                         << "Class " << curIndex << ": " << std::string(int(std::floor(elem * 10 + 0.5f)), '*')
                         << std::endl;
        ++curIndex;
    }

    int predictedDigit = std::max_element(prob.begin(), prob.end()) - prob.begin();
    return digit == predictedDigit;
}

//!
//! \brief Initializes members of the params struct using the command line args
//!
samplesCommon::OnnxSampleParams initializeSampleParams(const samplesCommon::Args& args)
{
    samplesCommon::OnnxSampleParams params;
    if (args.dataDirs.empty()) //!< Use default directories if user hasn't provided directory paths
    {
        params.dataDirs.push_back("data/mnist/");
        params.dataDirs.push_back("data/samples/mnist/");
    }
    else //!< Use the data directory provided by the user
    {
        params.dataDirs = args.dataDirs;
    }
    params.onnxFileName = "numplate_recg_10hwc_removed_sparsetodense.onnx";
    params.inputTensorNames.push_back("input:0");
    params.outputTensorNames.push_back("d_predictions:0");
    params.int8 = args.runInInt8;
    params.fp16 = args.runInFp16;
    return params;
}

//!
//! \brief Prints the help information for running this sample
//!
void printHelpInfo()
{
    std::cout << "Usage: ./sample_dynamic_reshape [-h or --help] [-d or --datadir=<path to data directory>]"
              << std::endl;
    std::cout << "--help, -h      Display help information" << std::endl;
    std::cout << "--datadir       Specify path to a data directory, overriding the default. This option can be used "
                 "multiple times to add multiple directories. If no data directories are given, the default is to use "
                 "(data/samples/mnist/, data/mnist/)"
              << std::endl;
    std::cout << "--int8          Run in Int8 mode." << std::endl;
    std::cout << "--fp16          Run in FP16 mode." << std::endl;
}

int listFiles( string path,  std::vector<string> &files)
{
   DIR *dir;
   struct dirent *ent;
   if ((dir = opendir (path.c_str())) != NULL) {
      /* print all the files and directories within directory */
      while ((ent = readdir (dir)) != NULL) {          
          files.push_back(path+ent->d_name);          
      }
      closedir (dir);
   } else {
      /* could not open directory */
      perror ("");
      return 1;
   }
   return 0;
}

//gdb --args ./platerecg_debug --useDLACore=-1 --fp16 --datadir=/usr/src/tensorrt/data/platerect/ --saveEngine=/usr/src/tensorrt/samples/Recognition loadEngine=/usr/src/tensorrt/samples/Recognition
//./platerecg --useDLACore=-1 --fp16 --datadir=/usr/src/tensorrt/data/platerect/ --saveEngine=/usr/src/tensorrt/samples/Recognition loadEngine=/usr/src/tensorrt/samples/Recognition
int main(int argc, char** argv)
{
    //string path="/usr/src/tensorrt/samples/Recognition/val";
    //std::vector<string> files;
    //listFiles( path, files );
    std::ifstream infile("/usr/src/tensorrt/samples/Recognition_DynamicBatch/val");
    std::vector<string> files;
    std::string line;
    while (std::getline(infile, line))
    {
        files.push_back(line);	    
    }
    samplesCommon::Args args;
    bool argsOK = samplesCommon::parseArgs(args, argc, argv);
    if (!argsOK)
    {
        sample::gLogError << "Invalid arguments" << std::endl;
        printHelpInfo();
        return EXIT_FAILURE;
    }
    if (args.help)
    {
        printHelpInfo();
        return EXIT_SUCCESS;
    }

    auto sampleTest = sample::gLogger.defineTest(gSampleName, argc, argv);

    sample::gLogger.reportTestStart(sampleTest);

    Recognition platerecg{initializeSampleParams(args)};

    if (!platerecg.build())
    {
        return sample::gLogger.reportFail(sampleTest);
    }
    if (!platerecg.prepare())
    {
        return sample::gLogger.reportFail(sampleTest);
    }
    if (!platerecg.infer(files))
    {
        return sample::gLogger.reportFail(sampleTest);
    }

    return sample::gLogger.reportPass(sampleTest);
}