app.py

import gradio as gr
import yfinance as yf
from prophet import Prophet
import pandas as pd
from datetime import datetime
import plotly.graph_objects as go
import plotly.express as px

def download_data(ticker, start_date='2010-01-01'):
    """ 주식 데이터를 다운로드하고 포맷을 조정하는 함수 """
    data = yf.download(ticker, start=start_date)
    if data.empty:
        raise ValueError(f"No data returned for {ticker}")
    data.reset_index(inplace=True)
    if 'Adj Close' in data.columns:
        data = data[['Date', 'Adj Close']].copy()
        data.rename(columns={'Date': 'ds', 'Adj Close': 'y'}, inplace=True)
        data['ds'] = pd.to_datetime(data['ds'])  # Ensure 'ds' is datetime type
    else:
        raise ValueError("Expected 'Adj Close' in columns")
    return data

def predict_future_prices(ticker, periods=1825):
    data = download_data(ticker)
    
    # Prophet 모델 생성 및 학습
    model = Prophet(daily_seasonality=False, weekly_seasonality=True, yearly_seasonality=True)
    model.fit(data)
    
    # 미래 데이터 프레임 생성 및 예측
    future = model.make_future_dataframe(periods=periods, freq='D')
    forecast = model.predict(future)
    
    # 예측 결과 그래프 생성
    fig_main = go.Figure()
    fig_main.add_trace(go.Scatter(x=data['ds'], y=data['y'], mode='lines', name='Actual (Black)', line=dict(color='black')))
    fig_main.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat'], mode='lines', name='Forecast (Blue)'))
    
    # 연간 및 주간 계절성 그래프 생성
    # 임시로 matplotlib 플롯 반환을 비활성화하고 결과만 반환
    fig_seasonal = model.plot_components(forecast)
    forecast['ds'] = pd.to_datetime(forecast['ds'])  # Revert to datetime type to avoid AttributeError
    
    fig_yearly = px.line(x=pd.to_datetime(forecast['ds']), y=forecast['yearly'], labels={'x': 'Date', 'y': 'Yearly Trend'})
    fig_weekly = px.line(x=pd.to_datetime(forecast['ds']), y=forecast['weekly'], labels={'x': 'Date', 'y': 'Weekly Trend'})

    return fig_main, forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']], fig_yearly, fig_weekly

# Gradio 인터페이스 설정 및 실행
with gr.Blocks() as app:
    with gr.Row():
        ticker_input = gr.Textbox(value="AAPL", label="Enter Stock Ticker for Forecast")
        periods_input = gr.Number(value=1825, label="Forecast Period (days)")
        forecast_button = gr.Button("Generate Forecast")
        
    forecast_chart = gr.Plot(label="Forecast Chart")
    forecast_data = gr.Dataframe(label="Forecast Data")
    yearly_chart = gr.Plot(label="Yearly (Monthly) Trend Chart")
    weekly_chart = gr.Plot(label="Weekly Trend Chart")
    
    forecast_button.click(
        fn=predict_future_prices,
        inputs=[ticker_input, periods_input],
        outputs=[forecast_chart, forecast_data, yearly_chart, weekly_chart]
    )

app.launch()