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

from __future__ import annotations

from typing import Any, Optional

import pandas as pd

from dragon_rule_catalog import classify_entry_reason, classify_exit_reason


def _event_payload(row: pd.Series, side: str, layer: str, reason: str) -> dict[str, object]:
    return {
        "date": row.name.date().isoformat(),
        "side": side,
        "layer": layer,
        "reason": reason,
        "close": float(row["close"]),
        "a1": float(row["a1"]),
        "b1": float(row["b1"]),
        "c1": float(row["c1"]),
        "kdj_buy": bool(row["kdj_buy"]),
        "kdj_sell": bool(row["kdj_sell"]),
        "ql_buy": bool(row["ql_buy"]),
        "ql_sell": bool(row["ql_sell"]),
    }


def _trade_payload(engine: Any, row: pd.Series, reason: str) -> dict[str, object]:
    return {
        "buy_date": engine.context.entry_date.isoformat() if engine.context.entry_date else "",
        "buy_price": engine.context.entry_price,
        "buy_reason": engine.context.entry_reason,
        "sell_date": row.name.date().isoformat(),
        "sell_price": float(row["close"]),
        "sell_reason": reason,
        "holding_days": engine._holding_days(row.name.date()),
        "return_pct": (
            float(row["close"]) / engine.context.entry_price - 1
            if engine.context.entry_price
            else None
        ),
    }


def _enrich_reason_metadata(engine: Any, events_df: pd.DataFrame, trades_df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    if not events_df.empty:
        event_decisions = [
            engine._build_decision(
                side=str(row["side"]),
                layer=str(row["layer"]),
                reason=str(row["reason"]),
            )
            for _, row in events_df.iterrows()
        ]
        events_df["reason_layer"] = [d.reason.layer.value if d.reason is not None else "unknown" for d in event_decisions]
        events_df["reason_family"] = [d.reason.family.value if d.reason is not None else "unknown" for d in event_decisions]
        events_df["reason_code"] = [d.reason.code if d.reason is not None else "" for d in event_decisions]

    if not trades_df.empty:
        buy_meta = trades_df["buy_reason"].map(classify_entry_reason)
        sell_meta = trades_df["sell_reason"].map(classify_exit_reason)
        trades_df["buy_reason_layer"] = [meta.layer.value for meta in buy_meta]
        trades_df["buy_reason_family"] = [meta.family.value for meta in buy_meta]
        trades_df["buy_reason_code"] = [meta.code for meta in buy_meta]
        trades_df["sell_reason_layer"] = [meta.layer.value for meta in sell_meta]
        trades_df["sell_reason_family"] = [meta.family.value for meta in sell_meta]
        trades_df["sell_reason_code"] = [meta.code for meta in sell_meta]

    return events_df, trades_df


def run_compat_execution(engine: Any, df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Execute the legacy-compatible strategy loop.

    The engine object is expected to provide the same internal hooks as
    DragonRuleEngine. This keeps the behavior path unchanged while physically
    decoupling execution runtime from strategy rule declarations.
    """

    engine.context = engine.context.__class__()
    events: list[dict[str, object]] = []
    trades: list[dict[str, object]] = []
    prev_row: Optional[pd.Series] = None

    for _, row in df.iterrows():
        engine._record_cross_counters(row)
        engine._update_position_counters(row)
        engine._update_pending_states(row)
        just_bought = False

        if (not engine.context.in_position) or bool(row["kdj_buy"] or row["ql_buy"]):
            action, reason = engine._buy_decision(row, prev_row)
            if action == "BUY":
                engine._post_real_buy(row, reason)
                just_bought = True
                events.append(_event_payload(row, side="BUY", layer="real_trade", reason=reason))
            elif action == "AUX_BUY":
                engine._post_aux_buy(row)
                events.append(_event_payload(row, side="BUY", layer="aux_signal", reason=reason))

        state_aux_sell_candidate = (not engine.context.in_position) and engine._should_emit_state_aux_sell(row)
        if not just_bought and (engine.context.in_position or bool(row["kdj_sell"] or row["ql_sell"]) or state_aux_sell_candidate):
            if engine.context.in_position and bool(row["kdj_sell"] or row["ql_sell"]):
                engine.context.sell_signal_count += 1
                if bool(row["kdj_sell"]):
                    engine.context.kdj_sell_signal_count += 1
                if bool(row["ql_sell"]):
                    engine.context.ql_sell_signal_count += 1
                if float(row["b1"]) < 0:
                    engine.context.b1_negative_sell_count += 1
            action, reason = engine._sell_decision(row, prev_row)
            if action == "SELL":
                trades.append(_trade_payload(engine=engine, row=row, reason=reason))
                engine.context.first_exit_checked = True
                events.append(_event_payload(row, side="SELL", layer="real_trade", reason=reason))
                engine._post_real_sell(row, reason)
            elif action == "AUX_SELL":
                if engine.context.in_position:
                    engine.context.first_exit_checked = True
                engine._post_aux_sell(row, reason)
                events.append(_event_payload(row, side="SELL", layer="aux_signal", reason=reason))

        prev_row = row

    events_df = pd.DataFrame(events)
    trades_df = pd.DataFrame(trades)
    return _enrich_reason_metadata(engine=engine, events_df=events_df, trades_df=trades_df)