import gradio as gr
import yfinance as yf
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import plotly.graph_objects as go
from transformers import pipeline
import torch
from textblob import TextBlob
import nltk
from ta.trend import SMAIndicator, MACD
from ta.momentum import RSIIndicator

# Download NLTK data
try:
    nltk.download('punkt', quiet=True)
    nltk.download('averaged_perceptron_tagger', quiet=True)
except Exception as e:
    print(f"NLTK download warning (not critical): {e}")

class StockAnalysisApp:
    def __init__(self):
        try:
            self.sentiment_analyzer = pipeline("sentiment-analysis", 
                                            device=0 if torch.cuda.is_available() else -1)
        except:
            self.sentiment_analyzer = pipeline("sentiment-analysis", device=-1)
    
    def get_stock_data(self, ticker, period='1y'):
        """Fetch stock data"""
        try:
            stock = yf.Ticker(ticker)
            data = stock.history(period=period)
            return data, stock.info
        except Exception as e:
            return None, None

    def create_chart(self, data, ticker):
        """Create interactive stock chart"""
        try:
            fig = go.Figure()
            
            # Candlestick chart
            fig.add_trace(go.Candlestick(
                x=data.index,
                open=data['Open'],
                high=data['High'],
                low=data['Low'],
                close=data['Close'],
                name=ticker
            ))
            
            # Add SMAs
            sma20 = SMAIndicator(close=data['Close'], window=20).sma_indicator()
            sma50 = SMAIndicator(close=data['Close'], window=50).sma_indicator()
            
            fig.add_trace(go.Scatter(x=data.index, y=sma20,
                                   name='SMA20',
                                   line=dict(color='orange')))
            
            fig.add_trace(go.Scatter(x=data.index, y=sma50,
                                   name='SMA50',
                                   line=dict(color='blue')))
            
            fig.update_layout(
                title=f'{ticker} Stock Price',
                yaxis_title='Price',
                template='plotly_dark',
                xaxis_rangeslider_visible=False
            )
            
            return fig
        except Exception as e:
            return None

    def get_technical_analysis(self, data):
        """Generate technical analysis"""
        try:
            current_price = data['Close'].iloc[-1]
            prev_price = data['Close'].iloc[-2]
            price_change = ((current_price - prev_price) / prev_price) * 100
            
            # Calculate indicators
            rsi = RSIIndicator(close=data['Close']).rsi().iloc[-1]
            macd = MACD(close=data['Close'])
            macd_line = macd.macd().iloc[-1]
            signal_line = macd.macd_signal().iloc[-1]
            
            analysis = f"""
            Technical Analysis Summary:
            
            Current Price: ${current_price:.2f}
            Daily Change: {price_change:.2f}%
            
            Technical Indicators:
            - RSI (14): {rsi:.2f} ({'Overbought' if rsi > 70 else 'Oversold' if rsi < 30 else 'Neutral'})
            - MACD: {macd_line:.2f}
            - Signal Line: {signal_line:.2f}
            - MACD Status: {'Bullish' if macd_line > signal_line else 'Bearish'}
            
            Volume Analysis:
            - Current Volume: {int(data['Volume'].iloc[-1]):,}
            - Avg Volume (20D): {int(data['Volume'].rolling(20).mean().iloc[-1]):,}
            """
            return analysis
        except Exception as e:
            return f"Error in technical analysis: {str(e)}"

    def process_query(self, message, history):
        """Process chat queries"""
        try:
            message = message.strip()
            
            # Extract potential stock ticker
            words = message.split()
            ticker = None
            for word in words:
                if word.isupper() and 1 < len(word) <= 5:
                    ticker = word
                    break
            
            if ticker:
                data, info = self.get_stock_data(ticker)
                if data is not None:
                    analysis = self.get_technical_analysis(data)
                    return analysis
            
            # General queries
            message_lower = message.lower()
            if "help" in message_lower:
                return """I can help you with:
                1. Stock Analysis (e.g., "Analyze AAPL")
                2. Technical Indicators (e.g., "What's RSI?")
                3. Market Information (e.g., "Tell me about TSLA")
                
                You can also use the interface above to:
                - View stock charts
                - Get detailed technical analysis
                - See price predictions
                - Track multiple stocks"""
            
            return "Please provide a stock ticker or ask for help to see what I can do."
            
        except Exception as e:
            return f"Error processing query: {str(e)}"

def create_ui():
    """Create the complete Gradio interface"""
    app = StockAnalysisApp()
    
    def analyze_stock(ticker, period):
        try:
            data, info = app.get_stock_data(ticker, period)
            if data is None:
                return None, "Error fetching stock data"
            
            chart = app.create_chart(data, ticker)
            analysis = app.get_technical_analysis(data)
            
            return chart, analysis
        except Exception as e:
            return None, f"Error: {str(e)}"
    
    # Create the interface
    with gr.Blocks(theme=gr.themes.Soft()) as interface:
        gr.Markdown("""
        # Stock Analysis Dashboard
        Enter a stock ticker and select analysis period to get started.
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                ticker_input = gr.Textbox(
                    label="Stock Ticker",
                    placeholder="e.g., AAPL",
                    value="AAPL"
                )
                period_input = gr.Dropdown(
                    choices=["1mo", "3mo", "6mo", "1y", "2y", "5y"],
                    value="1y",
                    label="Analysis Period"
                )
                analyze_button = gr.Button("Analyze Stock")
            
            with gr.Column(scale=2):
                with gr.Tab("Chart"):
                    chart_output = gr.Plot()
                with gr.Tab("Analysis"):
                    analysis_output = gr.Textbox(
                        label="Technical Analysis",
                        lines=10
                    )
        
        gr.Markdown("---")
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Chat with AI Assistant")
                chatbot = gr.ChatInterface(
                    app.process_query,
                    examples=[
                        "Analyze AAPL",
                        "What's the trend for TSLA?",
                        "Help"
                    ]
                )
        
        # Set up event handlers
        analyze_button.click(
            analyze_stock,
            inputs=[ticker_input, period_input],
            outputs=[chart_output, analysis_output]
        )
    
    return interface

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