import gradio as gr
import yfinance as yf
import pandas as pd
import matplotlib.pyplot as plt
import requests
import json
import numpy as np
from datetime import datetime, timedelta

# --- 1. Lightweight Market Data Functions ---
def get_market_data_safely(symbol):
    """Safely fetch market data with better error handling"""
    yf_sym = symbol.upper()
    if not yf_sym.startswith("^"):
        yf_sym = f"{yf_sym}.NS"
    
    try:
        # Fetch data with minimal memory usage
        df = yf.download(
            yf_sym, 
            period="1mo", 
            interval="1d", 
            progress=False,
            show_errors=False,
            threads=False
        )
        
        if df.empty:
            return None, f"No data found for symbol '{symbol}'"
        
        # Reset index to avoid formatting issues
        df = df.reset_index()
        
        return df, None
        
    except Exception as e:
        return None, f"Error fetching data: {str(e)}"

def calculate_technical_indicators(df):
    """Calculate basic technical indicators"""
    try:
        # Moving averages
        df['MA5'] = df['Close'].rolling(window=5).mean()
        df['MA20'] = df['Close'].rolling(window=20).mean()
        df['MA50'] = df['Close'].rolling(window=50).mean()
        
        # RSI calculation
        delta = df['Close'].diff()
        gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
        rs = gain / loss
        df['RSI'] = 100 - (100 / (1 + rs))
        
        # Bollinger Bands
        df['BB_Middle'] = df['Close'].rolling(window=20).mean()
        bb_std = df['Close'].rolling(window=20).std()
        df['BB_Upper'] = df['BB_Middle'] + (bb_std * 2)
        df['BB_Lower'] = df['BB_Middle'] - (bb_std * 2)
        
        return df
    except Exception as e:
        print(f"Technical indicator error: {e}")
        return df

def get_market_sentiment(symbol, price_data):
    """Generate market sentiment based on technical analysis"""
    try:
        latest_price = float(price_data['Close'].iloc[-1])
        prev_price = float(price_data['Close'].iloc[-2]) if len(price_data) > 1 else latest_price
        change_pct = ((latest_price - prev_price) / prev_price) * 100 if prev_price != 0 else 0
        
        # RSI analysis
        rsi = price_data['RSI'].iloc[-1] if 'RSI' in price_data.columns and not pd.isna(price_data['RSI'].iloc[-1]) else 50
        
        # Moving average analysis
        ma20 = price_data['MA20'].iloc[-1] if 'MA20' in price_data.columns and not pd.isna(price_data['MA20'].iloc[-1]) else latest_price
        ma50 = price_data['MA50'].iloc[-1] if 'MA50' in price_data.columns and not pd.isna(price_data['MA50'].iloc[-1]) else latest_price
        
        # Sentiment scoring
        sentiment_score = 0
        signals = []
        
        # Price momentum
        if change_pct > 2:
            sentiment_score += 2
            signals.append("Strong positive momentum")
        elif change_pct > 0:
            sentiment_score += 1
            signals.append("Positive momentum")
        elif change_pct < -2:
            sentiment_score -= 2
            signals.append("Strong negative momentum")
        elif change_pct < 0:
            sentiment_score -= 1
            signals.append("Negative momentum")
        
        # RSI analysis
        if rsi > 70:
            sentiment_score -= 1
            signals.append("Overbought (RSI > 70)")
        elif rsi < 30:
            sentiment_score += 1
            signals.append("Oversold (RSI < 30)")
        elif 40 <= rsi <= 60:
            signals.append("Neutral RSI")
        
        # Moving average analysis
        if latest_price > ma20 > ma50:
            sentiment_score += 2
            signals.append("Bullish trend (above MAs)")
        elif latest_price < ma20 < ma50:
            sentiment_score -= 2
            signals.append("Bearish trend (below MAs)")
        elif latest_price > ma20:
            sentiment_score += 1
            signals.append("Short-term bullish")
        
        # Determine overall sentiment
        if sentiment_score >= 3:
            overall_sentiment = "🟢 BULLISH"
        elif sentiment_score >= 1:
            overall_sentiment = "🟡 CAUTIOUSLY BULLISH" 
        elif sentiment_score <= -3:
            overall_sentiment = "🔴 BEARISH"
        elif sentiment_score <= -1:
            overall_sentiment = "🟠 CAUTIOUSLY BEARISH"
        else:
            overall_sentiment = "⚪ NEUTRAL"
        
        return {
            'sentiment': overall_sentiment,
            'score': sentiment_score,
            'signals': signals,
            'rsi': rsi,
            'change_pct': change_pct
        }
        
    except Exception as e:
        return {
            'sentiment': '⚪ NEUTRAL',
            'score': 0,
            'signals': [f"Analysis error: {str(e)}"],
            'rsi': 50,
            'change_pct': 0
        }

def generate_comprehensive_analysis(question, symbol, market_data):
    """Generate detailed analysis without heavy AI models"""
    try:
        # Calculate technical indicators
        market_data = calculate_technical_indicators(market_data)
        
        # Get basic metrics
        latest_price = float(market_data['Close'].iloc[-1])
        prev_price = float(market_data['Close'].iloc[-2]) if len(market_data) > 1 else latest_price
        change = latest_price - prev_price
        change_pct = (change / prev_price) * 100 if prev_price != 0 else 0
        
        # Get 52-week high/low
        high_52w = float(market_data['High'].max()) if 'High' in market_data.columns else latest_price
        low_52w = float(market_data['Low'].min()) if 'Low' in market_data.columns else latest_price
        
        # Get volume info
        avg_volume = int(market_data['Volume'].mean()) if 'Volume' in market_data.columns else 0
        latest_volume = int(market_data['Volume'].iloc[-1]) if 'Volume' in market_data.columns else 0
        
        # Get sentiment analysis
        sentiment_data = get_market_sentiment(symbol, market_data)
        
        # Build comprehensive analysis
        analysis = f"📊 **COMPREHENSIVE ANALYSIS FOR {symbol.upper()}**\n\n"
        
        # Price information
        analysis += f"💰 **CURRENT METRICS**\n"
        analysis += f"• Price: ₹{latest_price:.2f} ({change_pct:+.2f}%)\n"
        analysis += f"• Change: {'+' if change > 0 else ''}{change:.2f}\n"
        analysis += f"• 52W Range: ₹{low_52w:.2f} - ₹{high_52w:.2f}\n"
        analysis += f"• Position in Range: {((latest_price - low_52w) / (high_52w - low_52w) * 100):.1f}%\n"
        
        if avg_volume > 0:
            volume_ratio = latest_volume / avg_volume
            analysis += f"• Volume: {latest_volume:,} ({volume_ratio:.1f}x avg)\n"
        
        # Technical analysis
        analysis += f"\n🔍 **TECHNICAL ANALYSIS**\n"
        analysis += f"• Overall Sentiment: {sentiment_data['sentiment']}\n"
        analysis += f"• RSI (14): {sentiment_data['rsi']:.1f}\n"
        
        # Moving averages
        if 'MA20' in market_data.columns and not pd.isna(market_data['MA20'].iloc[-1]):
            ma20 = market_data['MA20'].iloc[-1]
            analysis += f"• 20-day MA: ₹{ma20:.2f} ({((latest_price/ma20-1)*100):+.1f}%)\n"
        
        if 'MA50' in market_data.columns and not pd.isna(market_data['MA50'].iloc[-1]):
            ma50 = market_data['MA50'].iloc[-1]
            analysis += f"• 50-day MA: ₹{ma50:.2f} ({((latest_price/ma50-1)*100):+.1f}%)\n"
        
        # Key signals
        analysis += f"\n📈 **KEY SIGNALS**\n"
        for signal in sentiment_data['signals']:
            analysis += f"• {signal}\n"
        
        # Question-specific advice
        analysis += f"\n💡 **SPECIFIC ADVICE FOR YOUR QUESTION**\n"
        question_lower = question.lower()
        
        if any(word in question_lower for word in ["buy", "invest", "purchase", "entry"]):
            if sentiment_data['score'] >= 2:
                analysis += "🟢 **FAVORABLE FOR BUYING**\n"
                analysis += "• Technical indicators suggest upward momentum\n"
                analysis += "• Consider buying on minor dips\n"
                analysis += "• Set stop-loss below recent support\n"
            elif sentiment_data['score'] <= -2:
                analysis += "🔴 **CAUTION ADVISED**\n"
                analysis += "• Wait for reversal signals\n"
                analysis += "• Consider dollar-cost averaging\n"
                analysis += "• Look for support levels\n"
            else:
                analysis += "🟡 **NEUTRAL - WAIT AND WATCH**\n"
                analysis += "• Mixed signals present\n"
                analysis += "• Wait for clearer direction\n"
                analysis += "• Consider systematic investment\n"
                
        elif any(word in question_lower for word in ["sell", "exit", "book", "profit"]):
            if sentiment_data['score'] >= 2:
                analysis += "🟡 **PARTIAL PROFIT BOOKING**\n"
                analysis += "• Trend is positive, avoid full exit\n"
                analysis += "• Consider booking 25-50% profits\n"
                analysis += "• Trailing stop-loss recommended\n"
            elif sentiment_data['score'] <= -2:
                analysis += "🔴 **CONSIDER SELLING**\n"
                analysis += "• Negative momentum building\n"
                analysis += "• Protect capital with stop-loss\n"
                analysis += "• Review position sizing\n"
            else:
                analysis += "🟡 **HOLD POSITION**\n"
                analysis += "• No clear sell signal\n"
                analysis += "• Monitor key levels\n"
                analysis += "• Maintain stop-loss\n"
                
        elif any(word in question_lower for word in ["trend", "direction", "outlook", "forecast"]):
            if sentiment_data['score'] >= 1:
                analysis += "📈 **UPWARD TREND LIKELY**\n"
                analysis += "• Positive technical setup\n"
                analysis += "• Higher probability of gains\n"
                analysis += "• Watch for resistance levels\n"
            elif sentiment_data['score'] <= -1:
                analysis += "📉 **DOWNWARD PRESSURE**\n"
                analysis += "• Bearish technical setup\n"
                analysis += "• Potential for further decline\n"
                analysis += "• Look for support levels\n"
            else:
                analysis += "➡️ **SIDEWAYS MOVEMENT**\n"
                analysis += "• Range-bound trading likely\n"
                analysis += "• Buy support, sell resistance\n"
                analysis += "• Wait for breakout\n"
        
        # Risk management
        analysis += f"\n⚠️ **RISK MANAGEMENT**\n"
        analysis += f"• Stop Loss: ₹{latest_price * 0.95:.2f} (-5%)\n"
        analysis += f"• Target 1: ₹{latest_price * 1.05:.2f} (+5%)\n"
        analysis += f"• Target 2: ₹{latest_price * 1.10:.2f} (+10%)\n"
        
        # Market context
        analysis += f"\n🌍 **MARKET CONTEXT**\n"
        analysis += f"• Consider overall market sentiment\n"
        analysis += f"• Check sector performance\n"
        analysis += f"• Review company fundamentals\n"
        analysis += f"• Monitor global cues\n"
        
        analysis += f"\n📅 **Analysis Date:** {datetime.now().strftime('%Y-%m-%d %H:%M')}\n"
        analysis += f"⚠️ **Disclaimer:** This analysis is for educational purposes only. Always consult financial advisors before investing."
        
        return analysis
        
    except Exception as e:
        return f"❌ Error generating analysis: {str(e)}"

# --- 2. Main Q&A Function ---
def analyze_market_question(question, symbol):
    if not question.strip():
        return "❌ Please enter a question about the market or stock."
    
    if not symbol.strip():
        return "❌ Please enter a valid stock symbol."
    
    # Get market data
    market_data, error = get_market_data_safely(symbol)
    if error:
        return f"❌ {error}"
    
    # Generate comprehensive analysis
    return generate_comprehensive_analysis(question, symbol, market_data)

# --- 3. Technical Chart Function ---
def plot_technical_chart(symbol):
    if not symbol.strip():
        fig, ax = plt.subplots(figsize=(12, 8))
        ax.text(0.5, 0.5, "Please enter a stock symbol", ha='center', va='center', fontsize=14)
        ax.axis('off')
        return fig
    
    market_data, error = get_market_data_safely(symbol)
    if error:
        fig, ax = plt.subplots(figsize=(12, 8))
        ax.text(0.5, 0.5, f"Error: {error}", ha='center', va='center', wrap=True, fontsize=12)
        ax.axis('off')
        return fig
    
    try:
        # Calculate technical indicators
        market_data = calculate_technical_indicators(market_data)
        
        # Create subplots
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10), height_ratios=[3, 1])
        
        # Main price chart
        ax1.plot(market_data['Date'], market_data['Close'], label='Close Price', linewidth=2, color='blue')
        
        # Moving averages
        if 'MA20' in market_data.columns:
            ax1.plot(market_data['Date'], market_data['MA20'], label='MA20', alpha=0.7, color='orange')
        if 'MA50' in market_data.columns:
            ax1.plot(market_data['Date'], market_data['MA50'], label='MA50', alpha=0.7, color='red')
        
        # Bollinger Bands
        if all(col in market_data.columns for col in ['BB_Upper', 'BB_Lower']):
            ax1.fill_between(market_data['Date'], market_data['BB_Upper'], market_data['BB_Lower'], 
                           alpha=0.2, color='gray', label='Bollinger Bands')
        
        ax1.set_title(f"{symbol.upper()} - Technical Analysis Chart", fontsize=16, fontweight='bold')
        ax1.set_ylabel("Price (₹)", fontsize=12)
        ax1.legend(loc='upper left')
        ax1.grid(True, alpha=0.3)
        
        # RSI subplot
        if 'RSI' in market_data.columns:
            ax2.plot(market_data['Date'], market_data['RSI'], color='purple', linewidth=2)
            ax2.axhline(y=70, color='r', linestyle='--', alpha=0.7, label='Overbought (70)')
            ax2.axhline(y=30, color='g', linestyle='--', alpha=0.7, label='Oversold (30)')
            ax2.fill_between(market_data['Date'], 30, 70, alpha=0.1, color='blue')
            ax2.set_title("RSI (14)", fontsize=12)
            ax2.set_ylabel("RSI", fontsize=10)
            ax2.set_ylim(0, 100)
            ax2.legend(loc='upper right', fontsize=8)
            ax2.grid(True, alpha=0.3)
        
        # Format dates
        plt.xticks(rotation=45)
        plt.tight_layout()
        
        return fig
        
    except Exception as e:
        fig, ax = plt.subplots(figsize=(12, 8))
        ax.text(0.5, 0.5, f"Chart Error: {str(e)}", ha='center', va='center', fontsize=12)
        ax.axis('off')
        return fig

# --- 4. Pattern Detection ---
def detect_chart_patterns(image, timeframe):
    if image is None:
        return "📷 Please upload a chart image first."
    
    analysis = f"📊 **PROFESSIONAL CHART PATTERN ANALYSIS**\n"
    analysis += f"⏱️ **Timeframe:** {timeframe}\n\n"
    
    # Pattern categories based on timeframe
    if timeframe in ["1m", "5m", "15m"]:
        analysis += "🚀 **SCALPING PATTERNS** (High Frequency)\n\n"
        patterns = [
            "• **Flag & Pennant** - Quick continuation patterns (5-15 minutes)",
            "• **Wedges** - Reversal patterns with tight stops",
            "• **Support/Resistance** - Key levels for quick bounces",
            "• **Breakout Patterns** - Volume confirmation essential"
        ]
    elif timeframe in ["1h", "4h"]:
        analysis += "📈 **SWING TRADING PATTERNS** (Medium Term)\n\n"
        patterns = [
            "• **Head & Shoulders** - Classic reversal pattern",
            "• **Double Top/Bottom** - Strong reversal signals",
            "• **Triangles** - Continuation or reversal",
            "• **Cup & Handle** - Bullish continuation pattern",
            "• **Rectangle** - Consolidation pattern"
        ]
    else:
        analysis += "📊 **POSITION TRADING PATTERNS** (Long Term)\n\n"
        patterns = [
            "• **Inverse Head & Shoulders** - Major trend reversal",
            "• **Ascending/Descending Triangle** - Trend continuation",
            "• **Symmetrical Triangle** - Neutral pattern",
            "• **Channel Patterns** - Trend channels",
            "• **Rounding Bottom** - Long-term accumulation"
        ]
    
    for pattern in patterns:
        analysis += f"{pattern}\n"
    
    analysis += f"\n🎯 **PATTERN ANALYSIS CHECKLIST**\n"
    analysis += f"✅ **Volume Confirmation** - Patterns need volume support\n"
    analysis += f"✅ **Breakout Direction** - Wait for decisive breakout\n"
    analysis += f"✅ **Risk-Reward Ratio** - Minimum 1:2 ratio recommended\n"
    analysis += f"✅ **Market Context** - Align with overall trend\n"
    analysis += f"✅ **Stop Loss Placement** - Below/above pattern boundaries\n\n"
    
    # Timeframe-specific advice
    analysis += f"⚡ **{timeframe.upper()} TRADING STRATEGY**\n"
    
    if timeframe in ["1m", "5m", "15m"]:
        analysis += f"• **Entry:** Quick entries on pattern completion\n"
        analysis += f"• **Stop Loss:** Tight stops (0.5-1% of price)\n"
        analysis += f"• **Target:** Quick 1-2% moves\n"
        analysis += f"• **Time Horizon:** Minutes to hours\n"
        analysis += f"• **Risk Level:** HIGH - Requires constant monitoring\n"
    elif timeframe in ["1h", "4h"]:
        analysis += f"• **Entry:** Wait for confirmation candles\n"
        analysis += f"• **Stop Loss:** Moderate stops (2-3% of price)\n"
        analysis += f"• **Target:** 5-10% moves expected\n"
        analysis += f"• **Time Horizon:** Days to weeks\n"
        analysis += f"• **Risk Level:** MEDIUM - Check 2-3 times daily\n"
    else:
        analysis += f"• **Entry:** Patient entries at key levels\n"
        analysis += f"• **Stop Loss:** Wider stops (5-8% of price)\n"
        analysis += f"• **Target:** 15-30% moves possible\n"
        analysis += f"• **Time Horizon:** Weeks to months\n"
        analysis += f"• **Risk Level:** LOW - Weekly monitoring sufficient\n"
    
    analysis += f"\n🔍 **WHAT TO LOOK FOR IN YOUR CHART:**\n"
    analysis += f"1. **Clear Pattern Formation** - Well-defined shapes\n"
    analysis += f"2. **Volume Characteristics** - Decreasing in pattern, increasing on breakout\n"
    analysis += f"3. **Price Action** - Respect of pattern boundaries\n"
    analysis += f"4. **Breakout Quality** - Strong momentum move\n"
    analysis += f"5. **Follow-through** - Sustained move after breakout\n\n"
    
    analysis += f"📚 **PATTERN RELIABILITY RANKING:**\n"
    analysis += f"🥇 **Most Reliable:** Head & Shoulders, Double Top/Bottom\n"
    analysis += f"🥈 **Good Reliability:** Triangles, Flags, Wedges  \n"
    analysis += f"🥉 **Moderate Reliability:** Rectangles, Channels\n\n"
    
    analysis += f"⚠️ **IMPORTANT REMINDERS:**\n"
    analysis += f"• No pattern is 100% reliable\n"
    analysis += f"• Always use proper risk management\n"
    analysis += f"• Practice with paper trading first\n"
    analysis += f"• Consider market sentiment and news\n"
    analysis += f"• Pattern failure is also a valid signal\n\n"
    
    analysis += f"🎓 **NEXT STEPS:**\n"
    analysis += f"1. Identify the pattern in your chart\n"
    analysis += f"2. Mark entry, stop-loss, and target levels\n"
    analysis += f"3. Wait for confirmation before entering\n"
    analysis += f"4. Monitor volume for validation\n"
    analysis += f"5. Execute your plan with discipline\n"
    
    return analysis

# --- 5. Lightweight Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft(), title="AI Trading Assistant") as demo:
    gr.Markdown("""
    # 🧠 AI Trading Assistant for Indian Markets
    
    **⚡ Optimized for Performance** | Professional Technical Analysis & Market Insights
    
    **📋 Supported Symbols:**
    - **Large Cap:** RELIANCE, TCS, INFY, HDFC, ICICIBANK, HINDUNILVR
    - **Mid Cap:** BAJAJFINSV, ASIANPAINT, MARUTI, SUNPHARMA
    - **Indices:** ^NSEI (Nifty 50), ^BSESN (Sensex)
    - **Sectors:** IT, BANKING, PHARMA, AUTO stocks
    """)
    
    with gr.Tab("🤖 Market Q&A"):
        gr.Markdown("### 💬 Get Professional Market Analysis")
        
        with gr.Row():
            symbol_in = gr.Textbox(
                value="RELIANCE", 
                label="📊 Stock Symbol",
                placeholder="Enter: RELIANCE, TCS, INFY, ^NSEI, etc.",
                scale=2
            )
        
        question_in = gr.Textbox(
            lines=3, 
            label="❓ Your Trading Question",
            placeholder="Ask specific questions like:\n• Should I buy RELIANCE now?\n• What's the trend for TCS?\n• Is HDFC good for long-term?",
            scale=3
        )
        
        ask_btn = gr.Button("🔍 Get Professional Analysis", variant="primary", size="lg")
        answer_out = gr.Textbox(label="📊 Professional Analysis", lines=25, max_lines=30)
        
        gr.Markdown("""
        **💡 Pro Tips for Better Analysis:**
        - Be specific about your trading timeframe
        - Mention if you're looking to buy, sell, or hold
        - Ask about specific technical levels
        - Include your risk tolerance in questions
        """)
        
        ask_btn.click(
            analyze_market_question, 
            inputs=[question_in, symbol_in], 
            outputs=answer_out
        )
    
    with gr.Tab("📈 Technical Charts"):
        gr.Markdown("### 📊 Advanced Technical Analysis Charts")
        
        with gr.Row():
            tech_symbol = gr.Textbox(
                value="RELIANCE", 
                label="📊 Stock Symbol",
                placeholder="Enter symbol for technical analysis",
                scale=2
            )
            chart_btn = gr.Button("📈 Generate Advanced Chart", variant="primary", scale=1)
        
        chart_plot = gr.Plot()
        
        gr.Markdown("""
        **📊 Chart Features:**
        - Price action with volume analysis
        - Moving averages (20-day, 50-day)
        - RSI indicator for momentum
        - Bollinger Bands for volatility
        - Support and resistance levels
        """)
        
        chart_btn.click(
            plot_technical_chart, 
            inputs=[tech_symbol], 
            outputs=chart_plot
        )
    
    with gr.Tab("🔍 Pattern Analysis"):
        gr.Markdown("### 📸 Professional Chart Pattern Recognition")
        
        with gr.Row():
            img = gr.Image(type="pil", label="📷 Upload Your Chart Screenshot", scale=2)
            timeframe = gr.Dropdown(
                choices=["1m", "5m", "15m", "1h", "4h", "1d", "1w"], 
                value="1h", 
                label="⏱️ Trading Timeframe",
                scale=1
            )
        
        detect_btn = gr.Button("🔍 Analyze Chart Patterns", variant="primary", size="lg")
        pattern_out = gr.Textbox(label="📊 Professional Pattern Analysis", lines=30, max_lines=35)
        
        gr.Markdown("""
        **📸 Chart Upload Tips:**
        - Use clear, high-resolution screenshots
        - Include volume bars if possible
        - Mark important levels you see
        - Specify the time period shown
        """)
        
        detect_btn.click(
            detect_chart_patterns, 
            inputs=[img, timeframe], 
            outputs=pattern_out
        )
    
    gr.Markdown("""
    ---
    **⚠️ Important Disclaimers:**
    - This tool provides educational analysis only
    - Past performance doesn't guarantee future results  
    - Always do your own research before investing
    - Consider consulting certified financial advisors
    - Markets are subject to risks - invest wisely
    
    **🔄 Last Updated:** June 2025 | **💡 Optimized for Hugging Face Spaces**
    """)

# Launch the application
if __name__ == "__main__":
    demo.launch()