import os, sys
sys.path.append(os.path.split(os.path.realpath(__file__))[0] + "/../lib")
from cnstream import *
from cnstream_cpptest import *

# in order to determine whether the python function is called by cpp
init_called = False
execute_called = False
postproc_params = None
obj_init_called = False
obj_execute_called = False
obj_postproc_params = None
received_input_shapes = None
received_output_shape = None

class CustomPostproc(Postproc):
    def __init__(self):
        Postproc.__init__(self)

    def init(self, params):
        global init_called
        global postproc_params 
        init_called = True
        postproc_params = params
        return True

    def execute(self, net_outputs, input_shapes, finfo):
        global execute_called
        execute_called = True
        global received_input_shapes
        global received_output_shape
        received_input_shapes = input_shapes
        received_output_shape = net_outputs[0].shape

class CustomObjPostproc(ObjPostproc):
    def __init__(self):
        ObjPostproc.__init__(self)

    def init(self, params):
        global obj_init_called
        global obj_postproc_params 
        obj_init_called = True
        obj_postproc_params = params
        return True

    def execute(self, net_outputs, input_shapes, finfo, obj):
        global obj_execute_called
        obj_execute_called = True
        global received_input_shapes
        global received_output_shape
        received_input_shapes = input_shapes
        received_output_shape = net_outputs[0].shape

class TestPostproc:
    def test_postproc(self):
        params = {'pyclass_name' : 'test.postproc_test.CustomPostproc', 'param' : 'value'}
        assert cpptest_pypostproc(params)
        # test cpp call python init function success
        assert init_called
        # test custom parameters from cpp pass to python success
        assert postproc_params['param'] == 'value'
        # test cpp call python execute function success
        assert execute_called
        # check I/O shapes
        expected_input_shapes = [[4, 160, 40, 4]]
        expected_output_shape = (20, 1, 84)
        assert expected_input_shapes == received_input_shapes
        assert expected_output_shape == received_output_shape

    def test_obj_postproc(self):
        params = {'pyclass_name' : 'test.postproc_test.CustomObjPostproc', 'param' : 'value'}
        assert cpptest_pyobjpostproc(params)
        # test cpp call python init function success
        assert obj_init_called
        # test custom parameters from cpp pass to python success
        assert obj_postproc_params['param'] == 'value'
        # test cpp call python execute function success
        assert obj_execute_called
        # check I/O shapes
        expected_input_shapes = [[4, 160, 40, 4]]
        expected_output_shape = (20, 1, 84)
        assert expected_input_shapes == received_input_shapes
        assert expected_output_shape == received_output_shape