fetch_data.py

"""
fetch_data.py
==============
Data acquisition pipeline for the Regime-Aware Multi-Factor Strategy.

This script transparently records how raw input data are obtained:
  1. S&P 500 constituent list from Wikipedia (current composition)
  2. Monthly adjusted close prices from Yahoo Finance via yfinance
  3. Fama-French 5-factor data from Kenneth R. French's Data Library

Usage:
    python fetch_data.py

Outputs (all written to ./data/):
    - sp500_monthly_prices.xlsx
    - F-F_Research_Data_5_Factors_2x3.csv
    - SP500_Benchmark_Returns.csv  (derived from ^GSPC)

Note on survivorship bias:
    The fetched S&P 500 list reflects the *current* composition. Stocks that
    were dropped from the index during the backtest period are not included.
    For a production-grade backtest, use CRSP/Compustat historical constituent
    data. This script is provided for reproducibility demonstration.
"""

from __future__ import annotations

import os
import re
import zipfile
from io import BytesIO

import pandas as pd
import requests
import yfinance as yf
from tqdm import tqdm


DATA_DIR = "data"
SP500_URL = "https://en.wikipedia.org/wiki/List_of_S%26P_500_companies"
FF5_URL = (
    "https://mba.tuck.dartmouth.edu/pages/faculty/ken.french/"
    "ftp/F-F_Research_Data_5_Factors_2x3_CSV.zip"
)

START_DATE = "2010-01-01"
END_DATE = "2025-12-31"


def ensure_data_dir() -> None:
    os.makedirs(DATA_DIR, exist_ok=True)


def fetch_sp500_tickers() -> list[str]:
    """Pull current S&P 500 constituents from Wikipedia."""
    print("[1/3] Fetching S&P 500 constituent list from Wikipedia...")
    tables = pd.read_html(SP500_URL)
    tickers = tables[0]["Symbol"].tolist()
    # Yahoo Finance uses hyphens for class-B shares (e.g. BRK.B -> BRK-B)
    tickers = [t.replace(".", "-") for t in tickers]
    print(f"        Found {len(tickers)} tickers.")
    return tickers


def download_prices(tickers: list[str], batch_size: int = 50) -> pd.DataFrame:
    """Download monthly adjusted-close prices from Yahoo Finance in batches."""
    print(f"[2/3] Downloading prices from Yahoo Finance ({START_DATE} ~ {END_DATE})...")
    print(f"        Processing in batches of {batch_size} to avoid rate limits.")

    all_frames: list[pd.DataFrame] = []
    for i in range(0, len(tickers), batch_size):
        batch = tickers[i : i + batch_size]
        print(f"        Batch {i // batch_size + 1}/{(len(tickers) - 1) // batch_size + 1}: {batch[0]} ... {batch[-1]}")
        try:
            data = yf.download(
                batch,
                start=START_DATE,
                end=END_DATE,
                interval="1d",
                auto_adjust=True,
                progress=False,
                threads=True,
            )
            if data.empty:
                continue
            # Extract adjusted close; handle single-ticker case
            if isinstance(data.columns, pd.MultiIndex):
                adj_close = data["Adj Close"]
            else:
                adj_close = data.to_frame(name=batch[0]) if len(batch) == 1 else data["Adj Close"]
            all_frames.append(adj_close)
        except Exception as exc:
            print(f"        [WARN] Batch failed: {exc}")
            continue

    if not all_frames:
        raise RuntimeError("No price data could be downloaded.")

    combined = pd.concat(all_frames, axis=1)
    # Resample to month-end, then reset index to month-start for consistency
    monthly = combined.resample("ME").last()
    monthly.index = monthly.index + pd.offsets.MonthBegin(1)
    monthly.index.name = "Date"
    monthly = monthly.sort_index()
    return monthly


def save_sp500_prices(df: pd.DataFrame) -> str:
    """Save monthly prices to Excel, matching the format expected by the strategy."""
    path = os.path.join(DATA_DIR, "sp500_monthly_prices.xlsx")
    df.reset_index().to_excel(path, index=False, engine="openpyxl")
    print(f"        Saved -> {path}  ({df.shape[0]} rows x {df.shape[1]} cols)")
    return path


def download_ff5_factors() -> str:
    """Download FF5 monthly factors from Kenneth French's data library."""
    print("[3/3] Downloading Fama-French 5-Factor data...")
    r = requests.get(FF5_URL, timeout=30)
    r.raise_for_status()

    with zipfile.ZipFile(BytesIO(r.content)) as zf:
        # The CSV inside is named F-F_Research_Data_5_Factors_2x3.csv
        csv_name = [n for n in zf.namelist() if n.endswith(".csv")][0]
        with zf.open(csv_name) as f:
            lines = f.read().decode("utf-8").splitlines()

    # Find the header row (contains "Mkt-RF" and "SMB")
    header_idx = None
    for idx, line in enumerate(lines):
        if "Mkt-RF" in line and "SMB" in line:
            header_idx = idx
            break

    if header_idx is None:
        raise ValueError("Could not locate FF5 header row in downloaded CSV.")

    # Write raw CSV preserving the original format (comments + data)
    raw_path = os.path.join(DATA_DIR, "F-F_Research_Data_5_Factors_2x3.csv")
    with open(raw_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))
    print(f"        Saved -> {raw_path}")
    return raw_path


def derive_benchmark_from_prices(prices_df: pd.DataFrame) -> None:
    """Derive S&P 500 benchmark monthly returns from ^GSPC if available, else skip."""
    print("[Extra] Deriving benchmark returns from ^GSPC...")
    try:
        gs = yf.download("^GSPC", start=START_DATE, end=END_DATE, interval="1d", progress=False)
        if gs.empty:
            print("        [SKIP] ^GSPC data not available.")
            return
        adj = gs["Adj Close"] if isinstance(gs.columns, pd.MultiIndex) else gs
        monthly = adj.resample("ME").last()
        monthly.index = monthly.index + pd.offsets.MonthBegin(1)
        returns = monthly.pct_change().dropna()
        returns.name = "SP500_Monthly_Return"
        bench = returns.reset_index()
        bench.columns = ["Date", "SP500_Monthly_Return"]
        path = os.path.join(DATA_DIR, "SP500_Benchmark_Returns.csv")
        bench.to_csv(path, index=False)
        print(f"        Saved -> {path}")
    except Exception as exc:
        print(f"        [SKIP] Could not derive benchmark: {exc}")


def main() -> None:
    ensure_data_dir()

    tickers = fetch_sp500_tickers()
    prices = download_prices(tickers)
    save_sp500_prices(prices)

    download_ff5_factors()
    derive_benchmark_from_prices(prices)

    print("\n[Done] All raw data have been written to ./data/")
    print("       You can now run:  python main_update.py")


if __name__ == "__main__":
    main()