#!/usr/bin/env python3
"""
Financial Sentiment Analysis - Enhanced Ensemble Gradio Demo for Hugging Face Space
Yerel uygulamayla tam uyumlu versiyon
"""

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
        }
    
    def extract_keywords(self, text: str) -> Dict[str, float]:
        """Extract and score keywords from text"""
        text_lower = text.lower()
        found_keywords = {}
        
        # Check bullish keywords
        for keyword, weight in self.bullish_keywords.items():
            if keyword in text_lower:
                found_keywords[keyword] = weight
        
        # Check bearish keywords  
        for keyword, weight in self.bearish_keywords.items():
            if keyword in text_lower:
                found_keywords[keyword] = -weight  # Negative for bearish
                
        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"""
        
        keywords = self.extract_keywords(text)
        if not keywords:
            return model_probabilities, "No significant keywords found"
        
        # Calculate keyword score
        keyword_score = sum(keywords.values())
        
        # Get model's confidence
        max_prob = np.max(model_probabilities)
        
        # Apply rules only if model confidence is low
        if max_prob < confidence_threshold:
            adjustment_strength = 0.3  # How much to adjust
            
            if keyword_score > 0.5:  # Strong bullish keywords
                # Boost bullish probability
                model_probabilities[2] += adjustment_strength
                model_probabilities[0] -= adjustment_strength * 0.5
                model_probabilities[1] -= adjustment_strength * 0.5
                rule_msg = f"Bullish keywords detected (score: {keyword_score:.2f}), boosting bullish probability"
                
            elif keyword_score < -0.5:  # Strong bearish keywords
                # Boost bearish probability
                model_probabilities[0] += adjustment_strength
                model_probabilities[1] -= adjustment_strength * 0.5
                model_probabilities[2] -= adjustment_strength * 0.5
                rule_msg = f"Bearish keywords detected (score: {keyword_score:.2f}), boosting bearish probability"
            else:
                rule_msg = f"Mixed signals (score: {keyword_score:.2f}), no adjustment"
        else:
            rule_msg = f"High model confidence ({max_prob:.2%}), rules not applied"
        
        # Normalize probabilities
        model_probabilities = np.maximum(model_probabilities, 0)
        model_probabilities = model_probabilities / np.sum(model_probabilities)
        
        return model_probabilities, rule_msg

# Initialize rule engine
rule_engine = SentimentRuleEngine()

class EnsembleFinancialPredictor:
    """Yerel uygulamayla tam uyumlu ensemble predictor"""
    
    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 📈"]
        
        # Yerel uygulamayla aynı model isimleri ve yapılandırması
        self.model_info = {
            "distilbert": {
                "name": "DistilBERT (Fast)",
                "repo_id": "codealchemist01/financial-sentiment-distilbert",
                "description": "Hızlı ve verimli model (87.96% doğruluk)"
            },
            "balanced": {
                "name": "Balanced Model",
                "repo_id": "codealchemist01/financial-sentiment-improved",  # Improved model as balanced
                "description": "Dengeli performans modeli"
            },
            "advanced": {  # YERELDEKİ GİBİ "advanced" ismi
                "name": "BERT-Large (Advanced)",
                "repo_id": "codealchemist01/financial-sentiment-bert-large",
                "description": "En gelişmiş model (85.85% doğruluk)"
            }
        }
        
        # Yerel uygulamayla AYNI ensemble weights
        self.ensemble_weights = {
            "smart_ensemble": {"distilbert": 0.3, "advanced": 0.7},  # ADVANCED ismi!
            "all_models": {"distilbert": 0.2, "balanced": 0.3, "advanced": 0.5}
        }
        
        self.load_models()
    
    def load_models(self):
        """Load all available models from Hugging Face Hub"""
        loaded_models = []
        
        for model_key, model_info in self.model_info.items():
            try:
                logger.info(f"Loading {model_info['name']} from {model_info['repo_id']}")
                
                tokenizer = AutoTokenizer.from_pretrained(model_info["repo_id"])
                model = AutoModelForSequenceClassification.from_pretrained(model_info["repo_id"])
                model.to(self.device)
                model.eval()
                
                self.tokenizers[model_key] = tokenizer
                self.models[model_key] = model
                loaded_models.append(model_info["name"])
                
                logger.info(f"✅ {model_info['name']} loaded successfully")
                    
            except Exception as e:
                logger.error(f"❌ Error loading {model_info['name']}: {e}")
        
        logger.info(f"🎯 Total loaded models: {len(loaded_models)}")
        return loaded_models
    
    def predict_single_model(self, text, model_key):
        """Predict with 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):
        """Predict using ensemble approach - yerel uygulamayla aynı mantık"""
        if not text.strip():
            return "Please enter some text to analyze.", {}, ""
        
        try:
            # Get predictions from all models
            model_predictions = {}
            model_details = []
            
            if ensemble_type == "smart_ensemble":
                # Use best performing combination: DistilBERT + BERT-large (ADVANCED)
                weights = self.ensemble_weights["smart_ensemble"]
                models_to_use = ["distilbert", "advanced"]  # ADVANCED ismi!
            elif ensemble_type == "all_models":
                # Use all three models
                weights = self.ensemble_weights["all_models"]
                models_to_use = ["distilbert", "balanced", "advanced"]
            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
            
            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 model 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_info[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
            
            # Store original probabilities
            original_probabilities = ensemble_probabilities.copy()
            
            # 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_info[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 predictor
try:
    predictor = EnsembleFinancialPredictor()
    available_models = list(predictor.models.keys())
    gpu_info = f"🚀 **Models loaded:** {len(available_models)} models on {predictor.device}"
except Exception as e:
    gpu_info = f"❌ **Error loading models:** {str(e)}"
    predictor = None
    available_models = []

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

# Example texts - yerel uygulamayla aynı
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.", "advanced", True],  # ADVANCED ismi!
    ["Crypto market rebounds strongly after recent dip.", "smart_ensemble", True]
]

# Create Gradio interface - yerel uygulamayla aynı stil
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 System</h1>
        <h3>🚀 3-Model Ensemble - Advanced AI Analysis</h3>
        <p>{gpu_info}</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(
                label="📝 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"),
                        ("⚖️ Balanced Model", "balanced"),
                        ("🔥 BERT-Large (Advanced)", "advanced")  # ADVANCED ismi!
                    ],
                    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
    gr.Examples(
        examples=examples,
        inputs=[text_input, model_selection, use_rules],
        outputs=[result_output, prob_output, comparison_output],
        fn=analyze_sentiment,
        cache_examples=False
    )
    
    # Model information - yerel uygulamayla aynı
    gr.HTML("""
    <div class="model-info">
        <h4>🤖 Ensemble System Information</h4>
        <ul>
            <li><strong>🧠 Smart Ensemble:</strong> DistilBERT + BERT-Large (79.7% average accuracy)</li>
            <li><strong>🎯 All Models:</strong> DistilBERT + Balanced + BERT-Large (79.1% average accuracy)</li>
            <li><strong>⚡ DistilBERT:</strong> Fast and efficient (87.96% accuracy)</li>
            <li><strong>⚖️ Balanced Model:</strong> Optimized for balanced performance</li>
            <li><strong>🔥 BERT-Large:</strong> Most advanced model (85.85% accuracy)</li>
        </ul>
        <p><em>💡 Tip: Smart Ensemble provides the best balance of accuracy and performance!</em></p>
        <p><em>🤖 Rule Engine: Keyword-based post-processing improves accuracy on financial texts</em></p>
    </div>
    """)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )