import gradio as gr
import yfinance as yf
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import warnings
import time
import gc
import os
import torch
from datetime import datetime, timedelta
from typing import Optional, Dict, Any, Tuple

warnings.filterwarnings('ignore')

# Environment optimizations for Hugging Face Spaces
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
os.environ['HF_HUB_DISABLE_PROGRESS_BARS'] = '1'
os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
os.environ['TRANSFORMERS_VERBOSITY'] = 'error'

if torch.cuda.is_available():
    torch.cuda.empty_cache()
torch.set_num_threads(min(4, os.cpu_count() or 1))

class FastAIStockAnalyzer:
    """Optimized AI Stock Analyzer for Gradio with robust error handling"""
    
    def __init__(self):
        self.context_length = 32
        self.prediction_length = 7
        self.device = "cpu"
        self.model_cache = {}
        
    def fetch_stock_data(self, symbol: str, period: str = "6mo") -> Tuple[Optional[pd.DataFrame], Optional[Dict]]:
        """Fetch stock data with error handling"""
        try:
            ticker = yf.Ticker(symbol)
            data = ticker.history(period=period, interval="1d", 
                                actions=False, auto_adjust=True, 
                                back_adjust=False, repair=False)
            
            if data.empty:
                return None, None
                
            try:
                info = {
                    'longName': ticker.info.get('longName', symbol),
                    'sector': ticker.info.get('sector', 'Unknown'),
                    'marketCap': ticker.info.get('marketCap', 0)
                }
            except:
                info = {'longName': symbol, 'sector': 'Unknown', 'marketCap': 0}
            
            return data, info
            
        except Exception as e:
            return None, None
    
    def load_chronos_tiny(self) -> Tuple[Optional[Any], str]:
        """Load Chronos model with caching and fallback"""
        model_key = "chronos_tiny"
        
        if model_key in self.model_cache:
            return self.model_cache[model_key], "chronos"
        
        try:
            from chronos import ChronosPipeline
            
            # Try primary loading method
            pipeline = ChronosPipeline.from_pretrained(
                "amazon/chronos-t5-tiny",
                device_map="cpu",
                torch_dtype=torch.float32,
                low_cpu_mem_usage=True,
                trust_remote_code=True
            )
            
            self.model_cache[model_key] = pipeline
            return pipeline, "chronos"
            
        except ImportError:
            # Chronos not available
            return None, None
        except Exception as e:
            # Try fallback loading method
            try:
                pipeline = ChronosPipeline.from_pretrained(
                    "amazon/chronos-t5-tiny",
                    device_map="auto",
                    torch_dtype=torch.float32
                )
                self.model_cache[model_key] = pipeline
                return pipeline, "chronos"
            except:
                return None, None
    
    def load_moirai_small(self) -> Tuple[Optional[Any], str]:
        """Load Moirai model with updated method and fallbacks"""
        model_key = "moirai_small"
        
        if model_key in self.model_cache:
            return self.model_cache[model_key], "moirai"
        
        try:
            from uni2ts.model.moirai import MoiraiForecast, MoiraiModule
            
            # Method 1: Try the standard approach
            try:
                module = MoiraiModule.from_pretrained(
                    "Salesforce/moirai-1.0-R-small",
                    device_map="cpu",
                    torch_dtype=torch.float32,
                    trust_remote_code=True,
                    low_cpu_mem_usage=True
                )
                
                model = MoiraiForecast(
                    module=module,
                    prediction_length=self.prediction_length,
                    context_length=self.context_length,
                    patch_size="auto",
                    num_samples=10,
                    target_dim=1,
                    feat_dynamic_real_dim=0,
                    past_feat_dynamic_real_dim=0
                )
                
                self.model_cache[model_key] = model
                return model, "moirai"
                
            except Exception as e1:
                # Method 2: Try newer version
                try:
                    module = MoiraiModule.from_pretrained(
                        "Salesforce/moirai-1.1-R-small",
                        device_map="cpu",
                        torch_dtype=torch.float32,
                        trust_remote_code=True
                    )
                    
                    model = MoiraiForecast(
                        module=module,
                        prediction_length=self.prediction_length,
                        context_length=self.context_length,
                        patch_size="auto",
                        num_samples=5,  # Reduced for stability
                        target_dim=1,
                        feat_dynamic_real_dim=0,
                        past_feat_dynamic_real_dim=0
                    )
                    
                    self.model_cache[model_key] = model
                    return model, "moirai"
                    
                except Exception as e2:
                    # Method 3: Minimal configuration
                    try:
                        module = MoiraiModule.from_pretrained("Salesforce/moirai-1.0-R-small")
                        model = MoiraiForecast(
                            module=module,
                            prediction_length=7,
                            context_length=32,
                            patch_size="auto",
                            num_samples=5,
                            target_dim=1
                        )
                        
                        self.model_cache[model_key] = model
                        return model, "moirai"
                        
                    except Exception as e3:
                        return None, None
                        
        except ImportError:
            # uni2ts not available
            return None, None
        except Exception as e:
            return None, None
    
    def predict_chronos_fast(self, pipeline: Any, data: np.ndarray) -> Optional[Dict]:
        """Fast Chronos prediction with error handling"""
        try:
            context_data = data[-self.context_length:]
            context = torch.tensor(context_data, dtype=torch.float32).unsqueeze(0)
            
            with torch.no_grad():
                forecast = pipeline.predict(
                    context=context,
                    prediction_length=self.prediction_length,
                    num_samples=10,
                    temperature=1.0,
                    top_k=50,
                    top_p=1.0
                )
            
            forecast_array = forecast[0].numpy()
            predictions = {
                'mean': np.median(forecast_array, axis=0),
                'q10': np.quantile(forecast_array, 0.1, axis=0),
                'q90': np.quantile(forecast_array, 0.9, axis=0),
                'std': np.std(forecast_array, axis=0)
            }
            
            return predictions
            
        except Exception as e:
            return None
    
    def predict_moirai_fast(self, model: Any, data: np.ndarray) -> Optional[Dict]:
        """Fast Moirai prediction with enhanced error handling"""
        try:
            from gluonts.dataset.common import ListDataset
            
            # Prepare dataset with minimal configuration
            dataset = ListDataset([{
                "item_id": "stock",
                "start": "2023-01-01",
                "target": data[-self.context_length:].tolist()
            }], freq='D')
            
            # Create predictor with safe parameters
            predictor = model.create_predictor(
                batch_size=1,
                num_parallel_samples=5  # Further reduced for stability
            )
            
            # Generate forecast
            forecasts = list(predictor.predict(dataset))
            
            if not forecasts:
                return None
                
            forecast = forecasts[0]
            
            predictions = {
                'mean': forecast.mean,
                'q10': forecast.quantile(0.1),
                'q90': forecast.quantile(0.9),
                'std': np.std(forecast.samples, axis=0) if hasattr(forecast, 'samples') else np.zeros(7)
            }
            
            return predictions
            
        except Exception as e:
            return None
    
    def generate_simple_prediction(self, data: np.ndarray) -> Dict:
        """Fallback prediction method using simple statistical models"""
        try:
            # Simple moving average with trend
            recent_data = data[-30:]  # Last 30 days
            short_ma = np.mean(recent_data[-7:])   # 7-day average
            long_ma = np.mean(recent_data[-21:])   # 21-day average
            
            # Calculate trend
            trend = (short_ma - long_ma) / long_ma if long_ma != 0 else 0
            
            # Generate predictions
            current_price = data[-1]
            predictions = []
            
            for i in range(7):
                # Simple trend projection with some noise
                predicted_price = current_price * (1 + trend * (i + 1) * 0.1)
                predictions.append(predicted_price)
            
            predictions = np.array(predictions)
            
            return {
                'mean': predictions,
                'q10': predictions * 0.95,  # 5% lower
                'q90': predictions * 1.05,  # 5% higher
                'std': np.full(7, np.std(recent_data) * 0.5)
            }
            
        except Exception:
            # Ultimate fallback - flat prediction
            current_price = data[-1]
            return {
                'mean': np.full(7, current_price),
                'q10': np.full(7, current_price * 0.98),
                'q90': np.full(7, current_price * 1.02),
                'std': np.full(7, 0.01)
            }

# Initialize analyzer globally for caching
analyzer = FastAIStockAnalyzer()

def analyze_stock(stock_symbol, model_choice, investment_amount, progress=gr.Progress()):
    """Main analysis function with comprehensive error handling and fallbacks"""
    
    try:
        progress(0.1, desc="Fetching stock data...")
        
        # Validate input
        if not stock_symbol or stock_symbol.strip() == "":
            return (
                "❌ Error: Please enter a valid stock symbol.",
                None,
                "❌ Invalid Input",
                "N/A",
                "N/A"
            )
        
        # Fetch data
        stock_data, stock_info = analyzer.fetch_stock_data(stock_symbol.upper())
        
        if stock_data is None or len(stock_data) < 50:
            return (
                f"❌ Error: Insufficient data for {stock_symbol.upper()}. Please check the stock symbol or try a different one.",
                None,
                "❌ Data Error",
                "N/A",
                "N/A"
            )
        
        current_price = stock_data['Close'].iloc[-1]
        company_name = stock_info.get('longName', stock_symbol) if stock_info else stock_symbol
        
        progress(0.3, desc="Loading AI model...")
        
        # Determine model type and load
        model_type = "chronos" if "Chronos" in model_choice else "moirai"
        model = None
        model_name = ""
        prediction_method = None
        
        if model_type == "chronos":
            model, loaded_type = analyzer.load_chronos_tiny()
            model_name = "Amazon Chronos Tiny"
            prediction_method = "chronos"
        else:
            model, loaded_type = analyzer.load_moirai_small()
            model_name = "Salesforce Moirai Small"
            prediction_method = "moirai"
            
            # Fallback to Chronos if Moirai fails
            if model is None:
                progress(0.4, desc="Moirai unavailable, switching to Chronos...")
                model, loaded_type = analyzer.load_chronos_tiny()
                model_name = "Amazon Chronos Tiny (Fallback)"
                prediction_method = "chronos"
        
        # If both models fail, use simple prediction
        if model is None:
            progress(0.5, desc="Using statistical fallback method...")
            model_name = "Statistical Trend Model (Fallback)"
            prediction_method = "simple"
        
        progress(0.6, desc="Generating predictions...")
        
        # Generate predictions based on available method
        predictions = None
        
        if prediction_method == "chronos" and model is not None:
            predictions = analyzer.predict_chronos_fast(model, stock_data['Close'].values)
        elif prediction_method == "moirai" and model is not None:
            predictions = analyzer.predict_moirai_fast(model, stock_data['Close'].values)
        
        # Use simple prediction if AI models fail
        if predictions is None:
            predictions = analyzer.generate_simple_prediction(stock_data['Close'].values)
            model_name = "Statistical Trend Model (AI Models Unavailable)"
        
        progress(0.8, desc="Calculating investment scenarios...")
        
        # Analysis results
        mean_pred = predictions['mean']
        final_pred = mean_pred[-1]
        week_change = ((final_pred - current_price) / current_price) * 100
        
        # Decision logic
        if week_change > 5:
            decision = "🟢 STRONG BUY"
            explanation = "Model expects significant gains!"
        elif week_change > 2:
            decision = "🟢 BUY"
            explanation = "Model expects moderate gains"
        elif week_change < -5:
            decision = "🔴 STRONG SELL"
            explanation = "Model expects significant losses"
        elif week_change < -2:
            decision = "🔴 SELL"
            explanation = "Model expects losses"
        else:
            decision = "⚪ HOLD"
            explanation = "Model expects stable prices"
        
        # Create analysis text
        analysis_text = f"""
# 🎯 {company_name} ({stock_symbol.upper()}) Analysis

## 🤖 RECOMMENDATION: {decision}
**{explanation}**
*Powered by {model_name}*

## 📊 Key Metrics
- **Current Price**: ${current_price:.2f}
- **7-Day Prediction**: ${final_pred:.2f} ({week_change:+.2f}%)
- **Confidence Level**: {min(100, max(50, 70 + abs(week_change) * 1.5)):.0f}%
- **Analysis Method**: {model_name}

## 💰 Investment Scenario (${investment_amount:,.0f})
- **Shares**: {investment_amount/current_price:.2f}
- **Current Value**: ${investment_amount:,.2f}
- **Predicted Value**: ${investment_amount + ((final_pred - current_price) * (investment_amount/current_price)):,.2f}
- **Profit/Loss**: ${((final_pred - current_price) * (investment_amount/current_price)):+,.2f} ({week_change:+.2f}%)

## ⚠️ Important Disclaimers
- **This is for educational purposes only**
- **Not financial advice - consult professionals**
- **AI predictions can be wrong - invest responsibly**
- **Past performance ≠ future results**
"""
        
        progress(0.9, desc="Creating visualizations...")
        
        # Create chart
        fig = go.Figure()
        
        # Historical data (last 30 days)
        recent = stock_data.tail(30)
        fig.add_trace(go.Scatter(
            x=recent.index, 
            y=recent['Close'],
            mode='lines', 
            name='Historical Price',
            line=dict(color='blue', width=2)
        ))
        
        # Predictions
        future_dates = pd.date_range(
            start=stock_data.index[-1] + pd.Timedelta(days=1),
            periods=7, 
            freq='D'
        )
        
        fig.add_trace(go.Scatter(
            x=future_dates, 
            y=mean_pred,
            mode='lines+markers', 
            name='Prediction',
            line=dict(color='red', width=3),
            marker=dict(size=8)
        ))
        
        # Confidence bands
        if 'q10' in predictions and 'q90' in predictions:
            fig.add_trace(go.Scatter(
                x=future_dates.tolist() + future_dates[::-1].tolist(),
                y=predictions['q90'].tolist() + predictions['q10'][::-1].tolist(),
                fill='toself', 
                fillcolor='rgba(255,0,0,0.1)',
                line=dict(color='rgba(255,255,255,0)'),
                name='Confidence Range', 
                showlegend=True
            ))
        
        fig.update_layout(
            title=f"{stock_symbol.upper()} - Stock Forecast ({model_name})",
            xaxis_title="Date", 
            yaxis_title="Price ($)",
            height=500, 
            showlegend=True,
            template="plotly_white"
        )
        
        progress(1.0, desc="Analysis complete!")
        
        # Create summary metrics
        try:
            day_change = stock_data['Close'].iloc[-1] - stock_data['Close'].iloc[-2]
            day_change_pct = (day_change / stock_data['Close'].iloc[-2]) * 100
        except:
            day_change_pct = 0
        
        current_metrics = f"${current_price:.2f} ({day_change_pct:+.2f}%)"
        prediction_metrics = f"${final_pred:.2f} ({week_change:+.2f}%)"
        
        return (
            analysis_text,
            fig,
            decision,
            current_metrics,
            prediction_metrics
        )
        
    except Exception as e:
        # Ultimate error handler
        error_msg = f"""
# ❌ Analysis Error

**Something went wrong during the analysis:**

- **Error**: {str(e)[:200]}...
- **Stock**: {stock_symbol}
- **Model**: {model_choice}

## 🔧 Try these solutions:
1. **Check stock symbol** - Make sure it's valid (e.g., AAPL, GOOGL)
2. **Try different model** - Switch between Chronos and Moirai
3. **Refresh and try again** - Temporary server issues
4. **Use popular stocks** - AAPL, MSFT, GOOGL work best

## 📞 Still having issues?
This may be due to Hugging Face Spaces resource limitations or model availability.
"""
        
        return (
            error_msg,
            None,
            "❌ Error",
            "N/A",
            "N/A"
        )

# Create Gradio interface with enhanced UI
with gr.Blocks(
    theme=gr.themes.Soft(),
    title="⚡ Fast AI Stock Predictor",
    css="""
    footer {visibility: hidden}
    .gradio-container {max-width: 1200px; margin: auto;}
    .main-header {text-align: center; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 10px; margin-bottom: 20px;}
    .disclaimer {background-color: #fff3cd; border: 1px solid #ffeaa7; padding: 15px; border-radius: 8px; margin: 10px 0;}
    """
) as demo:
    
    gr.HTML("""
    <div class="main-header">
        <h1 style="margin: 0; font-size: 2.5em;">⚡ AI Stock Predictor</h1>
        <p style="margin: 10px 0 0 0; font-size: 1.2em;"><strong>🤖 Powered by Amazon Chronos & Salesforce Moirai</strong></p>
        <p style="margin: 5px 0 0 0; opacity: 0.9;">Advanced AI models for stock price forecasting</p>
    </div>
    """)
    
    gr.HTML("""
    <div class="disclaimer">
        <strong>⚠️ IMPORTANT DISCLAIMER:</strong> This tool is for educational purposes only. 
        Not financial advice. AI predictions can be wrong. Always consult financial professionals 
        before making investment decisions. Only invest what you can afford to lose.
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1, min_width=300):
            gr.HTML("<h3>🎯 Analysis Configuration</h3>")
            
            stock_input = gr.Dropdown(
                choices=["AAPL", "GOOGL", "MSFT", "TSLA", "AMZN", "META", "NFLX", "NVDA", "ORCL", "CRM"],
                value="AAPL",
                label="📈 Select Stock Symbol",
                allow_custom_value=True,
                info="Choose popular stocks or enter any valid symbol"
            )
            
            model_input = gr.Radio(
                choices=["🚀 Chronos (Fast & Reliable)", "🎯 Moirai (Advanced)"],
                value="🚀 Chronos (Fast & Reliable)",
                label="🤖 AI Model Selection",
                info="Chronos: Faster, more stable | Moirai: More sophisticated (may fallback to Chronos)"
            )
            
            investment_input = gr.Slider(
                minimum=500,
                maximum=100000,
                value=5000,
                step=500,
                label="💰 Investment Amount ($)",
                info="Amount for profit/loss calculation"
            )
            
            analyze_btn = gr.Button(
                "🚀 Analyze Stock Now", 
                variant="primary", 
                size="lg",
                scale=1
            )
            
            gr.HTML("<br>")
            
            # Quick stats
            with gr.Group():
                gr.HTML("<h4>📊 Quick Metrics</h4>")
                current_price_display = gr.Textbox(
                    label="Current Price", 
                    interactive=False,
                    container=True
                )
                prediction_display = gr.Textbox(
                    label="7-Day Prediction", 
                    interactive=False,
                    container=True
                )
                decision_display = gr.Textbox(
                    label="AI Recommendation", 
                    interactive=False,
                    container=True
                )
            
        with gr.Column(scale=2, min_width=600):
            gr.HTML("<h3>📊 Analysis Results</h3>")
            
            analysis_output = gr.Markdown(
                value="""
# 👋 Welcome to AI Stock Predictor!

**Ready to analyze stocks with cutting-edge AI?**

🎯 **How to use:**
1. Select a stock symbol (or enter your own)
2. Choose AI model (Chronos recommended for beginners)
3. Set investment amount for scenario analysis
4. Click "Analyze Stock Now"

💡 **Tips:**
- Try popular stocks like AAPL, GOOGL, MSFT first
- Chronos model is faster and more reliable
- Analysis takes 30-60 seconds for first-time model loading

⚡ **Click the button to get started!**
""",
                container=True
            )
    
    with gr.Row():
        chart_output = gr.Plot(
            label="📈 Stock Price Chart & AI Predictions",
            container=True,
            show_label=True
        )
    
    # Event handlers
    analyze_btn.click(
        fn=analyze_stock,
        inputs=[stock_input, model_input, investment_input],
        outputs=[
            analysis_output, 
            chart_output, 
            decision_display,
            current_price_display,
            prediction_display
        ],
        show_progress=True
    )
    
    # Examples section
    gr.HTML("<h3>🎭 Try These Examples</h3>")
    gr.Examples(
        examples=[
            ["AAPL", "🚀 Chronos (Fast & Reliable)", 5000],
            ["TSLA", "🎯 Moirai (Advanced)", 10000],
            ["GOOGL", "🚀 Chronos (Fast & Reliable)", 2500],
            ["MSFT", "🎯 Moirai (Advanced)", 7500],
            ["NVDA", "🚀 Chronos (Fast & Reliable)", 3000],
        ],
        inputs=[stock_input, model_input, investment_input],
        label="Click any example to load it:",
        examples_per_page=5
    )
    
    # Footer
    gr.HTML("""
    <div style="text-align: center; padding: 20px; margin-top: 30px; border-top: 1px solid #eee;">
        <p><strong>🤖 AI Stock Predictor</strong> | Built with ❤️ using Gradio & Hugging Face</p>
        <p style="font-size: 12px; color: #666;">
            Powered by Amazon Chronos & Salesforce Moirai | 
            Educational Tool - Not Financial Advice
        </p>
    </div>
    """)

# Launch configuration
if __name__ == "__main__":
    # Enable queue using the new method
    #demo.queue(max_size=20)  # Optional: set max queue size
    
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        max_threads=10
    )