"""M12 教学图：三种权重，以及等权组合的波动怎样随只数下降。

左图是一组给定的 α 和协方差，不是样本外回测。右图是常数相关公式。
只做多解用活跃集枚举，不调用 scipy 或 cvxpy。
"""

from __future__ import annotations

import math
from decimal import Decimal
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from quant.portfolio.allocation import AllocationPolicy, sleeve_to_account
from quant.risk.volatility import inverse_volatility_weights, target_volatility_scale

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

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


def covariance(vols: np.ndarray, corr: np.ndarray) -> np.ndarray:
    scale = np.outer(vols, vols)
    return scale * corr


def contributions(weights: np.ndarray, cov: np.ndarray) -> tuple[float, np.ndarray]:
    marginal = cov @ weights
    variance = float(weights @ marginal)
    vol = math.sqrt(variance)
    rc = weights * marginal / vol
    return vol, rc / vol


def equal_corr(n: int, rho: float) -> np.ndarray:
    corr = np.full((n, n), rho)
    np.fill_diagonal(corr, 1.0)
    return corr


def risk_parity(vols: np.ndarray, corr: np.ndarray) -> np.ndarray:
    # 从逆波动出发，按风险贡献偏离来改权重。常数相关时它应停在逆波动。
    cov = covariance(vols, corr)
    weights = (1.0 / vols) / np.sum(1.0 / vols)
    for _ in range(200):
        _, share = contributions(weights, cov)
        weights = weights * (1.0 / share)
        weights = weights / weights.sum()
    return weights


def mean_variance(alpha: np.ndarray, cov: np.ndarray, risk_aversion: float) -> np.ndarray:
    # 目标是 w'α − λ w'Σw，没有预算约束。λ 决定总仓位，现金补足 1 以内的差额。
    direction = np.linalg.solve(cov, alpha)
    return direction / (2.0 * risk_aversion)


def fully_invested(alpha: np.ndarray, cov: np.ndarray, risk_aversion: float) -> np.ndarray:
    # 强制权重和为 1。γ 是预算约束的乘子，λ 不再被归一化消掉。
    ones = np.ones(len(alpha))
    precision_alpha = np.linalg.solve(cov, alpha)
    precision_ones = np.linalg.solve(cov, ones)
    gamma = (precision_alpha.sum() - 2.0 * risk_aversion) / precision_ones.sum()
    return np.linalg.solve(cov, alpha - gamma * ones) / (2.0 * risk_aversion)


def diversified_vol(n: int, sigma: float, rho: float) -> float:
    return sigma * math.sqrt(rho + (1.0 - rho) / n)


def _bound_states(n: int) -> list[tuple[int, ...]]:
    # 0 压在下界，1 在内部，2 压在上限。三只股票只有 27 种，可以穷尽。
    states: list[tuple[int, ...]] = []
    for mask in range(3**n):
        flags: list[int] = []
        cursor = mask
        for _ in range(n):
            flags.append(cursor % 3)
            cursor //= 3
        states.append(tuple(flags))
    return states


def kkt_mean_variance(
    weights: np.ndarray,
    alpha: np.ndarray,
    cov: np.ndarray,
    risk_aversion: float,
    upper: float,
) -> float:
    """最大化 w'α − λ w'Σw 的违反量。自由变量梯度应相等；下界上的梯度不能更高。"""
    gradient = alpha - 2.0 * risk_aversion * (cov @ weights)
    free = (weights > 1e-10) & (weights < upper - 1e-10)
    lower = weights <= 1e-10
    capped = weights >= upper - 1e-10
    if not np.any(free):
        return 1.0
    gamma = float(np.mean(gradient[free]))
    gap = float(np.max(np.abs(gradient[free] - gamma)))
    if np.any(lower):
        gap = max(gap, float(np.max(gradient[lower] - gamma)))
    if np.any(capped):
        gap = max(gap, float(np.max(gamma - gradient[capped])))
    gap = max(gap, abs(float(weights.sum()) - 1.0))
    gap = max(gap, float(np.max(np.maximum(-weights, 0.0))))
    gap = max(gap, float(np.max(weights - upper)))
    return max(gap, 0.0)


def long_only_mean_variance(
    alpha: np.ndarray,
    cov: np.ndarray,
    risk_aversion: float,
    upper: float,
) -> np.ndarray:
    """在边界上重解。把无约束负数截成 0 再归一，过不了上面的梯度检验。"""
    best_w: np.ndarray | None = None
    best_obj = -np.inf
    for state in _bound_states(len(alpha)):
        free = [i for i, flag in enumerate(state) if flag == 1]
        fixed = np.array([upper if flag == 2 else 0.0 for flag in state])
        budget = 1.0 - float(fixed.sum())
        if budget < -1e-12 or (free and budget <= 1e-12):
            continue
        if not free:
            if abs(budget) > 1e-10:
                continue
            weights = fixed.copy()
        else:
            size = len(free)
            system = np.zeros((size + 1, size + 1))
            system[:size, :size] = 2.0 * risk_aversion * cov[np.ix_(free, free)]
            system[:size, size] = 1.0
            system[size, :size] = 1.0
            rhs = np.zeros(size + 1)
            rhs[:size] = alpha[free] - 2.0 * risk_aversion * (cov[free] @ fixed)
            rhs[size] = budget
            try:
                solved = np.linalg.solve(system, rhs)
            except np.linalg.LinAlgError:
                continue
            interior = solved[:size]
            if np.any(interior <= 1e-12) or np.any(interior >= upper - 1e-12):
                continue
            weights = fixed.copy()
            weights[free] = interior
        gap = kkt_mean_variance(weights, alpha, cov, risk_aversion, upper)
        if gap >= 1e-8:
            continue
        obj = float(weights @ alpha - risk_aversion * (weights @ cov @ weights))
        if obj > best_obj:
            best_obj = obj
            best_w = weights.copy()
    if best_w is None:
        raise RuntimeError("long-only mean-variance has no KKT point")
    return best_w


def kkt_min_variance(weights: np.ndarray, cov: np.ndarray, upper: float) -> float:
    """最小化 w'Σw。压在 0 的名字，边际方差不能低于仍留在组合里的名字。"""
    marginal = cov @ weights
    free = (weights > 1e-10) & (weights < upper - 1e-10)
    lower = weights <= 1e-10
    capped = weights >= upper - 1e-10
    if not np.any(free):
        return 1.0
    gamma = float(np.mean(marginal[free]))
    gap = float(np.max(np.abs(marginal[free] - gamma)))
    if np.any(lower):
        gap = max(gap, float(np.max(gamma - marginal[lower])))
    if np.any(capped):
        gap = max(gap, float(np.max(marginal[capped] - gamma)))
    gap = max(gap, abs(float(weights.sum()) - 1.0))
    gap = max(gap, float(np.max(np.maximum(-weights, 0.0))))
    gap = max(gap, float(np.max(weights - upper)))
    return max(gap, 0.0)


def long_only_min_variance(cov: np.ndarray, upper: float = 1.0) -> np.ndarray:
    """只做多最小方差。无约束解里的空头要在边界上重解，不能直接改成 0。"""
    best_w: np.ndarray | None = None
    best_var = np.inf
    for state in _bound_states(cov.shape[0]):
        free = [i for i, flag in enumerate(state) if flag == 1]
        fixed = np.array([upper if flag == 2 else 0.0 for flag in state])
        budget = 1.0 - float(fixed.sum())
        if budget < -1e-12 or (free and budget <= 1e-12):
            continue
        if not free:
            if abs(budget) > 1e-10:
                continue
            weights = fixed.copy()
        else:
            size = len(free)
            system = np.zeros((size + 1, size + 1))
            system[:size, :size] = cov[np.ix_(free, free)]
            system[:size, size] = -1.0
            system[size, :size] = 1.0
            rhs = np.zeros(size + 1)
            rhs[:size] = -(cov[free] @ fixed)
            rhs[size] = budget
            try:
                solved = np.linalg.solve(system, rhs)
            except np.linalg.LinAlgError:
                continue
            interior = solved[:size]
            if np.any(interior <= 1e-12) or np.any(interior >= upper - 1e-12):
                continue
            weights = fixed.copy()
            weights[free] = interior
        if kkt_min_variance(weights, cov, upper) >= 1e-8:
            continue
        variance = float(weights @ cov @ weights)
        if variance < best_var:
            best_var = variance
            best_w = weights.copy()
    if best_w is None:
        raise RuntimeError("long-only minimum variance has no KKT point")
    return best_w


def diversification_multiplier(weights: np.ndarray, corr: np.ndarray) -> float:
    """IDM = 1 / sqrt(w' C w)。相关低于 1 且权重和为 1 时，这个倍数大于 1。"""
    quad = float(weights @ corr @ weights)
    return 1.0 / math.sqrt(quad)


def print_weights(label: str, weights: np.ndarray) -> None:
    text = " ".join(f"{value:.4f}" for value in weights)
    print(label, text, "sum", f"{float(weights.sum()):.4f}")


def main() -> None:
    vols = np.array([0.10, 0.20, 0.30])
    inv = (1.0 / vols) / np.sum(1.0 / vols)
    print_weights("inv_vol", inv)
    for rho in (0.0, 0.5):
        cov = covariance(vols, equal_corr(3, rho))
        vol, share = contributions(inv, cov)
        erc = risk_parity(vols, equal_corr(3, rho))
        print(f"rho_{rho:.1f}", "vol", f"{vol:.4f}", "share", " ".join(f"{s:.4f}" for s in share))
        print_weights(f"erc_{rho:.1f}", erc)

    uneven = np.array(
        [
            [1.0, 0.10, 0.60],
            [0.10, 1.0, 0.10],
            [0.60, 0.10, 1.0],
        ]
    )
    print("uneven_min_eig", f"{float(np.linalg.eigvalsh(uneven)[0]):.4f}")
    cov_u = covariance(vols, uneven)
    vol_u, share_u = contributions(inv, cov_u)
    erc_u = risk_parity(vols, uneven)
    _, share_e = contributions(erc_u, cov_u)
    print("uneven_inv_vol", f"{vol_u:.4f}", "share", " ".join(f"{s:.4f}" for s in share_u))
    print_weights("uneven_erc", erc_u)
    print("uneven_erc_share", " ".join(f"{s:.4f}" for s in share_e))

    alpha = np.array([0.06, 0.04, 0.02])
    shocked = np.array([0.09, 0.04, 0.02])
    cov = covariance(vols, equal_corr(3, 0.30))
    print("mv_min_eig", f"{float(np.linalg.eigvalsh(cov)[0]):.6f}")
    base = mean_variance(alpha, cov, 3.0)
    jump = mean_variance(shocked, cov, 3.0)
    print_weights("mv_base", base)
    print_weights("mv_shock", jump)
    print_weights("mv_base_mix", base / base.sum())
    print_weights("mv_shock_mix", jump / jump.sum())
    invested = fully_invested(alpha, cov, 3.0)
    print_weights("mv_sum1", invested)
    gradient = alpha - 2.0 * 3.0 * (cov @ invested)
    print("sum1_grad_span", f"{float(gradient.max() - gradient.min()):.6f}")
    min_var = np.linalg.solve(cov, np.ones(3))
    min_var = min_var / min_var.sum()
    print_weights("min_var", min_var)

    # 活跃集留下的点要过 KKT。截断再归一的点故意留在对照里，它的梯度差应明显大于 1e-8。
    long_only = long_only_mean_variance(alpha, cov, 3.0, upper=1.0)
    cap50 = long_only_mean_variance(alpha, cov, 3.0, upper=0.50)
    long_min = long_only_min_variance(cov)
    naive = np.maximum(base, 0.0)
    naive = naive / naive.sum()
    expected = {
        "long_only_sum1": np.array([56.0 / 57.0, 1.0 / 57.0, 0.0]),
        "long_only_cap50": np.array([0.5, 245.0 / 564.0, 37.0 / 564.0]),
        "long_only_minvar": np.array([17.0 / 19.0, 2.0 / 19.0, 0.0]),
        "naive_clip": np.array([70.0 / 73.0, 3.0 / 73.0, 0.0]),
    }
    solved = {
        "long_only_sum1": long_only,
        "long_only_cap50": cap50,
        "long_only_minvar": long_min,
        "naive_clip": naive,
    }
    for label, weights in solved.items():
        text = " ".join(f"{value:.10f}" for value in weights)
        upper = 0.50 if label == "long_only_cap50" else 1.0
        if label == "long_only_minvar":
            gap = kkt_min_variance(weights, cov, upper)
        else:
            gap = kkt_mean_variance(weights, alpha, cov, 3.0, upper)
        mu = float(weights @ alpha)
        variance = float(weights @ cov @ weights)
        obj = mu - 3.0 * variance
        print(label, text)
        print(
            label, "mu", f"{mu:.10f}", "vol", f"{math.sqrt(variance):.10f}",
            "obj", f"{obj:.10f}", "kkt", f"{gap:.3e}",
        )
        if not np.allclose(weights, expected[label], atol=1e-12):
            raise RuntimeError(f"{label} drifted from the closed form")
        if label == "naive_clip":
            if gap < 1e-8:
                raise RuntimeError("clipped weights unexpectedly satisfy KKT")
        elif gap >= 1e-8:
            raise RuntimeError(f"{label} KKT gap {gap}")
    long_obj = float(long_only @ alpha - 3.0 * (long_only @ cov @ long_only))
    naive_obj = float(naive @ alpha - 3.0 * (naive @ cov @ naive))
    print("naive_obj_loss", f"{long_obj - naive_obj:.10f}")
    naive_g = alpha - 6.0 * (cov @ naive)
    positive = naive > 1e-10
    span = float(naive_g[positive].max() - naive_g[positive].min())
    print("naive_grad_span", f"{span:.10f}")

    idm_corr = equal_corr(3, 0.25)
    idm_w = np.full(3, 1.0 / 3.0)
    quad = float(idm_w @ idm_corr @ idm_w)
    idm = diversification_multiplier(idm_w, idm_corr)
    # 三个子系统都已在目标波动上时，IDM 就是想加的杠杆。上限 1 会把大于 1 的倍数截掉。
    combined = 0.10 * math.sqrt(quad)
    scale = target_volatility_scale(combined, 0.10, max_scale=1.0)
    print("idm", f"{idm:.10f}", "quad", f"{quad:.10f}")
    print("idm_scale", f"{scale:.4f}", "wanted", f"{0.10 / combined:.10f}")

    market = np.array([0.60, 0.40])
    pair_vol = np.array([0.20, 0.15])
    pair_cov = covariance(pair_vol, equal_corr(2, 0.30))
    implied = 2.5 * (pair_cov @ market)
    print("implied", f"{implied[0]:.4f}", f"{implied[1]:.4f}")
    # Ω 取成观点的先验方差时，后验利差是先验利差和观点的平均。
    tau = 0.05
    view = np.array([[1.0, -1.0]])
    omega = view @ (tau * pair_cov) @ view.T
    prior_var = tau * pair_cov
    gain = prior_var @ view.T @ np.linalg.inv(view @ prior_var @ view.T + omega)
    posterior = implied + gain @ (np.array([0.04]) - view @ implied)
    print("posterior", f"{float(posterior[0]):.6f}", f"{float(posterior[1]):.6f}")
    cluster_var = np.array([0.04, 0.01])
    cluster_w = (1.0 / cluster_var)
    cluster_w = cluster_w / cluster_w.sum()
    print("hrp_two", f"{cluster_w[0]:.4f}", f"{cluster_w[1]:.4f}")

    one_way = 0.5 * float(np.abs(base - jump).sum())
    print("one_way_from_shock", f"{one_way:.4f}")
    for side, name in ((0.0015, "us"), (0.0025, "hk")):
        print(name, "rebalance_cost", f"{2.0 * one_way * side:.6f}")

    series = pd.Series(vols, index=["A", "B", "C"])
    capped = inverse_volatility_weights(series, max_weight=0.40)
    print_weights("cap_0.40", capped.to_numpy(dtype=float))
    try:
        inverse_volatility_weights(series, max_weight=0.10)
    except Exception as exc:
        print("cap_0.10", type(exc).__name__)

    policy = AllocationPolicy(
        strategy_budget_pct=Decimal("0.30"),
        max_position_pct=Decimal("0.10"),
        max_order_value_pct=Decimal("0.05"),
        rebalance_buffer_pct=Decimal("0.01"),
        max_cost_skip_pct=Decimal("0.02"),
        max_daily_child_orders=4,
        lot_size=Decimal("1"),
    )
    sleeve = {name: Decimal(f"{weight:.4f}") for name, weight in zip("ABC", inv)}
    account = sleeve_to_account(sleeve, policy)
    print(
        "account",
        " ".join(f"{float(account[name]):.4f}" for name in "ABC"),
        "sum",
        f"{float(sum(account.values(), Decimal(0))):.4f}",
    )
    capped_sleeve = sleeve_to_account({"A": Decimal("0.80"), "B": Decimal("0.20")}, policy)
    print(
        "no_renorm",
        f"{float(capped_sleeve['A']):.4f}",
        f"{float(capped_sleeve['B']):.4f}",
    )

    for rho_sb, label in ((-0.2, "neg"), (0.4, "pos")):
        pair_vol = np.array([0.16, 0.06])
        pair = covariance(pair_vol, equal_corr(2, rho_sb))
        weights = (1.0 / pair_vol) / np.sum(1.0 / pair_vol)
        port, _ = contributions(weights, pair)
        scale = target_volatility_scale(port, 0.10, max_scale=1.0)
        print(
            label, "w", f"{weights[0]:.4f}", f"{weights[1]:.4f}",
            "vol", f"{port:.4f}", "scale", f"{scale:.4f}",
            "wanted", f"{0.10 / port:.4f}",
        )

    for n in (1, 10, 30, 40, 60):
        row = " ".join(f"{diversified_vol(n, 0.40, rho):.4f}" for rho in (0.0, 0.3, 0.6))
        print("breadth", n, row)
    print("limit", " ".join(f"{0.40 * math.sqrt(rho):.4f}" for rho in (0.0, 0.3, 0.6)))

    names = ["低波动", "中波动", "高波动"]
    methods = {
        "等权": np.full(3, 1.0 / 3.0),
        "逆波动": inv,
        "无约束": base,
        "α 加 3 点": jump,
    }
    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.8))
    ax = axes[0]
    x = np.arange(len(names))
    width = 0.18
    colors = ["#1f4e79", "#5b8e7d", "#c4a35a", "#a33b3b"]
    for index, (label, weights) in enumerate(methods.items()):
        ax.bar(x + (index - 1.5) * width, weights * 100.0, width, label=label, color=colors[index])
    ax.set_xticks(x, names)
    ax.set_ylabel("权重（%）")
    ax.axhline(0.0, color="#888888", linewidth=0.8)
    ax.set_title("无约束解可以超过 1，也可以为负")
    ax.legend(frameon=False, fontsize=8)
    ax = axes[1]
    counts = np.arange(1, 61)
    for rho, color in ((0.0, "#5b8e7d"), (0.3, "#1f4e79"), (0.6, "#a33b3b")):
        curve = [diversified_vol(int(n), 0.40, rho) * 100.0 for n in counts]
        ax.plot(counts, curve, color=color, label=f"相关 {rho:.1f}")
    ax.set_xlabel("等权只数")
    ax.set_ylabel("组合波动（%）")
    ax.set_title("单只波动 40% 时，只数降不掉共同部分")
    ax.legend(frameon=False)
    fig.suptitle("M12 权重与分散：公式，不是回测")
    fig.tight_layout()
    fig.savefig(OUT, dpi=140)
    print("wrote", OUT)


if __name__ == "__main__":
    main()