"""
Streamlit Web App for RNN Text Classification (HuggingFace Spaces)
Run with: streamlit run app_hf.py
"""
import streamlit as st
import torch
import numpy as np
import os
import sys

# Add current directory to path for imports
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# Import model and inference functions
from model import SimpleRNN

# Inference functions (inlined for HuggingFace compatibility)
def load_model(checkpoint_path, device):
    """Load trained model from checkpoint"""
    checkpoint = torch.load(checkpoint_path, map_location=device)
    
    model_config = checkpoint['model_config']
    vocab = checkpoint['vocab']
    num_classes = checkpoint['num_classes']
    
    model = SimpleRNN(
        vocab_size=len(vocab),
        num_classes=num_classes,
        **model_config
    ).to(device)
    
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    
    return model, vocab, num_classes


def preprocess_text(text, vocab, max_length=128):
    """Preprocess text for inference"""
    tokens = text.lower().split()
    sequence = [vocab.get(token, vocab['<UNK>']) for token in tokens]
    
    if len(sequence) > max_length:
        sequence = sequence[:max_length]
    else:
        sequence = sequence + [vocab['<PAD>']] * (max_length - len(sequence))
    
    return torch.tensor([sequence], dtype=torch.long)


def predict(model, text, vocab, device, max_length=128):
    """Make prediction on a single text"""
    model.eval()
    
    # Preprocess
    input_tensor = preprocess_text(text, vocab, max_length).to(device)
    
    # Predict
    with torch.no_grad():
        output = model(input_tensor)
        probabilities = torch.softmax(output, dim=1)
        predicted_class = torch.argmax(output, dim=1).item()
        confidence = probabilities[0][predicted_class].item()
    
    # Convert to numpy array safely
    probs_tensor = probabilities[0].cpu().detach()
    probs_list = probs_tensor.tolist()
    probs_array = np.array(probs_list, dtype=np.float32)
    
    return predicted_class, confidence, probs_array

# Page configuration
st.set_page_config(
    page_title="RNN Text Classifier",
    page_icon="🤖",
    layout="wide"
)

# Title
st.title("🤖 RNN Text Classification")
st.markdown("Classify text using trained Simple RNN models")

# Sidebar for model selection
st.sidebar.header("Model Selection")
dataset_choice = st.sidebar.selectbox(
    "Choose a model:",
    ["Emotion Classifier", "AG News Classifier"],
    help="Select which model to use for classification"
)

# Map selection to dataset name
dataset_map = {
    "Emotion Classifier": "emotion",
    "AG News Classifier": "ag_news"
}
dataset_name = dataset_map[dataset_choice]

# Class labels
emotion_labels = ['anger', 'fear', 'joy', 'love', 'sadness', 'surprise']
ag_news_labels = ['World', 'Sports', 'Business', 'Science/Technology']

labels_map = {
    "emotion": emotion_labels,
    "ag_news": ag_news_labels
}

# Load model (cached to avoid reloading)
@st.cache_resource
def load_cached_model(dataset_name):
    """Load model with caching"""
    # Use CPU for HuggingFace Spaces (GPU not always available)
    device = torch.device('cpu')
    checkpoint_path = f'checkpoints/best_model_{dataset_name}.pt'
    
    # Try different paths (HuggingFace might have different structure)
    possible_paths = [
        checkpoint_path,
        f'/{checkpoint_path}',  # Absolute path
        f'./{checkpoint_path}',  # Relative path
    ]
    
    model_path = None
    for path in possible_paths:
        if os.path.exists(path):
            model_path = path
            break
    
    if model_path is None:
        return None, None, None, None
    
    try:
        model, vocab, num_classes = load_model(model_path, device)
        return model, vocab, num_classes, device
    except Exception as e:
        import traceback
        error_msg = f"Error loading model: {str(e)}\n{traceback.format_exc()}"
        print(error_msg)  # Print for debugging
        return None, None, None, None

# Load model
model, vocab, num_classes, device = load_cached_model(dataset_name)

if model is None:
    st.error(f"❌ Model not found! Please ensure `checkpoints/best_model_{dataset_name}.pt` exists.")
    st.info("""
    **For HuggingFace Spaces:**
    1. Upload your trained model files to the `checkpoints/` folder in your repository
    2. Files should be named: `best_model_emotion.pt` and `best_model_ag_news.pt`
    3. Push to your HuggingFace Space repository
    """)
else:
    st.sidebar.success(f"✅ {dataset_choice} loaded successfully")
    
    # Main content area
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.subheader("📝 Enter Text to Classify")
        
        # Text input
        if dataset_name == "emotion":
            default_text = "I am feeling so happy and excited today!"
            placeholder = "Enter your text here... (e.g., 'I am feeling so happy today!')"
        else:
            default_text = "The stock market reached new highs today as investors cheered strong earnings reports."
            placeholder = "Enter your text here... (e.g., 'The stock market reached new highs today')"
        
        # Use example text if set from button click
        initial_value = st.session_state.get('example_text', default_text)
        
        user_text = st.text_area(
            "Text Input:",
            value=initial_value,
            height=150,
            placeholder=placeholder,
            help="Enter the text you want to classify",
            key="text_input"
        )
        
        # Clear example text after using it
        if 'example_text' in st.session_state:
            del st.session_state.example_text
        
        # Predict button
        predict_button = st.button("🔍 Classify", type="primary", use_container_width=True)
    
    with col2:
        st.subheader("ℹ️ About")
        if dataset_name == "emotion":
            st.markdown("""
            **Emotion Classifier**
            
            Classifies text into 6 emotions:
            - 😠 Anger
            - 😨 Fear
            - 😊 Joy
            - ❤️ Love
            - 😢 Sadness
            - 😲 Surprise
            """)
        else:
            st.markdown("""
            **AG News Classifier**
            
            Classifies news articles into 4 categories:
            - 🌍 World
            - ⚽ Sports
            - 💼 Business
            - 🔬 Science/Technology
            """)
    
    # Prediction
    if predict_button and user_text.strip():
        with st.spinner("Analyzing text..."):
            try:
                predicted_class, confidence, probabilities = predict(
                    model, user_text, vocab, device
                )
                
                labels = labels_map[dataset_name]
                
                # Display results
                st.markdown("---")
                st.subheader("📊 Classification Results")
                
                # Main result
                col1, col2, col3 = st.columns([1, 2, 1])
                with col2:
                    st.markdown(f"### Predicted: **{labels[predicted_class]}**")
                    st.markdown(f"### Confidence: **{confidence:.1%}**")
                
                # Progress bar for confidence
                st.progress(confidence)
                
                # All probabilities
                st.markdown("#### All Class Probabilities:")
                
                # Display as columns
                cols = st.columns(len(labels))
                for i, (label, prob) in enumerate(zip(labels, probabilities)):
                    with cols[i]:
                        is_predicted = i == predicted_class
                        color = "🟢" if is_predicted else "⚪"
                        st.markdown(f"**{color} {label}**")
                        st.progress(prob)
                        st.caption(f"{prob:.1%}")
                
                # Detailed breakdown
                with st.expander("📈 Detailed Breakdown"):
                    import pandas as pd
                    # Ensure probabilities is a numpy array or list
                    if hasattr(probabilities, 'tolist'):
                        probs_list = probabilities.tolist()
                    else:
                        probs_list = list(probabilities)
                    
                    df = pd.DataFrame({
                        "Class": labels,
                        "Probability": probs_list,
                        "Predicted": [i == predicted_class for i in range(len(labels))]
                    })
                    df = df.sort_values("Probability", ascending=False)
                    st.dataframe(df, use_container_width=True, hide_index=True)
                
            except Exception as e:
                st.error(f"Error during prediction: {str(e)}")
                st.info("Make sure the text is not empty and contains valid characters.")
    
    elif predict_button:
        st.warning("⚠️ Please enter some text to classify!")
    
    # Example texts
    st.markdown("---")
    st.subheader("💡 Example Texts")
    
    if dataset_name == "emotion":
        examples = [
            "I am feeling so happy and excited today!",
            "I'm really scared about what might happen tomorrow.",
            "I love spending time with my family.",
            "I feel so sad and lonely right now.",
            "I'm surprised by how well this turned out!",
            "I'm so angry about what happened yesterday."
        ]
    else:
        examples = [
            "The stock market reached new highs today as investors cheered strong earnings reports.",
            "LeBron James scored 40 points to lead his team to victory in the championship game.",
            "Scientists discover new planet that could potentially support life.",
            "Global leaders meet to discuss climate change and environmental policies."
        ]
    
    # Display examples
    example_cols = st.columns(2)
    for i, example in enumerate(examples):
        with example_cols[i % 2]:
            if st.button(f"📋 Use Example", key=f"example_{i}", use_container_width=True):
                st.session_state.example_text = example
                st.rerun()
            st.caption(f'"{example}"')
    
    # Footer
    st.markdown("---")
    st.caption("Built with Simple RNN (PyTorch) | Model trained on HuggingFace datasets")