# %%
# Install necessary packages if not already installed
# pip install gradio yfinance prophet plotly matplotlib

import gradio as gr
import pandas as pd
import yfinance as yf
from datetime import datetime
from prophet import Prophet
import plotly.express as px
import plotly.graph_objects as go
import matplotlib.pyplot as plt
import numpy as np

# Functions for calculating indicators (SMA, RSI, etc.) and generating trading signals
# (Reuse the code you've already written for technical indicators and forecasting)

def calculate_sma(df, window):
    return df['Close'].rolling(window=window).mean()

def calculate_macd(df):
    short_ema = df['Close'].ewm(span=12, adjust=False).mean()
    long_ema = df['Close'].ewm(span=26, adjust=False).mean()
    macd = short_ema - long_ema
    signal = macd.ewm(span=9, adjust=False).mean()
    return macd, signal

def calculate_rsi(df):
    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
    rsi = 100 - (100 / (1 + rs))
    return rsi

def calculate_bollinger_bands(df):
    middle_bb = df['Close'].rolling(window=20).mean()
    upper_bb = middle_bb + 2 * df['Close'].rolling(window=20).std()
    lower_bb = middle_bb - 2 * df['Close'].rolling(window=20).std()
    return middle_bb, upper_bb, lower_bb

def calculate_stochastic_oscillator(df):
    lowest_low = df['Low'].rolling(window=14).min()
    highest_high = df['High'].rolling(window=14).max()
    slowk = ((df['Close'] - lowest_low) / (highest_high - lowest_low)) * 100
    slowd = slowk.rolling(window=3).mean()
    return slowk, slowd

def generate_trading_signals(df):

    # Calculate Simple Moving Averages (SMA)
    df['SMA_50'] = calculate_sma(df, 50)
    df['SMA_200'] = calculate_sma(df, 200)

    # Calculate other technical indicators
    df['RSI'] = calculate_rsi(df)
    df['MiddleBB'], df['UpperBB'], df['LowerBB'] = calculate_bollinger_bands(df)
    df['SlowK'], df['SlowD'] = calculate_stochastic_oscillator(df)

    # Generate trading signals
    df['SMA_Signal'] = np.where(df['SMA_50'] > df['SMA_200'], 1, 0)
    
    macd, signal = calculate_macd(df)
    df['MACD_Signal'] = np.where((macd > signal.shift(1)) & (macd.shift(1) < signal), 1, 0)
    
    df['RSI_Signal'] = np.where(df['RSI'] < 30, 1, 0)
    df['RSI_Signal'] = np.where(df['RSI'] > 70, -1, df['RSI_Signal'])
    
    df['BB_Signal'] = np.where(df['Close'] < df['LowerBB'], 1, 0)
    df['BB_Signal'] = np.where(df['Close'] > df['UpperBB'], -1, df['BB_Signal'])
    
    df['Stochastic_Signal'] = np.where((df['SlowK'] < 20) & (df['SlowD'] < 20), 1, 0)
    df['Stochastic_Signal'] = np.where((df['SlowK'] > 80) & (df['SlowD'] > 80), -1, df['Stochastic_Signal'])

    # Summing the values of each individual signal column
    df['Combined_Signal'] = df[['SMA_Signal', 'MACD_Signal', 'RSI_Signal', 'BB_Signal', 'Stochastic_Signal']].sum(axis=1)






# %%
import plotly.graph_objects as go

def plot_combined_signals(df, ticker):
    # Create a figure
    fig = go.Figure()

    # Add closing price trace
    fig.add_trace(go.Scatter(
        x=df.index, y=df['Close'], 
        mode='lines', 
        name='Closing Price', 
        line=dict(color='lightcoral', width=2)
    ))

    # Add buy signals
    buy_signals = df[df['Combined_Signal'] >= 2]
    fig.add_trace(go.Scatter(
        x=buy_signals.index, y=buy_signals['Close'], 
        mode='markers', 
        marker=dict(symbol='triangle-up', size=10, color='lightgreen'), 
        name='Buy Signal'
    ))

    # Add sell signals
    sell_signals = df[df['Combined_Signal'] <= -2]
    fig.add_trace(go.Scatter(
        x=sell_signals.index, y=sell_signals['Close'], 
        mode='markers', 
        marker=dict(symbol='triangle-down', size=10, color='lightsalmon'), 
        name='Sell Signal'
    ))

    # Add combined signal trace
    fig.add_trace(go.Scatter(
        x=df.index, y=df['Combined_Signal'], 
        mode='lines', 
        name='Combined Signal', 
        line=dict(color='deepskyblue', width=2), 
        yaxis='y2'
    ))

    # Update layout for secondary y-axis
    fig.update_layout(
        title=f'{ticker}: Stock Price and Combined Trading Signal (Last 60 Days)',
        xaxis=dict(title='Date', gridcolor='gray', gridwidth=0.5),
        yaxis=dict(title='Price', side='left', gridcolor='gray', gridwidth=0.5),
        yaxis2=dict(title='Combined Signal', overlaying='y', side='right', showgrid=False),
        plot_bgcolor='black',
        paper_bgcolor='black',
        font=dict(color='white'),
        legend=dict(x=0.01, y=0.99, bgcolor='rgba(0,0,0,0)'),
        hovermode='x unified'
    )

    return fig


# %%
def stock_analysis(ticker, start_date, end_date):
    # Download stock data from Yahoo Finance
    df = yf.download(ticker, start=start_date, end=end_date)
    
    # Run your existing trading signals and indicators here
    generate_trading_signals(df)
    
    # Last 60 days
    df_last_60 = df.tail(60)
    
    # Plot trading signals using the improved function
    fig_signals = plot_combined_signals(df_last_60, ticker)
    
    
    # Combine the figures into HTML output
    return fig_signals


# %%
# Define Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Stock Market Analysis App")

    ticker_input = gr.Textbox(label="Enter Stock Ticker (e.g., AAPL, NVDA)", value="NVDA")
    start_date_input = gr.Textbox(label="Start Date (YYYY-MM-DD)", value="2022-01-01")
    end_date_input = gr.Textbox(label="End Date (YYYY-MM-DD)", value=str(datetime.now().date()))

    # Create a submit button that runs the stock analysis function
    button = gr.Button("Analyze Stock")
    
    # Outputs: Display results, charts
    signals_output = gr.Plot(label="Trading Signals")

    # Link button to function
    button.click(stock_analysis, inputs=[ticker_input, start_date_input, end_date_input], outputs=[signals_output, forecast_output])

# Launch the interface
demo.launch()