import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
# как же я это все ненавижу

@st.cache_resource
def load_model():
    model_path = "HoiAlice/bert-paper-classifier-arxiv"
    inference_tokenizer = AutoTokenizer.from_pretrained(model_path)
    inference_model = AutoModelForSequenceClassification.from_pretrained(model_path)
    return pipeline("text-classification", model=inference_model, tokenizer=inference_tokenizer, top_k=None)

def top_pct(preds, threshold: float = 0.9):
    """Возвращает топ предсказаний, пока их суммарная вероятность не превысит threshold"""
    if not preds:
        return []
    preds = sorted(preds, key=lambda x: -x["score"])
    cum_score = 0
    for i, item in enumerate(preds):
        cum_score += item["score"]
        if cum_score >= threshold:
            break
    return preds[:(i+1)]

def format_predictions(preds) -> str:
    """Форматирует предсказания для вывода"""
    if not preds:
        return "Нет результатов"
    return "\n".join([f"{i+1}. {item['label']} (score {item['score']:.2f})" for i, item in enumerate(preds)])

# Интерфейс Streamlit
st.set_page_config(page_title="Классификатор научных статей", page_icon="📚")
st.title("📚 Классификатор научных статей по тематикам")
st.write("Введите текст абстракта статьи, и модель определит наиболее подходящие тематики:")

# Поле для ввода текста
abstract = st.text_area(
    "Текст абстракта:", 
    height=200,
    placeholder="Введите текст научного абстракта здесь..."
)

# Слайдер для выбора порога уверенности
threshold = st.slider(
    "Порог уверенности (суммарная вероятность тематик):", 
    min_value=0.5, 
    max_value=1.0, 
    value=0.9, 
    step=0.05
)

if st.button("Определить тематики"):
    if not abstract.strip():
        st.warning("Пожалуйста, введите текст абстракта")
    else:
        with st.spinner("Загружаем модель... (это может занять некоторое время при первом запуске)"):
            classifier = load_model()
            
        if classifier is not None:
            with st.spinner("Анализируем текст..."):
                try:
                    # Получаем предсказания
                    predictions = classifier(abstract)[0]
                    # Фильтруем по порогу
                    top_predictions = top_pct(predictions, threshold)
                    
                    # Выводим результаты
                    st.subheader("Результаты классификации:")
                    st.text(format_predictions(top_predictions))
                    
                    # Визуализация в виде столбчатой диаграммы
                    st.subheader("Визуализация:")
                    chart_data = {p['label']: p['score'] for p in top_predictions}
                    st.bar_chart(chart_data)
                    
                except Exception as e:
                    st.error(f"Произошла ошибка при анализе текста: {str(e)}")

# Добавляем пояснения в сайдбар
with st.sidebar:
    st.markdown("""
    ## О сервисе
    Этот сервис использует модель уже обученную языковую модель для выделения тематики статьи по ее описанию.
    
    ### Как использовать:
    1. Введите текст абстракта в поле ввода
    2. Отрегулируйте порог уверенности (по умолчанию 0.90)
    3. Нажмите кнопку "Определить тематики"
    
    ### Техническая информация:
    - Используемый датасет: [arXiv papers](https://www.kaggle.com/datasets/neelshah18/arxivdataset)
    - Модель: [BiomedNLP-PubMedBERT](https://huggingface.co/microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract)
                
    ### Благодарности:
    - Большое спасибо Диму Сафину, Ekaterina Zabairachnaya и Андрею Панферову за содержательную критику и активную поддержку,
                без вас бы я уже ливнул бы отседова. 
    """
    )