import itertools
from typing import Optional

import optuna
from ldnode.calculator.schemas import (
    HyperparameterObjectiveFunction,
    HyperparameterSearchScope,
)


def planning_hyperparameter_search(
    study_id: str,
    n_trials: int,
    objective_functions: list[HyperparameterObjectiveFunction],
    space: dict[str, HyperparameterSearchScope],
    space_size: int,
    use_grid: bool = False,
    existing_unique_keys: Optional[set[str]] = None,
    max_trials: Optional[int] = None,
    study_tag: Optional[str] = None,
    load_if_exists: bool = True,
) -> tuple[optuna.Study, int]:
    """规划超参数搜索"""
    # 优化ID
    study_name: str = f"hyperparameter-search:{study_id}"
    if study_tag:
        study_name = f"{study_name}:{study_tag}"
    # 优化方向
    directions: list[str] = [
        objective["direction"] for objective in objective_functions
    ]
    # 生成优化计划
    if use_grid:
        # 降级为网格搜索
        study = optuna.create_study(
            study_name=study_name,
            directions=directions,
            load_if_exists=load_if_exists,
        )
        # 仅生成“剩余网格”: 过滤历史/已尝试组合
        existing = existing_unique_keys or set()
        limit = max_trials if max_trials is not None else n_trials
        candidates = [scope["candidate_values"] for scope in space.values()]
        keys = [key for key in space.keys()]
        enqueued = 0
        for candidate in itertools.product(*candidates):
            unique_key = ",".join(f"{value}" for value in candidate)
            if unique_key in existing:
                continue
            study.enqueue_trial(
                {keys[index]: value for index, value in enumerate(candidate)}
            )
            enqueued += 1
            if enqueued >= limit:
                break
        return study, enqueued
    else:
        # 贝叶斯-TPE 采样优化
        effective_trials = max_trials if max_trials is not None else n_trials
        sampler = optuna.samplers.TPESampler(
            n_startup_trials=min(max(int(effective_trials / 10), 10), effective_trials),
        )
        return optuna.create_study(
            study_name=study_name,
            directions=directions,
            sampler=sampler,
            load_if_exists=load_if_exists,
        ), effective_trials
