#!/usr/bin/env ruby

require 'benchmark'

PATH = ARGV[0]
if PATH == nil || File.file?(PATH)
    puts "Need a source directory to start the benchmark"
    exit(1)
end

DEVSHM = "/dev/shm"
if File.file?(DEVSHM)
    puts "No tmpfs at #{DEVSHM} available"
    exit(2)
end

TITLE = File.basename(PATH).sub('.', '')
N_MEMVALS = 10

# formats: { format_label => {:file_extension, :compression_command, :decompression_command}, ... }
formats = {
    'Lzop'  => {:extension => 'lzo', :compress => 'lzop -U', :decompress => 'lzop -Ud'},
    'Gzip'  => {:extension => 'gz',  :compress => 'gzip',    :decompress => 'gunzip'},
    'Bzip2' => {:extension => 'bz2', :compress => 'bzip2',   :decompress => 'bunzip2'},
    '7-Zip' => {:extension => '7z',
        :compress => '/usr/lib/p7zip/7zr a -r -mmt=off',
        :decompress => '/usr/lib/p7zip/7zr x -y -mmt=off'},
    'Xz'    => {:extension => 'xz',  :compress => 'xz',      :decompress => 'xz -d'},
}

# tee(str):
#
# Prints %str to the standard output and to the file #{TITLE}-summary.txt.
def tee(str = '')
    if @tee_file == nil
        @tee_file = File.open("#{TITLE}-summary.txt", 'w')
    end
    puts str
    @tee_file.puts str
end

# Array#n_values(n):
#
# Return an array with %n values. Two last elements are respectively the
# average value and the last value if the array has more than %n values,
# otherwise the last value is repeated till %n.
class Array
    def n_values(n)
        arr = self
        if arr.count > n
            tmp = arr[(n-1)..arr.count]
            last_val = tmp.reduce(:+) / tmp.count
            arr[n-2] = last_val
            arr[n-1] = tmp.last
        end
        arr = arr[0, n]
        if arr.count < n
            count = i = arr.count
            while i < n
                arr[i] = arr[count - 1]
                i = i + 1
            end
        end
        arr
    end
end

# bm_process(command):
#
# Returns the time and memory (as array) needed for %command to perform.
def bm_process(command)
    p = 0
    mem = Array.new

    t = Thread.new do
        while true do
            next if p == 0
            File.open("/proc/#{p}/status").each do |line|
                mem << line.sub(/.*:\s+/, '').to_i if line =~ /^VmRSS:.*/
            end
            sleep 0.5
        end
    end

    time = Benchmark.realtime do
        p = fork do
            exec(command)
        end
        Process.waitpid(p)
        t.exit
    end

    tee "\tCPU time (s): #{time}"
    if mem.count > 1
        tee "\tMax/Average RSS memory (kB): #{mem.max}/#{mem.drop(1).reduce(:+) / mem.drop(1).count}"
    else
        tee "\tRSS memory (kB): #{mem[0]}"
    end

    return time, mem
end

# Start the benchmarking
memory = []
cpu_time = [['Format Type', 'Compression Time (s)', 'Decompression Time (s)']]
file_size = [['Format Type', 'File Size (KiB)']]

tee "Benchmarking #{formats.count} compression tools on '#{TITLE}'..."

begin
    tmppath = DEVSHM + "/" + TITLE + ".tmp"

    # Create tar archive
    `tar cf #{tmppath}.tar #{PATH}`

    # Display number of files
    tee "Files and directories: " + `tar tf #{tmppath}.tar | wc -l`

    # Display archive file size
    filesize = File.size?("#{tmppath}.tar")
    file_size << ["Pristine", filesize]
    tee "File size (KiB): " + (filesize / 1024).to_s
    tee

    formats.each do |k, v|
        compress_command = "#{v[:compress]} #{tmppath}.tar"
        decompress_command = "#{v[:decompress]} #{tmppath}.tar.#{v[:extension]}"
        if k == '7-Zip'
            compress_command = "#{v[:compress]} #{tmppath}.tar.7z #{tmppath}.tar"
            decompress_command = "#{v[:decompress]} -o#{File.dirname(tmppath)} #{tmppath}.tar.7z"
        end

        # Time consumed by CPU for compressing
        tee "#{k}"
        compression = bm_process(compress_command)

        # Display archive file size
        filesize = File.size?("#{tmppath}.tar.#{v[:extension]}")
        tee "\tFile size (KiB): " + (filesize / 1024).to_s

        # Time consumed by CPU for decompressing
        tee "#{k} (dec)"
        decompression = bm_process(decompress_command)

        tee

        memory << [k, compression[1].n_values(N_MEMVALS)].flatten
        memory << ["#{k} (dec)", decompression[1].n_values(N_MEMVALS)].flatten
        cpu_time << [k, compression[0], decompression[0]]
        file_size << [k, filesize]
    end
ensure
    `rm -rf #{tmppath}.tar*`
end

# Write CSV files
require 'csv'

memory = memory.transpose
CSV.open("#{TITLE}-memory.csv", 'w') do |writer|
    memory.each do |row|
        writer << row
    end
end

CSV.open("#{TITLE}-cpu-time.csv", 'w') do |writer|
    cpu_time.each do |row|
        writer << row
    end
end

CSV.open("#{TITLE}-file-size.csv", 'w') do |writer|
    file_size.each do |row|
        writer << row
    end
end

# Retrieve charts through Google Chart API
require 'gchart'

def max_value(v)
    v + v * 5 / 100
end

def max_value_f(v)
    if v < 1
        v = (v * 10.0).ceil / 10.0
    else
        v = v.ceil
    end
    v
end

def y_axis_values(max)
    Array.new(6) {|i| ((i * 0.2) * max).to_i}
end

def y_axis_values_f(max, dec = 4.0)
    dec = dec.to_f
    Array.new(6) {|i| ((i * 0.2) * max * 10 ** dec).to_i / 10 ** dec}
end

def markers(labels, step)
    markers = Array.new(labels.count) do |i|
        "A#{labels[i]},666666,#{i},#{(i+1)*step},12"
    end
end

# Generate Memory Chart
puts "Writing '#{TITLE}-memory.png'..."
bar_colors = ['76A4FB', 'FFCC33', '9FF670', '990066', 'BBCCED', '224499', '70AE4F']
chart_data = memory.drop(1).transpose
chart_legend = memory[0]
max_value = chart_data.flatten.max
max_value = max_value + max_value * 15 / 100
markers = markers(chart_legend, (N_MEMVALS - 2) / (chart_legend.count).to_f).join('|')
markers = markers + "|R,CCCCCC,0,#{(N_MEMVALS-2)/N_MEMVALS.to_f-0.02},#{(N_MEMVALS-1)/N_MEMVALS.to_f+0.005}" \
    if cpu_time.drop(1).transpose.drop(1).flatten.max > N_MEMVALS / 2 # snapshot timeout in bm_process is 1/2 second
Gchart.line(
    :title => 'RSS Memory (MiB)',
    :size => '700x325',
    :encoding => 'text',
    :data => chart_data,
    :axis_with_labels => 'x,y',
    :axis_labels => [['0', "\u221E"], y_axis_values_f(max_value / 1024.0, 1.0)],
    :min_value => 0,
    :max_value => max_value,
    :axis_range => [[0, max_value]],
    :legend => chart_legend,
    :grid_lines => '0,10,0,0',
    :bar_width_and_spacing => {:width => 'a', :spacing => 20, :group_spacing => 0},
    :bar_colors => bar_colors[0..(chart_data.count)],
    :thickness => ([3] * chart_data.count).join('|'),
    :new_markers => markers,
    :format => 'file',
    :filename => "#{TITLE}-memory.png")

# Generate CPU Time Chart
puts "Writing '#{TITLE}-cpu-time.png'..."
chart_array = cpu_time.drop(1).transpose
max_value = max_value_f(chart_array[1].max) + max_value_f(chart_array[2].max)
Gchart.bar(
    :grouped => false,
    :title => 'CPU Time (s)',
    :size => '350x245',
    :encoding => 'text',
    :data => [chart_array[1], chart_array[2]],
    :axis_with_labels => 'x,y',
    :axis_labels => [chart_array[0], y_axis_values_f(max_value)],
    :min_value => 0,
    :max_value => max_value,
    :axis_range => [[0, max_value]],
    :legend => ['Compression', 'Decompression'],
    :grid_lines => '0,10,0,0',
    :bar_width_and_spacing => {:width => 'a', :spacing => 20, :group_spacing => 0},
    :bar_colors => ['3072F3', 'FF9900'],
    :new_markers => 'N**,000000,0,-1,12,,c|N**,000000,1,-1,12&chdlp=t',
    :format => 'file',
    :filename => "#{TITLE}-cpu-time.png")

# Generate File Size Chart
puts "Writing '#{TITLE}-file-size.png'..."
chart_array = file_size.drop(1).transpose
chart_array[1] = chart_array[1].collect {|i| i/1024/1024.0}
max_value = max_value_f(chart_array[1].max)
Gchart.bar(
    :title => 'File Size (MiB)',
    :size => '350x245',
    :encoding => 'text',
    :data => chart_array[1],
    :axis_with_labels => 'x,y',
    :axis_labels => [chart_array[0], y_axis_values_f(max_value, 2.0)],
    :min_value => 0,
    :max_value => max_value,
    :axis_range => [[0, max_value]],
    :grid_lines => '0,10,0,0',
    :bar_width_and_spacing => {:width => 'a', :spacing => 20, :group_spacing => 0 },
    :bar_colors => '3072F3',
    :new_markers => 'N**,000000,0,-1,12',
    :format => 'file',
    :filename => "#{TITLE}-file-size.png")

# Mashup the charts into one image
puts "Writing '#{TITLE}-3charts.png'..."
img1 = "#{TITLE}-file-size.png"
img2 = "#{TITLE}-cpu-time.png"
img3 = "#{TITLE}-memory.png"
img_output = "#{TITLE}-3charts.png"
begin
    `convert #{img1} #{img2} +append #{img_output}`
    `convert #{img3} #{img_output} -append #{img_output}`
rescue
    puts "The command line tool 'convert' from ImageMagick is not installed"
    exit(3)
end