import time
import tvm
from tvm import te
import tvm.relay as relay
from tvm.contrib.download import download_testdata
from tensorflow import keras
import numpy as np
from tvm.contrib import graph_executor
import json
import time
######################################################################
# Load a test image
# ------------------
# A single cat dominates the examples!
from PIL import Image
from matplotlib import pyplot as plt
from tensorflow.keras.applications.resnet50 import preprocess_input
import subprocess
import sys
try:
    partition_point = int(sys.argv[1])
except:
    partition_point = 72

import subprocess
img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true"
img_path = download_testdata(img_url, "cat.png", module="data")
img = Image.open(img_path).resize((224, 224))
# input preprocess
data = np.array(img)[np.newaxis, :].astype("float32")
data = preprocess_input(data).transpose([0, 3, 1, 2])
print("input_1", data.shape)

# compile the model
target = "cuda"

dev = tvm.cuda(0)

batch_size = 10
lib_path = './model_build/'
full_lib_file = 'local_lib_full_{}.tar'.format(str(partition_point))
#front_lib_file = 'local_lib_front_{}.tar'.format(str(partition_point))
#back_lib_file = 'local_lib_back_{}.tar'.format(str(partition_point))

full_lib = tvm.runtime.load_module(lib_path + full_lib_file)
#front_lib = tvm.runtime.load_module(lib_path + front_lib_file)
#back_lib = tvm.runtime.load_module(lib_path + back_lib_file)

total_time = 0
full_model = graph_executor.GraphModule(full_lib["default"](dev))
#front_model = graph_executor.GraphModule(front_lib["default"](dev))
#back_model = graph_executor.GraphModule(back_lib["default"](dev))


full_time_start = time.time()
usage = []
for i in range(batch_size):
    input_data = tvm.nd.array(data)
    full_model.set_input('input_1', input_data)
    full_model.run()
    usage.append(int(subprocess.check_output("nvidia-smi | grep 3019MiB | grep -o -E '[0-9]+%' | tail -1", shell=True).decode().split('%')[0]))
    full_out = full_model.get_output(0).asnumpy()
    #full_model.set_input('input_1', input_data)
    #full_model.run()
    #full_out = full_model.get_output(0).asnumpy()

print(sum(usage) / len(usage), '%')
full_time = time.time() - full_time_start
print(subprocess.check_output('nvidia-smi', shell=True).decode())
 
#partition_time_start = time.time()

# front_time = 0
# back_time = 0
# for i in range(batch_size):
#     front_start = time.time()
# 
#     input_data = tvm.nd.array(data)
#     front_model.set_input('input_1', input_data)
#     front_model.run()
#     out_deploy = front_model.get_output(0).asnumpy()
#     
#     front_time += time.time() - front_start
# 
#     back_start = time.time()
#     input_data = tvm.nd.array(out_deploy)
#     back_model.set_input('input_1', input_data)
#     back_model.run()
#     partition_out = back_model.get_output(0).asnumpy()
# 
#     back_time += time.time() - back_start
# 
# partition_time = time.time() - partition_time_start


synset_url = "".join(
    [
        "https://gist.githubusercontent.com/zhreshold/",
        "4d0b62f3d01426887599d4f7ede23ee5/raw/",
        "596b27d23537e5a1b5751d2b0481ef172f58b539/",
        "imagenet1000_clsid_to_human.txt",
    ]
)
synset_name = "imagenet1000_clsid_to_human.txt"
synset_path = download_testdata(synset_url, synset_name, module="data")
with open(synset_path) as f:
    synset = eval(f.read())


top1_full = np.argmax(full_out)
#top1_partition= np.argmax(partition_out)
print("full model  (local)    : {:.8f}s".format(full_time))
#print("front model (local)    : {:.8f}s".format(front_time))
#print("back model  (local)    : {:.8f}s".format(back_time))
#print("partition model (local): {:.8f}s".format(partition_time))
print("Relay top-1 Full id: {}, class name: {}".format(top1_full, synset[top1_full]))
#print("Relay top-1 Partition id: {}, class name: {}".format(top1_partition, synset[top1_partition]))
# confirm correctness with keras output
# keras_out = keras_resnet50.predict(data.transpose([0, 2, 3, 1]))
# top1_keras = np.argmax(keras_out)
# print("Keras top-1 id: {}, class name: {}".format(top1_keras, synset[top1_keras]))
