"""M18 教学图：除息后的三种收益，以及价位、印花税和假设价差的基点。

公司行动调用仓库的总收益指数。价位和汇率是教学输入，不是某一天的行情。
"""

from __future__ import annotations

from decimal import Decimal
from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd

from quant.backtest.money import whole_lots
from quant.data.corporate_actions import return_factors

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

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


def main() -> None:
    lots = whole_lots(Decimal("250"), Decimal("100"))
    print("lots", lots, "left", Decimal("250") - lots)

    index = pd.to_datetime(["2024-01-02", "2024-01-03"])
    close = pd.Series([100.0, 98.0], index=index)
    actions = pd.DataFrame(
        {"ex_date": [index[1]], "kind": ["dividend"], "value": [2.0]}
    )
    factor = float(return_factors(close, actions).iloc[1])
    raw = 98.0 / 100.0 - 1.0
    after_tax = (98.0 + 2.0 * 0.70) / 100.0 - 1.0
    print("raw", f"{raw:.6f}")
    print("tr_factor", f"{factor:.6f}")
    print("after_tax", f"{after_tax:.6f}")

    split_close = pd.Series([100.0, 50.0], index=index)
    split_actions = pd.DataFrame(
        {"ex_date": [index[1]], "kind": ["split"], "value": [2.0]}
    )
    split_factor = float(return_factors(split_close, split_actions).iloc[1])
    print("split_factor", f"{split_factor:.6f}")

    notional = 20_000.0
    stamp_side = notional * 0.001
    print("stamp_side", f"{stamp_side:.2f}")
    print("stamp_round_trip_bps", f"{2.0 * stamp_side / notional * 10_000:.1f}")
    print("tick_bps", f"{0.005 / 0.40 * 10_000:.1f}")

    implied = 400.0 / 7.80
    gap = 52.0 / implied - 1.0
    print("implied_usd", f"{implied:.6f}")
    print("gap", f"{gap:.6f}")

    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.8))
    ax = axes[0]
    labels = ["未复权", "总收益", "预提 30% 后"]
    heights = [raw * 100.0, (factor - 1.0) * 100.0, after_tax * 100.0]
    ax.bar(labels, heights, color=["#a33b3b", "#1f4e79", "#c4a35a"])
    ax.axhline(0.0, color="#333333", linewidth=0.8)
    ax.set_ylabel("一日收益（%）")
    ax.set_title("收盘 100 到 98，股息 2")
    ax = axes[1]
    costs = [0.005 / 0.40 * 10_000, 20.0, gap * 10_000]
    ax.bar(["一个价位", "买卖印花税", "假设价差"], costs, color=["#a33b3b", "#1f4e79", "#c4a35a"])
    ax.set_ylabel("基点")
    ax.set_title("价 0.40、税率 0.1%、400 港元对 52 美元")
    fig.suptitle("M18：公司行动与摩擦；右图价差不是可成交的套利")
    fig.tight_layout()
    fig.savefig(OUT, dpi=140)
    print("wrote", OUT)


if __name__ == "__main__":
    main()