"""
参数搜索空间定义。
"""

from dataclasses import dataclass, field

import optuna


@dataclass
class SearchSpace:
    """参数搜索空间。"""

    params: list[dict] = field(default_factory=list)
    # 每个元素: {"name": str, "type": "int"|"float"|"categorical",
    #            "low": float, "high": float, "step": float, "choices": list}

    def suggest(self, trial: optuna.Trial) -> dict:
        """从 trial 中采样一组参数。"""
        result = {}
        for p in self.params:
            name = p["name"]
            kind = p["type"]
            if kind == "int":
                low = p.get("low", 0)
                high = p.get("high", 100)
                step = p.get("step", 1)
                result[name] = trial.suggest_int(name, int(low), int(high), step=int(step))
            elif kind == "float":
                low = p.get("low", 0.0)
                high = p.get("high", 1.0)
                result[name] = trial.suggest_float(name, float(low), float(high))
            elif kind == "categorical":
                choices = p.get("choices", [])
                result[name] = trial.suggest_categorical(name, choices)
        return result


# ── 预置搜索空间 ──────────────────────────────────────────

sma_cross_space = SearchSpace(params=[
    {"name": "fast", "type": "int", "low": 2, "high": 30, "step": 1},
    {"name": "slow", "type": "int", "low": 15, "high": 120, "step": 5},
])

rsi_revert_space = SearchSpace(params=[
    {"name": "oversold", "type": "int", "low": 10, "high": 45, "step": 1},
    {"name": "overbought", "type": "int", "low": 55, "high": 90, "step": 1},
])

momentum_breakout_space = SearchSpace(params=[
    {"name": "lookback", "type": "int", "low": 10, "high": 60, "step": 5},
    {"name": "exit_period", "type": "int", "low": 5, "high": 30, "step": 1},
])

factor_cross_space = SearchSpace(params=[
    {"name": "buy_threshold", "type": "float", "low": -10.0, "high": 10.0},
    {"name": "sell_threshold", "type": "float", "low": -10.0, "high": 10.0},
])

# 名称 → 空间映射
SPACES = {
    "sma_cross": sma_cross_space,
    "rsi_revert": rsi_revert_space,
    "momentum_breakout": momentum_breakout_space,
    "factor_cross": factor_cross_space,
}