#!/usr/bin/env python3
"""
optuna_momentum_tpe_space.py — 모멘텀 Optuna **연속(TPE) 탐색 공간**
==============================================================================
기존 ``_momentum_grids()`` / ``suggest_categorical`` 경로는 **건드리지 않음**.

사용:
    python3 kis_trader/backtest/param_search_optuna.py \\
        --strategy momentum --mode tpe --trials 200 ...

TPE 가 trial 을 쌓으며 좋은 구간을 좁힘 (메뉴판 이산 스윕이 아님).
``skip_hts_scan_dupes`` 는 KR=False / US=True 고정.
금액(슬롯·손절액·min_price) 축은 TPE 에 없음 — 특히 US 고정 유니버스는 DB 슬롯만.

축 정책 (2026-08-16):
  - RSI(mom_rsi_min/max) 제거: 국장 스위치 기본 OFF → 엔진이 값을 안 보므로 낭비 축.
  - EMA 추세 필터 추가: 전략 정체성(추세 확인). use_ema_filter(ON/OFF) + fast/slow 기간.
"""
from __future__ import annotations

from typing import Any, Dict, List

import optuna

from kis_trader.backtest.optuna_tpe_common import (
    ORDERBOOK_TPE_AXIS_KEYS,
    RATCHET_TPE_AXIS_KEYS,
    WHIPSAW_TPE_AXIS_KEYS,
    r1,
    r2,
    suggest_orderbook_entry_tpe,
    suggest_ratchet_tiers_pct,
    suggest_whipsaw_tpe,
)
from kis_trader.backtest.param_search_momentum import _momentum_combo_grid_valid

# 연속 축 키 목록 (JSON/힌트용)
# 매수 시간창(time_*) 은 TPE 제외 — 운영 DB 고정(다단래칫·청산축만 탐색).
# 래칫: 숫자축(on/n/gain/cut) + 조립 문자열 ratchet_tiers
_MOMENTUM_BASE_KEYS: List[str] = [
    # EMA 추세 필터 — 전략 정체성 (ON/OFF + 기간)
    "use_ema_filter",
    "ema_fast_period",
    "ema_slow_period",
    # 거래량·진입 조건
    "mom_vol_mult",
    "mom_vol_win",
    "e_min_chg_pct",
    # 손익 청산
    "tp_pct",
    "tp_max_pct",
    "sl_pct",
    # 눌림·숄더 패턴
    "shoulder_min_high",
    "shoulder_cut_pct",
    "trail_pct",
    "trail_arm_pct",
    # 운영 제어
    "max_hold_bars",
    "cooldown_min",
    "max_daily",
    "max_daily_chg",
    "mom_max_from_open_pct",
    "min_margin",
    # 추격·눌림 타이밍
    "chase_lookback_min",
    "pullback_lookback_min",
    "pullback_min_pct",
    "pullback_max_pct",
    # 셋업 거래량
    "setup_vol_max_mult",
    "setup_bear_bars_min",
]
MOMENTUM_TPE_AXIS_KEYS: List[str] = (
    list(_MOMENTUM_BASE_KEYS)
    + list(RATCHET_TPE_AXIS_KEYS)
    + list(ORDERBOOK_TPE_AXIS_KEYS)
    + list(WHIPSAW_TPE_AXIS_KEYS)
)


def momentum_tpe_axis_keys() -> List[str]:
    return list(MOMENTUM_TPE_AXIS_KEYS)


def suggest_momentum_params_tpe(
    trial: optuna.Trial,
    *,
    market: str = "KR",
) -> Dict[str, Any]:
    """
    연속 구간 suggest — TPE 가 좁혀 감.
    UI% 단위 (evaluate / apply 경로와 동일).

    market=US: 금액(슬롯·손절액·min_price) 축 없음 — 고정 유니버스·DB 슬롯만.

    EMA 추세 필터:
      - use_ema_filter: categorical(True/False) — TPE 가 ON/OFF 를 함께 탐색
      - EMA ON 일 때만 fast/slow 기간이 의미 있음.
        OFF 시 _momentum_combo_grid_valid 가 fast=9/slow=21 만 허용 → TrialPruned.
    RSI(mom_rsi_min/max) 제거 이유:
      국장 MOMENTUM_USE_RSI_FILTER 기본 OFF. 엔진이 값을 안 읽어 빈 축이 됨.
    """
    mk = (market or "KR").strip().upper() or "KR"
    combo: Dict[str, Any] = {}

    # ── EMA 추세 필터 (전략 정체성) ────────────────────────────────────────────
    combo["use_ema_filter"] = trial.suggest_categorical("use_ema_filter", [True, False])
    combo["ema_fast_period"] = trial.suggest_int("ema_fast_period", 5, 20)
    combo["ema_slow_period"] = trial.suggest_int("ema_slow_period", 15, 60)
    if combo["ema_slow_period"] <= combo["ema_fast_period"]:
        raise optuna.TrialPruned("ema_slow <= ema_fast")
    # EMA OFF 시 기간 축은 의미 없음 — _momentum_combo_grid_valid 에서 추가 검증


    combo["mom_vol_mult"] = r2(trial.suggest_float("mom_vol_mult", 1.0, 10.0, step=0.1))
    combo["mom_vol_win"] = trial.suggest_int("mom_vol_win", 3, 10)
    # HTS K: 전일종가 대비 최소등락(%) — 실매 앵커 0.2 포함 연속 텀 (step 0.1)
    combo["e_min_chg_pct"] = r1(
        trial.suggest_float("e_min_chg_pct", 0.0, 3.0, step=0.1),  # 확장: 2.0→3.0
    )

    combo["tp_pct"] = r2(trial.suggest_float("tp_pct", 1.5, 15.0, step=0.1))
    combo["tp_max_pct"] = r2(trial.suggest_float("tp_max_pct", 2.0, 20.0, step=0.5))
    if combo["tp_max_pct"] + 1e-9 < combo["tp_pct"]:
        raise optuna.TrialPruned("tp_max < tp")

    combo["sl_pct"] = r2(trial.suggest_float("sl_pct", 0.5, 8.0, step=0.1))  # 확장: 1.0→0.5, 6.0→8.0

    combo["shoulder_min_high"] = r2(
        trial.suggest_float("shoulder_min_high", 0.3, 5.0, step=0.1),
    )
    combo["shoulder_cut_pct"] = r2(
        trial.suggest_float("shoulder_cut_pct", 0.05, 0.8, step=0.05),
    )
    combo["trail_pct"] = r2(trial.suggest_float("trail_pct", 0.0, 4.0, step=0.1))
    combo["trail_arm_pct"] = r2(
        trial.suggest_float("trail_arm_pct", 0.0, 5.0, step=0.1),
    )

    combo["max_hold_bars"] = trial.suggest_int("max_hold_bars", 0, 300, step=5)  # 확장: 180→300
    combo["cooldown_min"] = r1(trial.suggest_float("cooldown_min", 0.0, 30.0, step=1.0))
    combo["max_daily"] = trial.suggest_int("max_daily", 5, 100, step=5)
    combo["max_daily_chg"] = r1(
        trial.suggest_float("max_daily_chg", 15.0, 60.0, step=1.0),
    )
    combo["mom_max_from_open_pct"] = r1(
        trial.suggest_float("mom_max_from_open_pct", 15.0, 60.0, step=1.0),
    )
    combo["min_margin"] = r2(trial.suggest_float("min_margin", 0.1, 2.0, step=0.1))

    combo["chase_lookback_min"] = trial.suggest_int("chase_lookback_min", 5, 20)
    combo["pullback_lookback_min"] = trial.suggest_int("pullback_lookback_min", 8, 30)
    combo["pullback_min_pct"] = r2(
        trial.suggest_float("pullback_min_pct", 0.1, 1.0, step=0.1),
    )
    combo["pullback_max_pct"] = r2(
        trial.suggest_float("pullback_max_pct", 1.0, 10.0, step=0.1),  # 확장: 6.0→10.0
    )
    if combo["pullback_min_pct"] >= combo["pullback_max_pct"]:
        raise optuna.TrialPruned("pullback min >= max")

    combo["setup_vol_max_mult"] = r2(
        trial.suggest_float("setup_vol_max_mult", 0.5, 3.0, step=0.1),
    )
    combo["setup_bear_bars_min"] = trial.suggest_int("setup_bear_bars_min", 0, 3)

    # 래칫: 숫자축 → \"5:2,10:1\" 조립 (빈칸=OFF)
    combo.update(
        suggest_ratchet_tiers_pct(
            trial,
            off_token="",
            n_max=3,
            gain_low=2.0,
            gain_high=15.0,
            gain_step=0.5,
            cut_low=0.5,
            cut_high=3.0,
            cut_step=0.1,
        ),
    )

    # 운영 고정 — 그리드 스윕 금지 (금액·슬롯·min_price 축 없음)
    if mk == "US":
        combo["skip_hts_scan_dupes"] = True  # HTS 없음
    else:
        combo["skip_hts_scan_dupes"] = False

    combo.update(suggest_orderbook_entry_tpe(trial))
    combo.update(suggest_whipsaw_tpe(trial))

    if not _momentum_combo_grid_valid(combo):
        raise optuna.TrialPruned("momentum invalid combo")
    return combo