#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2024/9/4
file:       general.py
version:
doc:        公共的通用配置.
"""
import argparse
import os
import sys

from corekit.callbacks.base_callbacks import BaseCallbacks
from corekit.utils import WORK_DIR
from corekit.utils.results_collection import BaseResultsCollection

# 工程路径, ip地址
DEFAULAT_WORK_DIR = WORK_DIR
DEFAULT_IP_ADDRESS = '127.0.0.1'
os.environ['work_dir'] = DEFAULAT_WORK_DIR
os.environ['ip_address'] = DEFAULT_IP_ADDRESS

DEFAULT_CONFIG: dict = {
    'work_dir': DEFAULAT_WORK_DIR,
    'ip_address': DEFAULT_IP_ADDRESS,
    # 训练参数
    "alg": "PPO",
    "engine": "demo",
    "connect": False,  # 是否需要使用通信服务
    "ip": "127.0.0.1",  # 仅当 connect 为 True 时生效
    "port": 5555,  # 仅当 connect 为 True 时生效
    "protocol": "zmq",  # 仅当 connect 为 True 时生效
    "num_workers": 1,
    "num_agents": 2,
    "train_batch_size": 10000,
    "training_mode": "local",  # local:本地训练, normal:本地训练(不可调试), cluster:集群训练, 需在配置文件中指定集群head节点ip,port
    "api": "tune",  # 测试分布式新增
    "framework": "tf",
    "lr": 5e-05,
    "no_done_at_end": False,
    "gamma": 0.95,
    "lambda": 0.9,
    "sgd_minibatch_size": 256,
    "store_model_interval": 1,
    "store_model_num": 100,
    "load_pretrained": False,
    "pretrained_path": os.path.join(DEFAULAT_WORK_DIR, 'results', 'demo', 'pretrained_model'),
    "rule_engine": False,
    "camp": "Red",
    "callbacks": BaseCallbacks,
    "results_collection": BaseResultsCollection(),  # 结果汇总实例化
    # 停止条件
    "stop_iters": 100,
    "stop_steps": 1000000,
    "stop_timesteps": 5000000,
    "stop_reward": 10000000,
    # 硬件资源
    "num_envs_per_worker": 1,
    "num_cpus_per_worker": 1,
    "num_gpus_per_worker": 0,
    "num_cpus": 10,
    "num_gpus": 0,
    # ckpt
    "checkpoint_config": {
        "num_to_keep": 3, "checkpoint_frequency": 1, "checkpoint_score_attribute": 'training_iteration'
    },
    # "enable_connectors": False,
    # "output": os.path.join(DEFAULAT_WORK_DIR, 'results', 'demo', 'offline'),
    # "input_": None,
    # "create_env_on_local_worker": True  # 当设置了离线数据源时, 将该参数设为 True, 甚至可以不再设置 env
    # 日志
    "log_path": os.path.join(WORK_DIR, "log")  # 日志路径
}  # 训练默认设置


def get_general_command_args(
    engine: str = 'demo', alg: str = 'PPO', parser: argparse.ArgumentParser = None
) -> argparse.ArgumentParser:
    """获得通用参数的命令行设置
    :param engine: str, 引擎名, 需保证engine/ 和 envs/ 下的引擎名一致
    :param alg: str, 算法名
    :param parser: argparse.ArgumentParser
    :return:
    """
    if parser is None:
        parser = argparse.ArgumentParser()
    # 外协意见, ray 多机多卡模式时, os.getcwd()可能会不可靠, '127.0.0.1'不可用, 应设置一个全局变量来保证工作路径和ip正确
    parser.add_argument("--work_dir", type=str, help='工程路径', default=None)
    parser.add_argument("--ip_address", type=str, help='ip地址', default=None)
    # 训练参数
    parser.add_argument("--alg", type=str, help='强化学习训练算法', default=alg)
    parser.add_argument("--engine", type=str, help='仿真引擎', default=engine)
    parser.add_argument("--connect", type=bool, help='是否需要通信服务', default=None)
    parser.add_argument("--ip", type=str, help=f'ip地址, 仅当 connect 为 True 时生效', default=None)
    parser.add_argument("--port", type=int, help=f'通信端口, 仅当 connect 为 True 时生效', default=None)
    parser.add_argument("--protocol", type=str, help=f'通信协议, 仅当 connect 为 True 时生效',
                        choices=['zmq', 'udp', 'grpc'], default='zmq' if 'win' in sys.platform else None)
    parser.add_argument("--num_workers", type=int, help='训练的worker数量', default=None)
    parser.add_argument("--num_agents", type=int, help='红方智能体个数', default=None)
    parser.add_argument("--train_batch_size", type=int, help='训练批次大小',
                        default=10000 if 'win' in sys.platform else None)
    # local: 本地训练(可调试), normal: 本地训练(不可调试), cluster: 集群训练, 需在配置文件中指定集群head节点ip, port.
    parser.add_argument("--training_mode", type=str, help='ray训练模式设置',
                        choices=['local', 'normal', 'cluster'], default=None)
    parser.add_argument("--api", type=str, help='使用哪个API进行ray分布式训练',
                        choices=['tune', 'built', 'train'], default=None)
    parser.add_argument("--framework", type=str, help='深度学习框架',
                        choices=['tf', 'tf2', 'torch'], default=None)
    parser.add_argument("--lr", type=float, help='学习率', default=None)
    parser.add_argument("--no_done_at_end", help='是否终止', type=bool, default=None)
    parser.add_argument("--gamma", type=float, help='折扣因子', default=None)
    parser.add_argument("--lambda", type=float, help='GAE参数', default=None)
    parser.add_argument("--sgd_minibatch_size", type=int, help='随机梯度最小批次大小', default=None)
    parser.add_argument("--store_model_interval", type=int, help='每训练多少次保存一次模型', default=None)
    parser.add_argument("--store_model_num", type=int, help='一共保存最新多少个模型', default=None)
    parser.add_argument("--load_pretrained", type=bool, help='是否加载预训练模型', default=None)
    parser.add_argument("--pretrained_path", type=str, help='预训练模型加载路径', default=None)
    parser.add_argument("--rule_engine", type=bool, help='是否采用规则引擎', default=None)
    parser.add_argument("--camp", type=str, help='待训练智能体所处的阵营', choices=['red', 'blue'], default=None)
    # 停止条件
    parser.add_argument("--stop_iters", type=int, help='训练终止的迭代次数',
                        default=100 if 'win' in sys.platform else None)
    parser.add_argument("--stop_steps", type=int, help='每一局最大采样步数', default=None)
    parser.add_argument("--stop_timesteps", type=int, help='训练完成最大采样步数', default=None)
    parser.add_argument("--stop_reward", type=float, help='训练完成最大奖励值', default=None)
    # 硬件资源
    parser.add_argument("--num_envs_per_worker", type=int, help='每个进程的环境数', default=None)
    parser.add_argument("--num_cpus_per_worker", type=int, help='每个进程的CPU数量', default=None)
    parser.add_argument("--num_gpus_per_worker", type=int, help='每个进程的GPU数量', default=None)
    parser.add_argument("--num_cpus", type=int, help='硬件设备的CPU数量', default=None)
    parser.add_argument("--num_gpus", type=int, help='硬件设备的GPU数量', default=None)
    # 日志路径
    parser.add_argument("--log_path", type=str, help='日志路径', default=None)
    return parser
