package LogLikeTable;

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

use LogSpace;
use IntIterator;

@ISA = qw(Exporter);

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

sub new {
    my $obj = shift;
    my $pkg = ref($obj) || $obj;
    my $self = { _permitted=>\%fields, %fields };
    bless $self, $pkg;
    $self->range([]);
    $self->data({});
    if (@_>0) { $self->params(shift) }
    $self->rangesValid(0);
    $self;
}

sub clone {
    my ($obj) = @_;
    my $self = $obj->new($obj->params);
    my ($key,$val);
    while (($key,$val) = each %{$obj->data}) { $self->data->{$key} = $val }
    $self->rangesValid(0);
    $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) = @_;
    unless (!defined($p) || ref($p)) { confess "Not an ARRAY ref" }
    if (!defined($self->params) || ($self->params>0 && @$p!=$self->params)) { confess "Bad number of parameters" }
}

sub params2key {
    my ($self,$p) = @_;
    $self->checkParamArray($p);
    foreach (@$p) { $_ = $_+0 }    # ensure all numeric
    my $pkey = join " ",@$p;
    $pkey;
}

sub key2params {
    my ($self,$pkey) = @_;
    my @p = split /\s+/,$pkey;
    $self->checkParamArray(\@p);
    \@p;
}

sub getEntryRef {
    my ($self,$pRef) = @_;
    if (!defined $self->params) { $self->params(scalar(@$pRef)) }
    my $pKey = $self->params2key($pRef);
    unless (exists $self->data->{$pKey}) {
	$self->rangesValid(0);
	$self->data->{$pKey} = undef;
    }
    \$self->data->{$pKey};
}

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

sub hasEntryAt {
    my ($self,$p) = @_;
    exists ($self->data->{$self->params2key($p)});
}

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

# input/new methods

sub newFromHandle {
    my ($pkg,$handle,@args) = @_;
    my $self = $pkg->new(@args);
    my $entry;
    my $i;
    while ($entry = <$handle>) {
#	if (++$i%100==0) { warn "... 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);
}

sub newFromArray {
    my ($pkg,$array,@args) = @_;
    my $self = $pkg->new(@args);
    my $entry;
    my $i;
    foreach $entry (@$array) {
#	if (++$i%100==0) { warn "... 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 newFromFunction {
    my ($pkg,$rangeArray,$function) = @_;
    my $self = $pkg->new(scalar @$rangeArray);
    my @maxIndex = map(@$_ - 1,@$rangeArray);
    my $index = IntIterator->new(\@maxIndex);
    for (;!$index->done;$index->next) {
	my (@q,$i);
	for ($i=0;$i<@$rangeArray;$i++) { $q[$i] = $rangeArray->[$i]->[$index->pos->[$i]] }
	$self->addEntry(\@q,&$function(\@q));
    }
    $self;
}

# methods for maintaining consistency of table

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

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

sub calculateRanges {
    my ($self) = @_;
    unless (keys(%{$self->data}) > 0) { confess "Need at least one data point" }
    my $p;
    for ($p=0;$p<$self->params;$p++) {
	my %rangeHash = ();
	foreach (keys %{$self->data}) {
	    $rangeHash{$self->key2params($_)->[$p]} = 1;
	}
	$self->range->[$p] = [sort {$a<=>$b} keys %rangeHash];
    }
    $self->rangesValid(1);
}

sub paramRange {
    my ($self,$param) = @_;
    $self->checkParamIndex($param);
    unless ($self->rangesValid) { $self->calculateRanges }
    $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,$val);
    while (($key,$val) = each %{$self->data}) {
	my $p = $self->key2params($key);
	$self->data->{$key} = &$function($self,$p,$val);
    }
}

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) || keys(%{$self->data})==0) { return undef }
    my $modeKey;
    my $modeVal = undef;
    my ($key,$val);
    while (($key,$val) = each %{$self->data}) {
	if (!defined($modeVal) || $val > $modeVal) {
	    $modeKey = $key;
	    $modeVal = $val;
	}
    }
    $self->key2params($modeKey);
}

sub combine {
    my ($pkg,$table1,$newParamIndex1,$table2,$newParamIndex2) = @_;
    if (!defined($table1->params) || !defined($table2->params)) { confess "Can't combine empty tables" }
    $table1->checkParamArray($newParamIndex1);
    $table2->checkParamArray($newParamIndex2);
    my (@oldParamIndex,@sharedIndex1,@sharedIndex2,$i,$j);
    for ($i=0;$i<@$newParamIndex1;$i++) { $oldParamIndex[$newParamIndex1->[$i]] = $i + 1 }
    for ($i=0;$i<@$newParamIndex2;$i++) { $oldParamIndex[$newParamIndex2->[$i]] = $i + $table1->params + 1 }
    foreach (@oldParamIndex) { (--$_ < 0) && confess "Table resulting from combine operation has undefined columns" }
    foreach $i (@$newParamIndex1) { if (($j=grep($_==$i,@$newParamIndex1))>1) { confess "Can't map $j parameters onto one" } }
    foreach $i (@$newParamIndex2) { if (($j=grep($_==$i,@$newParamIndex2))>1) { confess "Can't map $j parameters onto one" } }
    for ($i=0;$i<@$newParamIndex1;$i++) {
	for ($j=0;$j<@$newParamIndex2;$j++) {
	    if ($newParamIndex1->[$i] == $newParamIndex2->[$j]) {
		push @sharedIndex1,$i;
		push @sharedIndex2,$j;
	    }
	}
    }
    my $result = $pkg->new(scalar(@oldParamIndex));
    my ($key1,$val1);
    while (($key1,$val1) = each %{$table1->data}) {
	my $p1 = $table1->key2params($key1);
	my @sharedp1 = @{$p1}[@sharedIndex1];
	my ($key2,$val2);
      COMBINE_LOOP:	while (($key2,$val2) = each %{$table2->data}) {
	    my $p2 = $table2->key2params($key2);
	    my @sharedp2 = @{$p2}[@sharedIndex2];
	    for ($i=0;$i<@sharedp1;$i++) { if ($sharedp1[$i]!=$sharedp2[$i]) { next COMBINE_LOOP } }
	    my @newp = (@$p1,@$p2)[@oldParamIndex];
	    $result->addEntry(\@newp,$val1);
	    $result->addEntry(\@newp,$val2);
	}
    }
    $result;
}

sub absorb {
    my ($self,$table2) = @_;
    my ($key,$val);
    while (($key,$val) = each %{$table2->data}) {
	$self->addEntry($table2->key2params($key),$val);
    }
    $self;
}

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

sub byKey {
    my ($a,$b) = @_;
    my ($a1,$othera) = split /\s+/,$a,2;
    my ($b1,$otherb) = split /\s+/,$b,2;
    $a1<=>$b1 || ($othera=~/\S/ && $otherb=~/\S/ ? byKey($othera,$otherb) : 0);
}

sub print {
    my $self = shift;
    my $handle = @_ ? shift : \*STDOUT;
    my $condition = @_ ? shift : sub {1};
    my $key;
    foreach $key (sort {byKey($a,$b)} keys %{$self->data}) {
	my $p = $self->key2params($key);
	print $handle (join "\t",(@$p,$self->data->{$key}))."\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 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 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(scalar @reducedParamList);
    my ($key,$val);
    while (($key,$val) = each %{$self->data}) {
	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 = map(1,1..$self->params);
    $self->marginalise(\@flag)->entryAt();
}

sub normalise {
    my ($self) = @_;
    my $lognorm = $self->integrate;
    $self->evaluateFunction(sub { LogSpace::sum($_[2],-$lognorm) });
}

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(scalar(@$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;
