from openpy_fxts.mdits.mdls_dl.base_tf import tkl, K


class Time2Vec(tkl.Layer):

    def __init__(self, output_dim=None, **kwargs):
        self.output_dim = output_dim
        super(Time2Vec, self).__init__(**kwargs)

    def build(self, input_shape):
        self.W = self.add_weight(
            name='W',
            shape=(input_shape[-1], self.output_dim),
            initializer='uniform',
            trainable=True)

        self.P = self.add_weight(
            name='P',
            shape=(input_shape[1], self.output_dim),
            initializer='uniform',
            trainable=True
        )
        self.w = self.add_weight(
            name='w',
            shape=(input_shape[1], 1),
            initializer='uniform',
            trainable=True
        )

        self.p = self.add_weight(
            name='p',
            shape=(input_shape[1], 1),
            initializer='uniform',
            trainable=True
        )
        super(Time2Vec, self).build(input_shape)

    def call(self, x):
        original = self.w * x + self.p
        sin_trans = K.sin(K.dot(x, self.W) + self.P)

        return K.concatenate([sin_trans, original], -1)
