import argparse
import yaml

def get_config(args):
    path = getattr(args, 'config_file', None)
    if path:
        with open(path, 'r', encoding='utf-8') as f:
            cfg = yaml.safe_load(f) or {}

        for k, v in cfg.items():
            setattr(args, k, v)

def get_args():
    """
    The configuration parser for common hyperparameters of all environment.
    Please reach each `scripts/train/<env>_runner.py` file to find private hyperparameters
    only used in <env>.

    Prepare parameters:
        --algorithm_name <algorithm_name>
            specifiy the algorithm, including `["rmappo", "mappo", "rmappg", "mappg", "trpo"]`
        --experiment_name <str>
            an identifier to distinguish different experiment.
        --seed <int>
            set seed for numpy and torch
        --cuda
            by default True, will use GPU to train; or else will use CPU;
        --cuda_deterministic
            by default, make sure random seed effective. if set, bypass such function.
        --n_training_threads <int>
            number of training threads working in parallel. by default 1
        --n_rollout_threads <int>
            number of parallel envs for training rollout. by default 32
        --n_eval_rollout_threads <int>
            number of parallel envs for evaluating rollout. by default 1
        --n_render_rollout_threads <int>
            number of parallel envs for rendering, could only be set as 1 for some environments.
        --num_env_steps <int>
            number of env steps to train (default: 10e6)
        --user_name <str>
            [for wandb usage], to specify user's name for simply collecting training data.
        --use_wandb
            [for wandb usage], by default True, will log date to wandb server. or else will use tensorboard to log data.

    Env parameters:
        --env_name <str>
            specify the name of environment
        --use_obs_instead_of_state
            [only for some env] by default False, will use global state; or else will use concatenated local obs.

    Replay Buffer parameters:
        --episode_length <int>
            the max length of episode in the buffer.

    Network parameters:
        --share_policy
            by default True, all agents will share the same network; set to make training agents use different policies.
        --use_centralized_V
            by default True, use centralized training mode; or else will decentralized training mode.
        --stacked_frames <int>
            Number of input frames which should be stack together.
        --hidden_size <int>
            Dimension of hidden layers for actor/critic networks
        --layer_N <int>
            Number of layers for actor/critic networks
        --use_ReLU
            by default True, will use ReLU. or else will use Tanh.
        --use_popart
            by default True, use PopArt to normalize rewards.
        --use_valuenorm
            by default True, use running mean and std to normalize rewards.
        --use_feature_normalization
            by default True, apply layernorm to normalize inputs.
        --use_orthogonal
            by default True, use Orthogonal initialization for weights and 0 initialization for biases. or else, will use xavier uniform inilialization.
        --gain
            by default 0.01, use the gain # of last action layer
        --use_naive_recurrent_policy
            by default False, use the whole trajectory to calculate hidden states.
        --use_recurrent_policy
            by default, use Recurrent Policy. If set, do not use.
        --recurrent_N <int>
            The number of recurrent layers ( default 1).
        --data_chunk_length <int>
            Time length of chunks used to train a recurrent_policy, default 10.

    Optimizer parameters:
        --lr <float>
            learning rate parameter,  (default: 5e-4, fixed).
        --critic_lr <float>
            learning rate of critic  (default: 5e-4, fixed)
        --opti_eps <float>
            RMSprop optimizer epsilon (default: 1e-5)
        --weight_decay <float>
            coefficience of weight decay (default: 0)

    PPO parameters:
        --ppo_epoch <int>
            number of ppo epochs (default: 15)
        --use_clipped_value_loss
            by default, clip loss value. If set, do not clip loss value.
        --clip_param <float>
            ppo clip parameter (default: 0.2)
        --num_mini_batch <int>
            number of batches for ppo (default: 1)
        --entropy_coef <float>
            entropy term coefficient (default: 0.01)
        --use_max_grad_norm
            by default, use max norm of gradients. If set, do not use.
        --max_grad_norm <float>
            max norm of gradients (default: 0.5)
        --use_gae
            by default, use generalized advantage estimation. If set, do not use gae.
        --gamma <float>
            discount factor for rewards (default: 0.99)
        --gae_lambda <float>
            gae lambda parameter (default: 0.95)
        --use_proper_time_limits
            by default, the return value does consider limits of time. If set, compute returns with considering time limits factor.
        --use_huber_loss
            by default, use huber loss. If set, do not use huber loss.
        --use_value_active_masks
            by default True, whether to mask useless data in value loss.
        --huber_delta <float>
            coefficient of huber loss.

    PPG parameters:
        --aux_epoch <int>
            number of auxiliary epochs. (default: 4)
        --clone_coef <float>
            clone term coefficient (default: 0.01)

    Run parameters：
        --use_linear_lr_decay
            by default, do not apply linear decay to learning rate. If set, use a linear schedule on the learning rate

    Save & Log parameters:
        --save_interval <int>
            time duration between contiunous twice models saving.
        --log_interval <int>
            time duration between contiunous twice log printing.

    Eval parameters:
        --use_eval
            by default, do not start evaluation. If set`, start evaluation alongside with training.
        --eval_interval <int>
            time duration between contiunous twice evaluation progress.
        --eval_episodes <int>
            number of episodes of a single evaluation.

    Render parameters:
        --save_gifs
            by default, do not save render video. If set, save video.
        --use_render
            by default, do not render the env during training. If set, start render. Note: something, the environment has internal render process which is not controlled by this hyperparam.
        --render_episodes <int>
            the number of episodes to render a given env
        --ifi <float>
            the play interval of each rendered image in saved video.

    Pretrained parameters:
        --model_dir <str>
            by default None. set the path to pretrained model.
    """
    parser = argparse.ArgumentParser(
        description="onpolicy", formatter_class=argparse.RawDescriptionHelpFormatter
    )

    # prepare parameters
    parser.add_argument("--algorithm_name", type=str, default="mappo")
    parser.add_argument(
        "--trainer_algorithm_name",
        type=str,
        default=None,
        choices=["rmappo", "mappo", "qmix", "vdn", "hatt_qmix"],
        help="internal trainer implementation; algorithm_name is used as experiment label",
    )

    parser.add_argument(
        "--experiment_name",
        type=str,
        default="check",
        help="an identifier to distinguish different experiment.",
    )
    parser.add_argument("--seed", type=int, default=1, help="Random seed for numpy/torch")
    parser.add_argument(
        "--cuda",
        type=int,
        default=0,
        help="by default True, will use GPU to train; or else will use CPU;",
    )
    parser.add_argument(
        "--cuda_deterministic",
        action="store_false",
        default=True,
        help="by default, make sure random seed effective. if set, bypass such function.",
    )
    parser.add_argument(
        "--n_training_threads",
        type=int,
        default=2,
        help="Number of torch threads for training",
    )
    parser.add_argument(
        "--n_rollout_threads",
        type=int,
        default=5,
        help="Number of parallel envs for training rollouts",
    )
    parser.add_argument(
        "--n_eval_rollout_threads",
        type=int,
        default=2,
        help="Number of parallel envs for evaluating rollouts",
    )
    parser.add_argument(
        "--n_render_rollout_threads",
        type=int,
        default=1,
        help="Number of parallel envs for rendering rollouts",
    )
    parser.add_argument(
        "--num_env_steps",
        type=int,
        default=10e6,
        help="Number of environment steps to train (default: 10e6)",
    )
    parser.add_argument(
        "--user_name",
        type=str,
        default="marl",
        help="[for wandb usage], to specify user's name for simply collecting training data.",
    )

    # env parameters
    parser.add_argument("--env_name", type=str, default="MyEnv", help="specify the name of environment")
    parser.add_argument(
        "--use_obs_instead_of_state",
        action="store_true",
        default=False,
        help="Whether to use global state or concatenated obs",
    )

    # replay buffer parameters
    parser.add_argument("--episode_length", type=int, default=200, help="Max length for any episode")

    parser.add_argument(
        "--explore_action_mode",
        type=str,
        default="nearest",
        choices=["nearest", "sector", "survey"],
        help="exploration action semantics: nearest target oracle, sector directions, or explicit survey tasks",
    )
    parser.add_argument(
        "--num_explore_actions",
        type=int,
        default=1,
        help="number of directional exploration actions when explore_action_mode=sector",
    )
    parser.add_argument(
        "--explore_step_size",
        type=float,
        default=None,
        help="sector exploration waypoint distance; defaults to 25%% of map_size",
    )
    parser.add_argument(
        "--explore_random_jitter",
        type=float,
        default=0.0,
        help="fraction of sector width used as random angular jitter for sector exploration",
    )
    parser.add_argument(
        "--sector_redecision_on_waypoint",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="reactivate a sector explorer after reaching its current waypoint",
    )
    parser.add_argument(
        "--explore_event_redecision",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="reactivate all living explorers when a new delivery task is discovered",
    )
    parser.add_argument(
        "--sector_min_leg_mobility_steps",
        type=float,
        default=1.0,
        help="mask sector actions whose nominal in-map leg is shorter than this many mobility steps",
    )
    parser.add_argument(
        "--delivery_tasks_initially_discovered",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="make every delivery task visible at reset without changing task/action dimensions",
    )
    parser.add_argument(
        "--coalition_task_feature_mode",
        choices=("compact", "full"),
        default="compact",
        help=(
            "task feature schema for coalition delivery: compact preserves the "
            "historical 6-D observation; full exposes all capability requirements"
        ),
    )
    parser.add_argument(
        "--mask_explore_actions",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="mask all Explore action slots while preserving the checkpoint action dimension",
    )
    parser.add_argument(
        "--idle_time_cost_mode",
        choices=("legacy", "uniform"),
        default="legacy",
        help=(
            "legacy keeps the shaped reward's action-type cost; uniform "
            "charges the same per-step time cost to Idle and non-Idle agents"
        ),
    )
    parser.add_argument(
        "--explore_observer_relative_geometry",
        action=argparse.BooleanOptionalAction,
        default=False,
        help=(
            "describe sector candidates with each observer's executable goal "
            "and leg length, without enabling searched-area coverage memory"
        ),
    )
    parser.add_argument(
        "--explore_coverage_features",
        action=argparse.BooleanOptionalAction,
        default=False,
        help=(
            "add searched-area coverage to sector candidates; this also "
            "enables their observer-relative executable geometry"
        ),
    )
    parser.add_argument(
        "--explore_coverage_grid_resolution",
        type=int,
        default=20,
        help="cells per map axis used by optional searched-area coverage features",
    )
    # network parameters
    parser.add_argument(
        "--share_policy",
        action="store_false",
        default=False,
        help="Whether agent share the same policy",
    )
    parser.add_argument(
        "--use_centralized_V",
        action="store_false",
        default=True,
        help="Whether to use centralized V function",
    )
    parser.add_argument(
        "--stacked_frames",
        type=int,
        default=1,
        help="Dimension of hidden layers for actor/critic networks",
    )
    parser.add_argument(
        "--use_stacked_frames",
        action="store_true",
        default=False,
        help="Whether to use stacked_frames",
    )
    parser.add_argument(
        "--hidden_size",
        type=int,
        default=64,
        help="Dimension of hidden layers for actor/critic networks",
    )
    parser.add_argument(
        "--layer_N",
        type=int,
        default=1,
        help="Number of layers for actor/critic networks",
    )
    parser.add_argument("--use_ReLU", action="store_false", default=True, help="Whether to use ReLU")
    parser.add_argument(
        "--use_popart",
        action="store_true",
        default=False,
        help="by default False, use PopArt to normalize rewards.",
    )
    parser.add_argument(
        "--use_valuenorm",
        action="store_false",
        default=True,
        help="by default True, use running mean and std to normalize rewards.",
    )
    parser.add_argument(
        "--use_feature_normalization",
        action="store_false",
        default=True,
        help="Whether to apply layernorm to the inputs",
    )
    parser.add_argument(
        "--use_orthogonal",
        action="store_false",
        default=True,
        help="Whether to use Orthogonal initialization for weights and 0 initialization for biases",
    )
    parser.add_argument("--gain", type=float, default=0.01, help="The gain # of last action layer")

    # recurrent parameters
    parser.add_argument(
        "--use_naive_recurrent_policy",
        action="store_true",
        default=False,
        help="Whether to use a naive recurrent policy",
    )
    parser.add_argument(
        "--use_recurrent_policy",
        action="store_false",
        default=False,
        help="use a recurrent policy",
    )
    parser.add_argument("--recurrent_N", type=int, default=1, help="The number of recurrent layers.")
    parser.add_argument(
        "--data_chunk_length",
        type=int,
        default=10,
        help="Time length of chunks used to train a recurrent_policy",
    )

    # optimizer parameters
    parser.add_argument("--lr", type=float, default=5e-4, help="learning rate (default: 5e-4)")
    parser.add_argument(
        "--critic_lr",
        type=float,
        default=5e-4,
        help="critic learning rate (default: 5e-4)",
    )
    parser.add_argument(
        "--opti_eps",
        type=float,
        default=1e-5,
        help="RMSprop optimizer epsilon (default: 1e-5)",
    )
    parser.add_argument("--weight_decay", type=float, default=0)

    # ppo parameters
    parser.add_argument("--ppo_epoch", type=int, default=15, help="number of ppo epochs (default: 15)")
    parser.add_argument(
        "--use_clipped_value_loss",
        action="store_false",
        default=True,
        help="by default, clip loss value. If set, do not clip loss value.",
    )
    parser.add_argument(
        "--clip_param",
        type=float,
        default=0.2,
        help="ppo clip parameter (default: 0.2)",
    )
    parser.add_argument(
        "--ppo_target_kl",
        type=float,
        default=0.0,
        help="stop remaining PPO minibatch updates when approx KL exceeds this value; <=0 disables",
    )
    parser.add_argument(
        "--ppo_max_clip_fraction",
        type=float,
        default=0.0,
        help="stop remaining PPO minibatch updates when clipped sample fraction exceeds this value; <=0 disables",
    )
    parser.add_argument(
        "--num_mini_batch",
        type=int,
        default=1,
        help="number of batches for ppo (default: 1)",
    )
    parser.add_argument(
        "--entropy_coef",
        type=float,
        default=0.01,
        help="entropy term coefficient (default: 0.01)",
    )
    parser.add_argument(
        "--value_loss_coef",
        type=float,
        default=1,
        help="value loss coefficient (default: 0.5)",
    )
    parser.add_argument(
        "--use_max_grad_norm",
        action="store_false",
        default=True,
        help="by default, use max norm of gradients. If set, do not use.",
    )
    parser.add_argument(
        "--max_grad_norm",
        type=float,
        default=10.0,
        help="max norm of gradients (default: 0.5)",
    )
    parser.add_argument(
        "--use_gae",
        action="store_false",
        default=True,
        help="use generalized advantage estimation",
    )
    parser.add_argument(
        "--gamma",
        type=float,
        default=0.99,
        help="discount factor for rewards (default: 0.99)",
    )
    parser.add_argument(
        "--gae_lambda",
        type=float,
        default=0.95,
        help="gae lambda parameter (default: 0.95)",
    )
    parser.add_argument(
        "--use_proper_time_limits",
        action="store_true",
        default=False,
        help="compute returns taking into account time limits",
    )
    parser.add_argument(
        "--use_huber_loss",
        action="store_false",
        default=True,
        help="by default, use huber loss. If set, do not use huber loss.",
    )
    parser.add_argument(
        "--use_value_active_masks",
        action="store_false",
        default=True,
        help="by default True, whether to mask useless data in value loss.",
    )
    parser.add_argument(
        "--use_policy_active_masks",
        action="store_false",
        default=True,
        help="by default True, whether to mask useless data in policy loss.",
    )
    parser.add_argument("--huber_delta", type=float, default=10.0, help=" coefficience of huber loss.")

    # Recurrent value decomposition (QMIX / VDN / HATT-QMIX)
    parser.add_argument("--qmix_agent_hidden_dim", type=int, default=128)
    parser.add_argument("--qmix_mixer_hidden_dim", type=int, default=64)
    parser.add_argument("--qmix_hypernet_hidden_dim", type=int, default=128)
    parser.add_argument("--qmix_buffer_size", type=int, default=128, help="episode replay capacity")
    parser.add_argument("--qmix_batch_size", type=int, default=8)
    parser.add_argument("--qmix_sequence_length", type=int, default=32)
    parser.add_argument("--qmix_burn_in", type=int, default=8)
    parser.add_argument(
        "--qmix_use_recurrent_agent",
        type=lambda value: str(value).lower() in ("1", "true", "yes", "y"),
        default=True,
        help="use a GRU in the per-agent utility network",
    )
    parser.add_argument(
        "--qmix_transition_mode",
        choices=["physical_legacy", "physical_constrained", "event_smdp"],
        default="physical_legacy",
        help="replay/target semantics for event-driven coalition actions",
    )
    parser.add_argument(
        "--qmix_hatt_sequence_chunk_size",
        type=int,
        default=8,
        help=(
            "number of recurrent time steps whose memoryless HATT backbone and "
            "candidate decoder are evaluated together"
        ),
    )
    parser.add_argument("--qmix_learning_starts", type=int, default=16, help="episodes before updates")
    parser.add_argument(
        "--qmix_train_interval",
        type=int,
        default=1,
        help="completed episodes between update groups",
    )
    parser.add_argument("--qmix_updates_per_train", type=int, default=1)
    parser.add_argument(
        "--qmix_target_update_interval",
        type=int,
        default=200,
        help="completed episodes between hard target-network updates",
    )
    parser.add_argument(
        "--qmix_optimizer", choices=["rmsprop", "adam"], default="rmsprop"
    )
    parser.add_argument("--qmix_rmsprop_alpha", type=float, default=0.99)
    parser.add_argument(
        "--qmix_max_abs_q",
        type=float,
        default=10000.0,
        help="fail fast when joint online/target Q exceeds this magnitude; <=0 disables",
    )
    parser.add_argument(
        "--qmix_double_q", type=lambda value: str(value).lower() in ("1", "true", "yes", "y"),
        default=True,
    )
    parser.add_argument("--qmix_epsilon_start", type=float, default=1.0)
    parser.add_argument("--qmix_epsilon_finish", type=float, default=0.05)
    parser.add_argument("--qmix_epsilon_anneal_steps", type=int, default=500000)
    parser.add_argument("--qmix_reward_scale", type=float, default=1.0)
    parser.add_argument(
        "--mappo_team_mean_reward",
        type=lambda value: str(value).lower() in ("1", "true", "yes", "y"),
        default=False,
        help="train MAPPO/HATT-MAPPO with the team-mean reward broadcast to all agents",
    )

    # HATT utility head auxiliary loss
    parser.add_argument(
        "--utility_loss_coef",
        type=float,
        default=0.01,
        help="coefficient for the auxiliary utility-head MSE loss (v3.2d+)",
    )

    # run parameters
    parser.add_argument(
        "--use_linear_lr_decay",
        action="store_true",
        default=False,
        help="use a linear schedule on the learning rate",
    )
    # save parameters
    parser.add_argument(
        "--save_interval",
        type=int,
        default=1,
        help="time duration between contiunous twice models saving.",
    )
    parser.add_argument(
        "--save_episode_checkpoints",
        action="store_true",
        default=False,
        help="also save tagged actor/critic checkpoints at each save interval",
    )

    # log parameters
    parser.add_argument(
        "--log_interval",
        type=int,
        default=5,
        help="time duration between contiunous twice log printing.",
    )

    # eval parameters
    parser.add_argument(
        "--use_eval",
        action="store_true",
        default=False,
        help="by default, do not start evaluation. If set`, start evaluation alongside with training.",
    )
    parser.add_argument(
        "--eval_interval",
        type=int,
        default=25,
        help="time duration between contiunous twice evaluation progress.",
    )
    parser.add_argument(
        "--eval_episodes",
        type=int,
        default=32,
        help="number of episodes of a single evaluation.",
    )
    parser.add_argument(
        "--max_train_episodes",
        type=int,
        default=None,
        help="optional upper bound on runner training episodes",
    )
    parser.add_argument(
        "--eval_seed",
        type=int,
        default=10000,
        help="first seed in the fixed evaluation scenario suite",
    )
    parser.add_argument(
        "--eval_deterministic",
        action="store_true",
        default=True,
        help="use greedy actions during evaluation; set false in YAML to sample actions",
    )
    parser.add_argument(
        "--training_eval_cuda",
        type=int,
        default=None,
        help="optional CUDA device for periodic training evaluation subprocesses",
    )
    parser.add_argument(
        "--eval_progress_file",
        type=str,
        default=None,
        help="internal atomic per-episode evaluation progress file",
    )

    # render parameters
    parser.add_argument(
        "--save_gifs",
        action="store_true",
        default=False,
        help="by default, do not save render video. If set, save video.",
    )
    parser.add_argument(
        "--use_render",
        action="store_true",
        default=False,
        help="by default, do not render the env during training. If set, start render. Note: something, the environment has internal render process which is not controlled by this hyperparam.",
    )
    parser.add_argument(
        "--render_episodes",
        type=int,
        default=5,
        help="the number of episodes to render a given env",
    )
    parser.add_argument(
        "--ifi",
        type=float,
        default=0.1,
        help="the play interval of each rendered image in saved video.",
    )

    # PCFA baseline parameters
    parser.add_argument(
        "--use_pcfa",
        action="store_true",
        default=False,
        help="use the Project Manager-oriented Coalition Formation Algorithm baseline instead of learning",
    )
    parser.add_argument(
        "--pcfa_comm_range",
        type=float,
        default=-1.0,
        help="communication range for PCFA consensus; -1 means global communication",
    )
    parser.add_argument(
        "--pcfa_lambda",
        type=float,
        default=0.01,
        help="distance decay coefficient in PCFA fitness",
    )
    parser.add_argument(
        "--pcfa_consensus_loops",
        type=int,
        default=-1,
        help="number of PCFA consensus loops; -1 means num_agents",
    )
    parser.add_argument(
        "--pcfa_max_coalition_size",
        type=int,
        default=3,
        help="maximum coalition size selected by PCFA",
    )
    parser.add_argument(
        "--pcfa_min_coalition_size",
        type=int,
        default=1,
        help="minimum coalition size selected by PCFA",
    )
    parser.add_argument(
        "--pcfa_match_metric",
        type=str,
        default="cosine",
        choices=["cosine", "dot"],
        help="metric used to match agent capabilities with task requirements",
    )
    parser.add_argument(
        "--pcfa_resupply_battery_threshold",
        type=float,
        default=0.2,
        help="DelivCoal PCFA battery ratio below which an active agent resupplies",
    )
    parser.add_argument(
        "--pcfa_resupply_load_threshold",
        type=float,
        default=0.05,
        help="DelivCoal PCFA payload ratio below which an active agent resupplies",
    )
    parser.add_argument(
        "--pcfa_priority_weight",
        type=float,
        default=1.0,
        help="DelivCoal PCFA task-priority weight in project-manager ordering",
    )
    parser.add_argument(
        "--pcfa_urgency_weight",
        type=float,
        default=1.0,
        help="DelivCoal PCFA deadline-pressure weight in project-manager ordering",
    )
    parser.add_argument(
        "--pcfa_explore_distance_weight",
        type=float,
        default=0.15,
        help="DelivCoal PCFA penalty for long observer-relative Explore legs",
    )

    # pretrained parameters
    parser.add_argument(
        "--model_dir",
        type=str,
        default=None,
        help="by default None. set the path to pretrained model.",
    )

    return parser