"""
경제 지표 데이터 처리 관련 함수
"""

import pandas as pd
import yfinance as yf
import FinanceDataReader as fdr
from fredapi import Fred
from statsmodels.tsa.holtwinters import ExponentialSmoothing

def get_economic_data(start_date, end_date, fred_api_key):
    """
    경제 지표와 시장 데이터를 수집하는 함수
    """

    # 경제 지표 데이터 다운로드 (FRED API)
    fred = Fred(api_key=fred_api_key)

    # 국채 수익률
    DGS = pd.concat([
        fred.get_series('DGS2', start_date, end_date),
        fred.get_series('DGS5', start_date, end_date),
        fred.get_series('DGS10', start_date, end_date)
    ], axis=1)
    DGS.columns = ['2-year', '5-year', '10-year']

    # 기타 경제 지표
    T10Y2Y = fdr.DataReader('FRED:T10Y2Y', start_date, end_date)
    VIX = fdr.DataReader('FRED:VIXCLS', start_date, end_date)
    Unemployment_Rate = fdr.DataReader('FRED:UNRATE', start_date, end_date)
    CPI = fdr.DataReader('FRED:CPIAUCSL', start_date, end_date)
    FEDFUNDS = fdr.DataReader('FRED:FEDFUNDS', start_date, end_date)
    GDP = pd.DataFrame(fred.get_series('GDP', start_date, end_date), columns=['GDP'])

    # 주요 주식 지수
    index_tickers = {
        "^DJI": "DJI Close", 
        "NDAQ": "NDAQ Close",
        "^GSPC": "SPX Close",
        "^RUT": "RUT Close"
    }

    index_data = {}
    for ticker, name in index_tickers.items():
        df = yf.download(ticker, start=start_date, end=end_date)
        df = df[['Close']].rename(columns={'Close': name})
        index_data[name] = df

    Index_data = pd.concat(index_data.values(), axis=1)

    # 산업별 ETF
    sectors = {
        "VDE": "Energy",
        "MXI": "Materials",
        "VIS": "Industrials",
        "VCR": "Consumer Cyclical",
        "XLP": "Consumer Staples",
        "VHT": "Health Care",
        "XLF": "Financials",
        "VGT": "Information Technology",
        "VOX": "Communication Services",
        "XLU": "Utilities",
        "VNQ": "Real Estate"
    }

    sector_data = {}
    for etf, sector_name in sectors.items():
        df = yf.download(etf, start=start_date, end=end_date)
        df.rename(columns={'Close': f'{sector_name} Close'}, inplace=True)
        sector_data[sector_name] = df[[f'{sector_name} Close']]

    ETF_data = pd.concat(sector_data.values(), axis=1)

    # 데이터 보간 및 예측
    def interpolate_and_forecast(df, col_name, stock_end_date):
        df = df.resample('D').asfreq().interpolate()
        forecast_steps = (pd.to_datetime(stock_end_date) - df.index[-1]).days
        if forecast_steps > 0:
            forecast_df = pd.DataFrame(index=pd.date_range(df.index[-1] + pd.Timedelta(days=1), stock_end_date))
            model = ExponentialSmoothing(df[col_name], trend='add').fit()
            forecast_df[col_name] = model.forecast(steps=forecast_steps)
            df = pd.concat([df, forecast_df])
        return df

    # 모든 경제 및 시장 지표 결합
    econ_df = (
        interpolate_and_forecast(DGS, '2-year', end_date)
        .join(interpolate_and_forecast(T10Y2Y, 'T10Y2Y', end_date), how='left')
        .join(interpolate_and_forecast(VIX, 'VIXCLS', end_date), how='left')
        .join(interpolate_and_forecast(Unemployment_Rate, 'UNRATE', end_date), how='left')
        .join(interpolate_and_forecast(CPI, 'CPIAUCSL', end_date), how='left')
        .join(interpolate_and_forecast(FEDFUNDS, 'FEDFUNDS', end_date), how='left')
        .join(interpolate_and_forecast(GDP, 'GDP', end_date), how='left')
    )

    # 멀티인덱스 처리
    if isinstance(Index_data.columns, pd.MultiIndex):
        Index_data.columns = Index_data.columns.get_level_values(1)
    if isinstance(ETF_data.columns, pd.MultiIndex):
        ETF_data.columns = ETF_data.columns.get_level_values(1)
        
    econ_df = econ_df.join(Index_data, how='left').join(ETF_data, how='left')
    econ_df.fillna(method='ffill', inplace=True)
    econ_df.fillna(method='bfill', inplace=True)
    
    return econ_df