import sys
from analysis_plots.parse import parse_info
from paper_plots import paper_conf as conf

import gflags
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import rc

from paper_plots.stat_util import pooled_sd

FLAGS = gflags.FLAGS

gflags.DEFINE_string('path', '/Users/Michael/Documents/deep_config/paper_results/baseline_run_1/,'
                             '/Users/Michael/Documents/deep_config/paper_results/baseline_run_2/,'
                             '/Users/Michael/Documents/deep_config/paper_results/baseline_run_3/',
                     'path to input')
gflags.DEFINE_string('prefix', 'dba_train_tiny_14400,'
                               'default_train_tiny_14400,'
                               'random_train_tiny_14400,'
                               'online_train_simple_tiny_14400,'
                               'no_pretrain_train_complex_tiny_14400,'
                               'online_train_complex_tiny_14400'
                     , 'experiment_prefix')

#RL-pre,RL-online,RL-pre-perturb

gflags.DEFINE_string('label', 'DBA,No indexing,Random,RL-pre,RL-online,RL-pre-perturb',
                     'experiment label')
gflags.DEFINE_string('export_path', '/Users/Michael/Documents/tf-paper/img/',
                     'export dir')

gflags.DEFINE_integer('queries', 10,
                      'experiment label')
gflags.DEFINE_integer('clients', 10, 'clients')

rc('text', usetex=True)
rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})
params = {'text.latex.preamble': [r'\usepackage{upgreek}', r'\usepackage{amsmath}']}
plt.rcParams.update(params)


def main(argv):
    try:
        argv = FLAGS(argv)
    except gflags.FlagsError as e:
        print('%s\\nUsage: %s ARGS\\n%s' % (e, sys.argv[0], FLAGS))

    path = FLAGS.path.split(',')
    prefix = FLAGS.prefix.split(',')
    label = FLAGS.label.split(',')
    export_path = FLAGS.export_path

    clients = FLAGS.clients
    queries = FLAGS.queries

    iterations = 3
    # TODO das mit paths ist noch falsch
    plot_queries(export_path, label, path, prefix, clients, queries, 0, 'train_', iterations)


def plot_queries(export_path, label, path, experiment_variant, clients, queries, offset, result_prefix, iterations):
    ind = np.arange(queries - 5)
    width = 0.15
    # Iterate over all prefixes. For each prefix, compute all queries

    for k in range(len(experiment_variant)):
        iteration_means = np.zeros((iterations, (queries - 5)))
        iteration_stds = np.zeros((iterations, (queries - 5)))
        sample_sizes = np.zeros((iterations, (queries - 5)))

        for iteration in range(iterations):
            update, query = parse_info(path[iteration] + '/' + experiment_variant[k], clients, queries, False, offset)

            for i in range(len(query) - 5):
                # print('Length of query content' + str(len(query[i])))
                scaled_up = np.asarray(query[i] * 1000)  # in ms
                iteration_means[iteration][i] = np.median(scaled_up)
                # print('Mean value for query {} is {}'.format( offset + i + 1, iteration_means[iteration][i]))
                iteration_stds[iteration][i] = np.std(scaled_up)
                sample_sizes[iteration][i] = len(scaled_up)

        means = np.mean(iteration_means, axis=0)
        stds = []

        for query in range(queries - 5):
            # Das hier sind jetzt die einzelnen stds fuer die query entlang der iteration achse
            std_list = list(iteration_stds[:, query])
            sizes = list(sample_sizes[:, query])

            stds.append(pooled_sd(sd_list=std_list, sample_sizes=sizes, num_samples=iterations))

        plt.bar(ind + k * width, means, yerr=stds, bottom=0, width=width,
                color=conf.greyscales[k], hatch=conf.hatches[k])

    plt.tick_params(axis='both', which='major', labelsize=conf.font_size)
    plt.xlabel('Training query id', fontsize=conf.font_size)
    plt.xticks(ind, ('Q1', 'Q2', 'Q3', 'Q4', 'Q5'))
    plt.ylabel('Median execution time (ms)', fontsize=conf.font_size)
    #plt.ylim(0, 6000)
    plt.legend(label, loc='upper left')
    plt.yscale('log')
    plt.tight_layout()

    plt.savefig('/Users/Michael/Documents/deep_config/analysis/paper_img/train_query_bars_first.pdf', format="pdf", bbox_inches="tight")
    plt.savefig(export_path + result_prefix + 'query_bars_first.pdf', format="pdf", bbox_inches="tight")
    plt.clf()
    plt.cla()
    plt.close()

    for k in range(len(experiment_variant)):
        iteration_means = np.zeros((iterations, (queries - 5)))
        iteration_stds = np.zeros((iterations, (queries - 5)))
        sample_sizes = np.zeros((iterations, (queries - 5)))

        for iteration in range(iterations):
            update, query = parse_info(path[iteration] + '/' + experiment_variant[k], clients, queries, False, offset)

            for i in range(5):
                j = i + 5
                scaled_up = np.asarray(query[j] * 1000)  # in ms

                #median or mean?
                iteration_means[iteration][i] = np.median(scaled_up)
                # print('Mean value for query {} is {}'.format(offset + j + 1, iteration_means[iteration][i]))
                iteration_stds[iteration][i] = np.std(scaled_up)
                sample_sizes[iteration][i] = len(scaled_up)

        means = np.mean(iteration_means, axis=0)
        stds = []

        for query in range(queries - 5):
            std_list = list(iteration_stds[:, query])
            sizes = list(sample_sizes[:, query])

            stds.append(pooled_sd(sd_list=std_list, sample_sizes=sizes, num_samples=iterations))

        plt.bar(ind + k * width, means, yerr=stds, bottom=0, width=width,
                color=conf.greyscales[k], hatch=conf.hatches[k])

    plt.xlabel('Training query id', fontsize=conf.font_size)
   # plt.ylim(0, 6000)
    plt.tick_params(axis='both', which='major', labelsize=conf.font_size)
    plt.xticks(ind, ('Q6', 'Q7', 'Q8', 'Q9', 'Q10'))
    plt.ylabel('Median execution time (ms)', fontsize=conf.font_size)
   # plt.legend(label, loc='upper left')
    plt.yscale('log')
    plt.tight_layout()

    plt.savefig('/Users/Michael/Documents/deep_config/analysis/paper_img/train_query_bars_second.pdf', format="pdf", bbox_inches="tight")

    plt.savefig(export_path + result_prefix + 'query_bars_second.pdf', format="pdf", bbox_inches="tight")
    plt.clf()
    plt.cla()
    plt.close()


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