import gradio as gr
import torch
import os
from transformers import AutoModelForSequenceClassification, AutoTokenizer


MODEL_ID = "Batuka0901/mongolian_news_classifier"
MAX_LENGTH = 128
HF_TOKEN = os.environ.get("HF_TOKEN")

CATEGORY_DESC = {
    "байгал орчин": "Байгаль орчин, экологи, цаг уур, ногоон хөгжил",
    "боловсрол":    "Сургуули, их дээд сургууль, багш, сурагч, хичээл",
    "спорт":        "Тэмцээн, уралдаан, шигшээ баг, медаль, аварга",
    "технологи":    "Шинэ технологи, программ хангамж, интернет, инновац",
    "улс төр":      "Засгийн газар, сонгууль, хууль тогтоомж, бодлого",
    "урлаг соёл":   "Урлаг, дуу хөгжим, кино, уран зохиол, соёлын арга хэмжээ",
    "хууль":        "Шүүх, хэрэг, дүрэм журам, хуулийн байгууллага",
    "эдийн засаг":  "Санхүү, худалдаа, үйлдвэрлэл, инфляц, ДНБ",
    "эрүүл мэнд":   "Эмнэлэг, эм, өвчин, эмч, эрүүл мэндийн бодлого",
}


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, token=HF_TOKEN).to(device).eval()


id2label = {int(k): v for k, v in model.config.id2label.items()}

def classify(text: str):
    text = text.strip()
    if not text:
        return "", {}, ""

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

    with torch.no_grad():
        probs = torch.nn.functional.softmax(model(**inputs).logits, dim=-1)[0]

    pred_id    = probs.argmax().item()
    pred_label = id2label[pred_id]
    confidence = probs[pred_id].item()

    scores = {id2label[i]: float(probs[i]) for i in range(len(id2label))}
    desc = CATEGORY_DESC.get(pred_label, "")
    detail = f"Ангилал: {pred_label}  |  Итгэлцэл: {confidence:.1%}\n{desc}"

    return pred_label, scores, detail

EXAMPLES = [
    ["Өнөөдөр хөлбөмбөгийн тэмцээнд Монголын шигшээ баг хожлоо."],
    ["Эдийн засгийн сайд өнөөдөр мэдээлэл хийлээ."],
    ["Шинэ технологийн компани байгуулагдлаа."],
    ["Эмнэлэгт шинэ тоног төхөөрөмж суурилуулжээ."],
    ["Багш нарын цалин нэмэгдэнэ."],
    ["Нийслэлийн агаарын бохирдол буурлаа."],
]

CSS = """
#title { text-align: center; margin-bottom: 4px; }
#result_label { font-size: 1.3em; font-weight: bold; }
footer { display: none !important; }
"""

with gr.Blocks(css=CSS, title="Монгол мэдээний ангилагч") as demo:

    gr.Markdown("## Монгол мэдээний гарчиг ангилагч", elem_id="title")

    with gr.Row():
        with gr.Column(scale=3):
            txt_input = gr.Textbox(
                label="Мэдээний гарчиг",
                placeholder="Мэдээний гарчгийг энд бичнэ үү...",
                lines=3,
            )
            with gr.Row():
                btn_clear  = gr.Button("Арилгах")
                btn_submit = gr.Button("Ангилах", variant="primary")

            gr.Examples(examples=EXAMPLES, inputs=txt_input)

        with gr.Column(scale=2):
            out_label  = gr.Textbox(label="Ангилал", elem_id="result_label")
            out_scores = gr.Label(label="Ангилал бүрийн магадлал", num_top_classes=9)
            out_detail = gr.Textbox(label="Дэлгэрэнгүй", lines=2)

    btn_submit.click(fn=classify, inputs=txt_input, outputs=[out_label, out_scores, out_detail])
    btn_clear.click(fn=lambda: ("", {}, ""), outputs=[out_label, out_scores, out_detail])
    txt_input.submit(fn=classify, inputs=txt_input, outputs=[out_label, out_scores, out_detail])

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