#!/usr/bin/env python3
"""
Financial Sentiment Analysis - Enhanced Ensemble Gradio Demo for Hugging Face Space
3-Model Ensemble System with Rule Engine
"""

import gradio as gr
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import logging
import re
from typing import Dict, List, Tuple

# Logging setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class SentimentRuleEngine:
    """Rule-based post-processing for sentiment analysis"""
    
    def __init__(self):
        # Strong bullish keywords with weights
        self.bullish_keywords = {
            'soaring': 0.9, 'skyrocketing': 0.9, 'surging': 0.9, 'exploding': 0.9,
            'excellent': 0.8, 'outstanding': 0.8, 'exceptional': 0.8, 'amazing': 0.8,
            'breakthrough': 0.8, 'revolutionary': 0.8, 'record-breaking': 0.9,
            'all-time high': 0.9, 'new high': 0.8, 'moon': 0.8, 'rocket': 0.8,
            'mooning': 0.9, 'rocketing': 0.8, 'booming': 0.7, 'thriving': 0.7,
            'up 10%': 0.8, 'up 15%': 0.9, 'up 20%': 0.9, 'gained 10%': 0.8,
            'rose 15%': 0.8, 'jumped 20%': 0.9, 'spiked': 0.8, 'surged': 0.8,
            'rising': 0.6, 'climbing': 0.6, 'gaining': 0.6, 'growing': 0.6,
            'strong': 0.5, 'solid': 0.5, 'robust': 0.5, 'healthy': 0.5,
            'positive': 0.4, 'optimistic': 0.5, 'bullish': 0.8, 'rally': 0.7,
            'beat': 0.7, 'exceeded': 0.7, 'outperformed': 0.7, 'success': 0.6,
            'profit': 0.3, 'earnings': 0.2, 'revenue': 0.2, 'growth': 0.5
        }
        
        # Strong bearish keywords with weights
        self.bearish_keywords = {
            'crashing': 0.9, 'plummeting': 0.9, 'collapsing': 0.9, 'tanking': 0.9,
            'disaster': 0.8, 'terrible': 0.8, 'awful': 0.8, 'horrible': 0.8,
            'crisis': 0.7, 'recession': 0.8, 'bankruptcy': 0.9, 'failed': 0.7,
            'down 10%': 0.8, 'down 15%': 0.9, 'down 20%': 0.9, 'lost 10%': 0.8,
            'fell 15%': 0.8, 'dropped 20%': 0.9, 'plunged': 0.8, 'tumbled': 0.7,
            'falling': 0.6, 'declining': 0.6, 'dropping': 0.6, 'losing': 0.6,
            'weak': 0.5, 'poor': 0.5, 'bad': 0.4, 'negative': 0.4,
            'bearish': 0.8, 'selloff': 0.7, 'sell-off': 0.7, 'correction': 0.6,
            'missed': 0.6, 'disappointed': 0.6, 'concerns': 0.4, 'worried': 0.5
        }
        
        # Neutral keywords that should reduce extreme predictions
        self.neutral_keywords = {
            'mixed': 0.7, 'uncertain': 0.6, 'unclear': 0.6, 'sideways': 0.8,
            'flat': 0.7, 'stable': 0.5, 'unchanged': 0.8, 'waiting': 0.6,
            'consolidating': 0.7, 'range-bound': 0.8, 'choppy': 0.7
        }
    
    def extract_keywords(self, text: str) -> Dict[str, float]:
        """Extract sentiment keywords and their weights from text"""
        text_lower = text.lower()
        found_keywords = {'bullish': [], 'bearish': [], 'neutral': []}
        
        # Check for bullish keywords
        for keyword, weight in self.bullish_keywords.items():
            if keyword in text_lower:
                found_keywords['bullish'].append((keyword, weight))
        
        # Check for bearish keywords  
        for keyword, weight in self.bearish_keywords.items():
            if keyword in text_lower:
                found_keywords['bearish'].append((keyword, weight))
        
        # Check for neutral keywords
        for keyword, weight in self.neutral_keywords.items():
            if keyword in text_lower:
                found_keywords['neutral'].append((keyword, weight))
        
        return found_keywords
    
    def apply_rules(self, text: str, model_probabilities: np.ndarray, 
                   confidence_threshold: float = 0.7) -> Tuple[np.ndarray, str]:
        """Apply rule-based post-processing to model predictions"""
        
        original_probs = model_probabilities.copy()
        adjusted_probs = model_probabilities.copy()
        
        # Extract keywords
        keywords = self.extract_keywords(text)
        
        # Calculate keyword scores
        bullish_score = sum(weight for _, weight in keywords['bullish'])
        bearish_score = sum(weight for _, weight in keywords['bearish']) 
        neutral_score = sum(weight for _, weight in keywords['neutral'])
        
        explanation_parts = []
        
        # Apply adjustments based on keyword scores
        if bullish_score > 0.5:
            # Boost bullish probability
            boost = min(0.3, bullish_score * 0.2)
            adjusted_probs[2] += boost  # Bullish
            adjusted_probs[0] = max(0.05, adjusted_probs[0] - boost/2)  # Bearish
            adjusted_probs[1] = max(0.05, adjusted_probs[1] - boost/2)  # Neutral
            explanation_parts.append(f"Bullish keywords detected (score: {bullish_score:.2f})")
        
        if bearish_score > 0.5:
            # Boost bearish probability
            boost = min(0.3, bearish_score * 0.2)
            adjusted_probs[0] += boost  # Bearish
            adjusted_probs[2] = max(0.05, adjusted_probs[2] - boost/2)  # Bullish
            adjusted_probs[1] = max(0.05, adjusted_probs[1] - boost/2)  # Neutral
            explanation_parts.append(f"Bearish keywords detected (score: {bearish_score:.2f})")
        
        if neutral_score > 0.5:
            # Boost neutral probability
            boost = min(0.2, neutral_score * 0.15)
            adjusted_probs[1] += boost  # Neutral
            adjusted_probs[0] = max(0.05, adjusted_probs[0] - boost/2)  # Bearish
            adjusted_probs[2] = max(0.05, adjusted_probs[2] - boost/2)  # Bullish
            explanation_parts.append(f"Neutral keywords detected (score: {neutral_score:.2f})")
        
        # Normalize probabilities
        adjusted_probs = adjusted_probs / np.sum(adjusted_probs)
        
        # Create explanation
        if explanation_parts:
            explanation = "Applied: " + ", ".join(explanation_parts)
        else:
            explanation = "No significant keywords detected"
        
        return adjusted_probs, explanation

# Initialize rule engine
rule_engine = SentimentRuleEngine()

class FinancialSentimentEnsemble:
    """Ensemble model for financial sentiment analysis using Hugging Face models"""
    
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.models = {}
        self.tokenizers = {}
        self.label_names = ["Bearish 📉", "Neutral ⚖️", "Bullish 📈"]
        
        # Hugging Face model configurations
        self.model_configs = {
            "distilbert": {
                "name": "DistilBERT (Fast)",
                "repo_id": "codealchemist01/financial-sentiment-distilbert",
                "description": "Fast and efficient model"
            },
            "bert_large": {
                "name": "BERT-Large (Advanced)", 
                "repo_id": "codealchemist01/financial-sentiment-bert-large",
                "description": "Most advanced model"
            },
            "improved": {
                "name": "Improved Model",
                "repo_id": "codealchemist01/financial-sentiment-improved", 
                "description": "Enhanced model with advanced training"
            }
        }
        
        # Ensemble weights for different combinations
        self.ensemble_weights = {
            "smart_ensemble": {"distilbert": 0.3, "bert_large": 0.7},
            "all_models": {"distilbert": 0.2, "improved": 0.3, "bert_large": 0.5}
        }
        
        self.load_models()
    
    def load_models(self):
        """Load models from Hugging Face Hub"""
        loaded_models = []
        
        for model_key, config in self.model_configs.items():
            try:
                logger.info(f"Loading {config['name']} from {config['repo_id']}")
                
                tokenizer = AutoTokenizer.from_pretrained(config["repo_id"])
                model = AutoModelForSequenceClassification.from_pretrained(config["repo_id"])
                model.to(self.device)
                model.eval()
                
                self.tokenizers[model_key] = tokenizer
                self.models[model_key] = model
                loaded_models.append(config["name"])
                
                logger.info(f"✅ {config['name']} loaded successfully")
                
            except Exception as e:
                logger.error(f"❌ Error loading {config['name']}: {e}")
        
        logger.info(f"🎯 Total loaded models: {len(loaded_models)}")
        return loaded_models
    
    def predict_single_model(self, text, model_key):
        """Get prediction from a single model"""
        if model_key not in self.models:
            return None, f"Model {model_key} not available"
        
        try:
            tokenizer = self.tokenizers[model_key]
            model = self.models[model_key]
            
            inputs = tokenizer(
                text,
                return_tensors="pt",
                truncation=True,
                padding=True,
                max_length=512
            )
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            with torch.no_grad():
                outputs = model(**inputs)
                probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
                probabilities = probabilities.cpu().numpy()[0]
            
            return probabilities, None
            
        except Exception as e:
            return None, f"Error in {model_key}: {str(e)}"
    
    def predict_ensemble(self, text, ensemble_type="smart_ensemble", use_rules=True):
        """Make ensemble prediction"""
        if not text.strip():
            return "Please enter some text to analyze.", {}, ""
        
        try:
            # Determine which models to use
            if ensemble_type == "smart_ensemble":
                weights = self.ensemble_weights["smart_ensemble"]
                models_to_use = ["distilbert", "bert_large"]
            elif ensemble_type == "all_models":
                weights = self.ensemble_weights["all_models"] 
                models_to_use = ["distilbert", "improved", "bert_large"]
            else:
                # Single model prediction
                models_to_use = [ensemble_type]
                weights = {ensemble_type: 1.0}
            
            # Get predictions from each model
            ensemble_probabilities = np.zeros(3)
            total_weight = 0
            model_predictions = {}
            model_details = []
            
            for model_key in models_to_use:
                if model_key in self.models:
                    probabilities, error = self.predict_single_model(text, model_key)
                    if probabilities is not None:
                        weight = weights.get(model_key, 1.0)
                        ensemble_probabilities += probabilities * weight
                        total_weight += weight
                        
                        # Store individual results
                        predicted_class = np.argmax(probabilities)
                        confidence = probabilities[predicted_class]
                        model_predictions[model_key] = {
                            "prediction": self.label_names[predicted_class],
                            "confidence": float(confidence),
                            "probabilities": probabilities.tolist()
                        }
                        
                        model_details.append(
                            f"**{self.model_configs[model_key]['name']}:** "
                            f"{self.label_names[predicted_class]} ({confidence:.2%})"
                        )
            
            if total_weight == 0:
                return "No models available for prediction.", {}, ""
            
            # Normalize ensemble probabilities
            ensemble_probabilities = ensemble_probabilities / total_weight
            
            # Apply rule-based post-processing if enabled
            rule_explanation = ""
            if use_rules:
                ensemble_probabilities, rule_explanation = rule_engine.apply_rules(
                    text, ensemble_probabilities, confidence_threshold=0.7
                )
            
            # Get final prediction
            predicted_class = np.argmax(ensemble_probabilities)
            confidence = ensemble_probabilities[predicted_class]
            
            # Create detailed results
            if len(models_to_use) > 1:
                result_text = f"**🎯 Ensemble Prediction:** {self.label_names[predicted_class]}\\n"
                result_text += f"**🔥 Ensemble Confidence:** {confidence:.2%}\\n\\n"
                
                result_text += "**🤖 Individual Model Results:**\\n"
                for detail in model_details:
                    result_text += f"- {detail}\\n"
                result_text += "\\n"
            else:
                result_text = f"**🎯 Prediction:** {self.label_names[predicted_class]}\\n"
                result_text += f"**🔥 Confidence:** {confidence:.2%}\\n\\n"
            
            # Show rule engine effects if applied
            if use_rules and rule_explanation:
                result_text += f"**🤖 Rule Engine:** {rule_explanation}\\n\\n"
            
            result_text += "**📊 Final Probabilities:**\\n"
            
            # Create probability dictionary for gradio
            prob_dict = {}
            for i, (label, prob) in enumerate(zip(self.label_names, ensemble_probabilities)):
                prob_dict[label] = float(prob)
                result_text += f"- {label}: {prob:.2%}\\n"
            
            # Create model comparison details
            comparison_details = ""
            if len(model_predictions) > 1:
                comparison_details = "**🔍 Model Comparison:**\\n"
                for model_key, pred_data in model_predictions.items():
                    comparison_details += f"\\n**{self.model_configs[model_key]['name']}:**\\n"
                    for i, (label, prob) in enumerate(zip(self.label_names, pred_data['probabilities'])):
                        comparison_details += f"  - {label}: {prob:.2%}\\n"
            
            return result_text, prob_dict, comparison_details
            
        except Exception as e:
            logger.error(f"Prediction error: {e}")
            return f"Error during prediction: {str(e)}", {}, ""

# Initialize ensemble model
try:
    ensemble = FinancialSentimentEnsemble()
    available_models = list(ensemble.models.keys())
    gpu_info = f"🚀 **Models loaded:** {len(available_models)} models on {ensemble.device}"
except Exception as e:
    gpu_info = f"❌ **Error loading models:** {str(e)}"
    ensemble = None
    available_models = []

def analyze_sentiment(text, model_selection, use_rules):
    """Main analysis function"""
    if ensemble is None:
        return "Models not loaded. Please check the error above.", {}, ""
    
    return ensemble.predict_ensemble(text, model_selection, use_rules)

# Example texts for testing
examples = [
    ["Tesla stock is soaring after excellent Q3 earnings report! 🚀", "smart_ensemble", True],
    ["The market is showing mixed signals today, uncertain direction.", "smart_ensemble", True], 
    ["Major selloff expected as inflation concerns grow. Bearish outlook.", "all_models", True],
    ["Apple announces new iPhone with revolutionary features!", "distilbert", False],
    ["Economic indicators suggest potential recession ahead.", "bert_large", True],
    ["Crypto market rebounds strongly after recent dip.", "smart_ensemble", True]
]

# Create Gradio interface
with gr.Blocks(
    title="Financial Sentiment Analysis - Ensemble System",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        max-width: 1000px !important;
        margin: auto !important;
    }
    .header {
        text-align: center;
        padding: 20px;
        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
        color: white;
        border-radius: 10px;
        margin-bottom: 20px;
    }
    .model-info {
        background-color: #f8f9fa;
        padding: 15px;
        border-radius: 8px;
        margin: 10px 0;
    }
    """
) as demo:
    
    gr.HTML(f"""
    <div class="header">
        <h1>📈 Financial Sentiment Analysis Ensemble</h1>
        <h3>Advanced AI-powered sentiment analysis for financial texts using an ensemble of 3 fine-tuned models</h3>
        <p>{gpu_info}</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(
                label="📝 Enter Financial Text to Analyze",
                placeholder="Enter financial news, tweets, or market commentary...",
                lines=4
            )
            
            with gr.Row():
                model_selection = gr.Dropdown(
                    choices=[
                        ("🧠 Smart Ensemble (Recommended)", "smart_ensemble"),
                        ("🎯 All Models Ensemble", "all_models"),
                        ("⚡ DistilBERT (Fast)", "distilbert"),
                        ("🔥 BERT-Large (Advanced)", "bert_large"),
                        ("🚀 Improved Model", "improved")
                    ],
                    value="smart_ensemble",
                    label="🤖 Model Selection"
                )
                
                use_rules = gr.Checkbox(
                    label="🤖 Rule-Based Enhancement",
                    value=True,
                    info="Apply keyword-based post-processing"
                )
            
            analyze_btn = gr.Button("🔍 Analyze Sentiment", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            result_output = gr.Textbox(
                label="📊 Analysis Results",
                lines=12,
                interactive=False
            )
            
            prob_output = gr.Label(
                label="📈 Probability Distribution",
                num_top_classes=3
            )
    
    with gr.Row():
        comparison_output = gr.Textbox(
            label="🔍 Model Comparison Details",
            lines=8,
            interactive=False,
            visible=True
        )
    
    # Event handlers
    analyze_btn.click(
        fn=analyze_sentiment,
        inputs=[text_input, model_selection, use_rules],
        outputs=[result_output, prob_output, comparison_output]
    )
    
    # Examples section
    gr.Examples(
        examples=examples,
        inputs=[text_input, model_selection, use_rules],
        outputs=[result_output, prob_output, comparison_output],
        fn=analyze_sentiment,
        cache_examples=False,
        label="💡 Try these examples:"
    )
    
    # Model information
    gr.HTML("""
    <div class="model-info">
        <h4>🤖 Ensemble System Information</h4>
        <ul>
            <li><strong>🧠 Smart Ensemble:</strong> DistilBERT + BERT-Large (Best balance of speed and accuracy)</li>
            <li><strong>🎯 All Models:</strong> DistilBERT + Improved + BERT-Large (Maximum consensus)</li>
            <li><strong>⚡ DistilBERT:</strong> Fast and efficient model optimized for real-time analysis</li>
            <li><strong>🔥 BERT-Large:</strong> Most advanced model with deep contextual understanding</li>
            <li><strong>🚀 Improved Model:</strong> Enhanced with advanced training techniques</li>
        </ul>
        <p><em>💡 Tip: Smart Ensemble provides the best balance of accuracy and performance!</em></p>
        <p><em>🤖 Rule Engine: Applies keyword-based post-processing to improve accuracy on financial texts</em></p>
    </div>
    """)

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