from os import path
from make_graphs import mpl, rcParams, plt
import json
import argparse
import numpy as np

# tors100 mwm-delay2T win-size1 win-delay2 iters1 avg-rate~176Mbps HULL mwm-top5_degs_graph.json


GRAPH_KEYS = {"HULL": "HULL", "DCTCP": "pFabric", "VL2": "VL2"}


def get_graph_key(graph_name):
    for key in GRAPH_KEYS:
        if key in graph_name:
            return GRAPH_KEYS[key]
    raise Exception("Unknown graph key (%s)"%graph_name)


def get_graphs_dict(graph_names):
    return {get_graph_key(name): json.load(open(name)) for name in graph_names}


def get_normalize_graph_name(graph_name):
    new_name = graph_name
    for k in GRAPH_KEYS:
        new_name = graph_name.replace(k, "")
    return new_name

#dist_color, cent_color, chopin_color = "r", "g", "b"
#dist_color, cent_color, chopin_color = "cyan", "steelblue", "b"
#dist_color, cent_color, chopin_color = "gold", "teal", "steelblue"
dist_color, cent_color, chopin_color = "gray", "steelblue", "navy"

def make_bars_graph_from_degree_graphs(graph_names, output_path, degree_index=3):
    # tors80 mwm-delay2T win-size1 win-delay2 iters1 avg-rate~172Mbps DCTCP mwm-top5_degs_graph
    print graph_names
    graphs = get_graphs_dict(graph_names)
    centralized = []
    distributed = []
    chopin = []
    traffic_dists = GRAPH_KEYS.values()
    for traffic_dist in traffic_dists:
        centralized.append(graphs[traffic_dist]["centralized"][degree_index])
        distributed.append(graphs[traffic_dist]["dist_only"][degree_index])
        chopin.append(graphs[traffic_dist]["chopin"][degree_index])
    N = 3
    ind = np.arange(N)  # the x locations for the groups
    width = 0.27  # the width of the bars

    fig = plt.figure()
    ax = fig.add_subplot(111)
    rects1 = ax.bar(ind, distributed, width, color=dist_color, label="distributed",
                    yerr=[(0,0,0), tuple(chopin[i]-distributed[i] for i in range(3))])
    rects2 = ax.bar(ind+width, centralized, width, color=cent_color, label="centralized",
                    yerr=[(0,0,0), tuple(chopin[i]-centralized[i] for i in range(3))])
    rects3 = ax.bar(ind+width*2, chopin, width, color=chopin_color, label="chopin")
    ax.plot([],[],"k", label="improvement")

    ax.set_ylabel('Optical Throughput Ratio')
    ax.set_xticks(ind + width)
    ax.set_xticklabels(traffic_dists)
    ax.legend((rects1[0], rects2[0], rects3[0]), ('distributed', 'centralized', 'chopin'))
    ax.set_ylim(0,1)
    lg = plt.legend(bbox_to_anchor=(1.05, 1.0), loc='upper left')

    def autolabel(rects):
        for i in range(len(rects)):
            rect = rects[i]
            c_rect = rects3[i]
            h = rect.get_height()
            h3 = c_rect.get_height()
            ax.text(rect.get_x() + rect.get_width() / 2., 1.05 * h3, '%.1f%%' % float((h3-h)*100.0/h),
                    ha='center', va='bottom')

    autolabel(rects1)
    autolabel(rects2)
    #autolabel(rects3)
    
    plt.savefig(path.join(output_path, path.basename(get_normalize_graph_name(graph_names[0])) + "_deg_comb_imp.pdf"), bbox_extra_artists=(lg,), bbox_inches='tight')
    plt.show()
    plt.close()


def make_bars_graph_from_reconf_graphs(graph_names, output_path, epoch_index=1):
    # tors80 mwm-delay2T win-size1 win-delay2 iters1 max-deg4 avg-rate~172Mbps DCTCP mwm-top5_reconf_graph
    # print graph_names
    graphs = get_graphs_dict(graph_names)
    centralized = []
    distributed = []
    chopin = []
    traffic_dists = GRAPH_KEYS.values()
    for traffic_dist in traffic_dists:
        centralized.append(graphs[traffic_dist]["centralized"][epoch_index])
        distributed.append(graphs[traffic_dist]["dist_only"][epoch_index])
        chopin_values = [graphs[traffic_dist][k][epoch_index] for k in graphs[traffic_dist] if
                         k.startswith("dist_") and k != "dist_only"]
        chopin.append(max(chopin_values))
    
    N = 3
    ind = np.arange(N)  # the x locations for the groups
    width = 0.27  # the width of the bars
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    rects1 = ax.bar(ind, distributed, width, color=dist_color, label="distributed")
    rects2 = ax.bar(ind + width, centralized, width, color=cent_color, label="centralized")
    rects3 = ax.bar(ind + width * 2, chopin, width, color=chopin_color, label="chopin")
    
    ax.set_ylabel('Reconfiguration Ratio')
    ax.set_xticks(ind + width)
    ax.set_xticklabels(traffic_dists)
    ax.legend((rects1[0], rects2[0], rects3[0]), ('distributed', 'centralized', 'chopin'))
    ax.set_ylim(0, 1)
    lg = plt.legend(bbox_to_anchor=(1.05, 1.0), loc='upper left')
    
    plt.savefig(path.join(output_path, path.basename(get_normalize_graph_name(graph_names[0])) + "_comb_reconf.pdf"),
                bbox_extra_artists=(lg,), bbox_inches='tight')
    plt.show()
    plt.close()


def main():
    parser = argparse.ArgumentParser(
        description="""Make graphs """,
        epilog="""
    """)
    parser.add_argument('-g', '--graphs', help='<Required> graph jsons', required=True, nargs=3)#, action='append'
    parser.add_argument('--output_path', default="results",
                        help='output path (default: results)')
    parser.add_argument('--func', default="improvement",
                        help='output path (default: improvement)')
    args = parser.parse_args()
    if args.func == 'improvement':
        make_bars_graph_from_degree_graphs(args.graphs, args.output_path)
    elif args.func == 'reconf':
        make_bars_graph_from_reconf_graphs(args.graphs, args.output_path)
    else:
        print "unmatched function"


if __name__ == "__main__":
    main()