"""
Created on Wed Aug 19 03:47:54 2020

@author: manal khallaayoune
"""

import numpy as np
import GPy
from GPy.kern.src.stationary import Stationary


class ExpSin(Stationary):
    """
    Exponential of Sinus kernel: 
        Product of 1D Exponential of Sinus kernels

    .. math::

        &k(x,x')_i = \sigma^2 \prod_{j=1}^{dimension} \exp \\bigg( - \\frac{ \sin ( x_{i,j}-x_{i,j}' ) ^2}{2 \ell_j^2} \\bigg)
        
        &x,x' \in \mathcal{M}_{n,dimension}
        
        &k \in \mathcal{M}_{n,n}

    """
    def __init__(self, input_dim, variance=1., lengthscale=None, ARD=False, active_dims=None, name='ExpSin'):
        super(ExpSin, self).__init__(input_dim, variance, lengthscale, ARD, active_dims, name)

    def K_of_r(self, dist):
        n = dist.shape[2]
        s = 0
        for k in range(n):
            s+= np.sin(dist[:,:,k])**2
        return self.variance * np.exp(-s/(2*self.lengthscale**2))
    
    def K(self, X, X2):
        dist = X[:,None,:]-X2[None,:,:]
        return self.K_of_r(dist)

    def dK_dr(self,dist,dimX):
        K = self.K_of_r(dist)
        d1 = dist[:,:,dimX]
        return -K*np.sin(d1)*np.cos(d1)/self.lengthscale**2
    
    def dK_dX(self, X, X2, dimX):
        dist = X[:,None,:]-X2[None,:,:]
        dK_dr = self.dK_dr(dist,dimX)
        return dK_dr
    
    def dK_dX2(self,X,X2,dimX2):
        return -self.dK_dX(X,X2, dimX2)
    
    def dK2_dXdX2(self, X, X2, dimX, dimX2):
        dist = X[:,None,:]-X2[None,:,:]
        K = self.K_of_r(dist)
        dK_dX = self.dK_dX(X,X2,dimX)
        lengthscale2inv = np.ones((X.shape[1]))/(self.lengthscale**2)
        l1 = lengthscale2inv[dimX]
        l2 = lengthscale2inv[dimX2]
        d1 = dist[:,:,dimX]
        d2 = dist[:,:,dimX2]
        s1 = np.sin(d1)
        c1 = np.cos(d1)
        return (dimX!=dimX2)*dK_dX*l2*np.sin(d2)*np.cos(d2) + (dimX==dimX2)*l1*(dK_dX*s1*c1+K*c1**2-K*s1**2)
    
