#!/usr/bin/env python3
"""Script for resuming training from command line."""
import click

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


@click.group()
def cli():  # noqa: D103
    pass


@cli.command()
@click.argument('from_dir')
@click.option(
    '--from_epoch',
    default='last',
    help='Index of epoch to restore from. Can be "first", "last" or a number. '
    'Not applicable when snapshot_mode="last"')
@click.option(
    '--log_dir',
    default=None,
    help='Path to save the log snapshot. If not specified, will be the same '
    'as from_dir.')
def resume(from_dir, from_epoch, log_dir):
    """Resume from command line."""
    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()
