import argparse

from tsp import TSP_Bench_ALL
from tsp import TSP_Bench_PATH

from model.my_model import MyModel

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument('-e', '--easy', action='store_true', help='Benchmark all easy level TSP')
    parser.add_argument('-m', '--medium', action='store_true', help='Benchmark all medium level TSP')
    parser.add_argument('-d', '--difficult', action='store_true', help='Benchmark all difficult level TSP')

    args = parser.parse_args()

    if (args.easy):
        TSP_Bench_PATH("./data/easy/", MyModel, timeout=60)
    if (args.medium):
    	TSP_Bench_PATH("./data/medium/", MyModel, timeout=180)
    if (args.difficult):
        TSP_Bench_PATH("./data/difficult/", MyModel, timeout=300)

    if( (not args.easy) and (not args.medium) and (not args.difficult) ):
        TSP_Bench_ALL('./', MyModel)