"""
Alpha against a benchmark (QQQ by default): the yardstick for "does this
strategy add anything a Nasdaq-100 index fund does not?"

    from research.lab.alpha import bench_returns, alpha_stats
    b = bench_returns(r.index)                  # QQQ total return, daily
    s = alpha_stats(r, b)                       # alpha, beta, t, IR, ...

alpha/beta come from the daily regression of excess returns (over the
T-bill) on the benchmark's excess return; the alpha t-stat uses Newey-West
standard errors (10 lags). IR is the annualized mean over the standard
deviation of the active return r - b. QQQ comes from
data/research/long_history.pkl (research/lab/fetch_long_history.py).
"""

from __future__ import annotations

import os
import pickle

import numpy as np
import pandas as pd

from research.lab.levered_book import tbill

LONG_HISTORY = os.path.join("data", "research", "long_history.pkl")


def bench_returns(index: pd.DatetimeIndex, ticker: str = "QQQ") -> pd.Series:
    with open(LONG_HISTORY, "rb") as fh:
        px = pickle.load(fh)["close"][ticker].dropna()
    return px.reindex(index.union(px.index)).ffill().pct_change(fill_method=None).reindex(index)


def _nw_se(x: np.ndarray, e: np.ndarray, lags: int = 10) -> np.ndarray:
    """Newey-West standard errors of OLS coefficients."""
    n = len(e)
    xtx_inv = np.linalg.inv(x.T @ x)
    u = x * e[:, None]
    s = u.T @ u
    for l in range(1, lags + 1):
        w = 1 - l / (lags + 1)
        g = u[l:].T @ u[:-l]
        s += w * (g + g.T)
    return np.sqrt(np.diag(xtx_inv @ s @ xtx_inv))


def alpha_stats(r: pd.Series, b: pd.Series, rf: pd.Series | None = None) -> dict:
    df = pd.DataFrame({"r": r, "b": b}).dropna()
    if rf is None:
        rf = tbill(df.index)
    d = rf.reindex(df.index).fillna(0.0) / 252
    y, xb = (df["r"] - d).to_numpy(), (df["b"] - d).to_numpy()
    X = np.column_stack([np.ones(len(y)), xb])
    coef, *_ = np.linalg.lstsq(X, y, rcond=None)
    se = _nw_se(X, y - X @ coef)
    active = df["r"] - df["b"]
    yrs = len(df) / 252
    cagr = lambda x: (1 + x).prod() ** (1 / yrs) - 1          # noqa: E731
    yr_r = (1 + df["r"]).groupby(df.index.year).prod()
    yr_b = (1 + df["b"]).groupby(df.index.year).prod()
    return {"alpha": coef[0] * 252, "t": coef[0] / se[0], "beta": coef[1],
            "ir": active.mean() / active.std() * np.sqrt(252),
            "cagr": cagr(df["r"]), "bench_cagr": cagr(df["b"]),
            "yrs_beat": float((yr_r > yr_b).mean())}


def fmt_alpha(name: str, s: dict) -> str:
    return (f"{name:44s} alpha {s['alpha']*100:6.1f}%/yr t {s['t']:5.2f} beta {s['beta']:4.2f} "
            f"IR {s['ir']:5.2f} cagr {s['cagr']*100:5.1f}% vs {s['bench_cagr']*100:5.1f}% "
            f"yrs beat {s['yrs_beat']*100:3.0f}%")