"""M21 教学图：截断不变性，以及三本账净值差的瀑布拆解。

价格、分红、费率和净值都是教学输入，不是行情，也不是券商对账单。
左图比较「把全样本分位数用在过去」和「只向前滚动」的信号差。
右图把影子账到成本调整账的差额拆开，残差单独留下。
"""

from __future__ import annotations

import math
from decimal import ROUND_CEILING, ROUND_DOWN, Decimal
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

from quant.backtest.fees import FeeSchedule, FixedFill, apply_fixed_fills
from quant.core.types import Side

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

OUT = Path(__file__).resolve().parent / "truncation_and_bridge.png"
CENT = Decimal("0.01")
# 第 6 天是 1 起计的第 6 个价格。截断样本只含这一天及更早的价格。
DAY6 = 6


def percentile_rank(sample: np.ndarray) -> np.ndarray:
    # 分母用样本长度。全样本分位会把后面的价格算进今天的秩。
    n = len(sample)
    return np.array([(sample <= x).sum() / n for x in sample], dtype=float)


def expanding_rank(prices: np.ndarray) -> np.ndarray:
    # 每一天的分母只含当天和更早的价格，所以截掉未来不会改过去。
    out = np.empty(len(prices), dtype=float)
    for i in range(len(prices)):
        out[i] = percentile_rank(prices[: i + 1])[-1]
    return out


def truncation_bias(signal_fn, prices: np.ndarray) -> np.ndarray:
    full = signal_fn(prices)
    bias = np.empty(len(prices), dtype=float)
    for i in range(len(prices)):
        short = signal_fn(prices[: i + 1])
        bias[i] = full[i] - short[-1]
    return bias


def mean_deviation(prices: np.ndarray) -> np.ndarray:
    # 减去喂入样本的均值。样本一变长，同一天的偏差就变。
    return prices - prices.mean()


def fee_cents(amount: Decimal) -> Decimal:
    # 费用向上取到分，避免把成本估小。这里的成交额都能整除，取整不改变数字。
    return amount.quantize(CENT, rounding=ROUND_CEILING)


def whole_lots(shares: Decimal, lot_size: Decimal) -> Decimal:
    # 不够一手的尾数不能下单。向下取整，不把现金凑成碎股。
    if shares <= 0:
        return Decimal(0)
    lots = (shares / lot_size).to_integral_value(rounding=ROUND_DOWN)
    return lots * lot_size


def teaching_fee(notional: Decimal, bps: Decimal = Decimal("10")) -> Decimal:
    if notional == 0:
        return Decimal(0)
    return fee_cents(notional * bps / Decimal(10000))


def golden_fixture() -> list[dict[str, Decimal]]:
    """五天夹具。信号在收盘形成，权重到下一交易日才成交，引擎里不再 shift。"""
    closes = [
        Decimal("100"),
        Decimal("110"),
        Decimal("100"),
        Decimal("110"),
        Decimal("121"),
    ]
    # 除息日现金只付给开盘前已经持有的股份。当天收盘才买入的不拿这笔分红。
    dividends = [Decimal(0), Decimal(0), Decimal("10"), Decimal(0), Decimal(0)]
    tri = [Decimal(1)]
    for i in range(1, len(closes)):
        factor = (closes[i] + dividends[i]) / closes[i - 1]
        tri.append(tri[-1] * factor)
    next_weight = [Decimal(0)]
    day_return = [Decimal(0)]
    for i in range(1, len(closes)):
        ret = tri[i] / tri[i - 1] - 1
        day_return.append(ret)
        next_weight.append(Decimal("0.50") if ret > 0 else Decimal(0))
    cash = Decimal("10000.00")
    shares = Decimal(0)
    rows: list[dict[str, Decimal]] = []
    for i, price in enumerate(closes):
        # 先结当天分红，再按「上一收盘决定的权重」用未复权价成交。
        div_cash = shares * dividends[i]
        cash += div_cash
        weight = Decimal(0) if i == 0 else next_weight[i - 1]
        nav_pre = cash + shares * price
        target_notional = weight * nav_pre
        raw_shares = target_notional / price if price else Decimal(0)
        target_shares = whole_lots(raw_shares, Decimal(10))
        delta = target_shares - shares
        notional = abs(delta) * price
        fee = teaching_fee(notional)
        if delta > 0:
            cash -= notional + fee
        elif delta < 0:
            cash += notional - fee
        shares = target_shares
        nav = cash + shares * price
        rows.append(
            {
                "close": price,
                "div": dividends[i],
                "tri": tri[i],
                "ret": day_return[i],
                "weight": weight,
                "raw_shares": raw_shares,
                "shares": shares,
                "fee": fee,
                "div_cash": div_cash,
                "cash": cash,
                "nav": nav,
            }
        )
    return rows


def bridge_levels() -> dict[str, Decimal]:
    # 四项加残差是恒等式。残差不拿去填平任何一项。
    shadow = Decimal("101200.00")
    signal = Decimal("-80.00")
    lot = Decimal("-150.00")
    execution = Decimal("-40.00")
    cost = Decimal("-90.00")
    epsilon = Decimal("0.20")
    actual = shadow + signal + lot + execution + cost + epsilon
    base = Decimal("100000")
    return {
        "shadow": shadow,
        "signal": signal,
        "lot": lot,
        "exec": execution,
        "cost": cost,
        "epsilon": epsilon,
        "actual": actual,
        "gap": actual - shadow,
        "gap_bps": (actual - shadow) / base * Decimal(10000),
        "eps_bps": epsilon / base * Decimal(10000),
        "alarm_bps": Decimal("12.00") / base * Decimal(10000),
    }


def sample_std(values: list[float]) -> float:
    arr = np.array(values, dtype=float)
    return float(arr.std(ddof=1))


def pearson(a: list[float], b: list[float]) -> float:
    return float(np.corrcoef(np.array(a), np.array(b))[0, 1])


def main() -> None:
    prices = np.arange(10, 20, dtype=float)
    pct_bias = truncation_bias(percentile_rank, prices)
    mean_bias = truncation_bias(mean_deviation, prices)
    causal_bias = truncation_bias(expanding_rank, prices)
    i = DAY6 - 1
    print("day6_pct_trunc", f"{percentile_rank(prices[:DAY6])[-1]:.6f}")
    print("day6_pct_full", f"{percentile_rank(prices)[i]:.6f}")
    print("day6_pct_bias", f"{pct_bias[i]:.6f}")
    print("day6_mean_trunc", f"{mean_deviation(prices[:DAY6])[-1]:.6f}")
    print("day6_mean_full", f"{mean_deviation(prices)[i]:.6f}")
    print("day6_mean_bias", f"{mean_bias[i]:.6f}")
    print("day6_causal_bias", f"{causal_bias[i]:.6f}")
    print("pct_curve", " ".join(f"{x:.6f}" for x in pct_bias))
    print("mean_curve", " ".join(f"{x:.6f}" for x in mean_bias))
    print("causal_curve", " ".join(f"{x:.6f}" for x in causal_bias))

    leaked = np.array([11.0, 100.0, 11.0])
    # 长表按日期交错时，不分组的 shift 会把另一只的价格当成自己的上一笔。
    print("bug1_wrong", f"{11.0 / leaked[1] - 1:.6f}")
    print("bug1_right", f"{11.0 / 10.0 - 1:.6f}")
    # 暖机期被写成 0 收益时，会把真正持仓日的波动拉低，而不是报错。
    held = [0.10, -0.10, 0.10, -0.10]
    print("bug2_wrong_std", f"{sample_std([0.0, 0.0, *held]):.6f}")
    print("bug2_right_std", f"{sample_std(held):.6f}")
    raw_w = [-2.0, -1.0, 0.5]
    kept = [max(x, 0.0) for x in raw_w]
    wrong_w = [x / sum(kept) for x in kept]
    # 名额按原组合的只数保留。熊市里只剩一只要正动量时，空位是现金，不重新归一成 1。
    print("bug4_wrong", f"{wrong_w[-1]:.6f}")
    print("bug4_right", f"{(1 / len(raw_w)):.6f}")
    cash_free, pending_sell, px, wanted = 1000, 1000, 20, 80
    wrong_shares = wanted if cash_free + pending_sell >= wanted * px else 0
    right_shares = min(wanted, cash_free // px)
    print("bug5_wrong_shares", str(wrong_shares))
    print("bug5_right_shares", str(right_shares))
    signal_path = [1.0, 2.0, 3.0]
    print("bug8_leaked_pos", f"{signal_path[1]:.6f}")
    print("bug8_ic_leaked", f"{pearson([0.01, -0.02, 0.03], [0.01, -0.02, 0.03]):.6f}")
    causal_ic = pearson([0.2, -0.1, 0.0], [0.01, -0.02, 0.04])
    print("bug8_ic_causal", f"{causal_ic:.6f}")
    print("bug10_leaked_corr", f"{1.0:.6f}")
    sr = 0.01 / 0.04 * math.sqrt(12)
    sr_double = (0.01 - 0.04 / 12) / 0.04 * math.sqrt(12)
    sr_day = 0.0004 / 0.01 * math.sqrt(252)
    sr_month_scale = 0.0004 / 0.01 * math.sqrt(12)
    print("bug7_sharpe", f"{sr:.6f}")
    print("bug7_double_rf", f"{sr_double:.6f}")
    print("bug7_daily", f"{sr_day:.6f}")
    print("bug7_wrong_scale", f"{sr_month_scale:.6f}")
    float_sum = sum(0.01 for _ in range(100))
    dec_sum = Decimal("0.01") * 100
    print("bug9_float", format(float_sum, ".16f"))
    print("bug9_decimal", f"{dec_sum:.2f}")

    for n, row in enumerate(golden_fixture(), start=1):
        print(
            f"day{n}",
            f"close {row['close']:.2f}",
            f"div {row['div']:.2f}",
            f"tri {row['tri']:.6f}",
            f"ret {row['ret']:.6f}",
            f"w {row['weight']:.2f}",
            f"raw {row['raw_shares']:.6f}",
            f"shares {row['shares']:.0f}",
            f"fee {row['fee']:.2f}",
            f"div_cash {row['div_cash']:.2f}",
            f"cash {row['cash']:.2f}",
            f"nav {row['nav']:.2f}",
        )
    # 除息日收盘才买到的 50 股，不能再领 10 元现金分红。总收益指数里已经有这 10 元。
    print("bug6_extra_div", f"{Decimal('50') * Decimal('10'):.2f}")

    fees = FeeSchedule(
        commission_bps=Decimal("10"),
        min_commission=Decimal("0"),
        platform_fee=Decimal("0"),
        label="teaching-not-a-broker-schedule",
    )
    fills = (
        FixedFill(Side.BUY, Decimal("50"), Decimal("100")),
        FixedFill(Side.SELL, Decimal("50"), Decimal("110")),
    )
    nav_1x = apply_fixed_fills(Decimal("10000"), fills, fees, Decimal("110"))
    nav_2x = apply_fixed_fills(Decimal("10000"), fills, fees.scaled(Decimal("2")), Decimal("110"))
    print("nav_1x", f"{nav_1x:.2f}")
    print("nav_2x", f"{nav_2x:.2f}")
    print("nav_drop", f"{nav_2x - nav_1x:.2f}")

    levels = bridge_levels()
    for key in (
        "shadow",
        "signal",
        "lot",
        "exec",
        "cost",
        "epsilon",
        "actual",
        "gap",
        "gap_bps",
        "eps_bps",
        "alarm_bps",
    ):
        value = levels[key]
        money = {"shadow", "signal", "lot", "exec", "cost", "epsilon", "actual", "gap"}
        places = 2 if key in money else 6
        print(key, f"{value:.{places}f}")
    drag = 2 * Decimal("0.40") * Decimal("0.0010")
    s4_drag = 2 * Decimal("2.90") * Decimal("0.0005")
    print("drag_teach", f"{drag:.6f}")
    print("drag_s4_commission", f"{s4_drag:.6f}")
    print("s4_cost_gap_pp", f"{15.77 - 14.84:.2f}")
    loss = Decimal("0.96") * Decimal("0.97") * Decimal("1.01") - 1
    print("loss_window", f"{loss:.6f}")
    print("unfilled_rate", f"{Decimal('2000') / Decimal('8000'):.6f}")

    fig, axes = plt.subplots(1, 2, figsize=(11.6, 5.0))
    days = np.arange(1, 11)
    ax = axes[0]
    ax.plot(days, pct_bias, color="#a33b3b", marker="o", label="未来分位数偏差")
    ax.plot(days, causal_bias, color="#1f4e79", linewidth=2, label="向前滚动偏差")
    ax.scatter([DAY6], [pct_bias[i]], color="#a33b3b", s=60, zorder=3)
    ax.annotate(
        "第6天 −0.400000",
        xy=(DAY6, pct_bias[i]),
        xytext=(DAY6 + 0.3, pct_bias[i] - 0.18),
        fontsize=9,
        color="#a33b3b",
    )
    ax.axhline(0.0, color="#333333", linewidth=0.6)
    ax.set_xticks(days)
    ax.set_xlabel("交易日（教学价格 10 到 19）")
    ax.set_ylabel("全样本信号减截断信号")
    ax.set_title("截断不变性：分位泄漏 vs 零偏差")
    ax.legend(frameon=False, loc="lower right")

    ax = axes[1]
    labels = ["影子账", "信号", "整手", "让价", "费用未成交", "残差", "成本调整账"]
    shadow = float(levels["shadow"])
    deltas = [
        float(levels["signal"]),
        float(levels["lot"]),
        float(levels["exec"]),
        float(levels["cost"]),
        float(levels["epsilon"]),
    ]
    running = shadow
    bottoms = [100700.0]
    heights = [shadow - 100700.0]
    colors = ["#1f4e79"]
    for delta in deltas:
        if delta >= 0:
            bottoms.append(running)
            heights.append(delta)
            colors.append("#5b8e7d")
        else:
            bottoms.append(running + delta)
            heights.append(-delta)
            colors.append("#a33b3b")
        running += delta
    bottoms.append(100700.0)
    heights.append(float(levels["actual"]) - 100700.0)
    colors.append("#1f4e79")
    xpos = np.arange(len(labels))
    ax.bar(xpos, heights, bottom=bottoms, color=colors, width=0.72)
    cursor = shadow
    for x, delta in enumerate(deltas, start=1):
        ax.plot([x - 1.36, x - 0.36], [cursor, cursor], color="#666666", linewidth=0.7)
        cursor += delta
    notes = [
        "101200.00",
        "−80.00",
        "−150.00",
        "−40.00",
        "−90.00",
        "+0.20",
        "100840.20",
    ]
    for x, bottom, height, text in zip(xpos, bottoms, heights, notes, strict=True):
        # 标在柱顶之上，避免黑字压在红柱里。
        ax.text(x, bottom + height + 14, text, ha="center", va="bottom", fontsize=8)
    ax.set_ylim(100700, 101480)
    ax.set_xticks(xpos, labels, fontsize=8)
    ax.set_ylabel("教学净值")
    ax.set_title("月初净值 100000 的四项拆解")
    fig.suptitle("M21：截断偏差与三本账瀑布；数字是教学输入，不是历史回测")
    fig.tight_layout()
    fig.savefig(OUT, dpi=140)
    print("wrote", OUT.name)


if __name__ == "__main__":
    main()