import gradio as gr
import yfinance as yf
import plotly.graph_objects as go
import pandas as pd

def get_interval(period):
    if period in ['1mo', '3mo']:
        return '1d'
    elif period in ['6mo', '1y']:
        return '5d'
    elif period in ['2y', '5y', '10y', 'ytd']:
        return '1wk'
    else:
        return '1mo'

def get_moving_average_windows(period):
    if period in ['1mo', '3mo']:
        return 5, 20  # Shorter windows for shorter periods
    elif period in ['6mo', '1y']:
        return 20, 50
    else:
        return 50, 100  # Longer windows for longer periods

def plot_stock_data(ticker, period):
    interval = get_interval(period)
    short_window, long_window = get_moving_average_windows(period)

    data = yf.Ticker(ticker).history(period=period, interval=interval)

    candlestick = go.Candlestick(x=data.index, open=data['Open'], high=data['High'], low=data['Low'], close=data['Close'], name='Candlestick')
    ma_short = go.Scatter(x=data.index, y=data['Close'].rolling(window=short_window).mean(), mode='lines', name=f'{short_window}-Day MA')
    ma_long = go.Scatter(x=data.index, y=data['Close'].rolling(window=long_window).mean(), mode='lines', name=f'{long_window}-Day MA')
    dividends = data[data['Dividends'] > 0]
    dividend_events = go.Scatter(x=dividends.index, y=dividends['Close'], mode='markers', name='Dividends', marker=dict(color='orange', size=10), hovertext=['Dividend: $' + str(div) for div in dividends['Dividends']], hoverinfo="text")

    fig = go.Figure(data=[candlestick, ma_short, ma_long, dividend_events])
    fig.update_layout(title=f'{ticker} Stock Data - {period}', xaxis_rangeslider_visible=False)

    return fig

interface = gr.Interface(
    fn=plot_stock_data,
    inputs=["text", gr.Dropdown(['1mo', '3mo', '6mo', '1y', '2y', '5y', '10y', 'ytd', 'max'])],
    outputs=gr.Plot(label="Stock Data Graph"),
    title="Historical Stock Data Visualization",
    description="Visualize historical stock data with candlestick graphs, including dynamically adjusted moving averages and dividend events."
)

interface.launch(share=False, debug=True)