import gradio as gr
import torch
import re
import unicodedata
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch.nn.functional as F

# ⚠️ REPLACE WITH YOUR ACTUAL REPO NAME from Step 1
REPO_NAME = "ArabicNewsAnalyzer/arabic-dialect-identifier-msa-lev" 

# --- Preprocessing Logic (Exactly as trained) ---
PREPROCESS_CONFIG = {
    'remove_urls'           : True,
    'remove_mentions'       : True,
    'remove_hashtags'       : True,
    'remove_emojis'         : True,
    'remove_punctuation'    : True,
    'remove_latin'          : True,
    'remove_digits'         : True,
    'normalize_alef'        : True,
    'normalize_yeh'         : True,
    'normalize_teh_marbuta' : True,
    'remove_tatweel'        : True,
    'remove_diacritics'     : True,
    'remove_extra_spaces'   : True,
    'min_token_length'      : 2,
}

ARABIC_DIACRITICS_PATTERN = re.compile(r'[\u0610-\u061A\u064B-\u065F\u0670\u06D6-\u06DC\u06DF-\u06E4\u06E7\u06E8\u06EA-\u06ED]')
TATWEEL_PATTERN           = re.compile(r'\u0640')
ALEF_PATTERN              = re.compile(r'[\u0622\u0623\u0625\u0671]')
YEH_PATTERN               = re.compile(r'\u0649')
TEH_MARBUTA_PATTERN       = re.compile(r'\u0629')

def is_emoji(char):
    try:
        return unicodedata.category(char) in ('So', 'Sm') or ord(char) in range(0x1F600, 0x1FFFF)
    except TypeError:
        return False

def preprocess(text: str, cfg: dict) -> str:
    if not isinstance(text, str):
        return ''
    if cfg['remove_urls']:
        text = re.sub(r'https?://\S+|www\.\S+', ' ', text)
    if cfg['remove_mentions']:
        text = re.sub(r'@\w+', ' ', text)
    if cfg['remove_hashtags']:
        text = re.sub(r'#\w+', ' ', text)
    if cfg['remove_emojis']:
        text = ''.join(c for c in text if not is_emoji(c))
    if cfg['remove_punctuation']:
        text = re.sub(r'[!\"#$%&\'()*+,\-./:;<=>?@\[\\\]^_`{|}~،؛؟«»]', ' ', text)
    if cfg['remove_latin']:
        text = re.sub(r'[a-zA-Z]', ' ', text)
    if cfg['remove_digits']:
        text = re.sub(r'[0-9٠-٩]', ' ', text)
    if cfg['normalize_alef']:
        text = ALEF_PATTERN.sub('\u0627', text)
    if cfg['normalize_yeh']:
        text = YEH_PATTERN.sub('\u064A', text)
    if cfg['normalize_teh_marbuta']:
        text = TEH_MARBUTA_PATTERN.sub('\u0647', text)
    if cfg['remove_tatweel']:
        text = TATWEEL_PATTERN.sub('', text)
    if cfg['remove_diacritics']:
        text = ARABIC_DIACRITICS_PATTERN.sub('', text)
    if cfg['remove_extra_spaces']:
        text = re.sub(r'\s+', ' ', text).strip()
    return text

# --- Load Model & Tokenizer ---
tokenizer = AutoTokenizer.from_pretrained(REPO_NAME)
model = AutoModelForSequenceClassification.from_pretrained(REPO_NAME)
model.eval()

# Check available device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

def predict_dialect(text):
    if not text.strip():
        return {"Please enter some text": 1.0}
        
    # Preprocess matching the training exactly
    clean_text = preprocess(text, PREPROCESS_CONFIG)
    
    # Optional guard if preprocessing erased the whole text
    if len(clean_text.split()) < PREPROCESS_CONFIG['min_token_length']:
         return {"Error: Input too short after cleaning (URLs/Mentions removed).": 1.0}

    # Tokenize
    inputs = tokenizer(
        clean_text, 
        max_length=128, 
        padding='max_length', 
        truncation=True, 
        return_tensors="pt"
    )
    
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    # Infer
    with torch.no_grad():
        outputs = model(**inputs)
        # Convert logits to probabilities
        probabilities = F.softmax(outputs.logits, dim=-1).squeeze(0)
    
    # Build dictionary output {Label Name: Probability (0-1.0)} for Gradio component
    confidences = {model.config.id2label[i]: float(probabilities[i]) for i in range(model.config.num_labels)}
    
    # Return as dict. Gradio formats this nicely as bars!
    return confidences

# --- Gradio Interface ---
iface = gr.Interface(
    fn=predict_dialect,
    inputs=gr.Textbox(
        lines=3, 
        placeholder="Enter Arabic text here...", 
        label="Arabic Text"
    ),
    outputs=gr.Label(num_top_classes=len(model.config.id2label), label="Predicted Dialect Confidence"),
    title="Arabic Dialect Identification (Two Regions)",
    description="Fine-tuned UBC-NLP/MARBERTv2 to detect regional Arabic dialects. (Preprocessing removes punctuation, Latin chars, URLs, Emojis, and Normalizes text exactly as in training).",
    examples=[
        ["شلونك شخبارك يا خوي"], 
        ["وش حالك وش علومك؟"], 
        ["ازيك يا عم عامل ايه؟"],
    ]
)

iface.launch()