# =============================================================================
#  Project   : anomaly-atlas
#  File      : src/anomaly_atlas/stats/spa.py
#  Purpose   : White Reality Check & Hansen SPA over a rule-return matrix
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Data src  : hfmarketdata.io (sole data source)
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64)
#  License   : All rights reserved (research code)
# =============================================================================
"""Data-snooping corrections over a searched universe of rules.

Inputs are a (T days × N rules) matrix of rule returns. H0: no rule has
positive expected return — max_k E[f_k] <= 0.

* White (2000) Reality Check: max-statistic over the centered stationary
  bootstrap (Politis-Romano 1994).
* Hansen (2005) SPA: studentized statistic with the recentering threshold,
  less sensitive to poor/irrelevant rules in the universe.

Validated on synthetic ground truth (§8.1 gate: pure noise must not
survive; a planted profitable rule must).
"""

from __future__ import annotations

import numpy as np


def stationary_bootstrap_indices(
    n: int, mean_block: float, n_boot: int, seed: int = 42
) -> np.ndarray:
    """(n_boot, n) index matrix from the Politis-Romano stationary bootstrap.

    Geometric block lengths with mean `mean_block`, circular wrapping —
    resamples preserve short-range dependence in expectation.
    """
    rng = np.random.default_rng(seed)
    p = 1.0 / mean_block
    idx = np.empty((n_boot, n), dtype=np.int64)
    for b in range(n_boot):
        t = 0
        while t < n:
            start = rng.integers(0, n)
            length = min(int(rng.geometric(p)), n - t)
            idx[b, t : t + length] = (start + np.arange(length)) % n
            t += length
    return idx


def reality_check(
    x: np.ndarray, n_boot: int = 500, mean_block: float = 5.0, seed: int = 42
) -> dict:
    """White's Reality Check p-value for max_k mean(x_k) > 0.

    x: (T, N) rule-return matrix (NaN rows dropped listwise).
    """
    x = np.asarray(x, dtype=float)
    x = x[np.isfinite(x).all(axis=1)]
    t_len, n_rules = x.shape
    if t_len < 30 or n_rules == 0:
        return {"p": float("nan"), "best_rule": None, "v_stat": float("nan")}
    means = x.mean(axis=0)
    v = np.sqrt(t_len) * means.max()
    idx = stationary_bootstrap_indices(t_len, mean_block, n_boot, seed)
    centered = x - means  # White: bootstrap distribution of centered means
    v_boot = np.empty(n_boot)
    for b in range(n_boot):
        v_boot[b] = np.sqrt(t_len) * centered[idx[b]].mean(axis=0).max()
    p = float((np.sum(v_boot >= v) + 1) / (n_boot + 1))
    return {"p": p, "best_rule": int(means.argmax()), "v_stat": float(v),
            "best_mean_daily": float(means.max())}


def spa_test(
    x: np.ndarray, n_boot: int = 500, mean_block: float = 5.0, seed: int = 42
) -> dict:
    """Hansen's SPA p-value (consistent variant) for max_k mean(x_k) > 0."""
    x = np.asarray(x, dtype=float)
    x = x[np.isfinite(x).all(axis=1)]
    t_len, n_rules = x.shape
    if t_len < 30 or n_rules == 0:
        return {"p": float("nan"), "best_rule": None}
    means = x.mean(axis=0)
    idx = stationary_bootstrap_indices(t_len, mean_block, n_boot, seed)
    boot_means = np.empty((n_boot, n_rules))
    for b in range(n_boot):
        boot_means[b] = x[idx[b]].mean(axis=0)
    omega = np.sqrt(t_len) * boot_means.std(axis=0, ddof=1)
    omega = np.maximum(omega, 1e-12)
    t_stat = float((np.sqrt(t_len) * means / omega).max())
    # Hansen recentering: rules with sufficiently negative means contribute 0
    thresh = -omega / np.sqrt(t_len) * np.sqrt(2.0 * np.log(np.log(max(t_len, 3))))
    center = np.where(means >= thresh, means, 0.0)
    t_boot = np.empty(n_boot)
    for b in range(n_boot):
        z = np.sqrt(t_len) * (boot_means[b] - center) / omega
        t_boot[b] = max(z.max(), 0.0)
    p = float((np.sum(t_boot >= max(t_stat, 0.0)) + 1) / (n_boot + 1))
    rule_t = np.sqrt(t_len) * means / omega
    t95 = float(np.percentile(t_boot, 95))
    return {"p": p, "best_rule": int(rule_t.argmax()), "t_stat": t_stat,
            "rule_t": rule_t.tolist(), "t95": t95,
            "n_step1_survivors": int((rule_t >= t95).sum()) if np.isfinite(t95) else 0}


def deflated_sharpe(
    sr: float, t_len: int, skew: float, kurt: float,
    n_trials: int, sr_variance: float,
) -> dict:
    """Bailey & López de Prado (2014) Deflated Sharpe Ratio.

    `sr` is the per-period (e.g. daily) Sharpe of the BEST rule; `sr_variance`
    the variance of Sharpe estimates across the searched universe; `kurt` is
    Pearson kurtosis (normal = 3). Returns the expected max Sharpe under
    pure selection (`sr0`) and DSR = P[true SR > 0 | selection].
    """
    from math import sqrt

    from scipy.stats import norm

    if n_trials < 2 or sr_variance <= 0 or t_len < 10:
        return {"sr0": float("nan"), "dsr": float("nan")}
    gamma = 0.5772156649015329
    z1 = norm.ppf(1.0 - 1.0 / n_trials)
    z2 = norm.ppf(1.0 - 1.0 / (n_trials * np.e))
    sr0 = sqrt(sr_variance) * ((1.0 - gamma) * z1 + gamma * z2)
    denom = sqrt(max(1.0 - skew * sr + (kurt - 1.0) / 4.0 * sr**2, 1e-12))
    dsr = float(norm.cdf((sr - sr0) * sqrt(t_len - 1.0) / denom))
    return {"sr0": float(sr0), "dsr": dsr}