"""M01 教学图：SPY 日收益肥尾，以及 AAPL 对 SPY 的一元回归。

数据是动量 Plus 实验室的复权收盘价种子，不是新的回测业绩。
"""

from __future__ import annotations

import math
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False

ROOT = Path(__file__).resolve().parents[3]
SEED = ROOT / "Labs" / "momentum_plus_lab" / "seed_data"
OUT = Path(__file__).resolve().parent / "spy_tails_and_capm.png"
# 回归用常数年化 4% 扣成日超额。直方图仍用总收益，两种口径在正文里分开。
RF_ANNUAL = 0.04


def norm_cdf(x: float) -> float:
    return 0.5 * (1.0 + math.erf(x / math.sqrt(2.0)))


def norm_ppf(p: float) -> float:
    """二分法求标准正态分位数，避免依赖 scipy。"""
    lo, hi = -12.0, 12.0
    for _ in range(80):
        mid = 0.5 * (lo + hi)
        if norm_cdf(mid) < p:
            lo = mid
        else:
            hi = mid
    return 0.5 * (lo + hi)


def load_adj(name: str) -> pd.Series:
    path = SEED / "spy.csv" if name == "SPY" else SEED / "_cache" / f"{name}.csv"
    frame = pd.read_csv(path, parse_dates=["session"])
    series = frame.set_index("session")["adj_close"].astype(float).sort_index()
    series.name = name
    return series


def ols(y: np.ndarray, x: np.ndarray) -> dict[str, float]:
    design = np.column_stack([np.ones(len(x)), x])
    beta_hat, *_ = np.linalg.lstsq(design, y, rcond=None)
    resid = y - design @ beta_hat
    dof = len(y) - design.shape[1]
    sigma2 = float(resid @ resid) / dof
    xtx_inv = np.linalg.inv(design.T @ design)
    se = np.sqrt(np.diag(xtx_inv) * sigma2)
    t_stat = beta_hat / se
    ss_tot = float(((y - y.mean()) ** 2).sum())
    r2 = 1.0 - float(resid @ resid) / ss_tot
    return {
        "alpha": float(beta_hat[0]),
        "beta": float(beta_hat[1]),
        "se_alpha": float(se[0]),
        "se_beta": float(se[1]),
        "t_alpha": float(t_stat[0]),
        "t_beta": float(t_stat[1]),
        "r2": r2,
        "n": float(len(y)),
        "dof": float(dof),
    }


def newey_west_t(y: np.ndarray, x: np.ndarray, maxlags: int) -> tuple[float, float]:
    """Bartlett 核的 HAC t 值。得分用残差乘回归元，与 OLS 三明治一致。"""
    design = np.column_stack([np.ones(len(x)), x])
    beta_hat, *_ = np.linalg.lstsq(design, y, rcond=None)
    resid = y - design @ beta_hat
    scores = design * resid[:, None]
    n_obs = len(y)
    meat = scores.T @ scores
    for lag in range(1, maxlags + 1):
        weight = 1.0 - lag / (maxlags + 1.0)
        gamma = scores[lag:].T @ scores[:-lag]
        meat = meat + weight * (gamma + gamma.T)
    meat = meat / n_obs
    bread = np.linalg.inv(design.T @ design / n_obs)
    cov = bread @ meat @ bread / n_obs
    se = np.sqrt(np.diag(cov))
    t_stat = beta_hat / se
    return float(t_stat[0]), float(t_stat[1])


def jarque_bera(sample: np.ndarray) -> tuple[float, float, float]:
    """矩估计的 Jarque–Bera。偏度和 Pearson 峰度都除以 n，正态峰度用 3。"""
    center = sample - sample.mean()
    m2 = float(np.mean(center**2))
    skew = float(np.mean(center**3) / m2**1.5)
    kurt = float(np.mean(center**4) / m2**2)
    n_obs = len(sample)
    return n_obs / 6.0 * (skew**2 + (kurt - 3.0) ** 2 / 4.0), skew, kurt


def adf_t_stat(level: np.ndarray) -> tuple[float, float]:
    """带常数项、不加滞后差分的回归 t。原假设是滞后水平的系数为 0。

    这不是 MacKinnon 的 p 值。adfuller 还会选滞后，两边的数字不必相等。
    """
    diff = np.diff(level)
    lagged = level[:-1]
    design = np.column_stack([np.ones(len(diff)), lagged])
    coef, *_ = np.linalg.lstsq(design, diff, rcond=None)
    resid = diff - design @ coef
    dof = len(diff) - design.shape[1]
    sigma2 = float(resid @ resid) / dof
    se = math.sqrt(sigma2 * float(np.linalg.inv(design.T @ design)[1, 1]))
    return float(coef[1]), float(coef[1] / se)


def eg_residual_t(left: np.ndarray, right: np.ndarray) -> tuple[float, float]:
    """Engle–Granger 第一步的斜率，以及残差上同一套 ADF 回归的 t。

    协整临界值比普通 t 更严。这个 t 只核对两步法的算术，不能拿 1.96 当拒绝线。
    """
    design = np.column_stack([np.ones(len(right)), right])
    coef, *_ = np.linalg.lstsq(design, left, rcond=None)
    resid = left - design @ coef
    _rho, t_stat = adf_t_stat(resid)
    return float(coef[1]), t_stat


def library_reference() -> None:
    """手写结果始终打印。scipy 与 statsmodels 未安装时只打印 skipped，不改上面的数。"""
    toy = np.array([-0.04, -0.01, 0.00, 0.01, 0.02])
    stat, skew, kurt = jarque_bera(toy)
    print("jb", f"{stat:.6f}", "skew", f"{skew:.6f}", "kurt", f"{kurt:.6f}")
    level = np.array([1.0, 1.1, 0.9, 1.2, 0.8, 1.3])
    rho, t_adf = adf_t_stat(level)
    print("adf_rho", f"{rho:.6f}", "adf_t", f"{t_adf:.6f}")
    left = np.array([2.0, 3.9, 6.1, 8.0, 9.8, 12.2])
    right = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
    slope, t_coint = eg_residual_t(left, right)
    print("coint_slope", f"{slope:.6f}", "coint_t", f"{t_coint:.6f}")
    y = np.array([0.015, -0.018, 0.030, -0.004])
    x = np.array([0.01, -0.02, 0.02, 0.00])
    fit = ols(y, x)
    # 点估计与第 9 节相同。这里多印 t，用来和第 10 节那次「HAC 把 t 压低」对照。
    print("ols_t", f"{fit['t_alpha']:.6f}", f"{fit['t_beta']:.6f}")
    t_alpha, t_beta = newey_west_t(y, x, maxlags=1)
    print("hac_L1_t", f"{t_alpha:.6f}", f"{t_beta:.6f}")
    try:
        from scipy import stats

        scipy_stat, scipy_p = stats.jarque_bera(toy)
        print("scipy_jb", f"{float(scipy_stat):.6f}", f"{float(scipy_p):.6f}")
    except ImportError:
        print("scipy_jb", "skipped")
    try:
        import statsmodels.api as sm
        from statsmodels.tsa.stattools import adfuller, coint

        fit = sm.OLS(y, sm.add_constant(x)).fit(cov_type="HAC", cov_kwds={"maxlags": 1})
        print("statsmodels_hac_t", f"{float(fit.tvalues[0]):.6f}", f"{float(fit.tvalues[1]):.6f}")
        adf_out = adfuller(level, maxlag=0, regression="c", autolag=None)
        print("statsmodels_adf_t", f"{float(adf_out[0]):.6f}")
        coint_out = coint(left, right, trend="c", autolag=None)
        print("statsmodels_coint_t", f"{float(coint_out[0]):.6f}")
    except ImportError:
        print("statsmodels_hac", "skipped")
        print("statsmodels_adf", "skipped")
        print("statsmodels_coint", "skipped")


def main() -> None:
    library_reference()
    spy = load_adj("SPY")
    aapl = load_adj("AAPL")
    px = pd.concat([spy, aapl], axis=1, join="inner").dropna()
    ret = px.pct_change().dropna()
    spy_r = ret["SPY"]
    print("sample", spy_r.index.min().date(), spy_r.index.max().date(), "n", len(spy_r))

    mean = float(spy_r.mean())
    std = float(spy_r.std(ddof=1))
    z = (spy_r - mean) / std
    m2 = float(((spy_r - mean) ** 2).mean())
    m3 = float(((spy_r - mean) ** 3).mean())
    m4 = float(((spy_r - mean) ** 4).mean())
    n_obs = len(spy_r)
    moment_skew = m3 / m2**1.5
    moment_kurt = m4 / m2**2
    adj = math.sqrt(n_obs * (n_obs - 1)) / (n_obs - 2)
    pandas_skew = adj * moment_skew
    # pandas 超额峰度：对正态为 0。这里用 Fisher 调整的常用形式。
    excess = moment_kurt - 3.0
    print("mean", mean, "std", std, "ann_vol", std * math.sqrt(252))
    print("moment_skew", moment_skew, "pandas_skew", pandas_skew)
    print("pearson_kurt", moment_kurt, "excess_kurt", excess)
    print("lag1_autocorr", float(spy_r.autocorr(lag=1)))

    for threshold in (-0.03, -0.05):
        actual = int((spy_r <= threshold).sum())
        z_cut = (threshold - mean) / std
        expected = n_obs * norm_cdf(z_cut)
        print("tail", threshold, "actual", actual, "normal_expected", expected, "z", z_cut)

    day = pd.Timestamp("2020-03-16")
    r_day = float(spy_r.loc[day])
    vol_2019 = float(spy_r.loc["2019-01-01":"2019-12-31"].std(ddof=1))
    z_2019 = r_day / vol_2019
    print("spy_2020_03_16", r_day, "vol_2019", vol_2019, "ann_2019", vol_2019 * math.sqrt(252))
    print("z_vs_2019", z_2019, "z_vs_full", r_day / std)

    y = (ret["AAPL"] - RF_ANNUAL / 252.0).to_numpy()
    x = (ret["SPY"] - RF_ANNUAL / 252.0).to_numpy()
    fit = ols(y, x)
    lags = int(math.floor(4.0 * (len(y) / 100.0) ** (2.0 / 9.0)))
    t_a, t_b = newey_west_t(y, x, lags)
    print("ols", fit)
    print("nw_lags", lags, "t_alpha", t_a, "t_beta", t_b)
    print("bonferroni_z", norm_ppf(1.0 - 0.001 / 2.0))
    print("p_t_2_1", 2.0 * (1.0 - norm_cdf(2.1)))
    print("lo_se_5y", math.sqrt((1.0 + 0.5 * 1.0**2) / 5.0))

    toy = np.array([-0.04, -0.01, 0.00, 0.01, 0.02])
    tm = float(toy.mean())
    tm2 = float(((toy - tm) ** 2).mean())
    tm3 = float(((toy - tm) ** 3).mean())
    tm4 = float(((toy - tm) ** 4).mean())
    print("toy_mean", tm, "m2", tm2, "skew", tm3 / tm2**1.5, "kurt", tm4 / tm2**2)
    hx = np.array([0.01, -0.02, 0.02, 0.00])
    hy = np.array([0.015, -0.018, 0.030, -0.004])
    print("hand_ols", ols(hy, hx))

    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.6))
    ax = axes[0]
    bins = np.linspace(-0.08, 0.08, 65)
    ax.hist(spy_r.clip(-0.08, 0.08), bins=bins, density=True, color="#1f4e79", alpha=0.75)
    grid = np.linspace(-0.08, 0.08, 400)
    normal = np.exp(-0.5 * ((grid - mean) / std) ** 2) / (std * math.sqrt(2.0 * math.pi))
    ax.plot(grid, normal, color="#c1121f", lw=1.8, label="同均值、同方差的正态密度")
    ax.set_title("SPY 日收益：中间更尖，两尾更厚")
    ax.set_xlabel("日简单收益（图中截断在 ±8%，统计用全样本）")
    ax.set_ylabel("密度")
    ax.legend(frameon=False, fontsize=8)
    ax.grid(True, alpha=0.3)

    ax = axes[1]
    ax.scatter(x, y, s=6, alpha=0.25, color="#1f4e79", rasterized=True)
    line_x = np.linspace(float(np.min(x)), float(np.max(x)), 50)
    ax.plot(line_x, fit["alpha"] + fit["beta"] * line_x, color="#c1121f", lw=1.8)
    ax.set_title(f"超额收益回归：β={fit['beta']:.2f}，R²={fit['r2']:.2f}")
    ax.set_xlabel("SPY 日超额收益（无风险利率按 4% 均摊）")
    ax.set_ylabel("AAPL 日超额收益")
    ax.grid(True, alpha=0.3)
    fig.suptitle(
        f"M01 种子数据 {spy_r.index.min().date()} 至 {spy_r.index.max().date()}（复权收盘价）",
        fontsize=12,
    )
    fig.tight_layout()
    fig.savefig(OUT, dpi=140, bbox_inches="tight")
    print("wrote", OUT)


if __name__ == "__main__":
    main()