import yfinance as yf
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from prophet import Prophet
import gradio as gr
from ta.trend import SMAIndicator, EMAIndicator, MACD
from ta.momentum import RSIIndicator

# Function for trading strategy signals
def trading_signals(df, strategy="day trading"):
    signals = ""
    df = df.copy()
    
    try:
        if strategy.lower() == "day trading":
            # Use Bollinger Bands and RSI for intraday trades
            df['RSI'] = RSIIndicator(df['Close'], window=14).rsi()
            df['EMA_9'] = EMAIndicator(df['Close'], window=9).ema_indicator()
            df['EMA_21'] = EMAIndicator(df['Close'], window=21).ema_indicator()
            df['Signal'] = (df['EMA_9'] > df['EMA_21']).astype(int)
            
            signals = f"Day Trading Signals: {'Buy' if df['Signal'].iloc[-1] else 'Sell'}"
        
        elif strategy.lower() == "swing trading":
            # Swing Trading: SMA crossover
            df['SMA_50'] = SMAIndicator(df['Close'], window=50).sma_indicator()
            df['SMA_200'] = SMAIndicator(df['Close'], window=200).sma_indicator()
            df['Signal'] = (df['SMA_50'] > df['SMA_200']).astype(int)
            
            signals = f"Swing Trading Signals: {'Bullish' if df['Signal'].iloc[-1] else 'Bearish'}"
        
        elif strategy.lower() == "momentum trading":
            # Momentum Trading: RSI + MACD
            df['RSI'] = RSIIndicator(df['Close'], window=14).rsi()
            macd = MACD(df['Close'])
            df['MACD'] = macd.macd()
            df['MACD_signal'] = macd.macd_signal()
            
            # Simplified momentum condition
            df['Momentum'] = ((df['RSI'] > 50) & (df['MACD'] > df['MACD_signal'])).astype(int)
            
            signals = f"Momentum Trading Signal: {'Strong Momentum' if df['Momentum'].iloc[-1] else 'Weak Momentum'}"
        
        elif strategy.lower() == "scalping":
            # Scalping: EMA crossover on 1-min data
            df['EMA_3'] = EMAIndicator(df['Close'], window=3).ema_indicator()
            df['EMA_9'] = EMAIndicator(df['Close'], window=9).ema_indicator()
            df['Signal'] = (df['EMA_3'] > df['EMA_9']).astype(int)
            
            signals = f"Scalping Signal: {'Buy' if df['Signal'].iloc[-1] else 'Sell'}"
        
        else:
            signals = "Invalid Strategy"
        
    except Exception as e:
        signals = f"Error in trading signals: {e}"
    
    return signals

# Main function to fetch data, generate trends, and trading signals
def stock_analysis_with_trading(ticker, period="6mo", interval="1d", analysis_type="closing price trend", strategy=None):
    try:
        # Fetch historical stock data with enhanced error handling
        stock = yf.Ticker(ticker)
        
        # Try different methods to fetch data
        try:
            df = stock.history(period=period, interval=interval)
        except Exception as initial_error:
            # Fallback to different periods or intervals
            try:
                df = stock.history(period="1y", interval="1d")
            except Exception as fallback_error:
                return f"Error fetching data: {initial_error} | Fallback Error: {fallback_error}", None
        
        if df is None or df.empty:
            return f"No data found for ticker: {ticker}. Check ticker symbol and network connection.", None
        
        if analysis_type.lower() == "closing price trend":
            # Plot Closing Price Trend
            plt.figure(figsize=(10, 6))
            plt.plot(df.index, df['Close'], label='Closing Price')
            plt.title(f"Closing Price Trend for {ticker}")
            plt.xlabel("Date")
            plt.ylabel("Closing Price (USD)")
            plt.grid()
            plt.legend()
            plot_path = f"{ticker}_closing_price_trend.png"
            plt.savefig(plot_path)
            plt.close()
            return f"Closing Price Analysis for {ticker}", plot_path
        
        elif analysis_type.lower() == "forecast":
            # Forecast with Prophet
            forecast_df = df.reset_index()[['Date', 'Close']]
            forecast_df.columns = ['ds', 'y']
            model = Prophet()
            model.fit(forecast_df)
            future = model.make_future_dataframe(periods=30)
            forecast = model.predict(future)
            
            plt.figure(figsize=(10, 6))
            model.plot(forecast)
            plt.title(f"Forecasted Closing Price for {ticker}")
            plt.xlabel("Date")
            plt.ylabel("Price (USD)")
            plot_path = f"{ticker}_forecast_trend.png"
            plt.savefig(plot_path)
            plt.close()
            return "Forecast Generated Successfully", plot_path
        
        elif strategy:
            signals = trading_signals(df, strategy)
            return signals, None
        
        else:
            return "Invalid analysis type or strategy provided.", None
    
    except Exception as e:
        return f"Comprehensive Error: {e}", None

# Gradio Interface
def chat_with_agent(ticker, period, interval, query, strategy):
    response, plot_path = stock_analysis_with_trading(ticker, period, interval, query, strategy)
    return response, plot_path if plot_path else None

# Define Gradio Interface
iface = gr.Interface(
    fn=chat_with_agent,
    inputs=[
        gr.Textbox(label="Stock Ticker (e.g., AAPL, MSFT)"),
        gr.Textbox(label="Time Period (e.g., 6mo, 1d)", value="6mo"),
        gr.Textbox(label="Data Interval (e.g., 1d, 1h)", value="1d"),
        gr.Textbox(label="Query (e.g., closing price trend, forecast)", value="closing price trend"),
        gr.Textbox(label="Trading Strategy (optional)", value=""),
    ],
    outputs=[
        gr.Textbox(label="Analysis Result or Trading Signal"),
        gr.Image(label="Trend or Forecast Visualization")
    ],
    title="Stock Analysis & Trading Strategy Agent",
    description="Analyze stock trends, forecast prices, and generate trading signals using technical indicators."
)

# Launch the interface
iface.launch(share=True)