#!/usr/bin/perl

use Time::HiRes qw(time);
use Statistics::Descriptive;
use GD::Graph::lines;
use GD::Graph::colour;
use GD::Graph::Data;
use Data::Dumper;


sub doit
{
    if (defined($pre_start_cmd)) {
        $local_pre_start_cmd = $pre_start_cmd;
        &debug("before substitutions pre_start_cmd is [$local_pre_start_cmd] ...\n");
        $local_pre_start_cmd =~ s/%THREADS%/$threads/g;
        $local_pre_start_cmd =~ s/%COUNT%/$count/g;
        $local_pre_start_cmd =~ s/%BLOCKSIZE%/$blocksize/g;
        &debug("after substitutions pre_start_cmd is [$local_pre_start_cmd] ...\n");
        system($local_pre_start_cmd);
    }
    $start_time = time;
    if (defined($post_start_cmd)) {
        $local_post_start_cmd = $post_start_cmd;
        &debug("before substitutions post_start_cmd is [$local_post_start_cmd] ...\n");
        $local_post_start_cmd =~ s/%THREADS%/$threads/g;
        $local_post_start_cmd =~ s/%COUNT%/$count/g;
        $local_post_start_cmd =~ s/%BLOCKSIZE%/$blocksize/g;
        &debug("after substitutions post_start_cmd is [$local_post_start_cmd] ...\n");
        system($local_post_start_cmd);
    }

    for ($thread=0; $thread<$threads; $thread++) {
        $local_cmd = $cmd;
        $local_cmd =~ s/%THREAD%/$thread/g;
        $local_cmd =~ s/%COUNT%/$count/g;
        $local_cmd =~ s/%BLOCKSIZE%/$blocksize/g;
        if (!defined($pid=fork)) {
            die "fork failed: $!\n";
        } elsif ($pid != 0 ) { 
        } else {
            exec $local_cmd;
            die "exec failed ($!)\n";
        }
    }

    while (wait != -1) {
    }
    
    if (defined($pre_stop_cmd)) {
        $local_pre_stop_cmd = $pre_stop_cmd;
        &debug("before substitutions pre_stop_cmd is [$local_pre_stop_cmd] ...\n");
        $local_pre_stop_cmd =~ s/%THREADS%/$threads/g;
        $local_pre_stop_cmd =~ s/%COUNT%/$count/g;
        $local_pre_stop_cmd =~ s/%BLOCKSIZE%/$blocksize/g;
        &debug("after substitutions pre_stop_cmd is [$local_pre_stop_cmd] ...\n");
        system($local_pre_stop_cmd);
    }
    $end_time = time;
    if (defined($post_stop_cmd)) {
        $local_post_stop_cmd = $post_stop_cmd;
        &debug("before substitutions post_stop_cmd is [$local_post_stop_cmd] ...\n");
        $local_post_stop_cmd =~ s/%THREADS%/$threads/g;
        $local_post_stop_cmd =~ s/%COUNT%/$count/g;
        $local_post_stop_cmd =~ s/%BLOCKSIZE%/$blocksize/g;
        &debug("before substitutions post_stop_cmd is [$local_post_stop_cmd] ...\n");
        system($local_post_stop_cmd);
    }

    if ($blocksize =~ /(.*)K$/) {
        $blocksize_bytes = $1 * 1024;
    } elsif ($blocksize =~ /(.*)M$/) {
        $blocksize_bytes = $1 * 1024 * 1024;
    } else {
        $blocksize_bytes = $blocksize;
    }

    
    $rate = $count * $blocksize_bytes * $threads / ($end_time - $start_time);
}

@tests = (
    { 
        desc=>"CPU bandwidth",
        pre_start_cmd=>undef,
        post_start_cmd=>undef,
        pre_stop_cmd=>undef,
        post_stop_cmd=>undef,
        count=>10000,
        threads=>10,
	blocksizes=>[qw/1K 2K 4K 8K 16K 32K 64K 128K 256K 512K 1M 2M 4M 8M/],
        cmd=>"dd if=/dev/zero bs=%BLOCKSIZE% count=%COUNT% of=/dev/null 2>/dev/null" 
    },
    { 
        desc=>"network bandwidth",
        pre_start_cmd=>undef,
        pre_start_cmd=>"for ((THREAD=0; \$THREAD<%THREADS%; THREAD++)); do
                            ssh -fn macaroni \"nc -lp \$((8000+\$THREAD)) > /dev/null\" &
                        done
                        #  allow listeners to settle
                        sleep 2",
        post_start_cmd=>undef,
        pre_stop_cmd=>undef,
        post_stop_cmd=>"ssh -n macaroni killall nc 2>/dev/null",
        count=>1000,
        threads=>4,
	blocksizes=>[qw/1K 2K 4K 8K 16K 32K 64K 128K 256K 512K 1M/],
        cmd=>"dd if=/dev/zero bs=%BLOCKSIZE% count=%COUNT% 2>/dev/null | nc -q 0 macaroni \$((8000+%THREAD%))"
    },
    {
        desc=>"disk write bandwidth",
        pre_start_cmd=>undef,
        post_start_cmd=>undef,
        pre_stop_cmd=>"sync",
        post_stop_cmd=>undef,
        count=>10000,
        threads=>2,
        blocksizes=>[qw/1K 2K 4K 8K 16K 32K 64K 128K 256K 512K 1M/],
        cmd=>"dd if=/dev/zero bs=%BLOCKSIZE% count=%COUNT% of=/dev/vg0/scratch2 2>/dev/null"
    },
    {
        desc=>"disk read bandwidth",
        pre_start_cmd=>undef,
        post_start_cmd=>undef,
        pre_stop_cmd=>undef,
        post_stop_cmd=>undef,
        count=>1600,
        threads=>1,
	blocksizes=>[qw/1K 2K 4K 8K 16K 32K 64K 128K 256K 512K 1M/],
    }
);


#warn Dumper(\@tests);
#exit;

sub main
{
    $repetitions = 20;
    GD::Graph::colour::read_rgb("/usr/share/X11/rgb.txt") or die "Cannot read colours from rgb.txt";

    for ($i=0; $i<$#tests+1; $i++) {
        $desc = $tests[$i]{desc};
        $pre_start_cmd = $tests[$i]{pre_start_cmd};
        $post_start_cmd = $tests[$i]{post_start_cmd};
        $pre_stop_cmd = $tests[$i]{pre_stop_cmd};
        $post_stop_cmd = $tests[$i]{post_stop_cmd};
        $count = $tests[$i]{count};
        $threads = $tests[$i]{threads};
        $cmd = $tests[$i]{cmd};
        @blocksizes = @{$tests[$i]{blocksizes}};

        print "$desc\n", "=" x length($desc), "\n\n";
        $already_worked_out_divider = 0;
        for ($j=0; $j<$#blocksizes+1; $j++) {
            $blocksize = $blocksizes[$j];
	    # starting to do stats on the rates
	    $stat = Statistics::Descriptive::Full->new();
            for ($k=0; $k<$repetitions; $k++) {
               &doit;
	       $stat->add_data($rate);
            }
	    $mean = $stat->mean();
            if (!$already_worked_out_divider) {
                 ($divider, $units) = &get_divider_and_units($mean);
                 $already_worked_out_divider = 1;
            }
            $std  = $stat->standard_deviation();
            $max  = $stat->min();
            $min  = $stat->max();

	    printf "blocksize: %4s mean rate: %6.3f %3s, std dev: %6.3f %3s, min: %6.3f %3s, max: %6.3f %3s\n",
	            $blocksize,$mean/$divider,$units,$std/$divider,$units,$min/$divider,$units,$max/$divider,$units;

            #  Assemble data in 2D array for plotting
            $data[0][$j] = $blocksizes[$j];
            $data[1][$j] = $mean/$divider;

            # stop doing the stats and clear the array of data
            undef $stat;
            
        }
        print "\n";

        $my_graph = new GD::Graph::lines();
        $my_graph->set( 
        	x_label => 'blocksize',
        	y_label => "rate ($units)",
        	title => $desc,
        	#y_max_value => 10,
        	#y_min_value => -5,
        	#y_tick_number => 3,
        	#y_label_skip => 1,
        	zero_axis_only => 0,
        	long_ticks => 1,
        	x_ticks => 0,
        	dclrs => [ qw(red) ],
        	line_types => [ 1 ],
        	line_type_scale => 8,
        	legend_marker_width => 24,
        	line_width => 3,
        	show_values => 0,
        	transparent => 0,
        ) or warn $my_graph->error;
        
        $my_graph->set_legend($desc);
        $my_graph->plot(\@data)->png or die $my_graph->error;
	$ext = "png";
	open(OUT, ">$desc.$ext") or die "Cannot open $name.$ext for write: $!";
	binmode OUT;
	print OUT $my_graph->gd->$ext();
	close OUT;
        undef @data;
    }
}

sub get_divider_and_units
{
    my ($value) = @_;
    my ($units, $divider);

    if ($value > 1024*1024*1024) {
        $divider = 1024*1024*1024;
        $units = "GB/s";
    } elsif ($rate > 1024*1024) {
        $divider = 1024*1024;
        $units = "MB/s";
    } elsif ($rate > 1024) {
        $divider = 1024;
        $units = "KB/s";
    } else {
        $divider = 1;
        $units = "B/s";
    }

    return($divider, $units);

    #printf "%4s %6.3f %s\n", $blocksize, $rate, $units;
}

sub debug
{
    # printf @_;
}

&main;