import ktc
import onnx
from PIL import Image
import numpy as np
import sys

# Preprocess (change filenames)
optimized_onnx_path = sys.argv[1]
img_folder = "/docker_mount/"
onnx_folder = "/docker_mount/"
onnx_file = onnx_folder + optimized_onnx_path
test_img_list = [
    img_folder + "test_no_yawn0.jpg",
    img_folder + "test_no_yawn1.jpg",
    img_folder + "test_no_yawn2.jpg",
    img_folder + "test_no_yawn3.jpg",
]


# IP Evaluation
print("== Section 1. IP Evaluation ==")
km = ktc.ModelConfig(32770, "0001", "520", onnx_path = onnx_file)
eval_result = km.evaluate()
print(eval_result)

# E2E Simulator Check (Floating Point)
print("== Section 2. E2E Simulator Check (Floating Point) ==")
def preprocess(input_file):
    image = Image.open(input_file)
    image = image.convert("L")
    img_data = np.array(image.resize((320, 320), Image.BILINEAR)) / 255
    y = np.expand_dims(img_data, axis = 2)
    return y

#input_data = [preprocess(test_img_list[0])]
input_data = [preprocess(test_img_list[0])]
inf_results = ktc.kneron_inference(input_data, onnx_file = onnx_file, input_names=["input0"])
print(inf_results)

# Quantization
print("== Section 3. Quantization ==")
input_images = [
    preprocess(test_img_list[0]),
    preprocess(test_img_list[1]),
    preprocess(test_img_list[2]),
    preprocess(test_img_list[3]),
]

input_mapping = {"input0": input_images}
bie_path = km.analysis(input_mapping, output_bie = None, threads = 6)
#print(bie_path)

# E2E Simulator Check (Fixed Point)
print("== Section 4. E2E Simulator Check (Fixed Point) ==")
radix = ktc.get_radix(input_images)
fixed_results = ktc.kneron_inference(input_data, bie_file=bie_path, input_names=["input0"], radix=radix)
print(fixed_results)

# Batch compile
print("== Section 5. Batch compile ==")
compile_result = ktc.compile([km])

# E2E simulator (hardware)
print("== Section 6. E2E Simulator Check (Hardware) ==")
hw_results = ktc.kneron_inference(input_data, nef_file=compile_result, radix=radix, platform=520)
print(hw_results)