import streamlit as st
import pandas as pd
import numpy as np
import re
import nltk
import pickle
import os
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# ========== STREAMLIT SETUP ==========
st.set_page_config(
    page_title="🎬 IMDB Sentiment Analysis",
    page_icon="🎬",
    layout="centered"
)

st.markdown("""
<style>
    .prediction-box {
        padding: 25px;
        border-radius: 15px;
        margin: 25px 0;
        text-align: center;
        font-size: 1.4rem;
        font-weight: bold;
    }
    .positive { background-color: #d4edda; color: #155724; border: 2px solid #c3e6cb; }
    .negative { background-color: #f8d7da; color: #721c24; border: 2px solid #f5c6cb; }
    .header { text-align: center; margin-bottom: 30px; }
    .header h1 { color: #1f77b4; margin-bottom: 5px; }
    .header p { color: #666; font-size: 1.1rem; }
</style>
""", unsafe_allow_html=True)

st.markdown('<div class="header"><h1>🎬 IMDB Sentiment Analysis</h1><p>Powered by Naive Bayes • No dataset upload required</p></div>', unsafe_allow_html=True)

# Status indicator
status = st.empty()
status.markdown('<span class="status-badge status-training">⏳ Loading model...</span>', unsafe_allow_html=True)

# ========== NLTK SETUP (Safe for HF Spaces) ==========
@st.cache_resource
def setup_nltk():
    try:
        nltk.data.find('tokenizers/punkt')
        nltk.data.find('corpora/stopwords')
    except LookupError:
        nltk.download('punkt', quiet=True, download_dir='/tmp/nltk')
        nltk.download('stopwords', quiet=True, download_dir='/tmp/nltk')
        nltk.data.path.append('/tmp/nltk')

setup_nltk()

# ========== TEXT PREPROCESSING ==========
def preprocess_text(text):
    text = re.sub(r'<.*?>', '', text)  # Remove HTML tags
    text = text.lower()
    text = re.sub(r'[^a-zA-Z\s]', '', text)  # Keep only letters/spaces
    text = ' '.join(text.split())  # Remove extra whitespace
    
    # Stopwords removal
    from nltk.corpus import stopwords
    stop_words = set(stopwords.words('english'))
    words = text.split()
    words = [w for w in words if w not in stop_words]
    
    # Stemming
    from nltk.stem.porter import PorterStemmer
    stemmer = PorterStemmer()
    words = [stemmer.stem(w) for w in words]
    
    return ' '.join(words)

# ========== DATASET LOADING (Auto-download fallback) ==========
@st.cache_data
def load_dataset():
    """Load IMDB dataset - tries HF Datasets first, falls back to embedded mini-dataset"""
    try:
        # Try loading from Hugging Face Datasets (no CSV needed!)
        from datasets import load_dataset
        dataset = load_dataset("imdb", split="train[:1200]")
        df = pd.DataFrame({
            "review": dataset["text"],
            "sentiment": ["positive" if l == 1 else "negative" for l in dataset["label"]]
        })
        status.markdown('<span class="status-badge status-ready">✅ Using IMDB dataset (1,200 samples)</span>', unsafe_allow_html=True)
        return df
    except Exception as e:
        # Fallback: Embedded mini-dataset (always works!)
        status.markdown('<span class="status-badge status-ready">✅ Using embedded dataset (100 samples)</span>', unsafe_allow_html=True)
        mini_data = """review,sentiment
"A wonderful little production.",positive
"This is a very strange movie.",negative
"Very good!",positive
"Terrible acting",negative
"Brilliant direction",positive
"Awful dialogue",negative
"Fantastic performances",positive
"Boring plot",negative
"Emotional rollercoaster",positive
"Predictable ending",negative
"Visually stunning",positive
"Confusing storyline",negative
"Powerful message",positive
"Uninspired acting",negative
"Perfect pacing",positive
"Dragged on forever",negative
"Masterpiece!",positive
"Total garbage",negative
"Captivating from start to finish",positive
"Couldn't finish it",negative
"Oscar-worthy",positive
"Should be banned",negative
"Beautiful cinematography",positive
"Horrible soundtrack",negative
"Thought-provoking",positive
"Mind-numbingly dull",negative
"Instant classic",positive
"Instant regret",negative
"Flawless execution",positive
"Full of flaws",negative
"Left me speechless",positive
"Left me sleeping",negative
"Pure joy",positive
"Pure torture",negative
"Will watch again",positive
"Will never recover",negative
"Exceeded expectations",positive
"Failed completely",negative
"Hidden gem",positive
"Overhyped trash",negative
"Emotional depth",positive
"Emotionally vacant",negative
"Perfect casting",positive
"Miscast disaster",negative
"Artistic triumph",positive
"Artistic failure",negative
"Hauntingly beautiful",positive
"Hauntingly bad",negative
"Unforgettable",positive
"Instantly forgettable",negative
"Chills down my spine",positive
"Chills from boredom",negative
"Standing ovation",positive
"Walking out early",negative
"Cinematic poetry",positive
"Cinematic crime",negative
"Must see",positive
"Must skip",negative
"Life-changing",positive
"Waste of life",negative
"Perfection",positive
"Imperfect mess",negative
"Brilliant",positive
"Brainless",negative
"Captivating",positive
"Captivity",negative
"Enchanting",positive
"Enraging",negative
"Exhilarating",positive
"Exhausting",negative
"Mesmerizing",positive
"Mind-numbing",negative
"Riveting",positive
"Repulsive",negative
"Soul-stirring",positive
"Soul-crushing",negative
"Transcendent",positive
"Transcendently bad",negative
"Unmissable",positive
"Unwatchable",negative
"Visceral",positive
"Vapid",negative
"Wow",positive
"Ugh",negative
"Yes!",positive
"No!",negative
"Amazing movie!",positive
"Waste of time",negative
"Emotional rollercoaster",positive
"Predictable ending",negative
"Visually stunning",positive
"Confusing storyline",negative
"Powerful message",positive
"Uninspired acting",negative
"Perfect pacing",positive
"Dragged on forever",negative
"""
        from io import StringIO
        return pd.read_csv(StringIO(mini_data))

# ========== MODEL TRAINING (Cached) ==========
@st.cache_resource
def train_model():
    df = load_dataset()
    
    # Preprocess reviews
    df['clean_review'] = df['review'].apply(preprocess_text)
    
    # Vectorize
    vectorizer = CountVectorizer(max_features=1000)
    X = vectorizer.fit_transform(df['clean_review'])
    
    # Encode labels
    y = df['sentiment'].map({'positive': 1, 'negative': 0}).values
    
    # Split data CORRECTLY (fixes notebook bug!)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
    
    # Train model
    model = MultinomialNB()
    model.fit(X_train, y_train)
    
    # Calculate accuracy
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    
    return model, vectorizer, accuracy

# Load model (shows status during first run)
try:
    model, vectorizer, accuracy = train_model()
    status.empty()  # Clear status after success
except Exception as e:
    status.error(f"❌ Error loading model: {str(e)}")
    st.stop()

# ========== UI ==========
st.subheader("📝 Enter Your Movie Review")
user_input = st.text_area(
    "Type your review below:",
    height=120,
    placeholder="Example: 'This movie was absolutely fantastic! The acting was superb...'"
)

if st.button("🔍 Analyze Sentiment", type="primary", use_container_width=True):
    if not user_input.strip():
        st.warning("⚠️ Please enter a review first!")
    else:
        with st.spinner("Analyzing sentiment..."):
            # Preprocess & predict
            clean_text = preprocess_text(user_input)
            X = vectorizer.transform([clean_text])
            pred = model.predict(X)[0]
            proba = model.predict_proba(X)[0]
            confidence = max(proba) * 100
            
            # Display result
            if pred == 1:
                st.markdown(f"""
                <div class="prediction-box positive">
                    😊 POSITIVE REVIEW<br>
                    <span style="font-size: 1.1rem; font-weight: normal;">Confidence: {confidence:.1f}%</span>
                </div>
                """, unsafe_allow_html=True)
            else:
                st.markdown(f"""
                <div class="prediction-box negative">
                    😞 NEGATIVE REVIEW<br>
                    <span style="font-size: 1.1rem; font-weight: normal;">Confidence: {confidence:.1f}%</span>
                </div>
                """, unsafe_allow_html=True)

# Examples section
st.markdown("---")
st.subheader("💡 Try These Examples")
cols = st.columns(2)
with cols[0]:
    if st.button("👍 Positive Example"):
        st.session_state.user_input = "This movie was absolutely fantastic! The acting was superb and the plot kept me engaged throughout."
        st.rerun()
with cols[1]:
    if st.button("👎 Negative Example"):
        st.session_state.user_input = "Terrible waste of time. Poor acting, boring storyline, and awful special effects."
        st.rerun()

# Footer with accuracy
st.markdown("<hr>", unsafe_allow_html=True)
st.markdown(
    f"<div style='text-align: center; color: #888; font-size: 0.9rem;'>"
    f"✨ Model accuracy: {accuracy*100:.1f}% • No dataset upload required • Built with Streamlit"
    "</div>",
    unsafe_allow_html=True
)