app.py

import gradio as gr
import yfinance as yf
from prophet import Prophet
from sklearn.linear_model import LinearRegression
from neuralprophet import NeuralProphet
import pandas as pd
from datetime import datetime
import plotly.graph_objects as go

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']]
        data.rename(columns={'Date': 'ds', 'Adj Close': 'y'}, inplace=True)
    else:
        raise ValueError("Expected 'Adj Close' in columns")
    return data

def predict_future_prices(ticker, periods=1825):
    data = download_data(ticker)
    
    # Prophet 모델 생성 및 학습
    model_prophet = Prophet(daily_seasonality=False, weekly_seasonality=False, yearly_seasonality=True)
    model_prophet.fit(data)
    
    # 미래 데이터 프레임 생성 및 예측
    future = model_prophet.make_future_dataframe(periods=periods, freq='D')
    forecast_prophet = model_prophet.predict(future)
    
    # Linear Regression 모델 생성 및 학습
    model_lr = LinearRegression()
    X = pd.to_numeric(pd.Series(range(len(data))))
    y = data['y'].values
    model_lr.fit(X.values.reshape(-1, 1), y)
    
    # 미래 데이터 프레임 생성 및 예측
    future_dates = pd.date_range(start=data['ds'].iloc[-1], periods=periods+1, freq='D')[1:]
    future_lr = pd.DataFrame({'ds': future_dates})
    future_lr['ds'] = future_lr['ds'].dt.strftime('%Y-%m-%d')
    X_future = pd.to_numeric(pd.Series(range(len(data), len(data) + len(future_lr))))
    future_lr['yhat'] = model_lr.predict(X_future.values.reshape(-1, 1))

# NeuralProphet 모델 생성 및 학습
    model_np = NeuralProphet()
    metrics = model_np.fit(data, freq='D')
    future_np = model_np.make_future_dataframe(data, periods=periods)
    forecast_np = model_np.predict(future_np)

# 예측 결과 그래프 생성
    forecast_prophet['ds'] = forecast_prophet['ds'].dt.strftime('%Y-%m-%d')
    forecast_np['ds'] = pd.to_datetime(forecast_np['ds']).dt.strftime('%Y-%m-%d')
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=forecast_prophet['ds'], y=forecast_prophet['yhat'], mode='lines', name='Prophet Forecast (Blue)', line=dict(color='blue')))
    fig.add_trace(go.Scatter(x=future_lr['ds'], y=future_lr['yhat'], mode='lines', name='Linear Regression Forecast (Red)', line=dict(color='red')))
    fig.add_trace(go.Scatter(x=forecast_np['ds'], y=forecast_np['yhat1'], mode='lines', name='NeuralProphet Forecast (Green)', line=dict(color='green')))
    fig.add_trace(go.Scatter(x=data['ds'], y=data['y'], mode='lines', name='Actual (Black)', line=dict(color='black')))

    return fig, forecast_prophet[['ds', 'yhat', 'yhat_lower', 'yhat_upper']], future_lr[['ds', 'yhat']], forecast_np[['ds', 'yhat1']]


# 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_prophet = gr.Dataframe(label="Prophet Forecast Data")
    forecast_data_lr = gr.Dataframe(label="Linear Regression Forecast Data")
    forecast_data_np = gr.Dataframe(label="NeuralProphet Forecast Data")
    
    forecast_button.click(
        fn=predict_future_prices,
        inputs=[ticker_input, periods_input],
        outputs=[forecast_chart, forecast_data_prophet, forecast_data_lr, forecast_data_np]
    )

app.launch()