# Copyright © 2020 Arm Ltd and Contributors. All rights reserved.
# SPDX-License-Identifier: MIT

"""
Object detection demo that takes a video stream from a device, runs inference
on each frame producing bounding boxes and labels around detected objects,
and displays a window with the latest processed frame.
"""

import os
import sys

script_dir = os.path.dirname(__file__)
sys.path.insert(1, os.path.join(script_dir, '..', 'common'))

import cv2
from argparse import ArgumentParser
from ssd import ssd_processing, ssd_resize_factor
from yolo import yolo_processing, yolo_resize_factor
from utils import dict_labels, Profiling
from cv_utils import init_video_stream_capture, preprocess, draw_bounding_boxes
import style_transfer


def get_model_processing(model_name: str, video: cv2.VideoCapture, input_data_shape: tuple):
    """
    Gets model-specific information such as model labels and decoding and processing functions.
    The user can include their own network and functions by adding another statement.

    Args:
        model_name: Name of type of supported model.
        video: Video capture object, contains information about data source.
        input_data_shape: Contains shape of model input layer, used for scaling bounding boxes.

    Returns:
        Model labels, decoding and processing functions.
    """
    if model_name == 'ssd_mobilenet_v1':
        return ssd_processing, ssd_resize_factor(video)
    elif model_name == 'yolo_v3_tiny':
        return yolo_processing, yolo_resize_factor(video, input_data_shape)
    else:
        raise ValueError(f'{model_name} is not a valid model name')


def main(args):

    enable_profile = args.profiling_enabled == "true"
    action_profiler = Profiling(enable_profile)
    action_profiler.profiling_start()

    if args.tflite_delegate_path is not None:
        from network_executor_tflite import TFLiteNetworkExecutor as NetworkExecutor
        exec_input_args = (args.model_file_path, args.preferred_backends, args.tflite_delegate_path)
    else:
        from network_executor import ArmnnNetworkExecutor as NetworkExecutor
        exec_input_args = (args.model_file_path, args.preferred_backends)

    executor = NetworkExecutor(*exec_input_args)
    action_profiler.profiling_stop_and_print_us("Executor initialization")

    action_profiler.profiling_start()
    video = init_video_stream_capture(args.video_source)
    action_profiler.profiling_stop_and_print_us("Video initialization")
    model_name = args.model_name
    process_output, resize_factor = get_model_processing(args.model_name, video, executor.get_shape())
    labels = dict_labels(args.label_path, include_rgb=True)

    if all(element is not None for element in [args.style_predict_model_file_path,
                                               args.style_transfer_model_file_path,
                                               args.style_image_path, args.style_transfer_class]):
        style_image = cv2.imread(args.style_image_path)
        action_profiler.profiling_start()
        style_transfer_executor = style_transfer.StyleTransfer(args.style_predict_model_file_path,
                                                               args.style_transfer_model_file_path,
                                                               style_image, args.preferred_backends,
                                                               args.tflite_delegate_path)
        action_profiler.profiling_stop_and_print_us("Style Transfer Executor initialization")

    while True:
        frame_present, frame = video.read()
        frame = cv2.flip(frame, 1)  # Horizontally flip the frame
        if not frame_present:
            raise RuntimeError('Error reading frame from video stream')

        action_profiler.profiling_start()
        if model_name == "ssd_mobilenet_v1":
            input_data = preprocess(frame, executor.get_data_type(), executor.get_shape(), True)
        else:
            input_data = preprocess(frame, executor.get_data_type(), executor.get_shape(), False)

        output_result = executor.run([input_data])
        if not enable_profile:
            print("Running inference...")
        action_profiler.profiling_stop_and_print_us("Running inference...")
        detections = process_output(output_result)
        if all(element is not None for element in [args.style_predict_model_file_path,
                                                   args.style_transfer_model_file_path,
                                                   args.style_image_path, args.style_transfer_class]):
            action_profiler.profiling_start()
            frame = style_transfer.create_stylized_detection(style_transfer_executor, args.style_transfer_class,
                                                             frame, detections, resize_factor, labels)
            action_profiler.profiling_stop_and_print_us("Running Style Transfer")
        else:
            draw_bounding_boxes(frame, detections, resize_factor, labels)
        cv2.imshow('PyArmNN Object Detection Demo', frame)
        if cv2.waitKey(1) == 27:
            print('\nExit key activated. Closing video...')
            break
    video.release(), cv2.destroyAllWindows()


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--video_source', type=int, default=0,
                        help='Device index to access video stream. Defaults to primary device camera at index 0')
    parser.add_argument('--model_file_path', required=True, type=str,
                        help='Path to the Object Detection model to use')
    parser.add_argument('--model_name', required=True, type=str,
                        help='The name of the model being used. Accepted options: ssd_mobilenet_v1, yolo_v3_tiny')
    parser.add_argument('--label_path', required=True, type=str,
                        help='Path to the labelset for the provided model file')
    parser.add_argument('--preferred_backends', type=str, nargs='+', default=['CpuAcc', 'CpuRef'],
                        help='Takes the preferred backends in preference order, separated by whitespace, '
                             'for example: CpuAcc GpuAcc CpuRef. Accepted options: [CpuAcc, CpuRef, GpuAcc]. '
                             'Defaults to [CpuAcc, CpuRef]')
    parser.add_argument('--tflite_delegate_path', type=str,
                        help='Enter TensorFlow Lite Delegate file path (.so file). If not entered,'
                             'will use armnn executor')
    parser.add_argument('--profiling_enabled', type=str,
                        help='[OPTIONAL] Enabling this option will print important ML related milestones timing'
                             'information in micro-seconds. By default, this option is disabled.'
                             'Accepted options are true/false.')
    parser.add_argument('--style_predict_model_file_path', type=str,
                        help='Path to the style prediction model to use')
    parser.add_argument('--style_transfer_model_file_path', type=str,
                        help='Path to the style transfer model to use')
    parser.add_argument('--style_image_path', type=str,
                        help='Path to the style image to create stylized frames')
    parser.add_argument('--style_transfer_class', type=str,
                        help='A class to transform its style')

    args = parser.parse_args()
    main(args)