import numpy as np



from src.impl.tf import get_tf_output

from src.impl.pytorch import get_torch_output

from src.impl.mnn import get_mnn_output



import tensorflow as tf





def get_output(tensor, operator, variable=None):

    if operator == 'avg_pool':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'max_pool':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'bias_add':

        if variable is None:

            variable = tf.random.normal([np.asarray(tensor).shape[-1]], stddev=0.01).numpy()

            return get_tf_output(tensor, operator, variable).numpy(), \

                get_torch_output(tensor, operator, variable).numpy(), \

                get_mnn_output(tensor, operator, variable).read(), variable

        else:

            return get_tf_output(tensor, operator, variable).numpy(), \

                get_torch_output(tensor, operator, variable).numpy(), \

                get_mnn_output(tensor, operator, variable).read(), None

    elif operator == 'conv2d':

        if variable is None:

            variable = tf.random.normal([3, 3, tensor.shape[-3], 32], stddev=0.01).numpy()

            np.save('./data/variable', variable)

            return get_tf_output(tensor, operator, variable).numpy(), \

                get_torch_output(tensor, operator, variable).detach().numpy(), \

                get_mnn_output(tensor, operator, variable).read(), variable

        else:

            return get_tf_output(tensor, operator, variable).numpy(), \

                get_torch_output(tensor, operator, variable).detach().numpy(), \

                get_mnn_output(tensor, operator, variable).read(), None

    elif operator == 'softmax':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'batch_normalization':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'relu':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'reduce_mean':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'reduce_max':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'reduce_sum':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'reduce_min':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'sigmoid':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'tanh':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'abs':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'square':

        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \

               get_mnn_output(tensor, operator).read(), None

    elif operator == 'dense':

        if variable is None:

            variable = tf.random.normal([16, 10], stddev=0.01).numpy()

            return get_tf_output(tensor, operator, variable).numpy(), \

                   get_torch_output(tensor, operator, variable).numpy(), \

                   get_mnn_output(tensor, operator, variable).read(), variable

        else:

            return get_tf_output(tensor, operator, variable).numpy(), \

                   get_torch_output(tensor, operator, variable).numpy(), \

                   get_mnn_output(tensor, operator, variable).read(), None





def get_output_float16(tensor, operator, variable):

    tensor_16 = tensor.astype(np.float16)

    return get_output(tensor_16, operator, variable)





def get_output_float32(tensor, operator):

    tensor_32 = tensor.astype(np.float32)

    return get_output(tensor_32, operator)





def get_output_float32_with_variable(tensor, operator, variable):

    tensor_32 = tensor.astype(np.float32)

    return get_output(tensor_32, operator, variable)