"""
Universe acquisition, cached OHLCV downloads, momentum ranking, and cointegration pair discovery.

Requires: pandas, numpy, yfinance, tqdm, statsmodels; optional pyarrow for Parquet cache.

Run from repository root::

    .venv/bin/python universe_scanner.py
    .venv/bin/python universe_scanner.py --universe ndx100 --lookback-years 5
"""

from __future__ import annotations

import argparse
import itertools
import os
import time
import urllib.request
from io import StringIO
from pathlib import Path
from typing import Any, Literal

_USER_AGENT = (
    "Mozilla/5.0 (compatible; universe_scanner/1.0; +https://github.com/) "
    "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
)

import numpy as np
import pandas as pd

try:
    import yfinance as yf
except ImportError as e:  # pragma: no cover
    raise ImportError("yfinance is required: pip install yfinance") from e

try:
    from statsmodels.tsa.stattools import coint
except ImportError as e:  # pragma: no cover
    raise ImportError("statsmodels is required: pip install statsmodels") from e

try:
    from statsmodels.stats.multitest import multipletests
except ImportError as e:  # pragma: no cover
    raise ImportError("statsmodels is required for multiple testing correction: pip install statsmodels") from e

try:
    from tqdm import tqdm
except ImportError as e:  # pragma: no cover
    raise ImportError("tqdm is required: pip install tqdm") from e


# Default cache directory under this repo
_REPO_ROOT = Path(__file__).resolve().parent
DEFAULT_DATA_DIR = _REPO_ROOT / "data" / "universe_cache"


def _to_yahoo_symbol(symbol: str) -> str:
    """Wikipedia often uses BRK.B; Yahoo uses BRK-B."""
    return symbol.strip().upper().replace(".", "-")


def get_sp100_tickers(universe: Literal["sp100", "ndx100"] = "sp100") -> list[str]:
    """
    Scrape current S&P 100 or Nasdaq 100 constituents from Wikipedia.

    Parameters
    ----------
    universe
        ``"sp100"`` — S&P 100; ``"ndx100"`` — Nasdaq 100.
    """
    url = (
        "https://en.wikipedia.org/wiki/S%26P_100"
        if universe == "sp100"
        else "https://en.wikipedia.org/wiki/Nasdaq-100"
    )
    req = urllib.request.Request(url, headers={"User-Agent": _USER_AGENT})
    with urllib.request.urlopen(req, timeout=60) as resp:
        html = resp.read().decode("utf-8", errors="replace")
    try:
        tables = pd.read_html(StringIO(html))
    except ImportError as e:
        raise ImportError(
            "pandas.read_html needs an HTML parser. Install one of: pip install lxml html5lib beautifulsoup4"
        ) from e
    tickers: list[str] = []
    for tbl in tables:
        sym_col = None
        for c in tbl.columns:
            cl = str(c).strip().lower()
            if cl in ("symbol", "ticker", "symbols"):
                sym_col = c
                break
        if sym_col is None:
            continue
        raw = tbl[sym_col].dropna().astype(str)
        for s in raw:
            s = s.strip()
            if not s or s.lower() in ("ticker", "symbol", "company"):
                continue
            tickers.append(_to_yahoo_symbol(s.split("(")[0].strip()))
        if len(tickers) >= 50:
            break

    out = sorted(set(tickers))
    if len(out) < 20:
        raise RuntimeError(
            f"Parsed too few tickers ({len(out)}) from {url}. Wikipedia layout may have changed."
        )
    return out


def _cache_path(
    data_dir: Path,
    ticker: str,
    timeframe: str,
    lookback_years: float,
    use_parquet: bool,
) -> Path:
    safe = ticker.replace("/", "-").replace("\\", "-")
    ext = "parquet" if use_parquet else "csv"
    ytag = str(lookback_years).replace(".", "_")
    return data_dir / f"{safe}_{timeframe}_{ytag}y.{ext}"


def _read_ohlcv(path: Path) -> pd.DataFrame:
    if path.suffix.lower() == ".parquet":
        return pd.read_parquet(path)
    return pd.read_csv(path, index_col=0, parse_dates=True)


def _write_ohlcv(df: pd.DataFrame, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.suffix.lower() == ".parquet":
        df.to_parquet(path)
    else:
        df.to_csv(path)


def _cache_is_fresh(
    path: Path,
    *,
    max_age_hours: float = 24.0,
    max_data_stale_calendar_days: int = 7,
) -> bool:
    """
    Treat cache as fresh if file is recent enough OR last bar is not too old
    (weekends/holidays tolerated).
    """
    if not path.exists():
        return False
    age_sec = time.time() - path.stat().st_mtime
    if age_sec <= max_age_hours * 3600:
        return True
    try:
        df = _read_ohlcv(path)
    except Exception:
        return False
    if df.empty:
        return False
    last = pd.Timestamp(df.index.max()).normalize()
    today = pd.Timestamp.now().normalize()
    gap = (today - last).days
    return gap <= max_data_stale_calendar_days


def download_and_cache_data(
    tickers: list[str],
    *,
    timeframe: str = "1d",
    lookback_years: float = 5.0,
    data_dir: Path | str | None = None,
    prefer_parquet: bool = True,
    max_age_hours: float = 24.0,
    auto_adjust: bool = True,
) -> dict[str, pd.DataFrame]:
    """
    Download OHLCV via yfinance with disk cache under ``data_dir``.

    Cache hit: file exists and is considered fresh (recent mtime or recent last bar).

    Parameters
    ----------
    timeframe
        yfinance interval, e.g. ``\"1d\"``, ``\"1wk\"``.
    lookback_years
        Passed as ``period=f\"{int(ceil(years))}y\"`` when using period-based fetch.
    """
    data_dir = Path(data_dir or DEFAULT_DATA_DIR)
    data_dir.mkdir(parents=True, exist_ok=True)

    use_parquet = prefer_parquet
    if use_parquet:
        try:
            import pyarrow  # noqa: F401
        except ImportError:
            use_parquet = False

    years_int = max(1, int(np.ceil(lookback_years)))
    period = "max" if years_int >= 25 else f"{years_int}y"
    out: dict[str, pd.DataFrame] = {}

    for t in tqdm(tickers, desc="download/cache", unit="ticker"):
        path = _cache_path(data_dir, t, timeframe, lookback_years, use_parquet)
        if _cache_is_fresh(path, max_age_hours=max_age_hours):
            try:
                df = _read_ohlcv(path)
                if not df.empty and "Close" in df.columns:
                    out[t] = df
                    continue
            except Exception:
                pass

        try:
            ydf = yf.download(
                t,
                period=period,
                interval=timeframe,
                auto_adjust=auto_adjust,
                progress=False,
                threads=False,
            )
        except Exception:
            ydf = pd.DataFrame()

        if ydf is None or ydf.empty:
            continue
        if isinstance(ydf.columns, pd.MultiIndex):
            ydf.columns = ydf.columns.droplevel(1)
        ydf = ydf.sort_index()
        ydf = ydf[~ydf.index.duplicated(keep="last")]
        try:
            _write_ohlcv(ydf, path)
        except Exception:
            alt = path.with_suffix(".csv")
            _write_ohlcv(ydf, alt)
            path = alt
        out[t] = ydf

    return out


def rank_momentum(
    data_dict: dict[str, pd.DataFrame],
    *,
    lookback_days: int = 252,
    skip_recent_days: int = 21,
    price_col: str = "Close",
) -> pd.DataFrame:
    """
    Cross-sectional momentum using AQR "12-minus-1" methodology.

    For each ticker, compute:
        Return = Price_{t-21} / Price_{t-252} - 1

    This explicitly skips the most recent ``skip_recent_days`` to avoid short-term
    mean reversion.
    """
    rows: list[dict[str, Any]] = []
    for ticker, df in data_dict.items():
        if df is None or df.empty or price_col not in df.columns:
            continue
        px = df[price_col].dropna().astype(np.float64)
        if len(px) < 2:
            continue

        # Need enough observations to compute both shifts at the last timestamp.
        min_obs = max(int(lookback_days), int(skip_recent_days)) + 1
        if len(px) < min_obs:
            # IPO / insufficient history: drop from momentum ranking.
            continue

        # AQR momentum at time t (last available close):
        #   Return = close.shift(skip) / close.shift(lookback) - 1
        aqr_mom = px.shift(int(skip_recent_days)) / px.shift(int(lookback_days)) - 1.0
        r = float(aqr_mom.iloc[-1])
        if not np.isfinite(r):
            continue

        rows.append(
            {
                "ticker": ticker,
                "cum_return": r,
                "n_bars": len(px),
                "end_date": px.index.max(),
            }
        )

    out = pd.DataFrame(rows)
    if out.empty:
        return out
    out = out.sort_values("cum_return", ascending=False).reset_index(drop=True)
    out["rank"] = np.arange(1, len(out) + 1)
    return out


def calculate_half_life(spread: pd.Series, *, min_obs: int = 20) -> float:
    """
    Approximate OU half-life via linear regression:

        s_t = c + phi * s_{t-1} + ε_t

    For mean reversion, ``phi`` should satisfy ``0 < phi < 1`` and:

        hl = -ln(2) / ln(phi)

    Returns half-life in **bars/days** (depending on the input frequency).
    """
    s = spread.dropna().astype(np.float64)
    if len(s) < min_obs:
        return float("nan")

    s_lag = s.shift(1).iloc[1:]
    s_t = s.iloc[1:]
    if len(s_lag) < min_obs:
        return float("nan")

    y = s_t.to_numpy()
    x = s_lag.to_numpy()
    X = np.column_stack((np.ones_like(x), x))  # [1, s_{t-1}]
    try:
        coef, _, _, _ = np.linalg.lstsq(X, y, rcond=None)
    except Exception:
        return float("nan")

    phi = float(coef[1])
    if not np.isfinite(phi) or phi <= 0 or phi >= 1:
        # Non-mean-reverting (or numerical failure) -> treat as invalid.
        return float("inf")

    hl = -float(np.log(2.0)) / float(np.log(phi))
    if not np.isfinite(hl):
        return float("nan")
    return float(hl)


def calculate_hurst_exponent(spread: pd.Series, *, max_lags: int = 20) -> float:
    """
    Estimate Hurst exponent using variance of lagged differences:

        tau(lag) = std(s(t+lag) - s(t))
        log(tau) = H * log(lag) + const

    Returns H in approximately [0, 1], where < 0.5 suggests mean reversion.
    """
    s = spread.dropna().astype(np.float64)
    if len(s) < max_lags + 5:
        return float("nan")

    lags: list[int] = list(range(2, int(max_lags) + 1))
    tau: list[float] = []
    used_lags: list[int] = []

    arr = s.to_numpy()
    for lag in lags:
        if lag >= len(arr):
            break
        diff = arr[lag:] - arr[:-lag]
        if diff.size < 5:
            continue
        v = np.var(diff, ddof=1)
        if v <= 0 or not np.isfinite(v):
            continue
        tau.append(float(np.sqrt(v)))
        used_lags.append(lag)

    if len(tau) < 3:
        return float("nan")

    log_lags = np.log(np.asarray(used_lags, dtype=np.float64))
    log_tau = np.log(np.asarray(tau, dtype=np.float64))
    try:
        slope, _intercept = np.polyfit(log_lags, log_tau, 1)
    except Exception:
        return float("nan")
    return float(slope)


def _ols_hedge_spread_residual(
    log_y: pd.Series,
    log_x: pd.Series,
) -> pd.Series:
    """
    Hedge ratio via OLS: log_y = alpha + beta * log_x + eps
    Spread is the residual eps.
    """
    df = pd.concat([log_y, log_x], axis=1, join="inner").dropna()
    if df.empty:
        return pd.Series(dtype=np.float64)

    y = df.iloc[:, 0].to_numpy(dtype=np.float64)
    x = df.iloc[:, 1].to_numpy(dtype=np.float64)
    X = np.column_stack((np.ones_like(x), x))
    try:
        coef, _, _, _ = np.linalg.lstsq(X, y, rcond=None)
    except Exception:
        return pd.Series(dtype=np.float64)

    alpha = float(coef[0])
    beta = float(coef[1])
    resid = y - (alpha + beta * x)
    return pd.Series(resid, index=df.index, dtype=np.float64)


def find_cointegrated_pairs(
    data_dict: dict[str, pd.DataFrame],
    *,
    p_value_threshold: float = 0.05,
    price_col: str = "Close",
    min_obs: int = 60,
    hurst_max: float = 0.45,
    half_life_min: float = 1.0,
    half_life_max: float = 15.0,
    hurst_max_lags: int = 20,
    top_k: int = 0,
) -> pd.DataFrame:
    """
    Discovery of tradable cointegrated pairs with diversification.

    Pipeline:
      1) Raw Engle–Granger p-values for all pairs
      2) Holm-Bonferroni multiple testing correction
      3) Tradability filter on the residual spread:
         - Hurst < ``hurst_max``
         - ``half_life_min`` < half-life < ``half_life_max``
      4) Greedy cluster-risk reduction: select pairs sorted by lowest adjusted
         p-value such that no ticker appears in more than one selected pair.
    """
    tickers = [t for t, df in data_dict.items() if df is not None and price_col in df.columns]
    pairs = list(itertools.combinations(sorted(tickers), 2))

    # Step A: raw p-values for all pairs (store even if they fail thresholds later).
    results: list[dict[str, Any]] = []
    for a, b in tqdm(pairs, desc="coint pairs", unit="pair"):
        da = data_dict[a][price_col].dropna()
        db = data_dict[b][price_col].dropna()
        aligned = pd.concat([da, db], axis=1, join="inner").dropna()
        aligned.columns = ["a", "b"]
        if len(aligned) < min_obs:
            continue

        # Cointegration is tested on log prices.
        la = np.log(aligned["a"].astype(np.float64))
        lb = np.log(aligned["b"].astype(np.float64))

        try:
            coint_t, pvalue, _crit = coint(la, lb, trend="c", autolag="AIC")
        except Exception:
            continue

        results.append(
            {
                "Ticker_A": a,
                "Ticker_B": b,
                "p_value": float(pvalue),
                "n_obs": int(len(aligned)),
                "coint_t": float(coint_t),
            }
        )

    if not results:
        return pd.DataFrame(
            columns=["Ticker_A", "Ticker_B", "Adjusted_P_Value", "Hurst", "Half_Life"]
        )

    # Step B: Holm-Bonferroni multiple testing correction.
    pvals = np.asarray([r["p_value"] for r in results], dtype=np.float64)
    try:
        _reject, pvals_adj, _alphacSidak, _alphacBonf = multipletests(
            pvals, alpha=float(p_value_threshold), method="holm"
        )
    except Exception:
        # If correction fails for some reason, return empty rather than wrong results.
        return pd.DataFrame(
            columns=["Ticker_A", "Ticker_B", "Adjusted_P_Value", "Hurst", "Half_Life"]
        )

    for r, padj in zip(results, pvals_adj):
        r["Adjusted_P_Value"] = float(padj)

    survived = [r for r in results if np.isfinite(r["Adjusted_P_Value"]) and r["Adjusted_P_Value"] < float(p_value_threshold)]
    if not survived:
        return pd.DataFrame(
            columns=["Ticker_A", "Ticker_B", "Adjusted_P_Value", "Hurst", "Half_Life"]
        )

    # Step C: tradability filter on the residual spread.
    tradable: list[dict[str, Any]] = []
    for r in tqdm(survived, desc="tradability filters", unit="pair"):
        a = str(r["Ticker_A"])
        b = str(r["Ticker_B"])
        da = data_dict[a][price_col].dropna()
        db = data_dict[b][price_col].dropna()
        aligned = pd.concat([da, db], axis=1, join="inner").dropna()
        aligned.columns = ["a", "b"]
        if len(aligned) < min_obs:
            continue

        log_a = np.log(aligned["a"].astype(np.float64))
        log_b = np.log(aligned["b"].astype(np.float64))

        spread = _ols_hedge_spread_residual(log_a, log_b)
        if spread.empty or spread.size < min_obs:
            continue

        hurst = calculate_hurst_exponent(spread, max_lags=hurst_max_lags)
        half_life = calculate_half_life(spread)

        if not np.isfinite(hurst):
            continue
        if not np.isfinite(half_life):
            continue

        if hurst < float(hurst_max) and float(half_life_min) < float(half_life) < float(half_life_max):
            tradable.append(
                {
                    "Ticker_A": a,
                    "Ticker_B": b,
                    "Adjusted_P_Value": float(r["Adjusted_P_Value"]),
                    "Hurst": float(hurst),
                    "Half_Life": float(half_life),
                }
            )

    if not tradable:
        return pd.DataFrame(
            columns=["Ticker_A", "Ticker_B", "Adjusted_P_Value", "Hurst", "Half_Life"]
        )

    # Step 4: Greedy selection (cluster-risk reduction).
    tradable_sorted = sorted(tradable, key=lambda x: x["Adjusted_P_Value"])
    final_portfolio: list[dict[str, Any]] = []
    used_tickers: set[str] = set()

    for r in tradable_sorted:
        a = str(r["Ticker_A"])
        b = str(r["Ticker_B"])
        if a in used_tickers or b in used_tickers:
            continue
        final_portfolio.append(r)
        used_tickers.add(a)
        used_tickers.add(b)
        if top_k and len(final_portfolio) >= int(top_k):
            break

    out = pd.DataFrame(final_portfolio)
    if out.empty:
        return pd.DataFrame(
            columns=["Ticker_A", "Ticker_B", "Adjusted_P_Value", "Hurst", "Half_Life"]
        )
    return out.sort_values("Adjusted_P_Value", ascending=True).reset_index(drop=True)


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--universe", choices=("sp100", "ndx100"), default="sp100")
    p.add_argument("--timeframe", default="1d")
    p.add_argument("--lookback-years", type=float, default=5.0)
    # AQR 12-minus-1 defaults:
    #   lookback_days = 252 (~12 months), skip_recent_days = 21 (~1 month)
    p.add_argument("--momentum-days", type=int, default=252, help="AQR momentum lookback (t-252).")
    p.add_argument("--momentum-skip", type=int, default=21, help="AQR momentum skip (t-21).")
    p.add_argument("--coint-p", type=float, default=0.05)
    p.add_argument(
        "--hurst-max",
        type=float,
        default=0.45,
        help="Tradability filter: keep pairs with Hurst < this value.",
    )
    p.add_argument(
        "--half-life-min",
        type=float,
        default=1.0,
        help="Tradability filter: keep pairs with half-life > this value (bars/days).",
    )
    p.add_argument(
        "--half-life-max",
        type=float,
        default=15.0,
        help="Tradability filter: keep pairs with half-life < this value (bars/days).",
    )
    p.add_argument(
        "--hurst-max-lags",
        type=int,
        default=20,
        help="Hurst estimation uses variance of lagged differences up to these lags.",
    )
    p.add_argument("--data-dir", type=str, default=str(DEFAULT_DATA_DIR))
    p.add_argument(
        "--max-tickers",
        type=int,
        default=0,
        help="If >0, only use the first N tickers (smoke test / faster local runs)",
    )
    args = p.parse_args()

    os.chdir(_REPO_ROOT)

    print("Fetching universe from Wikipedia …")
    tickers = get_sp100_tickers(universe=args.universe)
    if args.max_tickers and args.max_tickers > 0:
        tickers = tickers[: args.max_tickers]
    print(f"Universe ({args.universe}): {len(tickers)} tickers")

    print("Downloading / loading cache …")
    data_dict = download_and_cache_data(
        tickers,
        timeframe=args.timeframe,
        lookback_years=args.lookback_years,
        data_dir=Path(args.data_dir),
    )
    print(f"Loaded OHLCV for {len(data_dict)} names")

    mom = rank_momentum(
        data_dict,
        lookback_days=args.momentum_days,
        skip_recent_days=args.momentum_skip,
    )
    print("\n=== Top 10 AQR momentum (12-minus-1) ===")
    if mom.empty:
        print("(no data)")
    else:
        print(mom.head(10).to_string(index=False))

    pairs = find_cointegrated_pairs(
        data_dict,
        p_value_threshold=args.coint_p,
        hurst_max=args.hurst_max,
        half_life_min=args.half_life_min,
        half_life_max=args.half_life_max,
        hurst_max_lags=args.hurst_max_lags,
    )
    print("\n=== Top 10 cointegrated pairs (lowest p-value) ===")
    if pairs.empty:
        print("(no pairs under threshold — try looser --coint-p or longer history)")
    else:
        print(pairs.head(10).to_string(index=False))


if __name__ == "__main__":
    main()