import streamlit as st
import pandas as pd
import numpy as np
import nltk
from nltk.tokenize import sent_tokenize, word_tokenize
from transformers import AutoModelForCausalLM, AutoTokenizer


# Настройка конфигурации страницы Streamlit
st.set_page_config(
    page_title="Generate reviews",
    initial_sidebar_state="expanded"
)

# Заголовок приложения
st.title("Генератор отзывов на основе ИИ")
st.write("Создайте уникальные текстовые отзывы о различных местах на основе категорий, рейтинга и ключевых слов.")


def download_nltk_data():
    nltk.download('punkt')
    nltk.download('punkt_tab')


# Загрузка модели и токенизатора
# @st.cache_data()
@st.cache_resource
def get_model():
    # Загрузка модели
    model = AutoModelForCausalLM.from_pretrained('model')
    # Загрузка токенизатора
    tokenizer = AutoTokenizer.from_pretrained('model')
    return model, tokenizer


# Генерация отзыва
def gen_review(input_text):
    model, tokenizer = get_model()
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    output = model.generate(
        input_ids,
        max_length=200,
        num_return_sequences=1,
        no_repeat_ngram_size=2,
        do_sample=True,
        top_p=0.95,
        top_k=60,
        temperature=0.9,
        eos_token_id=tokenizer.eos_token_id,
    )
    return tokenizer.decode(output[0], skip_special_tokens=True)


def capitalize_and_punctuate(text):    
    # Разделяем текст на предложения
    text = text.split(":")[-1].strip()
    sentences = sent_tokenize(text)

    # Проверка последнего предложения
    last_sentence = sentences[-1]
    if not last_sentence.endswith('.'):
        sentences.pop()

    # Обрабатываем оставшиеся предложения
    corrected_sentences = []
    for sentence in sentences:
        words = word_tokenize(sentence)

        # Делаем первую букву первого слова заглавной
        if len(words) > 0:
            words[0] = words[0].capitalize()

        # Собираем обратно предложение
        corrected_sentence = ' '.join(words)
        corrected_sentences.append(corrected_sentence)

    # Объединяем все предложения в единый текст
    final_text = ' '.join(corrected_sentences)
    final_text = final_text.replace(' .', '.')
    
    return final_text


# Главная функция
def main():
    
    if 'btn_predict' not in st.session_state:
        st.session_state['btn_predict'] = False

    category = st.text_input("Категория:", value="Кондитерская")
    rating = st.slider("Рейтинг", 1, 5, 1)
    key_words = st.text_input("Ключевые слова", value="десерт, торт, цена")

    # Ввод новых параметров
    input_text = f"Категория: {category}; Рейтинг: {rating}; Ключевые слова: {key_words} -> Отзыв:"

    if st.button('Generate'):
        with st.spinner('Генерация отзыва...'):
            generated_text = gen_review(input_text)
            generated_text = capitalize_and_punctuate(generated_text)        
        st.success("Готово!")
        st.text(generated_text)


if __name__ == "__main__":
    download_nltk_data()
    main()