import numpy as np
from tensorflow.quantization import fake_quant_with_min_max_vars
from tensorflow.python.framework import graph_util
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()

model_path = "../run_lite/models/resnet_v2_50.tflite"
input_file = "../run_lite/models/resnet_v2_50.tflite.ms.bin"
output_file = "../run_lite/models/resnet_v2_50.tflite.ms.out"
# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path=model_path)
tensors = interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()
print(input_details)
output_details = interpreter.get_output_details()
print(str(output_details))
input_shape = (1, 224, 224, 3)
with tf.Session( ) as sess:
    # input = ((np.random.randint(1, 1, input_shape))).astype(dtype=np.float32)
    input = (((np.array(list(range(224*224*3)))/(224*224*3)/2).reshape(input_shape))).astype(dtype=np.float32)
    print("input_data: ", input)
    interpreter.set_tensor(input_details[0]['index'], input)
    interpreter.invoke()

    with open(input_file, 'wb') as fo:
        fo.write(input.astype(np.float32, copy=False))
    with open(output_file, 'w') as text_file:
        for output_detail in output_details:
            output_data = interpreter.get_tensor(output_detail['index'])
            flatten_data = np.squeeze(output_data).flatten()
            size = flatten_data.shape[0]
            # text = output_detail['name'] + ' ' + str(len(output_detail['shape'])) + ' '
            text = 'Softmax-114' + ' ' + str(len(output_detail['shape'])) + ' '
            for i in range(len(output_detail['shape'])):
                text += str(output_detail['shape'][i]) + ' '
            text_file.write(text + '\n')
            for i in range(size):
                text_file.write(str((flatten_data[i])) + ' ')
            text_file.write('\n')
            print("output_data: ", output_data)