from __future__ import absolute_import
import numpy as np
from .. import ndarray
from .. import stream
from ..context import get_current_context, DeviceGroup
G_NODE_ID = 0


class Op(object):
    """Node in a computation graph."""

    def __init__(self, op_type, inputs, ctx=None):
        """Constructor
            Instance variables
            ------------------
            self.inputs: the list of input nodes.
            self.const_attr: the add or multiply constant.
                e.g. self.const_attr=5 if this node is created by x+5.
            self.name: node name for debugging.
        """
        self.inputs = inputs
        self.raw_ctx = get_current_context() if ctx is None else DeviceGroup(ctx)
        self.ctx = ctx
        self.const_attr = None
        self.dtype = None
        self.inplace = False
        self.lazy_execution = False
        self.event = None
        self.op_type = op_type.__name__
        global G_NODE_ID
        self.id = G_NODE_ID
        G_NODE_ID = G_NODE_ID + 1
        self.name = self.op_type + str(self.id)
        self.desc = self.name + \
            '(' + ', '.join([inp.name for inp in inputs]) + ')'

    def __add__(self, other):
        """Adding two nodes return a new node."""
        from .AddElewise import add_op
        from .AddConst import addbyconst_op

        # here the operator does NOT specify context
        # please explicitly specify the context in gradients!!
        if isinstance(other, Op):
            new_node = add_op(self, other)
        else:
            # Add by a constant stores the constant in new node's const_attr
            # 'other' argument is a constant
            new_node = addbyconst_op(self, other)
        return new_node

    def __mul__(self, other):
        """Multiplying two nodes return a new node."""
        from .MultiplyElewise import mul_op
        from .MultiplyConst import mul_byconst_op

        if isinstance(other, Op):
            new_node = mul_op(self, other)
        else:
            # Mul by a constant stores the constant in new node's const_attr
            # 'other' argument is a constant
            new_node = mul_byconst_op(self, other)
        return new_node

    # Allow left-hand-side add and multiply.
    __radd__ = __add__
    __rmul__ = __mul__

    def __str__(self):
        """Allow print to display node name."""
        return self.name

    def compute(self, input_vals, output_val, stream_handle=None):
        """Given values of input nodes, compute the output value.
        Parameters
        ----------
        node: node that performs the compute.
        input_vals: values of input nodes.
        output_val: output value of the node, modified in-place.
        """
        raise NotImplementedError

    def gradient(self, output_grad):
        """Given output gradient, compute partial gradient to each input node.
        Parameters
        ----------
        node: node that performs the gradient.
        output_grad: output gradient summed from children nodes' contributions
        Returns
        -------
        A list of gradient contributions to each input node respectively.
        """
        raise NotImplementedError

    def infer_shape(self, input_shapes):
        """Given shapes of input nodes, compute shape of output node.
        Implementation note:
        It's simpler to treat shape of constants as (1,), so that constants can
        be stored as a numpy array too and you would need fewer special case
        handling.
        Parameters
        ----------
        node: node whose shape is being inferred.
        input_vals: shapes of input nodes.
        Returns
        -------
        A tuple representing the shape of output node.
        """
        raise NotImplementedError

    def add_transfer_op(self, src_node, dst_ctx, h2d_ops, d2h_ops):
        from .DataTransfer import datah2d_op, datad2h_op, datad2h_sparse_op

        def add_h2d(prev_node, cur_ctx):
            if prev_node not in h2d_ops:
                h2d_ops[prev_node] = datah2d_op(prev_node, cur_ctx)
            return h2d_ops[prev_node]

        def add_d2h(prev_node):
            from .EmbeddingLookUp import EmbeddingLookUp_Gradient
            if prev_node not in d2h_ops:
                if isinstance(prev_node, EmbeddingLookUp_Gradient):
                    d2h_ops[prev_node] = datad2h_sparse_op(prev_node)
                else:
                    d2h_ops[prev_node] = datad2h_op(prev_node)
                if prev_node.event is None:
                    # here we should ensure the computation complete before d2h
                    prev_node.event = stream.create_event_handle(prev_node.ctx)
            return d2h_ops[prev_node]
        src_ctx = src_node.ctx
        result = src_node
        if src_ctx != dst_ctx:
            if ndarray.is_gpu_ctx(dst_ctx):
                if ndarray.is_gpu_ctx(src_ctx):
                    assert False, 'Please use NCCL to P2P communicate!'
                else:
                    result = add_h2d(result, dst_ctx)
            else:
                result = add_d2h(result)
        return result

    def forward_hook(self, config):
        # disable inplace if not lazy execution
        # previously we use array reshape lazy callback to do this, which is deprecated (not efficient)
        if not self.lazy_execution:
            for node in self.inputs:
                node.inplace = False

        # insert data transfer op if needed
        input_ctxs = set([n.ctx for n in self.inputs])
        assert None not in input_ctxs, 'Inputs contexts should already be determined.'
        if self.ctx is None:
            self.ctx = config.context
        for i in range(len(self.inputs)):
            self.inputs[i] = self.add_transfer_op(
                self.inputs[i], self.ctx, config.h2d_ops, config.d2h_ops)
        self.on_gpu = ndarray.is_gpu_ctx(self.ctx)
        self.on_cpu = not self.on_gpu
        if self in config.eval_node_list and self.on_gpu and self.event is None:
            self.event = stream.create_event_handle(self.ctx)

    def backward_hook(self, config):
        pass

    def deduce_states(self, input_states, input_duplicates):
        assert len(input_states) == len(self.inputs)
        assert len(input_states) == len(input_duplicates)
        if len(input_states) == 1:
            return input_states[0], input_duplicates[0]
        else:
            assert all([x is None or x == (1, 1) for x in input_states])
            return None, 1