app.py

import os
import io

import numpy as np
import pandas as pd
import streamlit as st
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
from huggingface_hub import login

# ==========================
# НАСТРОЙКИ ПРИЛОЖЕНИЯ
# ==========================

st.set_page_config(
    page_title="Поиск постдока",
    layout="wide",
)

HF_TOKEN = os.getenv("HF_TOKEN")
HF_MERGED_REPO = os.getenv("HF_MERGED_REPO")
HF_EMB_REPO = os.getenv("HF_EMB_REPO")
MODEL_NAME = os.getenv("MODEL_NAME")

# Отображаемые заголовки (русские имена колонок)
COLUMN_LABELS_RU = {
    "№": "№",
    "score": "Сходство",
    "fio": "ФИО",
    "title": "Название диссертации",
    "author_org_short": "Организация",
    "dissertation_type": "Тип",
    "protection_year": "Год",
    "registration_number": "Регистрационный номер",
    "vak_link": "Ссылка ВАК",   # только в Excel
}

# Порядок колонок в выдаче (для интерфейса)
DISPLAY_COLUMNS = [
    "№",
    "score",
    "fio",
    "title",
    "author_org_short",
    "dissertation_type",
    "protection_year",
    "registration_number",
]

if HF_TOKEN is None:
    st.error(
        "Не найден секрет `HF_TOKEN`. "
        "Задайте его в Settings → Variables and secrets вашего Space."
    )
    st.stop()

# Логин в Hugging Face Hub (для приватных датасетов)
try:
    login(token=HF_TOKEN)
except Exception:
    pass


# ==========================
# ЗАГРУЗКА ДАННЫХ И МОДЕЛИ
# ==========================

@st.cache_data(show_spinner="Загрузка данных...")
def load_data():
    # 1) МЕТА-ДАННЫЕ
    ds_meta = load_dataset(
        HF_MERGED_REPO,
        split="train",
    )
    df_meta = ds_meta.to_pandas()
    df_meta["registration_number"] = df_meta["registration_number"].astype(str)
    df_meta = df_meta.set_index("registration_number", drop=False)

    # 2) ЭМБЕДДИНГИ
    ds_emb = load_dataset(
        HF_EMB_REPO,
        split="train",
    )
    df_emb = ds_emb.to_pandas()
    df_emb["registration_number"] = df_emb["registration_number"].astype(str)
    reg_nums = df_emb["registration_number"].tolist()

    emb_matrix = np.vstack(df_emb["embedding"].to_list()).astype("float32")
    norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
    emb_matrix = emb_matrix / np.maximum(norms, 1e-8)

    return df_meta, reg_nums, emb_matrix


@st.cache_resource(show_spinner="Загрузка модели...")
def load_model():
    model = SentenceTransformer(MODEL_NAME)
    return model


try:
    df_all, reg_nums, emb_matrix = load_data()
    model = load_model()
except Exception as e:
    st.error(f"Ошибка при загрузке данных или модели: {e}")
    st.stop()


# ==========================
# ЛОГИКА ПОИСКА
# ==========================

def search_core(query: str, top_k: int = 10):
    query = query.strip()
    if not query:
        return []

    query_text = "query: " + query

    q_emb = model.encode(
        query_text,
        normalize_embeddings=True,
    )

    scores = emb_matrix @ q_emb
    top_k = min(int(top_k), len(scores))
    top_idx = np.argsort(-scores)[:top_k]

    results = []
    for rank, idx in enumerate(top_idx, start=1):
        results.append(
            {
                "rank": rank,
                "registration_number": reg_nums[idx],
                "score": float(scores[idx]),
            }
        )

    return results


def extract_year(value):
    """Аккуратно вытаскиваем год защиты из поля protection_date."""
    if value is None:
        return None
    try:
        if not isinstance(value, str):
            if pd.isna(value):
                return None
            dt = pd.to_datetime(value)
            return str(dt.year)
    except Exception:
        pass

    s = str(value).strip()
    if len(s) >= 4 and s[:4].isdigit():
        return s[:4]
    return None


def build_result_df(results):
    rows = []

    for r in results:
        reg = r["registration_number"]
        score = r["score"]

        if reg in df_all.index:
            meta = df_all.loc[reg]
        else:
            meta = pd.Series({}, index=df_all.columns)

        protection_date = meta.get("protection_date", None)
        protection_year = extract_year(protection_date)

        row = {
            "№": r["rank"],
            "score": round(score, 4),
            "fio": meta.get("fio", None),
            "title": meta.get("title", None),
            "author_org_short": meta.get("author_org_short", None),
            "dissertation_type": meta.get("dissertation_type", None),
            "protection_year": protection_year,
            "registration_number": meta.get("registration_number", reg),
            "vak_link": meta.get("vak_link", ""),
        }

        rows.append(row)

    if not rows:
        return pd.DataFrame(columns=DISPLAY_COLUMNS + ["vak_link"])

    df_res = pd.DataFrame(rows)
    df_res = df_res[DISPLAY_COLUMNS + ["vak_link"]]
    return df_res


def run_search(query: str, top_k: int):
    query = query.strip()
    if not query:
        empty_df = pd.DataFrame(columns=DISPLAY_COLUMNS + ["vak_link"])
        empty_df_ru = empty_df.rename(columns=COLUMN_LABELS_RU)
        return empty_df_ru, None

    results = search_core(query, top_k)
    df_res = build_result_df(results)

    if df_res.empty:
        empty_df_ru = df_res.rename(columns=COLUMN_LABELS_RU)
        return empty_df_ru, None

    # --- Разделяем датафреймы для отображения и для Excel ---
    df_excel = df_res.copy()       # vak_link оставляем
    df_display = df_res.copy()

    # HTML-ссылка в колонке "Тип"
    def make_type_cell(row):
        t = row.get("dissertation_type", "")
        link = row.get("vak_link") or ""
        if isinstance(t, str) and t and link:
            return f'<a href="{link}" target="_blank">{t}</a>'
        return t

    df_display["dissertation_type"] = df_display.apply(make_type_cell, axis=1)
    df_display = df_display.drop(columns=["vak_link"])

    df_display_ru = df_display.rename(columns=COLUMN_LABELS_RU)
    df_excel_ru = df_excel.rename(columns=COLUMN_LABELS_RU)

    # Excel (с отдельной колонкой "Ссылка ВАК")
    output = io.BytesIO()
    with pd.ExcelWriter(output, engine="xlsxwriter") as writer:
        df_excel_ru.to_excel(writer, index=False)
    output.seek(0)

    return df_display_ru, output


# ==========================
# UI НА STREAMLIT
# ==========================

# Заголовок по центру
st.markdown(
    "<h1 style='text-align: center; margin-bottom: 0.5rem;'>Поиск постдока🎓</h1>",
    unsafe_allow_html=True,
)

# Форма: Ctrl+Enter = отправка
with st.form("search_form"):
    top_k = st.slider(
        "Сколько результатов показать",
        min_value=1,
        max_value=100,
        value=20,
        step=1,
    )

    query = st.text_area(
        "Введите запрос",
        height=120,
        placeholder="Например: пластификаторы для самоуплотняющихся бетонов",
        key="query",
    )

    # Кнопка по центру через три колонки
    c1, c2, c3 = st.columns([1, 1, 1])
    with c1:
        st.write("")
    with c2:
        do_search = st.form_submit_button("🔍 Поиск", type="primary",use_container_width=True )
    with c3:
        st.write("")

# Результаты
if do_search:
    with st.spinner("Идёт поиск по базе диссертаций..."):
        df_res_ru, excel_bytes = run_search(query, top_k)

    if df_res_ru.empty:
        st.warning("Ничего не найдено. Попробуйте изменить формулировку запроса.")
    else:
        st.success(f"Найдено записей: {len(df_res_ru)}")

        table_html = df_res_ru.to_html(
            escape=False,
            index=False,
            classes="result-table",
        )

        st.markdown(
            """
            <style>
            table.result-table {
                width: 100%;
                border-collapse: collapse;
            }
            table.result-table th {
                text-align: center !important;
                vertical-align: middle;
            }
            </style>
            """,
            unsafe_allow_html=True,
        )

        st.markdown(table_html, unsafe_allow_html=True)

        if excel_bytes is not None:
            st.download_button(
                label="💾 Скачать результаты в Excel",
                data=excel_bytes,
                file_name="search_results.xlsx",
                mime=(
                    "application/vnd.openxmlformats-officedocument."
                    "spreadsheetml.sheet"
                ),
            )
else:
    st.info("Введите запрос выше и нажмите кнопку «Поиск» "
            "или нажмите Ctrl+Enter в поле ввода.")

# Футер
st.markdown(
    "<p style='font-size: 0.8rem; text-align: right; color: gray;'>"
    "(с) Антон Лощилов, 2025"
    "</p>",
    unsafe_allow_html=True,
)