from typing import TypedDict, Any

from ldnode.calculator.schemas.enumerations import HyperTuningScopeArgName


class HyperparameterSearchScopeIndex(TypedDict):
    """ 超参数搜索范围索引 """
    arg_name: HyperTuningScopeArgName # 参数名称
    index: int # 参数索引


class HyperparameterSearchScope(TypedDict):
    """ 超参数搜索范围 """
    key_path: str # 键路径
    candidate_values: list[Any] # 候选值
    index: HyperparameterSearchScopeIndex # 索引


class HyperparameterObjectiveFunction(TypedDict):
    """ 超参数目标函数 """
    function: str # 函数
    direction: str # 方向


class HyperparameterSearchConfig(TypedDict):
    """ 超参数搜索配置 """
    space_size: int # 空间大小
    n_trails: int # 试验次数
    scopes: list[HyperparameterSearchScope] # 范围
    objective_functions: list[HyperparameterObjectiveFunction] # 目标函数
    
    @staticmethod
    def get_space(scopes: list[HyperparameterSearchScope]) -> dict[str, HyperparameterSearchScope]:
        """ 获取空间 """
        return {
            f"hyper_param_{i}": scope
            for i, scope in enumerate(scopes)
        }


class HyperparameterSearchResult(TypedDict):
    """ 超参数搜索结果 """
    unique_id: str # 唯一ID
    normalized_id: str # 归一化ID
    metrics: dict[str, Any] # 指标
    

class HyperparameterSearchTrialParam(TypedDict):
    """ 超参数搜索试验参数 """
    arg_name: HyperTuningScopeArgName # 参数名称
    index: int # 参数索引
    value: Any # 参数值
    

class HyperparameterSearchTrial(TypedDict):
    """ 超参数搜索试验 """
    number: int # 试验编号
    unique_key: str # 唯一ID
    normalized_key: str # 归一化ID
    params: list[HyperparameterSearchTrialParam] # 参数
    metrics: dict[str, Any] # 指标

    
class HyperparameterSearchSuggestion(TypedDict):
    """ 超参数搜索建议 """
    key_path: str # 键路径
    value: Any # 值
    index: HyperparameterSearchScopeIndex # 索引


class HyperparameterSearchProgress(TypedDict):
    """ 超参数搜索进度 """
    study_id: str # 优化ID
    is_success: bool # 是否成功
    error_message: str # 错误信息
    