import streamlit as st
import pandas as pd
import torch
import torch.nn as nn

import nltk
nltk.download('stopwords')
nltk.download('punkt_tab')
nltk.download('wordnet')

from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
stop_words = set(stopwords.words('russian'))

from nltk.tokenize import word_tokenize
import re
import string
from collections import Counter
import numpy as np
from typing import List
import time

from models.rest.model_lstm import LSTM_Word2Vec_Attention


# # Добавление цвета для значений F1-Macro
def color_high(val):
    color = 'lightgreen' if val > 0.80 else ''
    return f'background-color: {color}'

# Данные для первой таблицы
data1 = {
    "Модель": ["Линейная регрессия", "Дерево решений"],
    "F1-Macro": [0.2235, 0.1688]
}

# Данные для второй таблицы
data2 = {
    "Модель": ["Линейная регрессия", "Наивный Байес", "Деревья", "XGBoost", "CatBoost"],
    "F1-Macro": [0.7821, 0.7313, 0.8170, 0.7785, 0.7693]
}

# Создание DataFrame
df1 = pd.DataFrame(data1)
df2 = pd.DataFrame(data2)

st.title(":blue[_Классификация отзывов на рестораны_]")
st.image('images/53896.jpg')

# Отображение заголовков и таблиц
st.subheader("Сравнение моделей по метрике F1-Macro")

st.subheader("Начало. До понимания, что данные хуже чем можно было представить")
st.table(df1)

st.subheader("После всего... Работа классических ML-моделей")
df2 = df2.style.applymap(color_high, subset=['F1-Macro'])
st.table(df2)

# Загрузка модели и весов
@st.cache_data()
def load_model():
    hidden_size = 32
    vocab_size = 310  
    embedding_dim = 50
    model = LSTM_Word2Vec_Attention(hidden_size, vocab_size, embedding_dim)
    model.load_state_dict(torch.load('models/rest/model_weights_3000cl.pt', map_location=torch.device('cpu')))
    model.eval()
    return model

model = load_model()

# Предобработка текста
def preprocess_text(text: str) -> List[str]:
    text = text.lower()
    text = re.sub(r'<.*?>', '', text)
    text = re.sub(r"[…–”—]", "", text)
    text = re.sub(r'\d+', '', text)
    text = re.sub(r'\b[a-zA-Z]+\b', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    text = ''.join([c for c in text if c not in string.punctuation])
    tokens = word_tokenize(text, language='russian')
    tokens = [word for word in tokens if word not in stopwords.words('russian')]
    return tokens

# Словарь
# Предположим, что у вас есть словарь, который вы использовали для обучения модели
vocab_to_int = np.load('models/rest/vocab_to_int.npy', allow_pickle=True).item()

# Функция для преобразования текста в индексы
def text_to_indices(text: str) -> torch.Tensor:
    tokens = preprocess_text(text)
    indices = [vocab_to_int.get(token, 0) for token in tokens]  # 0 для неизвестных слов
    return torch.tensor([indices[:200] + [0]*(200 - len(indices))])  # Пэддинг до 200


st.title("Классификация отзывов о ресторанах")

st.image('images/grafic.jpg')

review = st.text_area("Введите ваш отзыв о ресторане:")

if st.button("Предсказать"):
    if review.strip() == "":
        st.write("Пожалуйста, введите отзыв.")
    else:
        start_time = time.time()
        input_data = text_to_indices(review)
        with torch.no_grad():
            output, attention_weights = model(input_data)
        prediction = torch.argmax(output, dim=1).item()
        end_time = time.time()
        elapsed_time = end_time - start_time
        st.write(f"Предсказанный класс: {prediction} (Положительный: 2, Отрицательный: 1, Нейтральный: 0)")
        st.write(f"Время предсказания: {elapsed_time:.3f} секунд")