from __future__ import absolute_import
import numpy as np
from .Node import Op
from ..gpu_links import one_hot


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

    def compute(self, input_vals, output_val, stream_handle=None):
        if self.on_cpu:
            inputs = input_vals[0].asnumpy().astype(np.int)
            res = np.eye(self.num_classes)[inputs.reshape(-1)]
            output_val[:] = res.reshape(
                list(inputs.shape) + [self.num_classes]).astype(np.float32)
        else:
            one_hot(input_vals[0], output_val, stream_handle)

    def gradient(self, output_grad):
        return [None]

    def infer_shape(self, input_shapes):
        assert len(input_shapes) == 1
        return tuple(list(input_shapes[0]) + [self.num_classes])


def one_hot_op(node, num_classes, ctx=None):
    """Creates a node that represents one hot.

    Parameters:
    ----
    node : Node
        The input Node.
    num_classes: int
        Number of classes.

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

    """
    return OneHotOp(node, num_classes, ctx=ctx)