#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2023/12/14
file:       settings.py
version:    v1
doc:        核心, 生成训练所需参数.

陈旭辉注:
训练参数分为两部分, 一部分是ray并行训练框架参数: ray_conf, 一部分是对战配置(包括环境参数配置, 智能体数量配置等): env_conf.
或许有方法可以将两部分参数封装起来作为一个参数, 但由于ray的不透明, 因此还是采用将参数分开的方法进行传入.

ray框架参数 ray_conf 是一个字典, 应包括 gpu数量, 训练批次大小, worker数量等配置, 按照默认就好.
对战配置参数 env_conf 较为自由, 其数据结构任意, 参数设置用来配置/控制env行为, 故只需要环境可以解析即可.

按照目前的构想, 前端可以选择配置一些参数, 这些参数中一部分可以归为ray框架参数, 一部分可以归为对战配置参数.
"""
import importlib
import inspect
import os
import sys
from typing import Any, List, Optional, Tuple, Type

import numpy as np
from ray.rllib.algorithms.registry import ALGORITHMS
from ray.train import CheckpointConfig

from corekit.config.engine import EngineConfig
from corekit.config.framework_and_exploration import valid_framework_and_exploration_by_alg
from corekit.config.general import DEFAULT_CONFIG, get_general_command_args
from corekit.env import BaseEnv
from corekit.utils import get_debug_logger
from corekit.utils.annotations import Deprecated

_logger = get_debug_logger()


@Deprecated(
    old='corekit.config.settings.get_all_config',
    new='corekit.config.base.BaseTaskConfig.get_all_config',
    help='旧API只能适用于ray.tune.run训练任务, 新API可对接其他任务',
    error=True
)
def get_all_config(engine: str = 'demo', alg: str = 'PPO', verbose=True) -> dict:
    """
    封装的接口, 获得 ray_conf, env_conf, stop_conf
    :param engine: str, 引擎名
    :param alg: str, 算法名, 形参优先级 > 命令行参数
    :param verbose: bool, 是否打印设置信息
    :return: ray_conf, env_conf, stop_conf, 3个字典
    """
    if engine is None:
        engine = 'demo'
    engine_path: str = f'envs.{engine}'
    try:
        env_module = importlib.import_module(f'{engine_path}')  # 导入引擎相关模块
    except (ImportError, ModuleNotFoundError) as _e:
        _logger.error(f'导入 envs.{engine} 模块失败请检查, 错误信息: {_e}')
        sys.exit(1)

    # 从环境模块
    class_lst: List[Tuple[str, Type[Any]]] = inspect.getmembers(env_module, inspect.isclass)
    engine_config_class: Optional[Type["EngineConfig"]] = None  # 引擎参数类
    env_class: Type["BaseEnv"] = ...  # 环境类
    for class_name, class_type in class_lst:
        if issubclass(class_type, EngineConfig):
            engine_config_class: Type["EngineConfig"] = class_type
        elif issubclass(class_type, BaseEnv):
            env_class: Type["BaseEnv"] = class_type
    assert issubclass(env_class, BaseEnv), \
        f'{engine_path}模块下没找到可训练环境类,训练环境类应继承自 corekit.env.BaseEnv'
    if not engine_config_class:
        _logger.warning(
            f'没有从{engine_path} 模块下找到引擎参数类,引擎参数类应继承自 corekit.config.engine.EngineConfig'
        )
        engine_config_class = EngineConfig

    # 不可设置的参数
    args: dict = {
        'as_test': True,
        'explore': True,
        'env_config': {
            # "info_callback": None, "done_callback": None, "shared_viewer": True
        },  # 下面这个字典似乎可以随意添加数据
        # 可选项, 部分算法采用默认值会使每一次训练的间隔特别长
        'learning_starts': 500,
        'target_network_update_freq': 100,
        'timesteps_per_iteration': 100,
        'min_time_s_per_reporting': 1
    }
    # 获取并更新可设置的参数
    args.update(get_settable_args(engine, alg, engine_config_class))
    args['env'] = env_class
    # from ray.tune import registry
    #
    # registry.register_env(env_class.__name__, env_class)  # 注册环境
    # args['env'] = env_class.__name__

    # 根据框架和算法, 检查/调整/更新参数字典args
    args: dict = valid_framework_and_exploration_by_alg(args)
    # 更新引擎训练所需的参数
    args = engine_config_class.update_engine_config(args)

    if verbose:
        _logger.info(f'仿真引擎engine={args["engine"]}')
        _logger.info(f'深度学习框架framework={args["framework"]}')
        _logger.info(f'训练算法={args["alg"]}')
        # _logger.info(f"探索策略={args['exploration_config']['type']}")
        _logger.info(f'广义优势估计参数lambda={args["lambda"]}')
        _logger.info(f'学习率lr={args["lr"]}')
        if 'train_batch_size' in args:
            _logger.info(f'训练批次大小train_batch_size={args["train_batch_size"]}')
        _logger.info(f'随机梯度下降最小批次大小sgd_minibatch_size={args["sgd_minibatch_size"]}')
        _logger.info(f'折扣因子gamma={args["gamma"]}')
        _logger.info(f'是否加载预训练模型load_pretrained={args["load_pretrained"]}')
        if args["load_pretrained"]:
            _logger.info(f'加载预训练模型, {args["pretrained_path"]}')

        _logger.info(f'并行进程数量num_workers={args["num_workers"]}')
        _logger.info(f'每个进程环境数量num_envs_per_worker={args["num_envs_per_worker"]}')
        _logger.info(f'每个进程GPU数量num_gpus_per_worker={args["num_gpus_per_worker"]}')

        _logger.info(f'结束条件stop_iters={args["stop_iters"]}')
        _logger.info(f'结束条件stop_reward={args["stop_reward"]}')
        if args['engine'] == 'external':
            _logger.info(f'制导率选择为{args["Law"]}')
            _logger.info(f'导弹类型为{args["MSL"]}')
        elif args['engine'] == 'kz':
            _logger.info(f'训练模式train_manner={args["train_manner"]}')
            _logger.info(f'对手策略expert_pool={args["expert_pool"]}')

    return args


def get_settable_args(
    engine: str = 'demo', alg: str = 'PPO', engine_config_class: Type["EngineConfig"] = EngineConfig
) -> dict:
    """核心实现, 获得可设置参数的值.
    :param engine: str, 引擎名, 需保证engine/ 和 envs/ 下的引擎名一致
    :param alg: str, 算法名
    :param engine_config_class: dict, 用来处理引擎独有参数的设置
    :return: args: dict, 最终设置
    """
    # 可设置项更新, 设置优先级 命令行参数 > 前端参数 > 默认参数 ================================
    # 默认参数
    args = DEFAULT_CONFIG  # 通用参数的默认值
    args.update(engine_config_class.get_default_config())  # 引擎独有参数的默认值
    # 前端参数
    frontend_args = engine_config_class.get_frontend_args()
    args.update(frontend_args)
    # 命令行参数
    command_args = get_command_args(engine, alg, engine_config_class)
    for key, value in command_args.items():
        if value is None:
            continue
        args[key] = value
    # 可设置项更新, 设置优先级 命令行参数 > 前端参数 > 默认参数 ================================

    # 特殊参数的补充/调整-----------------------------------------------------------------
    engine: str = args['engine']
    work_dir: str = args['work_dir']
    # 环境变量的更新/设置
    os.environ['work_dir'] = args['work_dir']  # TODO, Linux下只在当前进程及子进程中有效, ray在非local模式下无法使用
    os.environ['ip_address'] = args['ip_address']
    # 决定 ray.init()中的模式
    args['local_mode']: bool = 'local' in args['training_mode']
    # 使用ray.tune.run训练时产生的训练结果保存路径, ray框架默认路径为用户目录下的 ~/ray_results
    args['results_dir'] = os.path.join(work_dir, 'results', engine, 'train_results')
    if not os.path.exists(args['results_dir']):
        os.makedirs(args['results_dir'])
    # ckpt 和 model 的保存路径
    args['pretrained_path'] = os.path.join(work_dir, 'envs', engine, 'agents')
    args['model_path'] = os.path.join(work_dir, 'envs', engine, 'trained_model')

    # if args['jiasubi']:  # zl 特有
    #     args['train_batch_size'] = 2
    #     args['sgd_minibatch_size'] = 2
    #     args['scene_name'] = "抗巡航导弹多向饱和攻击-02027.scen"
    # 特殊参数的补充/调整-----------------------------------------------------------------
    return args


def get_command_args(
    engine: str = 'demo', alg: str = 'PPO', engine_config_class: Optional[Type["EngineConfig"]] = None
) -> dict:
    """获取命令行参数设置
    :param engine: str, 引擎名, 需保证engine/ 和 envs/ 下的引擎名一致
    :param alg: str, 算法名
    :param engine_config_class: 可选的引擎参数配置类, 用来添加引擎独有参数的命令行设置
    :return:
    """
    parser = get_general_command_args(engine, alg)

    # 添加引擎独有参数的命令行设置===========================================================================
    if issubclass(engine_config_class, EngineConfig):
        parser = engine_config_class.add_command_args(parser)

    command_parse = parser.parse_args()
    commandt_args = vars(command_parse)
    return commandt_args


def splited_args(args: dict) -> Tuple[dict, dict, dict]:
    """
    读取解析参数字典, 分割成3个字典 ray_conf, trainer_conf, stop_conf.

    :param args: dict, 所有参数的字典
    :return: ray_conf, env_conf, stop_conf: dict, dict, dict
    """
    alg: str = args['alg']

    if alg == 'TD3':
        _, default_trainer_config = ALGORITHMS['DDPG']()
        trainer, _ = ALGORITHMS['TD3']()
    else:
        trainer, default_trainer_config = ALGORITHMS[alg]()
    # from ray.tune.registry import get_trainable_cls
    # trainer = get_trainable_cls(alg)  # 这个函数内部会有warning, 日志很丑
    default_trainer_config = trainer.get_default_config()  # 获取训练器默认参数配置
    # Trainer所需的参数
    trainer_conf: dict = {key: args[key] for key in args if key in default_trainer_config}
    # 其他参数
    custom_conf: dict = {key: args[key] for key in args if key not in default_trainer_config}
    trainer_conf['env_config'] = custom_conf  # 回调中会使用

    # 所有的参数设置
    args['run_or_experiment'] = trainer
    args['checkpoint_config'] = CheckpointConfig(**args['checkpoint_config'])

    # 停止条件参数配置
    stop_conf = {
        "training_iteration": args['stop_iters'],  # int
        "timesteps_total": args['stop_timesteps'],  # int
        "episode_reward_mean": args['stop_reward'],  # float
    }

    return args, trainer_conf, stop_conf


if __name__ == '__main__':
    np.set_printoptions(suppress=True)
    pass
