#!/usr/bin/env python

'''
    This script can be used to download data from the UTIAS-Multiseason Dataset
    and the UTIAS In The Dark dataset.

    This code is built on the download script provided at: 
    http://robots.engin.umich.edu/nclt/
'''

from __future__ import print_function
import sys
import os
import subprocess
import argparse

base_dir = 'ftp://asrl3.utias.utoronto.ca/M600'

#dataset_dir = {'multiseason': '%s/UTIAS-Multiseason' % (base_dir), 
#                 'inthedark': '%s/UTIAS-In-The-Dark' % (base_dir)}

#dataset_name = {'multiseason': 'UTIAS Multiseason Dataset', 
#                  'inthedark': 'UTIAS In The Dark Dataset'}

num_trials = {'montreal': 10, 'suffield': 15, 'utiascircle': 2, 'utiasfield': 5, 'utiasstraight': 1, 'utiasday': 5, 'utiaswinter': 4}

def main (args):

    getopt = argparse.ArgumentParser(description='Download VT&R dataset')
    getopt.add_argument('--montreal', action='store_true',
            help='Download data form the Montreal Dataset')
    getopt.add_argument('--suffield', action='store_true',
            help='Download data form the Suffield Dataset')
    getopt.add_argument('--utiascircle', action='store_true',
            help='Download data form the UTIAS Circle Dataset')
    getopt.add_argument('--utiasfield', action='store_true',
            help='Download data form the UTIAS Field Dataset')
    getopt.add_argument('--utiasstraight', action='store_true',
            help='Download data form the UTIAS Straight Dataset')
    getopt.add_argument('--utiasday', action='store_true',
            help='Download data form the UTIAS Day Dataset')
    getopt.add_argument('--utiaswinter', action='store_true',
            help='Download data form the UTIAS Winter Dataset')
    getopt.add_argument('--all', action='store_true',
            help='Download all trials form the given dataset')
    getopt.add_argument('--trials', type=int, nargs='+',
            help='Download trials listed as integers, e.g. --trials 2 21 35')
    getopt.add_argument('--trials_range', type=int, nargs=2,
            help='Download trials in range (inclusive) given as two integers, \
                  e.g. --trials_range 3 20')

    args = getopt.parse_args()

    if args.montreal:
        dataset = 'montreal'
    elif args.suffield:
        dataset = 'suffield'
    elif args.utiascircle:
        dataset = 'utiascircle'
    elif args.utiasfield:
        dataset = 'utiasfield'
    elif args.utiasstraight:
        dataset = 'utiasstraight'
    elif args.utiasday:
        dataset = 'utiasday'
    elif args.utiaswinter:
        dataset = 'utiaswinter'
    else:
        print('Dataset has not been specified. Use --help for options.')
        return 0

    if not args.all and args.trials is None and args.trials_range is None:
        print('Runs have not been specified. Use --help for options.')
        return 0

    print('Downloading data from %s\n' % dataset)

    trials = []
    if args.all:
        trials = list(range(num_trials[dataset]))
    elif args.trials is not None:
        trials = args.trials
    elif args.trials_range is not None:
        trials = list(range(args.trials_range[0], args.trials_range[1] + 1))

    if len(trials) == 0:
        print('An empty or invalid range of runs were specified.' \
              ' Please choose a different set of runs.')
        return 0 

    print('Preparing to download the following trials: \n%s\n' % (trials))

    for trial in trials:

        if trial not in range(num_trials[dataset]):
            print('Trial %d not in dataset' % (trial))
            continue

        print('%s/%s_trial%s.zip' % (base_dir, dataset, str(trial)))

        cmd = ['wget', '--continue', 
               '%s/%s_trial%s.zip' % (base_dir, dataset, str(trial))]

        #cmd = ['wget', '--continue', 'ftp://asrl3.utias.utoronto.ca/M600/montreal_trial1.zip']

        subprocess.call(cmd)

    return 0

if __name__ == '__main__':
    sys.exit (main(sys.argv))