# Copyright (c) 2012 James Hensman
# Licensed under the BSD 3-clause license (see LICENSE.txt)

from kernpart import kernpart
import numpy as np

class symmetric(kernpart):
    """
    Symmetrical kernels

    :param k: the kernel to symmetrify
    :type k: kernpart
    :param transform: the transform to use in symmetrification (allows symmetry on specified axes)
    :type transform: A numpy array (D x D) specifiying the transform
    :rtype: kernpart

    """
    def __init__(self,k,transform=None):
        if transform is None:
            transform = np.eye(k.D)*-1.
        assert transform.shape == (k.D, k.D)
        self.transform = transform
        self.D = k.D
        self.Nparam = k.Nparam
        self.name = k.name + '_symm'
        self.k = k
        self._set_params(k._get_params())

    def _get_params(self):
        """return the value of the parameters."""
        return self.k._get_params()

    def _set_params(self,x):
        """set the value of the parameters."""
        self.k._set_params(x)

    def _get_param_names(self):
        """return parameter names."""
        return self.k._get_param_names()

    def K(self,X,X2,target):
        """Compute the covariance matrix between X and X2."""
        AX = np.dot(X,self.transform)
        if X2 is None:
            X2 = X
            AX2 = AX
        else:
            AX2 = np.dot(X2, self.transform)
        self.k.K(X,X2,target)
        self.k.K(AX,X2,target)
        self.k.K(X,AX2,target)
        self.k.K(AX,AX2,target)

    def dK_dtheta(self,dL_dK,X,X2,target):
        """derivative of the covariance matrix with respect to the parameters."""
        AX = np.dot(X,self.transform)
        if X2 is None:
            X2 = X
            ZX2 = AX
        else:
            AX2 = np.dot(X2, self.transform)
        self.k.dK_dtheta(dL_dK,X,X2,target)
        self.k.dK_dtheta(dL_dK,AX,X2,target)
        self.k.dK_dtheta(dL_dK,X,AX2,target)
        self.k.dK_dtheta(dL_dK,AX,AX2,target)


    def dK_dX(self,dL_dK,X,X2,target):
        """derivative of the covariance matrix with respect to X."""
        AX = np.dot(X,self.transform)
        if X2 is None:
            X2 = X
            ZX2 = AX
        else:
            AX2 = np.dot(X2, self.transform)
        self.k.dK_dX(dL_dK, X, X2, target)
        self.k.dK_dX(dL_dK, AX, X2, target)
        self.k.dK_dX(dL_dK, X, AX2, target)
        self.k.dK_dX(dL_dK, AX ,AX2, target)

    def Kdiag(self,X,target):
        """Compute the diagonal of the covariance matrix associated to X."""
        foo = np.zeros((X.shape[0],X.shape[0]))
        self.K(X,X,foo)
        target += np.diag(foo)

    def dKdiag_dX(self,dL_dKdiag,X,target):
        raise NotImplementedError

    def dKdiag_dtheta(self,dL_dKdiag,X,target):
        """Compute the diagonal of the covariance matrix associated to X."""
        raise NotImplementedError