#!/usr/bin/env python3
"""Script for resuming training from command line."""
# pylint: disable=no-name-in-module
import click

from garage.experiment import LocalTFRunner
from garage.experiment import run_experiment


@click.group()
def cli():  # noqa: D103
    """The main command group."""


@cli.command()
@click.argument('from_dir')
@click.option(
    '--log_dir',
    default=None,
    help='Log path for resumed experiment. If not specified, will be the same '
    'as from_dir.')
# pylint: disable=bad-docstring-quotes
@click.option('--from_epoch',
              default='last',
              help='When there are multiple snapshots, '
              'specify the index of epoch to restore from. '
              'Can be "first", "last" or a number. '
              'Not applicable when snapshot_mode="last"')
def resume(from_dir, from_epoch, log_dir):
    # pylint: disable=missing-param-doc, missing-type-doc
    """Resume from experiment saved in FROM_DIR."""
    if log_dir is None:
        log_dir = from_dir

    def run_task(_snapshot_config, _variant_data, _from_dir, _from_epoch):
        with LocalTFRunner(snapshot_config=_snapshot_config) as runner:
            runner.restore(from_dir=_from_dir, from_epoch=_from_epoch)
            runner.resume()

    run_experiment(run_task,
                   log_dir=log_dir,
                   resume_from_dir=from_dir,
                   resume_from_epoch=from_epoch)


if __name__ == '__main__':
    cli()
