make_plots

#! /usr/bin/env ruby

require 'optparse'
require 'json'
require 'erb'
require 'command_line_reporter'
require 'colorize'
require 'descriptive-statistics'

def error(message)
	STDERR.puts message.red
end

# Parse the command line arguments
optparse = OptionParser.new do |opts|
	opts.banner = "Usage: {$0} <result dir>..."

	opts.on('-h', '--help', 'Display this screen') do
		puts opts
		exit
	end
end
optparse.parse!

# A Plotter takes all experiment data in an Execo result dir and
# plots various results.
# A new Plotter is necessary for each result dir
class Plotter
	include CommandLineReporter
	
	attr_reader :computes, :title, :x_label, :y_label, :nodes, :roles, :range

	def initialize(dir)
		@result_dir		= dir
		begin
			json_path = File.join(dir, 'experiment.json')
			@experiment		= JSON.load(File.read(json_path))
			@experiment['nodes']['services'].each {
				|r, s| @experiment['nodes']['services'][r] = s[/\w+\-\d+/]
			}
			@experiment['nodes']['computes'].map!   { |c| c[/\w+\-\d+/] }
		rescue => details
			error("Could not load #{json_path}: #{details.to_s}")
		end

		# Store the hosts and their role
		@roles = @experiment['nodes']['services'].keys \
			+ @experiment['nodes']['computes']
		@nodes = @experiment['nodes']['services'].values \
			+ @experiment['nodes']['computes']
		@computes = @experiment['nodes']['computes']

		# Read the energy data once now
		@energy_data = {}

		Dir.foreach(File.join(dir, 'energy')) do |filename|
			next if ['.', '..'].include? filename

			host = filename[/(\w+\-\d+)\.dat/, 1]
			File.open(File.join(dir, 'energy', filename)) do |file|
				line = file.gets # header
				while line = file.gets do
					values = line.split ' '
					t = values[1].to_i

					@energy_data[t] = {} if @energy_data[t].nil?
					@energy_data[t][host] = values[2].to_f
				end
			end
		end

		if @energy_data.empty? then
			url = "https://intranet.grid5000.fr/supervision/lyon/wattmetre/formulaire.php"
			error("No data in #{File.join(dir, 'energy')}")
			error("Please download them from #{url}")
			exit(1)
		end

		puts "Read #{@energy_data.size} values for " \
			+ "#{@energy_data.first[1].keys.size} machines"

		@source			= nil
		@title			= nil
		@x_label			= nil
		@y_label			= nil
		@range			= nil

		@averages		= {}
		@stats			= {}
	end

	def role2host(role)
		if !@experiment['nodes']['services'][role].nil? then
			@experiment['nodes']['services'][role]
		else
			role
		end
	end

	# Write the headers for a CVS file
	#
	#  == Parameters:
	#  file: the open file to write to
	def write_headers(file)
		# Write the header
		file.write "Time\t"
		@roles.each { |n| file.write "\t,#{n}"}
		file.write "\n"
	end

	def error?(xp)
		!xp['error'].nil?
	end

	# Walks into a directory and looks for JSON files to extract
	# energy data.
	#
	# == Parameters:
	def walk()
		[ 'plots', 'scripts', 'csv' ].each do |dir|
			dir = File.join(@result_dir, dir)
			Dir.mkdir dir if !File.exist? dir
		end

		# Loop through all the Rally tasks, analyze those that didn't fail.
		benchmarks = @experiment['benchmarks']
		benchmarks.each_key do |bench_json|
			bench = benchmarks[bench_json]
			bench_name = File.basename(bench_json, '.json')

			if bench['error'].nil? then
				if generate_data(bench_name, bench) then
					generate_plot(bench_name, bench, 'raw')
					generate_plot(bench_name, bench, 'corrected')
				end
			else
				error("Task #{bench_json} failed")
			end
		end
	end

	# Extract metrics from the large CSV file and writes them to a smaller, task
	# specific CSV file.
	#
	# == Parameters:
	# filename::
	#   the basename of file to extract metrics from
	def generate_data(bench_name, bench)
		# Quickly check we have the energy metrics for this task
		if bench['idle_start'] < @energy_data.keys.first \
				or bench['idle_end'] > @energy_data.keys.last then
			error("The energy metrics do not cover the whole period of #{bench_name} task")
			return false
		end

		# This hash will contain all the raw indexed by timestamp
		raw = {}
		corrected = {}

		bench_energy = @energy_data.select {
			|k, v| (k > bench['idle_start']) and (k < bench['idle_end'])
		}
		@averages[bench_name] = {}
		@stats[bench_name]    = {}
		
		@nodes.each do |host|
			before = []
			after  = []

			# Get the raw series
			bench_energy.each_key do |timestamp|
				t = timestamp.to_i
				s = t - bench['idle_start']

				raw[s] = {} if raw[s].nil?
				raw[s][host] = bench_energy[timestamp][host]

				# Record the baseline before and after
				if t < bench['run_start'] then
					before << bench_energy[timestamp][host]
				end

				if t > bench['run_end'] then
					after << bench_energy[timestamp][host]
				end
			end

			# Compute some statistics
			@stats[bench_name][host] = {} if @stats[bench_name][host].nil?
			@stats[bench_name][host][:avg] = {}
			@stats[bench_name][host][:stddev] = {}

			stats = DescriptiveStatistics::Stats.new before
			@stats[bench_name][host][:avg][:before]		= stats.mean
			@stats[bench_name][host][:stddev][:before]	= stats.standard_deviation

			stats = DescriptiveStatistics::Stats.new after
			@stats[bench_name][host][:avg][:after]		= stats.mean
			@stats[bench_name][host][:stddev][:after]	= stats.standard_deviation

			stats = DescriptiveStatistics::Stats.new before + after
			@stats[bench_name][host][:avg][:all]		= stats.mean
			@stats[bench_name][host][:stddev][:all]	= stats.standard_deviation

			# Compute the corrected series by substracting
			# the average idle to the raw
			bench_energy.each_key do |timestamp|
				s = timestamp - bench['idle_start']

				corrected[s] = {} if corrected[s].nil?
				corrected[s][host] = bench_energy[timestamp][host] \
					- @stats[bench_name][host][:avg][:before]
			end
		end

		# Write the raw values to a first file
		name = bench_name + '_raw.csv'
		path = File.join(@result_dir, 'csv', name)
		file = File.new(path, 'w')
		write_headers(file)

		raw.each do |t, row|
			file.write t.to_s
			@roles.each { |role| file.write ",\t#{row[role2host(role)]}" }
			file.write "\n"
		end
		file.close()
		puts "Wrote #{path.green}"

		# Write the corrected values to another file
		name = bench_name + '_corrected.csv'
		path = File.join(@result_dir, 'csv', name)
		file = File.new(path, 'w')
		write_headers(file)

		corrected.each do |t, row|
			file.write t.to_s
			@roles.each { |role| file.write ",\t#{row[role2host(role)]}" }
			file.write "\n"
		end

		file.close
		puts "Wrote #{path.green}"

		true
	end

	def generate_plot(bench_name, bench, what)
		# Let's write a R script
		renderer = ERB.new(File.read('templates/energy.erb'))

		# Input/output
		@source = File.join(@result_dir, 'csv', "#{bench_name}_#{what}.csv")
		@output = File.join(@result_dir, 'plots', "#{bench_name}_#{what}.pdf")

		# Values
		@title	= bench_name
		@x_label	= 'Time (sec)'
		@y_label	= 'Energy (watt)'

		# Range
		if what == 'raw' then
			@range = [85, 125]
		else
			#@range = [-15, 30]
		end

		# Straight lines to show start/end
		@idle_start	= 0
		@idle_end	= bench['idle_end']	- bench['idle_start']
		@run_start	= bench['run_start']	- bench['idle_start']
		@run_end		= bench['run_end']	- bench['idle_start']

		# Write the script
		dir = File.join(@result_dir, 'scripts')

		script = renderer.result(binding)
		path = File.join(dir, "#{bench_name}_#{what}.plt")
		file = File.new(path, 'w')
		file.write(script)
		file.close

		# Run it
		if system "/usr/bin/env gnuplot #{file.path}" then
			puts "Executed #{path.green} successfully"
		else
			puts "Execution of #{path.green} failed"
		end
	end

	def print_averages(node)
		max_len = @experiment['benchmarks'].select {
			|name, xp| !error? xp }.keys.map {|b| File.basename(b).size
		}.reduce { |max, current|
			if current > max then
				current
			else
				max
			end
		}

		header :title => "Idle power comsumption (average and stddev) before and aftear each task, for #{node}", :align => 'center'
		table(:border => true) do
			row do
				column 'Benchmark', :width => max_len
				column 'Before', :align => 'right', :width => 14
				column 'After', :align => 'right', :width => 14
				column 'All', :align => 'right', :width => 14
			end
			@stats.each do |b, bench|
				avg_before		= bench[node][:avg][:before].round(2)
				avg_after		= bench[node][:avg][:after].round(2)
				avg_all			= bench[node][:avg][:all].round(2)
				stddev_before	= bench[node][:stddev][:before].round(2)
				stddev_after	= bench[node][:stddev][:after].round(2)
				stddev_all		= bench[node][:stddev][:all].round(2)

				row do
					column b, :align => 'left'
					column "#{avg_before} (±#{stddev_before})"
					column "#{avg_after} (±#{stddev_after})"
					column "#{avg_all} (±#{stddev_all})"
				end
			end
		end
	end
end

ARGV.each do |dir|
	plotter = Plotter.new(dir)
	plotter.walk
	#plotter.print_averages plotter.nodes[0]
end