#!/usr/bin/env python
import numpy as np
import matplotlib.mlab as mlab
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import seaborn
import os

seaborn.set()

# Helper function to attach labels to the bars
def autolabel(subplot, bars):
    for bar in bars:
        height = round(bar.get_height(), 1)
        subplot.text(bar.get_x() + bar.get_width()/2., 1.05*height, '%s' % height, ha='center', va='bottom')

# Helper function to strip away the "opencl/" prefix and other stuff from all the names
def stripNames (category, names):
    stripped_names = []
    for elem in names:
        elem = elem.replace("opencl/", "")
        elem = elem.replace("/ocl", "")
        elem = elem.replace("/OpenCL", "")
        stripped_names.append(elem)
    return stripped_names

# List containing all the platform names
platforms = ['cpu', 'gpu-primary', 'gpu-secondary']

# List containing the categories of the benchamrks
categories = ['short', 'long']

averageTotal = {}
stderrTotal = {}

# Import data from file
for platform in platforms:
    averageTotal[platform] = np.genfromtxt('results/' + platform + '/average.csv', dtype=None, delimiter=',', names=['name', 'time', 'power'])
    stderrTotal[platform] = np.genfromtxt('results/' + platform + '/stderr.csv', dtype=None, delimiter=',', names=['name', 'time', 'power'])

# Create the folder where to store the charts, if it doesn't already exist
if not os.path.exists('charts'):
    os.makedirs('charts')

# Dict that contains the indexes(entry in the data files) of the benchmarks taking into account their category
indexes = {'short': [], 'long': []}

# We cycle on the benchmarks results for the cpu (since usually they are the ones with the longest run time) and divide in two groups the benchmarks, putting the relative indexes in two lists
for i in range(0,17):
    record = averageTotal['cpu'].take(i)
    if record['time'] < 30:
        indexes['short'].append(i)
    else:
        indexes['long'].append(i)


# We instantiate a new dict to contain the average and the std err of the measurements on the various platforms
average = {'cpu': {}, 'gpu-primary': {}, 'gpu-secondary': {}}
stderr = {'cpu': {}, 'gpu-primary': {}, 'gpu-secondary': {}}

for platform in platforms:
    for category in categories:
        average[platform][category] = averageTotal[platform].take(indexes[category])
        stderr[platform][category] = stderrTotal[platform].take(indexes[category])

# Generate an array as placeholder for the x axis (we need to pass from a list to an array to take advantage of range)
x = {}
x['short'] = np.arange(len(indexes['short']))
x['long'] = np.arange(len(indexes['long']))

# Strip from the names of the benchmarks unwanted parts, taking advantage of the helper function defined before
x_names_stripped = {'short': [], 'long': []}
for category in categories:
    x_names = averageTotal['cpu'].take(indexes[category])['name'].tolist()
    x_names_stripped[category] = stripNames(category, x_names)

# Helper function that create the plots of the execution time, takes as parameter the category of the benchmark we want to plot
def plotTimes (category):

    # Initialization of the figures
    fig, ax = plt.subplots(figsize=(20, 10))

    # Create the bar plot for the time values
    time_cpu_bars = ax.bar(x[category]-0.3, average['cpu'][category]['time'], width=0.2, color='b', align='edge', yerr=stderr['cpu'][category]['time'])
    time_gpu_primary_bars = ax.bar(x[category], average['gpu-primary'][category]['time'], width=0.2, color='r', align='center', yerr=stderr['gpu-primary'][category]['time'])
    time_gpu_secondary_bars = ax.bar(x[category]+0.3, average['gpu-secondary'][category]['time'], width=-0.2, color='g', align='edge', yerr=stderr['gpu-secondary'][category]['time'])

    # Change the labels of the x axis to contain the names of the benchmarks
    ax.set_xticks(x[category])
    ax.set_xticklabels(x_names_stripped[category])
    ax.set_title('Execution time of the various benchmarks expressed in seconds')
    ax.set_xlabel('Benchmark')
    ax.set_ylabel('seconds')

    # Add some patches as legend of the colors used for the various benchmarks
    red_patch = mpatches.Patch(color='blue', label='Execution time for cpu')
    blue_patch = mpatches.Patch(color='red', label='Execution time for gpu(4 core)')
    green_patch = mpatches.Patch(color='green', label='Execution time for gpu(2 core)')

    ax.legend(handles=[red_patch, blue_patch, green_patch])

    # Invoke the helper function to attach labels for each bar
    autolabel(ax, time_cpu_bars)
    autolabel(ax, time_gpu_primary_bars)
    autolabel(ax, time_gpu_secondary_bars)

    # Save the obtained plot on file
    plt.savefig('charts/times-' + category + '.pdf')

# Helper function that create the plots of the execution time, takes as parameter the category of the benchmark we want to plot
def plotPower (category):

    # Create a new figure
    fig, bx = plt.subplots(figsize=(20, 10))

    # Create the bar plot for the power values
    power_cpu_bars = bx.bar(x[category]-0.3, average['cpu'][category]['power'], width=0.2, color='b', align='edge', yerr=stderr['cpu'][category]['power'])
    power_gpu_primary_bars = bx.bar(x[category], average['gpu-primary'][category]['power'], width=0.2, color='r', align='center', yerr=stderr['gpu-primary'][category]['power'])
    power_gpu_secondary_bars = bx.bar(x[category]+0.3, average['gpu-secondary'][category]['power'], width=-0.2, color='g', align='edge', yerr=stderr['gpu-secondary'][category]['power'])

    # Change the labels of the x axis to contain the names of the benchmarks
    bx.set_xticks(x[category])
    bx.set_xticklabels(x_names_stripped[category])
    bx.set_title('Power consumption of the various benchmarks expressed in mWatt/hour')
    bx.set_xlabel('Benchmark')
    bx.set_ylabel('mWatt/hour')

    # Add some patches as legend of the colors used for the various benchmarks
    red_patch = mpatches.Patch(color='blue', label='Power consumption for cpu')
    blue_patch = mpatches.Patch(color='red', label='Power consumption for gpu(4 core)')
    green_patch = mpatches.Patch(color='green', label='Execution time for gpu(2 core)')

    bx.legend(handles=[red_patch, blue_patch, green_patch])

    # Invoke the helper function to attach labels one for each bar
    autolabel(bx, power_cpu_bars)
    autolabel(bx, power_gpu_primary_bars)
    autolabel(bx, power_gpu_secondary_bars)

    # Save the obtained plot on file
    plt.savefig('charts/power-' + category + '.pdf')

# Invoke the helper functions to plot the data
for category in categories:
    plotTimes(category)
    plotPower(category)