/**
 *  \brief main function to run fix Inferencer
 *  \author Mingzhe Jiang
 *  \copyright 2020 Kneron Inc. All right reserved.
 */


#include "KL520WrapperFactory.h"
#include "KL720WrapperFactory.h"
#include "KL530WrapperFactory.h"
#include "KL630WrapperFactory.h"
#include "KL730WrapperFactory.h"
#include "KL540WrapperFactory.h"
#include <math.h>

//example code, in case users want to use dump and parser functions in this header
//#include "run_fix_inferencer.h"
#include "iostream"

using namespace std;
using std::string;
using std::vector;
using std::unordered_map;
using std::unique_ptr;
using std::shared_ptr;
using namespace dynasty::inferencer;
using namespace dynasty::inferencer::fix_point;


/**
 * This is an example of how to use dynamic fixed point inference.
 * The inferencer output will contain all layer dump if dump level >0, else only the output layer
 * user might need to write his own dump function for different dump purpose, can take run_fix_reference as an example
 * configration parser is also up to user's choice. Can be json or command line params.
 * @param argc
 * @param argv
 * @return
 */

int main(int argc, char *argv[]) {
    {
            bool cuda_en = false;   //enable cuda if cuda library is installed. don't enable if no cuda ata ll
            string model_file_540 = "res/KL540/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/output/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2.piano.kdp540.scaled.onnx";
            string model_file_530 = "res/KL530/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/output/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2.piano.kdp530.scaled.onnx";
            string model_file_630 = "res/KL630/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/output/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2.piano.kdp630.scaled.onnx";
            string model_file_730 = "res/KL730/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/output/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2.piano.kdp730.scaled.onnx";
            string model_file_720 = "res/conv_8bit_3x3_base/output/conv_8bit_3x3_base.piano.kdp720.scaled.onnx";
            string model_file_520 = "res/conv_8bit_3x3_base/output/conv_8bit_3x3_base.piano.kdp520.scaled.onnx";
            string radixFile_520 = "res/conv_8bit_3x3_base/output/conv_8bit_3x3_base.piano.kdp520.scaled.onnx.json";
            string radixFile_720 = "res/conv_8bit_3x3_base/output/conv_8bit_3x3_base.piano.kdp720.scaled.onnx.json";
            string radixFile_540 = "res/KL540/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/output/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2.piano.kdp540.scaled.onnx.json";
            string radixFile_530 = "res/KL530/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/output/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2.piano.kdp530.scaled.onnx.json";
            string radixFile_630 = "res/KL630/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/output/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2.piano.kdp630.scaled.onnx.json";
            string radixFile_730 = "res/KL730/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/output/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2.piano.kdp730.scaled.onnx.json";

            //get only output node with dump_level = 0; get all layer dump with dump_level = 2
            int dump_level = 0;
            const float EPSILON = 0.8;  // Max difference tolerance between 520 result and 720 result

            unordered_map<string, string> inference_input_540;
            string input_name_540 = "0";
            string input_path_540 = "res/KL540/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/input/simulator_input/test_input.txt";
            inference_input_540[input_name_540] = input_path_540;

            // Inference with 540: cuda disabled
            InferencerUniquePtr<float> inferencer_540 = kl540::InferencerImplWrapperFactory::create(model_file_540, radixFile_540, cuda_en);

            unordered_map<std::string, std::vector<float>> float_inputs;
            inferencer_540->convertInput(inference_input_540, float_inputs);

             unordered_map<std::string, std::vector<float>> float_outputs_540 = inferencer_540->Inference(float_inputs, dump_level);
             unordered_map<std::string, std::vector<int>> fix_outputs_540 = inferencer_540->ConvertFloatToInt(float_outputs_540, dump_level);

             auto builder = kl540::InferencerImplWrapperFactory::GetBuilder(inferencer_540.get());
             auto input_dim_540 = builder->GetInputDimensions();
             for(auto& item: input_dim_540) {
                 cout << "input node name: " << item.first << ", shape: ";
                 for(auto num: item.second) cout<< num <<", ";
                 cout << endl;
             }
             cout << "inference with 540 completed" << endl;

            unordered_map<string, string> inference_input_630;
            string input_name_630 = "0";
            string input_path_630 = "res/KL630/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/input/simulator_input/test_input.txt";
            inference_input_630[input_name_630] = input_path_630;

            // Inference with 630: cuda disabled
            InferencerUniquePtr<float> inferencer_630 = kl630::InferencerImplWrapperFactory::create(model_file_630, radixFile_630, cuda_en);

            inferencer_630->convertInput(inference_input_630, float_inputs);

             unordered_map<std::string, std::vector<float>> float_outputs_630 = inferencer_630->Inference(float_inputs, dump_level);
             unordered_map<std::string, std::vector<int>> fix_outputs_630 = inferencer_630->ConvertFloatToInt(float_outputs_630, dump_level);

             builder = kl630::InferencerImplWrapperFactory::GetBuilder(inferencer_630.get());
             auto input_dim_630 = builder->GetInputDimensions();
             for(auto& item: input_dim_630) {
                 cout << "input node name: " << item.first << ", shape: ";
                 for(auto num: item.second) cout<< num <<", ";
                 cout << endl;
             }
             cout << "inference with 630 completed" << endl;

            unordered_map<string, string> inference_input_730;
            string input_name_730 = "0";
            string input_path_730 = "res/KL730/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/input/simulator_input/test_input.txt";
            inference_input_730[input_name_730] = input_path_730;

            // Inference with 730: cuda disabled
            InferencerUniquePtr<float> inferencer_730 = kl730::InferencerImplWrapperFactory::create(model_file_730, radixFile_730, cuda_en);

            inferencer_730->convertInput(inference_input_730, float_inputs);

             unordered_map<std::string, std::vector<float>> float_outputs_730 = inferencer_730->Inference(float_inputs, dump_level);
             unordered_map<std::string, std::vector<int>> fix_outputs_730 = inferencer_730->ConvertFloatToInt(float_outputs_730, dump_level);

             builder = kl730::InferencerImplWrapperFactory::GetBuilder(inferencer_730.get());
             auto input_dim_730 = builder->GetInputDimensions();
             for(auto& item: input_dim_730) {
                 cout << "input node name: " << item.first << ", shape: ";
                 for(auto num: item.second) cout<< num <<", ";
                 cout << endl;
             }
             cout << "inference with 730 completed" << endl;

            unordered_map<string, string> inference_input_530;
            string input_name_530 = "0";
            string input_path_530 = "res/KL530/Conv1Ch_inputFM16x208x1_group1_kernel3x3_pads1_1_1_1_stride2x2_bias1_dilations2/input/simulator_input/test_input.txt";
            inference_input_530[input_name_530] = input_path_530;

            // Inference with 530: cuda disabled
            InferencerUniquePtr<float> inferencer_530 = kl530::InferencerImplWrapperFactory::create(model_file_530, radixFile_530, cuda_en);

            inferencer_530->convertInput(inference_input_530, float_inputs);

             unordered_map<std::string, std::vector<float>> float_outputs_530 = inferencer_530->Inference(float_inputs, dump_level);
             unordered_map<std::string, std::vector<int>> fix_outputs_530 = inferencer_530->ConvertFloatToInt(float_outputs_530, dump_level);

             builder = kl530::InferencerImplWrapperFactory::GetBuilder(inferencer_530.get());
             auto input_dim_530 = builder->GetInputDimensions();
             for(auto& item: input_dim_530) {
                 cout << "input node name: " << item.first << ", shape: ";
                 for(auto num: item.second) cout<< num <<", ";
                 cout << endl;
             }
             cout << "inference with 530 completed" << endl;



            // Set up 520/720 inputs
            unordered_map<string, string> inference_input;
            string input_name = "input_1_o0";
            string input_path = "res/conv_8bit_3x3_base/input/simulator_input/test_input.txt";
            inference_input[input_name] = input_path;

            // Inference with 720: cuda disabled
            InferencerUniquePtr<float> inferencer_720 = kl720::InferencerImplWrapperFactory::create(model_file_720, radixFile_720, cuda_en);

            inferencer_720->convertInput(inference_input, float_inputs);

             unordered_map<std::string, std::vector<float>> float_outputs_720 = inferencer_720->Inference(float_inputs, dump_level);
             unordered_map<std::string, std::vector<int>> fix_outputs_720 = inferencer_720->ConvertFloatToInt(float_outputs_720, dump_level);

             builder = kl720::InferencerImplWrapperFactory::GetBuilder(inferencer_720.get());
             auto input_dim_720 = builder->GetInputDimensions();
             for(auto& item: input_dim_720) {
                 cout << "input node name: " << item.first << ", shape: ";
                 for(auto num: item.second) cout<< num <<", ";
                 cout << endl;
             }
             cout << "inference with 720 completed" << endl;


             // Inference with 520: cuda disabled
            InferencerUniquePtr<float> inferencer_520 = kl520::InferencerImplWrapperFactory::create(model_file_520, radixFile_520, cuda_en);

            inferencer_520->convertInput(inference_input, float_inputs);

            unordered_map<std::string, std::vector<float>> float_outputs_520 = inferencer_520->Inference(float_inputs, dump_level);
            unordered_map<std::string, std::vector<int>> fix_outputs_520 = inferencer_520->ConvertFloatToInt(float_outputs_520, dump_level);

            builder = kl520::InferencerImplWrapperFactory::GetBuilder(inferencer_520.get());
            auto output_dim = builder->GetOutputDimensions();
            for(auto& item: output_dim) {
                cout << "output node name: " << item.first << ", shape: ";
                for(auto num: item.second) cout<< num <<", ";
                cout << endl;
            }
            cout << "inference with 520 completed" << endl;


            // compare results :
            if ( float_outputs_520.size() != float_outputs_720.size() ) {
                cout << "size is not equal" << endl;
            } else {
                auto AreSame=[EPSILON](float a, float b) {
                    return fabs(a - b) < EPSILON;
                };
                for (auto& item : float_outputs_720) {
                    vector<float>& result_720 = item.second;
                    vector<float>& result_520 = float_outputs_520[item.first];
                    cout << "....checking result for node: " << item.first << endl;
                    if ( result_520.size() != result_720.size()) {
                        cout << "size is not equal for : " << item.first;
                        exit(-1);
                    } else {
                        for(size_t i=0 ; i < result_720.size();i++) {
                            float r_520 = result_520[i];
                            float r_720 = result_720[i];
                            if (! AreSame(r_520, r_720) ) {
                                cout << "r_520: " << r_520 << " , r_720: " << r_720 << endl;
                                exit(-1);
                            }
                        }
                    }

                }
            }

            cout << "Done!" << endl;

        return 0;
    }
}