#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2024/4/30
file:       base_callbacks.py
version:
doc:        通用回调函数, 基本功能为保存模型.

由于该模块的调用在ray的框架内实现无法改动, 因此无法通过形参传参.

ray创建Trainer时会进行参数检查, 传入自定义参数会报错.

目前的尝试中只发现可以从2个地方传入一个dict用来描述所需参数, 分别是

config['env_config']: dict = {Any: Any}  # 可以添加自定义参数字典
config['evaluation_config'] = {
    'in_evaluation': True,  # 必须有
    Any: Any,  # 可以添加自定义参数字典
}
"""
import json
import os
import sys
from typing import Dict, List, Tuple
import ray
import numpy as np
from keras.models import load_model
from ray.rllib.algorithms.algorithm import Algorithm
from ray.rllib.algorithms.callbacks import DefaultCallbacks
from ray.rllib.evaluation.rollout_worker import RolloutWorker
from ray.rllib.evaluation.worker_set import WorkerSet
# import torch
from ray.rllib.policy.policy import Policy, PolicySpec
from ray.rllib.utils.checkpoints import get_checkpoint_info
from ray.rllib.utils.metrics.learner_info import LEARNER_INFO, LEARNER_STATS_KEY
from ray.rllib.algorithms.dqn.dqn_tf_policy import DQNTFPolicy
from corekit.hikframe import HikModel
from corekit.utils import DESKTOP, WORK_DIR
from corekit.utils import get_debug_logger
from corekit.utils.typings import PolicyID
from pathlib import Path

_logger = get_debug_logger()
# 以下路径为测试用
a_folder = os.path.join(DESKTOP, 'temp', 'a-folder')  # 用在 Algorithm 的API
# a_ckpt = os.path.join(DESKTOP, 'temp', 'a-ckpt')
# os.makedirs(a_ckpt, exist_ok=True)
# p_model = os.path.join(DESKTOP, 'temp', f'p-model')
p_ckpt = os.path.join(DESKTOP, 'temp', 'p-ckpt')  # 用在 Policy 的API


class BaseCallbacks(DefaultCallbacks):
    """通用回调函数, 基本在各种引擎中都可以添加, 功能为保存模型"""
    engine: str  # 引擎
    custom_conf: dict  # 自定义配置

    alg: str  # 算法
    camp: str  # 阵营, 决定KZ中训练哪个阵营的智能体, 其实没什么用
    checkpoint_path: str  # ckpt文件路径
    load_pretrained: bool  # 是否加载预训练模型
    policy_ids: List[PolicyID]
    pretrained_path: str  # 预训练模型路径
    rule_engine: bool  # 规则引擎触发按钮
    saved_path: str  # 模型存放路径
    stop_iters: int  # 训练迭代次数
    store_model_interval: int  # 模型保存间隔
    store_model_num: int  # 模型保存个数
    hikmodel: HikModel

    policy_select: str  # 选择哪个策略id
    policy: "Policy" = None  # 当前所用的策略
    model = None  # 当前所用的模型
    workers: WorkerSet = None
    local_worker: RolloutWorker = None  # worker

    labels = ["entropy", "total_loss", "policy_reward_mean"]
    labels_name = {'entropy': '熵', 'total_loss': '损失', 'policy_reward_mean': '奖励'}
    labels_threshold = {'entropy': 0.1, 'total_loss': 0.1, 'policy_reward_mean': 0.1}
    last_index_info: dict = {}  # 上一步的指标信息

    def __init__(self, legacy_callbacks_dict: Dict[str, callable] = None):
        super().__init__(legacy_callbacks_dict)
        self.episode_count: int = 0  # 当前回合数(不等于训练次数)
        self.model_index: int = 0  # 保存模型时的后缀序号
        # 灵弈智能体改造
        # self.policy_name_list: list = []
        # self.policy_info_list: list = []
        self.policy_dict = {}
        self.policy_name_dict = {}  # policyname to policy
        self.policy_model_index = {}  # policyname to model index
        self.policy_reward_mean = {}  # policyname to reward mean: 用于评价策略的模型任务完成优劣，默认按此保存最优模型
        self.use_loss_select_best_model = False  # 评价策略的模型损失函数，用户也可按此保存最优模型
        self.loss_history = {}  # 新增：存储每个策略的Loss历史
        self.policy_total_loss = {}  # policyname to total loss  用于评价策略的模型损失函数，用户也可按此保存最优模型，
                                     # 此时需用户修改函数self.save(..

    def on_checkpoint_loaded(
            self,
            *,
            algorithm: "Algorithm",
            **kwargs,
    ) -> None:
        """Callback run when an Algorithm has loaded a new state from a checkpoint.

        This method gets called at the end of `Algorithm.load_checkpoint()`.

        Args:
            algorithm: Reference to the Algorithm instance.
            kwargs: Forward compatibility placeholder.
        """
        _logger.info(f'调用 on_checkpoint_loaded, 逻辑待实现')
        pass

    def on_algorithm_init(
            self,
            *,
            algorithm: "Algorithm",
            **kwargs,
    ) -> None:
        """该方法只会在实例化trainer时执行一次.

        Callback run when a new Algorithm instance has finished setup.

        This method gets called at the end of Algorithm.setup() after all
        the initialization is done, and before actually training starts.

        Args:
            algorithm: Reference to the Algorithm instance.
            kwargs: Forward compatibility placeholder.
        """
        custom_conf: dict = algorithm.config['env_config']  # 自定义的环境配置
        # 使用的参数-------------------------
        self.engine: str = custom_conf['engine']  # 引擎
        self.custom_conf: dict = custom_conf  # 自定义配置

        self.alg: str = custom_conf["alg"]  # 算法
        self.camp: str = custom_conf["camp"]  # 阵营, 决定KZ中训练哪个阵营的智能体, 其实没什么用
        self.load_pretrained: bool = custom_conf["load_pretrained"]  # 是否加载预训练模型
        self.pretrained_path: str = custom_conf['pretrained_path']  # 预训练模型路径
        self.rule_engine: bool = custom_conf["rule_engine"]  # 规则引擎触发按钮
        self.saved_path: str = custom_conf["model_path"]  # 模型存放路径
        self.stop_iters: int = custom_conf['stop_iters']  # 训练迭代次数
        self.store_model_interval: int = custom_conf['store_model_interval']  # 模型保存间隔
        self.store_model_num: int = custom_conf.get('store_model_num', 3)  # 模型保存

        self.workers: WorkerSet = algorithm.workers
        # 下行代码返回值的类型是类型提示的父类
        self.local_worker: RolloutWorker = algorithm.workers.local_worker()  # type: ignore
        self.policy_ids: List[PolicyID] = list(self.local_worker.policy_dict.keys())  # list[str]
        # 使用的参数-------------------------

        _save_dir = os.path.join(WORK_DIR, 'results', 'hikmodel')
        self.hikmodel = HikModel(_save_dir)
        # 在hiklayout/conf.json中写入 engine, 平台组画图需要知道引擎
        conf_path = os.path.join(WORK_DIR, 'hiklayout', 'conf.json')
        with open(conf_path, 'r', encoding='utf-8') as fp:
            content: dict = json.load(fp)
            content['data']['engine_name'] = self.engine
        with open(conf_path, 'w', encoding='utf-8') as fp:
            json.dump(content, fp, indent=4, ensure_ascii=False)

        # 路径读取 --------------------------------------------
        self.checkpoint_path = self.pretrained_path
        if not os.path.exists(self.checkpoint_path):
            os.makedirs(self.checkpoint_path, exist_ok=True)
        if not os.path.exists(self.saved_path):
            os.makedirs(self.saved_path, exist_ok=True)
        # 路径读取 --------------------------------------------

        policy_select, policy = self.choose_policy_id(algorithm)
        self.policy_select: str = policy_select  # 举例 'learn', 'policy_1'
        self.policy = policy
        self.model = getattr(policy, 'model', None)

        # 灵弈智能体模型保存修改
        for policy_name in self.policy_ids:
            policy = algorithm.get_policy(policy_name)
            self.policy_dict[policy_name] = policy
            self.policy_model_index[policy_name] = 1

        # 2025-0620:
        if custom_conf["load_pretrained"] and 'policy_module_path' in custom_conf:
            for policy_name, policy in self.policy_dict.items():
                if policy_name in custom_conf["policy_module_path"]:
                    model_path = custom_conf["policy_module_path"][policy_name]
                    if os.path.exists(model_path):
                        model = policy.model
                        with policy._sess.graph.as_default():
                            with policy._sess.as_default():
                                weight_file = os.path.join(model_path, "weights.h5")
                                if os.path.exists(weight_file):
                                    model.base_model.load_weights(weight_file)
                                    _logger.info("加载预训练文件:%s" % weight_file)

    def load_custom_model(self, path):
        normalized_path = os.path.normpath(path)

        if not os.path.exists(normalized_path):
            raise FileNotFoundError(f"模型文件 {normalized_path} 不存在")

        if not normalized_path.lower().endswith(('.h5', '.hdf5')):
            raise ValueError("仅支持.h5或.hdf5格式文件")

        return load_model(normalized_path, compile=False)

    def _update_loss_history(self, result: dict):
        """从训练结果中提取并存储每个策略的Loss值"""
        # 提取Learner信息
        learner_info = result.get("info", {}).get("learner", {})

        # 单策略场景
        if "learner_stats" in learner_info:
            stats = learner_info["learner_stats"]
            if "total_loss" in stats:
                policy_id = "default"
                self.loss_history.setdefault(policy_id, []).append(stats["total_loss"])

        # 多策略场景
        elif isinstance(learner_info, dict):
            for policy_id, policy_info in learner_info.items():
                stats = policy_info.get("learner_stats", {})
                if "total_loss" in stats:
                    self.loss_history.setdefault(policy_id, []).append(stats["total_loss"])

    def _get_policy_loss(self, policy_id: str) -> float:
        """获取指定策略的最新Loss值"""
        if not self.loss_history.get(policy_id):
            return float('nan')  # 没有记录时返回NaN
        return self.loss_history[policy_id][-1] if self.loss_history[policy_id] else float('nan')

    def on_train_result(
            self,
            *,
            algorithm: "Algorithm",
            result: dict,
            **kwargs,
    ) -> None:
        """Called at the end of Algorithm.train().

        Args:
            algorithm: Current Algorithm instance.
            result: Dict of results returned from Algorithm.train() call.
                You can mutate this object to add additional metrics.
            kwargs: Forward compatibility placeholder.
        """
        super().on_train_result(algorithm=algorithm, result=result, **kwargs)
        self.check_init_done(algorithm=algorithm)

        if self.use_loss_select_best_model:
            self._update_loss_history(result)  # 更新Loss历史

        iteration: int = algorithm.iteration  # 从 1 开始
        # _logger.debug(f'第 {iteration} 次训练回调开始')

        self.save(algorithm, result)
        self.hikmodel_save()
        policy_index: dict = self.get_index_info(result=result)
        self.hikmodel_plot(policy_index, iter_count=iteration)
        self.hikmodel_rule(policy_index)

        if iteration == self.stop_iters:
            _logger.info("已到最大训练迭代次数:%d, 训练结束！" % iteration)
            self.hikmodel.write_rule_info({"rule_engine_info": []})  # 训练完成后清空规则引擎信息

            finally_dict = {}  # 平台组要求的最终字典
            for policy_id, policy_info in policy_index.items():
                if policy_id != self.policy_select:
                    continue  # 单策略曲线模式, 只画一个策略的指标
                # 其他指标信息是所有策略都有的
                for label, value in policy_info.items():
                    finally_dict[f'{policy_id}-{label}'] = value
            self.hikmodel.write_train_metrics(finally_dict)

    def choose_policy_id(self, algorithm: "Algorithm") -> Tuple[str, "Policy"]:
        """选择默认的策略id和策略"""
        if self.alg in ['ARS', 'ES']:
            policy_select: str = self.policy_ids[0]
            policy = algorithm.get_policy()
        else:
            policy_select: str = self.policy_ids[0]
            policy = algorithm.get_policy(policy_select)
        return policy_select, policy

    def check_init_done(self, *, algorithm: "Algorithm") -> None:
        """检查是否执行了 on_trainer_init 函数"""
        if (
                self.policy is None  # None 为默认, 代表之前没执行过 on_algorithm_init 函数
                and algorithm.config['env_config'].get('alg') in ['ARS', 'ES']  # 这2种训练算法似乎不会自动执行 on_trainer_init 函数
        ):
            self.on_algorithm_init(algorithm=algorithm)  # 手动调用初始化函数

    def get_index_info(self, *, result: dict) -> dict:
        """获得训练指标信息字典"""
        # 获得指标数据字典 policy_index
        self.policy_ids: List[PolicyID] = list(self.local_worker.policy_dict.keys())  # 其他回调器可能会增加策略

        policy_index: dict = {policy_id: {} for policy_id in self.policy_ids}  # 初始化
        # dict[policy_id, policy_info], 原始指标信息字典
        raw_index_info: dict = result['info'].get(LEARNER_INFO, {})  # 'ARS', 'ES' 算法没有这个东西
        # 有时某个policy没有训练结果指标(猜测是因为对应的智能体没有收集到足够的样本)
        for policy_id, policy_info in raw_index_info.items():
            if policy_id not in policy_index:
                continue  # 竟然还有 'default_policy' 策略
            for key, value in policy_info[LEARNER_STATS_KEY].items():
                if key not in self.labels:
                    continue
                policy_index[policy_id][key] = float(value)

        policy_reward_mean_info: dict = result.get('policy_reward_mean', {})  # 第1次训练时为空, 'ARS', 'ES' 算法没有这个东西
        for policy_id, reward_value in policy_reward_mean_info.items():
            if policy_id not in policy_index:
                policy_index[policy_id] = {}
            policy_index[policy_id]['policy_reward_mean'] = float(reward_value)
        return policy_index

    def save(self, algorithm: "Algorithm", result: dict) -> None:
        """保存训练过程产生的ckpt, model, weights文件
        :param algorithm: 算法对象
        :param result: dict, 训练结果
        :return:
        """
        iteration: int = algorithm.iteration  # 从 1 开始
        if iteration % self.store_model_interval != 0:
            return
        algorithm.workers.sync_weights()
        self.model_index = self.model_index + 1

        checkpoint_saved_path = os.path.join(self.checkpoint_path, "checkpoint", "checkpoint_" + str(self.model_index))
        if not os.path.exists(checkpoint_saved_path):
            os.makedirs(checkpoint_saved_path)
        algorithm.save_checkpoint(checkpoint_saved_path)

        # 计算策略平均奖励，如果没有存，则存储；否则判断平均奖励更大，则更新存储；
        reward_result = {}
        # 计算策略损失函数总损失，如果没有存，则存储；否则判断损失值越小，则更新存储；
        policy_losses = {}
        if self.use_loss_select_best_model:  # 获取当前所有策略的Loss值
            for policy_name in self.policy_dict:
                policy_losses[policy_name] = self._get_policy_loss(policy_name)
        else:
            if 'policy_reward_mean' in result:
                reward_result = result["policy_reward_mean"]

        # 将所有策略保存到指定文件夹里
        policy_model_path = {}
        if "policy_module_path" in algorithm.config.env_config:
            policy_model_path = algorithm.config.env_config["policy_module_path"]
        for policy_name, policy in self.policy_dict.items():
            if policy_name not in policy_model_path:
                continue
            module_path = policy_model_path[policy_name]  # 策略的模型路径，和智能体描述相关

            if not os.path.exists(module_path):
                os.makedirs(module_path)  # 策略的模型路径不存在，则创建路径

            # 存储的模型序号
            model_index = 1
            if policy_name in self.policy_model_index:
                model_index = self.policy_model_index[policy_name]
            else:
                self.policy_model_index[policy_name] = 1

            # 获取平均奖励，查看是否最大训练轮数里最大的
            is_better_model = False  # 是当前更好的模型
            if self.use_loss_select_best_model:
                if policy_name in policy_losses:
                    crt_loss = policy_losses[policy_name]
                    if policy_name not in self.policy_total_loss:
                        self.policy_total_loss[policy_name] = crt_loss
                        is_better_model = True
                    elif crt_loss != float('nan') and crt_loss < self.policy_total_loss[policy_name]:
                        self.policy_total_loss[policy_name] = crt_loss
                        is_better_model = True
                else:
                    _logger.info("策略:%s 第%d轮训练total_loss是Nan！" % (policy_name, model_index))
            else:
                if policy_name in reward_result:
                    crt_reward = reward_result[policy_name]
                    if policy_name not in self.policy_reward_mean:
                        self.policy_reward_mean[policy_name] = crt_reward
                        is_better_model = True
                    elif crt_reward and crt_reward > self.policy_reward_mean[policy_name]:
                        self.policy_reward_mean[policy_name] = crt_reward
                        is_better_model = True
                else:
                    _logger.info("策略:%s 第%d轮训练平均奖励是Nan！" % (policy_name, model_index))

            model = policy.model
            if isinstance(policy, ray.rllib.algorithms.sac.sac_tf_policy.SACTFModel):
                with policy._sess.graph.as_default():
                    with policy._sess.as_default():
                        if is_better_model:
                            model.action_model.base_model.save_weights(
                                os.path.join(module_path, 'weights.h5'))
                            model_file_path = os.path.join(module_path, 'model_%d_better.h5' % model_index)
                            model.action_model.base_model.save(model_file_path)

                            if self.use_loss_select_best_model:
                                _logger.info("本次训练保存截止目前最小损失值的神经网络模型文件:%s, 并更文件'model_best.h5'" % model_file_path)
                            else:
                                _logger.info("本次训练保存截止目前最大平均奖励的神经网络模型文件:%s, 并更文件'model_best.h5'" % model_file_path)
                            model_file_path = os.path.join(module_path, 'model_best.h5')
                            model.action_model.base_model.save(model_file_path)
                        else:
                            model_file_path = os.path.join(module_path, 'model_%d.h5' % model_index)
                            model.action_model.base_model.save(model_file_path)
                            _logger.info("本次训练保存神经网络模型文件:%s" % model_file_path)
            else:
                with policy._sess.graph.as_default():
                    with policy._sess.as_default():
                        if is_better_model:
                            model.base_model.save_weights(os.path.join(module_path, 'weights.h5'))
                            model_file_path = os.path.join(module_path, 'model_%d_better.h5' % model_index)
                            model.base_model.save(model_file_path, include_optimizer=True,  # 明确包含优化器
                                                  save_format='h5')  # 对 HDF5 格式需显式开启

                            if self.use_loss_select_best_model:
                                _logger.info("本次训练保存截止目前最小损失值的神经网络模型文件:%s, 并更文件'model_best.h5'" % model_file_path)
                            else:
                                _logger.info("本次训练保存截止目前最大平均奖励的神经网络模型文件:%s, 并更文件'model_best.h5'" % model_file_path)
                            model_file_path = os.path.join(module_path, 'model_best.h5')
                            model.base_model.save(model_file_path, include_optimizer=True,  # 明确包含优化器
                                                  save_format='h5')  # 对 HDF5 格式需显式开启
                        else:
                            model_file_path = os.path.join(module_path, 'model_%d.h5' % model_index)
                            model.base_model.save(model_file_path, include_optimizer=True,  # 明确包含优化器
                                                  save_format='h5')  # 对 HDF5 格式需显式开启
                            _logger.info("本次训练保存神经网络模型文件:%s" % model_file_path)
            model_index += 1
            self.policy_model_index[policy_name] = model_index

    def load(self, checkpoint_dir: str, algorithm: "Algorithm") -> None:
        checkpoint_info = get_checkpoint_info(checkpoint_dir)
        if checkpoint_info['type'] == 'Algorithm':
            # algorithm.restore() -> algorithm.load_checkpoint() -> self.on_checkpoint_loaded() -> None
            algorithm.restore(checkpoint_dir)

            _policy_dict: Dict[str, PolicySpec] = self.local_worker.policy_dict
            _ = None
        elif checkpoint_info['type'] == 'Policy':
            self.policy.from_checkpoint(checkpoint_dir)

    def hikmodel_save(self) -> None:
        """使用HikModel模块管理保存的ckpt, model, weights文件"""
        self.hikmodel.save_model(self.saved_path)
        # self.hikmodel.save_model(self.checkpoint_path)
        # _logger.info(f"保存模型：{self.hikmodel.save_dir}")

    def hikmodel_plot(self, policy_index: dict, iter_count: int, one: bool = True) -> None:
        """使用HikModel模块传递指标数据, 在网页端画图.

        :param policy_index: dict, 所有策略的指标信息字典, 包括一个单
        :param iter_count: int, 训练次数
        :param one: bool, 是否只画单一策略的指标, False时会画出所有策略的指标
        :return:
        """
        for policy_id, policy_info in policy_index.items():
            if one and policy_id != self.policy_select:
                continue  # 单策略曲线模式, 只画一个策略的指标

            # 其他指标信息是所有策略都有的
            for label, value in policy_info.items():
                label_cn = self.labels_name.get(label, '未知')
                self.hikmodel.visual_summary(
                    data={label: value}, step=iter_count,
                    name=f"{policy_id}---{label_cn}曲线"
                )
        pass

    def hikmodel_rule(self, policy_index: dict, one: bool = True) -> None:
        """使用HikModel进行弹窗操作(当各指标数据触发规则引擎条件)"""
        if not self.last_index_info:
            self.last_index_info = policy_index  # 记录指标信息
            return  # 无上一步指标, 无相对变化, 无需考虑

        # 检查各label的指标是否达到规则引擎触发阈值
        rule_status = {label: False for label in self.labels}  # 初始化
        for policy_id, policy_info in policy_index.items():
            if one and policy_id != self.policy_select:
                continue  # 单策略曲线模式, 只画一个策略的指标
            # 其他指标信息是所有策略都有的
            for label, value in policy_info.items():
                last_value: float = self.last_index_info[policy_id].get(label, None)  #
                if last_value is None or last_value == 0.:
                    continue  # policy_reward_mean 第2次训练才有, 第3次训练才能作差
                if abs((value - last_value) / last_value) <= self.labels_threshold[label]:
                    rule_status[label] = True

        if any(rule_status.values()):
            rule_msg: str = '当前训练过程中, '
            for label, value in rule_status.items():
                if not value:
                    continue  # 该指标没有达到规则引擎触发阈值, pass
                ps = f'变化率小于'
                label_cn: str = self.labels_name[label]
                label_threshold: float = self.labels_threshold[label]
                rule_msg += f'{label_cn}{ps}{label_threshold}'
            rule_msg += "规则引擎触发, 用户可中止训练, 重新加载想定、设计奖励、更换对手、修改超参数后, 继续训练！"
            self.hikmodel.write_rule_info({"rule_engine_info": rule_msg})

        self.last_index_info = policy_index  # 更新指标信息


if __name__ == '__main__':
    np.set_printoptions(suppress=True, threshold=sys.maxsize)

    pass
