"""M11 教学图：因子模型的方差拆分，以及 HML 与市场的滚动相关。

左图是给定暴露和方差的公式。右图用 Kenneth French 月频，不是个股回测。
"""

from __future__ import annotations

import math
import zipfile
from io import BytesIO
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

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[2]
FF5_ZIP = HERE.parent / "M07" / "_cache" / "F-F_Research_Data_5_Factors_2x3_CSV.zip"
SPY = ROOT / "Labs" / "momentum_plus_lab" / "seed_data" / "spy.csv"
OUT = HERE / "risk_split_and_corr.png"


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)
    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_ff() -> pd.DataFrame:
    raw = zipfile.ZipFile(BytesIO(FF5_ZIP.read_bytes())).read(
        "F-F_Research_Data_5_Factors_2x3.csv"
    )
    return monthly_block(raw.decode("latin1"), ",Mkt-RF")


def two_names() -> None:
    vol = np.array([0.30, 0.20])
    corr = np.array([[1.0, 0.5], [0.5, 1.0]])
    cov = np.outer(vol, vol) * corr
    weights = np.array([0.5, 0.5])
    marginal = cov @ weights
    port_var = float(weights @ marginal)
    port_vol = math.sqrt(port_var)
    rc = weights * marginal / port_vol
    print("two_var", f"{port_var:.4f}", "two_vol", f"{port_vol:.4f}")
    print("two_marginal", f"{marginal[0]:.4f}", f"{marginal[1]:.4f}")
    print("two_rc", f"{rc[0]:.4f}", f"{rc[1]:.4f}")
    print("two_share", f"{rc[0] / port_vol:.4f}", f"{rc[1] / port_vol:.4f}")


def factor_book() -> np.ndarray:
    # 暴露和特质波动是教学输入。市场波动 18% 不是 French 的估计。
    beta = np.array([0.8, 1.0, 1.4])
    factor_var = 0.18 ** 2
    specific = np.array([0.25, 0.20, 0.30]) ** 2
    weights = np.array([0.50, 0.30, 0.20])
    beta_p = float(weights @ beta)
    factor_piece = beta_p ** 2 * factor_var
    specific_piece = weights ** 2 * specific
    total = factor_piece + float(specific_piece.sum())
    print("beta_p", f"{beta_p:.4f}", "factor_var", f"{factor_piece:.6f}")
    print(
        "specific_var",
        f"{specific_piece[0]:.6f}",
        f"{specific_piece[1]:.6f}",
        f"{specific_piece[2]:.6f}",
    )
    print("book_var", f"{total:.6f}", "book_vol", f"{math.sqrt(total):.4f}")
    print("factor_share", f"{factor_piece / total:.4f}")
    return np.concatenate([np.array([factor_piece]), specific_piece])


def teaching_wls_and_eigen() -> None:
    """教学特征值不是种子 7 的样本协方差。收缩目标是这两个特征值的均值。"""
    returns = np.array([0.01, 0.03])
    specific = np.array([0.04, 0.01])
    weights = 1.0 / specific
    wls = float(weights @ returns / weights.sum())
    print("wls", f"{wls:.6f}", "equal", f"{float(returns.mean()):.6f}")
    eigs = np.array([0.0400, 0.0100])
    shrunk = 0.5 * eigs + 0.5 * float(eigs.mean())
    print(
        "eigen", f"{shrunk[0]:.4f}", f"{shrunk[1]:.4f}",
        "cond", f"{shrunk[0] / shrunk[1]:.4f}",
    )
    # (0.27/0.18)^2 只乘因子一块。特质平方和留在 0.022825。
    factor = 0.031117 * (0.27 / 0.18) ** 2
    total = factor + 0.022825
    print(
        "vra_factor", f"{factor:.6f}",
        "vra_var", f"{total:.6f}",
        "vra_vol", f"{math.sqrt(total):.4f}",
    )


def shrinkage_demo() -> None:
    rng = np.random.default_rng(7)
    panel = rng.normal(size=(4, 6))
    sample = np.cov(panel, rowvar=False, ddof=1)
    eigs = np.linalg.eigvalsh(sample)
    scale = np.sqrt(np.diag(sample))
    corr = sample / np.outer(scale, scale)
    off = corr[np.triu_indices(6, k=1)]
    rho = float(off.mean())
    target = rho * np.outer(scale, scale)
    np.fill_diagonal(target, np.diag(sample))
    # δ = 0.5 是指定强度，不是 Ledoit–Wolf 估计出来的最优收缩。
    shrunk = 0.5 * sample + 0.5 * target
    shrunk_eigs = np.linalg.eigvalsh(shrunk)
    print(
        "sample_min_eig", f"{float(eigs[0]):.6f}",
        "sample_rank", int(np.linalg.matrix_rank(sample, tol=1e-10)),
    )
    print("rho_bar", f"{rho:.4f}")
    print(
        "shrunk_min_eig", f"{float(shrunk_eigs[0]):.6f}",
        "shrunk_cond", f"{float(shrunk_eigs[-1] / shrunk_eigs[0]):.2f}",
    )


def lrv_ratio(series: pd.Series, lag: int) -> float:
    values = series.to_numpy(dtype=float)
    centered = values - values.mean()
    length = len(centered)
    gamma0 = float(centered @ centered) / length
    total = gamma0
    for lag_j in range(1, lag + 1):
        cov_j = float(centered[lag_j:] @ centered[:-lag_j]) / length
        weight = 1.0 - lag_j / (lag + 1.0)
        total += 2.0 * weight * cov_j
    return total / gamma0


def french_panel(data: pd.DataFrame) -> pd.Series:
    pair = data[["Mkt-RF", "HML"]].dropna()
    window = pair.loc["2012-01":"2021-12"]
    year = pair.loc["2022-01":"2022-12"]
    print("pre_n", len(window), "y2022_n", len(year))
    for label, frame in (("pre", window), ("y2022", year)):
        vol_m = float(frame["Mkt-RF"].std(ddof=1) * math.sqrt(12.0))
        vol_h = float(frame["HML"].std(ddof=1) * math.sqrt(12.0))
        corr = float(frame["Mkt-RF"].corr(frame["HML"]))
        print(label, "vol_m", f"{vol_m:.4f}", "vol_h", f"{vol_h:.4f}", "corr", f"{corr:.4f}")
    growth = year["Mkt-RF"] - year["HML"]
    compound = float((1.0 + growth).prod() - 1.0)
    print("growth_tilt_2022", f"{compound:.4f}")
    for column in ("Mkt-RF", "HML"):
        print(column, "c2022", f"{float((1.0 + year[column]).prod() - 1.0):.4f}")
    lag = math.floor(4.0 * (len(window) / 100.0) ** (2.0 / 9.0))
    print("nw_lag", lag, "lrv_over_var_mkt", f"{lrv_ratio(window['Mkt-RF'], lag):.4f}")
    print("lrv_over_var_hml", f"{lrv_ratio(window['HML'], lag):.4f}")
    crash = pair.loc["2008-09":"2009-03", "Mkt-RF"]
    covid = pair.loc["2020-02":"2020-03", "Mkt-RF"]
    print("mkt_200809_200903", f"{float((1.0 + crash).prod() - 1.0):.4f}", "n", len(crash))
    print("mkt_202002_202003", f"{float((1.0 + covid).prod() - 1.0):.4f}", "n", len(covid))
    roll = pair["Mkt-RF"].rolling(36).corr(pair["HML"]).dropna()
    loc_min = roll.idxmin()
    loc_max = roll.idxmax()
    print(
        "roll_min", f"{float(roll.min()):.4f}", str(loc_min.date()),
        "roll_max", f"{float(roll.max()):.4f}", str(loc_max.date()),
    )
    print("roll_202112", f"{float(roll.loc[pd.Timestamp('2021-12-01')]):.4f}")
    print("roll_202212", f"{float(roll.loc[pd.Timestamp('2022-12-01')]):.4f}")
    return roll


def spy_forecast() -> None:
    close = pd.read_csv(SPY, parse_dates=["session"]).set_index("session")["adj_close"]
    ret = close.pct_change().dropna()
    start = pd.Timestamp("2020-02-19")
    end = pd.Timestamp("2020-03-23")
    # 预测只用 start 之前的收益。区间收益用两端收盘，不含 start 当天的日内。
    history = ret.loc[:start].iloc[:-1]
    window = history.iloc[-60:]
    sample_vol = float(window.std(ddof=1) * math.sqrt(252.0))
    # 用开窗前 60 日的方差做初值，再只吃 start 之前的收益。λ = 0.94 是 RiskMetrics 的日频常数。
    var = float(history.iloc[:60].var(ddof=1))
    lam = 0.94
    for value in history.iloc[60:]:
        var = lam * var + (1.0 - lam) * float(value) ** 2
    ewma_vol = math.sqrt(var * 252.0)
    held = float(close.loc[end] / close.loc[start] - 1.0)
    after = ret.loc[start:end].iloc[1:]
    realized = float(after.std(ddof=1) * math.sqrt(252.0))
    print("spy_sample60", f"{sample_vol:.4f}", "spy_ewma", f"{ewma_vol:.4f}")
    print("spy_realized", f"{realized:.4f}", "n", len(after), "spy_path", f"{held:.4f}")


def pure_factor() -> None:
    # 行业哑变量已知。Δ = I 时，纯行业组合是 X(X'X)^{-1} 的对应列。
    exposure = np.array([[1.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
    column = np.linalg.inv(exposure.T @ exposure)[:, 0]
    weights = exposure @ column
    print("pure_w", f"{weights[0]:.1f}", f"{weights[1]:.1f}", f"{weights[2]:.1f}")


def var_counterexample() -> None:
    # 单只违约概率 4%，95% 分位仍是 0。两只独立时，0 损失的概率只有 0.9216。
    quiet = 0.96 ** 2
    print("var_one", 0, "var_two", 100, "p_both_quiet", f"{quiet:.4f}")


def main() -> None:
    two_names()
    pieces = factor_book()
    shrinkage_demo()
    teaching_wls_and_eigen()
    pure_factor()
    var_counterexample()
    spy_forecast()
    roll = french_panel(load_ff())
    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.8))
    ax = axes[0]
    labels = ["市场因子", "股票 1 特质", "股票 2 特质", "股票 3 特质"]
    colors = ["#1f4e79", "#5b8e7d", "#c4a35a", "#a33b3b"]
    ax.bar(labels, pieces / pieces.sum() * 100.0, color=colors)
    ax.set_ylabel("占组合方差的比例（%）")
    ax.set_title("三只股票：因子方差与特质方差")
    ax.tick_params(axis="x", labelrotation=15)
    ax = axes[1]
    ax.plot(roll.index, roll.to_numpy(), color="#1f4e79")
    ax.axhline(0.0, color="#888888", linewidth=0.8)
    ax.set_title("HML 与 Mkt-RF 的 36 个月相关")
    ax.set_ylabel("相关系数")
    fig.suptitle("M11 风险拆分是公式；相关来自 French 月频")
    fig.tight_layout()
    fig.savefig(OUT, dpi=140)
    print("wrote", OUT)


if __name__ == "__main__":
    main()