from .kernel import CombinationKernel


class AddKernel(CombinationKernel):
    """
    The add kernel that computes a covariance matrix by suming the covariance
    matrices of a list of kernels.

    :param sub_kernels: a list of kernels that are combined to compute a covariance matrix.
    :type sub_kernels: [Kernel]
    :param name: the name of the kernel. The name is used to access kernel parameters.
    :type name: str
    :param dtype: the data type for float point numbers.
    :type dtype: numpy.float32 or numpy.float64
    :param ctx: the mxnet context (default: None/current context).
    :type ctx: None or mxnet.cpu or mxnet.gpu
    """
    def __init__(self, sub_kernels, name='add', dtype=None, ctx=None):
        kernels = []
        for k in sub_kernels:
            if isinstance(k, AddKernel):
                for k2 in k.sub_kernels:
                    kernels.append(k2)
            else:
                kernels.append(k)
        super(AddKernel, self).__init__(
            sub_kernels=kernels, name=name, dtype=dtype, ctx=ctx)

    def _compute_K(self, F, X, X2=None, **kernel_params):
        """
        The internal interface for the actual covariance matrix computation.

        This function takes as an assumption: The prefix in the keys of
        *kernel_params* that corresponds to the name of the kernel has been
        removed. The dimensions of *X* and *X2* have been sliced according to
        *active_dims*.

        :param F: MXNet computation type <mx.sym, mx.nd>.
        :param X: the first set of inputs to the kernel.
        :type X: MXNet NDArray or MXNet Symbol
        :param X2: (optional) the second set of arguments to the kernel. If X2 is None, this computes a square covariance matrix of X. In other words,
            X2 is internally treated as X.
        :type X2: MXNet NDArray or MXNet Symbol
        :param \**kernel_params: the set of kernel parameters, provided as keyword arguments.
        :type **kernel_params: {str: MXNet NDArray or MXNet Symbol}
        :return: The covariance matrix.
        :rtype: MXNet NDArray or MXNet Symbol
        """
        K = self.sub_kernels[0].K(F=F, X=X, X2=X2, **kernel_params)
        for k in self.sub_kernels[1:]:
            K = K + k.K(F=F, X=X, X2=X2, **kernel_params)
        return K

    def _compute_Kdiag(self, F, X, **kernel_params):
        """
        The internal interface for the actual computation for the diagonal of the covariance matrix.

        This function takes as an assumption: The prefix in the keys of *kernel_params* that corresponds to the name of the kernel has been
        removed. The dimensions of *X* has been sliced according to *active_dims*.

        :param F: MXNet computation type <mx.sym, mx.nd>.
        :param X: the first set of inputs to the kernel.
        :type X: MXNet NDArray or MXNet Symbol
        :param \**kernel_params: the set of kernel parameters, provided as keyword arguments.
        :type \**kernel_params: {str: MXNet NDArray or MXNet Symbol}
        :return: The covariance matrix.
        :rtype: MXNet NDArray or MXNet Symbol
        """
        K = self.sub_kernels[0].Kdiag(F=F, X=X, **kernel_params)
        for k in self.sub_kernels[1:]:
            K = K + k.Kdiag(F=F, X=X, **kernel_params)
        return K
