import streamlit as st
import pickle
from collections import Counter
import re
from typing import Counter as TCounter, Tuple, List, Any

# --- 1. CONFIGURATION ---

# The n-gram range used during model training (e.g., (1, 4) for unigrams up to 4-grams)
N_GRAM_RANGE = (1, 4) 
# Filenames saved in the previous step
MODEL_FILE = 'src/emotion.pkl'
VECTORIZER_FILE = 'src/vectorizer.pkl'

# --- 2. LOAD ARTIFACTS ---

@st.cache_resource
def load_artifacts():
    """Loads the saved classifier model and DictVectorizer."""
    try:
        # Load the Model
        with open(MODEL_FILE, 'rb') as model_file:
            loaded_model = pickle.load(model_file)
        
        # Load the Vectorizer
        with open(VECTORIZER_FILE, 'rb') as vec_file:
            loaded_vectorizer = pickle.load(vec_file)
        
        return loaded_model, loaded_vectorizer
    except FileNotFoundError:
        st.error(f"Error: Required files ({MODEL_FILE} or {VECTORIZER_FILE}) not found. Please ensure they are in the correct directory.")
        return None, None

# --- 3. FEATURE EXTRACTION FUNCTION (CRITICAL!) ---

# This function MUST be identical to the one used during model training.
def create_feature(text: str, nrange: Tuple[int, int]) -> TCounter[str]:
    """Extracts n-gram and punctuation features from a text string."""
    text_features: List[str] = []
    lower_text = text.lower()
    
    # Word N-Grams
    text_alphanum = re.sub(r'[^a-z0-9#]', ' ', lower_text)
    token = text_alphanum.split() 

    for n in range(nrange[0], nrange[1] + 1):
        if n > 0 and n <= len(token):
            ngrams = [' '.join(token[i-n:i]) for i in range(n, len(token) + 1)]
            text_features.extend(ngrams)

    # Punctuation Features
    text_punc = re.sub(r'[a-z0-9]', ' ', lower_text)
    text_features.extend(text_punc.split())
    
    return Counter(text_features)

# --- 4. PREDICTION LOGIC ---

def predict_emotion(text: str, model: Any, vectorizer: Any) -> str:
    """Processes text and returns the predicted emotion label."""
    if not text:
        return "Please enter text for analysis."

    # 1. Convert Text to Features (Counter object)
    feature_counter = create_feature(text, N_GRAM_RANGE)
    
    # 2. Vectorization (DictVectorizer expects a list of dictionaries/Counters)
    # The saved vectorizer is used to ensure the features are ordered correctly.
    X_processed = vectorizer.transform([feature_counter]) 
    
    # 3. Make Prediction
    prediction = model.predict(X_processed)
    
    return prediction[0]

# --- 5. STREAMLIT UI ---

# Load the model artifacts
model, vectorizer = load_artifacts()

if model and vectorizer:
    st.title("💬 Text Emotion Detection App")
    st.markdown("Enter a sentence or text below to see the predicted emotion.")
    
    # User input area
    user_input = st.text_area("Enter your text here:", "")
    
    if st.button("Predict Emotion"):
        if user_input:
            # Perform prediction
            result = predict_emotion(user_input, model, vectorizer)
            
            # --- Result Display (Using visual cues) ---
            st.subheader("Analysis Result")
            
            # Simple color-coding based on common emotions
            if result == 'joy':
                st.success(f"Predicted Emotion: **{result.upper()}** 🎉")
            elif result in ['sadness', 'fear']:
                st.warning(f"Predicted Emotion: **{result.upper()}** 😟")
            elif result == 'anger':
                st.error(f"Predicted Emotion: **{result.upper()}** 😡")
            elif result == 'disgust':
                st.markdown(f"Predicted Emotion: **{result.upper()}** 🤢", unsafe_allow_html=True)
            else:
                st.info(f"Predicted Emotion: **{result.upper()}**")
        else:
            st.warning("Please enter some text to predict.")