import gradio as gr
import torch

from src.data.preprocessing_utils import DataPreprocessor

MODEL_FILEPATH = 'models/nnet/nnet.pt'
device = 'cuda' if torch.cuda.is_available() else 'cpu'
with open(MODEL_FILEPATH, 'rb') as file:
    clf = torch.load(file, map_location=device)

preprocessor = DataPreprocessor()

strings = {
    'nationality': 'Есть предпочтения по национальности',
    'families': 'Есть предпочтение семьям',
    'sex': 'Есть предпочтения по полу'
}

examples = [
    '''Просьба посредников не беспокоить. Ищем ОДНУ ДЕВУШКУ. 
Сдаётся в аренду на длительный срок светлая и уютная квартира - студия 
общей площадью 33м2, находящаяся на 4м этаже 5и этажного теплого кирпичного дома. Современный ремонт! 
Рядом в пешей доступности парк Красная Пресня (5 мин)/ Красногвардейские Пруды (2 мин)/ 
Москва-Сити (10 мин)! Магазины/кофейни/рестораны! 10 мин на машине до любой точки в центре города! 
В квартире есть вся необходимая для проживания мебель и техника. 
Строго без животных, строго Славян. Просмотр в любое время - ключи на руках.
    ''',
    '''Сдам на длительный срок семейной паре, только с гражданством РФ.
Квартира после косметического ремонта. Без мебели. 
Есть кухонная мебель и мебель в ванной комнате. 
Бытовая техника для проживания присутствует.
Оплата = аренда + счётчики (свет, вода).
    ''',
    '''В современном доме. Собственник без комиссии.
Закрытая территория. Доступ через охрану.
М Прокшино 10 мин пешком.
Без детей и животных.
Возможно без залога.
Счетчики и интернет включены в стоимость
    '''
]


def make_output_string(labels):
    output_list = []
    for label in strings.keys():
        if labels[label]:
            output_list.append(strings[label])
    if output_list:
        output_str = ', '.join(output_list).capitalize()
    else:
        output_str = 'Нет особенностей'
    return output_str


def predict_label(text):
    preprocessed_text = preprocessor.preprocess_texts([text])
    print(preprocessed_text)
    if preprocessed_text == [[]]:
        return 'Введите текст объявления!'
    labels = clf.predict_labels(preprocessed_text)
    output_str = make_output_string(labels)
    return output_str


demo = gr.Interface(fn=predict_label, inputs=[gr.Text(label="Текст объявления", lines=5)],
                    outputs=[gr.Textbox(label="Особенности объявления")],
                    examples=examples)

demo.launch()