Bithumb/deepcoin/ground_truth/gt_signal_rules.py

"""
GT 모델(entry/exit)을 규칙 스캔·발화 형식으로 일반화.

ZigZag trough/peak + BB 필터 등 GT 타점 생성 로직과 동일 파라미터를
rule_eval 스캔 프레임 컬럼(gt_*)으로 노출합니다.
"""

from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd

from config import (
    GT_BUY_BB_MAX,
    GT_BUY_MIN_SWING_PCT,
    GT_MIN_SWING_PCT,
    GT_PIVOT_ORDER,
    MATCH_PRIMARY_INTERVAL,
)
from deepcoin.ground_truth.ground_truth import build_zigzag_pivots


def _local_extrema_mask(
    series: pd.Series,
    order: int,
    kind: str,
) -> pd.Series:
    """
    국소 극값 boolean 마스크.

    Args:
        series: 가격 시리즈.
        order: 좌우 봉 수.
        kind: min | max.

    Returns:
        boolean Series (index=series.index).
    """
    arr = series.astype(float).values
    n = len(arr)
    out = np.zeros(n, dtype=bool)
    if n < 2 * order + 1:
        return pd.Series(out, index=series.index)
    for i in range(order, n - order):
        window = arr[i - order : i + order + 1]
        if kind == "min" and arr[i] <= window.min():
            out[i] = True
        elif kind == "max" and arr[i] >= window.max():
            out[i] = True
    return pd.Series(out, index=series.index)


def enrich_scan_frame_gt_signals(
    frame: pd.DataFrame,
    *,
    pivot_order: int = GT_PIVOT_ORDER,
    buy_swing_pct: float = GT_BUY_MIN_SWING_PCT,
    sell_swing_pct: float = GT_MIN_SWING_PCT,
    bb_max: float = GT_BUY_BB_MAX,
    causal: bool | None = None,
) -> pd.DataFrame:
    """
    스캔 프레임에 GT 모델 신호 컬럼을 추가합니다.

    GT_SIGNAL_CAUSAL=1 이면 t 시점까지 데이터만 사용 (운영 정합).

    Args:
        frame: m3 스캔 프레임 (Low, High, bb_pos).
        pivot_order: 피벗 반경.
        buy_swing_pct: 매수 ZigZag 스윙%.
        sell_swing_pct: 매도 ZigZag 스윙%.
        bb_max: BB 하단 필터.
        causal: None이면 config GT_SIGNAL_CAUSAL.

    Returns:
        gt_* 컬럼이 추가된 DataFrame.
    """
    from config import GT_SIGNAL_CAUSAL

    use_causal = GT_SIGNAL_CAUSAL if causal is None else causal
    if use_causal:
        from deepcoin.ground_truth.gt_signal_causal import (
            enrich_scan_frame_gt_signals_causal,
        )

        return enrich_scan_frame_gt_signals_causal(
            frame,
            pivot_order=pivot_order,
            buy_swing_pct=buy_swing_pct,
            sell_swing_pct=sell_swing_pct,
            bb_max=bb_max,
        )

    out = frame.copy()
    if "Low" not in out.columns or "High" not in out.columns:
        return out

    low = out["Low"].astype(float)
    high = out["High"].astype(float)
    out["gt_trough_local"] = _local_extrema_mask(low, pivot_order, "min").astype(int)
    out["gt_peak_local"] = _local_extrema_mask(high, pivot_order, "max").astype(int)

    df_ohlc = out[["Low", "High"]].copy()
    if "close" in out.columns:
        df_ohlc["close"] = out["close"]
    df_ohlc.index = out.index

    buy_pivots = build_zigzag_pivots(
        df_ohlc,
        min_swing_pct=buy_swing_pct,
        pivot_order=pivot_order,
    )
    sell_pivots = build_zigzag_pivots(
        df_ohlc,
        min_swing_pct=sell_swing_pct,
        pivot_order=pivot_order,
    )

    trough_z = pd.Series(0, index=out.index, dtype=int)
    for p in buy_pivots:
        if p.kind == "trough" and p.ts in trough_z.index:
            trough_z.loc[p.ts] = 1
    peak_z = pd.Series(0, index=out.index, dtype=int)
    for p in sell_pivots:
        if p.kind == "peak" and p.ts in peak_z.index:
            peak_z.loc[p.ts] = 1

    out["gt_trough_zigzag"] = trough_z
    out["gt_peak_zigzag"] = peak_z

    bb_ok = pd.Series(True, index=out.index)
    if "bb_pos" in out.columns:
        bb = pd.to_numeric(out["bb_pos"], errors="coerce")
        bb_ok = bb <= bb_max

    out["gt_buy_signal"] = ((out["gt_trough_zigzag"] == 1) & bb_ok).astype(int)
    out["gt_sell_signal"] = (out["gt_peak_zigzag"] == 1).astype(int)
    return out


def build_gt_model_rules() -> list[dict[str, Any]]:
    """
    GT entry/exit 명세와 동일한 스캔 규칙 후보.

    Returns:
        rule dict 리스트 (buy 2종 + sell 2종).
    """
    return [
        {
            "rule_id": "gt_model_buy_zigzag_bb",
            "side": "buy",
            "kind": "gt_model",
            "logic": "and",
            "conditions": [
                {"col": "gt_buy_signal", "op": "eq_int", "value": 1},
            ],
            "gt_spec": "trough_zigzag + bb_pos <= GT_BUY_BB_MAX",
        },
        {
            "rule_id": "gt_model_buy_trough_local",
            "side": "buy",
            "kind": "gt_model",
            "logic": "and",
            "conditions": [
                {"col": "gt_trough_local", "op": "eq_int", "value": 1},
                {"col": "bb_pos", "op": "lte", "value": GT_BUY_BB_MAX},
            ],
            "gt_spec": "local trough + bb filter",
        },
        {
            "rule_id": "gt_model_sell_zigzag_peak",
            "side": "sell",
            "kind": "gt_model",
            "logic": "and",
            "conditions": [
                {"col": "gt_sell_signal", "op": "eq_int", "value": 1},
            ],
            "gt_spec": "major swing peak (ZigZag)",
        },
        {
            "rule_id": "gt_model_sell_peak_local",
            "side": "sell",
            "kind": "gt_model",
            "logic": "and",
            "conditions": [
                {"col": "gt_peak_local", "op": "eq_int", "value": 1},
            ],
            "gt_spec": "local high extremum",
        },
    ]


def gt_signal_rule_ids() -> set[str]:
    """GT 일반화 규칙 ID 집합."""
    return {r["rule_id"] for r in build_gt_model_rules()}