import copy
import threading as th
import time
from typing import Any, Callable, Optional

import optuna

from ldnode.calculator.plans.optuna import planning_hyperparameter_search
from ldnode.calculator.schemas import (
    CancellationToken,
    HyperparameterSearchConfig,
    HyperparameterSearchMode,
    HyperparameterSearchProgress,
    HyperparameterSearchResult,
    HyperparameterSearchSuggestion,
    HyperparameterSearchTrial,
    HyperparameterSearchTrialParam,
    HyperTuningScopeArgName,
    ResultCollector,
)


class HyperparameterSearchStorage:
    """超参搜索存储"""

    def __init__(self, init_data: Optional[list[HyperparameterSearchResult]] = None):
        self.lock = th.Lock()
        self.unique_keys: set[str] = set(
            [result["unique_id"] for result in init_data] if init_data else []
        )
        self.cache_metrics: dict[str, dict[str, Any]] = (
            {result["normalized_id"]: result["metrics"] for result in init_data}
            if init_data
            else {}
        )
        self.finished = 0

    def try_get_cache_metrics(self, cache_key: str) -> dict[str, Any] | None:
        """是否存在缓存指标"""
        with self.lock:
            return self.cache_metrics.get(cache_key)

    def try_add(self, unique_key: str) -> bool:
        """尝试添加"""
        with self.lock:
            if unique_key in self.unique_keys:
                return False
            else:
                self.unique_keys.add(unique_key)
                self.finished += 1
                return True

    def append_metrics(self, cache_key: str, metrics: dict[str, Any]) -> None:
        """添加缓存指标"""
        with self.lock:
            self.cache_metrics[cache_key] = metrics

    def giveup_result(self) -> None:
        """放弃结果"""
        with self.lock:
            self.finished -= 1

    def get_finished(self) -> int:
        """获取已完成数量"""
        with self.lock:
            return self.finished

    def get_unique_count(self) -> int:
        """获取唯一组合数量"""
        with self.lock:
            return len(self.unique_keys)

    def get_unique_snapshot(self) -> set[str]:
        """获取唯一组合快照"""
        with self.lock:
            return set(self.unique_keys)


class HyperparameterOptimizer:
    """超参优化器"""

    def __init__(
        self,
        study_id: str,
        method: str,
        config: HyperparameterSearchConfig,
        parameters: dict[str, Any],
        test_begin_date: Optional[str],
        storage: HyperparameterSearchStorage,
        invoke_handler: Callable[[str, dict[str, Any], float], Any],
        cancel_token: CancellationToken,
        result_collector: ResultCollector,
    ):
        self.study_id = study_id
        self.method = method
        self.parameters = parameters
        self.test_begin_date = test_begin_date
        self.storage = storage
        self.invoke = invoke_handler
        self.cancel_token = cancel_token
        self.collector = result_collector
        self.n_trials = config["n_trails"]
        self.objective_functions = config["objective_functions"]
        self.space = HyperparameterSearchConfig.get_space(config["scopes"])
        self.space_size = config["space_size"]
        self.history_unique_count = self.storage.get_unique_count()
        self.is_infinite_space = self.space_size == -1
        self.remaining_space = (
            -1
            if self.is_infinite_space
            else max(self.space_size - self.history_unique_count, 0)
        )
        # 续接感知模式判断:
        # - 无穷空间: 永远贝叶斯
        # - 有穷空间: 本次预算可覆盖剩余空间 -> 网格
        self.mode = (
            HyperparameterSearchMode.BAYESIAN
            if self.is_infinite_space
            else (
                HyperparameterSearchMode.GRID
                if self.remaining_space > 0 and self.n_trials >= self.remaining_space
                else HyperparameterSearchMode.BAYESIAN
            )
        )
        self.total_suggest_num = 0
        self.repeated_suggest_num = 0
        self.force_grid_remaining = False
        self.tpe_round = 0
        self.round_total_suggest_num = 0
        self.round_repeated_suggest_num = 0
        self.round_new_unique_num = 0
        self.no_progress_rounds = 0
        self.max_no_progress_rounds = 3
        self.max_total_suggest = max(self.n_trials * 20, 200)
        self.state_lock = th.Lock()
        self.cached_trials: list[HyperparameterSearchTrial] = []
        self.trial_collect_lock = th.Lock()
        self.latest_collect_at = time.time()

    @property
    def repeated_rate(self) -> float:
        """重复率"""
        with self.state_lock:
            if self.total_suggest_num == 0:
                return 0.0
            return self.repeated_suggest_num / self.total_suggest_num

    @property
    def is_finished(self) -> bool:
        """是否完成"""
        return self.storage.get_finished() >= self.n_trials

    @property
    def is_space_exhausted(self) -> bool:
        """有穷空间是否已耗尽(含历史结果)"""
        if self.is_infinite_space:
            return False
        return self.storage.get_unique_count() >= self.space_size

    @property
    def round_repeated_rate(self) -> float:
        """当前轮次重复率"""
        with self.state_lock:
            if self.round_total_suggest_num == 0:
                return 0.0
            return self.round_repeated_suggest_num / self.round_total_suggest_num

    def begin_round(self) -> None:
        """开始新一轮搜索统计"""
        with self.state_lock:
            self.round_total_suggest_num = 0
            self.round_repeated_suggest_num = 0
            self.round_new_unique_num = 0

    def suggest_hyperparameters(
        self, trial: optuna.Trial
    ) -> list[HyperparameterSearchSuggestion]:
        """建议超参数集"""
        with self.state_lock:
            self.total_suggest_num += 1
            self.round_total_suggest_num += 1
        return [
            {
                "key_path": scope["key_path"],
                "value": trial.suggest_categorical(
                    hyper_name, scope["candidate_values"]
                ),
                "index": scope["index"],
            }
            for hyper_name, scope in self.space.items()
        ]

    def unique_and_normalize(
        self, hyperparameters: list[HyperparameterSearchSuggestion]
    ) -> tuple[str, str]:
        """计算超参数集唯一键和归一化键"""
        unique_key = ",".join(
            f"{hyperparameter['value']}" for hyperparameter in hyperparameters
        )
        # 归一化打分因子权重参数
        rank_weight_params = [
            hyperparameter
            for hyperparameter in hyperparameters
            if hyperparameter["index"]["arg_name"]
            == HyperTuningScopeArgName.RANK_WEIGHT
        ]
        sum_weight = sum(
            hyperparameter["value"] for hyperparameter in rank_weight_params
        )
        normalized_weight_key = ",".join(
            f"{round(hyperparameter['value'] / sum_weight if sum_weight != 0 else 0, 4)}"
            for hyperparameter in rank_weight_params
        )
        # 归一化打分因子参数
        rank_factor_params = [
            hyperparameter["value"]
            for hyperparameter in hyperparameters
            if hyperparameter["index"]["arg_name"]
            == HyperTuningScopeArgName.RANK_FACTOR
        ]
        normalized_factor_key = ",".join(sorted(rank_factor_params))
        # 其它键不需要归一化
        other_params = [
            hyperparameter
            for hyperparameter in hyperparameters
            if hyperparameter["index"]["arg_name"]
            not in [
                HyperTuningScopeArgName.RANK_WEIGHT,
                HyperTuningScopeArgName.RANK_FACTOR,
            ]
        ]
        other_params_key = ",".join(
            f"{hyperparameter['value']}" for hyperparameter in other_params
        )
        # 生成归一化键
        normalized_key = (
            f"{normalized_weight_key},{normalized_factor_key},{other_params_key}"
        )
        return unique_key, normalized_key

    def build_parameters(
        self, hyperparameters: list[HyperparameterSearchSuggestion]
    ) -> dict[str, Any]:
        """构建参数"""
        new_parameters = copy.deepcopy(self.parameters)
        for hyperparameter in hyperparameters:
            parameters = new_parameters
            for key in hyperparameter["key_path"][:-1]:
                if type(parameters) is dict:
                    parameters = parameters[key]
                elif type(parameters) is list:
                    parameters = parameters[int(key)]
            parameters[hyperparameter["key_path"][-1]] = hyperparameter["value"]
        return new_parameters

    def collect_trial(
        self,
        number: int,
        unique_key: str,
        normalized_key: str,
        params: list[HyperparameterSearchTrialParam],
        metrics: dict[str, Any],
    ) -> None:
        """收集试验结果"""
        trial: HyperparameterSearchTrial = {
            "number": number,
            "unique_key": unique_key,
            "normalized_key": normalized_key,
            "params": params,
            "metrics": metrics,
        }
        # 如果超过3秒则收集试验结果
        with self.trial_collect_lock:
            self.cached_trials.append(trial)
            current_time = time.time()
            if current_time - self.latest_collect_at > 3:
                self.collector.collect(
                    {"data_type": "trials", "data": self.cached_trials}
                )
                self.cached_trials = []
                self.latest_collect_at = current_time

    def objective(self, trial: optuna.Trial) -> float | tuple[float, ...]:
        """目标函数"""
        hyperparameters = self.suggest_hyperparameters(trial)
        unique_key, normalized_key = self.unique_and_normalize(hyperparameters)
        # 尝试添加唯一键, 如果失败则剪枝
        if not self.storage.try_add(unique_key):
            with self.state_lock:
                self.repeated_suggest_num += 1
                self.round_repeated_suggest_num += 1
            raise optuna.TrialPruned("组合超参数重复，剪枝")
        with self.state_lock:
            self.round_new_unique_num += 1
        # 计算指标
        if cache_metrics := self.storage.try_get_cache_metrics(normalized_key):
            # 命中缓存, 则直接赋予缓存结果
            metrics = cache_metrics
        else:
            # 构造计算参数
            parameters = self.build_parameters(hyperparameters)
            # 检查打分因子是否重复
            rank_factors = parameters.get("strategy_config", {}).get("rank_factors", [])
            rank_factor_keys = [
                rank_factor.get("factor") for rank_factor in rank_factors
            ]
            # 插入测试集起始日期
            if self.test_begin_date:
                parameters["test_begin_date"] = self.test_begin_date
            try:
                if len(rank_factor_keys) != len(set(rank_factor_keys)):
                    raise Exception("打分因子重复")
                # 调用评价函数计算
                try:
                    result: dict[str, Any] = self.invoke(self.method, parameters, 600)
                    if result.get("data_type", "error") == "error":
                        raise Exception(
                            result.get("data", f"{self.method} 返回数据格式错误")
                        )
                    if "data" not in result:
                        raise Exception("目标函数计算返回数据格式错误")
                    metrics = result["data"]
                    # 更新指标值
                    self.storage.append_metrics(normalized_key, metrics)
                except Exception as e:
                    raise Exception(f"目标函数计算失败: {e}")
            except Exception:
                # 打分因子重复或者其他异常 如果是贝叶斯应该放弃掉这组参数继续搜索 网格则直接拒绝掉这组参数
                if self.mode == HyperparameterSearchMode.BAYESIAN:
                    self.storage.giveup_result()
                raise
        # 提交指标计算结果
        self.collect_trial(
            trial.number,
            unique_key,
            normalized_key,
            [
                {
                    "arg_name": hyperparameter["index"]["arg_name"],
                    "index": hyperparameter["index"]["index"],
                    "value": hyperparameter["value"],
                }
                for hyperparameter in hyperparameters
            ],
            metrics,
        )
        # 返回目标值
        if len(self.objective_functions) == 1:
            return metrics[self.objective_functions[0]["function"]]
        else:
            return tuple(
                metrics[objective_function["function"]]
                for objective_function in self.objective_functions
            )

    def control_callback(
        self, study: optuna.Study, trial: optuna.trial.FrozenTrial
    ) -> None:
        """控制回调"""
        # 硬停止条件
        if (
            self.cancel_token.is_canceled()
            or self.is_finished
            or self.is_space_exhausted
        ):
            study.stop()
        # 生成进度
        progress: Optional[HyperparameterSearchProgress] = None
        match trial.state:
            case optuna.trial.TrialState.COMPLETE:
                progress = {
                    "study_id": self.study_id,
                    "is_success": True,
                    "error_message": "",
                }
            case optuna.trial.TrialState.FAIL:
                # 网格模式下 异常应该直接上报 贝叶斯模式下则放弃异常继续搜索
                if self.mode == HyperparameterSearchMode.GRID:
                    progress = {
                        "study_id": self.study_id,
                        "is_success": False,
                        "error_message": "目标函数计算失败",
                    }
        # 上报进度数据
        if progress:
            self.collector.collect({"data_type": "progress", "data": progress})

    def search(self, n_jobs: int = 9):
        """超参搜索"""
        try:
            # 开始搜索: 不使用固定轮次上限，完全由任务状态驱动
            while not self.is_finished and not self.is_space_exhausted:
                if self.cancel_token.is_canceled():
                    break
                if self.total_suggest_num >= self.max_total_suggest:
                    break
                remaining_budget = max(self.n_trials - self.storage.get_finished(), 0)
                if remaining_budget == 0:
                    break
                if self.is_infinite_space:
                    max_trials = remaining_budget
                else:
                    remaining_space_now = max(
                        self.space_size - self.storage.get_unique_count(), 0
                    )
                    max_trials = min(remaining_budget, remaining_space_now)
                if max_trials == 0:
                    break
                self.begin_round()
                with self.state_lock:
                    use_grid_this_round = (
                        self.mode == HyperparameterSearchMode.GRID
                        or self.force_grid_remaining
                    )
                    current_tpe_round = self.tpe_round
                study, planned_trials = planning_hyperparameter_search(
                    self.study_id,
                    max_trials,
                    self.objective_functions,
                    self.space,
                    self.space_size,
                    use_grid=use_grid_this_round,
                    existing_unique_keys=self.storage.get_unique_snapshot(),
                    max_trials=max_trials,
                    study_tag=(
                        "grid-fill"
                        if use_grid_this_round
                        else f"tpe-round-{current_tpe_round}"
                    ),
                    load_if_exists=False,
                )
                if planned_trials == 0:
                    break
                study.optimize(
                    self.objective,
                    n_trials=planned_trials,
                    n_jobs=n_jobs,
                    show_progress_bar=False,
                    catch=(Exception,),
                    callbacks=[self.control_callback],
                )
                if not use_grid_this_round:
                    with self.state_lock:
                        self.tpe_round += 1
                with self.state_lock:
                    round_new_unique_num = self.round_new_unique_num
                    if round_new_unique_num == 0:
                        self.no_progress_rounds += 1
                    else:
                        self.no_progress_rounds = 0
                    no_progress_rounds = self.no_progress_rounds
                # 贝叶斯每轮后根据效果决定: 重启一轮TPE，或切到网格补齐
                with self.state_lock:
                    force_grid_remaining = self.force_grid_remaining
                if (
                    self.mode == HyperparameterSearchMode.BAYESIAN
                    and not self.is_infinite_space
                    and not force_grid_remaining
                    and not self.cancel_token.is_canceled()
                    and not self.is_finished
                ):
                    remaining_space_now = max(
                        self.space_size - self.storage.get_unique_count(), 0
                    )
                    # 20%重复率触发重启判断，同时要求本轮新增过少，避免误判
                    new_unique_too_low = round_new_unique_num <= max(
                        1, int(planned_trials * 0.1)
                    )
                    if self.round_repeated_rate >= 0.2 and new_unique_too_low:
                        # 剩余空间较小时，直接转网格补齐；否则开始下一轮TPE
                        if remaining_space_now <= max(planned_trials, 64):
                            with self.state_lock:
                                self.force_grid_remaining = True
                    # 连续无进展轮次过多，强制切到网格补齐防止外层空转
                    if no_progress_rounds >= self.max_no_progress_rounds:
                        with self.state_lock:
                            self.force_grid_remaining = True
                # 无穷空间下连续无进展则退出，避免无限循环
                if (
                    self.is_infinite_space
                    and no_progress_rounds >= self.max_no_progress_rounds
                ):
                    break
            # 收集剩余试验结果
            if self.cached_trials:
                self.collector.collect({"data_type": "trials", "data": self.cached_trials})
        finally:
            # 无论正常结束还是异常，都必须关闭收集器，否则 collector.read() 将永远阻塞
            self.collector.close()
