from __future__ import absolute_import
from .Node import Op
from .._base import DNNL_LIB
from ..cpu_links import matrix_elementwise_multiply as\
    cpu_matrix_elementwise_multiply
from ..cpu_links import matrix_elementwise_multiply_by_const as\
    cpu_matrix_elementwise_multiply_by_const
from ..gpu_links import matrix_elementwise_multiply,\
    matrix_elementwise_multiply_by_const


class MulOp(Op):
    def __init__(self, node_A, node_B, ctx=None):
        super().__init__(MulOp, [node_A, node_B], ctx)

    def compute(self, input_vals, output_val, stream_handle=None):
        if self.on_cpu:
            if DNNL_LIB['DnnlMatrixElementwiseMultiply'] and input_vals[0].shape == input_vals[1].shape:
                cpu_matrix_elementwise_multiply(
                    input_vals[0], input_vals[1], output_val)
            elif DNNL_LIB['DnnlMatrixElementwiseMultiplyByConst'] and (input_vals[0].shape == (1,) or input_vals[1].shape == (1,)):
                if input_vals[1].shape == (1,):
                    const_val = input_vals[1].asnumpy()[0]
                    cpu_matrix_elementwise_multiply_by_const(
                        input_vals[0], const_val, output_val)
                elif input_vals[0].shape == (1,):
                    const_val = input_vals[0].asnumpy()[0]
                    cpu_matrix_elementwise_multiply_by_const(
                        input_vals[1], const_val, output_val)
            else:
                output_val[:] = input_vals[0].asnumpy() * \
                    input_vals[1].asnumpy()
        else:
            if input_vals[0].shape == input_vals[1].shape:
                matrix_elementwise_multiply(
                    input_vals[0], input_vals[1], output_val, stream_handle)
            else:
                if input_vals[1].shape == (1,):
                    const_val = input_vals[1].asnumpy()[0]
                    matrix_elementwise_multiply_by_const(
                        input_vals[0], const_val, output_val, stream_handle)
                elif input_vals[0].shape == (1,):
                    const_val = input_vals[0].asnumpy()[0]
                    matrix_elementwise_multiply_by_const(
                        input_vals[1], const_val, output_val, stream_handle)

    def gradient(self, output_grad):
        return [mul_op(self.inputs[1], output_grad, ctx=self.raw_ctx),
                mul_op(self.inputs[0], output_grad, ctx=self.raw_ctx)]

    def infer_shape(self, input_shapes):
        """Need to handle input_vals[0].shape != input_vals[1].shape"""
        assert len(input_shapes) == 2
        if input_shapes[0] == input_shapes[1]:
            output = input_shapes[0]
        else:
            if input_shapes[0] == (1,):
                output = input_shapes[1]
            elif input_shapes[1] == (1,):
                output = input_shapes[0]
            else:
                assert False, "can't do elementwise multiply between variables of different sizes."
        return output


def mul_op(node_A, node_B, ctx=None):
    """Make a new instance of matrixs elementwise multiplication and call the instance.

    Parameters:
    ----
    node_a : Node
        The Node to be multiplied.
    node_b : Node
        Another Node to be multiplied.

    Returns:
    ----
    A new Node instance created by Op.

    """
    return MulOp(node_A, node_B, ctx=ctx)