"""M03 教学图：同一条 SPY 均线，有无 shift(1)，以及均线长度的夏普。

现金收益为 0。无 shift 的曲线是前视，不是更好的策略。
"""

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 / "lookahead_and_plateau.png"
RF = 0.02


def load_month() -> pd.Series:
    px = pd.read_csv(SPY, parse_dates=["session"])
    close = px.set_index("session")["adj_close"].astype(float).sort_index()
    return close.resample("ME").last().dropna()


def sharpe(ret: pd.Series) -> float:
    r = ret.dropna()
    std = float(r.std(ddof=1))
    if std == 0.0:
        return float("nan")
    excess = float(r.mean()) - RF / 12.0
    return excess / std * math.sqrt(12.0)


def cagr(ret: pd.Series) -> float:
    r = ret.dropna()
    years = len(r) / 12.0
    return float((1.0 + r).prod() ** (1.0 / years) - 1.0)


def run(month: pd.Series, window: int, delay: bool) -> pd.Series:
    sma = month.rolling(window).mean()
    held = (month > sma).where(sma.notna())
    if delay:
        held = held.shift(1)
    frame = pd.concat({"mkt": month.pct_change(), "held": held}, axis=1).dropna()
    return frame["mkt"] * frame["held"].astype(float)


def main() -> None:
    month = load_month()
    delayed = run(month, 10, True)
    leaked = run(month, 10, False)
    common = pd.concat({"ok": delayed, "leak": leaked}, axis=1).dropna()
    print("window", common.index.min().date(), common.index.max().date(), "n", len(common))
    print("delayed_cagr", cagr(common["ok"]), "sharpe", sharpe(common["ok"]))
    print("leaked_cagr", cagr(common["leak"]), "sharpe", sharpe(common["leak"]))

    lengths = list(range(6, 19))
    scores = [sharpe(run(month, n, True)) for n in lengths]
    for n, score in zip(lengths, scores, strict=True):
        print("len", n, "sharpe", round(score, 3))

    stitched: list[pd.Series] = []
    for year in range(2020, 2027):
        train_end = pd.Timestamp(year=year, month=1, day=1)
        test_end = pd.Timestamp(year=year + 1, month=1, day=1)
        best_n, best_score = None, -1e9
        for n in lengths:
            series = run(month, n, True)
            train_start = train_end - pd.DateOffset(years=5)
            train = series[(series.index >= train_start) & (series.index < train_end)]
            if len(train) < 36:
                continue
            score = sharpe(train)
            if score > best_score:
                best_n, best_score = n, score
        if best_n is None:
            continue
        test = run(month, best_n, True)
        test = test[(test.index >= train_end) & (test.index < test_end)]
        stitched.append(test)
        print(
            "year", year, "chose", best_n,
            "train_sharpe", round(best_score, 3), "test_n", len(test),
        )
    if stitched:
        oos = pd.concat(stitched)
        bh = month.pct_change().reindex(oos.index)
        fixed = run(month, 10, True).reindex(oos.index)
        print("oos_cagr", cagr(oos), "oos_sharpe", sharpe(oos), "oos_n", len(oos))
        print("same_months_buyhold", cagr(bh), sharpe(bh))
        print("same_months_ma10", cagr(fixed), sharpe(fixed))

    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.6))
    ax = axes[0]
    ax.plot((1.0 + common["ok"]).cumprod(), color="#1f4e79", lw=1.8, label="shift(1)，下月才持有")
    ax.plot((1.0 + common["leak"]).cumprod(), color="#c1121f", lw=1.6, label="用当月收盘赚当月")
    ax.set_yscale("log")
    ax.set_title("10 个月均线：前视把净值抬高")
    ax.set_ylabel("净值（对数刻度）")
    ax.legend(frameon=False, fontsize=8)
    ax.grid(True, alpha=0.3)

    ax = axes[1]
    ax.plot(lengths, scores, color="#1f4e79", marker="o", lw=1.6)
    ax.axvline(10, color="#6c757d", ls="--", lw=1.0)
    ax.set_xlabel("均线月数")
    ax.set_ylabel(f"夏普（月频，rf={RF:.0%}）")
    ax.set_title("参数从 6 到 18：看高原，不看单点")
    ax.grid(True, alpha=0.3)
    fig.suptitle("M03 种子数据演示（现金收益 0，不是实盘预期）", fontsize=12)
    fig.tight_layout()
    fig.savefig(OUT, dpi=140, bbox_inches="tight")
    print("wrote", OUT)


if __name__ == "__main__":
    main()