# 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 = 0, 
                        queue_size = 2000,  
                        threads_per_node = 2,
                        single_output = False,
                        node_count = 1, 
                        connectionConfiguration=None, 
                        name = None,
                        bundle_size = 100):
    """
    Scoring tuples received from input :py:class:`stream` using the online scoring endpoint of WML referenced by the :py:class:`deployment_guid`.
    Mapping from input attributes to the models minning fields is done by the :py:class:`field_mapping` parameter.


    Args:
        stream (:py:class:`topology_ref:streamsx.topology.topology.Stream`):         
            input stream of data to be scored
        deployment_guid (str): 
            GUID of the WML deplyoment to be used to score the input data
        field_mapping (list|str): 
            list or JSON string representing the
            mappings between model fields and input tuple elements
            [{'model_field':str, 'is_mandatory':bool,'tuple_field':str},{...}]
        credentials (dict|str): 
            dict or JSON string of WML credentials to be used
            {'url':<cp4d_cluster_url>,'token':<token_of authenticated_user>,'instance_id': 'wml_local','version':'2.5.0'}
        space_guid (str): 
            GUID of the used deployment space
        expected_load (int, optional): 
            depricated, use new parameter bundle_size !
            The expected tuple throughput which is
            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, defaults to None
        queue_size (int, optional): 
            The internal buffer size after which back-pressure happens, defaults to 2000
        threads_per_node (int, optional): 
            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, defaults to 2 
        single_output (bool, optional): 
            optional field to define if error data and success data should be send 
            to one output (stream) or to 2 different outputs (streams), defaults to False
        node_count(int, optional):
            optional field giving the number of REST nodes which share the load, defaults to 1
        name (str, optional): 
            Give the resulting Streams operator a name of your choice
        bundle_size (int, optional): 
            optional field for setting maximum number of tuple to be sent in one scoring request
            this value has most impact on throughput performance, defaults to 100
        
    Returns:
        result_stream, error_stream(:py:class:`topology_ref:streamsx.topology.topology.Stream`, :py:class:`topology_ref:streamsx.topology.topology.Stream`):
            Returns streams for further processing depending on ''single_output'' setting.
            
        ''single_output'' = False
            :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
            
        ''single_output'' = True:
            :result_stream: all Tuples as received on input, extended by the scoring result field ''prediction'' of type dict
                                or by field ''prediction_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:
            raise Exception("wml_online_scoring() parameter 'field_mapping' is not a valid json")    
    else:
        raise Exception("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:
            raise Exception("wml_online_scoring() parameter 'credentials' is not a valid json")
            
    else:
        raise Exception("wml_online_scoring() parameter 'credentials' has to be either a 'dict' or a valid JSON string of credentials")

    # either expected_load(depricated) or bundle_size(defaults to 100) have to be set
    # For backward compatibility!!!
    #     if expected_load is set use this and not bundle_size
    _bundle_size = 0
    if expected_load is 0: 
        if isinstance(bundle_size, int):
            _bundle_size = bundle_size
        else:
            raise Exception("wml_online_scoring() parameter 'bundle_size' has to be an integer")
    elif not isinstance(expected_load,int):
        raise Exception("Depricated wml_online_scoring() parameter 'expected_load' has to be an integer, consider using the new parameter 'bundle_size'")
    else:
        print ("wml_online_scoring() parameter 'expected_load' is depricated but still supported, consider using the new parameter 'bundle_size'")


    # 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,
                            single_output = single_output, 
                            node_count = node_count,
                            connectionConfiguration = connectionConfiguration, 
                            name = name,
                            bundle_size = _bundle_size)

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

    if single_output:
        return _op.outputs[0]
    else:
        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,
                       single_output,
                       node_count, 
                       connectionConfiguration, 
                       name,
                       bundle_size):

        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
        params['single_output'] = single_output
        params['node_count'] = node_count
        params['bundle_size'] = bundle_size

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




