#
#  Copyright (C) 2016 - 2018 Intel Corporation.
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#  1. Redistributions of source code must retain the above copyright notice(s),
#     this list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright notice(s),
#     this list of conditions and the following disclaimer in the documentation
#     and/or other materials provided with the distribution.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER(S) ``AS IS'' AND ANY EXPRESS
#  OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
#  MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO
#  EVENT SHALL THE COPYRIGHT HOLDER(S) BE LIABLE FOR ANY DIRECT, INDIRECT,
#  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
#  OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
#  ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
import numpy.ma as ma
import os
from shutil import rmtree

files = ('alloctest_hbw.txt', 'alloctest_glibc.txt', 'alloctest_tbb.txt', 'alloctest_pmem.txt')
legend = ('avg hbw', 'avg glibc', 'avg tbb', 'avg pmem', 'first operation')
colors = ('red', 'green', 'blue', 'black')
first_operation_color = 'yellow'

threads_values = ('1', '2', '4', '8', '16', '18', '36')

small_sizes_values = ('1', '4', '16')
medium_sizes_values = ('64', '256')
big_sizes_values = ('1024', '4096', '16384')
sizes_values = (small_sizes_values, medium_sizes_values, big_sizes_values)
sizes_display = ('small', 'medium', 'big')

operations = ('Allocation', 'Free', 'Total')

iterations = 1000
output_directory = './plots/'

threads_index      = np.arange(len(threads_values))
small_sizes_index  = np.arange(len(small_sizes_values))
medium_sizes_index = np.arange(len(medium_sizes_values))
big_sizes_index    = np.arange(len(big_sizes_values))
sizes_index = (small_sizes_index, medium_sizes_index, big_sizes_index)

# 3D and 2D plots needs different width, so that bars do not cover each other in 3D
bar_width_2D = 0.2
bar_width_3D = 0.1

# return times (total, allocation, free, first allocation, first free) as columns 2 - 6
def return_times(entry):
    return entry[:,2], entry[:,3], entry[:,4], entry[:,5], entry[:,6]

# initialize axis with labels, ticks and values
def init_axis(fig, suptitle, subplot_projection, x_label, x_ticks, x_values, y_label, y_ticks, y_values, z_label):
    assert fig is not None
    fig.clf()
    fig.suptitle(suptitle)
    if subplot_projection:
        ax = fig.add_subplot(111, projection=subplot_projection)
    else:
        ax = fig.add_subplot(111)
    assert ax is not None
    if x_label:
        ax.set_xlabel(x_label)
        if x_ticks is not None and x_values is not None:
            ax.set_xticks(x_ticks)
            ax.set_xticklabels(x_values)
    if y_label:
        ax.set_ylabel(y_label)
        if y_ticks is not None and y_values is not None:
            ax.set_yticks(y_ticks)
            ax.set_yticklabels(y_values)
    if z_label:
        ax.set_zlabel(z_label)
    return ax

def draw_bar(ax, show_first_operation, x, y, time, init_time, draw_color):
    assert ax is not None
    if y is None:
        if show_first_operation is True:
            mask_time = ma.where(time>=init_time)
            mask_init_time = ma.where(init_time>=time)
            ax.bar(x[mask_time], time[mask_time], width=bar_width, color=draw_color)
            ax.bar(x, init_time, width=bar_width, color=first_operation_color)
            ax.bar(x[mask_init_time], time[mask_init_time], width=bar_width, color=draw_color)
        else:
            ax.bar(x, time, width=bar_width, color=draw_color)
    else:
        if show_first_operation is True:
            mask_time = ma.where(time>=init_time)
            mask_init_time = ma.where(init_time>=time)
            ax.bar(x[mask_time], (time - init_time)[mask_time], y[mask_time], bottom=init_time[mask_time], zdir='y', width=bar_width, color=draw_color)
            ax.bar(x[mask_init_time], (init_time - time)[mask_init_time], y[mask_init_time], bottom=time[mask_init_time], zdir='y', width=bar_width, color=first_operation_color)
            ax.bar(x[mask_init_time], time[mask_init_time], y[mask_init_time], zdir='y', width=bar_width, color=draw_color)
            ax.bar(x[mask_time], init_time[mask_time], y[mask_time], zdir='y', width=bar_width, color=first_operation_color)
        else:
            ax.bar(x, time, y, zdir='y', width=bar_width, color=draw_color)

def save_plot(show_first_operation, subfolder, filename):
    if show_first_operation:
        path = os.path.join(output_directory, "with_first_operation")
    else:
        path = os.path.join(output_directory, "without_first_operation")
    if subfolder:
        path = os.path.join(path, subfolder)
    if not os.path.exists(path):
        os.mkdir(path)
    plt.savefig(os.path.join(path, filename))
    print "Saved file %s" % filename

def load_data_from_files():
    data = []
    # load all files into data
    for f in files:
        assert os.path.isfile(f) is True
        data.append(np.loadtxt(f, comments='#'))
    return data

def set_bar_width_and_offsets(requested_width):
    bar_width = requested_width
    offsets = (0, bar_width, 2*bar_width, 3*bar_width)
    return bar_width, offsets

if os.path.exists(output_directory):
    rmtree(output_directory)
os.mkdir(output_directory)

data = load_data_from_files()

fig = plt.figure()

########################################
# Draw 3D plots (time, sizes, threads) #
########################################

bar_width, offsets = set_bar_width_and_offsets(bar_width_3D)

# for each size range (small, medium, big)
for size_values, size_index, size_display in zip(sizes_values, sizes_index, sizes_display):
    # show plots with and without first operation
    for show_first_operation in (True, False):
        # for each operation (allocation, free, total)
        for operation in operations:
            # add bar_width to each element of size_index
            ax = init_axis(fig, "%s time of %s sizes (%s iterations)" % (operation, size_display, iterations), '3d',
                           'size [kB]', size_index + (bar_width,) * len(size_index), size_values,
                           'threads', threads_index, threads_values,
                           'time [ms]')
            legend_data = []
            # for each allocator (hbw, glibc, tbb, pmem)
            for entry, offset, draw_color in zip(data, offsets, colors):
                # remove all rows where column 1 (size) is not in size_values (current size range)
                entry = entry[np.in1d(entry[:,1], np.array(size_values).astype(np.int))]
                # convert column 0 (threads values) to thread index
                threads_col = np.array([threads_values.index(str(n)) for n in map(int, entry[:,0])])
                # convert column 1 (sizes values) to size index
                size_col = [size_values.index(str(n)) for n in map(int, entry[:,1])]
                # add offset to size index so that bars display near each other
                size_col = np.array(size_col) + offset
                total_time_col, alloc_time_col, free_time_col, first_alloc_time_col, first_free_time_col = return_times(entry)
                if operation == 'Allocation':
                    draw_bar(ax, show_first_operation, size_col, threads_col, alloc_time_col, first_alloc_time_col, draw_color)
                elif operation == 'Free':
                    draw_bar(ax, show_first_operation, size_col, threads_col, free_time_col, first_free_time_col, draw_color)
                elif operation == 'Total':
                    draw_bar(ax, show_first_operation, size_col, threads_col, total_time_col, first_alloc_time_col+first_free_time_col, draw_color)
                legend_data.append(plt.Rectangle((0, 0), 1, 1, fc=draw_color))
            if show_first_operation is True:
                legend_data.append(plt.Rectangle((0, 0), 1, 1, fc=first_operation_color))
            ax.legend(legend_data,legend,loc='best')
            plt.grid()
            save_plot(show_first_operation, None, "%s_time_of_%s_sizes_%s_iterations.png" % (operation, size_display, iterations))

################################################
# Draw 2D plots (time, sizes, constant thread) #
################################################

bar_width, offsets = set_bar_width_and_offsets(bar_width_2D)

# for each size range (small, medium, big)
for size_values, size_index, size_display in zip(sizes_values, sizes_index, sizes_display):
    # show plots with and without first operation
    for show_first_operation in (True, False):
        for thread in threads_values:
            # for each operation (allocation, free, total)
            for operation in operations:
                # add bar_width to each element of size_index
                ax = init_axis(fig, "%s time of %s sizes with %s threads (%s operations)" % (operation, size_display, thread, iterations), None,
                               'size [kB]', size_index + (bar_width,) * len(size_index), size_values,
                               'time [ms]', None, None,
                               None)
                legend_data = []
                # for each allocator (hbw, glibc, tbb, pmem)
                for entry, offset, draw_color in zip(data, offsets, colors):
                    # remove all rows where column 1 (size) is not in size_values (current size range)
                    entry = entry[np.in1d(entry[:,1], np.array(size_values).astype(np.int))]
                    # remove all rows where column 0 (threads) is not equal to currently analyzed thread value
                    entry = entry[entry[:,0] == int(thread)]
                    # convert column 0 (threads values) to thread index
                    threads_col = np.array([threads_values.index(str(n)) for n in map(int, entry[:,0])])
                    # convert column 1 (size values) to size index
                    size_col = [size_values.index(str(n)) for n in map(int, entry[:,1])]
                    # add offset to size index so that bars display near each other
                    size_col = np.array(size_col) + offset
                    total_time_col, alloc_time_col, free_time_col, first_alloc_time_col, first_free_time_col = return_times(entry)
                    if operation == 'Allocation':
                        draw_bar(ax, show_first_operation, size_col, None, alloc_time_col, first_alloc_time_col, draw_color)
                    elif operation == 'Free':
                        draw_bar(ax, show_first_operation, size_col, None, free_time_col, first_free_time_col, draw_color)
                    elif operation == 'Total':
                        draw_bar(ax, show_first_operation, size_col, None, total_time_col, first_alloc_time_col+first_free_time_col, draw_color)
                    legend_data.append(plt.Rectangle((0, 0), 1, 1, fc=draw_color))
                if show_first_operation is True:
                    legend_data.append(plt.Rectangle((0, 0), 1, 1, fc=first_operation_color))
                ax.legend(legend_data,legend,loc='best')
                plt.grid()
                save_plot(show_first_operation, "constant_thread", "%s_time_of_%s_sizes_with_%s_threads_%s_iterations.png" % (operation, size_display, thread, iterations))

################################################
# Draw 2D plots (time, threads, constant size) #
################################################

# for each size range (small, medium, big)
for size_values, size_index, size_display in zip(sizes_values, sizes_index, sizes_display):
    # show plots with and without first operation
    for show_first_operation in (True, False):
        for size in size_values:
            # for each operation (allocation, free, total)
            for operation in operations:
                # add bar_width to each element of threads_index
                ax = init_axis(fig, "%s time of %s kB (%s operations)" % (operation, size, iterations), None,
                               'threads', threads_index + (bar_width,) * len(threads_index), threads_values,
                               'time [ms]', None, None,
                               None)
                legend_data = []
                # for each allocator (hbw, glibc, tbb, pmem)
                for entry, offset, draw_color in zip(data, offsets, colors):
                    # remove all rows where column 1 (size) is not in size_values (current size range)
                    entry = entry[np.in1d(entry[:,1], np.array(size_values).astype(np.int))]
                    # remove all rows where column 1 (size) is not equal to currently analyzed size value
                    entry = entry[entry[:,1] == int(size)]
                    # convert column 0 (threads values) to thread index
                    threads_col = np.array([threads_values.index(str(n)) for n in map(int, entry[:,0])])
                    # add offset to thread index so that bars display near each other
                    threads_col = np.array(threads_col) + offset
                    # convert column 1 (size values) to size index
                    size_col = [size_values.index(str(n)) for n in map(int, entry[:,1])]
                    total_time_col, alloc_time_col, free_time_col, first_alloc_time_col, first_free_time_col = return_times(entry)
                    if operation == 'Allocation':
                        draw_bar(ax, show_first_operation, threads_col, None, alloc_time_col, first_alloc_time_col, draw_color)
                    elif operation == 'Free':
                        draw_bar(ax, show_first_operation, threads_col, None, free_time_col, first_free_time_col, draw_color)
                    elif operation == 'Total':
                        draw_bar(ax, show_first_operation, threads_col, None, total_time_col, first_alloc_time_col+first_free_time_col, draw_color)
                    legend_data.append(plt.Rectangle((0, 0), 1, 1, fc=draw_color))
                if show_first_operation is True:
                    legend_data.append(plt.Rectangle((0, 0), 1, 1, fc=first_operation_color))
                ax.legend(legend_data,legend,loc='best')
                plt.grid()
                save_plot(show_first_operation, "constant_size", "%s_time_of_%s_kB_%s_operations.png" % (operation, size, iterations))