#
# Copyright (c) 2021 Intel Corporation
#
# 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 sys
sys.path.append("../../../../demos/common/python")

import grpc
import numpy as np
import classes
from tensorflow import make_tensor_proto, make_ndarray, make_tensor_proto
import datetime
import argparse
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc

parser = argparse.ArgumentParser(description='Do requests to ie_serving and tf_serving using images in binary format')
parser.add_argument('--images_list', required=False, default='input_images.txt', help='path to a file with a list of labeled images')
parser.add_argument('--grpc_address',required=False, default='localhost',  help='Specify url to grpc service. default:localhost')
parser.add_argument('--grpc_port',required=False, default=9000, help='Specify port to grpc service. default: 9000')
parser.add_argument('--input_name',required=False, default='image_bytes', help='Specify input tensor name. default: image_bytes')
parser.add_argument('--output_name',required=False, default='probabilities', help='Specify output name. default: probabilities')
parser.add_argument('--model_name', default='resnet', help='Define model name, must be same as is in service. default: resnet',
                    dest='model_name')
# If input numpy file has too few frames according to the value of iterations and the batch size, it will be
# duplicated to match requested number of frames
parser.add_argument('--batchsize', default=1, help='Number of images in a single request. default: 1',
                    dest='batchsize')
args = vars(parser.parse_args())

channel = grpc.insecure_channel("{}:{}".format(args['grpc_address'],args['grpc_port']))
stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
input_images = args.get('images_list')
with open(input_images) as f:
    lines = f.readlines()
batch_size = int(args.get('batchsize'))
while batch_size > len(lines):
    lines += lines
print('Start processing:')
print('\tModel name: {}'.format(args.get('model_name')))
print('\tImages list file: {}'.format(args.get('images_list')))

count = 0
matched = 0
processing_times = np.zeros((0),int)

batch_i = 0
image_data = []
labels = []
for line in lines:
    batch_i += 1
    path, label = line.strip().split(" ")
    with open(path, 'rb') as f:
        image_data.append(f.read())
    labels.append(label)
    if batch_i < batch_size:
        continue
    request = predict_pb2.PredictRequest()
    request.model_spec.name = args.get('model_name')
    request.inputs[args['input_name']].CopyFrom(make_tensor_proto(image_data, shape=[len(image_data)]))
    start_time = datetime.datetime.now()
    result = stub.Predict(request, 10.0) # result includes a dictionary with all model outputs
    end_time = datetime.datetime.now()
    if args['output_name'] not in result.outputs:
        print("Invalid output name", args['output_name'])
        print("Available outputs:")
        for Y in result.outputs:
            print(Y)
        exit(1)
    duration = (end_time - start_time).total_seconds() * 1000
    processing_times = np.append(processing_times,np.array([int(duration)]))
    output = make_ndarray(result.outputs[args['output_name']])
    nu = np.array(output)
    # for object classification models show imagenet class
    print('Batch: {}; Processing time: {:.2f} ms; speed {:.2f} fps'.format(count // batch_size, round(duration, 2), round(1000 / duration, 2)))
    for i in range(nu.shape[0]):
        offset = 0
        if nu.shape[1] == 1001:
            offset = 1 
        ma = np.argmax(nu[i]) - offset
        mark_message = ""
        if int(labels[i]) == ma:
            matched += 1
            mark_message = "; Correct match."
        else:
            mark_message = "; Incorrect match. Should be {} {}".format(label, classes.imagenet_classes[int(label)])
        count += 1
        print("\t", count, classes.imagenet_classes[ma], ma, mark_message)
    image_data = []
    labels = []
    batch_i = 0

latency = np.average(processing_times)
accuracy = matched / count

print("Overall accuracy=",accuracy*100,"%")
print("Average latency=",latency,"ms")
