#include "tflite_wrapper.hpp"
#include <iostream>
#include <opencv2/opencv.hpp>
#include "utils/log_util.h"

#ifdef __ANDROID__
#include <tensorflow/lite/delegates/gpu/delegate.h>
#else
#include "plat_delegate.h"
#endif

//https://www.tensorflow.org/lite/guide/inference#load_and_run_a_model_in_c
TfliteWrapper::TfliteWrapper(const char *modelPath) {
#ifdef DEBUG
    std::cout << "TfliteInfer()" << std::endl;
#endif
#ifdef __ANDROID__
    model = tflite::FlatBufferModel::BuildFromFile(modelPath);
    assert(model);
    tflite::ops::builtin::BuiltinOpResolver docResolver;
    tflite::InterpreterBuilder(*model, docResolver)(&interpreter);
    TfLiteGpuDelegateOptionsV2 options;
    options.is_precision_loss_allowed = 1;
    options.inference_preference =
            TFLITE_GPU_INFERENCE_PREFERENCE_FAST_SINGLE_ANSWER;
    options.inference_priority1 = TFLITE_GPU_INFERENCE_PRIORITY_MIN_LATENCY;
    options.inference_priority2 = TFLITE_GPU_INFERENCE_PRIORITY_MIN_MEMORY_USAGE;
    options.inference_priority3 = TFLITE_GPU_INFERENCE_PRIORITY_MAX_PRECISION;
//    options.experimental_flags = TFLITE_GPU_EXPERIMENTAL_FLAGS_GL_ONLY;
    options.max_delegated_partitions = 1;
    options.model_token = nullptr;
    options.serialization_dir = nullptr;
    delegate = TfLiteGpuDelegateV2Create(&options);
    assert(interpreter->ModifyGraphWithDelegate(delegate) == kTfLiteOk);
    interpreter->AllocateTensors();
#else
    model = TfLiteModelCreateFromFile(modelPath);
    TfLiteInterpreterOptions* options = TfLiteInterpreterOptionsCreate();
    delegate = (TfLiteDelegate *)create_delegate();
    TfLiteInterpreterOptionsAddDelegate(options, delegate);
    // Create the interpreter.
    interpreter = TfLiteInterpreterCreate(model, options);
    // Allocate tensors and populate the input tensor data.
    TfLiteInterpreterAllocateTensors(interpreter);
    TfLiteInterpreterOptionsDelete(options);
#endif
}

void TfliteWrapper::Invoke(const void *input_data, size_t input_data_size, void *output_data,
                         size_t output_data_size) {
//#ifdef DEBUG
//    int64 startTime = cv::getTickCount();
//#endif
#ifdef __ANDROID__
    auto* input = interpreter->typed_input_tensor<float>(0);
    memcpy(input, input_data, input_data_size);
    interpreter->Invoke();
    auto* output = interpreter->typed_output_tensor<float>(0);
    memcpy(output_data, output, output_data_size);
#else
    TfLiteTensor *input_tensor = TfLiteInterpreterGetInputTensor(interpreter, 0);
    const TfLiteTensor *output_tensor = TfLiteInterpreterGetOutputTensor(interpreter, 0);
    TfLiteTensorCopyFromBuffer(input_tensor, input_data, input_data_size);
    auto status = TfLiteInterpreterInvoke(interpreter);
    assert(status == TfLiteStatus::kTfLiteOk);
    TfLiteTensorCopyToBuffer(output_tensor, output_data, output_data_size);
#endif
//#ifdef DEBUG
//    LOGW("tflite", "Invoke time=%f", (cv::getTickCount() - startTime) / cv::getTickFrequency());
//#endif
}

TfliteWrapper::~TfliteWrapper() {
#ifdef DEBUG
    std::cout << "~TfliteInfer()" << std::endl;
#endif
#ifdef __ANDROID__
    TfLiteGpuDelegateV2Delete(delegate);
#else
    delete_delegate(delegate);
    TfLiteModelDelete(model);
    TfLiteInterpreterDelete(interpreter);
#endif
}

void TfliteWrapper::getInputShape(uint32_t *shape, size_t size) {
#ifdef __ANDROID__
    auto *data = interpreter->input_tensor(0)->dims->data;
    for (size_t i = 0; i < size; i++) {
        shape[i] = *(data + i);
    }
#else
    auto *inputTensorPtr = TfLiteInterpreterGetInputTensor(interpreter, 0);
    auto dimsNum = TfLiteTensorNumDims(inputTensorPtr);
    assert(size == dimsNum);
    for (int i = 0; i < dimsNum; i++) {
        shape[i] = TfLiteTensorDim(inputTensorPtr, i);
    }
#endif
}