from __future__ import annotations

import json
from pathlib import Path

import gradio as gr
import pandas as pd
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForSequenceClassification, AutoTokenizer

MODEL_ID = "proxy3d/multi-motions-28"
MAX_LENGTH = 64

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
model.to(DEVICE)
model.eval()

# Prefer the thresholds published with the model. Fall back to 0.5 if unavailable.
try:
    thresholds_path = hf_hub_download(
        repo_id=MODEL_ID,
        filename="thresholds.json",
        repo_type="model",
    )
    with open(thresholds_path, "r", encoding="utf-8") as f:
        THRESHOLDS = json.load(f)
except Exception:
    THRESHOLDS = {}

ID2LABEL = {
    int(k): v for k, v in model.config.id2label.items()
}
LABELS = [ID2LABEL[i] for i in range(len(ID2LABEL))]


def _predict_scores(text: str) -> list[tuple[str, float]]:
    text = (text or "").strip()
    if not text:
        return []

    encoded = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=MAX_LENGTH,
    )
    encoded = {k: v.to(DEVICE) for k, v in encoded.items()}

    with torch.inference_mode():
        logits = model(**encoded).logits[0]
        probs = torch.sigmoid(logits).detach().cpu().float().tolist()

    return [(ID2LABEL[i], float(probs[i])) for i in range(len(probs))]


def predict(text: str):
    scores = _predict_scores(text)

    if not scores:
        empty = pd.DataFrame(columns=["emotion", "score", "threshold", "active"])
        return (
            "Введите текст / Enter text.",
            empty,
            empty,
        )

    rows = []
    active = []
    for label, score in scores:
        threshold = float(THRESHOLDS.get(label, 0.5))
        is_active = score >= threshold
        rows.append(
            {
                "emotion": label,
                "score": score,
                "threshold": threshold,
                "active": "✓" if is_active else "",
            }
        )
        if is_active:
            active.append((label, score))

    rows.sort(key=lambda x: x["score"], reverse=True)
    top_df = pd.DataFrame(rows[:10])
    all_df = pd.DataFrame(rows)

    if active:
        active.sort(key=lambda x: x[1], reverse=True)
        active_text = ", ".join(f"**{label}** ({score:.3f})" for label, score in active)
        summary = f"### Активные эмоции / Active labels\n{active_text}"
    else:
        top_label, top_score = rows[0]["emotion"], rows[0]["score"]
        summary = (
            "### Активные эмоции / Active labels\n"
            "Ни один класс не превысил свой tuned threshold. "
            f"Максимальный score: **{top_label}** ({top_score:.3f})."
        )

    return summary, top_df, all_df


EXAMPLES = [
    ["Я боюсь опоздать на рейс."],
    ["Мне наконец ответили — какое облегчение."],
    ["Спасибо, это действительно очень помогло."],
    ["Я совсем не понимаю, почему это произошло."],
    ["I'm worried I'll miss my flight."],
    ["Thank you so much, this really helped me."],
    ["I can't believe this actually happened!"],
    ["I'm disappointed, but I understand the decision."],
]


with gr.Blocks(title="Multi-Motions 28 — EN/RU Emotion Classifier") as demo:
    gr.Markdown(
        """
# Multi-Motions 28

**English + Russian · 28 GoEmotions classes · multi-label emotion classification**

Введите английский или русский текст. Модель возвращает полный вектор confidence scores,
а также labels, прошедшие индивидуальные tuned thresholds.

**Model:** [proxy3d/multi-motions-28](https://huggingface.co/proxy3d/multi-motions-28)  
**Author:** Ilya Zelenskiy (proxy3d) / Илья Зеленский (proxy3d)  
**Telegram:** [t.me/greenruff](https://t.me/greenruff)  

**Communication Styles LLM:**  
[Article](https://iproxy3d.github.io/communication-styles-llm/) ·
[GitHub](https://github.com/iproxy3d/communication-styles-llm)
        """
    )

    text = gr.Textbox(
        label="Text / Текст",
        placeholder="Введите сообщение на русском или английском...",
        lines=4,
    )
    run = gr.Button("Analyze / Анализировать", variant="primary")

    summary = gr.Markdown()

    gr.Markdown("### Top-10 scores")
    top_table = gr.Dataframe(
        headers=["emotion", "score", "threshold", "active"],
        datatype=["str", "number", "number", "str"],
        interactive=False,
    )

    with gr.Accordion("All 28 scores / Все 28 классов", open=False):
        all_table = gr.Dataframe(
            headers=["emotion", "score", "threshold", "active"],
            datatype=["str", "number", "number", "str"],
            interactive=False,
        )

    gr.Examples(
        examples=EXAMPLES,
        inputs=text,
        label="Examples / Примеры",
    )

    gr.Markdown(
        """
---
### Notes

- Scores are sigmoid outputs for all 28 GoEmotions labels.
- This is a **multi-label** classifier: more than one emotion may be active.
- The `active` column uses the per-label thresholds published with the model.
- Native-Russian transfer was additionally evaluated on CEDR and SemEval-2025 RU.

For implementation details, benchmarks and usage examples, see the
[model page](https://huggingface.co/proxy3d/multi-motions-28).
        """
    )

    run.click(
        fn=predict,
        inputs=text,
        outputs=[summary, top_table, all_table],
    )
    text.submit(
        fn=predict,
        inputs=text,
        outputs=[summary, top_table, all_table],
    )

if __name__ == "__main__":
    demo.launch()