from __future__ import absolute_import
from .Node import Op
from .._base import DNNL_LIB
from ..cpu_links import matrix_elementwise_multiply_by_const as cpu_matrix_elementwise_multiply_by_const
from ..gpu_links import matrix_elementwise_multiply_by_const


class MulByConstOp(Op):
    def __init__(self, node_A, const_val, ctx=None):
        super().__init__(MulByConstOp, [node_A], ctx)
        self.const_attr = const_val
        self.desc = self.name + '(%s, %s)' % (node_A.name, str(const_val))

    def compute(self, input_vals, output_val, stream_handle=None):
        assert self.const_attr is not None
        if self.on_cpu:
            if DNNL_LIB['DnnlMatrixElementwiseMultiplyByConst']:
                cpu_matrix_elementwise_multiply_by_const(
                    input_vals[0], self.const_attr, output_val)
            else:
                output_val[:] = input_vals[0].asnumpy() * self.const_attr
        else:
            matrix_elementwise_multiply_by_const(
                input_vals[0], self.const_attr, output_val, stream_handle)

    def gradient(self, output_grad):
        return [self.const_attr * output_grad]

    def infer_shape(self, input_shapes):
        assert len(input_shapes) == 1
        return input_shapes[0]


def mul_byconst_op(node_A, const_val, ctx=None):
    """Make a new instance of MulByConstOp and call the instance.

    Parameters:
    ----
    node : Node
        The Node to be multiplied.
    const_val : scalar value
        The constant value to be mutiplied.

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

    """
    return MulByConstOp(node_A, const_val, ctx=ctx)