cyb50-quant/research/dragon/v2/dragon_predictive_break_experiments.py

from __future__ import annotations

from pathlib import Path

import pandas as pd

from dragon_indicators import DragonIndicatorConfig, DragonIndicatorEngine
from dragon_strategy import DragonRuleEngine
from dragon_strategy_config import StrategyConfig
from dragon_workbook import DragonWorkbook


def _find_workbook(base_dir: Path) -> Path:
    matches = sorted(base_dir.glob("*.xlsx"))
    if not matches:
        raise FileNotFoundError(f"No workbook found in {base_dir}")
    return matches[0]


def _load_workbook_events(workbook_path: Path) -> pd.DataFrame:
    workbook = DragonWorkbook(workbook_path)
    return pd.DataFrame([{"date": e.date.isoformat(), "side": e.side, "layer": e.layer} for e in workbook.split_layers()])


def _overlap(workbook_events: pd.DataFrame, strategy_events: pd.DataFrame, side: str, layer: str) -> tuple[int, int]:
    wb = set(workbook_events[(workbook_events["side"] == side) & (workbook_events["layer"] == layer)]["date"])
    st = set(strategy_events[(strategy_events["side"] == side) & (strategy_events["layer"] == layer)]["date"])
    return len(wb & st), len(st - wb)


def _run(label: str, config: StrategyConfig, workbook_events: pd.DataFrame, indicator_df: pd.DataFrame, first_date: str, last_date: str) -> dict[str, object]:
    events, trades = DragonRuleEngine(config=config).run(indicator_df)
    events = events[(events["date"] >= first_date) & (events["date"] <= last_date)].copy()
    trades = trades[
        (trades["buy_date"] >= first_date)
        & (trades["buy_date"] <= last_date)
        & (trades["sell_date"] >= first_date)
        & (trades["sell_date"] <= last_date)
    ].copy()

    buy_overlap, buy_extra = _overlap(workbook_events, events, "BUY", "real_trade")
    sell_overlap, sell_extra = _overlap(workbook_events, events, "SELL", "real_trade")
    predictive = trades[trades["sell_reason"] == "predictive_b1_break_exit"]

    return {
        "experiment": label,
        "short_b1_max": config.predictive_b1_break_short_b1_max,
        "long_b1_max": config.predictive_b1_break_long_b1_max,
        "trades": int(len(trades)),
        "avg_return": float(trades["return_pct"].mean()),
        "real_buy_overlap": int(buy_overlap),
        "real_sell_overlap": int(sell_overlap),
        "real_buy_extra": int(buy_extra),
        "real_sell_extra": int(sell_extra),
        "predictive_trade_count": int(len(predictive)),
    }


def main() -> None:
    base_dir = Path(__file__).resolve().parent
    workbook_events = _load_workbook_events(_find_workbook(base_dir))
    first_date = pd.to_datetime(workbook_events["date"]).min().date().isoformat()
    last_date = pd.to_datetime(workbook_events["date"]).max().date().isoformat()

    engine = DragonIndicatorEngine(DragonIndicatorConfig(start_date="2015-01-01", end_date="2026-01-31"))
    indicator_df = engine.compute(engine.fetch_daily_data())

    baseline = StrategyConfig()
    experiments = [
        ("baseline", baseline),
        ("short_b1_looser", baseline.with_updates(predictive_b1_break_short_b1_max=-0.11)),
        ("short_b1_tighter", baseline.with_updates(predictive_b1_break_short_b1_max=-0.15)),
        ("long_b1_looser", baseline.with_updates(predictive_b1_break_long_b1_max=-0.10)),
        ("long_b1_tighter", baseline.with_updates(predictive_b1_break_long_b1_max=-0.14)),
    ]

    rows = [_run(label, cfg, workbook_events, indicator_df, first_date, last_date) for label, cfg in experiments]
    df = pd.DataFrame(rows)
    base_row = df[df["experiment"] == "baseline"].iloc[0]
    for col in ["avg_return", "real_buy_overlap", "real_sell_overlap", "predictive_trade_count"]:
        df[f"delta_{col}"] = df[col] - base_row[col]

    df.to_csv(base_dir / "dragon_predictive_break_experiments.csv", index=False, encoding="utf-8-sig")

    lines = [
        "# Dragon Predictive Break Experiments",
        "",
        f"- baseline predictive trade count: `{int(base_row['predictive_trade_count'])}`",
        f"- baseline real BUY / SELL overlap: `{int(base_row['real_buy_overlap'])}` / `{int(base_row['real_sell_overlap'])}`",
        "",
        "## Experiment Summary",
    ]
    for _, row in df.iterrows():
        lines.append(
            f"- `{row['experiment']}`: predictive trades `{int(row['predictive_trade_count'])}`, "
            f"delta_avg_return `{row['delta_avg_return']:.2%}`, real BUY `{int(row['real_buy_overlap'])}`, real SELL `{int(row['real_sell_overlap'])}`"
        )

    (base_dir / "dragon_predictive_break_experiments.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()