#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
#

import os
import sys
import argparse

import numpy as np

from infer import TensorRTInfer
from image_batcher import ImageBatcher


def main(args):
    annotations = {}
    for line in open(args.annotations, "r"):
        line = line.strip().split(args.separator)
        if len(line) < 2 or not line[1].isnumeric():
            print(
                "Could not parse the annotations file correctly, make sure the correct separator is used"
            )
            sys.exit(1)
        annotations[os.path.basename(line[0])] = int(line[1])

    trt_infer = TensorRTInfer(args.engine)
    batcher = ImageBatcher(
        args.input,
        *trt_infer.input_spec(),
        max_num_images=args.num_images,
        preprocessor=args.preprocessor
    )
    top1 = 0
    top5 = 0
    total = 0
    for batch, images in batcher.get_batch():
        classes, scores, top = trt_infer.infer(batch, top=5)
        for i in range(len(images)):
            image = os.path.basename(images[i])
            if image not in annotations.keys():
                print(
                    "Image '{}' does not appear in the annotations file, please make sure all evaluated "
                    "images have a corresponding ground truth label".format(image)
                )
                sys.exit(1)
            if annotations[image] == classes[i]:
                top1 += 1
            if annotations[image] in top[0][i]:
                top5 += 1
            total += 1
            top1_acc = 100 * (top1 / total)
            top5_acc = 100 * (top5 / total)
            print(
                "Processing {} / {} : Top-1 {:0.1f}% , Top-5: {:0.1f}%     ".format(
                    total, batcher.num_images, top1_acc, top5_acc
                ),
                end="\r",
            )
    print()
    print("Top-1 Accuracy: {:0.3f}%".format(top1_acc))
    print("Top-5 Accuracy: {:0.3f}%".format(top5_acc))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-e", "--engine", help="The TensorRT engine to infer with")
    parser.add_argument(
        "-i",
        "--input",
        help="The input to infer, either a single image path, or a directory of images",
    )
    parser.add_argument(
        "-a",
        "--annotations",
        help="Set the file to use for classification ground truth annotations",
    )
    parser.add_argument(
        "-s",
        "--separator",
        default=" ",
        help="Separator to use between columns when parsing the annotations file, default: ' ' (space)",
    )
    parser.add_argument(
        "-p",
        "--preprocessor",
        default="V2",
        choices=["V1", "V1MS", "V2"],
        help="Select the image preprocessor to use, either 'V2', 'V1' or 'V1MS', default: V2",
    )
    parser.add_argument(
        "-n",
        "--num_images",
        default=5000,
        type=int,
        help="The maximum number of images to use for validation, default: 5000",
    )
    args = parser.parse_args()
    if not all([args.engine, args.input, args.annotations]):
        parser.print_help()
        print("\nThese arguments are required: --engine  --input and --annotations")
        sys.exit(1)

    main(args)