'''
Wavelet kernel
slice allows kernel operation on feature subset
active_dims is iterable of feature dimensions to extract
input_dim must equal dimension defined by active_dims
'''

import numpy as np
import tensorflow as tf

from .. import util
from . import kernel
from .kernel_extras import *


class WaveletSlice(kernel.Kernel):

    def __init__(self, input_dim, active_dims=None, shift=0, scale = 0.01,
                 white=0.01, input_scaling=False):
        if input_scaling:
            self.shift = tf.Variable(shift * tf.ones([input_dim]))
            self.scale = tf.Variable(scale * tf.ones([input_dim]))
        else:
            self.shift = tf.Variable([shift], dtype=tf.float32)
            self.scale = tf.Variable([scale], dtype=tf.float32)
        self.input_dim = input_dim
        self.active_dims = active_dims
        self.white = white

    def kernel(self, points1, points2=None):
        if points2 is None:
            points2 = points1
            white_noise = (self.white * util.eye(tf.shape(points1)[0]) +
                0.1 * self.white * tf.ones( [tf.shape(points1)[0], tf.shape(points1)[0]]))
        else:
            white_noise = 0.01 * self.white * tf.ones( [tf.shape(points1)[0], tf.shape(points2)[0]] )

        points1, points2 = dim_slice(self, points1, points2)

        def h(x):
            # Zhang wavelet
            #return tf.cos(1.75*x)*tf.exp(-0.5*x**2)
            # mexican hat wavelet
            return (1-x**2)*tf.exp(-0.5*x**2)

        kern1, kern2 = h((points1 - self.shift)/tf.exp(self.scale)), h((points2 - self.shift)/tf.exp(self.scale))
        kern1, kern2 = tf.reduce_prod(kern1, axis=1), tf.reduce_prod(kern2, axis=1)
        kern = tf.einsum('i,j->ij', kern1, kern2)
        return kern + white_noise

    def diag_kernel(self, points):
        def h(x):
            # Zhang wavelet
            return tf.cos(1.75*x)*tf.exp(-0.5*x**2)
            # mexican hat wavelet
            #return (1-x**2)*tf.exp(-0.5*x**2)

        points = dim_slice_diag(self, points)
        kern = tf.reduce_prod(h((points - self.shift)/tf.exp(self.scale)) , axis=1) **2
        return kern + self.white

    def get_params(self):
        return [self.shift, self.scale]

[((1704, 1738), 'tensorflow.einsum', 'tf.einsum', (['"""i,j->ij"""', 'kern1', 'kern2'], {}), "('i,j->ij', kern1, kern2)\n", (1713, 1738), True, 'import tensorflow as tf\n'), ((660, 698), 'tensorflow.Variable', 'tf.Variable', (['[shift]'], {'dtype': 'tf.float32'}), '([shift], dtype=tf.float32)\n', (671, 698), True, 'import tensorflow as tf\n'), ((724, 762), 'tensorflow.Variable', 'tf.Variable', (['[scale]'], {'dtype': 'tf.float32'}), '([scale], dtype=tf.float32)\n', (735, 762), True, 'import tensorflow as tf\n'), ((1628, 1657), 'tensorflow.reduce_prod', 'tf.reduce_prod', (['kern1'], {'axis': '(1)'}), '(kern1, axis=1)\n', (1642, 1657), True, 'import tensorflow as tf\n'), ((1659, 1688), 'tensorflow.reduce_prod', 'tf.reduce_prod', (['kern2'], {'axis': '(1)'}), '(kern2, axis=1)\n', (1673, 1688), True, 'import tensorflow as tf\n'), ((1472, 1493), 'tensorflow.exp', 'tf.exp', (['(-0.5 * x ** 2)'], {}), '(-0.5 * x ** 2)\n', (1478, 1493), True, 'import tensorflow as tf\n'), ((1874, 1890), 'tensorflow.cos', 'tf.cos', (['(1.75 * x)'], {}), '(1.75 * x)\n', (1880, 1890), True, 'import tensorflow as tf\n'), ((1889, 1910), 'tensorflow.exp', 'tf.exp', (['(-0.5 * x ** 2)'], {}), '(-0.5 * x ** 2)\n', (1895, 1910), True, 'import tensorflow as tf\n'), ((532, 552), 'tensorflow.ones', 'tf.ones', (['[input_dim]'], {}), '([input_dim])\n', (539, 552), True, 'import tensorflow as tf\n'), ((599, 619), 'tensorflow.ones', 'tf.ones', (['[input_dim]'], {}), '([input_dim])\n', (606, 619), True, 'import tensorflow as tf\n'), ((1539, 1557), 'tensorflow.exp', 'tf.exp', (['self.scale'], {}), '(self.scale)\n', (1545, 1557), True, 'import tensorflow as tf\n'), ((1585, 1603), 'tensorflow.exp', 'tf.exp', (['self.scale'], {}), '(self.scale)\n', (1591, 1603), True, 'import tensorflow as tf\n'), ((2089, 2107), 'tensorflow.exp', 'tf.exp', (['self.scale'], {}), '(self.scale)\n', (2095, 2107), True, 'import tensorflow as tf\n'), ((1017, 1034), 'tensorflow.shape', 'tf.shape', (['points1'], {}), '(points1)\n', (1025, 1034), True, 'import tensorflow as tf\n'), ((1202, 1219), 'tensorflow.shape', 'tf.shape', (['points1'], {}), '(points1)\n', (1210, 1219), True, 'import tensorflow as tf\n'), ((1224, 1241), 'tensorflow.shape', 'tf.shape', (['points2'], {}), '(points2)\n', (1232, 1241), True, 'import tensorflow as tf\n'), ((1086, 1103), 'tensorflow.shape', 'tf.shape', (['points1'], {}), '(points1)\n', (1094, 1103), True, 'import tensorflow as tf\n'), ((1108, 1125), 'tensorflow.shape', 'tf.shape', (['points1'], {}), '(points1)\n', (1116, 1125), True, 'import tensorflow as tf\n')]