from os import path
import glob
from make_graphs import mpl, rcParams, plt
import json
import argparse
import numpy as np
import math
import csv
from make_graphs import get_dist_args
from compute_tp import SeriesStats
from collections import defaultdict
# 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"
dist_color1, cent_color1, chopin_color1, opt_online_color1 = "olive", "steelblue", "navy", "purple"



def make_bars_graph_from_degree_graphs(graph_names, output_path, degree):
    # 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"][graphs[traffic_dist]["x"].index(degree)])
        distributed.append(graphs[traffic_dist]["dist_only"][graphs[traffic_dist]["x"].index(degree)])
        chopin.append(graphs[traffic_dist]["chopin"][graphs[traffic_dist]["x"].index(degree)])
    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)
    
    plt.savefig(path.join(output_path, path.basename(get_normalize_graph_name(graph_names[0])) + "_deg%d_comb_imp.pdf"%degree),
                bbox_extra_artists=(lg,), bbox_inches='tight')
    plt.show()
    plt.close()


def make_flow_size_graph_from_degree_graphs(graph_names, output_path, degree):
    # tors80 mwm-delay2T win-size1 win-delay2 iters1 avg-rate~172Mbps DCTCP mwm-top5_degs_graph
    print graph_names
    
    def flow_size_from_file_name(file_name):
        start = file_name.find("avg_flow~") + len("avg_flow~")
        end = file_name[start:].find("KB") + start
        return int(file_name[start:end])
    
    graphs = {flow_size_from_file_name(name): json.load(open(name)) for name in graph_names}
    centralized = []
    distributed = []
    chopin = []
    opt_online = []
    flow_sizes = sorted(graphs.keys())
    for flow_size in flow_sizes:
        centralized.append(graphs[flow_size]["centralized"][graphs[flow_size]["x"].index(degree)])
        distributed.append(graphs[flow_size]["dist_only"][graphs[flow_size]["x"].index(degree)])
        chopin.append(graphs[flow_size]["chopin"][graphs[flow_size]["x"].index(degree)])
        opt_online.append(graphs[flow_size]["opt_online"][graphs[flow_size]["x"].index(degree)])
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(flow_sizes, distributed, color=dist_color1, label="distributed")  # yerr=[(0,0,0), tuple(chopin[i]-distributed[i] for i in range(3))])
    ax.plot(flow_sizes, centralized, color=cent_color1, label="centralized")  # yerr=[(0,0,0), tuple(chopin[i]-centralized[i] for i in range(3))])
    #ax.plot(flow_sizes, chopin, color=chopin_color1, label="chopin")
    #ax.plot(flow_sizes, opt_online, color=opt_online_color1, label="opt_online")
    
    ax.set_ylabel('Optical Throughput Ratio')
    ax.set_ylim(0.4, 1)
    #ax.set_ylim(0, 1)
    ax.set_xscale('log')
    # ax.set_yscale('log')
    ax.set_xlabel('Mean Flow Size (KB)')
    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])) + "_deg%d_comb_flow_size.pdf"%degree),
        bbox_extra_artists=(lg,), bbox_inches='tight')
    plt.show()
    plt.close()


def make_flow_size_and_rate_graph_from_degree_graphs(graph_name_patterns, output_path, degree):
    # tors80 mwm-delay2T win-size1 win-delay2 iters1 avg-rate~172Mbps DCTCP mwm-top5_degs_graph
    print graph_name_patterns
    graph_names = set()
    for pattern in graph_name_patterns:
        graph_names |= set(glob.glob(pattern))
    graph_names = list(graph_names)

    def avg_rate_from_file_name(file_name):
        start = file_name.find("avg-rate~") + len("avg-rate~")
        end = file_name[start:].find("Mb") + start
        return int(file_name[start:end])
    
    def flow_size_from_file_name(file_name):
        start = file_name.find("avg_flow~") + len("avg_flow~")
        end = file_name[start:].find("KB") + start
        return int(file_name[start:end])
    
    def str_keys(d):
        return {str(k): d[k] for k in d}
    
    shapes = {"HULL": "o", "pFabric": "s", "VL2": "^"}
    
    graphs = {(flow_size_from_file_name(name), avg_rate_from_file_name(name), get_graph_key(name)): json.load(open(name)) for name in graph_names}
    json.dump(str_keys(graphs), open("mega_comb_graph_deg%d_input.json"%degree, "w"), indent=True, sort_keys=True)
    results = defaultdict(lambda: defaultdict(list))
    for k in graphs:
        dist = k[2]
        if graphs[k]["dist_only"][graphs[k]["x"].index(degree)] > graphs[k]["centralized"][graphs[k]["x"].index(degree)]:
            results[dist][True].append((k[0], k[1]))
        else:
            results[dist][False].append((k[0], k[1]))

    json.dump(results, open("mega_comb_graph_deg%d_output.json"%degree, "w"), indent=True, sort_keys=True)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    for dist in results:
        shape = shapes[dist]
        for cond in results[dist]:
            if cond:
                label = dist + " dist.>cent."
                color = "r"
            else:
                label = dist + " cent.>dist."
                color = "b"
            x = [p[0] for p in results[dist][cond]]
            y = [p[1] for p in results[dist][cond]]
            ax.scatter(x, y, color=color, marker=shape, label=label)
            
    ax.set_ylabel('Mean Host Rate (Mb/s)')
    # ax.set_ylim(0.4, 0.95)
    ax.set_ylim(0, 1020)
    ax.set_xlim(-5, 27000)
    #ax.set_xscale('log')
    #ax.set_yscale('log')
    ax.set_xlabel('Mean Flow Size (KB)')
    lg = plt.legend(bbox_to_anchor=(1.05, 1.0), loc='upper left')
    
    plt.savefig(
        "mega_comb_graph_deg%d.pdf" % degree,
        bbox_extra_artists=(lg,), bbox_inches='tight')
    plt.show()
    plt.close()


def make_flow_size_and_rate_graph_per_dist_from_degree_graphs(graph_name_patterns, output_path, degree):
    # tors80 mwm-delay2T win-size1 win-delay2 iters1 avg-rate~172Mbps DCTCP mwm-top5_degs_graph
    print graph_name_patterns
    graph_names = set()
    for pattern in graph_name_patterns:
        graph_names |= set(glob.glob(pattern))
    graph_names = list(graph_names)
    
    def avg_rate_from_file_name(file_name):
        start = file_name.find("avg-rate~") + len("avg-rate~")
        end = file_name[start:].find("Mb") + start
        return int(file_name[start:end])
    
    def flow_size_from_file_name(file_name):
        start = file_name.find("avg_flow~") + len("avg_flow~")
        end = file_name[start:].find("KB") + start
        return int(file_name[start:end])
    
    def str_keys(d):
        return {str(k): d[k] for k in d}
    
    shapes = {"HULL": "o", "pFabric": "s", "VL2": "^"}
    
    graphs = {
        (flow_size_from_file_name(name), avg_rate_from_file_name(name), get_graph_key(name)): json.load(open(name)) for
        name in graph_names}
    json.dump(str_keys(graphs), open("mega_comb_graph_deg%d_input.json" % degree, "w"), indent=True, sort_keys=True)
    results = defaultdict(lambda: defaultdict(list))
    for k in graphs:
        dist = k[2]
        if graphs[k]["dist_only"][graphs[k]["x"].index(degree)] > graphs[k]["centralized"][
            graphs[k]["x"].index(degree)]:
            results[dist][True].append((k[0], k[1]))
        else:
            results[dist][False].append((k[0], k[1]))
    
    json.dump(results, open("mega_comb_graph_deg%d_output.json" % degree, "w"), indent=True, sort_keys=True)
    for dist in results:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        shape = "x"
        for cond in results[dist]:
            if cond:
                label ="dist.>cent."
                color = "r"
            else:
                label = "cent.>dist."
                color = "b"
            x = [p[0] for p in results[dist][cond]]
            y = [p[1] for p in results[dist][cond]]
            ax.scatter(x, y, color=color, marker=shape, label=label)
    
        ax.set_ylabel('Mean Host Rate (Mb/s)')
        # ax.set_ylim(0.4, 0.95)
        ax.set_ylim(0, 1020)
        ax.set_xlim(-5, 27000)
        # ax.set_xscale('log')
        # ax.set_yscale('log')
        ax.set_xlabel('Mean Flow Size (KB)')
        lg = plt.legend(bbox_to_anchor=(1.05, 1.0), loc='upper left')
        
        plt.savefig(
            "mega_comb_graph_deg%d_%s.pdf" % (degree,dist),
            bbox_extra_artists=(lg,), bbox_inches='tight')
        plt.show()
        plt.close()


def make_flow_size_and_rate_heatmap_per_dist_from_degree_graphs(graph_name_patterns, output_path, degree):
    # tors80 mwm-delay2T win-size1 win-delay2 iters1 avg-rate~172Mbps DCTCP mwm-top5_degs_graph
    import matplotlib.colors
    
    print graph_name_patterns
    graph_names = set()
    for pattern in graph_name_patterns:
        graph_names |= set(glob.glob(pattern))
    graph_names = list(graph_names)
    print "\n".join(graph_names)
    map_size = (11, 6)
    map_square_width = (2000, 100)
    map_factors = (1000, 1)
    def avg_rate_from_file_name(file_name):
        start = file_name.find("avg-rate~") + len("avg-rate~")
        end = file_name[start:].find("Mb") + start
        return int(file_name[start:end])
    
    def flow_size_from_file_name(file_name):
        start = file_name.find("avg_flow~") + len("avg_flow~")
        end = file_name[start:].find("KB") + start
        return int(file_name[start:end])
    
    def str_keys(d):
        return {str(k): d[k] for k in d}
    
    def get_matrix(init_val=np.NaN):
        return [[init_val for i in range(map_size[0])] for j in range(map_size[1])]

    def get_index(f, width):
        i = np.round(f * 1.0 / width) -1
        if i < 0:
            return 0
        return int(i)

    cmap = plt.get_cmap('RdBu')
    cmap.set_bad(color='lightgrey')
    
    shapes = {"HULL": "o", "pFabric": "s", "VL2": "^"}
    
    graphs = {
        (flow_size_from_file_name(name), avg_rate_from_file_name(name), get_graph_key(name)): json.load(open(name)) for
        name in graph_names}
    json.dump(str_keys(graphs), open("mega_comb_heatmap_deg%d_input.json" % degree, "w"), indent=True, sort_keys=True)
    heatmap = defaultdict(get_matrix)
    for k in graphs:
        dist = k[2]
        flow, rate = get_index(k[0], map_square_width[0]), get_index(k[1], map_square_width[1])
        if flow >= map_size[0] or rate >= map_size[1]:
            continue
        val = graphs[k]["dist_only"][graphs[k]["x"].index(degree)] - graphs[k]["centralized"][graphs[k]["x"].index(degree)]
        #heatmap[dist][rate][flow] = val
        #continue
        if heatmap[dist][rate][flow] is np.NaN:
            heatmap[dist][rate][flow] = val
        elif math.fabs(heatmap[dist][rate][flow]) < math.fabs(val):
            heatmap[dist][rate][flow] = val
    
    json.dump(heatmap, open("mega_comb_heatmap_deg%d_output.json" % degree, "w"), indent=True, sort_keys=True)
    
    for dist in heatmap:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        pos = ax.imshow(heatmap[dist], cmap=cmap, interpolation='nearest',  vmin=-0.17, vmax=0.17)
        fig.colorbar(pos, ax=ax)
        ax.set_ylabel('Mean Host Rate (Mb/s)')
        #ax.set_yticklabels([str(i*100) for i in range(1, map_size[1]+1)], minor=True)
        plt.yticks(range(map_size[1]), [str(i*map_square_width[1]/map_factors[1]) for i in range(1, map_size[1]+1)], fontsize=14)
        ax.invert_yaxis()
        # ax.set_ylim(0.4, 0.95)
        #ax.set_ylim(0, 1020)
        #ax.set_xlim(-5, 27000)
        # ax.set_xscale('log')
        # ax.set_yscale('log')
        ax.set_xlabel('Mean Flow Size (MB)')
        plt.xticks(range(map_size[0]), [str(i*map_square_width[0]/map_factors[0]) for i in range(1, map_size[0]+1)], fontsize=12, rotation=45)
        #ax.set_xticklabels([str(i) for i in range(1, map_size[0]+1)], minor=True)
        lg = plt.legend(bbox_to_anchor=(1.05, 1.0), loc='upper left')
        #ax.minorticks_on()

        plt.savefig(
            "mega_comb_heatmap_deg%d_%s.pdf" % (degree, dist),
            bbox_extra_artists=(lg,), bbox_inches='tight')
        plt.show()
        plt.close()
    

def make_avg_rate_graph_from_degree_graphs(graph_names, output_path, degree):
    # tors80 mwm-delay2T win-size1 win-delay2 iters1 avg-rate~172Mbps DCTCP mwm-top5_degs_graph
    print graph_names
    
    def avg_rate_from_file_name(file_name):
        start = file_name.find("avg-rate~") + len("avg-rate~")
        end = file_name[start:].find("Mb") + start
        return int(file_name[start:end])
    
    graphs = {avg_rate_from_file_name(name): json.load(open(name)) for name in graph_names}
    centralized = []
    distributed = []
    chopin = []
    opt_online=[]
    avg_rates = sorted(graphs.keys())
    for rate in avg_rates:
        centralized.append(graphs[rate]["centralized"][graphs[rate]["x"].index(degree)])
        distributed.append(graphs[rate]["dist_only"][graphs[rate]["x"].index(degree)])
        chopin.append(graphs[rate]["chopin"][graphs[rate]["x"].index(degree)])
        opt_online.append(graphs[rate]["opt_online"][graphs[rate]["x"].index(degree)])
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(avg_rates, distributed, color=dist_color1, label="distributed")  # yerr=[(0,0,0), tuple(chopin[i]-distributed[i] for i in range(3))])
    ax.plot(avg_rates, centralized, color=cent_color1, label="centralized")  # yerr=[(0,0,0), tuple(chopin[i]-centralized[i] for i in range(3))])
    #ax.plot(avg_rates, chopin, color=chopin_color1, label="chopin")
    # ax.plot([],[],"k", label="improvement")
    #ax.plot(avg_rates, opt_online, color=opt_online_color1, label="opt_online")
    
    ax.set_ylabel('Optical Throughput Ratio')
    ax.set_ylim(0, 1)
    #ax.set_xscale('log')
    #ax.set_yscale('log')
    ax.set_xlabel('Mean Host Rate (Mb/s)')
    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])) + "_deg%d_comb_avg_rate.pdf"%degree),
                bbox_extra_artists=(lg,), bbox_inches='tight')
    plt.show()
    plt.close()


def make_peers_num_graph_from_degree_graphs(graph_names, output_path, degree):
    # tors80 mwm-delay2T win-size1 win-delay2 iters1 avg-rate~172Mbps DCTCP mwm-top5_degs_graph
    print graph_names
    
    graph_dicts = [json.load(open(name)) for name in graph_names]
    graphs = {d["_graph_params"]["conf"]["n_host_peers"]*2: d for d in graph_dicts} # multiplying by 2 to get corresponding peers
    centralized = []
    distributed = []
    chopin = []
    opt_online = []
    peers_nums = sorted(graphs.keys())
    for peers_num in peers_nums:
        centralized.append(graphs[peers_num]["centralized"][graphs[peers_num]["x"].index(degree)])
        distributed.append(graphs[peers_num]["dist_only"][graphs[peers_num]["x"].index(degree)])
        chopin.append(graphs[peers_num]["chopin"][graphs[peers_num]["x"].index(degree)])
        opt_online.append(graphs[peers_num]["opt_online"][graphs[peers_num]["x"].index(degree)])
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(peers_nums, distributed, color=dist_color1, label="distributed")  # yerr=[(0,0,0), tuple(chopin[i]-distributed[i] for i in range(3))])
    ax.plot(peers_nums, centralized, color=cent_color1, label="centralized")  # yerr=[(0,0,0), tuple(chopin[i]-centralized[i] for i in range(3))])
    #ax.plot(peers_nums, chopin, color=chopin_color1, label="chopin")
    # ax.plot([],[],"k", label="improvement")
    #ax.plot(peers_nums, opt_online, color=opt_online_color1, label="opt_online")
    
    ax.set_ylabel('Optical Throughput Ratio')
    ax.set_ylim(0, 1)
    # ax.set_xscale('log')
    # ax.set_yscale('log')
    ax.set_xlabel('Average Number of Peers per Host')
    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])) + "_deg%d_comb_peers_num.pdf"% degree),
        bbox_extra_artists=(lg,), bbox_inches='tight')
    plt.show()
    plt.close()


def make_bars_graph_from_reconf_graphs(graph_names, output_path, epoch_index=3):
    # 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="+")#, 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()
    degree = 4
    if args.func == 'improvement':
        make_bars_graph_from_degree_graphs(args.graphs, args.output_path, degree)
    elif args.func == 'reconf':
        make_bars_graph_from_reconf_graphs(args.graphs, args.output_path)
    elif args.func == 'flow_size':
        make_flow_size_graph_from_degree_graphs(args.graphs, args.output_path, degree)
    elif args.func == 'peers_num':
        make_peers_num_graph_from_degree_graphs(args.graphs, args.output_path, degree)
    elif args.func == 'avg_rate':
        make_avg_rate_graph_from_degree_graphs(args.graphs, args.output_path, degree)
    elif args.func == 'rate_and_flow':
        make_flow_size_and_rate_graph_from_degree_graphs(args.graphs, args.output_path, degree)
        make_flow_size_and_rate_graph_per_dist_from_degree_graphs(args.graphs, args.output_path, degree)
    elif args.func == 'heatmap':
        make_flow_size_and_rate_heatmap_per_dist_from_degree_graphs(args.graphs, args.output_path, degree)
    else:
        print "unmatched function"


if __name__ == "__main__":
    main()