"""CLI for rolling out trained policies."""
import json

import click

from .utils import initialize_raylab


@click.command()
@click.argument(
    "checkpoint",
    type=click.Path(exists=True, file_okay=True, dir_okay=False, resolve_path=True),
)
@click.option(
    "--agent", required=True, default=None, help="Name of the trainable class to run."
)
@click.option(
    "--env",
    default=None,
    help="Name of the environment to interact with. "
    "Optional; can be inferred from config.",
    show_default=True,
)
@click.option(
    "--eval-config/--train-config",
    default=True,
    help="Whether to use the evaluation_config for the agent.",
    show_default=True,
)
@click.option(
    "--render/--no-render",
    default=False,
    help="Whether to render each episode.",
    show_default=True,
)
@click.option(
    "--video-dir",
    type=click.Path(exists=False, file_okay=False, dir_okay=True, resolve_path=True),
    default=None,
    help="Specifies the directory into which videos of all episode "
    "rollouts will be stored.",
)
@click.option(
    "--steps",
    default=10000,
    help="Number of timesteps to roll out (overwritten by --episodes).",
    show_default=True,
)
@click.option(
    "--episodes",
    default=0,
    help="Number of complete episodes to roll out (overrides --steps).",
    show_default=True,
)
@click.option(
    "--out",
    type=click.Path(file_okay=True, dir_okay=False, resolve_path=True),
    default=None,
    help="Output filename.",
)
@click.option(
    "--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.",
)
@click.option(
    "--save-info/--ignore-info",
    default=False,
    help="Save the info field generated by the step() method, "
    "as well as the action, observations, rewards and done fields.",
    show_default=True,
)
@click.option(
    "--shelve/--no-shelve",
    default=False,
    help="Save rollouts into a python shelf file (will save each episode "
    "as it is generated). An output filename must be set using --out.",
    show_default=True,
)
@click.option(
    "--track-progress/--no-progress",
    default=False,
    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.",
    show_default=True,
)
@initialize_raylab
def rollout(checkpoint, agent, env, eval_config, **rollout_kwargs):
    """Wrap `rllib rollout` with customized options."""
    # pylint:disable=too-many-locals
    import ray
    from ray.rllib.rollout import RolloutSaver
    from ray.rllib.rollout import rollout as rllib_rollout

    from raylab.utils.checkpoints import get_agent_from_checkpoint

    ray.init()
    agent = get_agent_from_checkpoint(
        checkpoint,
        agent,
        env,
        use_eval_config=eval_config,
        config_overrides=rollout_kwargs["config"],
    )

    with RolloutSaver(
        rollout_kwargs["out"],
        rollout_kwargs["shelve"],
        write_update_file=rollout_kwargs["track_progress"],
        target_steps=rollout_kwargs["steps"],
        target_episodes=rollout_kwargs["episodes"],
        save_info=rollout_kwargs["save_info"],
    ) as saver:
        rllib_rollout(
            agent,
            env,
            rollout_kwargs["steps"],
            rollout_kwargs["episodes"],
            saver,
            not rollout_kwargs["render"],
            rollout_kwargs["video_dir"],
        )
