#!/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 슬롯만.
"""
from __future__ import annotations

from typing import Any, Dict, List

import optuna

from kis_trader.backtest.optuna_tpe_common import (
    RATCHET_TPE_AXIS_KEYS,
    r1,
    r2,
    suggest_ratchet_tiers_pct,
)
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] = [
    "mom_rsi_min",
    "mom_rsi_max",
    "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)


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 슬롯만.
    """
    mk = (market or "KR").strip().upper() or "KR"
    combo: Dict[str, Any] = {}

    combo["mom_rsi_min"] = trial.suggest_int("mom_rsi_min", 45, 65)
    combo["mom_rsi_max"] = trial.suggest_int("mom_rsi_max", 70, 100)
    if combo["mom_rsi_min"] >= combo["mom_rsi_max"]:
        raise optuna.TrialPruned("rsi_min >= rsi_max")

    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, 2.0, step=0.1),
    )

    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", 1.0, 6.0, step=0.1))

    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, 180, step=5)
    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, 6.0, step=0.1),
    )
    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.3, 1.5, 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

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