import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm
from curl_cffi import requests
import yfinance as yf

# —————————————————————————————————————————————
# 1. 글로벌 지수 데이터 조회
# —————————————————————————————————————————————
def fetch_market_data(tickers, start_date, end_date, session):
    """
    yfinance를 이용해 글로벌 지수(종가) 데이터를 내려받습니다.
    :param tickers: 조회할 심볼 리스트 (예: ['^GSPC', '^KS11'])
    :param start_date: 조회 시작일 (YYYY-MM-DD)
    :param end_date: 조회 종료일 (YYYY-MM-DD)
    :param session: curl_cffi requests 세션
    :return: 날짜 인덱스에 맞춰진 DataFrame (columns=['S&P500','KOSPI'])
    """
    raw = yf.download(
        tickers,
        start=start_date,
        end=end_date,
        session=session,
        repair=True
    )['Close']
    raw.columns = ['S&P500', 'KOSPI']
    # 인덱스를 열로 올리고 이름 변경
    df = raw.reset_index().rename(columns={'index': 'Date'})
    return df

# —————————————————————————————————————————————
# 2. 로컬 엑셀 데이터 로드 및 전처리
# —————————————————————————————————————————————
def load_and_preprocess_excel(path, sheet, header_row=5):
    """
    out.xlsx의 PRICE_DATA 시트를 읽어온 뒤,
    필요한 열만 남기고 날짜/숫자 타입으로 변환합니다.
    :param path: 엑셀 파일 경로
    :param sheet: 시트 이름
    :param header_row: 실제 데이터 헤더가 시작되는 행(0-based)
    :return: 전처리된 DataFrame (columns=['구분','시장가격','기초지수'])
    """
    df = pd.read_excel(path, sheet_name=sheet, header=header_row)
    # 비어 있는 열 제거
    df = df.dropna(axis=1, how='all')
    # 핵심 열 선택
    df = df[['구분', '시장가격', '기초지수(WISE 화장품)']]
    # 타입 변환
    df['구분'] = pd.to_datetime(df['구분'])
    df['시장가격'] = pd.to_numeric(df['시장가격'], errors='coerce')
    df['기초지수(WISE 화장품)'] = pd.to_numeric(df['기초지수(WISE 화장품)'], errors='coerce')
    # 2023-01-01 이후 데이터만 필터
    df = df[df['구분'] >= '2023-01-01'].copy()
    # ETF close 컬럼 이름 정리
    df.rename(columns={'기초지수(WISE 화장품)': 'tiger화장품'}, inplace=True)
    # 인덱스를 날짜로 설정
    df.set_index('구분', inplace=True)
    return df

# —————————————————————————————————————————————
# 3. 데이터 병합 및 결측치 처리
# —————————————————————————————————————————————
def merge_and_fill(etf_df, market_df):
    """
    ETF 데이터와 글로벌 지수 데이터를 날짜 기준으로 병합,
    선형 보간 및 전후 평균(fill) 방식으로 결측치를 채웁니다.
    :param etf_df: ETF DataFrame (index=날짜)
    :param market_df: 지수 DataFrame (컬럼에 Date, S&P500, KOSPI)
    :return: 병합 및 결측치 처리된 DataFrame
    """
    market_df = market_df.copy()
    market_df['Date'] = pd.to_datetime(market_df['Date'])
    market_df.set_index('Date', inplace=True)

    # 병합
    df = etf_df.merge(
        market_df[['S&P500', 'KOSPI']],
        left_index=True, right_index=True,
        how='left'
    )

    # 1) 선형 보간
    df[['S&P500','KOSPI']] = df[['S&P500','KOSPI']].interpolate(method='linear')

    # 2) 전후 평균(fill) 방식 보완
    for col in ['S&P500','KOSPI']:
        df[col] = df[col].fillna((df[col].ffill() + df[col].bfill())/2)

    return df

# —————————————————————————————————————————————
# 4. 지표 계산 (상대수익률 및 이동평균)
# —————————————————————————————————————————————
def calculate_indicators(df):
    """
    누적 상대수익률과 5일 이동평균을 계산해 컬럼으로 추가합니다.
    :param df: 병합된 DataFrame
    :return: 지표 컬럼이 추가된 DataFrame
    """
    # 기준일 대비 누적 수익률
    df['tiger화장품_상대'] = df['tiger화장품'] / df['tiger화장품'].iloc[0]
    df['S&P500_상대'] = df['S&P500'] / df['S&P500'].iloc[0]
    df['KOSPI_상대']   = df['KOSPI'] / df['KOSPI'].iloc[0]

    # 5일 단순 이동평균
    for col in ['tiger화장품','S&P500','KOSPI']:
        df[f'{col}_SMA'] = df[col].rolling(window=5).mean()

    return df

# —————————————————————————————————————————————
# 5. 차트 그리기 & 저장
# —————————————————————————————————————————————
def plot_and_save(df, output_path, font_path):
    """
    한글 폰트 설정 후, 종가 추이와 누적 수익률 비교 차트를
    상/하단 서브플롯으로 생성하고 PNG로 저장합니다.
    :param df: 지표가 포함된 DataFrame (index=날짜)
    :param output_path: 저장할 파일 경로 (예: 'tiger_kbeauty.png')
    :param font_path: 시스템의 한글 폰트 파일 경로
    """
    # 한글 폰트 설정
    prop = fm.FontProperties(fname=font_path)
    plt.rcParams['font.family'] = prop.get_name()
    plt.rcParams['axes.unicode_minus'] = False

    fig, (ax1, ax2) = plt.subplots(
        2, 1, figsize=(12, 8), sharex=True
    )

    # 상단: 종가 추이
    ax1.plot(df.index, df['tiger화장품'], label='TIGER K-뷰티 ETF Close')
    ax1.plot(df.index, df['S&P500'], label='S&P 500 Close')
    ax1.plot(df.index, df['KOSPI'], label='KOSPI Close')
    ax1.set_ylabel('종가')
    ax1.set_title('2023년 이후 지수 및 ETF 종가 추이')
    ax1.legend()
    ax1.grid(True)

    # 하단: 누적 수익률
    ax2.plot(df.index, df['tiger화장품_상대'], label='ETF 누적 수익률')
    ax2.plot(df.index, df['S&P500_상대'], label='S&P500 누적 수익률')
    ax2.plot(df.index, df['KOSPI_상대'], label='KOSPI 누적 수익률')
    ax2.set_ylabel('누적 수익률')
    ax2.set_title('2023년 이후 상대 누적 수익률 비교')
    ax2.legend()
    ax2.grid(True)

    plt.xlabel('날짜')
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.savefig(output_path, dpi=400)
    plt.close(fig)


# —————————————————————————————————————————————
# 6. 실행부
# —————————————————————————————————————————————
def main():
    # 파라미터 설정
    start_date = "2023-01-02"
    end_date   = "2025-07-20"
    tickers    = ['^GSPC', '^KS11']
    excel_path = "out.xlsx"
    sheet_name = 'PRICE_DATA'
    font_path  = '/usr/share/fonts/truetype/nanum/NanumGothic.ttf' # matplotlib에서 한글 깨짐 방지
    output_png = 'tiger_kbeauty_comparison.png'

    # 세션 생성 (Chrome User-Agent 모방)
    session = requests.Session(impersonate="chrome") # yfinace에서 too many request error가 뜰시 해결하는 용도 

    # 1) 글로벌 지수 조회
    market_df = fetch_market_data(tickers, start_date, end_date, session)

    # 2) 로컬 엑셀 데이터 로드 및 전처리
    etf_df = load_and_preprocess_excel(excel_path, sheet_name) # ks.228790는 yfinace에서 조회 불가, https://www.tigeretf.com/ko/product/search/detail/index.do?ksdFund=KR7228790002에서 다운 로드 가능 ㅇ
    # 3) 데이터 병합 및 결측치 처리
    merged_df = merge_and_fill(etf_df, market_df)

    # 4) 지표 계산
    result_df = calculate_indicators(merged_df)

    # 5) 차트 그리기 및 저장
    plot_and_save(result_df, output_png, font_path)


if __name__ == "__main__":
    main()