src/data/hf_loader.py

import os
import urllib.request
import pandas as pd
import numpy as np
from typing import List, Optional, Tuple

class HuggingFaceETFLoader:
    """
    Data Loader for ETF price data (Hugging Face / yfinance fallback).
    Downloads and caches parquet files locally in `data/raw/` for offline reproducibility.
    """

    DEFAULT_UNIVERSE = [
        "SPY", "QQQ", "IWM", "MDY",
        "XLK", "XLF", "XLE", "XLV", "XLY", "XLP", "XLI", "XLB", "XLU", "XLC", "XLRE",
        "SMH", "XBI", "KRE", "ITB",
        "TLT", "IEF", "SHY", "LQD", "HYG", "TIP",
        "GLD", "SLV", "USO", "DBA",
        "EEM", "EFA", "FXI", "EWJ",
        "MTUM", "USMV", "QUAL", "IWD", "IWF"
    ]

    def __init__(self, cache_dir: str = "data/raw"):
        self.cache_dir = cache_dir
        os.makedirs(self.cache_dir, exist_ok=True)

    def load_prices(self, tickers: Optional[List[str]] = None) -> pd.DataFrame:
        """Load ETF prices DataFrame. Uses cached parquet if available, otherwise downloads via yfinance."""
        local_path = os.path.join(self.cache_dir, "hf_prices.parquet")
        
        if os.path.exists(local_path):
            print(f"[ETFLoader] Reading cached dataset from {local_path}...")
            df = pd.read_parquet(local_path)
            df["date"] = pd.to_datetime(df["date"])
            return df
        
        return self._fetch_yfinance_prices(tickers)

    def _fetch_yfinance_prices(self, tickers: Optional[List[str]] = None) -> pd.DataFrame:
        """Fetch daily OHLCV dataset using yfinance and save to local parquet cache."""
        import yfinance as yf
        tickers = tickers or self.DEFAULT_UNIVERSE
        print(f"[ETFLoader] Fetching {len(tickers)} tickers via yfinance...")
        
        try:
            data = yf.download(tickers, start="2014-01-01", end="2024-01-01", group_by="ticker", auto_adjust=False, progress=False)
            records = []
            
            for t in tickers:
                if t in data.columns.levels[0]:
                    sub = data[t].dropna(how="all").reset_index()
                    if len(sub) > 0:
                        sub["symbol"] = t
                        sub.rename(columns={
                            "Date": "date", "Open": "open", "High": "high", "Low": "low",
                            "Close": "close", "Adj Close": "adjusted_close", "Volume": "volume"
                        }, inplace=True)
                        records.append(sub)
            
            if records:
                df_res = pd.concat(records, ignore_index=True)
                df_res["date"] = pd.to_datetime(df_res["date"])
                if "adjusted_close" not in df_res.columns:
                    df_res["adjusted_close"] = df_res["close"]
                
                # Save cache
                df_res.to_parquet(os.path.join(self.cache_dir, "hf_prices.parquet"))
                print(f"[ETFLoader] Cached {len(df_res)} price rows to {self.cache_dir}/hf_prices.parquet")
                return df_res
        except Exception as e:
            print(f"[ETFLoader] Remote download failed: {e}")

        # Fallback to synthetic offline dataset generator
        return self.generate_synthetic_prices(tickers)

    @staticmethod
    def generate_synthetic_prices(tickers: List[str], num_days: int = 1000) -> pd.DataFrame:
        """Generate realistic synthetic OHLCV data for offline CI / local smoke tests."""
        print("[ETFLoader] Generating synthetic price data for offline testing...")
        dates = pd.date_range("2020-01-01", periods=num_days, freq="B")
        records = []
        np.random.seed(42)
        
        for t in tickers:
            price = 100.0 + np.random.uniform(-10, 10)
            volatility = 0.015 + np.random.uniform(0.005, 0.01)
            prices = [price]
            
            for _ in range(num_days - 1):
                ret = np.random.normal(0.0004, volatility)
                price *= np.exp(ret)
                prices.append(price)
            
            df_t = pd.DataFrame({
                "date": dates,
                "symbol": t,
                "open": prices,
                "high": [p * (1 + abs(np.random.normal(0, 0.005))) for p in prices],
                "low": [p * (1 - abs(np.random.normal(0, 0.005))) for p in prices],
                "close": prices,
                "adjusted_close": prices,
                "volume": np.random.uniform(1e6, 1e7, size=num_days)
            })
            records.append(df_t)
            
        return pd.concat(records, ignore_index=True)