import gradio as gr
import numpy as np
import pickle
import re
import nltk
from nltk.corpus import stopwords

nltk.download('stopwords', quiet=True)

# ─── Load baseline models (always available) ───────────────────────────────────
with open('logistic_regression.pkl', 'rb') as f:
    lr_model = pickle.load(f)

with open('tfidf_vectorizer.pkl', 'rb') as f:
    tfidf = pickle.load(f)

# ─── Try loading BiLSTM ────────────────────────────────────────────────────────
LSTM_AVAILABLE = False
try:
    import tensorflow as tf
    lstm_model = tf.keras.models.load_model('best_bilstm.keras')
    with open('tokenizer.pkl', 'rb') as f:
        tokenizer = pickle.load(f)
    from tf_keras.preprocessing.sequence import pad_sequences
    LSTM_AVAILABLE = True
    print("BiLSTM loaded successfully!")
except Exception as e:
    print(f"BiLSTM not loaded: {e}")

# ─── Try loading pretrained transformer (Twitter-RoBERTa) ─────────────────────
# This is NOT trained by the author — it's included as a reference point so visitors
# can see how a large pretrained model handles the same input. 
TRANSFORMER_AVAILABLE = False
try:
    from transformers import pipeline
    transformer_pipe = pipeline(
        "sentiment-analysis",
        model="cardiffnlp/twitter-roberta-base-sentiment-latest"
    )
    TRANSFORMER_AVAILABLE = True
    print("Transformer loaded successfully!")
except Exception as e:
    print(f"Transformer not loaded: {e}")

# ─── Preprocessing ─────────────────────────────────────────────────────────────
NEGATION_WORDS = {
    "not", "no", "never", "neither", "nobody", "nothing", "nowhere",
    "nor", "cannot", "can't", "won't", "don't", "doesn't", "didn't",
    "isn't", "aren't", "wasn't", "weren't", "hasn't", "haven't",
    "hadn't", "wouldn't", "shouldn't", "couldn't", "mustn't"
}
STOP_WORDS = set(stopwords.words('english')) - NEGATION_WORDS

def preprocess(text):
    text = text.lower()
    text = re.sub(r'http\S+|www\S+|https\S+', '', text)
    text = re.sub(r'@\w+', '', text)
    text = re.sub(r'#', '', text)
    text = re.sub(r'[^a-zA-Z\s]', '', text)
    tokens = [w for w in text.split() if w not in STOP_WORDS or w in NEGATION_WORDS]
    return ' '.join(tokens)

# ─── Prediction ────────────────────────────────────────────────────────────────
MAX_LEN = 50

MODEL_CHOICES = [
    "Logistic Regression Fast",
    "BiLSTM Deep Learning",
    "Twitter-RoBERTa Pretrained Transformer",
]

def predict_sentiment(text, model_choice):
    if not text.strip():
        return "Please enter some text.", "", "", ""

    cleaned = preprocess(text)

    if model_choice == "Twitter-RoBERTa Pretrained Transformer":
        if TRANSFORMER_AVAILABLE:
            # Transformer gets the RAW text, not the stopword-stripped version —
            # it needs full sentence context (that's the whole point of using it).
            # Pull ALL three class scores (negative/neutral/positive) instead of
            # just the single top label, so "neutral" doesn't get silently
            # collapsed into "positive."
            raw_out = transformer_pipe(text[:512], top_k=None)
            all_scores = raw_out[0] if isinstance(raw_out[0], list) else raw_out
            scores = {d['label'].lower(): d['score'] for d in all_scores}
            prob_pos = scores.get('positive', 0.0)
            prob_neg = scores.get('negative', 0.0)
            prob_neu = scores.get('neutral', 0.0)

            # Binary decision: compare positive vs negative directly, ignoring
            # which one is nominally "top" — this lets a negative lean surface
            # even on phrases the model scores mostly "neutral."
            label = 1 if prob_pos >= prob_neg else 0
            conf  = prob_pos if label == 1 else prob_neg

            model_used = "cardiffnlp/twitter-roberta-base (pretrained, not trained by author)"
            cleaned = (f"(raw text used — model scores: pos={prob_pos:.2f}, "
                       f"neu={prob_neu:.2f}, neg={prob_neg:.2f})")
        else:
            model_choice = "Logistic Regression Fast"  # fallback

    if model_choice != "Twitter-RoBERTa Pretrained Transformer":
        if model_choice == "BiLSTM Deep Learning" and LSTM_AVAILABLE:
            seq = pad_sequences(
                tokenizer.texts_to_sequences([cleaned]),
                maxlen=MAX_LEN, padding='post'
            )
            prob_pos = float(lstm_model.predict(seq, verbose=0)[0][0])
            label = 1 if prob_pos >= 0.5 else 0
            conf  = prob_pos if label == 1 else 1 - prob_pos
            model_used = "BiLSTM (GloVe embeddings)"
        else:
            if model_choice == "BiLSTM Deep Learning" and not LSTM_AVAILABLE:
                model_used = "Logistic Regression (BiLSTM unavailable, fallback)"
            else:
                model_used = "Logistic Regression (TF-IDF)"
            vec   = tfidf.transform([cleaned])
            prob  = lr_model.predict_proba(vec)[0]
            label = int(lr_model.predict(vec)[0])
            conf  = prob[label]

    sentiment  = "Positive" if label == 1 else "Negative"
    confidence = f"{conf * 100:.1f}%"
    cleaned_display = cleaned if cleaned else "(empty after preprocessing)"

    return sentiment, confidence, model_used, cleaned_display

# ─── Gradio UI ─────────────────────────────────────────────────────────────────
example_texts = [
    "I love this product! It works perfectly.",
    "This is the worst experience I have ever had.",
    "I can't believe how good this is!",
    "Not happy with this at all.",
    "Just okay, nothing special about it.",
    "I have had better.",
]

# Single accent color reused everywhere: title, submit button, focus rings.
ACCENT = "#FF7A45"
ACCENT_SOFT = "#FF7A4522"
BG = "#0b0b0d"
PANEL = "#151517"
BORDER = "#2a2a2e"
TEXT = "#f2f2f2"
TEXT_DIM = "#a8a8ac"

CSS = f"""
:root {{
    --accent: {ACCENT};
}}

.gradio-container {{
    background: {BG} !important;
    color: {TEXT} !important;
    font-size: 14px !important;
}}

/* ---------- Title ---------- */
.title-wrap {{
    text-align: center;
    padding: 6px 0 2px 0;
}}
.title-wrap h1 {{
    font-size: 1.6em !important;
    font-weight: 700 !important;
    background: linear-gradient(90deg, {ACCENT}, #ffb199);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin-bottom: 2px !important;
}}
.title-wrap p, .title-wrap strong {{
    color: {TEXT_DIM} !important;
    font-size: 0.9em !important;
    font-weight: 400 !important;
}}

hr {{
    border-color: {BORDER} !important;
    opacity: 0.6;
}}

/* ---------- Generic text sizing ---------- */
label, .label-wrap span, .gr-markdown, p, li {{
    font-size: 0.85em !important;
    color: {TEXT} !important;
}}

/* ---------- Inputs (textbox) ---------- */
textarea, input[type="text"] {{
    background: {PANEL} !important;
    color: {TEXT} !important;
    border: 1px solid {BORDER} !important;
    font-size: 0.85em !important;
    border-radius: 8px !important;
}}
textarea:focus, input[type="text"]:focus {{
    border-color: {ACCENT} !important;
    box-shadow: 0 0 0 2px {ACCENT_SOFT} !important;
}}
textarea::placeholder {{
    color: {TEXT_DIM} !important;
}}

/* ---------- Output boxes ---------- */
.result-box textarea {{
    font-size: 1em !important;
    font-weight: 600 !important;
    color: {ACCENT} !important;
}}

/* ---------- Radio (model choice) ---------- */
.model-radio {{
    background: transparent !important;
    border: none !important;
}}
.model-radio > label {{
    font-size: 0.85em !important;
    margin-bottom: 4px !important;
}}
.model-radio .wrap {{
    background: transparent !important;
    gap: 6px !important;
}}
.model-radio label.selected,
.model-radio label {{
    background: {PANEL} !important;
    color: {TEXT} !important;
    border: 1px solid {BORDER} !important;
    border-radius: 8px !important;
    padding: 6px 12px !important;
    font-size: 0.85em !important;
    font-weight: 400 !important;
}}
.model-radio input[type="radio"]:checked + span,
.model-radio label:has(input:checked) {{
    border-color: {ACCENT} !important;
    color: {ACCENT} !important;
}}
.model-radio span {{
    color: inherit !important;
    background: transparent !important;
}}

/* ---------- Submit button (accent) ---------- */
.submit-btn {{
    background: {ACCENT} !important;
    color: #1a0f0a !important;
    border: none !important;
    font-weight: 700 !important;
    font-size: 0.9em !important;
    border-radius: 8px !important;
}}
.submit-btn:hover {{
    filter: brightness(1.08);
}}

/* ---------- Example buttons ---------- */
.example-btn {{
    margin: 3px !important;
    padding: 5px 12px !important;
    font-size: 0.78em !important;
    background: transparent !important;
    color: {TEXT_DIM} !important;
    border: 1px solid {BORDER} !important;
    border-radius: 999px !important;
    cursor: pointer !important;
    transition: all 0.15s !important;
}}
.example-btn:hover {{
    background: {ACCENT_SOFT} !important;
    color: {ACCENT} !important;
    border-color: {ACCENT} !important;
}}
.example-label {{
    font-size: 0.78em !important;
    color: {TEXT_DIM} !important;
    margin-top: 10px !important;
}}

/* ---------- Panels / columns ---------- */
.panel-card {{
    background: {PANEL} !important;
    border: 1px solid {BORDER} !important;
    border-radius: 12px !important;
    padding: 14px !important;
}}
"""

with gr.Blocks(title="Sentiment Analysis · Jimma University", css=CSS, theme=gr.themes.Base()) as demo:

    with gr.Column(elem_classes="title-wrap"):
        gr.Markdown("""
# Sentiment Analysis of Social Media Posts
**Masters Capstone Project · Jimma University, Institute of Technology**

Classify tweet sentiment as **Positive** or **Negative** using traditional ML and deep learning.
        """)

    gr.HTML("<hr/>")

    with gr.Row():
        with gr.Column(scale=3, elem_classes="panel-card"):
            text_input = gr.Textbox(
                label="Enter a tweet or social media post",
                placeholder="e.g. I can't believe how amazing this is!",
                lines=4
            )

            model_choice = gr.Radio(
                choices=MODEL_CHOICES,
                value=MODEL_CHOICES[0],
                label="Choose model",
                elem_classes="model-radio"
            )

            submit_btn = gr.Button("Analyze Sentiment", elem_classes="submit-btn", size="lg")

            gr.Markdown("**Try an example:**", elem_classes="example-label")
            with gr.Row():
                for example in example_texts[:3]:
                    gr.Button(example, elem_classes="example-btn").click(
                        fn=lambda t=example: t,
                        outputs=text_input
                    )
            with gr.Row():
                for example in example_texts[3:]:
                    gr.Button(example, elem_classes="example-btn").click(
                        fn=lambda t=example: t,
                        outputs=text_input
                    )

        with gr.Column(scale=2, elem_classes="panel-card"):
            sentiment_out  = gr.Textbox(label="Predicted Sentiment",
                                         interactive=False, elem_classes="result-box")
            confidence_out = gr.Textbox(label="Confidence",
                                         interactive=False)
            model_used_out = gr.Textbox(label="Model Used",
                                         interactive=False)
            cleaned_out    = gr.Textbox(label="Text After Preprocessing",
                                         interactive=False)

    gr.HTML("<hr/>")

    gr.Markdown("""
### How it works
1. **Preprocessing** — URLs, @mentions, and punctuation are removed. Stopwords are filtered
   but *negation words* (not, never, don't…) are deliberately preserved.
2. **Logistic Regression** — TF-IDF features → fast classical classifier (test accuracy: **80.33%**)
3. **BiLSTM** — GloVe Twitter embeddings → Bidirectional LSTM → deep learning classifier (test accuracy: **80.71%**)
4. **Twitter-RoBERTa** — a large pretrained transformer (not trained by the author) shown for
   comparison. It sees full sentence context and generally handles implicit/comparative
   sentiment and mild sarcasm far better than the two models above —
   illustrating the accuracy/context trade-off discussed in the capstone report.

> *Capstone project — M.Sc. Data Science, Jimma University 2026*
    """)

    submit_btn.click(
        fn=predict_sentiment,
        inputs=[text_input, model_choice],
        outputs=[sentiment_out, confidence_out, model_used_out, cleaned_out]
    )

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