"""M16 教学图：论文表内夏普的复算，以及用当期收益当特征的泄漏。

左图只复算 Gu、Kelly 与 Xiu (2020) 表 7 已给出的月均和月波动。
右图是固定种子的教学序列，不是股票回测。不调用 sklearn。
"""

from __future__ import annotations

import math
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

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

OUT = Path(__file__).resolve().parent / "oos_and_leak.png"


def paper_sharpe(mean_m: float, std_m: float) -> float:
    # 表 7 的 SR 列等于月均除以月波动，再乘 √12。这里不另扣无风险利率。
    return mean_m / std_m * math.sqrt(12.0)


def r2_zero(y: np.ndarray, yhat: np.ndarray) -> float:
    # 分母是收益平方和，不是离差平方和。零预测是 GKX 式 (19) 的基准。
    sse = float(np.sum((y - yhat) ** 2))
    sst = float(np.sum(y ** 2))
    return 1.0 - sse / sst


def r2_centered(y: np.ndarray, yhat: np.ndarray) -> float:
    sse = float(np.sum((y - yhat) ** 2))
    sst = float(np.sum((y - y.mean()) ** 2))
    return 1.0 - sse / sst


def fit_line(x: np.ndarray, y: np.ndarray) -> tuple[float, float]:
    # 带截距的一元回归。x 已经只用到预测时点之前的观测。
    design = np.column_stack([np.ones(len(x)), x])
    coef, *_ = np.linalg.lstsq(design, y, rcond=None)
    return float(coef[0]), float(coef[1])


def walk_forward(y: np.ndarray, x: np.ndarray, min_train: int) -> float:
    preds = []
    actual = []
    for t in range(min_train, len(y)):
        # 训练窗停在 t 之前。x[t] 若等于 y[t]，这一行就是泄漏。
        a, b = fit_line(x[:t], y[:t])
        preds.append(a + b * x[t])
        actual.append(y[t])
    return r2_zero(np.array(actual), np.array(preds))


def purged(n: int, horizon: int, test: range) -> list[int]:
    # 标签 y_t 用到 r_{t+1} … r_{t+horizon}。与测试标签共用收益的训练行要删掉。
    test_returns = set()
    for t in test:
        test_returns.update(range(t + 1, t + horizon + 1))
    kept = []
    for t in range(n):
        if t in test:
            continue
        used = set(range(t + 1, t + horizon + 1))
        if used.isdisjoint(test_returns):
            kept.append(t)
    return kept


def triple_barrier_labels(
    close: np.ndarray,
    events: np.ndarray,
    sigma: float,
    k_pt: float,
    k_sl: float,
    horizon: int,
) -> dict[str, np.ndarray]:
    """谁先碰到上轨、下轨或持有期末，就按谁离场。波动率是事件日已经知道的，不用路径后半段。

    水平栏按栏位收益成交，越过栏位的那一跳不记进标签。同一根收盘若两边都穿过，先记下轨：
    只做多账户分不清日内顺序时，不能把已经打穿止损的日子记成止盈。
    """
    returns = []
    signs = []
    meta = []
    steps = []
    for event in events:
        entry = float(close[event])
        exited = False
        for step in range(1, horizon + 1):
            path = float(close[event + step]) / entry - 1.0
            # 先看下轨。收盘同时穿过上下轨时，日内谁先发生是不知道的。
            if path <= -k_sl * sigma:
                returns.append(-k_sl * sigma)
                signs.append(-1)
                meta.append(0)
                steps.append(step)
                exited = True
                break
            if path >= k_pt * sigma:
                returns.append(k_pt * sigma)
                signs.append(1)
                meta.append(1)
                steps.append(step)
                exited = True
                break
        if exited:
            continue
        path = float(close[event + horizon]) / entry - 1.0
        returns.append(path)
        if path > 0.0:
            signs.append(1)
        elif path < 0.0:
            signs.append(-1)
        else:
            signs.append(0)
        # 竖直栏到期仍算主策略的一笔，但没有碰到止盈，元标签是 0。
        meta.append(0)
        steps.append(horizon)
    return {
        "ret": np.array(returns),
        "sign": np.array(signs),
        "meta": np.array(meta),
        "step": np.array(steps),
    }


def meta_label_f1_precision(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, float]:
    """元标签是「会不会碰到止盈」。精确率是预测为 1 的信号里真正碰到上轨的比例。"""
    truth = np.asarray(y_true, dtype=int)
    pred = np.asarray(y_pred, dtype=int)
    tp = int(np.sum((pred == 1) & (truth == 1)))
    fp = int(np.sum((pred == 1) & (truth == 0)))
    fn = int(np.sum((pred == 0) & (truth == 1)))
    precision = tp / (tp + fp) if tp + fp else 0.0
    recall = tp / (tp + fn) if tp + fn else 0.0
    # 两边都是 0 时 F1 定义为 0，避免把「没有正类」除成 nan。
    f1 = 0.0 if precision + recall == 0.0 else 2.0 * precision * recall / (precision + recall)
    return {
        "tp": float(tp),
        "fp": float(fp),
        "fn": float(fn),
        "precision": precision,
        "recall": recall,
        "f1": f1,
    }


def teaching_barrier_path() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """八次只做多信号放在 0,5,...,35。持有窗 4 根，彼此不重叠，所以这里不用再清除标签。"""
    close = np.full(40, 100.0)
    for start in (0, 15, 25):
        close[start : start + 5] = [100.0, 101.0, 104.0, 104.0, 104.0]
    for start in (5, 20, 30):
        close[start : start + 5] = [100.0, 99.0, 98.0, 98.0, 98.0]
    close[10:15] = [100.0, 100.1, 100.2, 100.3, 100.5]
    close[35:40] = [100.0, 99.8, 99.6, 99.5, 99.0]
    events = np.arange(0, 40, 5)
    # 分数在事件日已经知道。它只决定仓位，不会把主模型改成做空。
    score = np.array([0.8, 0.2, 0.3, 0.7, 0.6, 0.9, 0.1, 0.4])
    return close, events, score


def long_only_meta_demo() -> None:
    close, events, score = teaching_barrier_path()
    # sigma、止盈倍数和止损倍数都在进场时给定。上轨是 +2σ，下轨是 −1σ。
    labeled = triple_barrier_labels(close, events, sigma=0.02, k_pt=2.0, k_sl=1.0, horizon=4)
    pred = (score >= 0.5).astype(int)
    stats = meta_label_f1_precision(labeled["meta"], pred)
    primary = labeled["ret"]
    kept = primary[pred == 1]
    # 仓位就是分数本身，落在 [0, 1]。低于 0.5 的信号缩小，而不是乘上 −1。
    sized = float(np.mean(score * primary))
    print("tb_ret", [f"{v:.6f}" for v in primary])
    print("tb_sign", [int(v) for v in labeled["sign"]])
    print("tb_meta", [int(v) for v in labeled["meta"]])
    print("tb_step", [int(v) for v in labeled["step"]])
    print("primary_win", f"{float(np.mean(primary > 0.0)):.6f}")
    print("primary_mean", f"{float(np.mean(primary)):.6f}")
    print("filtered_win", f"{float(np.mean(kept > 0.0)):.6f}")
    print("filtered_mean", f"{float(np.mean(kept)):.6f}")
    print("sized_mean", f"{sized:.6f}")
    print("tp", int(stats["tp"]), "fp", int(stats["fp"]), "fn", int(stats["fn"]))
    print("precision", f"{stats['precision']:.6f}")
    print("recall", f"{stats['recall']:.6f}")
    print("f1", f"{stats['f1']:.6f}")
    production_pipeline_reference(score.reshape(-1, 1), labeled["meta"])


def production_pipeline_reference(features: np.ndarray, labels: np.ndarray) -> None:
    """对照 Pipeline 与 TimeSeriesSplit。这个切分不清除重叠标签，三栏路径重叠时仍要先做清除。"""
    try:
        from lightgbm import LGBMClassifier
        from sklearn.model_selection import TimeSeriesSplit
        from sklearn.pipeline import Pipeline
        from sklearn.preprocessing import StandardScaler
    except ImportError:
        print("sklearn_lgbm", "skipped")
        return
    pipe = Pipeline(
        [
            ("scale", StandardScaler()),
            ("clf", LGBMClassifier(n_estimators=20, learning_rate=0.1, verbosity=-1)),
        ]
    )
    splitter = TimeSeriesSplit(n_splits=2)
    for fold, (train, test) in enumerate(splitter.split(features), start=1):
        pipe.fit(features[train], labels[train])
        prob = pipe.predict_proba(features[test])[:, 1]
        print("fold", fold, "n_test", len(test), "mean_prob", f"{float(prob.mean()):.6f}")


def main() -> None:
    long_only_meta_demo()
    rows = {
        "OLS-3": (0.94, 5.33, 0.61),
        "NN3": (2.12, 6.13, 1.20),
        "NN4": (2.26, 5.80, 1.35),
    }
    for name, (mean_m, std_m, published) in rows.items():
        got = paper_sharpe(mean_m, std_m)
        print(name, f"{got:.4f}", "published", f"{published:.2f}")
    print("timing_gap", f"{0.77 - 0.51:.2f}")
    print("sqrt_0.004", f"{math.sqrt(0.004):.6f}")
    print("false_ir", f"{math.sqrt(0.004) * math.sqrt(6000):.4f}")

    y = np.array([0.02, -0.01, 0.03, 0.01])
    yhat = np.array([0.01, 0.00, 0.01, 0.00])
    print("toy_sse", f"{float(np.sum((y - yhat) ** 2)):.6f}")
    print("toy_sumsq", f"{float(np.sum(y ** 2)):.6f}")
    print("toy_r2_zero", f"{r2_zero(y, yhat):.6f}")
    print("toy_r2_centered", f"{r2_centered(y, yhat):.6f}")

    rng = np.random.default_rng(7)
    # 240 期纯噪声。短样本上滞后特征的 R² 可以碰巧为正，这里把路径放长。
    series = rng.normal(0.0, 0.02, size=240)
    lag = np.roll(series, 1)
    lag[0] = 0.0
    honest = walk_forward(series, lag, 80)
    leaked = walk_forward(series, series, 80)
    print("honest_r2", f"{honest:.6f}")
    print("leaked_r2", f"{leaked:.6f}")
    nulls = []
    for seed in range(20):
        draw = np.random.default_rng(seed).normal(0.0, 0.02, size=240)
        lagged = np.roll(draw, 1)
        lagged[0] = 0.0
        nulls.append(walk_forward(draw, lagged, 80))
    print("null_mean_20", f"{float(np.mean(nulls)):.6f}")
    print("purged", purged(8, 2, range(5, 7)))

    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.8))
    ax = axes[0]
    names = list(rows)
    published = [rows[name][2] for name in names]
    recomputed = [paper_sharpe(*rows[name][:2]) for name in names]
    xpos = np.arange(len(names))
    ax.bar(xpos - 0.16, published, width=0.32, color="#1f4e79", label="表 7 的 SR")
    ax.bar(xpos + 0.16, recomputed, width=0.32, color="#c4a35a", label="月均/月波动×√12")
    ax.set_xticks(xpos, names)
    ax.set_ylabel("年化夏普")
    ax.set_title("价值加权十分位多空，1987–2016")
    ax.legend(frameon=False)
    ax = axes[1]
    ax.bar(["滞后特征", "当期收益当特征"], [honest, leaked], color=["#5b8e7d", "#a33b3b"])
    ax.axhline(0.0, color="#333333", linewidth=0.8)
    ax.set_ylabel("零基准样本外 R²")
    ax.set_title("240 期噪声，种子 7，训练至少 80 期")
    fig.suptitle("M16：论文表内夏普可以复算；当期收益不是特征")
    fig.tight_layout()
    fig.savefig(OUT, dpi=140)
    print("wrote", OUT)


if __name__ == "__main__":
    main()