kis_bot/kis_trader/backtest/universe_timeline.py

#!/usr/bin/env python3
"""
유니버스 타임라인 (전략 공통) — 스캔 시각(초)별 ``get_universe_at`` 재현 + EXIT 디바운스.

실매: 조건검색 변동 시 초단위 스냅샷(``target_candidates_history``) → 루프에서 그 시각
최신 유니버스. 백테(구): 1분 슬롯 strict 집계 → 깜빡임 EXIT·분봉 끝 집계로 실매와 어긋남.

이 모듈은 전략 무관(strategy_id 파라미터)으로, 모멘텀·돌파 등 모든 전략이 **하나의
``UniverseTimeline``** 로 실매(``get_universe_at``)와 초단위 정합을 맞추도록 통일한다.
(기존 ``MomentumUniverseTimeline`` 은 이 클래스의 alias 로 유지 — 모멘텀 코드 무변경)
"""
from __future__ import annotations

from bisect import bisect_right
from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional, Sequence, Set, Tuple


def resolve_universe_exit_debounce_sec(
    *,
    strategy_env_key: Optional[str] = None,
    default_when_no_grace: int = 0,
) -> int:
    """실매 ``CONDITION_EXIT_GRACE_SEC`` 와 동일 초로 스냅샷 EXIT 디바운스.

    history 가 장초·overnight 전환에서 raw 축소로 찍혀도, 실매 EXIT grace 처럼
    빠진 종목을 N초간 유지한다. 전략별 키가 있으면 그 값 우선(0=디바운스 OFF).
    """
    from kis_trader.utils.env import get_env_from_db, get_env_int

    if strategy_env_key:
        raw = get_env_from_db(strategy_env_key, None)
        if raw not in (None, ""):
            try:
                return max(0, int(float(raw)))
            except (TypeError, ValueError):
                pass
    grace = int(get_env_int("CONDITION_EXIT_GRACE_SEC", 0))
    if grace > 0:
        return grace
    return max(0, int(default_when_no_grace))


# strategy_id → (전략별 오버라이드 env 키, CONDITION_EXIT_GRACE=0 일 때 폴백)
# 웹·Optuna·resolve_*_universe 가 이 맵만 쓰면 debounce=0 특수케이스가 사라진다.
_STRATEGY_UNIVERSE_EXIT_DEBOUNCE: Dict[str, Tuple[Optional[str], int]] = {
    "SCALP": ("SCALP_UNIVERSE_EXIT_DEBOUNCE_SEC", 0),
    "MOMENTUM": ("MOMENTUM_UNIVERSE_EXIT_DEBOUNCE_SEC", 30),  # 모멘텀 기존 폴백 30 유지
    "BREAKOUT": ("BREAKOUT_UNIVERSE_EXIT_DEBOUNCE_SEC", 0),
    "SHORT": ("TAIL_UNIVERSE_EXIT_DEBOUNCE_SEC", 0),
    "TAIL": ("TAIL_UNIVERSE_EXIT_DEBOUNCE_SEC", 0),
    "RANGE_BREAK": ("RANGE_BREAK_UNIVERSE_EXIT_DEBOUNCE_SEC", 0),
}


def universe_exit_debounce_sec_for_strategy(strategy_id: str) -> int:
    """전략 공통 — history 슬롯·타임라인 EXIT 디바운스 초.

    우선순위: ``{STRATEGY}_UNIVERSE_EXIT_DEBOUNCE_SEC`` → ``CONDITION_EXIT_GRACE_SEC``
    → 전략별 폴백(모멘텀만 30, 나머지 0).
    """
    sid = str(strategy_id or "").strip().upper()
    key, dflt = _STRATEGY_UNIVERSE_EXIT_DEBOUNCE.get(sid, (None, 0))
    return resolve_universe_exit_debounce_sec(
        strategy_env_key=key,
        default_when_no_grace=dflt,
    )


def _event_time_to_key(event_time: str) -> str:
    et = str(event_time or "")
    if len(et) < 19:
        return ""
    return (
        et[0:4] + et[5:7] + et[8:10]
        + et[11:13] + et[14:16] + et[17:19]
    )


def _parse_scan_key(scan_key: str) -> datetime:
    return datetime.strptime(str(scan_key)[:14], "%Y%m%d%H%M%S")


def _add_minutes_to_scan_key(scan_key: str, minutes: int) -> str:
    dt = _parse_scan_key(scan_key)
    return (dt + timedelta(minutes=int(minutes))).strftime("%Y%m%d%H%M%S")


def _dedupe_codes_preserve_order(codes: Sequence[str]) -> List[str]:
    """스냅샷 내 종목 순서 유지 + 중복 제거 (HTS/DB insert 순)."""
    seen: Set[str] = set()
    out: List[str] = []
    for c in codes:
        cs = str(c).strip()
        if cs and cs not in seen:
            seen.add(cs)
            out.append(cs)
    return out


def normalize_universe_events(events: Sequence[Dict[str, Any]]) -> List[Tuple[str, List[str]]]:
    """iter_universe_events 결과 → (YYYYMMDDHHMMSS, codes) 시간순 리스트."""
    out: List[Tuple[str, List[str]]] = []
    for ev in events:
        et_key = _event_time_to_key(str(ev.get("event_time") or ""))
        if not et_key:
            continue
        codes = _dedupe_codes_preserve_order(
            [str(it["code"]) for it in ev.get("items", []) if it.get("code")]
        )
        out.append((et_key, codes))
    out.sort(key=lambda x: x[0])
    return out


def debounce_universe_snapshots(
    normalized: Sequence[Tuple[str, List[str]]],
    debounce_sec: int,
) -> List[Tuple[str, List[str]]]:
    """
    짧은 EXIT 후 재편입은 EXIT 로 치지 않음 (HTS 조건검색 노이즈 완화).
    각 스냅샷 시각에 **디바운스 적용 후** 유효 종목 리스트를 기록 (가나다 재정렬 없음).
    """
    if debounce_sec <= 0 or not normalized:
        return list(normalized)

    effective_order: List[str] = []
    effective_set: Set[str] = set()
    pending_exit: Dict[str, Tuple[str, datetime]] = {}
    out: List[Tuple[str, List[str]]] = []

    def _eff_remove(code: str) -> None:
        if code not in effective_set:
            return
        effective_set.discard(code)
        try:
            effective_order.remove(code)
        except ValueError:
            pass

    def _eff_add(code: str) -> None:
        if code in effective_set:
            return
        effective_set.add(code)
        effective_order.append(code)

    for et_key, codes in normalized:
        cur_list = _dedupe_codes_preserve_order([str(c) for c in codes if c])
        cur_set = set(cur_list)
        et_dt = _parse_scan_key(et_key)

        for code, (ex_key, ex_dt) in list(pending_exit.items()):
            if (et_dt - ex_dt).total_seconds() >= debounce_sec:
                if code not in cur_set:
                    _eff_remove(code)
                pending_exit.pop(code, None)

        for code in list(effective_order):
            if code not in cur_set and code not in pending_exit:
                pending_exit[code] = (et_key, et_dt)

        for code in cur_list:
            pending_exit.pop(code, None)
            _eff_add(code)

        out.append((et_key, list(effective_order)))

    return out


class UniverseTimeline:
    """스캔 시각별 유니버스 — 이진 탐색으로 O(log N) 조회 (전략 공통)."""

    __slots__ = ("_keys", "_codes", "_strict_avail")

    def __init__(
        self,
        snapshots: Sequence[Tuple[str, List[str]]],
        *,
        strict: bool = False,
        strict_lag_minutes: int = 0,
    ) -> None:
        keys: List[str] = []
        codes_list: List[List[str]] = []
        strict_avail: Dict[str, str] = {}
        prev_set: Set[str] = set()

        for et_key, codes in snapshots:
            keys.append(et_key)
            ordered = _dedupe_codes_preserve_order([str(c) for c in codes if c])
            cur_set = set(ordered)
            if strict:
                lag = max(0, int(strict_lag_minutes))
                for c in cur_set - prev_set:
                    strict_avail[c] = (
                        _add_minutes_to_scan_key(et_key, lag) if lag > 0 else et_key
                    )
                for c in prev_set - cur_set:
                    strict_avail.pop(c, None)
                prev_set = cur_set
            codes_list.append(ordered)

        self._keys = keys
        self._codes = codes_list
        self._strict_avail = strict_avail

    def codes_at(self, scan_key: str) -> List[str]:
        """scan_key(YYYYMMDDHHMMSS) 시점 유니버스 — 실매 get_universe_at 동일."""
        sk = str(scan_key or "")[:14]
        if not sk or not self._keys:
            return []
        idx = bisect_right(self._keys, sk) - 1
        if idx < 0:
            return []
        raw = list(self._codes[idx])
        if not self._strict_avail:
            return raw
        return [c for c in raw if sk >= self._strict_avail.get(c, "00000000000000")]

    @property
    def snapshot_count(self) -> int:
        return len(self._keys)


def build_universe_timeline(
    *,
    strategy_id: str,
    start_ymd: str,
    end_ymd: str,
    debounce_sec: int = 30,
    strict: bool = False,
    strict_lag_minutes: int = 1,
    history_source: str = "kiwoom",
) -> Optional[UniverseTimeline]:
    """이력 테이블(키움/LS) → 초단위 유니버스 타임라인.

    ``history_source``: ``kiwoom`` | ``ls`` — resolve/슬롯 dict 와 **동일 소스**여야
    scan_at 매수가 라벨과 어긋나지 않는다.
    """
    try:
        from kis_trader.database.db_manager import get_db
        from kis_trader.backtest.universe_history_source import (
            apply_ls_session_filter_to_start,
            normalize_universe_history_source,
            resolve_backtest_universe_history_source,
        )

        hs = resolve_backtest_universe_history_source(history_source)
        hs = normalize_universe_history_source(hs)
        ext = get_db()
        start_time = f"{start_ymd[:4]}-{start_ymd[4:6]}-{start_ymd[6:8]} 00:00:00"
        start_time = apply_ls_session_filter_to_start(start_time, source=hs)
        end_time = f"{end_ymd[:4]}-{end_ymd[4:6]}-{end_ymd[6:8]} 23:59:59"
        events = ext.iter_universe_events(
            strategy_id=strategy_id,
            start_time=start_time,
            end_time=end_time,
            preserve_insert_order=True,
            history_source=hs,
        )
        normalized = normalize_universe_events(events)
        if not normalized:
            return None
        debounced = debounce_universe_snapshots(normalized, debounce_sec)
        return UniverseTimeline(
            debounced,
            strict=strict,
            strict_lag_minutes=strict_lag_minutes,
        )
    except Exception:
        return None