import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import re
from typing import Dict
import textstat

# =======================
# إعداد النموذج
# =======================
MODEL_PATH = "GhadaAlothman/arabert_readability_3class"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
model.eval()

DIACRITICS = re.compile(r"[\u0610-\u061A\u064B-\u065F\u0670\u06D6-\u06ED]")
AR_LETTERS = r"[\u0600-\u06FF]"
SENT_SEP = re.compile(r"[\.!\?؟؛…]+")
WORD_RE = re.compile(fr"{AR_LETTERS}+")

# =======================
# دوال المساعدة
# =======================
def strip_diacritics(s: str):
    return DIACRITICS.sub("", s)

def normalize_arabic(s: str):
    return re.sub("[\u0622\u0623\u0625]", "ا", strip_diacritics(s)).replace("ى", "ي").replace("ة", "ه")

def split_sentences(text: str):
    return [p.strip() for p in SENT_SEP.split(text) if p.strip()]

def tokenize_words(text: str):
    return WORD_RE.findall(text)

def difficult_word(w: str, min_len: int = 6):
    return len(w) >= min_len

def compute_metrics(ar_text: str) -> Dict[str, float]:
    text_norm = normalize_arabic(ar_text)
    sents = split_sentences(text_norm)
    words = tokenize_words(text_norm)
    n_sents, n_words = max(len(sents), 1), max(len(words), 1)
    diff_count = sum(1 for w in words if difficult_word(w))
    try:
        osman = float(textstat.osman(ar_text))
    except Exception:
        osman = 0.0
    n_chars = sum(len(w) for w in words)
    ari_ar_score = round((4.71 * (n_chars / n_words)) + (0.5 * (n_words / n_sents)) - 21.43, 3)
    return {
        "Word count": n_words,
        "Sentence count": n_sents,
        "Character count": n_chars,
        "OSMAN_Score": round(osman, 3),
        "ARI_ArScore": ari_ar_score,
        "Difficult_Words_Count": diff_count,
        "Average_Sentence_Length_in_Words": round(n_words / n_sents, 3),
    }

# =======================
# دالة التنبؤ
# =======================
def analyze_text(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
        label_id = torch.argmax(probs, dim=1).item()
    label_map = {0: "سهل", 1: "متوسط", 2: "صعب"}
    label = label_map[label_id]
    stats = compute_metrics(text)
    return {"Predicted_Label": label, **stats}

# =======================
# واجهة Gradio
# =======================
demo = gr.Interface(
    fn=analyze_text,
    inputs=gr.Textbox(label="أدخل النص العربي هنا", lines=6),
    outputs=gr.JSON(label="نتائج التحليل"),
    api_name="predict",   # ← تأكدي إن هذا السطر موجود!
    title="Arabic Readability Analyzer",
    description="أداة ذكية لتقييم مقروئية النصوص العربية باستخدام نموذج AraBERT."
)

demo.launch()