import os
import sys
import argparse
import time  # Import time module for measuring execution time

PWD = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(1, os.path.join(PWD, '..'))

from utils.ExampleHelper import get_device_usb_speed_by_port_id
import kp
import cv2

SCPU_FW_PATH = os.path.join(PWD, '../../res/firmware/KL630/kp_firmware.tar')
MODEL_FILE_PATH = os.path.join(PWD, '../../res/models/KL630/yolov8n_face/models_630.nef')
IMAGE_FILE_PATH = r"C:\Users\Mantra\Desktop\KL630\face1.jpg"
LOOP_TIME = 10

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='KL630 Demo Generic Image Inference Example.')
    parser.add_argument('-p',
                        '--port_id',
                        help='Using specified port ID for connecting device (Default: port ID of first scanned Kneron '
                             'device)',
                        default=0,
                        type=int)
    args = parser.parse_args()

    usb_port_id = args.port_id

    try:
        if kp.UsbSpeed.KP_USB_SPEED_HIGH != get_device_usb_speed_by_port_id(usb_port_id=usb_port_id):
            print('\033[91m' + '[Warning] Device is not running at high speed.' + '\033[0m')
    except Exception as exception:
        print('Error: check device USB speed fail, port ID = \'{}\', error msg: [{}]'.format(usb_port_id,
                                                                                             str(exception)))
        exit(0)

    try:
        print('[Connect Device]')
        device_group = kp.core.connect_devices(usb_port_ids=[usb_port_id])
        print(' - Success')
    except kp.ApiKPException as exception:
        print('Error: connect device fail, port ID = \'{}\', error msg: [{}]'.format(usb_port_id,
                                                                                     str(exception)))
        exit(0)

    print('[Set Device Timeout]')
    kp.core.set_timeout(device_group=device_group, milliseconds=21000)  # Increased timeout to 15 seconds
    print(' - Success')

    try:
        print('[Upload Firmware]')
        kp.core.load_firmware_from_file(device_group=device_group,
                                        scpu_fw_path=SCPU_FW_PATH,
                                        ncpu_fw_path="")
        print(' - Success')
    except kp.ApiKPException as exception:
        print('Error: upload firmware failed, error = \'{}\''.format(str(exception)))
        exit(0)

    try:
        print('[Upload Model]')
        model_nef_descriptor = kp.core.load_model_from_file(device_group=device_group,
                                                            file_path=MODEL_FILE_PATH)
        print(' - Success')
    except kp.ApiKPException as exception:
        print('Error: upload model failed, error = \'{}\''.format(str(exception)))
        exit(0)

    print('[Read Image]')
    img = cv2.imread(filename=IMAGE_FILE_PATH)
    img_bgr565 = cv2.cvtColor(src=img, code=cv2.COLOR_BGR2BGR565)
    print(' - Success')

    generic_inference_input_descriptor = kp.GenericImageInferenceDescriptor(
        model_id=model_nef_descriptor.models[0].id,
        inference_number=0,
        input_node_image_list=[
            kp.GenericInputNodeImage(
                image=img_bgr565,
                image_format=kp.ImageFormat.KP_IMAGE_FORMAT_RGB565
            )
        ]
    )

    print("[Starting Inference Work]")
    print(' - Starting inference loop {} times'.format(LOOP_TIME))
    print(' - ', end='')
    total_inference_time = 0  # Initialize total inference time counter
    for i in range(LOOP_TIME):
        try:
            start_time = time.time()  # Record start time
            kp.inference.generic_image_inference_send(device_group=device_group,
                                                      generic_inference_input_descriptor=generic_inference_input_descriptor)
            generic_raw_result = kp.inference.generic_image_inference_receive(device_group=device_group)
            end_time = time.time()  # Record end time
            inference_time = end_time - start_time  # Calculate inference time
            total_inference_time += inference_time  # Accumulate total inference time
        except kp.ApiKPException as exception:
            print(' - Error: inference failed, error = {}'.format(exception))
            exit(0)
        print('.', end='', flush=True)
    print()

    avg_inference_time = total_inference_time / LOOP_TIME  # Calculate average inference time
    print('Average Inference Time: {:.4f} seconds'.format(avg_inference_time))  # Print average inference time