Bithumb/deepcoin/ground_truth/gt_signal_causal.py

"""
인과적(미래 미사용) GT 스타일 신호 — t봉 시점에 t 이하 데이터만 사용.

ZigZag/국소극값: pivot bar i-order 는 bar i 에서 확정 (i-order..i 구간만 관측).
"""

from __future__ import annotations

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,
)


def _confirmed_trough_mask(
    low: np.ndarray,
    order: int,
) -> np.ndarray:
    """
    bar i 에서 i-order 봉이 저점임을 확정 (low[i-order:i+1] 만 사용).

    Args:
        low: Low 가격 배열.
        order: pivot 반경(봉).

    Returns:
        길이 n, i 에 1이면 i 시점 매수 확인 신호.
    """
    n = len(low)
    out = np.zeros(n, dtype=np.int8)
    for i in range(2 * order, n):
        p = i - order
        seg = low[p - order : i + 1]
        if len(seg) == 0:
            continue
        if low[p] <= seg.min() + 1e-12:
            out[i] = 1
    return out


def _confirmed_peak_mask(
    high: np.ndarray,
    order: int,
) -> np.ndarray:
    """
    bar i 에서 i-order 봉이 고점임을 확정.

    Args:
        high: High 가격 배열.
        order: pivot 반경.

    Returns:
        i 시점 매도 확인 신호.
    """
    n = len(high)
    out = np.zeros(n, dtype=np.int8)
    for i in range(2 * order, n):
        p = i - order
        seg = high[p - order : i + 1]
        if len(seg) == 0:
            continue
        if high[p] >= seg.max() - 1e-12:
            out[i] = 1
    return out


def _zigzag_filter_causal(
    confirm: np.ndarray,
    prices: np.ndarray,
    min_swing_pct: float,
    kind: str,
    pivot_order: int = GT_PIVOT_ORDER,
) -> np.ndarray:
    """
    확정 피벗에 ZigZag 최소 스윙% 필터 (인과적, 순차 갱신).

    Args:
        confirm: bar i 에 확정 플래그.
        prices: pivot 가격 (i-order 위치의 low/high).
        pivot_indices: confirm==1 인 bar index.
        min_swing_pct: 최소 스윙 %.
        kind: trough | peak.

    Returns:
        zigzag 통과 시점에 1.
    """
    n = len(confirm)
    out = np.zeros(n, dtype=np.int8)
    order = int(pivot_order)
    last_kind: str | None = None
    last_price = 0.0
    min_ratio = min_swing_pct / 100.0

    for i in range(n):
        if confirm[i] != 1:
            continue
        p = i - order
        if p < 0:
            continue
        price = float(prices[p])
        if last_kind is None:
            out[i] = 1
            last_kind = kind
            last_price = price
            continue
        if kind == last_kind:
            if kind == "trough" and price < last_price:
                out[i - 1] = 0
                out[i] = 1
                last_price = price
            elif kind == "peak" and price > last_price:
                out[i - 1] = 0
                out[i] = 1
                last_price = price
            continue
        move = abs(price - last_price) / max(last_price, 1e-9)
        if move >= min_ratio:
            out[i] = 1
            last_kind = kind
            last_price = price
    return out


def enrich_scan_frame_gt_signals_causal(
    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,
) -> pd.DataFrame:
    """
    인과적 GT 신호 컬럼 (gt_*). t 시점 신호는 데이터 index<=t 만 사용.

    Args:
        frame: m3 스캔 프레임.
        pivot_order: 확정 지연(봉).
        buy_swing_pct: 매수 ZigZag 스윙%.
        sell_swing_pct: 매도 ZigZag 스윙%.
        bb_max: BB 하단 필터.

    Returns:
        gt_* 컬럼 추가 DataFrame.
    """
    out = frame.copy()
    if "Low" not in out.columns or "High" not in out.columns:
        return out

    low = out["Low"].astype(float).values
    high = out["High"].astype(float).values
    n = len(low)

    trough_conf = _confirmed_trough_mask(low, pivot_order)
    peak_conf = _confirmed_peak_mask(high, pivot_order)

    trough_z = _zigzag_filter_causal(
        trough_conf, low, buy_swing_pct, "trough", pivot_order=pivot_order
    )
    peak_z = _zigzag_filter_causal(
        peak_conf, high, sell_swing_pct, "peak", pivot_order=pivot_order
    )

    out["gt_trough_local"] = trough_conf
    out["gt_peak_local"] = peak_conf
    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"] = (pd.Series(trough_z, index=out.index) == 1) & bb_ok
    out["gt_buy_signal"] = out["gt_buy_signal"].astype(int)
    out["gt_sell_signal"] = pd.Series(peak_z, index=out.index).astype(int)
    out["gt_signal_causal"] = 1
    return out