import argparse
import json
EXAMPLE_USAGE = """
Example Usage via RLlib CLI:
    rllib rollout /tmp/ray/checkpoint_dir/checkpoint-0 --run DQN
    --env CartPole-v0 --steps 1000000 --out rollouts.pkl

Example Usage via executable:
    ./rollout.py /tmp/ray/checkpoint_dir/checkpoint-0 --run DQN
    --env CartPole-v0 --steps 1000000 --out rollouts.pkl
"""

def create_parser(parser_creator = None):
    #parser = argparse.ArgumentParser("Ray training with custom IG environment")

    ## parser for rollouts
    parser_creator = parser_creator or argparse.ArgumentParser
    parser = parser_creator(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="Roll out a reinforcement learning agent "
                    "given a checkpoint.",
        epilog=EXAMPLE_USAGE)

    parser.add_argument(
        "--checkpoint", default='' ,type=str, help="Checkpoint from which to roll out.")
    required_named = parser.add_argument_group("required named arguments")
    required_named.add_argument(
        "--run",
        type=str,
        required=True,
        help="The algorithm or model to train. This may refer to the name "
             "of a built-on algorithm (e.g. RLLib's DQN or PPO), or a "
             "user-defined trainable function or class registered in the "
             "tune registry.")
    required_named.add_argument(
        "--env", type=str, help="The gym environment to use.")
    parser.add_argument(
        "--no-render",
        default=False,
        action="store_const",
        const=True,
        help="Suppress rendering of the environment.")
    parser.add_argument(
        "--monitor",
        default=False,
        action="store_true",
        help="Wrap environment in gym Monitor to record video. NOTE: This "
             "option is deprecated: Use `--video-dir [some dir]` instead.")
    parser.add_argument(
        "--video-dir",
        type=str,
        default=None,
        help="Specifies the directory into which videos of all episode "
             "rollouts will be stored.")
    parser.add_argument(
        "--steps",
        default=20000,
        help="Number of timesteps to roll out (overwritten by --episodes).")
    parser.add_argument(
        "--episodes",
        default=0,
        help="Number of complete episodes to roll out (overrides --steps).")
    parser.add_argument("--out", default=None, help="Output filename.")
    parser.add_argument(
        "--config",
        default="{}",
        type=json.loads,
        help="Algorithm-specific configuration (e.g. env, hyperparams). "
             "Gets merged with loaded configuration from checkpoint file and "
             "`evaluation_config` settings therein.")
    parser.add_argument(
        "--save-info",
        default=False,
        action="store_true",
        help="Save the info field generated by the step() method, "
             "as well as the action, observations, rewards and done fields.")
    parser.add_argument(
        "--use-shelve",
        default=False,
        action="store_true",
        help="Save rollouts into a python shelf file (will save each episode "
             "as it is generated). An output filename must be set using --out.")
    parser.add_argument(
        "--track-progress",
        default=False,
        action="store_true",
        help="Write progress to a temporary file (updated "
             "after each episode). An output filename must be set using --out; "
             "the progress file will live in the same folder.")

    # save and restore file management
    parser.add_argument(
        "--policy-dir", type=str, help="folder name of the policy.", default="")
    parser.add_argument(
        "--experiment", type=str, help="chosen experiment to reload.", default="")
    parser.add_argument(
        "--ncheckpoint", type=str, help="chosen checkpoint to reload.", default="")

    parser.add_argument(
        "--heuristic-policy", type=bool, help="chosen checkpoint to reload.", default=False)

    parser.add_argument(
        "--static-targets", type=bool, help="chosen checkpoint to reload.", default=False)

    parser.add_argument(
        "--video_dir", type=str, help="chosen folder to save video.", default="")

    parser.add_argument(
        "--horizon", type=int, help="limit of timesteps.", default=40)


    ### Old arguments needs a cleanup

    parser.add_argument("--scenario", type=str, default="simple_spread_assigned",
                        choices=['simple', 'simple_speaker_listener',
                                 'simple_crypto', 'simple_push',
                                 'simple_tag', 'simple_spread', 'simple_adversary', 'simple_spread_assigned',
                                 'matlab_simple_spread_assigned','matlab_simple_spread_assigned_hardcoll', 'matlab_simple_spread_assigned_checkpoints'],
                        help="name of the scenario script")
    parser.add_argument("--max-episode-len", type=int, default=100,
                        help="maximum episode length")
    parser.add_argument("--num-episodes", type=int, default=60000,
                        help="number of episodes")
    parser.add_argument("--num-adversaries", type=int, default=0,
                        help="number of adversaries")
    parser.add_argument("--good-policy", type=str, default="maddpg",
                        help="policy for good agents")
    parser.add_argument("--adv-policy", type=str, default="maddpg",
                        help="policy of adversaries")

    # Core training parameters
    parser.add_argument("--lr", type=float, default=1e-3,
                        help="learning rate for Adam optimizer")
    parser.add_argument("--gamma", type=float, default=0.99,
                        help="discount factor")
    # NOTE: 1 iteration = sample_batch_size * num_workers timesteps * num_envs_per_worker
    parser.add_argument("--sample-batch-size", type=int, default=25,
                        help="number of data points sampled /update /worker")
    parser.add_argument("--train-batch-size", type=int, default=1024,
                        help="number of data points /update")
    parser.add_argument("--n-step", type=int, default=1,
                        help="length of multistep value backup")
    parser.add_argument("--num-units", type=int, default=128,
                        help="number of units in the mlp")
    parser.add_argument("--replay-buffer", type=int, default=1000000,
                        help="size of replay buffer in training")
    parser.add_argument("--seed", type=int, default=100,
                        help="initialization seed for the network weights")

    # Checkpoint
    parser.add_argument("--checkpoint-freq", type=int, default = 10, #75,
                        help="save model once every time this many iterations are completed")
    parser.add_argument("--local-dir", type=str, default="./ray_results",
                        help="path to save checkpoints")
    parser.add_argument("--restore", type=str, default=None,
                        help="directory in which training state and model are loaded")
    parser.add_argument("--in-evaluation", type=bool, default=False, help="trigger evaluation procedure")

    # Parallelism
    #parser.add_argument("--num-workers", type=int, default=0)
    #parser.add_argument("--num-envs-per-worker", type=int, default=1)
    #parser.add_argument("--num-gpus", type=int, default=0)

    parser.add_argument("--num-workers", type=int, default=0)  #0
    parser.add_argument("--num-envs-per-worker", type=int, default=1)  #1
    parser.add_argument("--num-gpus", type=int, default=0)  #0
    #parser.add_argument("--num-cpus-per-worker", type=int, default=1)
    parser.add_argument("--num-gpus-per-worker", type=int, default=0)  #0

    # From the ppo
    parser.add_argument("--stop-iters", type=int, default=100)
    parser.add_argument("--stop-timesteps", type=int, default=160000000)
    # parser.add_argument("--stop-reward", type=float, default=7.99)

    # For rollouts
    parser.add_argument("--stop-iters-rollout", type=int, default=1)
    parser.add_argument("--nagents", type=int, default=1)
    parser.add_argument("--ntargets", type=int, default=1)
    parser.add_argument("--nrobots", type=int, default=1)

    # mode of hand-engineered comm. policy (-1 no hand-engineered)
    parser.add_argument("--mode", type=int, default=-1)
    parser.add_argument("--test", type=int, default=0, choices = [0,1], help="whether we want to test the policy or not")
    parser.add_argument("--test-env", type=int, default=0, choices = [0,1], help="whether we want to act in the test environment or not")
    parser.add_argument("--deterministic", type=int, default=1, choices=[0, 1],
                        help="enable exploration or not during execution")

    return parser