"""M04 教学图：SPY 波动率目标，以及杠杆对几何增长的拖累。

波动率用截至昨日的 EWMA。权重上限 1.5，空余仓位的现金收益为 0。
"""

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]
SPY = ROOT / "Labs" / "momentum_plus_lab" / "seed_data" / "spy.csv"
OUT = Path(__file__).resolve().parent / "vol_target_and_growth.png"
RF = 0.02
TARGET = 0.12


def load_close() -> pd.Series:
    frame = pd.read_csv(SPY, parse_dates=["session"])
    return frame.set_index("session")["adj_close"].astype(float).sort_index()


def drawdown_halve(close: pd.Series) -> dict[str, float]:
    """10 个月均线；回撤用本月开始前的账户净值，到 10% 才把仓位减半。"""
    month = close.resample("ME").last()
    ret = month.pct_change()
    sma = month.rolling(10).mean()
    # 暖机期的 NaN 不能当成「已跌破均线」。shift(1) 让本月收益用上月收盘的信号。
    weight = (month > sma).where(sma.notna()).shift(1)
    nav = 1.0
    peak = 1.0
    halved = 0
    out: list[float] = []
    for session in weight.dropna().index:
        scale = 0.5 if (1.0 - nav / peak) >= 0.10 else 1.0
        halved += int(scale < 1.0)
        step = float(weight.loc[session]) * float(ret.loc[session]) * scale
        nav *= 1.0 + step
        peak = max(peak, nav)
        out.append(step)
    series = pd.Series(out)
    years = len(series) / 12.0
    cagr = float((1.0 + series).prod() ** (1.0 / years) - 1.0)
    sharpe = (float(series.mean()) - RF / 12.0) / float(series.std(ddof=1))
    sharpe *= math.sqrt(12.0)
    curve = (1.0 + series).cumprod()
    mdd = float((curve / curve.cummax() - 1.0).min())
    return {"cagr": cagr, "sharpe": sharpe, "mdd": mdd, "months_half": float(halved)}


def performance(ret: pd.Series, rf: float) -> dict[str, float]:
    r = ret.dropna()
    years = len(r) / 252.0
    cagr = float((1.0 + r).prod() ** (1.0 / years) - 1.0)
    std = float(r.std(ddof=1))
    sharpe = (float(r.mean()) - rf / 252.0) / std * math.sqrt(252.0)
    nav = (1.0 + r).cumprod()
    mdd = float((nav / nav.cummax() - 1.0).min())
    return {"cagr": cagr, "vol": std * math.sqrt(252.0), "sharpe": sharpe, "mdd": mdd}


def main() -> None:
    excess = 0.06
    sigma = 0.15
    full = excess / sigma**2
    print("kelly_full", full, "half", full / 2)
    f = full
    true = 0.03
    growth = f * true - (f**2) * (sigma**2) / 2
    print("mis_specified_excess_growth", growth)
    mu, sig = 0.12, 0.30
    print("geom_1x", mu - sig**2 / 2)
    print("geom_2x", 2 * mu - (2 * sig) ** 2 / 2)

    close = load_close()
    ret = close.pct_change().dropna()
    # span=60 的 EWMA 用到当日收益。shift(1) 后，今日仓位只用昨日及更早的波动。
    vol = ret.ewm(span=60).std(bias=False) * math.sqrt(252.0)
    weight = (TARGET / vol).clip(upper=1.5).shift(1)
    # 教学图允许 1.5，用来看上限。账户政策若把杠杆钉在 1，部署时要再截一次。
    capped = (TARGET / vol).clip(upper=1.0).shift(1)
    managed = (weight * ret).dropna()
    no_lev = (capped * ret).reindex(managed.index).dropna()
    held = ret.reindex(managed.index)
    print("sample", managed.index.min().date(), managed.index.max().date(), "n", len(managed))
    used = weight.reindex(managed.index)
    print(
        "weight_mean", float(used.mean()),
        "share_above_1", float((used > 1.0).mean()),
        "share_at_cap", float((used >= 1.5 - 1e-12).mean()),
    )
    print("buyhold", performance(held, RF))
    print("managed", performance(managed, RF))
    print("cap1", performance(no_lev, RF))
    for label, start, end in (
        ("2020", "2020-01-01", "2020-12-31"),
        ("2022", "2022-01-01", "2022-12-31"),
    ):
        print(label, "bh", performance(held.loc[start:end], RF))
        print(label, "vt", performance(managed.loc[start:end], RF))

    rng = np.random.default_rng(7)
    kelly = 0.55 - 0.45
    n_paths, n_bets = 1000, 1000
    for scale, name in ((0.5, "half"), (1.0, "full"), (2.0, "double")):
        f_bet = scale * kelly
        wins = rng.random((n_paths, n_bets)) < 0.55
        step = np.where(wins, 1.0 + f_bet, 1.0 - f_bet)
        wealth = np.cumprod(step, axis=1)
        terminal = wealth[:, -1]
        peak = np.maximum.accumulate(wealth, axis=1)
        mdd = (wealth / peak - 1.0).min(axis=1)
        print(
            name,
            "median", float(np.median(terminal)),
            "p10", float(np.quantile(terminal, 0.1)),
            "mdd_median", float(np.median(mdd)),
        )

    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.6))
    ax = axes[0]
    ax.plot((1.0 + held).cumprod(), color="#6c757d", lw=1.3, label="买入持有")
    ax.plot((1.0 + managed).cumprod(), color="#1f4e79", lw=1.6, label="目标波动 12%，上限 1.5")
    ax.set_yscale("log")
    ax.set_title("SPY：仓位用昨日波动，现金收益为 0")
    ax.set_ylabel("净值（对数刻度）")
    ax.legend(frameon=False, fontsize=8)
    ax.grid(True, alpha=0.3)

    ax = axes[1]
    lev = np.linspace(0.0, 3.0, 200)
    growth_curve = lev * mu - (lev * sig) ** 2 / 2.0
    ax.plot(lev, growth_curve * 100, color="#1f4e79", lw=1.8)
    ax.axvline(mu / sig**2, color="#c1121f", ls="--", lw=1.0)
    ax.set_xlabel("杠杆（无风险利率按 0）")
    ax.set_ylabel("近似几何超额（%）")
    ax.set_title("μ=12%，σ=30%：顶点在凯利杠杆")
    ax.grid(True, alpha=0.3)
    fig.suptitle("M04 教学图（右图不是历史回测）", fontsize=12)
    fig.tight_layout()
    fig.savefig(OUT, dpi=140, bbox_inches="tight")
    print("wrote", OUT)
    overlay = drawdown_halve(close)
    print(
        "ma10_halve",
        "cagr", round(overlay["cagr"], 4),
        "sharpe", round(overlay["sharpe"], 2),
        "mdd", round(overlay["mdd"], 4),
        "months_half", overlay["months_half"],
    )


if __name__ == "__main__":
    main()