package LogLikeGrid;

use vars qw($AUTOLOAD @ISA);
use Exporter;
use Carp;
use strict;

use LogSpace;
use IntIterator;

@ISA = qw(Exporter);

my %fields = (
	      params => undef,
	      data => undef,
	      keyRange => undef,
	      range => undef,
	      rangeLookup => undef,
	      multiplier => undef,
	      combineEntriesRule => sub { LogSpace::sum(@_) },
	      );

sub new {
    my ($obj,$range) = @_;
    my $pkg = ref($obj) || $obj;
    my $self = { _permitted=>\%fields, %fields };
    bless $self, $pkg;
    if ($range==0) { confess $range }
    $self->params(scalar(@$range));
    $self->range([]);
    $self->rangeLookup([]);
    $self->data([]);
    $self->multiplier([]);
    my $mul = 1;
    my $i = 0;
    for ($i=0;$i<@$range;$i++) {
	push @{$self->range}, [sort {$a<=>$b} @{$range->[$i]}];
	push @{$self->multiplier}, $mul;
	$mul *= @{$range->[$i]};
	my $j;
	for ($j=0;$j<@{$range->[$i]};$j++) { $self->rangeLookup->[$i]->{$range->[$i]->[$j]} = $j }
    }
    push @{$self->multiplier}, $mul;
    $self->keyRange($mul);
    $self->data->[$mul-1];
    $self;
}

sub clone {
    my ($obj) = @_;
    my $self = $obj->new($obj->range);
    my $i;
    for ($i=0;$i<@{$obj->data};$i++) { $self->data->[$i] = $obj->data->[$i] }
    $self;
}

sub AUTOLOAD {
    my $self = shift;
    my $type = ref($self) || confess "Don't know about $self";
    my $name = $AUTOLOAD;

    # don't propagate DESTROY messages...

    $name =~ /::DESTROY/ && return;

    $name =~ s/.*://; # get only the bit we want
    unless (exists $self->{'_permitted'}->{$name} ) { confess "In type $type, can't access $name" }

    if (@_) { return $self->{$name} = shift }
    else { return $self->{$name} }
}

# data access methods

sub checkParamIndex {
    my ($self,$param) = @_;
    if (!defined($self->params) || $param<0 || $param>=$self->params) { confess "Bad parameter index" }
}

sub checkParamArray {
    my ($self,$p) = @_;
    if (!defined($self->params) || ($self->params>0 && @$p!=$self->params)) { confess "Bad number of parameters" }
}

sub params2index {
    my ($self,$p) = @_;
    $self->checkParamArray($p);
    my (@index,$i);
    for ($i=0;$i<$self->params;$i++) {
	my $range = $self->range->[$i];
	my $j;
	for ($j=0;$j<@$range;$j++) { last if $range->[$j] <= $p->[$i] }
	confess "Not a recognised grid location (@$p)" if $range->[$j] != $p->[$i];
	$index[$i] = $j;
    }
    \@index;
}

sub index2params {
    my ($self,$index) = @_;
    $self->checkParamArray($index);
    my (@param,$p);
    for ($p=0;$p<$self->params;$p++) {
	my $range = $self->paramRange($p);
	if ($index->[$p]<0 || $index->[$p] >= @$range) { confess "Index for parameter $p (value = $index->[$p]) is no good" }
	$param[$p] = $range->[$index->[$p]];
    }
    \@param;
}

sub params2key {
    my ($self,$p) = @_;
    $self->checkParamArray($p);
    my ($key,$i);
    for ($i=0;$i<$self->params;$i++) {
	my $range = $self->range->[$i];
	my $j;
	if (exists $self->rangeLookup->[$i]->{$p->[$i]}) { $j = $self->rangeLookup->[$i]->{$p->[$i]} }
	else { confess "Not a recognised grid location (@$p)" }
	$key += $self->multiplier->[$i] * $j;
    }
    $key;
}

sub key2params {
    my ($self,$pKey) = @_;
    ($pKey>=0 && $pKey<$self->keyRange) || confess "Key out of range";
    my (@p,$i);
    for ($i=0;$i<$self->params;$i++) {
	$p[$i] = $self->range->[$i]->[$pKey % $self->multiplier->[$i+1]];
	$pKey = int ($pKey / $self->multiplier->[$i+1]);
    }
    \@p;
}

sub getEntryRef {
    my ($self,$pRef) = @_;
    if (!defined $self->params) { $self->params(scalar(@$pRef)) }
    \$self->data->[$self->params2key($pRef)];
}

sub addEntry {
    my ($self,$pRef,$newVal) = @_;
    my $valRef = $self->getEntryRef($pRef);
    if (defined $$valRef) { $newVal = &{$self->combineEntriesRule}($newVal,$$valRef) }
    $$valRef = $newVal;
}

sub entryAt {
    my ($self,$p) = @_;
    $self->data->[$self->params2key($p)];
}

# IO methods

sub newFromHandle {
    my ($pkg,$handle,@args) = @_;
    my (@rangeHash,$entry,$i);
    my @data = <$handle>;
    foreach $entry (@data) {
	$entry = LogSpace::convertText($entry);
	my @p = split /\s+/,$entry;
	for ($i=0;$i<$#p;$i++) { $rangeHash[$i]->{$p[$i]} = 1 }
    }
    my @range;
    for ($i=0;$i<@rangeHash;$i++) { $range[$i] = [sort {$a<=>$b} keys %{$rangeHash[$i]}] }
    my $self = $pkg->new(\@range,@args);
    $i=0;
    foreach $entry (@data) {
	if (++$i%100==0) { print "... processing line $i ...\n" }
	$entry = LogSpace::convertText($entry);
	my @p = split /\s+/,$entry;
	if (@p<1 || $entry =~ /[^(\d\s\.\-+e)]/) { confess "Bad input format" }
	$self->addEntry([@p[0..$#p-1]],$p[$#p]);
    }
    $self;
}

sub newFromFile {
    my ($pkg,$filename,@args) = @_;
    local *file;
    open file, "$filename" or confess "Couldn't open $filename: $!";
    $pkg->newFromHandle(\*file,@args);
}

# methods for maintaining consistency of table

sub isDefinedEverywhere {
    my ($self) = @_;
    foreach (@{$self->data}) {
	if (!defined $_) { return 0 }
    }
    1;
}

sub checkDefinedEverywhere {
    my ($self) = @_;
    unless ($self->isDefinedEverywhere) { confess "Need a defined entry everywhere" }
}

sub paramRange {
    my ($self,$param) = @_;
    $self->checkParamIndex($param);
    $self->range->[$param];
}

sub checkParamsInRange {
    my ($self,$p) = @_;
    $self->checkParamArray($p);
    my $i;
    for ($i=0;$i<$self->params;$i++) {
	my $range = $self->paramRange($i);
	if ($p->[$i]<$range->[0] || $p->[$i]>$range->[@$range-1]) { confess "Parameters out of range" }
    }
}

# methods for operating on the table

sub evaluateFunction {
    my ($self,$function) = @_;
    my $key;
    for ($key=0;$key<$self->keyRange;$key++) {
	my $valRef = $self->getEntryRef($key);
	my $p = $self->key2params($key);
	$$valRef = &$function($self,$p,$$valRef);
    }
}

sub subVolume {
    my ($self) = @_;
    if (!defined $self->params) { return undef }
    my $volume = 1;
    my $i;
    for ($i=0;$i<$self->params;$i++) {
	my $range = $self->paramRange($i);
	if (@$range>1) { $volume *= ($range->[@$range-1] - $range->[0]) }
    }
    $volume;
}

sub findMode {
    my ($self) = @_;
    if (!defined $self->params) { return undef }
    my $modeKey;
    my $modeVal = undef;
    my $key;
    for ($key=0;$key<$self->keyRange;$key++) {
	my $val = $self->data->[$key];
	if (!defined($modeVal) || $val > $modeVal) {
	    $modeKey = $key;
	    $modeVal = $val;
	}
    }
    key2params($modeKey);
}

sub combine {
    my ($pkg,$table1,$table2) = @_;
    if (!defined($table1->params)) { return $table2->clone }
    elsif (!defined($table2->params)) { return $table1->clone }
    my $result = $pkg->new($table1->params + $table2->params);
    my $key1;
    for ($key1=0;$key1<$table1->keyRange;$key1++) {
	my $val1 = $table1->data->[$key1];
	my $p1 = $table1->key2params($key1);
	my $key2;
	for ($key2=0;$key2<$table2->keyRange;$key2++) {
	    my $val2 = $table2->data->[$key2];
	    my $p2 = $table2->key2params($key2);
	    $result->addEntry([@$p1,@$p2],$val1);
	    $result->addEntry([@$p1,@$p2],$val2);
	}
    }
    $result;
}

sub absorb {
    my ($self,$table2) = @_;
    my $key;
    for ($key=0;$key<$self->keyRange;$key++) {
	my $val = $self->data->[$key];
	$self->addEntry($table2->key2params($key),$val);
    }
}

sub mergeTables {
    my ($self,@tables) = @_;
    my $result = $self->clone;
    foreach (@tables) { $result->absorb($_) }
    $result;
}

sub print {
    my $self = shift;
    my $handle = @_ ? shift : \*STDOUT;
    my $condition = @_ ? shift : sub {1};
    my $key;
    for ($key=0;$key<$self->keyRange;$key++) {
	my $val = $self->data->[$key];
	my $p = $self->key2params($key);
	print $handle (join "\t",(@$p,$val))."\n" if &$condition($self,$p);
    }
}

sub formatDescriptionText {
    my $text;
    $text .= "Each line represents a data point, formatted as follows (n = no. of parameters):\n";
    $text .= "\n";
    $text .= " param_1 param_2 ... param_n log-likelihood\n";
    $text .= "\n";
    $text;
}

# grid methods

sub newIndex {
    my ($self) = @_;
    my @max = ();
    my $p;
    for ($p=0;$p<$self->params;$p++) { $max[$p] = $self->paramRange($p) - 1 }
    IntIterator->new(\@max);
}

sub checkGrid {
    my ($self) = @_;
    my $index = $self->newIndex;
    my $missing = undef;
    for (;!$index->done;$index->next) {
	my $p = $self->index2params($index->pos);
	unless ($self->hasEntryAt($p)) { $missing .= "Missing data at (".params2key($p).")\n" }
    }
    if ($missing) { confess "$missing\nData points missing" }
    $self->checkDefinedEverywhere;
}

sub marginalise {
    my ($self,$integrateFlag) = @_;
    $self->checkParamArray($integrateFlag);
    my @reducedParamList = ();
    my @logDeltas;
    my ($i,$j);
    for ($i=0;$i<$self->params;$i++) {
	my $range = $self->paramRange($i);
	if ($integrateFlag->[$i]) {
	    if (@$range==1) {
		$logDeltas[$i]->{$range->[0]} = 0;
	    } else {
		for ($j=0;$j<@$range;$j++) {
		    my $delta = 0;
		    $delta += $range->[$j+1] - $range->[$j] if $j < @$range-1;
		    $delta += $range->[$j] - $range->[$j-1] if $j > 0;
		    $logDeltas[$i]->{$range->[$j]} = log($delta/2);
		}
	    }
	} else {
	    push @reducedParamList, $i;
	}
    }
    my $result = $self->new([@{$self->range}[@reducedParamList]]);
    my $key;
    for ($key=0;$key<$self->keyRange;$key++) {
	my $val = $self->data->[$key];
	my $p = $self->key2params($key);
	my $resultVal = $result->getEntryRef([@{$p}[@reducedParamList]]);
	my $logDelta = 0;
	for ($i=0;$i<$self->params;$i++) { $logDelta += $logDeltas[$i]->{$p->[$i]} if $integrateFlag->[$i] }
	if (defined $$resultVal) { $$resultVal = LogSpace::logsumexp($$resultVal,LogSpace::sum($logDelta,$val)) }
	else { $$resultVal = LogSpace::sum($logDelta,$val) }
    }
    $result;
}

sub integrate {
    my ($self) = @_;
    my (@flag,$i);
    for ($i=0;$i<$self->params;$i++) { $flag[$i] = 1 }
    $self->marginalise(\@flag)->entryAt();
}

sub interpolate {
    my ($self,$p) = @_;
    $self->checkParamArray($p);
    my @mix;
    my @pos;
    my $i;
    for ($i=0;$i<$self->params;$i++) {
	my $range = $self->paramRange($i);
	if (@$range<2) { confess "Need more than one point to interpolate" }
	my $j;
	for ($j=0;$j<@$range-2;$j++) { last if $p->[$i] <= $range->[$j] }
	$mix[$i] = ($range->[$j+1] - $p->[$i]) / ($range->[$j+1] - $range->[$j]);
	$pos[$i] = $j;
    }
    my $positive = -$LogSpace::INFTY;
    my $negative = -$LogSpace::INFTY;
    my $cornerIndex = $self->newIndex;
    foreach (@{$cornerIndex->max}) { $_ = 1 }
  INTERPOLATION_CORNER:    for (;!$cornerIndex->done;$cornerIndex->next) {
      my $mix = 1;
      for ($i=0;$i<$self->params;$i++) {
	  $mix *= $cornerIndex->pos->[$i]==0 ? $mix[$i] : 1-$mix[$i];
	  next INTERPOLATION_CORNER unless $mix;
      }
      my $entry = $self->entryAt($self->index2params(IntIterator::addArrays(\@pos,$cornerIndex->pos)));
      if ($mix>0) { $positive = LogSpace::logsumexp($positive,LogSpace::sum(LogSpace::safelog($mix),$entry)) }
      else { $negative = LogSpace::logsumexp($negative,LogSpace::sum(LogSpace::safelog(-$mix),$entry)) }
  }
    
    $positive>=$negative ? LogSpace::logdiffexp($positive,$negative) : -$LogSpace::INFTY;
}

sub transform {
    my ($self,$pFunction,$qRangeArray) = @_;
    my $result = $self->new($qRangeArray);
    my @maxIndex;
    my $i;
    for ($i=0;$i<@$qRangeArray;$i++) { $maxIndex[$i] = @{$qRangeArray->[$i]} - 1 }
    my $index = IntIterator->new(\@maxIndex);
    for (;!$index->done;$index->next) {
	my @q;
	for ($i=0;$i<@$qRangeArray;$i++) { $q[$i] = $qRangeArray->[$i]->[$index->pos->[$i]] }
	$result->addEntry(\@q,$self->interpolate(&$pFunction(\@q)));
    }
    $result;
}

1;
