"""M06 教学图：French 动量十分位的平均收益与多空累计。

PRIOR 是形成月之前第 12 个月到第 2 个月。只读月频块。
原始 zip 放在 _cache，不入库。
"""

from __future__ import annotations

import io
import math
import urllib.request
import zipfile
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

BASE = "https://mba.tuck.dartmouth.edu/pages/faculty/ken.french/ftp/"
CACHE = Path(__file__).resolve().parent / "_cache"
OUT = Path(__file__).resolve().parent / "momentum_deciles.png"
FF3 = (
    Path(__file__).resolve().parents[1]
    / "M05"
    / "_cache"
    / "F-F_Research_Data_Factors_CSV.zip"
)


def fetch(name: str) -> bytes:
    path = CACHE / name
    if not path.exists():
        CACHE.mkdir(parents=True, exist_ok=True)
        path.write_bytes(urllib.request.urlopen(BASE + name, timeout=60).read())
    return path.read_bytes()


def monthly_block(text: str, marker: str) -> pd.DataFrame:
    lines = text.splitlines()
    start = next(i for i, line in enumerate(lines) if marker in line)
    # 标题行下一行才是表头；`,Mkt-RF` 这种行本身就是表头。
    if lines[start].startswith(","):
        header_line = lines[start]
        data_start = start + 1
    else:
        header_line = lines[start + 1]
        data_start = start + 2
    header = [part.strip() for part in header_line.split(",")]
    rows: list[list[str]] = []
    for line in lines[data_start:]:
        if not line.strip():
            break
        parts = [part.strip() for part in line.split(",")]
        if len(parts[0]) != 6:
            break
        rows.append(parts)
    frame = pd.DataFrame(rows, columns=header)
    frame = frame.rename(columns={header[0]: "date"})
    for column in frame.columns[1:]:
        frame[column] = pd.to_numeric(frame[column]) / 100.0
    frame["date"] = pd.to_datetime(frame["date"], format="%Y%m")
    return frame.set_index("date")


def load_ff3() -> pd.DataFrame:
    raw = zipfile.ZipFile(FF3).read("F-F_Research_Data_Factors.csv")
    return monthly_block(raw.decode("latin1"), ",Mkt-RF")


def nw_lag(n: int) -> int:
    return int(math.floor(4.0 * (n / 100.0) ** (2.0 / 9.0)))


def mean_t(series: pd.Series) -> tuple[float, float, float]:
    values = series.dropna().to_numpy(dtype=float)
    n = len(values)
    mean = float(values.mean())
    std = float(values.std(ddof=1))
    t_stat = mean / (std / math.sqrt(n))
    demeaned = values - mean
    lag = nw_lag(n)
    gamma0 = float(demeaned @ demeaned) / n
    hac = gamma0
    for k in range(1, lag + 1):
        gamma = float(demeaned[k:] @ demeaned[:-k]) / n
        hac += 2.0 * (1.0 - k / (lag + 1.0)) * gamma
    t_nw = mean / math.sqrt(hac / n)
    return mean * 12.0, t_stat, t_nw


def ols_hac(y: np.ndarray, x: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """含截距。经典标准误用残差自由度 T-K。HAC 用 Bartlett 核。"""
    n, k = x.shape
    coef = np.linalg.lstsq(x, y, rcond=None)[0]
    resid = y - x @ coef
    sigma2 = float(resid @ resid) / (n - k)
    xtx_inv = np.linalg.inv(x.T @ x)
    classic = np.sqrt(np.diag(sigma2 * xtx_inv))
    lag = nw_lag(n)
    scores = x * resid[:, None]
    meat = scores.T @ scores / n
    for step in range(1, lag + 1):
        gamma = scores[step:].T @ scores[:-step] / n
        weight = 1.0 - step / (lag + 1.0)
        meat += weight * (gamma + gamma.T)
    scale = xtx_inv * n
    var = scale @ meat @ scale / n
    hac = np.sqrt(np.diag(var))
    return coef, classic, hac


def grs(excess: pd.DataFrame, factors: pd.DataFrame) -> float:
    """Gibbons-Ross-Shanken。残差协方差除以 T-K-1，因子协方差 ddof=1。"""
    joined = excess.join(factors, how="inner").dropna()
    y = joined[excess.columns].to_numpy(dtype=float)
    f = joined[factors.columns].to_numpy(dtype=float)
    t_obs, n_port = y.shape
    k = f.shape[1]
    design = np.column_stack([np.ones(t_obs), f])
    coef = np.linalg.lstsq(design, y, rcond=None)[0]
    resid = y - design @ coef
    sigma = resid.T @ resid / (t_obs - k - 1)
    alpha = coef[0]
    mu = f.mean(axis=0)
    omega = np.cov(f, rowvar=False, ddof=1)
    sharpe = float(mu @ np.linalg.solve(omega, mu))
    alpha_term = float(alpha @ np.linalg.solve(sigma, alpha))
    return (t_obs - n_port - k) / n_port * alpha_term / (1.0 + sharpe)


def panel_sort_and_fama_macbeth() -> None:
    """宽表按行做截面。特征是形成日已知的，收益列已经是下一期，不要再 shift。

    股票那一维用矩阵乘完，不再写第二层 Python 循环。两期三列与第 5 节手算是同一组数。
    """
    char = pd.DataFrame([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]])
    fwd = pd.DataFrame([[0.01, 0.02, 0.06], [0.03, 0.03, 0.00]])
    char_d = char.sub(char.mean(axis=1), axis=0)
    fwd_d = fwd.sub(fwd.mean(axis=1), axis=0)
    gamma = (char_d * fwd_d).sum(axis=1) / (char_d**2).sum(axis=1)
    slope = gamma.to_numpy(dtype=float)
    # 设计矩阵每一期都是 [1, x]。einsum 的 t 维一次解完，截距和斜率分开。
    ones = np.ones((char.shape[0], char.shape[1], 1))
    design = np.concatenate([ones, char.to_numpy(dtype=float)[..., None]], axis=2)
    target = fwd.to_numpy(dtype=float)
    xtx = np.einsum("tnk,tnj->tkj", design, design)
    xty = np.einsum("tnk,tn->tk", design, target)
    coef = np.linalg.solve(xtx, xty[..., None])[..., 0]
    s = float(slope.std(ddof=1))
    t_stat = float(slope.mean() / (s / math.sqrt(len(slope))))
    print("fm_slope", f"{slope[0]:.6f}", f"{slope[1]:.6f}")
    print("fm_mean", f"{float(slope.mean()):.6f}")
    print("fm_t", f"{t_stat:.6f}")
    print("fm_matrix_intercept", f"{coef[0, 0]:.6f}", f"{coef[1, 0]:.6f}")
    print("fm_matrix_slope", f"{coef[0, 1]:.6f}", f"{coef[1, 1]:.6f}")
    # 另一张四列宽表只做高低各半，避免把上面的三列手算改成另一套斜率。
    wide_x = pd.DataFrame([[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0]])
    wide_y = pd.DataFrame([[0.00, 0.01, 0.02, 0.04], [0.03, 0.02, 0.01, 0.00]])
    rank = wide_x.rank(axis=1)
    n_names = wide_x.shape[1]
    high = wide_y.where(rank > n_names / 2).mean(axis=1)
    low = wide_y.where(rank <= n_names / 2).mean(axis=1)
    hml = (high - low).to_numpy(dtype=float)
    print("hml", f"{hml[0]:.6f}", f"{hml[1]:.6f}")
    print("hml_mean", f"{float(hml.mean()):.6f}")
    try:
        import statsmodels.api as sm

        fitted = []
        for row in range(len(target)):
            model = sm.OLS(target[row], sm.add_constant(char.to_numpy(dtype=float)[row])).fit()
            fitted.append(float(model.params[1]))
        print("statsmodels_fm_slope", f"{fitted[0]:.6f}", f"{fitted[1]:.6f}")
    except ImportError:
        print("statsmodels_fm", "skipped")


def toy_numbers() -> None:
    factor = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
    ret = np.array([0.01, -0.02, 0.03, 0.00, 0.04])
    left = pd.Series(factor).rank()
    right = pd.Series(ret).rank()
    # 默认环境没有 scipy。Spearman 就是秩的 Pearson，不要调用 method="spearman"。
    ic = float(left.corr(right))
    print("spearman", float(ic))
    mean, std, n = 0.03, 0.12, 120
    icir_m = mean / std
    print("icir_m", icir_m, "icir_ann", icir_m * math.sqrt(12), "t", icir_m * math.sqrt(n))
    print("ir_50_m", 0.03 * math.sqrt(50 * 12), "ir_500_m", 0.03 * math.sqrt(500 * 12))
    print("ir_500_y", 0.03 * math.sqrt(500))


def main() -> None:
    panel_sort_and_fama_macbeth()
    toy_numbers()
    port_text = zipfile.ZipFile(io.BytesIO(fetch("10_Portfolios_Prior_12_2_CSV.zip"))).read(
        "10_Portfolios_Prior_12_2.csv"
    ).decode("latin1")
    value = monthly_block(port_text, "Value Weight Returns -- Monthly")
    equal = monthly_block(port_text, "Equal Weighted Returns -- Monthly")
    mom_text = zipfile.ZipFile(io.BytesIO(fetch("F-F_Momentum_Factor_CSV.zip"))).read(
        "F-F_Momentum_Factor.csv"
    ).decode("latin1")
    mom = monthly_block(mom_text, ",Mom")
    ff = load_ff3()
    print("vw", value.index.min().date(), value.index.max().date(), len(value))
    print("ew", equal.index.min().date(), equal.index.max().date(), len(equal))
    cols = list(value.columns)
    for label, frame in (("vw", value), ("ew", equal)):
        spread = frame[cols[-1]] - frame[cols[0]]
        ann, t_stat, t_nw = mean_t(spread)
        print(
            label,
            "hl_ann", round(ann, 4),
            "t", round(t_stat, 2),
            "t_nw", round(t_nw, 2),
            "ac1", round(float(spread.autocorr(1)), 3),
        )
        means = [f"{float(frame[col].mean() * 12):.4f}" for col in cols]
        print(label, "decile_ann", means)
    joined = value.join(ff[["Mkt-RF", "SMB", "HML", "RF"]], how="inner")
    y = (joined[cols[-1]] - joined[cols[0]]).to_numpy(dtype=float)
    x = np.column_stack([
        np.ones(len(joined)),
        joined[["Mkt-RF", "SMB", "HML"]].to_numpy(dtype=float),
    ])
    coef, classic, hac = ols_hac(y, x)
    print(
        "vw_alpha_m", round(float(coef[0]), 5),
        "t", round(float(coef[0] / classic[0]), 2),
        "t_nw", round(float(coef[0] / hac[0]), 2),
        "betas", [round(float(v), 3) for v in coef[1:]],
    )
    excess = value.sub(ff["RF"], axis=0).dropna()
    factors = ff[["Mkt-RF", "SMB", "HML"]].reindex(excess.index).dropna()
    print("grs", round(grs(excess[cols], factors), 3), "T", len(excess), "lag", nw_lag(len(y)))
    for stamp in ("2009-03-01", "2009-04-01"):
        row = value.loc[stamp]
        print(
            stamp[:7],
            "vw_hl", round(float(row[cols[-1]] - row[cols[0]]), 4),
            "mom", round(float(mom.loc[stamp].iloc[0]), 4),
        )
    mom_ann, mom_t, mom_nw = mean_t(mom.iloc[:, 0])
    print("mom", round(mom_ann, 4), round(mom_t, 2), round(mom_nw, 2))
    print("mom_ac1", f"{float(mom.iloc[:, 0].autocorr(1)):.4f}")
    for start, end, name in (
        ("1963-07", "1990-12", "y1963_1990"),
        ("1963-07", "2026-08", "y1963_end"),
        ("1991-01", "2025-12", "y1991_2025"),
    ):
        block = value.loc[start:end]
        ann, t_stat, t_nw = mean_t(block[cols[-1]] - block[cols[0]])
        print(name, round(ann, 4), round(t_stat, 2), round(t_nw, 2), "n", len(block))

    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.6))
    ax = axes[0]
    xpos = np.arange(1, 11)
    vw_means = value.mean().to_numpy() * 12 * 100
    ew_means = equal.reindex(value.index).mean().to_numpy() * 12 * 100
    ax.bar(xpos - 0.18, vw_means, width=0.36, color="#1f4e79", label="市值加权")
    ax.bar(xpos + 0.18, ew_means, width=0.36, color="#2a9d8f", label="等权")
    ax.set_xticks(xpos)
    ax.set_xlabel("过去 12 到 2 个月收益，低到高")
    ax.set_ylabel("算术平均月收益 × 12（%）")
    ax.legend(frameon=False, fontsize=8)
    ax.grid(True, axis="y", alpha=0.3)
    ax = axes[1]
    spread_vw = (1.0 + value[cols[-1]] - value[cols[0]]).cumprod()
    spread_ew = (1.0 + equal[cols[-1]] - equal[cols[0]]).cumprod()
    ax.plot(spread_vw, color="#1f4e79", lw=1.4, label="市值加权 高减低")
    ax.plot(spread_ew, color="#2a9d8f", lw=1.2, label="等权 高减低")
    ax.set_yscale("log")
    ax.set_ylabel("累计倍数（对数刻度）")
    ax.legend(frameon=False, fontsize=8)
    ax.grid(True, alpha=0.3)
    fig.suptitle("M06 French 动量十分位（月频，未扣成本）", fontsize=12)
    fig.tight_layout()
    fig.savefig(OUT, dpi=140, bbox_inches="tight")
    print("wrote", OUT)


if __name__ == "__main__":
    main()