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

from sentence_transformers import SentenceTransformer
from huggingface_hub import from_pretrained_keras

# load pre-trained model
emb_model = SentenceTransformer('distiluse-base-multilingual-cased-v2')
EMB_DIM = 512

# load model for first classification with defining classes
feedback_class_model = from_pretrained_keras('vitiugin/loop_feedback')
feedback_labels = ['Request', 'Thanks', 'Question', 'Opinion', 'Concern']

# load model for second classification with defining classes
area_class_model = from_pretrained_keras('vitiugin/loop_area')
area_labels = ['cross-cutting', 'education', 'food security', 'governance', 'health', 'protection', 'shelter', 'wash']


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):
    '''
    process(str, str) -> Dict, Dict
    Function fits texts (fedback and organization title) into predefined classes
    '''
    train_text, train_org = process(input_text, input_org) #tokenization

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

    # getting scores from classification model
    feedback_scores = feedback_class_model.predict([X_text, X_org])
    area_scores = area_class_model.predict([X_text, X_org])

    # create dict with classification-based probabilities
    feedback_scores = {feedback_labels[num]: feedback_scores[0][0][num] for num in range(len(feedback_labels))}
    area_scores = {area_labels[num]: area_scores[0][0][num] for num in range(len(area_labels))}

    return feedback_scores, area_scores

demo = gr.Interface(
    fn=feedback_classification, #define type of the app
    inputs=[gr.Textbox(placeholder='Enter a feedback text'), 
            gr.Textbox(placeholder='Enter a title of organization')], # define the interface fields
    outputs=['label', 'label'], # output interface
    examples=[['Thank you, but next time just send us more fresh vegetables!', 'Red Cross']]) # definition of data example

demo.launch()