# privacy_tuner/core/param_utils.py

import numpy as np
from typing import Dict, Any, Optional, List


def _normalize_spec(spec: Any) -> Dict:
    """
    将各种格式的参数规格统一转换为字典格式。
    支持:
        - 列表: 转为 {'type': 'categorical', 'choices': spec}
        - 元组 (low, high, type): 转为 {'type': type, 'range': (low, high)}
        - 字典: 原样返回
    """
    if isinstance(spec, list):
        return {'type': 'categorical', 'choices': spec}
    elif isinstance(spec, tuple) and len(spec) == 3:
        low, high, typ = spec
        return {'type': typ, 'range': (low, high)}
    elif isinstance(spec, dict):
        return spec
    else:
        raise ValueError(f"Invalid parameter spec: {spec}")


def is_valid_combo(params: Dict[str, Any], space_def: Dict[str, Any]) -> bool:
    """
    检查一组参数是否满足所有依赖条件。
    只检查 params 中存在的参数，忽略未出现的参数。
    """
    for param_name in params:
        if param_name not in space_def:
            return False
        norm_def = _normalize_spec(space_def[param_name])
        deps = norm_def.get('dependencies', {})
        for dep_param, allowed_values in deps.items():
            if dep_param not in params:
                return False
            if params[dep_param] not in allowed_values:
                return False
    return True


def sample_valid_combo(space_def: Dict[str, Any], max_trials: int = 1000) -> Optional[Dict[str, Any]]:
    """
    随机采样一个满足所有依赖的合法参数组合。
    组合可以只包含部分参数（只要满足依赖即可）。
    """
    normalized = {name: _normalize_spec(spec) for name, spec in space_def.items()}

    for _ in range(max_trials):
        candidate = {}
        remaining = set(normalized.keys())

        while remaining:
            # 找出当前可采样的参数：其所有依赖都已采样且依赖条件满足
            ready = []
            for name in remaining:
                norm = normalized[name]
                deps = norm.get('dependencies', {})
                satisfied = True
                for dep_name, allowed in deps.items():
                    if dep_name not in candidate:
                        satisfied = False
                        break
                    if candidate[dep_name] not in allowed:
                        satisfied = False
                        break
                if satisfied:
                    ready.append(name)

            if not ready:
                # 无法继续，返回当前已采样的组合
                break

            # 从 ready 中随机选择一个参数进行采样
            name = np.random.choice(ready)
            norm = normalized[name]
            if norm['type'] == 'int':
                low, high = norm['range']
                candidate[name] = np.random.randint(low, high + 1)
            elif norm['type'] == 'float':
                low, high = norm['range']
                candidate[name] = np.random.uniform(low, high)
            else:  # categorical
                candidate[name] = np.random.choice(norm['choices'])

            remaining.remove(name)

        # 返回当前候选（即使还有剩余参数未采样）
        return candidate

    return None