text: #  Copyright 2018 Amazon.com, Inc. or its affiliates. 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.
  #  A copy of the License is located at
  #  
  #      http://www.apache.org/licenses/LICENSE-2.0
  #  
  #  or in the "license" file accompanying this file. This file 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 numpy as np
  from google.protobuf import json_format
  from grpc.beta import implementations
  from tf_container.run import logger as _logger
  from tensorflow import make_tensor_proto
  from tensorflow.core.example import example_pb2, feature_pb2
  from tensorflow.core.framework import tensor_pb2
  from tensorflow.python.saved_model.signature_constants import DEFAULT_SERVING_SIGNATURE_DEF_KEY, PREDICT_INPUTS
  from tensorflow_serving.apis import get_model_metadata_pb2
  from tensorflow_serving.apis import predict_pb2, classification_pb2, inference_pb2, regression_pb2
  from tensorflow_serving.apis import prediction_service_pb2
  
  
  REGRESSION = 'tensorflow/serving/regression'
  CLASSIFY = 'tensorflow/serving/classify'
  INFERENCE = 'tensorflow/serving/inference'
  PREDICT = 'tensorflow/serving/predict'
  GENERIC_MODEL_NAME = "generic_model"
  
  
  class GRPCProxyClient(object):
      def __init__(self, tf_serving_port, host='localhost', request_timeout=10.0, model_name=GENERIC_MODEL_NAME,
                   input_tensor_name=PREDICT_INPUTS, signature_name=DEFAULT_SERVING_SIGNATURE_DEF_KEY):
          self.tf_serving_port = tf_serving_port
          self.host = host
          self.request_timeout = request_timeout
          self.model_name = model_name
          self.input_tensor_name = input_tensor_name
          self.signature_name = signature_name
          self.request_fn_map = {PREDICT: self.predict,
                                 CLASSIFY: self.classification,
                                 # TODO: implement inference and regression tf serving apis
                                 INFERENCE: self._raise_not_implemented_exception,
                                 REGRESSION: self._raise_not_implemented_exception,
                                 }
          self.prediction_type = None
          self.input_type_map = {}
  
      def parse_request(self, serialized_data):
          request_fn_map = {
              PREDICT: lambda: predict_pb2.PredictRequest(),
              INFERENCE: lambda: inference_pb2.MultiInferenceRequest(),
              CLASSIFY: lambda: classification_pb2.ClassificationRequest(),
              REGRESSION: lambda: regression_pb2.RegressionRequest()
          }
  
          request = request_fn_map[self.prediction_type]()
          request.ParseFromString(serialized_data)
  
          return request
  
      def request(self, data):
          request_fn = self.request_fn_map[self.prediction_type]
          return request_fn(data)
  
      def cache_prediction_metadata(self):
          channel = implementations.insecure_channel(self.host, self.tf_serving_port)
          stub = prediction_service_pb2.beta_create_PredictionService_stub(channel)
          request = get_model_metadata_pb2.GetModelMetadataRequest()
  
          request.model_spec.name = self.model_name
          request.metadata_field.append('signature_def')
          result = stub.GetModelMetadata(request, self.request_timeout)
  
          _logger.info('---------------------------Model Spec---------------------------')
          _logger.info(json_format.MessageToJson(result))
          _logger.info('----------------------------------------------------------------')
  
          signature_def = result.metadata['signature_def']
          signature_map = get_model_metadata_pb2.SignatureDefMap()
          signature_map.ParseFromString(signature_def.value)
  
          serving_default = signature_map.ListFields()[0][1]['serving_default']
          serving_inputs = serving_default.inputs
  
          self.input_type_map = {key: serving_inputs[key].dtype for key in serving_inputs.keys()}
          self.prediction_type = serving_default.method_name
  
      def predict(self, data):
  
          channel = implementations.insecure_channel(self.host, self.tf_serving_port)
          stub = prediction_service_pb2.beta_create_PredictionService_stub(channel)
  
          request = self._create_predict_request(data)
  
          result = stub.Predict(request, self.request_timeout)
  
          return result
  
      def _create_predict_request(self, data):
          # Send request
          # See prediction_service.proto for gRPC request/response details.
  
          if isinstance(data, predict_pb2.PredictRequest):
              return data
  
          request = predict_pb2.PredictRequest()
          request.model_spec.name = self.model_name
          request.model_spec.signature_name = self.signature_name
  
          input_map = self._create_input_map(data)
  
          for k, v in input_map.items():
              try:
                  request.inputs[k].CopyFrom(v)
              except:
                  raise ValueError("""Unsupported request data format: {}.
                  Valid formats: tensor_pb2.TensorProto and predict_pb2.PredictRequest""".format(type(data)))
  
          return request
  
      def classification(self, data):
          channel = implementations.insecure_channel(self.host, self.tf_serving_port)
          stub = prediction_service_pb2.beta_create_PredictionService_stub(channel)
  
          request = self._create_classification_request(data)
  
          result = stub.Classify(request, self.request_timeout)
  
          return result
  
      def _create_classification_request(self, data):
          if isinstance(data, classification_pb2.ClassificationRequest):
              return data
  
          request = classification_pb2.ClassificationRequest()
          request.model_spec.name = self.model_name
          request.model_spec.signature_name = self.signature_name
  
          feature_dict_list = self._create_feature_dict_list(data)
  
          examples = [_create_tf_example(feature_dict) for feature_dict in feature_dict_list]
  
          request.input.example_list.examples.extend(examples)
  
          return request
  
      def _create_feature_dict_list(self, data):
          """
          Parses the input data and returns a [dict<string, iterable>] which will be used to create the tf examples.
          If the input data is not a dict, a dictionary will be created with the default predict key PREDICT_INPUTS
  
          Examples:
              input                                   => output
              {'age': 39., 'workclass': 'Private'}    => [{'age': 39., 'workclass': 'Private'}]
              [{'age': 39., 'workclass': 'Private'}]  => [{'age': 39., 'workclass': 'Private'}]
              [{'age': 39., 'workclass': 'Private'}, {'age': 39., 'workclass':'Public'}]
                                                      => [{'age': 39., 'workclass': 'Private'},
                                                          {'age': 39., 'workclass': 'Public'}]
              [1, 2, 'string']                        => [{PREDICT_INPUTS: [1, 2, 'string']}]
              42                                      => [{PREDICT_INPUTS: [42]}]
  
  
          Args:
              data: request data. Can be an instance of float, int, str, map, or any iterable object.
  
  
          Returns: a dict[string, iterable] that will be used to create the tf example
  
          """
          if isinstance(data, dict):
              return [data]
          if hasattr(data, '__iter__'):
              if all(isinstance(x, dict) for x in data):
                  return data
              return [{self.input_tensor_name: data}]
          return [{self.input_tensor_name: [data]}]
  
      def _raise_not_implemented_exception(self, data):
          raise NotImplementedError('This prediction service type is not supported by SageMaker yet')
  
      def _create_input_map(self, data):
          """
          Parses the input data and returns a dict<string, TensorProto> which will be used to create the predict request.
          If the input data is not a dict, a dictionary will be created with the default predict key PREDICT_INPUTS
  
          input.
  
          Examples:
              input                                   => output
              {'inputs': tensor_proto}                => {'inputs': tensor_proto}
              tensor_proto                            => {PREDICT_INPUTS: tensor_proto}
              [1,2,3]                                 => {PREDICT_INPUTS: tensor_proto(1,2,3)}
          Args:
              data: request data. Can be any instance of dict<string, tensor_proto>, tensor_proto or any array like data.
  
          Returns:
              dict<string, tensor_proto>
  
  
          """
          msg = """Unsupported request data format: {}.
  Valid formats: tensor_pb2.TensorProto, dict<string,  tensor_pb2.TensorProto> and predict_pb2.PredictRequest"""
  
          if isinstance(data, dict):
              if all(isinstance(v, tensor_pb2.TensorProto) for k, v in data.items()):
                  return data
              raise ValueError(msg.format(data))
  
          if isinstance(data, tensor_pb2.TensorProto):
              return {self.input_tensor_name: data}
  
          try:
              # TODO: tensorflow container supports prediction requests with ONLY one tensor as input
              input_type = self.input_type_map.values()[0]
              ndarray = np.asarray(data)
              tensor_proto = make_tensor_proto(values=ndarray, dtype=input_type, shape=ndarray.shape)
              return {self.input_tensor_name: tensor_proto}
          except:
              raise ValueError(msg.format(data))
  
  
  def _create_tf_example(feature_dict):
      """
      Creates a tf example protobuf message given a feature dict. The protobuf message is defined here
          https://github.com/tensorflow/serving/blob/master/tensorflow_serving/apis/input.proto#L19
      Args:
          feature_dict (dict of str -> feature): feature can be any of the following:
            int, strings, unicode object, float, or list of any of the previous types.
  
      Returns:
          a tf.train.Example including the features
      """
  
      def _create_feature(feature):
          feature_list = feature if isinstance(feature, list) else [feature]
  
          # Each feature can be exactly one kind:
          # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/example/feature.proto#L76
  
          feature_type = type(feature_list[0])
          if feature_type == int:
              return feature_pb2.Feature(int64_list=feature_pb2.Int64List(value=feature_list))
          elif feature_type == str:
              return feature_pb2.Feature(bytes_list=feature_pb2.BytesList(value=feature_list))
          elif feature_type == unicode:
              return feature_pb2.Feature(bytes_list=feature_pb2.BytesList(value=map(lambda x: str(x), feature_list)))
          elif feature_type == float:
              return feature_pb2.Feature(float_list=feature_pb2.FloatList(value=feature_list))
          else:
              message = """Unsupported request data format: {}, {}.
                              Valid formats: float, int, str any object that implements __iter__
                                             or classification_pb2.ClassificationRequest"""
              raise ValueError(message.format(feature, type(feature)))
  
      features = {k: _create_feature(v) for k, v in feature_dict.items()}
      return example_pb2.Example(features=feature_pb2.Features(feature=features))
  