#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2024/9/3
file:       framework_and_exploration.py
version:
doc:        检查/调整/更新/适配 不同算法的相关设置.
"""
from corekit.utils import get_debug_logger
from corekit.utils.annotations import Deprecated

_logger = get_debug_logger()


@Deprecated(
    old='corekit.config.framework_and_exploration.valid_framework_and_exploration_by_alg',
    new='corekit.config.train.TrainTaskConfig._valid_framework_and_exploration_by_alg',
    help='旧API只能适用于ray.tune.run训练任务, 新API可对接其他任务',
    error=True
)
def valid_framework_and_exploration_by_alg(args: dict) -> dict:
    """检查/调整/更新/适配 20种学习算法的设置

    不同的算法对参数的设置要求不一, 主要设置可分为5个方面:
    0) 参数要求; 1) 框架要求; 2) 动作/状态空间要求; 3) 探索策略要求; 4) 模型要求;

    本函数仅调整 1, 3, 4, 5方面, 2)的调整仅给出空间类型,
    """
    # 1) 框架要求的知识, 只能使用torch框架, 无法在tf框架下运行
    framework_demand: set = {
        # ray限制必须使用torch框架
        'DDPPO', 'DREAMER', 'MBMPO', 'QMIX', 'RNNSAC', 'R2D2', 'BanditLinTS', 'BanditLinUCB',
        # ray无框架限制, 但tf框架下没跑通, 只能在torch下运行
        'APEX_DDPG', 'DDPG', 'TD3', 'APPO', 'MAML', 'SAC'  #
    }
    # 2) 动作/状态空间要求知识
    act_space_demand: dict = {
        'MultiDiscrete': {
            'A2C', 'A3C', 'BC', 'MARWIL', 'PG', 'PPO', 'APPO', 'IMPALA',
            'ARS', 'ES',  # 'DREAMER'
            'BanditLinTS', 'BanditLinUCB',  # 其他动作空间会轻松爆内存
            # 'QMIX'
            # 'CQL'
            # 'MBMPO',
            'MAML',
        },  # 是否支持多维离散空间的算法
        'Discrete': {
            'A2C', 'A3C', 'BC', 'MARWIL', 'PG', 'PPO', 'APPO', 'IMPALA', 'APEX',
            'DQN', 'SimpleQ', 'ARS', 'ES', 'SAC', 'DDPPO', 'R2D2',
            # 'BanditLinTS', 'BanditLinUCB',
            'MBMPO',
            'RNNSAC',
        },  # 是否支持 Discrete 空间
        'Box': {
            'A2C', 'A3C', 'BC', 'MARWIL', 'PG', 'PPO', 'APPO', 'IMPALA',
            'BanditLinTS', 'BanditLinUCB', 'APEX_DDPG', 'DDPG', 'TD3',  # 'DREAMER'
        },  # 是否支持 Box 空间
    }
    # 3) 探索策略要求知识
    # from ray.rllib.agents.mbmpo.utils import MBMPOExploration
    exploration_demand: dict = {
        'default': 'StochasticSampling',
        'APEX_DDPG': 'StochasticSampling',
        'DDPG': 'StochasticSampling',
        'TD3': 'GaussianNoise',
        'BanditLinTS': 'UpperConfidenceBound',
        'BanditLinUCB': 'ThompsonSampling',
        # 'MBMPO': MBMPOExploration
    }  # 可选的探索策略 StochasticSampling, EpsilonGreedy, ThompsonSampling

    if args['engine'] == 'missile':
        if args['alg'] != 'MADDPG':
            args['alg'] = 'MADDPG'
            _logger.warning(f'missile 只进行 MADDPG 算法训练, 已自动修正')

    alg: str = args['alg']

    # 3) 探索策略调整
    exploration_config = {"type": exploration_demand.get(alg, 'StochasticSampling')}
    args['exploration_config'] = exploration_config

    # 0) 参数要求调整
    if alg == 'PPO':
        # 调整与否都不影响能否跑通
        args['num_sgd_iter'] = 60
        args['lambda'] = 0.9
    elif alg == 'A3C':
        args['lambda'] = 0.9
    elif alg == 'DREAMER':
        args['num_workers'] = 0
    elif alg == 'DDPPO':
        args['rollout_fragment_length'] = args.pop('train_batch_size')
        args['num_gpus'] = 0
        args['policies_to_train'] = None
    elif alg in ['BanditLinTS', 'BanditLinUCB', ]:
        args['timesteps_per_iteration'] = 100
        args['rollout_fragment_length'] = 1
        args['train_batch_size'] = 1
    elif alg == 'CQL':
        args['num_gpus'] = max(args['num_gpus'], 1)
        # from ray.rllib.models.tf.tf_action_dist import TFActionDistribution
        # args['model'] = {'custom_action_dist': TFActionDistribution}
        args['normalize_actions'] = True
        args['_use_beta_distribution'] = False
        # args['input'] = 'd4rl'
    elif alg == 'MBMPO':
        args['num_gpus'] = max(args['num_gpus'], 1)
    elif alg == 'MADDPG':
        if args['framework'] not in ['tf', 'tf2']:
            args['framework'] = 'tf2'
            _logger.warning(f'当前算法 {alg} 只支持 tf 框架, 已自动将框架从torch调整为tf')
        args['agent_id'] = 0  # 必须是整数
        args['learning_starts'] = 800
        args['use_local_critic'] = True
        # args['use_state_preprocessor'] = True
    elif alg == 'ARS':
        # args['rollouts_used'] = 300
        # args['num_rollouts'] = 300
        args['num_workers'] = 2
        args['sgd_stepsize'] = 0.99
    elif alg == 'ES':
        # args['train_batch_size'] = 300
        args['episodes_per_batch'] = 1  # 每个批次跑多少局
    elif alg == 'BC':
        args['framework'] = 'tf2'
        args['enable_connectors'] = True  # 必须使用 connectors
        args['_disable_execution_plan_api'] = False
        args['_enable_new_api_stack'] = False
        args['exploration_config'] = {}
    elif alg == 'MARWIL':
        args['num_gpus'] = min(args['num_gpus'], 1)
        args['learning_starts'] = 100
        args['_disable_execution_plan_api'] = False
    elif alg == 'TD3':  # 失败的, 老问题 else batch.policy_batches[pid] KeyError: 'policy_3';新问题训练10次就会爆内存(16g)
        # 现有方案是不采用TD3Trainer, 直接采用DDPGTrainer
        args.update({
            "learning_starts": 1000,
            "actor_hiddens": [64, 64],
            "critic_hiddens": [32, 32],
        })
    # elif alg == 'DREAMER':  # 失败的
    #     # from ray.rllib.agents.dreamer.dreamer_model import DreamerModel
    #     from corekit.custom_ann import DreamerTorch as AnnModel
    #     # act_space = gym.spaces.Box(low=0, high=1, shape=(5,))  # 动作空间维度
    #     dreamer_model = {
    #         "custom_model": AnnModel,
    #         # RSSM/PlaNET parameters
    #         "deter_size": 1908,
    #         "stoch_size": 1,
    #         # CNN Decoder Encoder
    #         "depth_size": 1,
    #         # General Network Parameters
    #         "hidden_size": 64,
    #         # Action STD
    #         "action_init_std": 5.0,
    #     }
    #     args['dreamer_model'] = dreamer_model

    # 1) 框架调整
    if alg in framework_demand and args['framework'] != 'torch':
        args['framework'] = 'torch'
        _logger.warning(f'当前算法 {alg} 只支持 torch 框架, 已自动将框架从tf调整为torch')

    # 2) 动作 / 状态空间要求
    act_space: str = 'MultiDiscrete'
    obs_space: str = 'Box'
    if alg in act_space_demand['MultiDiscrete']:
        act_space = 'MultiDiscrete'
    elif alg in act_space_demand['Discrete']:
        act_space = 'Discrete'
    elif alg in act_space_demand['Box']:
        act_space = 'Box'
    elif alg == 'QMIX':
        act_space = 'Tuple'
        obs_space = 'Tuple'
    args['obs_space'] = obs_space
    args['act_space'] = act_space

    # 4) 模型调整
    framework: str = args['framework']
    if framework in ['tf']:
        if alg in ['SAC', 'CQL']:
            from corekit.dnn.custom_ann import SACTFModel as AnnModel
        else:
            from corekit.dnn.custom_ann import SharedWeightsModel1 as AnnModel
    elif framework in ['tf2']:
        if alg in ['SAC', 'CQL']:
            from corekit.dnn.custom_ann import SACTFModel as AnnModel
        else:
            from corekit.dnn.custom_ann import TF2SharedWeightsModel as AnnModel
    elif framework == 'torch':
        if alg in ['SAC', 'CQL']:
            from corekit.dnn.custom_ann import SACTorch as AnnModel
        elif alg in ['RNNSAC']:
            from corekit.dnn.custom_ann import RNNSACTorch as AnnModel
        elif alg in ['R2D2']:
            from corekit.dnn.custom_ann import RNNModelTorch as AnnModel
        else:
            from corekit.dnn.custom_ann import SharedWeightsModelTorch as AnnModel
    else:
        raise ValueError(f'不支持的框架{framework}')
    args['ann_model'] = AnnModel

    if args['framework'] == 'tf2':
        args['eager_tracing'] = False  # tf2 模式下必须, 否则会失败
    if args.get('input_'):
        args['explore'] = False  # 使用离线数据, 没必要探索

    return args
