#
# Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
# SPDX-License-Identifier: MIT
#
import argparse
from pathlib import Path
from typing import Union

import tflite_runtime.interpreter as tflite
from PIL import Image
import numpy as np


def check_args(args: argparse.Namespace):
    """Check the values used in the command-line have acceptable values

    args:
      - args: argparse.Namespace

    returns:
      - None

    raises:
      - FileNotFoundError: if passed files do not exist.
      - IOError: if files are of incorrect format.
    """
    input_image_p = args.input_image
    if not input_image_p.suffix in (".png", ".jpg", ".jpeg"):
        raise IOError(
            "--input_image option should point to an image file of the "
            "format .jpg, .jpeg, .png"
        )
    if not input_image_p.exists():
        raise FileNotFoundError("Cannot find ", input_image_p.name)
    model_p = args.model_file
    if not model_p.suffix == ".tflite":
        raise IOError("--model_file should point to a tflite file.")
    if not model_p.exists():
        raise FileNotFoundError("Cannot find ", model_p.name)
    label_mapping_p = args.label_file
    if not label_mapping_p.suffix == ".txt":
        raise IOError("--label_file expects a .txt file.")
    if not label_mapping_p.exists():
        raise FileNotFoundError("Cannot find ", label_mapping_p.name)

    # check all args given in preferred backends make sense
    supported_backends = ["GpuAcc", "CpuAcc", "CpuRef"]
    if not all([backend in supported_backends for backend in args.preferred_backends]):
        raise ValueError("Incorrect backends given. Please choose from "\
            "'GpuAcc', 'CpuAcc', 'CpuRef'.")

    return None


def load_image(image_path: Path, model_input_dims: Union[tuple, list], grayscale: bool):
    """load an image and put into correct format for the tensorflow lite model

    args:
      - image_path: pathlib.Path
      - model_input_dims: tuple (or array-like). (height,width)

    returns:
      - image: np.array
    """
    height, width = model_input_dims
    # load and resize image
    image = Image.open(image_path).resize((width, height))
    # convert to greyscale if expected
    if grayscale:
        image = image.convert("LA")

    image = np.expand_dims(image, axis=0)

    return image


def load_delegate(delegate_path: Path, backends: list):
    """load the armnn delegate.

    args:
      - delegate_path: pathlib.Path -> location of you libarmnnDelegate.so
      - backends: list -> list of backends you want to use in string format

    returns:
      - armnn_delegate: tflite.delegate
    """
    # create a command separated string
    backend_string = ",".join(backends)
    # load delegate
    armnn_delegate = tflite.load_delegate(
        library=delegate_path,
        options={"backends": backend_string, "logging-severity": "info"},
    )

    return armnn_delegate


def load_tf_model(model_path: Path, armnn_delegate: tflite.Delegate):
    """load a tflite model for use with the armnn delegate.

    args:
      - model_path: pathlib.Path
      - armnn_delegate: tflite.TfLiteDelegate

    returns:
      - interpreter: tflite.Interpreter
    """
    interpreter = tflite.Interpreter(
        model_path=model_path.as_posix(), experimental_delegates=[armnn_delegate]
    )
    interpreter.allocate_tensors()

    return interpreter


def run_inference(interpreter, input_image):
    """Run inference on a processed input image and return the output from
    inference.

    args:
      - interpreter: tflite_runtime.interpreter.Interpreter
      - input_image: np.array

    returns:
      - output_data: np.array
    """
    # Get input and output tensors.
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    # Test model on random input data.
    interpreter.set_tensor(input_details[0]["index"], input_image)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]["index"])

    return output_data


def create_mapping(label_mapping_p):
    """Creates a Python dictionary mapping an index to a label.

    label_mapping[idx] = label

    args:
      - label_mapping_p: pathlib.Path

    returns:
      - label_mapping: dict
    """
    idx = 0
    label_mapping = {}
    with open(label_mapping_p) as label_mapping_raw:
        for line in label_mapping_raw:
            label_mapping[idx] = line
            idx += 1

    return label_mapping


def process_output(output_data, label_mapping):
    """Process the output tensor into a label from the labelmapping file. Takes
    the index of the maximum valur from the output array.

    args:
      - output_data: np.array
      - label_mapping: dict

    returns:
      - str: labelmapping for max index.
    """
    idx = np.argmax(output_data[0])

    return label_mapping[idx]


def main(args):
    """Run the inference for options passed in the command line.

    args:
      - args: argparse.Namespace

    returns:
      - None
    """
    # sanity check on args
    check_args(args)
    # load in the armnn delegate
    armnn_delegate = load_delegate(args.delegate_path, args.preferred_backends)
    # load tflite model
    interpreter = load_tf_model(args.model_file, armnn_delegate)
    # get input shape for image resizing
    input_shape = interpreter.get_input_details()[0]["shape"]
    height, width = input_shape[1], input_shape[2]
    input_shape = (height, width)
    # load input image
    input_image = load_image(args.input_image, input_shape, False)
    # get label mapping
    labelmapping = create_mapping(args.label_file)
    output_tensor = run_inference(interpreter, input_image)
    output_prediction = process_output(output_tensor, labelmapping)

    print("Prediction: ", output_prediction)

    return None


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--input_image", help="File path of image file", type=Path, required=True
    )
    parser.add_argument(
        "--model_file",
        help="File path of the model tflite file",
        type=Path,
        required=True,
    )
    parser.add_argument(
        "--label_file",
        help="File path of model labelmapping file",
        type=Path,
        required=True,
    )
    parser.add_argument(
        "--delegate_path",
        help="File path of ArmNN delegate file",
        type=Path,
        required=True,
    )
    parser.add_argument(
        "--preferred_backends",
        help="list of backends in order of preference",
        type=str,
        nargs="+",
        required=False,
        default=["CpuAcc", "CpuRef"],
    )
    args = parser.parse_args()

    main(args)