import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import math
sns.set(style="white", context="talk")
plt.rcdefaults()
DPI=300
BBOX_INCHES='tight'

def autolabel(ax, rect):
    """
    Attach  a text label above each bar displaying its height
    """
    height = rect.get_height()
    ax.text(rect.get_x() + rect.get_width()/2., 1.05*height,
        '%d' % int(round(height)),
        ha='center', va='bottom')

path = 'figures/cpu_matvec'
fig, ax = plt.subplots(figsize=(6, 3))
systems = ('XE (32T)', "S822LC (160T)")
mlfmm = (8.65e4, 4.84e4)
total = (1.2e5, 5.77e4)
x_pos = np.arange(len(systems))

ax.bar(x_pos, mlfmm, color='0.4', label='MLFMM')
ax.bar(x_pos, [i-j for i,j in zip(total, mlfmm)], color='0.8', bottom=mlfmm, label='Non-MLFMM')
ax.set_xticks(x_pos)
ax.set_xticklabels(systems)
ax.set_ylabel("Execution Time (ms)")
handles, labels = ax.get_legend_handles_labels()
ax.legend(handles, labels)
sns.despine(trim=True)
plt.savefig(path+'.pdf', bbox_inches=BBOX_INCHES)
plt.savefig(path+'.png', dpi=DPI, bbox_inches=BBOX_INCHES)

print path



path = 'figures/kernels'
fig, ax = plt.subplots()
systems = ['32T\n(1 XE)', "160T\n(1 S822LC)", "4 GPU\n(4 XK)" ,"4 GPU\n(1 S822LC)"]
x_pos = np.arange(len(systems))
p2m = (127.10, 72.10749, 7.73, 1.604)
m2m = (156.2506, 102.61091, 9.613814, 1.746476)
m2l = (189.615, 82.67791, 18.177774, 2.671025)
l2l = (91.5957, 101.56461, 20.215436, 2.611185)
l2p = (196.2115, 68.38529, 6.994, 1.395)
p2p = (1117.368, 590.4818, 90.619, 18.265)
total = [sum(i) for i in zip(p2m,m2m,m2l,l2l,l2p,p2p)]
p2m_ratio = [i/j for i,j in zip(p2m, total)]
m2m_ratio = [i/j for i,j in zip(m2m, total)]
m2l_ratio = [i/j for i,j in zip(m2l, total)]
l2l_ratio = [i/j for i,j in zip(l2l, total)]
l2p_ratio = [i/j for i,j in zip(l2p, total)]
p2p_ratio = [i/j for i,j in zip(p2p, total)]
ax.bar(x_pos, p2m_ratio, color='0.75', label='P2M', bottom=[sum(i) for i in zip(m2m_ratio, m2l_ratio, l2l_ratio, l2p_ratio,p2p_ratio)])
ax.bar(x_pos, m2m_ratio, color='0.50', label='M2M', bottom=[sum(i) for i in zip(m2l_ratio, l2l_ratio, l2p_ratio,p2p_ratio)])
ax.bar(x_pos, m2l_ratio, color='0.40',  label='M2L', bottom=[sum(i) for i in zip(l2l_ratio, l2p_ratio,p2p_ratio)])
ax.bar(x_pos, l2l_ratio, color='0.25', label='L2L', bottom=[sum(i) for i in zip(l2p_ratio,p2p_ratio)])
ax.bar(x_pos, l2p_ratio, color='0.15', label='L2P', bottom=p2p_ratio)
ax.bar(x_pos, p2p_ratio, color='0', label='P2P')
ax.set_xticks(x_pos)
ax.set_xticklabels(systems)
ax.set_ylabel("MLFMM Kernel Breakdown")
handles, labels = ax.get_legend_handles_labels()
ax.legend(handles, labels)
# plt.ylim([1, 1e4])
# ax.set_title('How fast do you want to go today?')
sns.despine(trim=True)
legend = plt.legend(frameon = 1)
frame = legend.get_frame()
# frame.set_facecolor('white')
# frame.set_edgecolor('white')
frame.set_alpha(0.9)
plt.savefig(path+'.pdf', bbox_inches=BBOX_INCHES)
plt.savefig(path+'.png', dpi=DPI, bbox_inches=BBOX_INCHES)
print path

path='figures/mlfmm_bw'
width=0.33
fig, ax = plt.subplots()
systems = ('1T\n(1 XE)', "32T\n(1 XE)", "1 GPU\n(1 XK)" ,"4 GPU\n(4 XK)", "16 GPU\n(16 XK)")
mlfmm = (1.50e6, 8.64e4, 2.783779e4, 7.01e3, 1.89e3)
num = (45,45,45,45,47)
x_pos = np.arange(len(systems))
rects = ax.bar([p-width/2 for p in x_pos], [i/j for i,j in zip(mlfmm,num)], color='darkgray', log=True, width=width)
autolabel(ax, rects[0])
autolabel(ax, rects[1])
autolabel(ax, rects[2])
autolabel(ax, rects[3])
autolabel(ax, rects[4])
ax.set_ylim(1,8e4)
ax2=ax.twinx()
rects = ax2.bar([p+width/2 for p in x_pos], [mlfmm[0] / i for i in mlfmm], color='lightgray', log=True, width=width)
autolabel(ax2, rects[1])
autolabel(ax2, rects[2])
autolabel(ax2, rects[3])
autolabel(ax2, rects[4])
ax2.set_ylim(1,1.5e3)
ax.set_xticks(x_pos)
ax.set_xticklabels(systems)
ax.set_ylabel("Per-MLFMM Execution Time (ms)")
ax2.set_ylabel("Speedup over Sequential")
plt.savefig(path+'.pdf', bbox_inches=BBOX_INCHES)
plt.savefig(path+'.png', dpi=DPI, bbox_inches=BBOX_INCHES)
print path



path="figures/mlfmm_minsky"
width = 0.33
fig, ax = plt.subplots()
systems = ('1T', "160T", "1 GPU" ,"4 GPU")
mlfmm = [1.25e6, 4.84e4, 5.22e3, 1.29e3]
num = (44,44,44,44)
x_pos = np.arange(len(systems))
rects = ax.bar([p-width/2 for p in x_pos], [i/j for i,j in zip(mlfmm,num)], color='darkgray', log=True, width=width)
autolabel(ax, rects[0])
autolabel(ax, rects[1])
autolabel(ax, rects[2])
autolabel(ax, rects[3])
ax.set_ylim(1,8e4)
ax2=ax.twinx()
rects = ax2.bar([p+width/2 for p in x_pos], [mlfmm[0] / i for i in mlfmm], color='lightgray', log=True, width=width)
autolabel(ax2, rects[1])
autolabel(ax2, rects[2])
autolabel(ax2, rects[3])
ax2.set_ylim(1,1.5e3)
ax.set_xticks(x_pos)
ax.set_xticklabels(systems)
ax.set_ylabel("Per-MLFMM Execution Time (ms)")
ax2.set_ylabel("Speedup over Sequential")
plt.savefig(path+'.pdf', bbox_inches=BBOX_INCHES)
plt.savefig(path+'.png', dpi=DPI, bbox_inches=BBOX_INCHES)
print path



path="figures/mlfmm"
width=0.33
WALLTIME="0.8"
SPEEDUP="0.5"
fig, axes = plt.subplots(2,1)
ax=axes[0]
ax.set_title("(a)")
systems = ('1T\n(1 XE)', "32T\n(1 XE)", "1 GPU\n(1 XK)" ,"4 GPU\n(4 XK)", "16 GPU\n(16 XK)")
mlfmm = (1.50e6, 8.64e4, 2.783779e4, 7.01e3, 1.89e3)
num = (45,45,45,45,47)
x_pos = np.arange(len(systems))
rects = ax.bar([p-width/2 for p in x_pos], [i/j for i,j in zip(mlfmm,num)], color=WALLTIME, log=True, width=width)
autolabel(ax, rects[0])
autolabel(ax, rects[1])
autolabel(ax, rects[2])
autolabel(ax, rects[3])
autolabel(ax, rects[4])
ax.set_ylim(1,2e5)

ax2=ax.twinx()
rects = ax2.bar([p+width/2 for p in x_pos], [mlfmm[0] / i for i in mlfmm], color=SPEEDUP, log=True, width=width)
autolabel(ax2, rects[1])
autolabel(ax2, rects[2])
autolabel(ax2, rects[3])
autolabel(ax2, rects[4])
ax2.set_ylim(1,5e3)
ax.set_xticks(x_pos)
ax.set_xticklabels(systems)
ax=axes[1]
ax.set_title("(b)")
systems = ('1T\n(1 S822LC)', "160T\n(1 S822LC)", "1 GPU\n(1 S822LC)" ,"4 GPU\n(1 S822LC)")
mlfmm = [1.25e6, 4.84e4, 5.22e3, 1.29e3]
num = (44,44,44,44)
x_pos = np.arange(len(systems))
rects = ax.bar([p-width/2 for p in x_pos], [i/j for i,j in zip(mlfmm,num)], color=WALLTIME, log=True, width=width)
autolabel(ax, rects[0])
autolabel(ax, rects[1])
autolabel(ax, rects[2])
autolabel(ax, rects[3])
ax.set_ylim(1,2e5)
ax2=ax.twinx()
rects = ax2.bar([p+width/2 for p in x_pos], [mlfmm[0] / i for i in mlfmm], color=SPEEDUP, log=True, width=width)
autolabel(ax2, rects[1])
autolabel(ax2, rects[2])
autolabel(ax2, rects[3])
ax2.set_ylim(1,5e3)
ax.set_xticks(x_pos)
ax.set_xticklabels(systems)
fig.text(-0.01, 0.5, "Per-MLFMM Execution Time (ms)", va='center', rotation='vertical')
fig.text(0.99, 0.5, "Speedup over Sequential", va='center', rotation='vertical')
fig.tight_layout()
plt.savefig(path+'.pdf', bbox_inches=BBOX_INCHES)
plt.savefig(path+'.png', dpi=DPI, bbox_inches=BBOX_INCHES)
print path