#
# http_infer_mlp
#
import sys
import numpy as np
import datetime
import argparse
import tritonclient.http as httpclient
from sklearn.datasets import make_moons
from sklearn.preprocessing import StandardScaler

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Sends requests via KServe REST API using images in numpy format. '
                                                 'It displays performance statistics and optionally the model accuracy')
    parser.add_argument('--http_address', required=False, default='localhost', help='Specify url to http service. default:localhost')
    parser.add_argument('--http_port', required=False, default=8000, help='Specify port to http service. default: 8000')
    parser.add_argument('--input_name', required=False, default='input', help='Specify input tensor name. default: input')
    parser.add_argument('--output_name', required=False, default='resnet_v1_50/predictions/Reshape_1',
                        help='Specify output name. default: resnet_v1_50/predictions/Reshape_1')
    parser.add_argument('--iterations', default=1000,
                        help='Number of requests iterations',
                        dest='iterations', type=int)
    parser.add_argument('--model_name', default='resnet', help='Define model name, must be same as is in service. default: resnet',
                        dest='model_name')
    parser.add_argument('--pipeline_name', default='', help='Define pipeline name, must be same as is in service',
                        dest='pipeline_name')
    parser.add_argument('--binary_data', default=False, action='store_true', help='Send input data in binary format', dest='binary_data')
    parser.add_argument('--tls', default=False, action='store_true', help='use TLS communication with gRPC endpoint')
    parser.add_argument('--server_cert', required=False, help='Path to server certificate', default=None)
    parser.add_argument('--client_cert', required=False, help='Path to client certificate', default=None)
    parser.add_argument('--client_key', required=False, help='Path to client key', default=None)

    args = vars(parser.parse_args())
    iterations = args.get('iterations')
    
    # Create the data
    X, y = make_moons(n_samples=iterations, noise=0.2, random_state=42)
    scaler = StandardScaler()
    X = scaler.fit_transform(X)

    address = "{}:{}".format(args['http_address'], args['http_port'])
    batch_size = 1

    if args['tls']:
        ssl_options = {
            'keyfile':args['client_key'],
            'cert_file':args['client_cert'],
            'ca_certs':args['server_cert']
        }
    else:
        ssl_options = None

    triton_client = httpclient.InferenceServerClient(
                url=address,
                ssl=args['tls'],
                ssl_options=ssl_options,
                verbose=False)

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

    print('Start processing:')
    print('\tModel name: {}'.format(args.get('pipeline_name') if bool(args.get('pipeline_name')) else args.get('model_name')))
    print('\tIterations: {}'.format(iterations))

    iteration = 0
    is_pipeline_request = bool(args.get('pipeline_name'))
    y_test = []
    num_correct = 0

    while iteration < iterations:
        inputs = []
        inputs.append(httpclient.InferInput(args['input_name'], [1,2], "FP32"))
        point = np.array([X[iteration]],dtype=np.float32)
        inputs[0].set_data_from_numpy(point)
        start_time = datetime.datetime.now()
        results = triton_client.infer(
            model_name=args.get('pipeline_name') if is_pipeline_request else args.get('model_name'),
            inputs=inputs)
        end_time = datetime.datetime.now()
        duration = (end_time - start_time).total_seconds() * 1000
        processing_times = np.append(processing_times, np.array([int(duration)]))
        output = results.as_numpy(args['output_name'])
        y_test.append(float(output[0][0]))
        if y[iteration] == round(float(output[0][0])):
            num_correct += 1
                                
        
        # for object classification models show imagenet class
        print('Iteration {}; Processing time: {:.2f} ms; speed {:.2f} fps'.format(iteration, round(np.average(duration), 2),
                                                                                  round(1000 * batch_size / np.average(duration), 2)
                                                                                  ))
        iteration += 1
        
        
        
    print('finished with {} out of {} correct'.format(num_correct, iterations))