"""
Functional Transformers

"""
import numpy as np
import tensorflow_hub as hub
from sklearn.preprocessing import (
    FunctionTransformer,
    SplineTransformer
)

UNISERVAL_SENT_ENCODER_URL = "https://tfhub.dev/google/universal-sentence-encoder-multilingual/3"

"""
Cyclical Feature Encoding
"""

def sin_transformer(period):
    """
    Sine Transformer
    """
    def sin_encoder(x, period=period):
        return np.sin(x / period * 2 * np.pi)
    return FunctionTransformer(sin_encoder)


def cos_transformer(period):
    """
    Consine Transformer
    """
    def cos_encoder(x, period=period):
        return np.cos(x / period * 2 * np.pi)
    return FunctionTransformer(cos_encoder)


def periodic_spline_transformer(period, n_splines=None, degree=3):
    """
    Period Spline Transformer
    """
    if n_splines is None:
        n_splines = period
    n_knots = n_splines + 1  # periodic and include_bias is True
    return SplineTransformer(
        degree=degree,
        n_knots=n_knots,
        knots=np.linspace(0, period, n_knots).reshape(n_knots, 1),
        extrapolation="periodic",
        include_bias=True,
    )


"""
Text Embedding

TODO: At a later point probably consolidate this into a transformers module
"""


def sentence_embedding_transformer(sentences, model=UNISERVAL_SENT_ENCODER_URL):
    """
    Sentence Embedding Transformer Network

    params:
        sentence (List[str]) - returns sentence embedding using
                                the embedding model
        model (tensorflow Model) - tensorflow hub model to be used

    returns:
        embedding of tensors (N, 512) - returns embedding of the tensors
    """
    model = hub.KerasLayer(model)
    embeddings = model(sentences)

    return embeddings