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

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable


@dataclass(frozen=True)
class CandidateSignal:
    date: str
    side: str
    source: str
    reason: str


@dataclass(frozen=True)
class LayeredSignal:
    date: str
    side: str
    layer: str
    source: str
    reason: str


class DragonStateMachine:
    """
    Baseline signal splitter.

    This intentionally does not encode the full Dragon trading rules yet.
    It only turns upstream candidate entry/exit signals into:
    - real trades
    - auxiliary bullish/bearish signals
    """

    def split(self, candidates: Iterable[CandidateSignal]) -> list[LayeredSignal]:
        state = "flat"
        layered: list[LayeredSignal] = []
        for signal in candidates:
            layer = "aux_signal"
            reason = signal.reason
            if state == "flat" and signal.side == "BUY":
                layer = "real_trade"
                state = "long"
            elif state == "long" and signal.side == "SELL":
                layer = "real_trade"
                state = "flat"
            elif state == "flat" and signal.side == "SELL":
                reason = "bearish_signal_after_exit"
            elif state == "long" and signal.side == "BUY":
                reason = "bullish_signal_while_holding"
            layered.append(
                LayeredSignal(
                    date=signal.date,
                    side=signal.side,
                    layer=layer,
                    source=signal.source,
                    reason=reason,
                )
            )
        return layered