# coding=utf-8
# Licensed Materials - Property of IBM
# Copyright IBM Corp. 2020

import os
import streamsx.spl.op
import streamsx.spl.types
from streamsx.topology.schema import CommonSchema, StreamSchema
from streamsx.spl.types import rstring
import json

PRIVATE_WML_TOOLKIT_NAMESPACE = "com.ibm.streams.wml"
PRIVATE_WML_TOOLKIT_LOCATION = "spl/toolkit.wml/com.ibm.streams.wml"

def _add_toolkit(topo):
    # specific private toolkit used in this wrapper
    # info.xml is an optional element -- so we don't provide for our private toolkit
    # toolkit.xml is generated during topology/application build -- so we don't care
    # So the only think is to provide the correct TK structure containing just the 
    # python file defining the primitive_operator, operator xml is also generated by topology/build
    # the toolkit is part of this streamsx.wml python package
    streamsx.spl.toolkit.add_toolkit(topo, os.path.dirname(os.path.realpath(__file__))+'/'+PRIVATE_WML_TOOLKIT_LOCATION)


def wml_online_scoring( stream, 
                        deployment_guid, 
                        #mapping_function,
                        field_mapping,
                        credentials, 
                        space_guid, 
                        expected_load = 1000, 
                        queue_size = 2000,  
                        threads_per_node = 2, 
                        connectionConfiguration=None, 
                        name = None):
    """Scores tuples received from input stream using the online scoring endpoint of WML referenced by the deployment_guid.
    Models can be created, trained and deployed in Watson Studio, using notebooks resp. other tools in Cloud Pak for Data.

    Args:
        stream              input stream
        deployment_guid     GUID of the WML deplyoment to be used to score the input data
        field_mapping       list (or JSOON array) of mappings between model fields and input tuple elements
                            [{'model_field':str, 'is_mandatory':bool,'tuple_field':str},{...}]
        credentials         dict or (JSON object) of WML credentials to be used
                            {'url':str,'token':str,'instance_id': str,'version':str}
        space_guid          GUID of the used deployment space
        expected_load       optional field, giving the expected tuple throughput
                            used to determine a maximum number of tuple to be sent in one request to 
                            WML online scoring, maximum number = expected_load / threads_per_node
                            this value has most impact on throughput performance
                            default:1000 
        queue_size          optional field to set the internal buffer size after which back-pressure happens
                            default:2000
        threads_per_node    optional field to set the number of threads used to process the received tuples
                            may increase throughput performance by using free resources while other thread(s) waiting
                            for the asynchronous result from online scoring
                            default: 2 
        name                optional field to give the resulting Streams operator a name of your choice
    Returns:
        tuple(result_stream, error_stream)
        
            result_stream: Tuples as received on input extended by the scoring result field 'predictions' of type dict
            error_stream : Tuples as received on input extended by the error indication 'scoring_error' of type dict
    """
    # add private toolkit
    topology = stream.topology
    _add_toolkit(topology)
    
    #check input parameter

    # field_mapping has to be either a list
    if isinstance(field_mapping,list):
        field_mapping_json = json.dumps(field_mapping)
    # may be also a JSON array
    elif isinstance(field_mapping,str):
        try:
            test_load=json.loads(field_mapping)
            field_mapping_json = field_mapping
        except:
            print ("wml_online_scoring() parameter 'field_mapping' is not a valid json")    
    else:
        print ("wml_online_scoring() parameter 'field_mapping' has to be either a 'list' or a valid JSON string of mappings")
    

    # credentials has to be either dict
    if isinstance(credentials,dict):
        credentials_json = json.dumps(credentials)
    # or a JSON string
    elif isinstance(credentials,str):
        try:
            test_load=json.loads(credentials)
            credentials_json = credentials
        except:
            print ("wml_online_scoring() parameter 'credentials' is not a valid json")
            
    else:
        print ("wml_online_scoring() parameter 'credentials' has to be either a 'dict' or a valid JSON string of credentials")



    # create instance of wrapper class
    _op = _WMLOnlineScoring(stream = stream, 
                            deployment_guid = deployment_guid, 
                            #mapping_function = mapping_function,
                            field_mapping = field_mapping_json,
                            credentials = credentials_json, 
                            space_guid = space_guid, 
                            expected_load = expected_load, 
                            queue_size = queue_size,  
                            threads_per_node = threads_per_node, 
                            connectionConfiguration = connectionConfiguration, 
                            name = name)

    # calling SPL operators will result anytime in schema based output streams
    # these need to be mapped back to the Python object Stream we received

    return _op.outputs[0],_op.outputs[1]




class _WMLOnlineScoring(streamsx.spl.op.Invoke):
    def __init__(self, stream, 
                       deployment_guid, 
                       #mapping_function,
                       field_mapping,
                       credentials, 
                       space_guid, 
                       expected_load, 
                       queue_size, 
                       threads_per_node, 
                       connectionConfiguration, 
                       name):

        topology = stream.topology
        kind="com.ibm.streams.wml::WMLOnlineScoring"
        inputs=[stream]
        schemas=[object,object]  #not [StreamSchema('tuple<blob __spl_po>'),StreamSchema('tuple<blob __spl_po>')]
        params = dict()
        params['deployment_guid'] = deployment_guid
        params['wml_credentials'] = credentials
        #params['mapping_function'] = mapping_function
        params['field_mapping'] = field_mapping,
        params['space_guid'] = space_guid
        params['expected_load'] = expected_load
        params['queue_size'] = queue_size
        params['threads_per_node'] = threads_per_node

        super(_WMLOnlineScoring, self).__init__(topology,kind,inputs,schemas,params,name)




