# ============================================================
# Hugging Face Spaces / Gradio
# Hassaniya Sentiment Analysis – MARBERT
# Multilingual UI: Français / العربية / English
# ============================================================

import re
import html

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


# ============================================================
# 1) MODEL CONFIGURATION
# ============================================================

# Les fichiers du modèle sont directement dans main/
MODEL_PATH = "."

ID2LABEL = {
    0: "Negative",
    1: "Neutral",
    2: "Positive",
}

LABELS = {
    "fr": {
        "Negative": "Négatif",
        "Neutral": "Neutre",
        "Positive": "Positif",
    },
    "ar": {
        "Negative": "سلبي",
        "Neutral": "محايد",
        "Positive": "إيجابي",
    },
    "en": {
        "Negative": "Negative",
        "Neutral": "Neutral",
        "Positive": "Positive",
    },
}

UI_TEXT = {
    "fr": {
        "empty": "Veuillez saisir un texte.",
        "confidence": "Confiance",
        "prediction": "Prédiction finale",
        "normalized": "Texte normalisé",
    },
    "ar": {
        "empty": "يرجى إدخال نص.",
        "confidence": "نسبة الثقة",
        "prediction": "التصنيف النهائي",
        "normalized": "النص بعد التنظيف",
    },
    "en": {
        "empty": "Please enter a text.",
        "confidence": "Confidence",
        "prediction": "Final prediction",
        "normalized": "Normalized text",
    },
}

COLORS = {
    "Positive": "#2563eb",
    "Neutral": "#f59e0b",
    "Negative": "#dc2626",
}


# ============================================================
# 2) HASSANIYA CLEANING
# ============================================================

ARABIC_REPLACEMENTS = {
    "أ": "ا",
    "إ": "ا",
    "آ": "ا",
    "ى": "ي",
    "ؤ": "و",
    "ئ": "ي",
    "ة": "ه",
    "گ": "ك",
    "ڨ": "ق",
    "چ": "ج",
}

DIACRITICS = re.compile(r"[ًٌٍَُِّْـ]")


def clean_hassaniya(text: str) -> str:
    text = str(text)

    text = re.sub(r"http\S+|www\S+", " ", text)
    text = re.sub(r"@\w+|#\w+", " ", text)
    text = re.sub(r"[a-zA-Z0-9]", " ", text)
    text = re.sub(DIACRITICS, "", text)

    for old, new in ARABIC_REPLACEMENTS.items():
        text = text.replace(old, new)

    text = re.sub(r"(.)\1{2,}", r"\1", text)
    text = re.sub(r"[^\u0600-\u06FF\s]", " ", text)
    text = re.sub(r"\s+", " ", text).strip()

    return text


# ============================================================
# 3) LOAD MODEL
# ============================================================

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

print(f"Loading model from: {MODEL_PATH}")
print(f"Using device: {device}")

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_PATH
)

model.to(device)
model.eval()

print("✅ Model loaded successfully")


# ============================================================
# 4) PREDICTION
# ============================================================

def predict_sentiment(text, language):
    lang = language if language in UI_TEXT else "fr"

    if not text or not str(text).strip():
        return f"""
        <div class="result-card">
            <h3 class="empty-message">
                {UI_TEXT[lang]["empty"]}
            </h3>
        </div>
        """

    cleaned_text = clean_hassaniya(text)

    if not cleaned_text:
        return f"""
        <div class="result-card">
            <h3 class="empty-message">
                {UI_TEXT[lang]["empty"]}
            </h3>
        </div>
        """

    inputs = tokenizer(
        cleaned_text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=96,
    )

    inputs = {
        key: value.to(device)
        for key, value in inputs.items()
    }

    with torch.inference_mode():
        outputs = model(**inputs)
        probs = torch.softmax(
            outputs.logits,
            dim=-1
        )[0].detach().cpu()

    scores = {
        ID2LABEL[i]: float(probs[i])
        for i in range(len(ID2LABEL))
    }

    best_label = max(scores, key=scores.get)
    best_score = scores[best_label] * 100
    best_color = COLORS[best_label]

    bars_html = ""

    for label in ["Positive", "Neutral", "Negative"]:
        percent = scores[label] * 100
        color = COLORS[label]
        display_label = LABELS[lang][label]

        bars_html += f"""
        <div class="bar-row">

            <div
                class="bar-label"
                style="color:{color} !important;"
            >
                {display_label}
            </div>

            <div class="bar-track">
                <div
                    class="bar-fill"
                    style="
                        width:{percent:.2f}%;
                        background:{color};
                    "
                ></div>
            </div>

            <div class="bar-value">
                {percent:.2f}%
            </div>

        </div>
        """

    safe_cleaned_text = html.escape(cleaned_text)

    return f"""
    <div class="result-card">

        <div
            class="final-prediction"
            style="border-color:{best_color};"
        >

            <div class="prediction-title">
                {UI_TEXT[lang]["prediction"]}
            </div>

            <div
                class="prediction-label"
                style="color:{best_color} !important;"
            >
                {LABELS[lang][best_label]}
            </div>

            <div class="confidence">
                {UI_TEXT[lang]["confidence"]}: {best_score:.2f}%
            </div>

        </div>

        <div class="bars-container">
            {bars_html}
        </div>

        <div class="normalized-box">

            <strong>
                {UI_TEXT[lang]["normalized"]}:
            </strong>

            <div>
                {safe_cleaned_text}
            </div>

        </div>

    </div>
    """


def clear_fields():
    return "", """
    <div class="result-card">
        <h3 class="result-placeholder">
            Result will appear here
        </h3>
    </div>
    """


# ============================================================
# 5) CSS
# ============================================================

custom_css = """
.gradio-container {
    background:
        linear-gradient(
            135deg,
            #006233 0%,
            #ffffff 47%,
            #d01c1f 100%
        ) !important;

    font-family: "Segoe UI", Arial, sans-serif !important;
    min-height: 100vh !important;
    color: #111827 !important;
    overflow-x: hidden !important;
}

#main-card {
    width: calc(100% - 40px);
    max-width: 1180px;
    margin: 28px auto;
    padding: 30px;
    box-sizing: border-box;
    border-radius: 26px;
    background: rgba(255,255,255,0.97);
    border: 3px solid #ffc400;
    box-shadow: 0 18px 45px rgba(0,0,0,0.18);
}

#header-box {
    text-align: center;
    padding: 30px;
    border-radius: 22px;
    background: linear-gradient(135deg, #006233, #008f4c);
    margin-bottom: 24px;
    border-bottom: 7px solid #ffc400;
}

#header-box h1 {
    margin: 0;
    font-size: 34px;
    font-weight: 800;
    color: #ffffff !important;
}

#header-box p {
    margin-top: 10px;
    margin-bottom: 0;
    font-size: 16px;
    color: #fff7cc !important;
}

#language-selector,
#language-selector label {
    color: #111827 !important;
}

#input-box textarea {
    direction: rtl !important;
    text-align: right !important;
    font-size: 18px !important;
    line-height: 1.8 !important;
    color: #111827 !important;
    -webkit-text-fill-color: #111827 !important;
    background: #ffffff !important;
    border: 2px solid #006233 !important;
    border-radius: 14px !important;
}

#input-box textarea::placeholder {
    color: #6b7280 !important;
    -webkit-text-fill-color: #6b7280 !important;
    opacity: 1 !important;
}

button.primary {
    background: linear-gradient(135deg, #006233, #008f4c) !important;
    color: #ffffff !important;
    border-radius: 14px !important;
    border: none !important;
    font-weight: 800 !important;
}

#clear-button {
    color: #111827 !important;
    background: #ffffff !important;
    border: 1px solid #d1d5db !important;
    border-radius: 14px !important;
    font-weight: 700 !important;
}

.result-card {
    width: 100%;
    box-sizing: border-box;
    background: #ffffff !important;
    color: #111827 !important;
    border-radius: 20px;
    padding: 24px;
    border: 2px solid #ffc400;
    box-shadow: 0 12px 30px rgba(0,0,0,0.12);
    min-height: 290px;
}

.final-prediction {
    text-align: center;
    padding: 20px;
    border: 3px solid;
    border-radius: 18px;
    margin-bottom: 24px;
    background: #f9fafb !important;
}

.prediction-title {
    font-size: 15px;
    color: #374151 !important;
    font-weight: 700;
    margin-bottom: 8px;
}

.prediction-label {
    font-size: 34px;
    font-weight: 900;
    margin-bottom: 8px;
}

.confidence {
    font-size: 16px;
    color: #111827 !important;
    font-weight: 700;
}

.bars-container {
    display: flex;
    flex-direction: column;
    gap: 16px;
}

.bar-row {
    display: grid;
    grid-template-columns: 110px minmax(120px, 1fr) 80px;
    gap: 12px;
    align-items: center;
}

.bar-label {
    font-weight: 800;
    font-size: 15px;
}

.bar-track {
    width: 100%;
    height: 24px;
    background: #e5e7eb !important;
    border-radius: 999px;
    overflow: hidden;
}

.bar-fill {
    height: 100%;
    border-radius: 999px;
}

.bar-value {
    font-weight: 800;
    color: #111827 !important;
    font-size: 14px;
    text-align: right;
}

.normalized-box {
    margin-top: 22px;
    padding: 14px 16px;
    border-radius: 14px;
    background: #f3f4f6 !important;
    color: #374151 !important;
    font-size: 13px;
}

.normalized-box strong {
    color: #374151 !important;
}

.normalized-box div {
    margin-top: 6px;
    direction: rtl;
    text-align: right;
    font-size: 15px;
    color: #111827 !important;
}

.empty-message,
.result-placeholder {
    text-align: center;
    color: #006233 !important;
    padding-top: 70px;
}

/* Examples */
.gradio-container table,
.gradio-container table td,
.gradio-container table th {
    color: #111827 !important;
    background: #ffffff !important;
}

.gradio-container table * {
    color: #111827 !important;
    -webkit-text-fill-color: #111827 !important;
}

footer {
    display: none !important;
}

/* ============================================================
   MOBILE
============================================================ */

@media (max-width: 768px) {

    .gradio-container {
        padding: 0 !important;
    }

    #main-card {
        width: calc(100% - 12px) !important;
        margin: 6px auto !important;
        padding: 10px !important;
        border-radius: 14px !important;
    }

    #header-box {
        padding: 17px 10px !important;
        border-radius: 13px !important;
        margin-bottom: 12px !important;
    }

    #header-box h1 {
        font-size: 21px !important;
    }

    #header-box p {
        font-size: 12px !important;
    }

    #main-row {
        flex-direction: column !important;
        flex-wrap: nowrap !important;
        gap: 12px !important;
    }

    #main-row > div {
        width: 100% !important;
        max-width: 100% !important;
        min-width: 0 !important;
        flex: 1 1 100% !important;
    }

    #input-box textarea {
        min-height: 145px !important;
        font-size: 16px !important;
    }

    #action-row {
        gap: 8px !important;
        flex-wrap: wrap !important;
    }

    #action-row button {
        min-height: 46px !important;
        flex: 1 1 145px !important;
    }

    .result-card {
        min-height: auto !important;
        padding: 14px !important;
    }

    .prediction-label {
        font-size: 26px !important;
    }

    .bar-row {
        grid-template-columns: 68px minmax(0,1fr) 54px !important;
        gap: 6px !important;
    }

    .bar-label {
        font-size: 12px !important;
    }

    .bar-value {
        font-size: 11px !important;
    }

    .bar-track {
        height: 17px !important;
    }

    .normalized-box {
        font-size: 12px !important;
    }
}
"""


# ============================================================
# 6) GRADIO INTERFACE
# ============================================================

with gr.Blocks(
    title="Hassaniya Sentiment Analysis"
) as demo:

    with gr.Column(elem_id="main-card"):

        gr.HTML("""
        <div id="header-box">

            <h1>
                Hassaniya Sentiment Analysis
            </h1>

            <p>
                تحليل المشاعر للنصوص الحسانية باستخدام MARBERT
            </p>

            <p>
                Analyse des sentiments · Sentiment Analysis
            </p>

        </div>
        """)

        language = gr.Radio(
            choices=[
                ("Français", "fr"),
                ("العربية", "ar"),
                ("English", "en"),
            ],
            value="fr",
            label="Language / Langue / اللغة",
            elem_id="language-selector",
        )

        with gr.Row(
            equal_height=True,
            elem_id="main-row",
        ):

            # INPUT
            with gr.Column(
                scale=6,
                min_width=300,
            ):

                text_input = gr.Textbox(
                    lines=10,
                    placeholder="اكتب نصا بالحسانية هنا...",
                    label=(
                        "Texte Hassaniya / "
                        "النص الحساني / "
                        "Hassaniya Text"
                    ),
                    value="أنا ماشي عجلان يحنين حد",
                    elem_id="input-box",
                )

                with gr.Row(elem_id="action-row"):

                    analyze_btn = gr.Button(
                        "🔍 Analyser / تحليل / Analyze",
                        variant="primary",
                    )

                    clear_btn = gr.Button(
                        "🗑️ Clear / Effacer / مسح",
                        elem_id="clear-button",
                    )

            # OUTPUT
            with gr.Column(
                scale=5,
                min_width=300,
            ):

                output = gr.HTML(
                    value="""
                    <div class="result-card">

                        <h3 class="result-placeholder">
                            Result will appear here
                        </h3>

                    </div>
                    """
                )

        # Examples
        gr.Examples(
            examples=[
                ["ذ شغل زينة حته"],
                ["أنا ماشي عجلان يحنين حد"],
                ["الموقع عادي كاع ماه اكبر من كدو"],
                ["اشغل ما عجبتني"],
                ["شكرا بارك الله فيكم"],
            ],
            inputs=[text_input],
            label="Examples / Exemples / أمثلة",
        )

        # Analyze
        analyze_btn.click(
            fn=predict_sentiment,
            inputs=[
                text_input,
                language,
            ],
            outputs=output,
        )

        # Press Enter
        text_input.submit(
            fn=predict_sentiment,
            inputs=[
                text_input,
                language,
            ],
            outputs=output,
        )

        # Clear
        clear_btn.click(
            fn=clear_fields,
            inputs=None,
            outputs=[
                text_input,
                output,
            ],
        )


# ============================================================
# 7) START HUGGING FACE SPACE
# ============================================================

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