# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Inference and metrics computation code using a loaded model.

Arguments:
    image : image to be inferenced
    classmap: classmap generated by training script

Returns:
    Network predictions
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function


import argparse
import json
import os

import keras
from keras.applications.imagenet_utils import preprocess_input
import numpy as np
from PIL import Image

import nvidia_tao_tf1.cv.common.logging.logging as status_logging
from nvidia_tao_tf1.cv.common.utils import check_tf_oom
from nvidia_tao_tf1.cv.multitask_classification.utils.model_io import load_model


def build_command_line_parser(parser=None):
    """Build a command line parser for inference."""
    if parser is None:
        parser = argparse.ArgumentParser(description="Multitask classification inference script")
    parser.add_argument("--model_path",
                        "-m",
                        type=str,
                        help="TLT model file")
    parser.add_argument("--image_dir",
                        "-i",
                        type=str,
                        help="inference image")
    parser.add_argument("--classmap",
                        "-cm",
                        type=str,
                        help="Class map file generated from train example")
    parser.add_argument("--key",
                        "-k",
                        default="",
                        type=str,
                        help="TLT model key")
    parser.add_argument("-r",
                        "--results_dir",
                        type=str,
                        default=None,
                        help="Path to results directory")

    return parser


def parse_command_line_arguments(args=None):
    """Parse command line arguments for inference."""
    parser = build_command_line_parser()
    return vars(parser.parse_known_args(args)[0])


@check_tf_oom
def inference(model_file, image_file=None, classmap=None, key=None, results_file=None):
    """Inference on an image using a pretrained model file.

    Args:
        model_file : .hdf5 keras model file containing weights and topology
        image_file : image to be inferenced
        classmap : path to json file containing classmap output generated from the train script
        results_file: Path to store predicted outputs (result.txt)
    Returns:
        None

    Log:
        Image Mode:
            print classifier output
        Directory Mode:
            Classifier accuracy for given KPI dataset
    """
    s_logger = status_logging.get_status_logger()
    s_logger.write(
        status_level=status_logging.Status.STARTED,
        message="Starting inference."
    )
    # Retrieve model using the pretrained file
    # set custom_object to arbitrary function to avoid not_found error.
    keras.backend.set_learning_phase(0)

    model = load_model(model_file, key=key)

    # extracting the data format parameter to detect input shape
    data_format = model.layers[1].data_format

    # Computing shape of input tensor
    image_shape = model.layers[0].input_shape[1:4]

    # Printing summary of retrieved model
    model.summary()

    # Setting input shape
    if data_format == "channels_first":
        image_height, image_width = image_shape[1:3]
    else:
        image_height, image_width = image_shape[0:2]

    if image_file is not None:
        # Open image and preprocessing
        image = Image.open(image_file)
        image = image.resize((image_width, image_height), Image.ANTIALIAS).convert('RGB')
        inference_input = preprocess_input(np.array(image).astype(np.float32).transpose(2, 0, 1))
        inference_input.shape = (1, ) + inference_input.shape

        # Keras inference
        raw_predictions = model.predict(inference_input, batch_size=1)

        with open(classmap, "r") as cm:
            class_dict = json.load(cm)
            cm.close()
        task_name = class_dict['tasks']
        class_map = class_dict['class_mapping']
        for idx, task in enumerate(task_name):
            pred = raw_predictions[idx].reshape(-1)
            print("Task {}:".format(task))
            print("Predictions: {}".format(pred))
            class_name = class_map[task][str(np.argmax(pred))]
            print("Class name = {}".format(class_name))
            print('********')
            # Log into results_file if exists
            if results_file:
                with open(results_file, "a") as f:
                    f.write("Task {}:\n".format(task))
                    f.write("Predictions: {}\n".format(pred))
                    f.write("Class name = {}\n".format(class_name))
                    f.write('********\n')
    s_logger.write(
        status_level=status_logging.Status.SUCCESS,
        message="Inference finished successfully."
    )


if __name__ == "__main__":
    arguments = parse_command_line_arguments()
    results_file = None
    # Create results directory and init status.json
    if arguments["results_dir"]:
        if not os.path.exists(arguments["results_dir"]):
            os.makedirs(arguments["results_dir"])
        status_file = os.path.join(arguments["results_dir"], "status.json")
        status_logging.set_status_logger(
            status_logging.StatusLogger(
                filename=status_file,
                is_master=True,
                verbosity=1,
                append=True
            )
        )
        results_file = os.path.join(arguments["results_dir"], "result.txt")

    try:
        inference(arguments['model_path'], image_file=arguments['image_dir'],
                  classmap=arguments['classmap'], key=arguments['key'],
                  results_file=results_file)
    except (KeyboardInterrupt, SystemExit):
        status_logging.get_status_logger().write(
            message="Inference was interrupted",
            verbosity_level=status_logging.Verbosity.INFO,
            status_level=status_logging.Status.FAILURE
        )
    except Exception as e:
        status_logging.get_status_logger().write(
            message=str(e),
            status_level=status_logging.Status.FAILURE
        )
        raise e
