#! /usr/bin/perl
# Copyright 2007 Christopher Lansdown
# http://www.chrislansdown.com/
# You may redistribute this program under the terms of the artistic license.

use strict;
use DBI;
use Getopt::Long;
Getopt::Long::Configure("bundling");

our $VERSION = "1";
our($DatabaseName, $Help, $Version, $Verbose, $Size, $NormalizeX, $NormalizeY, $NormalizeZ, $Title);

###################
# Option Defaults #
###################
$DatabaseName = `whoami`;
$Size = "800x600";

###################
# Get Our Options #
###################
GetOptions("database|d=s" => \$DatabaseName,
		   "help|h" => \$Help,
		   "normalize-x|x" => \$NormalizeX,
		   "normalize-y|y" => \$NormalizeY,
		   "normalize-z|z" => \$NormalizeZ,
		   "size|s=s" => \$Size,
		   "title|t=s" => \$Title,
		   "verbose|v" => \$Verbose,
		   "version" => \$Version
		  );

########################
# Positional Arguments #
########################
our($Column1, $Column2, $Column3, $OutputFile);
$Column1 = $ARGV[0];
$Column2 = $ARGV[1];
if(@ARGV == 3 || $ARGV[2] =~ m/\.png$/) {
	$OutputFile = $ARGV[2];
} else {
	$Column3 = $ARGV[2];
	$OutputFile = $ARGV[3];
}
if($OutputFile =~ m/\.png$/) {
	$OutputFile = substr($OutputFile, 0, length($OutputFile) - 4);
}

print "output name: $OutputFile\n" if $Verbose;

#########
# Usage #
#########
if($Help || @ARGV < 3) {
	print <<EOHELP;
usage: generate_heat_plot.pl [OPTIONS] column1 column2 [column3] filename

This script automates the process of generating heat plots to
visualize the correlation between two independent variables.

Note: the png extension on "filename" is optional for two columns, and
prohibited for three columns.

    -d, --database=DATABASE
        Connects to the database named DATABASE. defaults to
        the current user's name.

    -h, --help
        Prints this help and exits.

    -s, --size=WxH
        Configures the size of the output png file. defaults to
        800x600.

    -t, --title=TITLE
        Set the title for the plot.

    -v, --verbose
        Turns on verbose output.

    --version
        Prints the version and then exits.

    -x, --normalize-x
        Normalizes the count by the values in the x column.

    -y, --normalize-y
        Normalizes the count by the values in the y column.

    -z, --normalize-z
        Normalizes the count by the values in the z column.

EOHELP
  exit(0);
}


###########
# Version #
###########
if($Version) {
	print <<VERSION;
generate_heat_plot.pl version $VERSION
Copyright 2007 Christopher Lansdown
This program is released under the artistic license.
VERSION
}

###########################
# Connect to the database #
###########################
my $DatabaseConnection = DBI->connect("DBI:Pg:dbname=$DatabaseName") or 
  warn("Unable to connect to the database" . DBI->errstr . "\n");

###############################
# Determine what the query is #
###############################
my $Query;
if(!defined($Column3)) {
	$Query = qq{SELECT foo.$Column1, foo.$Column2, count(foo.id)@@@@@@@@ FROM response as foo WHERE foo.$Column1 IS NOT NULL AND foo.$Column2 IS NOT NULL GROUP BY foo.$Column1, foo.$Column2 ORDER BY foo.$Column1, foo.$Column2};
} else {
	$Query = qq{SELECT foo.$Column1, foo.$Column2, foo.$Column3, count(foo.id)@@@@@@@@ FROM response as foo WHERE foo.$Column1 IS NOT NULL AND foo.$Column2 IS NOT NULL AND foo.$Column3 IS NOT NULL GROUP BY foo.$Column1, foo.$Column2, foo.$Column3 ORDER BY foo.$Column1, foo.$Column2, foo.$Column3};
}

my $norm;
if($NormalizeX || $NormalizeY || $NormalizeZ) {
	$norm = "::float/(";
	if($NormalizeX) {
		$norm .= "(select count($Column1) from response where $Column1 = foo.$Column1)";
		if($NormalizeY || $NormalizeZ) {
			$norm .= "*";
		}
	}
	if($NormalizeY) {
		$norm .= "(select count($Column2) from response where $Column2 = foo.$Column2)";
	}
	if(($NormalizeX || $NormalizeY) && $NormalizeZ) {
		$norm .= "*";
	}
	if($NormalizeZ) {
		$norm .= "(select count($Column3) from response where $Column3 = foo.$Column3)";
	}
	$norm .= ")";
}
$Query =~ s{@@@@@@@@}{$norm}e;
print "Using Query: $Query\n" if $Verbose;

#####################
# Execute the query #
#####################
# Note: we use fetchall_arrayref because we're really unlikely to be
# trying to plot data sets that don't fit into memory. If you are, I
# suggest that you just buy more ram. Thanks to Windows Vista, it's
# getting pretty cheap these days.
my $sth = $DatabaseConnection->prepare($Query) or die DBI->errstr;
$sth->execute() or die DBI->errstr;
my $rows = $sth->fetchall_arrayref();

###########################################
# Calculate the steppings for each column #
###########################################
my(%data, @values);
foreach my $row (@$rows) {
	if($Column3) {
		$data{$row->[0]+0}->{$row->[1]+0}->{$row->[2]+0} = $row->[3];
	} else {
		$data{$row->[0]+0}->{$row->[1]+0} = $row->[2];
	}
	for(my $i = 0; $i < @$row - 1; $i++) {
		$values[$i]->{$row->[$i]} = 1;
	}
}
my @steppings;
for(my $i = 0; $i < @values; $i++) {
	my @v = sort keys %{$values[$i]};
	for(my $j = 1; $j < @v; $j++) {
		if($v[$j] - $v[$j-1] != 0) {
			$steppings[$i] = min($steppings[$i], abs($v[$j] - $v[$j-1]));
		}
	}
}
print "calculated steppings: " . join(", ", @steppings)."\n" if $Verbose;

###########################
# Calculate the min & max #
###########################
my ($x_min, $x_max, $y_min, $y_max, $z_min, $z_max) = (2**31, 0, 2**31, 0, 2**31, 0);
foreach my $xkey (keys %data) {
	$x_min = min($xkey, $x_min);
	$x_max = max($xkey, $x_max);
	foreach my $ykey (keys %{$data{$xkey}||{}}) {
		$y_min = min($ykey, $y_min);
		$y_max = max($ykey, $y_max);
		if($Column3) {
			foreach my $zkey (keys %{$data{$xkey}->{$ykey}||{}}) {
				$z_min = min($zkey, $z_min);
				$z_max = max($zkey, $z_max);
			}
		}
	}
}

my($x_step, $y_step, $z_step) = @steppings;

###############################
# Write out the Complete data #
###############################
my($biggest, %biggest);
unless($Column3) {
	$biggest = 0.1;
	open(my $fh, ">", "/tmp/input.data") or die $!;
	foreach my $x ($x_min..$x_max) {
		for(my $y = $y_min; $y <= $y_max; $y += $y_step) {
			print $fh "$x $y " . ($data{$x}->{$y} + 0) . "\n";
			$biggest = max($biggest, $data{$x}->{$y});
		}
		print $fh "\n";
	}
	close($fh);
} else {
	foreach my $x ($x_min..$x_max) {
		$biggest{$x} = 0.1;
		open(my $fh, ">", "/tmp/input$x.data") or die $!;
		for(my $y = $y_min; $y <= $y_max; $y += $y_step) {
			for(my $z = $z_min; $z <= $z_max; $z += $z_step) {
				print $fh "$y $z " . ($data{$x}->{$y}->{$z} + 0) . "\n";
				$biggest{$x} = max($biggest{$x}, $data{$x}->{$y}->{$z});
			}
			print $fh "\n";
		}
		close($fh);
	}

}

#####################
# Generate the plot #
#####################
unless($Column3) {
	open(my $gp, ">", "/tmp/heat_plot.gnuplot");
	print $gp <<PLOT;
set pm3d map
set terminal png truecolor size 800,600
set output "$OutputFile.png"
set title "$Title"
set cbrange [0:$biggest]
splot "/tmp/input.data" title "" with pm3d
PLOT
	close($gp);
	system("gnuplot", "/tmp/heat_plot.gnuplot") == 0 or die $@;
} else {
	foreach my $x ($x_min..$x_max) {
		open(my $gp, ">", "/tmp/heat_plot$x.gnuplot");
		print $gp <<PLOT;
set pm3d map
set terminal png truecolor size 800,600
set output "$OutputFile$x.png"
set title "$Title"
set cbrange[0:$biggest{$x}]
splot "/tmp/input$x.data" title "" with pm3d
PLOT
		close($gp);
		system("gnuplot", "/tmp/heat_plot$x.gnuplot") == 0 or die $@;
	}
}

####################
# Create Animation #
####################
if($Column3) {
	open(my $fh, ">", "$OutputFile.html");

	# Header 
	print $fh <<EOJS;
<script language="JavaScript">
window.currentFrame = 0;
window.${OutputFile}Values = new Array();
window.${OutputFile}Animation = new Array();
EOJS
	my $i = 0;
	foreach my $x ($x_min..$x_max) {
		print $fh qq{${OutputFile}Values[$i] = "$x";\n};
		print $fh qq{(${OutputFile}Animation[$i] = new Image(800,600)).src = "$OutputFile$x.png";\n};
		$i++;
	}
	print $fh <<EOJS;
function ${OutputFile}Play() {
    var f = document.getElementById('${OutputFile}fps');
    var fps = parseFloat(f.options[f.selectedIndex].value);
    ${OutputFile}Timer = setInterval(${OutputFile}Next, 1000/fps);
}
function ${OutputFile}Pause() {
    clearInterval(${OutputFile}Timer);
}
function ${OutputFile}Next() {
    window.currentFrame = (window.currentFrame + 1) % $i;
    document.getElementById('${OutputFile}Img').src = ${OutputFile}Animation[window.currentFrame].src;
    document.getElementById('${OutputFile}Title').innerHTML = window.${OutputFile}Values[window.currentFrame];
}
function ${OutputFile}Previous() {
    window.currentFrame = (window.currentFrame - 1) % $i;
    document.getElementById('${OutputFile}Img').src = ${OutputFile}Animation[window.currentFrame].src;
    document.getElementById('${OutputFile}Title').innerHTML = window.${OutputFile}Values[window.currentFrame];
}

</script>
<div class="plot">
  <center id="${OutputFile}Title">&nbsp;</center>
  <img id="${OutputFile}Img" width="800" height="600" src="${OutputFile}$x_min.png" />
  <div class="toolbar">
    <input type="button" value="|<" onClick="${OutputFile}Previous();" title="frame back">
    <input type="button" value=">" onClick="${OutputFile}Play();" title="play">
    <input type="button" value="||" onClick="${OutputFile}Pause();" title="pause">
    <input type="button" value=">|" onClick="${OutputFile}Next();" title="frame advance">
FPS: <select id="${OutputFile}fps">
<option>30</option>
<option>24</option>
<option>20</option>
<option>15</option>
<option selected>10</option>
<option>7</option>
<option>5</option>
<option>4</option>
<option>3</option>
<option>2</option>
<option>1</option>
<option>0.5</option>
<option>0.2</option>
<option>0.1</option>
</select>
  </div>
</div>
EOJS

	close($fh);
}


############
# Clean up #
############
unlink(glob("/tmp/input*.data"));
unlink(glob("/tmp/heat_plot*.gnuplot"));







sub min {
	if(!defined($_[0])) { return $_[1] }
	return $_[0] > $_[1] ? $_[1] : $_[0];
}

sub max {
	return $_[0] > $_[1] ? $_[0] : $_[1];
}