#!/usr/bin/env python

import argparse
import numpy as np

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

def megabytes(x, pos):
    return '%1.f' % (x/(1024*1024))

def get_memory_utilization_figure(allocated, requested):
    memory_fig = plt.figure()
    memory_ax = memory_fig.add_subplot(111)

    memory_ax.plot(allocated, label='allocated memory', rasterized=True)
    memory_ax.plot(requested, label='requested memory', rasterized=True)

    memory_ax.set_ylim(ymin=0)

    memory_ax.set_xticks([])
    memory_ax.set_xlabel('Application timeline')

    memory_ax.yaxis.set_major_formatter(FuncFormatter(megabytes))
    memory_ax.get_yaxis().tick_left()

    return memory_fig

def get_memory_utilization_index(allocated, requested):
    memory_fig = plt.figure()
    memory_ax = memory_fig.add_subplot(111)

    memory_ax.plot(np.true_divide(requested, allocated), rasterized=True)

    memory_ax.set_xticks([])
    memory_ax.set_xlabel('Application timeline')

    memory_ax.get_yaxis().tick_left()
    memory_ax.set_ylabel('Memory utilization index')

    return memory_fig

if __name__=='__main__':
    parser = argparse.ArgumentParser(
            description="parse dmmlib's memory traces")
    parser.add_argument("stat_arrays_file", type=argparse.FileType('r'),
                    help="compressed statistics file")
    args = parser.parse_args()
    npzfile = np.load(args.stat_arrays_file)
    allocated = npzfile['allocated']
    requested = npzfile['requested']
    memory_fig = get_memory_utilization_figure(allocated, requested)
    memory_fig.savefig(
            '{}_mem_usage.pdf'
            .format(args.stat_arrays_file.name.split('.')[0]),
            dpi=150
            )
    memory_fig = get_memory_utilization_index(allocated, requested)
    memory_fig.savefig(
            '{}_mem_utilization_index.pdf'
            .format(args.stat_arrays_file.name.split('.')[0]),
            dpi=150
            )
    args.stat_arrays_file.close()