#!/usr/bin/env python3
#
import sys
import json
import gzip
import numpy as np
import matplotlib.pylab as plt
#
def load_cpu_stats(path):
    with gzip.GzipFile(path, 'r') as f:
        return json.load(f)
    #
#
def calculate_cpu_usage(initial, current):
    """
    Calculation adapted from: https://stackoverflow.com/questions/23367857/accurate-calculation-of-cpu-usage-given-in-percentage-in-linux/
    """
    #
    initial_idle = initial['idle'] + initial['iowait']
    current_idle = current['idle'] + current['iowait']
    #
    initial_non_idle = initial['user'] + initial['nice'] + initial['system'] + initial['irq'] + initial['softirq'] + initial['steal']
    current_non_idle = current['user'] + current['nice'] + current['system'] + current['irq'] + current['softirq'] + current['steal']
    #
    initial_total = initial_idle + initial_non_idle
    current_total = current_idle + current_non_idle
    #
    return (current_non_idle-initial_non_idle)/(current_total-initial_total)
#
def calculate_cpu_usage_continuous(stats):
	cpu_usages = []
	for i in range(len(stats)-1):
		cpu_usages.append(calculate_cpu_usage(stats[i], stats[i+1]))
	#
	return cpu_usages
#
def parse_range_list(range_list_str):
        '''
        Take an input like '1-3,5,7-10' and return a list like [1,2,3,5,7,8,9,10].
        '''
        #
        range_strings = range_list_str.split(',')
        all_items = []
        for range_str in range_strings:
                if '-' in range_str:
                        # in the form '12-34'
                        range_ends = [int(x) for x in range_str.split('-')]
                        assert(len(range_ends) == 2)
                        all_items.extend(range(range_ends[0], range_ends[1]+1))
                else:
                        # just a number
                        all_items.append(int(range_str))
                #
        #
        return all_items
#
cpu_data = load_cpu_stats(sys.argv[1])
#
unused_cpus = [str(x) for x in parse_range_list('128-129,138-139,18-19,148-149,28-29,158-159,38-39,48-49,58-59,68-69,78-79,98-99,108-109,118-119')]
numa_sets = [[0, 80, 1, 81], [10, 90, 11, 91], [20, 100, 21, 101], [30, 110, 31, 111], [40, 120, 41, 121], [50, 130, 51, 131], [60, 140, 61, 141], [70, 150, 71, 151], [2, 82, 3, 83], [12, 92, 13, 93], [22, 102, 23, 103], [32, 112, 33, 113], [42, 122, 43, 123], [52, 132, 53, 133], [62, 142, 63, 143], [72, 152, 73, 153], [4, 84, 5, 85], [14, 94, 15, 95], [24, 104, 25, 105], [34, 114, 35, 115], [44, 124, 45, 125], [54, 134, 55, 135], [64, 144, 65, 145], [74, 154, 75, 155], [6, 86, 7, 87], [16, 96, 17, 97], [26, 106, 27, 107], [36, 116, 37, 117], [46, 126, 47, 127], [56, 136, 57, 137], [66, 146, 67, 147], [76, 156, 77, 157], [8, 88, 9, 89]]
#
timestamps = (np.array(cpu_data['timestamps'][1:]) + np.array(cpu_data['timestamps'][:-1])) / 2
cpu_usages = {int(cpu): np.array(calculate_cpu_usage_continuous(cpu_data['stats']['cpus'][cpu])) for cpu in cpu_data['stats']['cpus'] if cpu not in unused_cpus}
tor_usages = [np.sum([cpu_usages[cpu] for cpu in x], axis=0)/4 for x in numa_sets]
print(len(cpu_usages))
print(len(tor_usages))
#
plt.figure()
#
for cpu in tor_usages:
	plt.plot(timestamps, cpu*100)
#
plt.xlabel('Time (s)')
plt.ylabel('CPU Usage (Average over 4 cores)')
plt.show()