import typing
from collections.abc import Sequence
from typing import Any, Dict, List, Optional, Callable, Iterator

# import numpy as np
from ConfigSpace import ConfigurationSpace, Float, Integer, Categorical, EqualsCondition


class HyperparameterSpace:
    """基于 ConfigSpace 的增强版超参数空间管理器"""

    def __init__(
        self,
        name: str = 'HyperparameterSpace',
        seed: Optional[int] = None,
        space: Optional[ConfigurationSpace] = None,
    ):
        """初始化超参数空间

        Parameters
        ----------
        - `name` : `str`, optional
            - 空间名称, by default 'HyperparameterSpace'
        - `seed` : `int`, optional
            - 随机种子, by default None
        """
        self.space = ConfigurationSpace(name=name, seed=seed) if space is None else space
        self._dynamic_fn: Optional[
            Callable[[int, typing.Mapping[str, Any]], typing.Mapping[str, Any]]
        ] = None
        self._sample_index = 0
        self._param_maps: Dict[str, Dict[str, Dict[str, Any]]] = {}

    def add_uniform(
        self,
        name: str,
        low: float,
        high: float,
        condition: Optional[typing.Mapping[str, Any]] = None,
    ):
        """添加均匀分布参数

        Parameters
        ----------
        - `name` : `str`
            - 参数名称
        - `low` : `float`
            - 下界
        - `high` : `float`
            - 上界
        - `condition` : `Dict[str, Any]`, optional
            - 条件依赖, by default None
        """
        param = Float(name, bounds=(low, high))
        self.space.add(param)

        if condition:
            self._add_condition(name, condition)
        return self

    def add_log_uniform(
        self,
        name: str,
        low: float,
        high: float,
        condition: Optional[typing.Mapping[str, Any]] = None,
    ):
        """添加对数均匀分布参数

        Parameters
        ----------
        - `name` : `str`
            - 参数名称
        - `low` : `float`
            - 下界
        - `high` : `float`
            - 上界
        - `condition` : `Dict[str, Any]`, optional
            - 条件依赖, by default None
        """
        param = Float(name, bounds=(low, high), log=True)
        self.space.add(param)

        if condition:
            self._add_condition(name, condition)
        return self

    def add_int_uniform(
        self,
        name: str,
        low: int,
        high: int,
        step: int = 1,
        condition: Optional[typing.Mapping[str, Any]] = None,
    ):
        """添加整型均匀分布参数

        Parameters
        ----------
        - `name` : `str`
            - 参数名称
        - `low` : `int`
            - 下界
        - `high` : `int`
            - 上界
        - `step` : `int`, optional
            - 步长, by default 1
        - `condition` : `Dict[str, Any]`, optional
            - 条件依赖, by default None
        """
        if step == 1:
            param = Integer(name, bounds=(low, high))
        else:
            # ConfigSpace 不直接支持 step，我们用 Categorical 模拟
            choices = list(range(low, high + 1, step))
            param = Categorical(name, choices)

        self.space.add(param)

        if condition:
            self._add_condition(name, condition)
        return self

    def add_categorical(
        self,
        name: str,
        choices: typing.Sequence[Any],
        condition: Optional[typing.Mapping[str, Any]] = None,
        weights: Sequence[float] | None = None,
    ):
        """添加类别型参数

        Parameters
        ----------
        - `name` : `str`
            - 参数名称
        - `choices` : `List[Any]`
            - 选择列表
        - `condition` : `Dict[str, Any]`, optional
            - 条件依赖, by default None
        """
        param = Categorical(name, choices, weights=weights)
        self.space.add(param)

        if condition:
            self._add_condition(name, condition)
        return self

    def add_combination(
        self,
        params: typing.Sequence[str],
        values: typing.Sequence[typing.Sequence[Any]],
        weights: Optional[typing.Mapping[str, float]] = None,
        condition: Optional[typing.Mapping[str, Any]] = None,
        combination_name: str = 'combination',
    ):
        """添加参数组合选择器

        将多个参数打包成组合，每次采样时选择一个组合。

        Parameters
        ----------
        - `params` : `list[str]`
            - 参数名称列表
        - `values` : `list[list[Any]]`
            - 参数值组合列表，每个子列表对应一个组合的参数值
        - `weights` : `Dict[str, float]`, optional
            - 权重映射，格式：{"Comb 0": weight1, "Comb 1": weight2, ...}, by default None
        - `condition` : `Dict[str, Any]`, optional
            - 条件依赖, by default None
        - `combination_name` : `str`, optional
            - 组合参数名称, by default 'combination'

        Examples
        --------
        ```python
        space.add_combination(
            params=['optimizer_type', 'lr', 'beta1'],
            values=[
                ['adam', 0.001, 0.9],
                ['sgd', 0.01, None],
                ['rmsprop', 0.002, None]
            ],
            combination_name='optimizer'
        )
        ```
        """
        if len(params) == 0:
            raise ValueError('参数列表不能为空')

        # 验证每个组合的参数数量是否匹配
        for i, value_combo in enumerate(values):
            if len(value_combo) != len(params):
                raise ValueError(
                    f'组合 {i} 的参数数量({len(value_combo)})与参数名数量({len(params)})不匹配'
                )

        # 生成组合选择项
        choices = [f'Comb {i}' for i in range(len(values))]

        # 处理权重：从字典转换为序列
        weight_sequence = None
        if weights:
            # 验证权重键是否匹配选择项
            for choice in choices:
                if choice not in weights:
                    raise ValueError(f"权重字典中缺少选择项 '{choice}' 的权重")
            weight_sequence = [weights[choice] for choice in choices]

        # 构造参数映射字典
        param_combinations = {}
        for i, value_combo in enumerate(values):
            combo_name = choices[i]
            param_combinations[combo_name] = dict(zip(params, value_combo))

        # 处理重复的组合名称
        original_name = combination_name
        counter = 1
        while combination_name in self._param_maps:
            combination_name = f'{original_name}_{counter}'
            counter += 1

        # 添加分类参数（修复权重类型问题）
        self.add_categorical(
            combination_name, choices, condition, weights=weight_sequence
        )

        # 存储映射关系
        self._param_maps[combination_name] = param_combinations

        return self

    def add_dynamic(
        self,
        fn_generator: Callable[[int, typing.Mapping[str, Any]], typing.Mapping[str, Any]],
    ):
        """添加动态参数生成器

        Parameters
        ----------
        - `fn_generator` : `Callable[[int, Dict[str, Any]], Dict[str, Any]]`
            - 动态参数生成函数，接受采样索引和当前参数，返回新参数
        """
        if self._dynamic_fn is not None:
            raise ValueError('动态参数生成器只能添加一次')
        self._dynamic_fn = fn_generator
        return self

    def _add_condition(self, param_name: str, condition: typing.Mapping[str, Any]):
        """添加条件依赖

        Parameters
        ----------
        - `param_name` : `str`
            - 子参数名称
        - `condition` : `Dict[str, Any]`
            - 条件字典，格式：{父参数名: 期望值}
        """
        if len(condition) != 1:
            raise ValueError('条件必须是包含单一键值对的字典')

        parent_param, required_value = list(condition.items())[0]

        if parent_param not in [hp.name for hp in self.space.values()]:
            raise ValueError(f"依赖的父参数 '{parent_param}' 不存在")

        parent_hp = self.space[parent_param]
        child_hp = self.space[param_name]

        cond = EqualsCondition(child_hp, parent_hp, required_value)
        self.space.add(cond)

    def sample(self) -> Dict[str, Any]:
        """从空间中采样一组参数

        Returns
        -------
        - `Dict[str, Any]`
            - 采样得到的参数字典
        """
        config = self.space.sample_configuration()
        result = dict(config)

        # 处理组合参数映射
        for map_name, param_map in self._param_maps.items():
            if map_name in result:
                combo_choice = result.pop(map_name)
                if combo_choice in param_map:
                    # 展开组合参数到结果中
                    combo_params = param_map[combo_choice]
                    for param_name, param_value in combo_params.items():
                        # 只添加非None的参数值
                        if param_value is not None:
                            result[param_name] = param_value

        # 应用动态生成器
        if self._dynamic_fn:
            dynamic_params = self._dynamic_fn(self._sample_index, result.copy())
            if isinstance(dynamic_params, dict):
                result.update(dynamic_params)

        self._sample_index += 1
        return result

    def random_search(self, n_trials: int) -> Iterator[Dict[str, Any]]:
        """执行随机搜索

        Parameters
        ----------
        - `n_trials` : `int`
            - 搜索试验次数

        Yields
        ------
        - `Dict[str, Any]`
            - 每次试验的参数字典
        """
        # 修复：移除了错误的size=1参数
        for _ in range(n_trials):
            yield self.sample()

    def reset_sample_index(self):
        """重置采样索引

        Returns
        -------
        - `HyperparameterSpace`
            - 返回self以支持链式调用
        """
        self._sample_index = 0
        return self

    def get_space_info(self) -> Dict[str, Any]:
        """获取空间信息

        Returns
        -------
        - `Dict[str, Any]`
            - 包含空间统计信息的字典
        """
        hyperparams = list(self.space.values())
        return {
            'name': self.space.name,
            'n_hyperparameters': len(hyperparams),
            'hyperparameter_names': [hp.name for hp in hyperparams],
            'n_param_maps': len(self._param_maps),
            'has_dynamic_fn': self._dynamic_fn is not None,
            'sample_index': self._sample_index,
        }

    def __repr__(self) -> str:
        info = self.get_space_info()
        return f"HyperparameterSpace(name='{info['name']}', n_hyperparameters={info['n_hyperparameters']}, n_param_maps={info['n_param_maps']})"

    def to_yaml(self, path: str):
        self.space.to_yaml(path)

    @staticmethod
    def from_yaml(path: str) -> 'HyperparameterSpace':
        space = ConfigurationSpace.from_yaml(path)
        return HyperparameterSpace(space=space)
