#!/usr/bin/env python3
"""
momentum_rr_crossval.py — 모멘텀 손익비(어깨·손절·익절) 스윕 → 파라서치 → 교차검증

1) 고정 진입(RSI·거래량=현 운영값) + 청산 축만 여러 프리셋 스윕 (학습 구간)
2) 1위 R:R 프리셋을 base 로 param_search (mode=rr) 실행
3) 탐색 상위 N개를 검증 구간(기본 2026-06-01)에서 재평가

실행:
  cd /home/hoon/kis_bot
  python3 kis_trader/backtest/momentum_rr_crossval.py
  python3 kis_trader/backtest/momentum_rr_crossval.py --train-start 2026-05-11 --train-end 2026-05-30 \\
      --oos-start 2026-06-01 --oos-end 2026-06-01 --search-top 5 --skip-search
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from datetime import datetime
from itertools import product
from typing import Any, Dict, List, Optional, Tuple

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(os.path.dirname(HERE))
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

import logging

logging.getLogger("TradeDB").setLevel(logging.WARNING)

from database import TradeDB
from kis_trader.backtest import scalping_backtest_common as sbc
from kis_trader.backtest.param_search_momentum import (
    _evaluate_momentum_chunk,
    _load_candles_for_search,
    _mom_fixed_defaults,
    _momentum_grids,
    _ui_to_engine_params,
    run_search,
)
from kis_trader.backtest.param_search_cli_common import apply_session_to_fixed
from kis_trader.engine import momentum_engine as me
from kis_trader.utils.env import get_env_float, get_env_from_db, get_env_int


def _results_dir() -> str:
    d = os.path.join(HERE, "results")
    os.makedirs(d, exist_ok=True)
    return d


def _load_universe(start: str, end: str) -> Tuple[Optional[Dict[str, List[str]]], str]:
    start_ymd = start.replace("-", "")
    end_ymd = end.replace("-", "")
    universe, source, _, _ = sbc.resolve_scalp_universe(
        start_ymd, end_ymd, use_saved_history=True, strategy_id="MOMENTUM",
    )
    if universe:
        return universe, source
    start_key = start_ymd + "0000"
    end_key = end_ymd + "2359"
    db = TradeDB()
    try:
        from kis_trader.backtest import momentum_backtest_common as mbc
        candles, _ = mbc.load_momentum_candles_by_code(db, start_key, end_key)
    finally:
        db.close()
    sim = me.build_universe_simulation_momentum(candles)
    return sim, "sim"


def _base_entry_ui(fixed: Dict[str, Any]) -> Dict[str, Any]:
    """현 운영 진입 축 고정 — 청산(R:R)만 스윕."""
    return {
        "mom_rsi_min": int(get_env_int("MOMENTUM_RSI_MIN", 55)),
        "mom_rsi_max": int(get_env_int("MOMENTUM_RSI_MAX", 90)),
        "mom_vol_mult": float(get_env_float("MOMENTUM_VOL_MULT", 1.5)),
        "mom_vol_win": float(fixed.get("mom_vol_win", 5)),
        "mom_time_end_hm": int(fixed.get("mom_time_end_hm", 1430)),
        "mom_max_from_open_pct": float(get_env_float("MOMENTUM_MAX_FROM_OPEN_PCT", 30.0)),
        "cooldown_min": float(fixed.get("cooldown_min", 1)),
        "max_daily": float(fixed.get("max_daily", 10)),
        "fee_rate": float(fixed.get("fee_rate", 0.015)),
        "sell_tax": float(fixed.get("sell_tax", 0.18)),
        "time_start_hm": int(fixed.get("time_start_hm", 830)),
        "time_end_hm": int(fixed.get("time_end_hm", 1530)),
        "min_price": int(fixed.get("min_price", 1000)),
        "max_loss_krw": float(fixed.get("max_loss_krw", 200_000)),
        "min_margin": float(fixed.get("min_margin", 0.2)),
        "use_defense_filters": bool(fixed.get("use_defense_filters", True)),
        "slot_money": float(get_env_float("MOMENTUM_SLOT_MONEY", 200_000)),
    }


def _rr_presets() -> List[Dict[str, Any]]:
    """
    청산 프리셋 (UI % 단위).
    shoulder_min_high >= 50 → 사실상 어깨 OFF (일반 장중 +50% 미도달).
    """
    raw = str(get_env_from_db("MOMENTUM_RR_PRESET_GRID_JSON", "") or "").strip()
    if raw:
        try:
            data = json.loads(raw)
            if isinstance(data, list) and data:
                return data
        except json.JSONDecodeError:
            pass
    return [
        {
            "name": "baseline_live",
            "sl_pct": 2.0, "tp_pct": 2.0, "tp_max_pct": 1.8,
            "shoulder_min_high": 0.5, "shoulder_cut_pct": 0.02,
        },
        {
            "name": "shoulder_off_tp25",
            "sl_pct": 2.0, "tp_pct": 2.5, "tp_max_pct": 2.5,
            "shoulder_min_high": 50.0, "shoulder_cut_pct": 0.3,
        },
        {
            "name": "wide_shoulder_sl15",
            "sl_pct": 1.5, "tp_pct": 2.5, "tp_max_pct": 2.5,
            "shoulder_min_high": 1.0, "shoulder_cut_pct": 0.25,
        },
        {
            "name": "delayed_shoulder_sl12",
            "sl_pct": 1.2, "tp_pct": 3.0, "tp_max_pct": 3.0,
            "shoulder_min_high": 1.2, "shoulder_cut_pct": 0.30,
        },
        {
            "name": "tp_only_sl15",
            "sl_pct": 1.5, "tp_pct": 2.0, "tp_max_pct": 2.0,
            "shoulder_min_high": 50.0, "shoulder_cut_pct": 0.2,
        },
        {
            "name": "loose_shoulder_sl18",
            "sl_pct": 1.8, "tp_pct": 2.2, "tp_max_pct": 2.2,
            "shoulder_min_high": 0.8, "shoulder_cut_pct": 0.20,
        },
    ]


def _run_backtest_period(
    candles: Dict[str, List[Dict]],
    universe: Optional[Dict[str, List[str]]],
    ui_params: Dict[str, Any],
    *,
    slot_money: float,
    max_stocks: int,
    total_budget: float,
    fee_rate: float,
    sell_tax: float,
) -> Dict[str, Any]:
    engine = _ui_to_engine_params(ui_params)
    engine["slot_money"] = float(slot_money)
    meta: Dict[str, Any] = {}
    trades = sbc.run_scalping_backtest_web_aligned(
        candles,
        engine,
        universe,
        slot_money=slot_money,
        fee_rate=fee_rate,
        sell_tax=sell_tax,
        max_stocks=max_stocks,
        total_budget_krw=total_budget,
        meta_out=meta,
        mode="momentum",
    )
    stats = sbc.summarize_trades(trades, total_budget_krw=total_budget)
    by_reason: Dict[str, Dict[str, float]] = {}
    for t in trades:
        r = str(t.get("sell_reason") or t.get("reason") or "?").split("(")[0].strip()
        pnl = float(t.get("pnl") or t.get("realized_pnl") or 0)
        bucket = by_reason.setdefault(r, {"n": 0, "pnl": 0.0})
        bucket["n"] += 1
        bucket["pnl"] += pnl
    wins = [t for t in trades if float(t.get("pnl") or t.get("realized_pnl") or 0) > 0]
    losses = [t for t in trades if float(t.get("pnl") or t.get("realized_pnl") or 0) <= 0]
    avg_win = (
        sum(float(t.get("pnl") or t.get("realized_pnl") or 0) for t in wins) / len(wins)
        if wins else 0.0
    )
    avg_loss = (
        sum(float(t.get("pnl") or t.get("realized_pnl") or 0) for t in losses) / len(losses)
        if losses else 0.0
    )
    return {
        "trades": len(trades),
        "total_pnl": round(stats.get("total_pnl", 0)),
        "win_rate": round(stats.get("win_rate", 0), 2),
        "pf": round(float(stats.get("pf") or 0), 3),
        "bot_pct": round(stats.get("bot_pct", 0), 3),
        "avg_win": round(avg_win),
        "avg_loss": round(avg_loss),
        "by_reason": by_reason,
        "params_ui": ui_params,
    }


def phase_rr_sweep(
    train_start: str,
    train_end: str,
    *,
    slot_money: float,
    max_stocks: int,
    total_budget: float,
) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
    fixed = _mom_fixed_defaults()
    apply_session_to_fixed(fixed, time_start_hm=830, time_end_hm=1530)
    entry_base = _base_entry_ui(fixed)

    print("\n" + "=" * 70)
    print(f"[Phase 1] R:R 프리셋 스윕 | 학습 {train_start} ~ {train_end}")
    print("=" * 70)

    candles = _load_candles_for_search(train_start, train_end, int(fixed["rsi_period"]))
    universe, src = _load_universe(train_start, train_end)
    print(f"  캔들 종목 {len(candles)} | 유니버스={src} 슬롯={len(universe or {})}")

    db = TradeDB()
    try:
        row = db.conn.execute("SELECT * FROM env_config ORDER BY id DESC LIMIT 1").fetchone()
        env_row = dict(row) if row else {}
    finally:
        db.close()
    fee_rate, sell_tax, _ = sbc.fee_and_slot_from_env(env_row, strategy="MOMENTUM")

    results: List[Dict[str, Any]] = []
    for preset in _rr_presets():
        ui = dict(entry_base)
        ui.update({k: v for k, v in preset.items() if k != "name"})
        ui["preset_name"] = preset.get("name", "unnamed")
        row = _run_backtest_period(
            candles, universe, ui,
            slot_money=slot_money, max_stocks=max_stocks, total_budget=total_budget,
            fee_rate=fee_rate, sell_tax=sell_tax,
        )
        row["preset_name"] = ui["preset_name"]
        results.append(row)
        print(
            f"  {ui['preset_name']:<22} "
            f"trades={row['trades']:3d} wr={row['win_rate']:5.1f}% "
            f"pnl={row['total_pnl']:>10,} pf={row['pf']:.2f} "
            f"avgW={row['avg_win']:>6,} avgL={row['avg_loss']:>7,}"
        )

    results.sort(key=lambda x: (x["total_pnl"], x["pf"]), reverse=True)
    best = results[0] if results else {}
    print(f"\n  ▶ 학습 1위 프리셋: {best.get('preset_name')} pnl={best.get('total_pnl'):,}")
    return results, best


def phase_crossval_presets(
    presets_results: List[Dict[str, Any]],
    oos_start: str,
    oos_end: str,
    top_k: int,
    *,
    slot_money: float,
    max_stocks: int,
    total_budget: float,
) -> List[Dict[str, Any]]:
    print("\n" + "=" * 70)
    print(f"[Phase 1b] R:R 프리셋 교차검증 (OOS) {oos_start} ~ {oos_end}")
    print("=" * 70)

    fixed = _mom_fixed_defaults()
    entry_base = _base_entry_ui(fixed)
    candles = _load_candles_for_search(oos_start, oos_end, int(fixed["rsi_period"]))
    universe, src = _load_universe(oos_start, oos_end)
    print(f"  캔들 종목 {len(candles)} | 유니버스={src}")

    db = TradeDB()
    try:
        row = db.conn.execute("SELECT * FROM env_config ORDER BY id DESC LIMIT 1").fetchone()
        env_row = dict(row) if row else {}
    finally:
        db.close()
    fee_rate, sell_tax, _ = sbc.fee_and_slot_from_env(env_row, strategy="MOMENTUM")

    oos_rows: List[Dict[str, Any]] = []
    for train_row in presets_results[:top_k]:
        ui = dict(entry_base)
        p = train_row.get("params_ui") or {}
        for k in ("sl_pct", "tp_pct", "tp_max_pct", "shoulder_min_high", "shoulder_cut_pct"):
            if k in p:
                ui[k] = p[k]
        ui["preset_name"] = train_row.get("preset_name", "?")
        row = _run_backtest_period(
            candles, universe, ui,
            slot_money=slot_money, max_stocks=max_stocks, total_budget=total_budget,
            fee_rate=fee_rate, sell_tax=sell_tax,
        )
        row["preset_name"] = ui["preset_name"]
        row["train_pnl"] = train_row.get("total_pnl")
        oos_rows.append(row)
        print(
            f"  {ui['preset_name']:<22} train={train_row.get('total_pnl'):>9,} "
            f"oos={row['total_pnl']:>9,} wr={row['win_rate']:.1f}% pf={row['pf']:.2f}"
        )
    return oos_rows


def _apply_rr_base_to_fixed(fixed: Dict[str, Any], best_ui: Dict[str, Any]) -> None:
    for k in ("sl_pct", "tp_pct", "tp_max_pct", "shoulder_min_high", "shoulder_cut_pct"):
        if k in best_ui:
            fixed[k] = best_ui[k]


def phase_param_search(
    train_start: str,
    train_end: str,
    mode: str,
    best_rr_ui: Dict[str, Any],
    *,
    slot_money: float,
    max_stocks: int,
    total_budget: float,
    top_n: int,
) -> Optional[str]:
    print("\n" + "=" * 70)
    print(f"[Phase 2] 파라서치 mode={mode} (R:R 1위 청산축 고정 폴백 반영)")
    print("=" * 70)
    os.environ["MOMENTUM_RR_SEARCH_BASE_JSON"] = json.dumps({
        k: best_rr_ui[k]
        for k in ("sl_pct", "tp_pct", "tp_max_pct", "shoulder_min_high", "shoulder_cut_pct")
        if k in best_rr_ui
    })
    run_search(
        train_start, train_end, mode, top_n,
        min_trades=5,
        min_win_rate=0.0,
        apply_rank=None,
        use_fallback_universe=False,
        slot_money=slot_money,
        max_stocks=max_stocks,
        total_budget_krw=total_budget,
        time_start_hm=830,
        time_end_hm=1530,
    )
    from kis_trader.backtest.param_search_momentum import _latest_json
    return _latest_json("search_momentum_")


def phase_crossval_search_json(
    json_path: Optional[str],
    oos_start: str,
    oos_end: str,
    top_k: int,
    *,
    slot_money: float,
    max_stocks: int,
    total_budget: float,
) -> List[Dict[str, Any]]:
    from kis_trader.backtest.param_search_momentum import _latest_json

    path = json_path or _latest_json("search_momentum_")
    if not path or not os.path.isfile(path):
        print("  ⚠️ search_momentum JSON 없음 — Phase 3 스킵")
        return []

    print("\n" + "=" * 70)
    print(f"[Phase 3] 파라서치 상위 {top_k} → OOS {oos_start}~{oos_end}")
    print(f"  JSON: {path}")
    print("=" * 70)

    with open(path, "r", encoding="utf-8") as f:
        doc = json.load(f)
    ranked = doc.get("top") or doc.get("ranked") or doc.get("results") or []
    if not ranked:
        print("  top/ranked 비어 있음")
        return []

    fixed = _mom_fixed_defaults()
    apply_session_to_fixed(fixed, time_start_hm=830, time_end_hm=1530)
    candles = _load_candles_for_search(oos_start, oos_end, int(fixed["rsi_period"]))
    universe, _ = _load_universe(oos_start, oos_end)

    db = TradeDB()
    try:
        row = db.conn.execute("SELECT * FROM env_config ORDER BY id DESC LIMIT 1").fetchone()
        env_row = dict(row) if row else {}
    finally:
        db.close()
    fee_rate, sell_tax, _ = sbc.fee_and_slot_from_env(env_row, strategy="MOMENTUM")

    oos_rows: List[Dict[str, Any]] = []
    for i, item in enumerate(ranked[:top_k], 1):
        merged = dict(item.get("merged_params") or item.get("params") or {})
        ui = {**fixed, **merged}
        row = _run_backtest_period(
            candles, universe, ui,
            slot_money=slot_money, max_stocks=max_stocks, total_budget=total_budget,
            fee_rate=fee_rate, sell_tax=sell_tax,
        )
        row["rank"] = i
        row["train_pnl"] = item.get("total_pnl")
        row["train_pf"] = item.get("pf")
        oos_rows.append(row)
        print(
            f"  #{i} train_pnl={item.get('total_pnl'):>9,} "
            f"oos={row['total_pnl']:>9,} trades={row['trades']} wr={row['win_rate']:.1f}% "
            f"sl={merged.get('sl_pct')} tp={merged.get('tp_pct')} "
            f"sh={merged.get('shoulder_min_high')}/{merged.get('shoulder_cut_pct')}"
        )
    return oos_rows


def main() -> None:
    ap = argparse.ArgumentParser(description="모멘텀 R:R 스윕 + 파라서치 + 교차검증")
    ap.add_argument("--train-start", default="2026-05-11")
    ap.add_argument("--train-end", default="2026-05-30")
    ap.add_argument("--oos-start", default="2026-06-01")
    ap.add_argument("--oos-end", default="2026-06-01")
    ap.add_argument("--slot-money", type=float, default=None)
    ap.add_argument("--max-stocks", type=int, default=20)
    ap.add_argument("--total-budget", type=float, default=2_000_000)
    ap.add_argument("--search-mode", default="rr", choices=["rr", "fast", "coarse"])
    ap.add_argument("--search-top", type=int, default=30)
    ap.add_argument("--cross-top", type=int, default=5)
    ap.add_argument("--skip-search", action="store_true")
    ap.add_argument("--skip-sweep", action="store_true")
    args = ap.parse_args()

    slot = float(args.slot_money or get_env_float("MOMENTUM_SLOT_MONEY", 200_000))
    t0 = time.time()
    report: Dict[str, Any] = {
        "generated_at": datetime.now().isoformat(timespec="seconds"),
        "train": [args.train_start, args.train_end],
        "oos": [args.oos_start, args.oos_end],
        "portfolio": {
            "slot_money": slot,
            "max_stocks": args.max_stocks,
            "total_budget_krw": args.total_budget,
        },
    }

    sweep_results: List[Dict[str, Any]] = []
    best_rr: Dict[str, Any] = {}
    if not args.skip_sweep:
        sweep_results, best_rr = phase_rr_sweep(
            args.train_start, args.train_end,
            slot_money=slot, max_stocks=args.max_stocks, total_budget=args.total_budget,
        )
        report["rr_sweep_train"] = sweep_results
        report["rr_sweep_oos"] = phase_crossval_presets(
            sweep_results, args.oos_start, args.oos_end, args.cross_top,
            slot_money=slot, max_stocks=args.max_stocks, total_budget=args.total_budget,
        )

    best_ui = (best_rr.get("params_ui") or {}) if best_rr else {}
    json_path: Optional[str] = None
    if not args.skip_search and best_ui:
        json_path = phase_param_search(
            args.train_start, args.train_end, args.search_mode, best_ui,
            slot_money=slot, max_stocks=args.max_stocks, total_budget=args.total_budget,
            top_n=args.search_top,
        )
        report["search_json"] = json_path

    report["search_oos"] = phase_crossval_search_json(
        json_path, args.oos_start, args.oos_end, args.cross_top,
        slot_money=slot, max_stocks=args.max_stocks, total_budget=args.total_budget,
    )

    out_path = os.path.join(
        _results_dir(),
        f"momentum_rr_crossval_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
    )
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)
    print(f"\n✅ 리포트 저장: {out_path}")
    print(f"⏱ 총 {time.time() - t0:.0f}초")


if __name__ == "__main__":
    main()