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

import numpy as np
import urllib.request
import cv2
import network_executor_tflite
import cv_utils


def style_transfer_postprocess(preprocessed_frame: np.ndarray, image_shape: tuple):
    """
        Resizes the output frame of style transfer network and changes the color back to original configuration

        Args:
            preprocessed_frame: A preprocessed frame after style transfer.
            image_shape: Contains shape of the original frame before preprocessing.

        Returns:
            Resizing factor to scale coordinates according to image_shape.
    """

    postprocessed_frame = np.squeeze(preprocessed_frame, axis=0)
    # select original height and width from image_shape
    frame_height = image_shape[0]
    frame_width = image_shape[1]
    postprocessed_frame = cv2.resize(postprocessed_frame, (frame_width, frame_height)).astype("float32") * 255
    postprocessed_frame = cv2.cvtColor(postprocessed_frame, cv2.COLOR_RGB2BGR)

    return postprocessed_frame


def create_stylized_detection(style_transfer_executor, style_transfer_class, frame: np.ndarray,
                              detections: list, resize_factor, labels: dict):
    """
        Perform style transfer on a detected class in a frame

        Args:
            style_transfer_executor: The style transfer executor
            style_transfer_class: The class detected to change its style
            frame: The original captured frame from video source.
            detections: A list of detected objects in the form [class, [box positions], confidence].
            resize_factor: Resizing factor to scale box coordinates to output frame size.
            labels: Dictionary of labels and colors keyed on the classification index.
    """
    for detection in detections:
        class_idx, box, confidence = [d for d in detection]
        label = labels[class_idx][0]
        if label.lower() == style_transfer_class.lower():
            # Obtain frame size and resized bounding box positions
            frame_height, frame_width = frame.shape[:2]
            x_min, y_min, x_max, y_max = [int(position * resize_factor) for position in box]

            # Ensure box stays within the frame
            x_min, y_min = max(0, x_min), max(0, y_min)
            x_max, y_max = min(frame_width, x_max), min(frame_height, y_max)

            # Crop only the detected object
            cropped_frame = cv_utils.crop_bounding_box_object(frame, x_min, y_min, x_max, y_max)

            # Run style_transfer on preprocessed_frame
            stylized_frame = style_transfer_executor.run_style_transfer(cropped_frame)

            # Paste stylized_frame on the original frame in the correct place
            frame[int(y_min)+1:int(y_max),  int(x_min)+1:int(x_max)] = stylized_frame

    return frame


class StyleTransfer:

    def __init__(self, style_predict_model_path: str, style_transfer_model_path: str,
                 style_image: np.ndarray, backends: list, delegate_path: str):
        """
        Creates an inference executor for style predict network, style transfer network,
        list of backends and a style image.

        Args:
            style_predict_model_path: model which is used to create a style bottleneck
            style_transfer_model_path: model which is used to create stylized frames
            style_image: an image to create the style bottleneck
            backends: List of backends to optimize network.
            delegate_path: tflite delegate file path (.so).
        """

        self.style_predict_executor = network_executor_tflite.TFLiteNetworkExecutor(style_predict_model_path, backends,
                                                                                    delegate_path)
        self.style_transfer_executor = network_executor_tflite.TFLiteNetworkExecutor(style_transfer_model_path,
                                                                                     backends,
                                                                                     delegate_path)
        self.style_bottleneck = self.run_style_predict(style_image)

    def get_style_predict_executor_shape(self):
        """
            Get the input shape of the initiated network.

            Returns:
                tuple: The Shape of the network input.
        """
        return self.style_predict_executor.get_shape()

    # Function to run create a style_bottleneck using preprocessed style image.
    def run_style_predict(self, style_image):
        """
            Creates bottleneck tensor for a given style image.

            Args:
                style_image: an image to create the style bottleneck

            Returns:
                style bottleneck tensor
        """
        # The style image has to be preprocessed to (1, 256, 256, 3)
        preprocessed_style_image = cv_utils.preprocess(style_image, self.style_predict_executor.get_data_type(),
                                                       self.style_predict_executor.get_shape(), True, keep_aspect_ratio=False)
        # output[0] is the style bottleneck tensor
        style_bottleneck = self.style_predict_executor.run([preprocessed_style_image])[0]

        return style_bottleneck

    # Run style transform on preprocessed style image
    def run_style_transfer(self, content_image):
        """
            Runs inference for given content_image and style bottleneck to create a stylized image.

            Args:
                content_image:a content image to stylize
        """
        # The content image has to be preprocessed to (1, 384, 384, 3)
        preprocessed_style_image = cv_utils.preprocess(content_image, np.float32,
                                                       self.style_transfer_executor.get_shape(), True, keep_aspect_ratio=False)

        # Transform content image. output[0] is the stylized image
        stylized_image = self.style_transfer_executor.run([preprocessed_style_image, self.style_bottleneck])[0]

        post_stylized_image = style_transfer_postprocess(stylized_image, content_image.shape)

        return post_stylized_image