import numpy as np
import itertools
from typing import Dict, Any, Tuple, List, Generator
from concurrent.futures import ProcessPoolExecutor, as_completed
import logging
import os
from tqdm import tqdm

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, _normalize_spec

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class GridSearchOptimizer(Optimizer):
    """
    网格搜索优化器，支持并行执行和进度条。
    """

    def __init__(self, n_jobs: int = 1, use_progress: bool = True):
        """
        参数:
            n_jobs: 并行工作进程数，-1表示使用所有CPU核心，1表示单进程
            use_progress: 是否显示进度条
        """
        self.n_jobs = n_jobs
        self.use_progress = use_progress

    def _generate_all_combinations(self, param_space: Dict[str, Any]) -> Generator[Dict[str, Any], None, None]:
        # 与原代码相同（见之前实现）
        param_names = list(param_space.keys())

        expanded_values = []
        for name in param_names:
            spec = param_space[name]
            if isinstance(spec, list):
                expanded_values.append(spec)
            elif isinstance(spec, tuple) and len(spec) == 3:
                low, high, typ = spec
                if typ == 'int':
                    expanded_values.append(list(range(int(low), int(high) + 1)))
                else:  # float
                    expanded_values.append(np.linspace(low, high, 10).tolist())
            elif isinstance(spec, dict):
                if spec['type'] == 'int':
                    low, high = spec['range']
                    expanded_values.append(list(range(int(low), int(high) + 1)))
                elif spec['type'] == 'float':
                    low, high = spec['range']
                    expanded_values.append(np.linspace(low, high, 10).tolist())
                else:  # categorical
                    expanded_values.append(spec['choices'])
            else:
                raise ValueError(f"Invalid param spec for {name}: {spec}")

        for combo in itertools.product(*expanded_values):
            params = dict(zip(param_names, combo))
            if is_valid_combo(params, param_space):
                yield params

    def _evaluate_one(self, params: Dict[str, Any],
                      original_data: np.ndarray,
                      mechanism: PrivacyMechanism,
                      risk_meter: RiskMeter,
                      utility_evaluator: UtilityEvaluator,
                      risk_constraint: Tuple[float, float]) -> Tuple[Dict[str, Any], float, float]:
        """评估单个参数组合，返回 (params, utility, risk) 或 (params, -inf, inf) 若不可行"""
        try:
            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 params, utility, risk
            else:
                return params, -np.inf, np.inf
        except Exception as e:
            logger.error(f"评估参数 {params} 时出错: {e}")
            return params, -np.inf, np.inf

    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]:

        # 生成所有合法组合
        all_params = list(self._generate_all_combinations(param_space))
        if len(all_params) > max_iter:
            all_params = all_params[:max_iter]

        best_params = None
        best_utility = -np.inf
        best_risk = None

        # 单进程执行
        if self.n_jobs == 1 or self.n_jobs == 0:
            iterator = all_params
            if self.use_progress:
                iterator = tqdm(all_params, desc="网格搜索")
            for params in iterator:
                _, utility, risk = self._evaluate_one(
                    params, original_data, mechanism, risk_meter, utility_evaluator, risk_constraint
                )
                if utility > best_utility:
                    best_utility = utility
                    best_params = params
                    best_risk = risk
        else:
            # 并行执行
            n_workers = self.n_jobs if self.n_jobs > 0 else os.cpu_count()
            logger.info(f"使用 {n_workers} 个进程并行执行网格搜索")
            with ProcessPoolExecutor(max_workers=n_workers) as executor:
                futures = []
                for params in all_params:
                    future = executor.submit(
                        self._evaluate_one,
                        params, original_data, mechanism, risk_meter, utility_evaluator, risk_constraint
                    )
                    futures.append(future)

                # 使用进度条包装 as_completed
                if self.use_progress:
                    futures_iter = tqdm(as_completed(futures), total=len(futures), desc="网格搜索")
                else:
                    futures_iter = as_completed(futures)

                for future in futures_iter:
                    params, utility, risk = future.result()
                    if utility > best_utility:
                        best_utility = utility
                        best_params = params
                        best_risk = risk

        if best_params is None:
            logger.warning("未找到满足约束的参数组合")
            return {}, -np.inf, np.inf

        logger.info(f"最优参数: {best_params}, 效用: {best_utility:.4f}, 风险: {best_risk:.4f}")
        return best_params, best_utility, best_risk