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

from __future__ import annotations

from pathlib import Path

import pandas as pd

from dragon_branch_configs import alpha_first_selective_veto_config
from dragon_strategy import DragonRuleEngine
from dragon_strategy_config import StrategyConfig


def _load_indicator_snapshot(base_dir: Path) -> pd.DataFrame:
    df = pd.read_csv(base_dir / "dragon_indicator_snapshot.csv", encoding="utf-8-sig")
    df["date"] = pd.to_datetime(df["date"])
    return df.set_index("date", drop=False)


def _load_true_trade_events(base_dir: Path) -> pd.DataFrame:
    return pd.read_csv(base_dir / "true_trade_events.csv", encoding="utf-8-sig")


def _profit_factor(series: pd.Series) -> float:
    gross_profit = series[series > 0].sum()
    gross_loss = -series[series < 0].sum()
    if gross_loss == 0:
        return float("inf") if gross_profit > 0 else 0.0
    return float(gross_profit / gross_loss)


def _holding_bucket(days: int) -> str:
    if days <= 5:
        return "00-05d"
    if days <= 10:
        return "06-10d"
    if days <= 20:
        return "11-20d"
    if days <= 40:
        return "21-40d"
    return "41d+"


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


def _run_branch(
    label: str,
    config: StrategyConfig,
    indicator_df: pd.DataFrame,
    workbook_events: pd.DataFrame,
    first_date: str,
    last_date: str,
) -> tuple[dict[str, object], pd.DataFrame]:
    engine = DragonRuleEngine(config)
    events, trades = engine.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()
    trades["holding_bucket"] = trades["holding_days"].astype(int).map(_holding_bucket)
    short = trades[trades["holding_bucket"].isin({"00-05d", "06-10d"})].copy()
    glued_short = short[short["buy_reason"] == "glued_buy"].copy()

    buy_overlap, buy_extra = _event_match(events, workbook_events, "BUY")
    sell_overlap, sell_extra = _event_match(events, workbook_events, "SELL")

    row = {
        "experiment": label,
        "trades": int(len(trades)),
        "win_rate": float((trades["return_pct"] > 0).mean()) if not trades.empty else float("nan"),
        "avg_return": float(trades["return_pct"].mean()) if not trades.empty else float("nan"),
        "median_return": float(trades["return_pct"].median()) if not trades.empty else float("nan"),
        "profit_factor": _profit_factor(trades["return_pct"]) if not trades.empty else float("nan"),
        "real_buy_overlap": int(buy_overlap),
        "real_buy_extra": int(buy_extra),
        "real_sell_overlap": int(sell_overlap),
        "real_sell_extra": int(sell_extra),
        "short_trade_count": int(len(short)),
        "short_avg_return": float(short["return_pct"].mean()) if not short.empty else float("nan"),
        "short_00_05d_avg_return": float(short[short["holding_bucket"] == "00-05d"]["return_pct"].mean()),
        "short_06_10d_avg_return": float(short[short["holding_bucket"] == "06-10d"]["return_pct"].mean()),
        "glued_short_count": int(len(glued_short)),
        "glued_short_avg_return": float(glued_short["return_pct"].mean()) if not glued_short.empty else float("nan"),
        "post_sell_rebound_short_count": int(short[short["buy_reason"] == "post_sell_rebound_buy"].shape[0]),
    }
    diff = trades[trades["holding_bucket"].isin({"00-05d", "06-10d"})][
        ["buy_date", "sell_date", "buy_reason", "sell_reason", "holding_days", "return_pct"]
    ].copy()
    diff["experiment"] = label
    return row, diff


def main() -> None:
    base_dir = Path(__file__).resolve().parent
    indicator_df = _load_indicator_snapshot(base_dir)
    workbook_events = _load_true_trade_events(base_dir)
    first_date = workbook_events["date"].min()
    last_date = workbook_events["date"].max()

    baseline = alpha_first_selective_veto_config()
    experiments = [
        ("baseline_alpha_first", baseline),
        ("disable_post_sell_rebound_buy", baseline.with_updates(disabled_rules={"post_sell_rebound_buy"})),
        (
            "glued_veto_hot_positive_b1",
            baseline.with_updates(
                glued_selective_hot_c1_min=40.0,
                glued_selective_hot_b1_min=0.10,
            ),
        ),
        (
            "glued_veto_low_weak_range",
            baseline.with_updates(
                glued_selective_low_c1_min=23.0,
                glued_selective_low_c1_max=28.0,
                glued_selective_low_b1_max=0.02,
            ),
        ),
        (
            "glued_veto_hot_and_low",
            baseline.with_updates(
                glued_selective_hot_c1_min=40.0,
                glued_selective_hot_b1_min=0.10,
                glued_selective_low_c1_min=23.0,
                glued_selective_low_c1_max=28.0,
                glued_selective_low_b1_max=0.02,
            ),
        ),
        (
            "glued_veto_hot_low_and_disable_post_sell",
            baseline.with_updates(
                disabled_rules={"post_sell_rebound_buy"},
                glued_selective_hot_c1_min=40.0,
                glued_selective_hot_b1_min=0.10,
                glued_selective_low_c1_min=23.0,
                glued_selective_low_c1_max=28.0,
                glued_selective_low_b1_max=0.02,
            ),
        ),
    ]

    rows: list[dict[str, object]] = []
    diffs: list[pd.DataFrame] = []
    for label, config in experiments:
        row, diff = _run_branch(label, config, indicator_df, workbook_events, first_date, last_date)
        rows.append(row)
        diffs.append(diff)

    result_df = pd.DataFrame(rows)
    baseline_row = result_df[result_df["experiment"] == "baseline_alpha_first"].iloc[0]
    for col in [
        "trades",
        "win_rate",
        "avg_return",
        "median_return",
        "profit_factor",
        "real_buy_overlap",
        "real_sell_overlap",
        "short_trade_count",
        "short_avg_return",
        "short_00_05d_avg_return",
        "short_06_10d_avg_return",
        "glued_short_count",
        "glued_short_avg_return",
        "post_sell_rebound_short_count",
    ]:
        result_df[f"delta_{col}"] = result_df[col] - baseline_row[col]

    diff_df = pd.concat(diffs, ignore_index=True)
    result_df.to_csv(base_dir / "dragon_short_holding_experiments.csv", index=False, encoding="utf-8-sig")
    diff_df.to_csv(base_dir / "dragon_short_holding_experiment_trades.csv", index=False, encoding="utf-8-sig")

    lines = [
        "# Dragon Short Holding Experiments",
        "",
        "- Baseline branch: `alpha_first_selective_veto`.",
        "- Goal: reduce the main short-holding drag with the smallest possible extra complexity.",
        "",
        "## Summary",
    ]
    for _, row in result_df.iterrows():
        lines.append(
            f"- `{row['experiment']}`: trades `{int(row['trades'])}`, avg_return `{row['avg_return']:.2%}`, "
            f"profit_factor `{row['profit_factor']:.2f}`, short_avg_return `{row['short_avg_return']:.2%}`, "
            f"`00-05d` `{row['short_00_05d_avg_return']:.2%}`, `06-10d` `{row['short_06_10d_avg_return']:.2%}`, "
            f"real BUY / SELL `{int(row['real_buy_overlap'])}/{int(row['real_sell_overlap'])}`"
        )

    lines.extend(["", "## Delta Vs Alpha-First Baseline"])
    for _, row in result_df[result_df["experiment"] != "baseline_alpha_first"].iterrows():
        lines.append(
            f"- `{row['experiment']}`: delta_avg_return `{row['delta_avg_return']:.2%}`, "
            f"delta_profit_factor `{row['delta_profit_factor']:.2f}`, delta_short_avg_return `{row['delta_short_avg_return']:.2%}`, "
            f"delta_glued_short_avg_return `{row['delta_glued_short_avg_return']:.2%}`, "
            f"real BUY / SELL `{int(row['real_buy_overlap'])}/{int(row['real_sell_overlap'])}`"
        )

    best = result_df[result_df["experiment"] != "baseline_alpha_first"].sort_values(
        ["avg_return", "profit_factor"], ascending=[False, False]
    ).head(1)
    if not best.empty:
        row = best.iloc[0]
        lines.extend(
            [
                "",
                "## Quant Judgment",
                f"- Best branch in this pack: `{row['experiment']}` with avg_return `{row['avg_return']:.2%}` and profit_factor `{row['profit_factor']:.2f}`.",
                "- Compare the winning branch to the audit: if glued short trades fall materially while overlap loss stays controlled, the next optimization should stay on the glued-entry side.",
                "- If disabling `post_sell_rebound_buy` contributes little, that family is secondary for this stage.",
            ]
        )

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


if __name__ == "__main__":
    main()