kis_bot/kis_trader/backtest/bt_candle_source.py

#!/usr/bin/env python3
"""
백테·Optuna ws_candles 소스 선택 — 실매 get_candles 와 동일.

읽기: 틱 메인 WS 한 칸, 없으면 kiwoom+rest 구멍, 그다음 kiwoom+rollup.
CANDLE_SOURCE=kis|kiwoom 이면 그 메인(+구멍). 빈값이면 LIVE_TICK_PROVIDER.
"""
from __future__ import annotations

import logging
import time
from collections import defaultdict
from typing import Any, Dict, List, Optional, Sequence, Tuple

from kis_trader.ws.candle_series import (
    dedupe_by_read_pairs,
    live_read_label,
    live_read_pairs,
    normalize_source_channel,
)

ReadPair = Tuple[str, str]

logger = logging.getLogger("bt_candle_source")

# 레거시 호환 이름 (실제 필터는 live_read_pairs)
BT_WS_CANDLE_SOURCES: Tuple[str, ...] = ("kiwoom", "kis", "rest", "rollup_1m")


def resolve_bt_candle_source_override() -> str:
    """'' = LIVE 병합, 'kis'|'kiwoom' = 그 메인(+키움 REST 구멍)."""
    from kis_trader.ws.candle_series import _provider_and_override
    _provider, override = _provider_and_override()
    return override


def live_candle_source_order() -> Tuple[str, ...]:
    """호환: source 문자열만. 실제 조회는 resolve_bt_read_pairs."""
    return tuple(dict.fromkeys(p[0] for p in live_read_pairs()))


def resolve_bt_read_pairs() -> Tuple[ReadPair, ...]:
    return live_read_pairs()


def resolve_bt_candle_source_order() -> Tuple[str, ...]:
    """호환 래퍼 — 신규 코드는 resolve_bt_read_pairs 사용."""
    return live_candle_source_order()


def resolve_bt_candle_source_label() -> str:
    return live_read_label()


def dedupe_candle_rows(
    rows: Sequence[Dict[str, Any]],
    source_order: Optional[Sequence[Any]] = None,
) -> List[Dict[str, Any]]:
    """candle_time 1행. source_order 가 쌍이면 그대로, 아니면 live_read_pairs."""
    pairs: Optional[Sequence[ReadPair]] = None
    if source_order:
        first = source_order[0]
        if isinstance(first, (tuple, list)) and len(first) >= 2:
            pairs = tuple((str(a), str(b)) for a, b in source_order)  # type: ignore[misc]
        else:
            # 레거시 source 문자열 → 정규화 후 순위
            mapped: List[ReadPair] = []
            for s in source_order:
                mapped.append(normalize_source_channel(str(s), ""))
            pairs = tuple(mapped)
    return dedupe_by_read_pairs(rows, pairs)


def _pairs_in_sql(pairs: Sequence[ReadPair]) -> Tuple[str, List[str]]:
    parts: List[str] = []
    params: List[str] = []
    for src, ch in pairs:
        parts.append("(source=%s AND channel=%s)")
        params.extend([src, ch])
    return " AND (" + " OR ".join(parts) + ")", params


def _tick_time_bounds(start_key: str, end_key: str) -> Tuple[str, str]:
    sk = str(start_key or "").strip()
    ek = str(end_key or "").strip()
    if len(sk) == 12:
        sk = sk + "00"
    if len(ek) == 12:
        ek = ek + "59"
    return sk, ek


def _ws_ticks_table(market: Optional[str]) -> str:
    mk = (market or "").strip().upper()
    if mk == "US":
        return "ws_ticks_us"
    return "ws_ticks"


def _ingest_flat_tick_rows(
    rows: List[Dict[str, Any]],
    target: Dict[str, List[Dict[str, Any]]],
) -> int:
    """SELECT 행 → 종목별 flat 리스트 (봉 pick 쓰레기검사용).

    ``breakout_tick_loader._ingest_tick_rows`` 는 minute_key 버킷 딕트를 만들지만,
    bt_candle_source 소비자 (``dedupe_by_read_pairs``) 는 flat 리스트를 원한다.
    → 같은 컬럼 (``recv_ts`` + ``_lag_sec``) 을 유지하되 리턴 형식만 다르게.
    """
    from kis_trader.backtest.breakout_tick_loader import _tick_row_lag_seconds

    n = 0
    for r in rows or []:
        code = str(r.get("code") or "").strip()
        if not code:
            continue
        tt = str(r.get("tick_time") or "")[:14]
        if len(tt) < 12:
            continue
        kw_lag = _tick_row_lag_seconds(r)
        # recv_ts: bar_is_garbage 가 wall-clock 정합 판정에 사용 (docs/정합성.md §9)
        # 문자열로 보존. None 이면 봉끝 폴백.
        _recv_raw = r.get("recv_ts")
        _recv_str = ""
        if _recv_raw is not None:
            try:
                _recv_str = (
                    _recv_raw.strftime("%Y-%m-%d %H:%M:%S")
                    if hasattr(_recv_raw, "strftime")
                    else str(_recv_raw)
                )
            except Exception:
                _recv_str = str(_recv_raw)
        target[code].append({
            "code": code,
            "tick_time": tt,
            "source": r.get("source") or "",
            "tick_time_raw": str(r.get("tick_time_raw") or ""),
            "_lag_sec": kw_lag,
            "recv_ts": _recv_str,
        })
        n += 1
    return n


def _load_ticks_by_code_bulk(
    db,
    start_key: str,
    end_key: str,
    *,
    market: Optional[str] = None,
    codes_filter: Optional[Sequence[str]] = None,
) -> Optional[Dict[str, List[Dict[str, Any]]]]:
    """쓰레기 검사용 틱 — 로더 통일 (2026-09-06 C안).

    이전(자체 SELECT): ``tick_time, source, tick_time_raw`` 만 → recv_ts/_lag_sec 미포함
    → ``bar_is_garbage`` 가 봉끝 폴백 → **22.8% 대량 컷** (실매와 다른 기준).

    현재: ``breakout_tick_loader._fetch_ws_ticks_day_rows`` 위임 →
    엔진 로더와 **동일 SELECT + 동일 lag 계산 + 동일 recv_ts** →
    쓰레기 판정도 wall-clock 정합.

    - KR → ``ws_ticks``  /  US → ``ws_ticks_us``
    - OFF 면 None (봉만 dedupe)
    - 리턴 형식은 flat ``Dict[code, List[tick]]`` (bt_candle_source 소비자 호환)
    """
    from kis_trader.engine.feed_fallback import candle_garbage_fallback_enabled

    if not candle_garbage_fallback_enabled():
        return None

    from kis_trader.backtest.breakout_tick_loader import (
        _candle_keys_to_tick_range,
        _fetch_ws_ticks_day_rows,
        _iter_tick_day_chunks,
        _ws_ticks_table as _bl_ws_ticks_table,
    )

    mk = (market or "").strip().upper()
    tt_start, tt_end = _candle_keys_to_tick_range(start_key, end_key)

    want: Optional[set] = None
    if codes_filter:
        want = {str(c).strip() for c in codes_filter if str(c).strip()}

    out: Dict[str, List[Dict[str, Any]]] = defaultdict(list)

    def _pull(table: str, mkt_str: str) -> int:
        total = 0
        for chunk_s, chunk_e in _iter_tick_day_chunks(tt_start, tt_end):
            try:
                rows = _fetch_ws_ticks_day_rows(table, mkt_str, chunk_s, chunk_e, want)
            except Exception as e:
                logger.warning(
                    "쓰레기검사 틱 로드 실패 (table=%s day=%s): %s",
                    table, chunk_s[:8], e,
                )
                continue
            total += _ingest_flat_tick_rows(rows, out)
        return total

    try:
        n_kr = n_us = 0
        if mk == "US":
            n_us = _pull(_bl_ws_ticks_table("US"), "US")
        elif mk == "KR":
            n_kr = _pull(_bl_ws_ticks_table("KR"), "KR")
        else:
            n_kr = _pull(_bl_ws_ticks_table("KR"), "KR")
            try:
                n_us = _pull(_bl_ws_ticks_table("US"), "US")
            except Exception:
                n_us = 0
        logger.info("📥 틱 bulk 쓰레기검사용: KR=%s US=%s 종목=%s", n_kr, n_us, len(out))
        return dict(out)
    except Exception as exc:
        logger.warning("틱 bulk 스킵(봉만): %s", exc)
        return None


def _ticks_for_bar_garbage(
    db,
    code: str,
    start_key: str,
    end_key: str,
    *,
    market: Optional[str] = None,
) -> Optional[List[Dict[str, Any]]]:
    """단일 종목 쓰레기 검사용 틱 — 로더 통일 (bulk 재사용, 2026-09-06 C안).

    이전엔 별도 SELECT (``tick_time, source, tick_time_raw`` 만) → recv_ts/_lag_sec 없음.
    이제 ``_load_ticks_by_code_bulk`` (breakout_tick_loader 위임) 를 코드 1개로 호출 →
    동일 정합.
    """
    from kis_trader.engine.feed_fallback import candle_garbage_fallback_enabled

    if not candle_garbage_fallback_enabled():
        return None
    if not str(code or "").strip():
        return None
    result = _load_ticks_by_code_bulk(
        db, start_key, end_key, market=market, codes_filter=[str(code).strip()],
    )
    if result is None:
        return None
    return result.get(str(code).strip(), [])


def list_ws_candle_codes(
    db,
    timeframe: int,
    start_key: str,
    end_key: str,
    *,
    market: Optional[str] = None,
) -> List[str]:
    """기간 내 종목 코드 — 선택된 (source,channel) 기준 DISTINCT."""
    pairs = resolve_bt_read_pairs()
    src_sql, src_params = _pairs_in_sql(pairs)
    mk = (market or "").strip().upper()
    if mk:
        rows = db.conn.execute(
            "SELECT DISTINCT code FROM ws_candles WHERE timeframe=%s AND market=%s "
            "AND candle_time >= %s AND candle_time <= %s"
            + src_sql
            + " ORDER BY code",
            [int(timeframe), mk, start_key, end_key, *src_params],
        ).fetchall()
    else:
        rows = db.conn.execute(
            "SELECT DISTINCT code FROM ws_candles WHERE timeframe=%s "
            "AND candle_time >= %s AND candle_time <= %s"
            + src_sql
            + " ORDER BY code",
            [int(timeframe), start_key, end_key, *src_params],
        ).fetchall()
    return [r["code"] for r in rows]


def fetch_ws_candles_for_code(
    db,
    code: str,
    timeframe: int,
    start_key: str,
    end_key: str,
    *,
    extra_select: str = "",
    peak_sel: str = "",
    market: Optional[str] = None,
    confirmed_only: bool = True,
) -> List[Dict[str, Any]]:
    """단일 종목·기간 봉 로드 — 메인 WS 우선, 구멍 kiwoom+rest(+rollup)."""
    pairs = resolve_bt_read_pairs()
    confirmed_sql = " AND is_confirmed=1" if confirmed_only else ""
    mk = (market or "").strip().upper()
    src_sql, src_params = _pairs_in_sql(pairs)
    cols = (
        f"candle_time, open, high, low, close, volume, source, channel"
        f"{peak_sel}{extra_select}"
    )
    if mk:
        rows = db.conn.execute(
            f"SELECT {cols} FROM ws_candles "
            "WHERE timeframe=%s AND code=%s AND market=%s "
            "AND candle_time >= %s AND candle_time <= %s"
            + confirmed_sql
            + src_sql
            + " ORDER BY candle_time ASC",
            [int(timeframe), code, mk, start_key, end_key, *src_params],
        ).fetchall()
    else:
        rows = db.conn.execute(
            f"SELECT {cols} FROM ws_candles "
            "WHERE timeframe=%s AND code=%s "
            "AND candle_time >= %s AND candle_time <= %s"
            + confirmed_sql
            + src_sql
            + " ORDER BY candle_time ASC",
            [int(timeframe), code, start_key, end_key, *src_params],
        ).fetchall()
    ticks = _ticks_for_bar_garbage(
        db, str(code), start_key, end_key, market=mk or None,
    )
    return dedupe_by_read_pairs(
        [dict(r) for r in rows], pairs, ticks=ticks, tf_min=int(timeframe),
        missing_policy="hole", live_cover=False,
    )


def fetch_ws_candles_by_code_bulk(
    db,
    timeframe: int,
    start_key: str,
    end_key: str,
    *,
    extra_select: str = "",
    peak_sel: str = "",
    market: Optional[str] = None,
    confirmed_only: bool = True,
    codes_filter: Optional[Sequence[str]] = None,
) -> Dict[str, List[Dict[str, Any]]]:
    """기간 전체 봉 1쿼리 + 틱 1~2쿼리 후 종목별 기존 쓰레기 검사.

    웜업용 ``fetch_ws_candles_for_code`` / ``fetch_ws_candles_warmup_before`` 는 그대로.
    """
    t0 = time.perf_counter()
    pairs = resolve_bt_read_pairs()
    confirmed_sql = " AND is_confirmed=1" if confirmed_only else ""
    mk = (market or "").strip().upper()
    src_sql, src_params = _pairs_in_sql(pairs)
    cols = (
        f"code, candle_time, open, high, low, close, volume, source, channel"
        f"{peak_sel}{extra_select}"
    )
    want: List[str] = []
    if codes_filter:
        want = [
            str(c).strip()
            for c in codes_filter
            if str(c).strip()
        ]
    in_sql = ""
    in_params: List[str] = []
    if want:
        in_sql = " AND code IN (" + ",".join(["%s"] * len(want)) + ")"
        in_params = want

    params: List[Any] = [int(timeframe)]
    mk_sql = ""
    if mk:
        mk_sql = " AND market=%s"
        params.append(mk)
    params.extend([start_key, end_key, *src_params, *in_params])

    rows = db.conn.execute(
        f"SELECT {cols} FROM ws_candles "
        "WHERE timeframe=%s"
        + mk_sql
        + " AND candle_time >= %s AND candle_time <= %s"
        + confirmed_sql
        + src_sql
        + in_sql
        + " ORDER BY code ASC, candle_time ASC",
        params,
    ).fetchall()

    by_code: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
    for r in rows or []:
        code = str(r.get("code") or "").strip()
        if not code:
            continue
        by_code[code].append(dict(r))

    ticks_map = _load_ticks_by_code_bulk(
        db, start_key, end_key, market=mk or None, codes_filter=want or None,
    )
    out: Dict[str, List[Dict[str, Any]]] = {}
    stats: Dict[str, Any] = {
        "slots": 0, "picked": 0, "hole": 0, "garbage_skip": 0,
        "pick_by": {}, "garbage_by": {}, "raw_by": {},
    }
    for code, raw in by_code.items():
        ticks = None if ticks_map is None else ticks_map.get(code, [])
        out[code] = dedupe_by_read_pairs(
            raw, pairs, ticks=ticks, tf_min=int(timeframe),
            missing_policy="hole", live_cover=False, stats=stats,
        )
    elapsed = time.perf_counter() - t0
    n_tick_rows = 0
    if ticks_map:
        n_tick_rows = sum(len(v) for v in ticks_map.values())
    slots = int(stats.get("slots") or 0) or 1
    pick_by = stats.get("pick_by") or {}
    garb_by = stats.get("garbage_by") or {}
    raw_by = stats.get("raw_by") or {}
    pick_txt = " ".join(
        "%s=%s(%.0f%%)" % (k, v, 100.0 * int(v) / max(1, int(stats.get("picked") or 1)))
        for k, v in sorted(pick_by.items(), key=lambda x: -int(x[1]))
    ) or "—"
    garb_txt = " ".join(
        "%s=%s" % (k, v) for k, v in sorted(garb_by.items(), key=lambda x: -int(x[1]))
    ) or "0"
    raw_txt = " ".join(
        "%s=%s" % (k, v) for k, v in sorted(raw_by.items(), key=lambda x: -int(x[1]))
    ) or "—"
    logger.info(
        "📂 캔들 bulk: 봉 %s행 · 틱 %s행 · 종목 %s · %.2fs (tf=%s market=%s)",
        len(rows or []),
        n_tick_rows,
        len(out),
        elapsed,
        int(timeframe),
        mk or "ALL",
    )
    logger.info(
        "📊 봉 pick 비율(1차→2차→REST): 슬롯 %s · pick %s · hole %s(%.1f%%) · "
        "쓰레기스킵 %s(슬롯대비 %.1f%%) | raw[%s] | pick[%s] | garbage[%s]",
        slots,
        stats.get("picked"),
        stats.get("hole"),
        100.0 * int(stats.get("hole") or 0) / slots,
        stats.get("garbage_skip"),
        100.0 * int(stats.get("garbage_skip") or 0) / slots,
        raw_txt,
        pick_txt,
        garb_txt,
    )
    return out


def fetch_ws_candles_warmup_before(
    db,
    code: str,
    timeframe: int,
    before_candle_time: str,
    limit: int,
    *,
    extra_select: str = "",
    peak_sel: str = "",
    confirmed_only: bool = True,
) -> List[Dict[str, Any]]:
    """기간 시작 이전 N봉 — prepend 웜업용 (오래된→최신)."""
    if limit <= 0 or not before_candle_time:
        return []
    pairs = resolve_bt_read_pairs()
    confirmed_sql = " AND is_confirmed=1" if confirmed_only else ""
    fetch_limit = max(int(limit) * max(len(pairs), 1), int(limit) + 50)
    src_sql, src_params = _pairs_in_sql(pairs)
    cols = (
        f"candle_time, open, high, low, close, volume, source, channel"
        f"{peak_sel}{extra_select}"
    )
    rows = db.conn.execute(
        f"SELECT {cols} FROM ws_candles "
        "WHERE timeframe=%s AND code=%s AND candle_time < %s"
        + confirmed_sql
        + src_sql
        + " ORDER BY candle_time DESC LIMIT %s",
        [int(timeframe), code, before_candle_time, *src_params, fetch_limit],
    ).fetchall()
    raw_rows = [dict(r) for r in reversed(rows)]
    t0 = str(raw_rows[0].get("candle_time") or "")[:12] if raw_rows else ""
    ticks = (
        _ticks_for_bar_garbage(db, str(code), t0, str(before_candle_time))
        if t0 else None
    )
    bars = dedupe_by_read_pairs(
        raw_rows, pairs,
        ticks=ticks,
        tf_min=int(timeframe),
        missing_policy="hole",
        live_cover=False,
    )
    return bars[-limit:] if len(bars) > limit else bars