"""
Historical market data loading and normalization for a multi-timeframe system.

- Daily bars: macro momentum filter (e.g. SPY/QQQ regime).
- Intraday bars (1h / 15m): micro mean-reversion / stat-arb style signals.

Dependencies: pandas, numpy, yfinance (install: pip install yfinance).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import numpy as np
import pandas as pd

try:
    import yfinance as yf
except ImportError as e:  # pragma: no cover
    yf = None  # type: ignore[assignment]
    _YF_IMPORT_ERROR = e
else:
    _YF_IMPORT_ERROR = None

# yfinance intraday history is vendor-limited (typically ~60 days for sub-daily).
IntradayInterval = Literal["15m", "30m", "60m", "1h", "1d"]


def _require_yfinance() -> None:
    if yf is None:
        raise ImportError(
            "yfinance is required. Install with: python -m pip install yfinance\n"
            f"Original error: {_YF_IMPORT_ERROR}"
        )


def _standardize_ohlcv(df: pd.DataFrame, *, prefer_adjusted: bool = True) -> pd.DataFrame:
    """
    Normalize yfinance output to lowercase OHLCV columns with a DatetimeIndex.

    Handles both single-ticker (flat columns) and occasional MultiIndex columns.
    """
    if df.empty:
        return pd.DataFrame()

    out = df.copy()
    # Single ticker: columns like Open, High, Low, Close, Adj Close, Volume
    if isinstance(out.columns, pd.MultiIndex):
        # Take first level if multiple tickers were requested accidentally
        if out.columns.nlevels > 1:
            out.columns = out.columns.droplevel(1) if out.columns.nlevels == 2 else out.columns.get_level_values(0)

    rename_map = {
        "Open": "open",
        "High": "high",
        "Low": "low",
        "Close": "close",
        "Adj Close": "adj_close",
        "Volume": "volume",
    }
    out = out.rename(columns={k: v for k, v in rename_map.items() if k in out.columns})

    if prefer_adjusted and "adj_close" in out.columns:
        out["close"] = out["adj_close"]

    required = ["open", "high", "low", "close", "volume"]
    missing = [c for c in required if c not in out.columns]
    if missing:
        raise ValueError(f"Expected OHLCV columns missing after normalize: {missing}. Got: {list(out.columns)}")

    out = out[required].astype(np.float64)
    out = out.sort_index()
    out = out[~out.index.duplicated(keep="last")]
    return out


@dataclass
class DataLoader:
    """
    Fetch and clean OHLCV history from Yahoo Finance.

    Example
    -------
    >>> loader = DataLoader()
    >>> daily = loader.fetch_daily("SPY", period="5y")
    >>> hourly = loader.fetch_intraday("AAPL", interval="1h", period="60d")
    """

    auto_adjust: bool = False
    """If True, pass auto_adjust=True to yfinance (we still prefer explicit adj_close when available)."""

    def fetch_daily(
        self,
        ticker: str,
        *,
        start: str | pd.Timestamp | None = None,
        end: str | pd.Timestamp | None = None,
        period: str | None = None,
    ) -> pd.DataFrame:
        """
        Load daily OHLCV for momentum / regime models.

        Provide either (start & end) or period (e.g. '5y', 'max'), not both if yfinance conflicts.
        """
        _require_yfinance()
        t = yf.Ticker(ticker)
        if period:
            raw = t.history(period=period, interval="1d", auto_adjust=self.auto_adjust)
        else:
            raw = t.history(start=start, end=end, interval="1d", auto_adjust=self.auto_adjust)
        out = _standardize_ohlcv(raw)
        out.index = pd.to_datetime(out.index).tz_localize(None)
        return out

    def fetch_intraday(
        self,
        ticker: str,
        *,
        interval: IntradayInterval = "1h",
        start: str | pd.Timestamp | None = None,
        end: str | pd.Timestamp | None = None,
        period: str | None = None,
    ) -> pd.DataFrame:
        """
        Load intraday OHLCV for mean-reversion / rolling z-score engines.

        Note: Yahoo limits intraday depth (often ~730 days for hourly; shorter for 15m).
        Prefer `period='60d'` or similar when using 15m bars.
        """
        _require_yfinance()
        # yfinance accepts '60m' for hourly in some versions; map 1h -> 60m for max compatibility
        iv = "60m" if interval in ("1h", "60m") else interval
        t = yf.Ticker(ticker)
        if period:
            raw = t.history(period=period, interval=iv, auto_adjust=self.auto_adjust)
        else:
            raw = t.history(start=start, end=end, interval=iv, auto_adjust=self.auto_adjust)
        out = _standardize_ohlcv(raw)
        out.index = pd.to_datetime(out.index)
        if out.index.tz is not None:
            out.index = out.index.tz_convert("UTC").tz_localize(None)
        return out

    def align_to_trading_days(
        self,
        intraday: pd.DataFrame,
        *,
        day_tz: str | None = "America/New_York",
    ) -> pd.DataFrame:
        """
        Optional: assign a calendar date column for merging intraday bars with daily regime.

        Intraday index is timestamp; we add `trade_date` as the session date in `day_tz`.
        """
        if intraday.empty:
            return intraday
        idx = intraday.index
        if day_tz:
            ts = idx.tz_localize("UTC") if idx.tz is None else idx.tz_convert("UTC")
            ts = ts.tz_convert(day_tz)
        else:
            ts = idx
        out = intraday.copy()
        out["trade_date"] = pd.to_datetime(ts.normalize().date)
        return out


if __name__ == "__main__":
    # Smoke test (requires network)
    _require_yfinance()
    dl = DataLoader()
    d = dl.fetch_daily("SPY", period="6mo")
    print(d.tail())
    print("rows", len(d))