import visualization as vis
from problems import Problems

import argparse
from argparse import RawTextHelpFormatter


def get_vis(plotter, func):
    if func == 0:
        plotter.plot_costs()
    elif func == 1:
        plotter.plot_pulls()
    elif func == 2:
        plotter.plot_cost_time()
    elif func == 3:
        plotter.plot_pulls_time()
    elif func == 4:
        plotter.plot_means_time()


# # Parse commandline arguments
parser = argparse.ArgumentParser()
parser = argparse.ArgumentParser(description='Driver for comparing cost weighted BAI algorithms.', formatter_class=RawTextHelpFormatter)

help = 'Pick problem(s):'
for i, problem in enumerate(Problems):
    help += '\n{} = {}'.format(i, problem.get_name())

parser.add_argument('-P', '--problems', help=help, nargs='*',
                    type=int, default=[i for i in range(len(Problems))])

help = 'Pick plots(s):\n\
        0: costs,\n\
        1: pulls,\n\
        2: costs_time,\n\
        3: pulls_time,\n\
        4: means_time'
parser.add_argument('-p', '--plots', help=help, nargs='*',
                    type=int, default=[i for i in range(len(Problems))])

args = parser.parse_args()

algos = ["KL-ULCB", "CR", "TAS", "CTAS", "CO"]
problems = []
for problem_num in args.problems:
    filename = Problems[problem_num].get_name()
    path = "./data/{}".format(filename)
    problems.append(path)

plotter = vis.Visualization(problems, algos)
for func in args.plots:
    get_vis(plotter, func)
