#!/usr/bin/python3

import time
import threading
import subprocess
import re
import sys
import os
import pickle
import gzip
import argparse

PROC_STAT_HEADERS = ('user', 'nice', 'system', 'idle', 'iowait', 'irq', 'softirq', 'steal', 'guest', 'guest_nice')
PROC_PID_STAT_HEADERS = ('pid', 'comm', 'state', 'ppid', 'pgrp', 'session', 'tty_nr', 'tpgid', 'flags', 'minflt',
                         'cminflt', 'majflt', 'cmajflt', 'utime', 'stime', 'cutime', 'cstime', 'priority', 'nice',
                         'num_threads', 'itrealvalue', 'starttime', 'vsize', 'rss', 'rsslim', 'startcode',
                         'endcode', 'startstack', 'kstkesp', 'kstkeip', 'signal', 'blocked', 'sigignore',
                         'sigcatch', 'wchan', 'nswap', 'cnswap', 'exit_signal', 'processor', 'rt_priority',
                         'policy', 'delayacct_blkio_ticks', 'guest_time', 'cguest_time', 'start_data', 'end_data',
                         'start_brk', 'arg_start', 'arg_end', 'env_start', 'env_end', 'exit_code')

pid_stat_regex = re.compile('(\(.*\)|\S+)', flags=re.DOTALL)

def read_proc_stat_file(f):
	f.seek(0)

	cpu_lines = []
	found_cpu_line = False
	for line in f:
		# only read as much of the file as we need
		if len(line) == 0:
			continue
		#
		if line[0:3] == 'cpu':
			cpu_lines.append(line)
			found_cpu_line = True
		elif found_cpu_line:
			break
		#
	#

	return cpu_lines
#

def parse_proc_stat_file(cpu_lines):
	stats = {'system': None, 'cpus': {}}

	for l in cpu_lines:
		l_split = l.split()
		cpu_index = int(l_split[0][3:]) if len(l_split[0][3:]) != 0 else None
		cpu_stats = {x[0]: int(x[1]) for x in zip(PROC_STAT_HEADERS, l_split[1:])}

		if cpu_index == None:
			stats['system'] = cpu_stats
		else:
			stats['cpus'][cpu_index] = cpu_stats
		#
	#

	return stats
#

def read_proc_pid_stat_file(f):
	f.seek(0)
	return f.read()
#

def parse_proc_pid_stat_file(contents):
	raw_stats = pid_stat_regex.findall(contents)
	proc_stats = {x[0]: x[1] for x in zip(PROC_PID_STAT_HEADERS, raw_stats)}

	for k in proc_stats:
		if k != 'comm' and k != 'state':
			proc_stats[k] = int(proc_stats[k])
		#
	#

	return proc_stats
#

def calculate_cpu_idle_stats(stats):
	idle = stats['idle'] + stats['iowait']
	non_idle = stats['user'] + stats['nice'] + stats['system'] + stats['irq'] + stats['softirq'] + stats['steal']

	return (idle, non_idle)
#

def calculate_core_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_non_idle) = calculate_cpu_idle_stats(initial)
	initial_total = initial_idle + initial_non_idle

	(current_idle, current_non_idle) = calculate_cpu_idle_stats(current)
	current_total = current_idle + current_non_idle

	clock_ticks = current_non_idle-initial_non_idle
	fraction = clock_ticks/(current_total-initial_total)
	return (clock_ticks, fraction)
#

def calculate_process_cpu_usage(process_initial, process_current, cpu_initial, cpu_current):
	process_initial_non_idle = process_initial['utime'] + process_initial['stime']
	process_current_non_idle = process_current['utime'] + process_current['stime']

	used_cores = (process_initial['processor'], process_current['processor'])
	core_totals = []
	for core in used_cores:
		(initial_idle, initial_non_idle) = calculate_cpu_idle_stats(cpu_initial[core])
		initial_total = initial_idle + initial_non_idle

		(current_idle, current_non_idle) = calculate_cpu_idle_stats(cpu_current[core])
		current_total = current_idle + current_non_idle

		core_totals.append(current_total-initial_total)
	#

	clock_ticks = process_current_non_idle-process_initial_non_idle
	fraction = clock_ticks/(sum(x**2 for x in core_totals)/sum(core_totals))
	return (clock_ticks, fraction)
#

def calculate_core_cpu_usage_continuous(stats):
	ticks = []
	fractions = []

	for i in range(len(stats)-1):
		(clock_ticks, fraction) = calculate_core_cpu_usage(stats[i], stats[i+1])
		ticks.append(clock_ticks)
		fractions.append(fraction)
	#

	return {'ticks': ticks, 'fractions': fractions}
#

def calculate_process_cpu_usage_continuous(process_stats, cpu_stats):
	ticks = []
	fractions = []

	assert all([len(process_stats) == len(cpu_stats[i]) for i in cpu_stats])

	for i in range(len(process_stats)-1):
		(clock_ticks, fraction) = calculate_process_cpu_usage(process_stats[i], process_stats[i+1], {core: cpu_stats[core][i] for core in cpu_stats}, {core: cpu_stats[core][i+1] for core in cpu_stats})
		ticks.append(clock_ticks)
		fractions.append(fraction)
	#

	return {'ticks': ticks, 'fractions': fractions}
#

def log_cpu_stats(path, interval, pids, stop_event):
	pids = sorted([int(pid) for pid in pids])
	tids = {pid: sorted([int(tid) for tid in os.listdir('/proc/{}/task'.format(pid))]) for pid in pids}

	stat_file = open('/proc/stat', 'r')
	pid_files = {pid: open('/proc/{}/stat'.format(pid), 'r') for pid in pids}
	tid_files = {pid: {tid: open('/proc/{}/task/{}/stat'.format(pid, tid), 'r') for tid in tids[pid]} for pid in pids}

	raw_stats = {'timestamps': [],
	             'timestamps_finished': [],
	             'system': [],
	             'process': {x: {'pid': [],
	                             'tid': {y: [] for y in tid_files[x]}} for x in pid_files}}

	# begin collecting data

	while not stop_event.is_set():
		start_time = time.time()
		raw_stats['timestamps'].append(start_time)

		t_0 = time.time()
		contents = read_proc_stat_file(stat_file)
		t_1 = time.time()
		raw_stats['system'].append((contents, t_1, t_1-t_0))

		for pid in pids:
			t_0 = time.time()
			contents = read_proc_pid_stat_file(pid_files[pid])
			t_1 = time.time()
			raw_stats['process'][pid]['pid'].append((contents, t_1, t_1-t_0))
			for tid in tids[pid]:
				t_0 = time.time()
				contents = read_proc_pid_stat_file(tid_files[pid][tid])
				t_1 = time.time()
				raw_stats['process'][pid]['tid'][tid].append((contents, t_1, t_1-t_0))
			#
		#

		finished_time = time.time()
		raw_stats['timestamps_finished'].append(finished_time)

		wait_time = max(0, interval-(time.time()-finished_time))
		stop_event.wait(wait_time)
	#

	# begin formatting data

	stats = {'timestamps': raw_stats['timestamps'],
	         'timestamps_finished': raw_stats['timestamps_finished'],
	         'cpu':{'system': [],
	                'id': {cpu: [] for cpu in parse_proc_stat_file(raw_stats['system'][0][0])['cpus'].keys()}},
	         'process': {pid: {'pid': [],
	                           'tid': {tid: [] for tid in tids[pid]}} for pid in pids}}

	for x in range(len(raw_stats['timestamps'])):
		current_stats = parse_proc_stat_file(raw_stats['system'][x][0])

		system_stats = current_stats['system']
		system_stats['read_time'] = raw_stats['system'][x][1]
		system_stats['read_duration'] = raw_stats['system'][x][2]
		stats['cpu']['system'].append(system_stats)

		for cpu in current_stats['cpus']:
			stats['cpu']['id'][cpu].append(current_stats['cpus'][cpu])
		#

		for pid in pids:
			pid_stats = parse_proc_pid_stat_file(raw_stats['process'][pid]['pid'][x][0])
			pid_stats['read_time'] = raw_stats['process'][pid]['pid'][x][1]
			pid_stats['read_duration'] = raw_stats['process'][pid]['pid'][x][2]
			stats['process'][pid]['pid'].append(pid_stats)

			for tid in tids[pid]:
				tid_stats = parse_proc_pid_stat_file(raw_stats['process'][pid]['tid'][tid][x][0])
				tid_stats['read_time'] = raw_stats['process'][pid]['tid'][tid][x][1]
				tid_stats['read_duration'] = raw_stats['process'][pid]['tid'][tid][x][2]
				stats['process'][pid]['tid'][tid].append(tid_stats)
			#
		#
	#

	with gzip.GzipFile(path, 'wb') as f:
		pickle.dump(stats, f, protocol=4)
	#
#

def load_cpu_stats(path):
	with gzip.GzipFile(path, 'rb') as f:
		return pickle.load(f)
	#
#

if __name__ == '__main__':
	stop_event = threading.Event()

	parser = argparse.ArgumentParser(description='Log CPU usage data and save as a gzipped pickle file.',
	                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
	parser.add_argument('--interval', type=float, default=1, help='interval for data collection')
	parser.add_argument('--pids', type=str, help='comma-separated list of processes to log')
	parser.add_argument('file_out', metavar='file-out', type=str, help='where to save the data')
	args = parser.parse_args()

	if args.pids != None:
		pids = [int(pid) for pid in args.pids.split(',')]
	else:
		pids = []
	#

	t = threading.Thread(target=log_cpu_stats, args=(args.file_out, args.interval, pids, stop_event))
	t.start()

	try:
		while t.is_alive():
			t.join(timeout=100)
		#
	except KeyboardInterrupt:
		stop_event.set()
		print()
	#

	t.join()
#