"""
How Am I Feeling - Emotion Classification Web App
Hugging Face Space version (runs on port 7860)
"""

from flask import Flask, render_template, request, jsonify
import pandas as pd
import numpy as np
from gensim.models import Word2Vec
from nltk.tokenize import word_tokenize
from tensorflow import keras
import pickle
import re
import os
import nltk

# Download NLTK data at startup
try:
    nltk.data.find('tokenizers/punkt')
except LookupError:
    nltk.download('punkt')
    nltk.download('punkt_tab')

# Initialize Flask app
app = Flask(__name__)

# Load models at startup
print("Loading emotion classification models...")

try:
    # Load Word2Vec model (optimized version)
    w2v_model = Word2Vec.load('models/word2vec_optimized.model')

    # Load neural network classifier (without compile to avoid optimizer issues)
    classifier = keras.models.load_model('models/best_model.keras', compile=False)

    # Load label encoder
    with open('models/label_encoder.pkl', 'rb') as f:
        label_encoder = pickle.load(f)

    # Get vector size
    vector_size = w2v_model.wv.vector_size

    print("All models loaded successfully!")
    MODELS_LOADED = True

except Exception as e:
    print(f"Error loading models: {e}")
    print("Please run the training scripts first (02-06.py files)")
    MODELS_LOADED = False

# Emotion colors and emojis for UI (20 emotions)
EMOTION_CONFIG = {
    'happiness': {
        'color': '#FFD700',
        'emoji': ':)',
        'gradient': 'linear-gradient(135deg, #FFD700, #FFA500)'
    },
    'sadness': {
        'color': '#4682B4',
        'emoji': ':(',
        'gradient': 'linear-gradient(135deg, #4682B4, #1E90FF)'
    },
    'fear': {
        'color': '#9370DB',
        'emoji': 'O_O',
        'gradient': 'linear-gradient(135deg, #9370DB, #8A2BE2)'
    },
    'anger': {
        'color': '#DC143C',
        'emoji': '>:(',
        'gradient': 'linear-gradient(135deg, #DC143C, #B22222)'
    },
    'disgust': {
        'color': '#8B4513',
        'emoji': 'X_X',
        'gradient': 'linear-gradient(135deg, #8B4513, #A0522D)'
    },
    'surprise': {
        'color': '#FFD700',
        'emoji': ':O',
        'gradient': 'linear-gradient(135deg, #FFD700, #FFFF00)'
    },
    'love': {
        'color': '#FF1493',
        'emoji': '<3',
        'gradient': 'linear-gradient(135deg, #FF1493, #FF69B4)'
    },
    'excitement': {
        'color': '#FF6347',
        'emoji': '!!!',
        'gradient': 'linear-gradient(135deg, #FF6347, #FF4500)'
    },
    'embarrassment': {
        'color': '#FF69B4',
        'emoji': '>///<',
        'gradient': 'linear-gradient(135deg, #FF69B4, #FF1493)'
    },
    'loneliness': {
        'color': '#708090',
        'emoji': '...',
        'gradient': 'linear-gradient(135deg, #708090, #778899)'
    },
    'anxiety': {
        'color': '#9370DB',
        'emoji': ':S',
        'gradient': 'linear-gradient(135deg, #9370DB, #7B68EE)'
    },
    'frustration': {
        'color': '#FF8C00',
        'emoji': '>:|',
        'gradient': 'linear-gradient(135deg, #FF8C00, #FFA500)'
    },
    'guilt': {
        'color': '#696969',
        'emoji': ':/',
        'gradient': 'linear-gradient(135deg, #696969, #808080)'
    },
    'disappointment': {
        'color': '#6A5ACD',
        'emoji': ':-/',
        'gradient': 'linear-gradient(135deg, #6A5ACD, #7B68EE)'
    },
    'jealousy': {
        'color': '#2E8B57',
        'emoji': ':-|',
        'gradient': 'linear-gradient(135deg, #2E8B57, #3CB371)'
    },
    'gratitude': {
        'color': '#98FB98',
        'emoji': '^_^',
        'gradient': 'linear-gradient(135deg, #98FB98, #90EE90)'
    },
    'pride': {
        'color': '#FFD700',
        'emoji': 'B)',
        'gradient': 'linear-gradient(135deg, #FFD700, #DAA520)'
    },
    'relief': {
        'color': '#87CEEB',
        'emoji': '*phew*',
        'gradient': 'linear-gradient(135deg, #87CEEB, #87CEFA)'
    },
    'hope': {
        'color': '#FFA07A',
        'emoji': ':)',
        'gradient': 'linear-gradient(135deg, #FFA07A, #FA8072)'
    },
    'confusion': {
        'color': '#D3D3D3',
        'emoji': '???',
        'gradient': 'linear-gradient(135deg, #D3D3D3, #C0C0C0)'
    },
}

# Preprocessing function
def preprocess_text(text):
    """Clean text the same way we did during training."""
    if pd.isna(text):
        return ""

    text = str(text).lower()
    text = re.sub(r'http\S+|www\S+|https\S+', '', text, flags=re.MULTILINE)
    text = re.sub(r'@\w+', '', text)
    text = re.sub(r'#\w+', '', text)

    harmful_punctuation = '"#$%&()*+-/:;<=>@[\\]^_`{|}~'
    text = text.translate(str.maketrans('', '', harmful_punctuation))
    text = re.sub(r'\s+', ' ', text).strip()

    return text

# Sentence to vector function
def sentence_to_vector(sentence):
    """Convert sentence to vector."""
    words = word_tokenize(sentence.lower())
    word_vectors = [w2v_model.wv[word] for word in words if word in w2v_model.wv]

    if len(word_vectors) == 0:
        return np.zeros(vector_size)

    return np.mean(word_vectors, axis=0)

# Prediction function
def predict_emotion(sentence, top_k=5):
    """
    Predict emotion for a sentence.
    Returns list of (emotion, confidence, config) tuples
    """
    if not MODELS_LOADED:
        return None

    # Preprocess
    cleaned = preprocess_text(sentence)

    if len(cleaned) == 0:
        return None

    # Convert to vector
    vector = sentence_to_vector(cleaned).reshape(1, -1)

    # Make prediction
    probs = classifier.predict(vector, verbose=0)[0]

    # Get top-k predictions
    top_indices = np.argsort(probs)[-top_k:][::-1]

    # Build results
    results = []
    for idx in top_indices:
        emotion = label_encoder.inverse_transform([idx])[0]
        confidence = float(probs[idx])

        # Get emotion config
        config = EMOTION_CONFIG.get(emotion, {
            'color': '#808080',
            'emoji': '?',
            'gradient': 'linear-gradient(135deg, #808080, #A9A9A9)'
        })

        results.append({
            'emotion': emotion,
            'confidence': confidence,
            'percentage': round(confidence * 100, 1),
            'emoji': config['emoji'],
            'color': config['color'],
            'gradient': config['gradient']
        })

    return results

# Routes
@app.route('/')
def index():
    """Main page"""
    return render_template('index.html', models_loaded=MODELS_LOADED)

@app.route('/analyze', methods=['POST'])
def analyze():
    """Analyze text and return emotion predictions"""
    data = request.json
    text = data.get('text', '')

    if not text.strip():
        return jsonify({'error': 'Please enter some text'}), 400

    if not MODELS_LOADED:
        return jsonify({'error': 'Models not loaded. Please run training scripts first.'}), 500

    predictions = predict_emotion(text, top_k=5)

    if predictions is None:
        return jsonify({'error': 'Unable to process text'}), 400

    return jsonify({
        'success': True,
        'text': text,
        'predictions': predictions
    })

@app.route('/health')
def health():
    """Health check endpoint"""
    return jsonify({
        'status': 'healthy',
        'models_loaded': MODELS_LOADED
    })

if __name__ == '__main__':
    # Ensure templates directory exists
    os.makedirs('templates', exist_ok=True)

    # Hugging Face Spaces runs on port 7860
    port = int(os.environ.get('PORT', 7860))

    print("\n" + "="*70)
    print("HOW AM I FEELING - Emotion Detection App")
    print("="*70)
    print(f"Starting server on port {port}")
    print("Enter your text to analyze your emotions!")
    print("="*70 + "\n")

    app.run(host='0.0.0.0', port=port, debug=False)