import numpy as np
import gradio as gr
import tensorflow as tf

from sentence_transformers import SentenceTransformer, util

emb_model = SentenceTransformer('distiluse-base-multilingual-cased-v2')
EMB_DIM = 512

feedback_class_model = tf.keras.models.load_model('models/feedback_class.keras')
area_class_model = tf.keras.models.load_model('models/area_class.keras')

def process(text_1, text_2):
    '''
    process(str, str) -> Union[List[torch.Tensor], numpy.ndarray, torch.Tensor]
    Function encodes texts from to embeddings.
    '''
    tokenized_text = emb_model.encode([text_1])
    tokenized_org = emb_model.encode([text_2])
    return tokenized_text, tokenized_org


def feedback_classification(input_text, input_org):
    train_text, train_org = process(input_text, input_org)

    X_text = tf.reshape(train_text, [-1, 1, EMB_DIM])
    X_org = tf.reshape(train_org, [-1, 1, EMB_DIM])

    feedback_scores = feedback_class_model.predict([X_text, X_org])
    #area_scores = area_class_model.predict([X_text, X_org])

    print(feedback_scores[0][0])

feedback_classification('Хотіли би подякувати усім співробітникам Gdynia Community Center!!!',
              'Community Center Gdynia | Danish Refugee Council (DRC) | Stowarzyszenie OVUM')