import tensorflow as tf
import numpy as np
import tensorflow_hub as hub
from tensorflow.keras.utils import register_keras_serializable

@register_keras_serializable()
class USE(tf.keras.layers.Layer):
  def __init__(self,
               encoder_layer,
               encoder_url="https://tfhub.dev/google/universal-sentence-encoder/4"
               #encoder_url="https://www.kaggle.com/models/google/universal-sentence-encoder/frameworks/TensorFlow2/variations/universal-sentence-encoder"
               , **kwargs):
    super(USE, self).__init__(**kwargs)
    self.encoder_layer=encoder_layer
    self.encoder_url=encoder_url

  def call(self, inputs):
    return self.encoder_layer(inputs)

  def get_config(self):
    config=super(USE, self).get_config()
    config['encoder_url']=self.encoder_url
    return config

  @classmethod
  def from_config(cls, config):
    encoder_url=config.pop('encoder_url')
    encoder_layer=hub.KerasLayer(encoder_url, input_shape=[], dtype=tf.string, trainable=False, name='USE')
    return cls(encoder_layer, encoder_url=encoder_url, **config)
  
encoder_layer=hub.KerasLayer("https://tfhub.dev/google/universal-sentence-encoder/4",
                            input_shape=[], ## The input is of variable length, hence the ip_length=[]
                             dtype=tf.string,
                             trainable=False,
                             name='USE')


class Model():
    def __init__(self, encoder_layer=encoder_layer):
        super(Model, self).__init__()
        self.encoder_layer=encoder_layer
        self.USE=USE(self.encoder_layer)
        self.Sentiment_model=tf.keras.models.load_model('Best_sentiment_model.keras', custom_objects={'USE':self.USE})
        self.Sentiment_model.trainable=False
        self.Emotion_model=tf.keras.models.load_model('model_emotion_lstm.keras')
        self.Emotion_model.trainable=False

    def predict(self, text_sentiment=None, text_emotion=None):
      sentiment = None
      emotion = None
    
      if text_sentiment is not None:
        sentiment = self.Sentiment_model.predict(text_sentiment)
        
      if text_emotion is not None:
        emotion = self.Emotion_model.predict(text_emotion)
    
    # If both were requested, return both
      if text_sentiment is not None and text_emotion is not None:
        return sentiment, emotion
    # If only sentiment was requested
      elif text_sentiment is not None:
        return sentiment
    # If only emotion was requested
      elif text_emotion is not None:
        return emotion
    # If neither was provided
      return None



if __name__=='__main__':
    print(tf.__version__)
    model=Model()
    res1, res2=model.predict(tf.convert_to_tensor(['I am happy'], dtype=tf.string))
    print(np.argmax(res1, axis=1), np.argmax(res2, axis=1))