# privacy_tuner/core/optimizers/bayesian.py

import numpy as np
from typing import Dict, Any, Tuple, List
from skopt import gp_minimize
from skopt.space import Real, Integer, Categorical
from skopt.utils import use_named_args

from ..optimizer import Optimizer
from ..privacy_mechanism import PrivacyMechanism
from ..risk_meter import RiskMeter
from ..utility_evaluator import UtilityEvaluator
from ..param_utils import is_valid_combo


class BayesianOptimizer(Optimizer):
    """
    基于高斯过程的贝叶斯优化器。
    适用于连续和混合参数空间，使用 skopt 的 gp_minimize 实现。
    支持参数依赖（通过 is_valid_combo 过滤非法组合）。

    参数:
        base_estimator: str，代理模型 ("GP", "RF", "ET", "GBRT")，默认 "GP"
        acq_func: str，采集函数 ("EI", "PI", "LCB")，默认 "EI"
        random_state: int，随机种子，默认 42
        n_initial_points: int，初始随机探索点数，默认 5
    """

    def __init__(self,
                 base_estimator: str = "GP",
                 acq_func: str = "EI",
                 random_state: int = 42,
                 n_initial_points: int = 5):
        self.base_estimator = base_estimator
        self.acq_func = acq_func
        self.random_state = random_state
        self.n_initial_points = n_initial_points

    def _convert_param_space(self, param_space: Dict[str, Any]) -> List:
        """
        将通用参数空间转换为 skopt 的 Space 维度列表。
        """
        dimensions = []
        self.param_names = []

        for name, spec in param_space.items():
            self.param_names.append(name)

            # 处理各种规格
            if isinstance(spec, tuple) and len(spec) == 3:
                low, high, typ = spec
                if typ == 'int':
                    dimensions.append(Integer(low, high, name=name))
                elif typ == 'float':
                    dimensions.append(Real(low, high, name=name))
                else:
                    raise ValueError(f"未知类型: {typ}")
            elif isinstance(spec, list):
                dimensions.append(Categorical(spec, name=name))
            elif isinstance(spec, dict):
                # 带依赖的参数定义
                if spec['type'] == 'int':
                    low, high = spec['range']
                    dimensions.append(Integer(low, high, name=name))
                elif spec['type'] == 'float':
                    low, high = spec['range']
                    dimensions.append(Real(low, high, name=name))
                else:  # categorical
                    dimensions.append(Categorical(spec['choices'], name=name))
            else:
                raise ValueError(f"参数空间格式不支持: {name}: {spec}")

        return dimensions

    def search(self,
               original_data: np.ndarray,
               mechanism: PrivacyMechanism,
               risk_meter: RiskMeter,
               utility_evaluator: UtilityEvaluator,
               param_space: Dict[str, Any],
               risk_constraint: Tuple[float, float],
               max_iter: int = 100) -> Tuple[Dict[str, Any], float, float]:
        """
        贝叶斯优化搜索。
        """
        # 转换参数空间
        dimensions = self._convert_param_space(param_space)

        # 定义目标函数（负效用，因为 gp_minimize 是最小化）
        @use_named_args(dimensions)
        def objective(**params):
            # 1. 检查合法性
            if not is_valid_combo(params, param_space):
                # 非法组合，返回极大惩罚
                return 1e10

            # 应用机制
            perturbed = mechanism.apply(original_data, params)
            # 计算风险
            risk = risk_meter.measure(original_data, perturbed, params=params)
            # 检查风险约束
            if risk_constraint[0] <= risk <= risk_constraint[1]:
                utility = utility_evaluator.evaluate(original_data, perturbed)
                return -utility
            else:
                return 1e10

        # 执行优化（使用 gp_minimize）
        result = gp_minimize(
            func=objective,
            dimensions=dimensions,
            n_calls=max_iter,
            n_initial_points=self.n_initial_points,
            random_state=self.random_state,
            acq_func=self.acq_func,
            base_estimator=self.base_estimator
        )

        # 提取最优参数
        best_params = dict(zip(self.param_names, result.x))

        # 计算最优参数对应的风险和效用（重新计算一次）
        perturbed = mechanism.apply(original_data, best_params)
        best_risk = risk_meter.measure(original_data, perturbed, params=best_params)
        best_utility = -result.fun  # 因为最小化负效用

        return best_params, best_utility, best_risk