import os

import tensorflow as tf
import tensorflow_hub as hub

from textattack.constraints.semantics.sentence_encoders import SentenceEncoder


class UniversalSentenceEncoder(SentenceEncoder):
    """ 
    Constraint using similarity between sentence encodings of x and x_adv where 
    the text embeddings are created using the Universal Sentence Encoder.
    """

    def __init__(self, threshold=0.8, large=False, metric="angular", **kwargs):
        super().__init__(threshold=threshold, metric=metric, **kwargs)
        if large:
            tfhub_url = "https://tfhub.dev/google/universal-sentence-encoder-large/5"
        else:
            tfhub_url = "https://tfhub.dev/google/universal-sentence-encoder/4"

        self.model = hub.load(tfhub_url)

    def encode(self, sentences):
        return self.model(sentences).numpy()
