import re
import string
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
from peft import LoraConfig, get_peft_model
from huggingface_hub import hf_hub_download

# ==========================================
# 0. Setup & Class Mapping
# ==========================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
LABELS = ["Negative", "Neutral", "Positive"]

# ==========================================
# 1. Model 1 Setup (CAMeLBERT + LoRA)
# ==========================================
M1_NAME = "CAMeL-Lab/bert-base-arabic-camelbert-mix-sentiment"
M1_REPO = "mahmoudmohammad/MARBERTv2-Sentiment_Classification"
M1_FILE = "best_stl_model.pth"

class SentimentSTL(nn.Module):
    def __init__(self, num_classes=3):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(M1_NAME)
        peft_config = LoraConfig(
            task_type="FEATURE_EXTRACTION",
            r=32,
            lora_alpha=64,
            lora_dropout=0.1,
            target_modules=["query", "value"]
        )
        self.peft_model = get_peft_model(self.encoder, peft_config)
        hidden_size = self.peft_model.config.hidden_size

        self.classifier = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(hidden_size, hidden_size // 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_size // 2, num_classes),
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.peft_model(input_ids=input_ids, attention_mask=attention_mask)
        hidden = outputs.last_hidden_state
        mask = attention_mask.unsqueeze(-1).float()
        mean_rep = (hidden * mask).sum(1) / mask.sum(1)
        cls_rep = hidden[:, 0, :]
        cls_rep = (cls_rep + mean_rep) / 2.0
        return self.classifier(cls_rep)

def preprocess_m1(text):
    if not isinstance(text, str):
        return str(text)
    text = re.sub(r'http\S+|www\.\S+', '', text)
    text = re.sub(r'[@]\S+', '', text)
    text = re.sub(r'\S+@\S+', ' ', text)
    text = re.sub(r'\d+|[٠١٢٣٤٥٦٧٨٩]+', '', text)
    text = text.replace('#', ' ').replace('_', ' ')
    text = re.sub("[إأآا]", "ا", text)
    text = re.sub("ى", "ي", text)
    text = re.sub("ؤ", "و", text)
    text = re.sub("ئ", "ي", text)
    text = re.sub("ة", "ه", text)
    arabic_punc = '`÷×؛«»<>()*&^%][ـ،/:".،,\'{}~¦+|"…""–ـ'
    eng_punc = string.punctuation.replace('!', '').replace('?', '')
    text = text.translate(str.maketrans('', '', arabic_punc + eng_punc))
    return re.sub(r'\s+', ' ', text).strip()

print("Loading Model 1 (CAMeLBERT + LoRA)...")
tokenizer_m1 = AutoTokenizer.from_pretrained(M1_NAME)
model_m1 = SentimentSTL(num_classes=3).to(device)
m1_weights_path = hf_hub_download(repo_id=M1_REPO, filename=M1_FILE)
model_m1.load_state_dict(torch.load(m1_weights_path, map_location=device))
model_m1.eval()

# ==========================================
# 2. Model 2 Setup (MARBERTv2 Sequence Classifier)
# ==========================================
M2_PATH = "ArabicNewsAnalyzer/MARBERTv2-Sentiment-ml128-bs32-error-fix-v6-aug"

def preprocess_m2(text: str) -> str:
    if not text or not isinstance(text, str):
        return ""
    text = str(text)
    text = re.sub(r"[\u064B-\u0652]", "", text)  # Tashkeel
    text = re.sub(r"\u0640", "", text)          # Tatweel
    text = re.sub(r"[\u0622\u0623\u0625]", "\u0627", text)
    text = re.sub(r"\u0649", "\u064A", text)
    text = re.sub(r"\u0629", "\u0647", text)
    return re.sub(r"\s+", " ", text).strip()

print("Loading Model 2 (MARBERTv2)...")
tokenizer_m2 = AutoTokenizer.from_pretrained(M2_PATH)
model_m2 = AutoModelForSequenceClassification.from_pretrained(M2_PATH).to(device)
model_m2.eval()

# ==========================================
# 3. Prediction & Ensemble Logic
# ==========================================
def predict_m1(text):
    cleaned = preprocess_m1(text)
    enc = tokenizer_m1(
        cleaned,
        add_special_tokens=True,
        max_length=256,
        padding='max_length',
        truncation=True,
        return_attention_mask=True,
        return_tensors='pt'
    )
    ids = enc['input_ids'].to(device)
    mask = enc['attention_mask'].to(device)
    
    with torch.no_grad():
        logits = model_m1(ids, mask)
        probs = F.softmax(logits, dim=1).squeeze().cpu().numpy()
    return probs

def predict_m2(text):
    cleaned = preprocess_m2(text)
    enc = tokenizer_m2(
        cleaned,
        return_tensors="pt",
        truncation=True,
        max_length=128,
        padding=True
    )
    enc = {k: v.to(device) for k, v in enc.items()}
    
    with torch.no_grad():
        logits = model_m2(**enc).logits
        probs = F.softmax(logits, dim=-1).squeeze().cpu().numpy()
    return probs

def ensemble_predict(text, weight_m1=0.5):
    if not text.strip():
        return {label: 0.0 for label in LABELS}

    weight_m2 = 1.0 - weight_m1

    # Get probability vectors [prob_neg, prob_neu, prob_pos]
    p1 = predict_m1(text)
    p2 = predict_m2(text)

    # Weighted average ensemble
    ensemble_probs = (weight_m1 * p1) + (weight_m2 * p2)

    return {
        LABELS[0]: float(ensemble_probs[0]),
        LABELS[1]: float(ensemble_probs[1]),
        LABELS[2]: float(ensemble_probs[2])
    }

# ==========================================
# 4. Minimal Gradio UI
# ==========================================
demo = gr.Interface(
    fn=ensemble_predict,
    inputs=[
        gr.Textbox(
            lines=4, 
            label="Arabic Text / النص العربي", 
            placeholder="اكتب النص هنا..."
        ),
        gr.Slider(
            minimum=0.0, 
            maximum=1.0, 
            value=0.5, 
            step=0.05, 
            label="Model 1 Weight (CAMeLBERT+LoRA vs MARBERTv2)"
        )
    ],
    outputs=gr.Label(label="Ensemble Sentiment Probabilities", num_top_classes=3),
    title="Arabic Sentiment Ensemble",
    description="Weighted average ensemble between Model 1 (CAMeLBERT+LoRA) and Model 2 (MARBERTv2).",
    examples=[
        ["هذا المطعم يقدم طعام سيء جدًا ومحروق، لا أنصح أحد بزيارته.", 0.5],
        ["أعلنت وزارة الصحة اليوم افتتاح ثلاث مستشفيات جديدة في العاصمة.", 0.5],
        ["بصراحة، الخدمة كانت ممتازة والموظفين غاية في الإحترام، أنصح بشدة!", 0.5]
    ]
)

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