/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \brief Example code on load and run TVM module.s
 * \file cpp_deploy.cc
 */
#include <dlpack/dlpack.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
#include <cstdio>
#include <fstream>
#include </usr/include/opencv2/opencv.hpp>
#include </usr/include/opencv2/highgui/highgui.hpp>  
#include </usr/include/opencv2/imgproc/imgproc.hpp> 
#include <iostream>
#include <typeinfo>
#include <algorithm>


using namespace cv;
void Verify(tvm::runtime::Module mod, std::string fname) {
  // Get the function from the module.
  tvm::runtime::PackedFunc f = mod.GetFunction(fname);
  ICHECK(f != nullptr);
  // Allocate the DLPack data structures.
  //
  // Note that we use TVM runtime API to allocate the DLTensor in this example.
  // TVM accept DLPack compatible DLTensors, so function can be invoked
  // as long as we pass correct pointer to DLTensor array.
  //
  // For more information please refer to dlpack.
  // One thing to notice is that DLPack contains alignment requirement for
  // the data pointer and TVM takes advantage of that.
  // If you plan to use your customized data container, please
  // make sure the DLTensor you pass in meet the alignment requirement.
  //
  DLTensor* x;
  DLTensor* y;
  int ndim = 1;
  int dtype_code = kDLFloat;
  int dtype_bits = 32;
  int dtype_lanes = 1;
  int device_type = kDLCPU;
  int device_id = 0;
  int64_t shape[1] = {10};
  TVMArrayAlloc(shape, ndim, dtype_code, dtype_bits, dtype_lanes, device_type, device_id, &x);
  TVMArrayAlloc(shape, ndim, dtype_code, dtype_bits, dtype_lanes, device_type, device_id, &y);
  for (int i = 0; i < shape[0]; ++i) {
    static_cast<float*>(x->data)[i] = i;
  }
  // Invoke the function
  // PackedFunc is a function that can be invoked via positional argument.
  // The signature of the function is specified in tvm.build
  f(x, y);
  // Print out the output
  for (int i = 0; i < shape[0]; ++i) {
    ICHECK_EQ(static_cast<float*>(y->data)[i], i + 1.0f);
  }
  LOG(INFO) << "Finish verification...";
  TVMArrayFree(x);
  TVMArrayFree(y);
}

void DeploySingleOp() {
  // Normally we can directly
  tvm::runtime::Module mod_dylib = tvm::runtime::Module::LoadFromFile("lib/test_addone_dll.so");
  LOG(INFO) << "Verify dynamic loading from test_addone_dll.so";
  Verify(mod_dylib, "addone");
  // For libraries that are directly packed as system lib and linked together with the app
  // We can directly use GetSystemLib to get the system wide library.
  LOG(INFO) << "Verify load function from system lib";
  tvm::runtime::Module mod_syslib = (*tvm::runtime::Registry::Get("runtime.SystemLib"))();
  Verify(mod_syslib, "addonesys");
}

void PreProcess(const Mat& image, Mat& image_blob)
{
	Mat input;
	image.copyTo(input);

	std::vector<Mat> channels, channel_p;
	split(input, channels);
	Mat R, G, B;
	B = channels.at(0);
	G = channels.at(1);
	R = channels.at(2);

	B = (B / 255. - 0.406) / 0.225;
	G = (G / 255. - 0.456) / 0.224;
	R = (R / 255. - 0.485) / 0.229;

	channel_p.push_back(R);
	channel_p.push_back(G);
	channel_p.push_back(B);

	Mat outt;
	merge(channel_p, outt);
	image_blob = outt;
}

void Mat_to_CHW(float *img_data, cv::Mat &frame)
{
    assert(img_data && !frame.empty());
    unsigned int volChl = 224 * 224;

    for(int c = 0; c < 3; ++c)
    {
        for (unsigned j = 0; j < volChl; ++j)
            img_data[c*volChl + j] = static_cast<float>(float(frame.data[j * 3 + c])/255.0);
    }

}

void DeployGraphExecutor() {
  LOG(INFO) << "Running graph executor...";
  // load in the library

  DLDevice dev{kDLROCM, 0};
  tvm::runtime::Module mod_factory = tvm::runtime::Module::LoadFromFile("lib/MobileNet_V2.so");
  
  // create the graph executor module
  using namespace std;
  tvm::runtime::Module gmod = mod_factory.GetFunction("default")(dev);
  cout<<"---------------"<<endl;

  tvm::runtime::PackedFunc set_input = gmod.GetFunction("set_input");
  tvm::runtime::PackedFunc get_output = gmod.GetFunction("get_output");
  tvm::runtime::PackedFunc run = gmod.GetFunction("run");
//  cv::Mat image = cv::imread("/home/linjq/ambulance.jpg");
  cv::Mat image = cv::imread("./cow.jpg");
  cv::Mat in_put;
  cv::Mat img_in;
  cv::resize(image, in_put, cv::Size(224, 224));
  //cv::cvtColor(frame, in_put, cv::COLOR_BGR2RGB);  

  float img_data[224*224*3];
//  PreProcess(in_put, img_in);
//  Mat_to_CHW(img_data, img_in);
  Mat_to_CHW(img_data, in_put);
  
  DLTensor* y;
  int out_ndim = 2;
  int64_t out_shape[2] = {1, 1000};
  int dtype_code = kDLFloat;
  int dtype_bits = 32;
  int dtype_lanes = 1;
  int device_type = kDLROCM;
  int device_id = 0;
  TVMArrayAlloc(out_shape, out_ndim, dtype_code, dtype_bits, dtype_lanes, device_type, device_id, &y);
  


  
  DLTensor* x;
  int ndim = 4;
  int64_t shape[4] = {1, 3 ,224, 224};
  TVMArrayAlloc(shape, ndim, dtype_code, dtype_bits, dtype_lanes, device_type, device_id, &x);
  memcpy(x->data,&img_data,3*224*224*sizeof(float));
  
  // set the right input
  set_input("data", x);
  // run the code
  run();
  //get the output
  
  get_output(0, y);

  float* result = new float[1000];
  TVMArrayCopyToBytes(y, result, 1000 * sizeof(float));

  float max_num = *max_element(result, result+1000);
//  int max_num_index = max_element(result,result+1000)-result;
  auto max_iter = std::max_element(result, result + 1000);
  auto max_num_index = std::distance(result, max_iter);
  cout<<"max_num:"<<max_num<<endl;
  cout<<"max_iter:"<<max_iter<<endl;
  cout<<"max_num_index:"<<max_num_index<<endl;
}

int main(void) {
  //DeploySingleOp();
  DeployGraphExecutor();
  return 0;
}