from pathlib import Path

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


MODEL_DIR = Path(__file__).resolve().parent / "model"
DEMO_LABELS = [
    "cs.CL",
    "math.NT",
    "astro-ph.GA",
    "q-bio.NC",
    "physics.optics",
    "econ.EM",
]


def make_text(title: str, abstract: str) -> str:
    title = title.strip()
    abstract = abstract.strip()
    if abstract:
        return f"Title: {title}\nAbstract: {abstract}"
    return f"Title: {title}"


def top_95(labels, probabilities):
    order = torch.argsort(probabilities, descending=True)
    result = []
    total = 0.0
    for idx in order:
        probability = float(probabilities[idx])
        result.append((labels[int(idx)], probability))
        total += probability
        if total >= 0.95:
            break
    return result


@st.cache_resource(show_spinner=False)
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, use_fast=False)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
    model.eval()
    labels = [model.config.id2label[i] for i in range(model.config.num_labels)]
    return tokenizer, model, labels


def predict(text: str):
    tokenizer, model, labels = load_model()
    encoded = tokenizer(
        text,
        truncation=True,
        max_length=384,
        padding=True,
        return_tensors="pt",
    )
    with torch.inference_mode():
        logits = model(**encoded).logits[0]
    probabilities = torch.softmax(logits, dim=-1)
    return top_95(labels, probabilities)


st.set_page_config(page_title="ArXiv Topic Classifier", page_icon=":books:", layout="centered")

st.title("ArXiv Topic Classifier")
st.caption("Классификатор тематики научной статьи по названию и abstract.")

title = st.text_input("Название статьи", placeholder="Attention Is All You Need")
abstract = st.text_area(
    "Abstract",
    height=180,
    placeholder="Вставьте abstract. Если его нет, модель использует только название.",
)

if st.button("Классифицировать", type="primary", use_container_width=True):
    if not title.strip() and not abstract.strip():
        st.warning("Введите название, abstract или оба поля.")
    else:
        st.subheader("Топ-95% тематик")
        for label, probability in predict(make_text(title, abstract)):
            st.progress(probability, text=f"{label}: {probability:.1%}")