#!/usr/local/bin/perl -w

use PropTree;

@fields = qw(seqname source feature start end score strand frame group);

@dummygfflist = (["dummy",".",".",0,0,0,".",".","",-1]);

$infinity = 10**9;

($prog = $0) =~ s#.*/(\S+)#$1#;

$usage .= "$prog - assemble GFF segments by dynamic programming\n";
$usage .= "\n";
$usage .= "Usage: $prog [-d<flags>] <model file> [<GFF files>]\n";
$usage .= "\n";
$usage .= "GFF files must be sorted by seqname and startpoint\n";
$usage .= "\n";
$usage .= "Model file format: too busy to document - sorry... email ihh\@sanger.ac.uk\n";
$usage .= "\n";
$usage .= "Use -d<flags> to set debug flags\n";
$usage .= "\n";

while (@ARGV) {
    last unless $ARGV[0] =~ /^-/;
    $opt = shift;
    if ($opt =~ s/^-d//) { while (($c = chop $opt) ne "") { $debug{$c} = 1 } }
    else { die "$usage\nUnknown option: $opt\n" }
}

@ARGV>=1 or die $usage;
$modelfile = shift;

$model = new_from_file PropTree($modelfile);

if (defined $model->{"initialise"}) { foreach (@{$model->{"initialise"}}) { eval $_ } }

defined($flushlen = $model->flushlen->[0]) or $flushlen = $infinity;

defined($starttag = $model->{"startstate"}->[0]) or $starttag = "start";
defined($endtag = $model->{"endstate"}->[0]) or $endtag = "end";

$startstate = taglookup($starttag,\%statenumber,\@statename);
$endstate = taglookup($endtag,\%statenumber,\@statename);

$declaredstates = 0;
if (defined $model->states) {
    foreach $tag (split /\s+/,"@{$model->states}") { taglookup($tag,\%statenumber,\@statename) }
    $declaredstates = 1;
}

foreach $link (@{$model->link}) {

    defined($fromtag = $link->{"from"}->[0]) or die "All links must have a <from> field";
    defined($totag = $link->{"to"}->[0]) or die "All links must have a <to> field";

    $from = taglookup(lc $fromtag,\%statenumber,\@statename);
    $to = taglookup(lc $totag,\%statenumber,\@statename);
    
    defined($lencostexpr = $link->{"lencost"}->[0]) or $lencostexpr = 0;

    defined($span = $link->{"span"}->[0]) or die "All links must have a <span> field";
    if ($span eq "free" || $span eq "lrfree") { ($ltied,$rtied,$gap) = (0,0,0) }
    elsif ($span eq "ltied" || $span eq "rfree") { ($ltied,$rtied,$gap) = (1,0,0) }
    elsif ($span eq "rtied" || $span eq "lfree") { ($ltied,$rtied,$gap) = (0,1,0) }
    elsif ($span eq "tied" || $span eq "lrtied") { ($ltied,$rtied,$gap) = (1,1,0) }
    elsif ($span eq "gap") { ($ltied,$rtied,$gap) = (0,0,1) }
    else { die "Illegal <span> field: $span" }

    my @lencost;
    if (defined($maxlen = $link->{"maxlen"}->[0])) {
	if (!defined($maxlenfrom[$from]) || $maxlen > $maxlenfrom[$from]) { $maxlenfrom[$from] = $maxlen }
	for ($x=0;$x<=$maxlen;$x++) { $lencost[$x] = eval $lencostexpr }
	$lencost = \@lencost;
    } else {
	$gap or die "Must specify a maximum length for non-gap links";
	$maxlen = $infinity;
	$lencost = 0;
    }
    
    if (defined($stack = $link->{"stack"}->[0])) {
	if ($stack eq "push") { $stackflag = +1 }
	elsif ($stack eq "pop") { $stackflag = -1 }
	else { die "Illegal <stack> field: $stack" }
    } else {
	$stackflag = 0;
    }

    if (defined($gfffilter = $link->{"gfffilter"}->[0])) { $gfffilter = substitute($gfffilter) }
    
    push @links, [$from,
		  $to,
		  $lencost,
		  $maxlen,
		  $ltied,
		  $rtied,
		  $gap,
		  $stackflag,
		  $gfffilter,
		  $link];
    
    $numberoflinks{$from,$to}++;

}

if ($debug{'m'}) {
    warn "* model: states=(@statename)\n";
    foreach $link (@links) { warn "* model: @$link\n" }
}

$numberoflinks{$startstate,$startstate} or die "Start state not padded";
$numberoflinks{$endstate,$endstate} or die "End state not padded";

$links = @links;
$states = keys %statenumber;

for ($state=0;$state<$states;$state++) {
    if (defined($ml = $maxlenfrom[$state])) {
	$maxlenfrom[$state] = $flushlen if $ml > $flushlen;
    } else {
	$maxlenfrom[$state] = 0;
    }
}

@stateinfo = map([],1..$states);

$gffcount = 0;
while (<>) {
    s/#.*//;
    next unless /\S/;
    chomp;
    $gff = [split(/\t/,$_,9)];
    
    ($seqname,$start) = @{$gff}[0,3];
    if (!defined($lastseqname) || $seqname ne $lastseqname || $start > $laststart + $flushlen) {
	flushgfflist();
	flushstates();
	$lastseqname = $gff->[0];
	if ($debug{'n'}) { warn "* seqname $lastseqname\tstart $start\n" }
    }
    $laststart = $start;
    
    if ($debug{'p'}) { warn "* seqname $lastseqname\tstart $start\n" }
    updatestack($start);
    updategfflist($start);
    updatestates($start);
    addtogfflist($gff,++$gffcount);
}

flushgfflist();
flushstates();

###################################################################################
# subroutines
###################################################################################

sub taglookup {
    my ($tag,$lookup,$name) = @_;
    unless (exists $lookup->{$tag}) {
	if ($declaredstates) { die "Undeclared state: $tag" }
	$name->[$lookup->{$tag} = keys %$lookup] = $tag;
    }
    $lookup->{$tag};
}

sub addtogfflist {
    my ($gff,$count) = @_;
    $gff->[9] = $count;
    if ($debug{'g'}) { warn "* addgff: ".join("\t",@$gff)."\n" }
    my $end = $gff->[4];
    my $i;
    for ($i=0;$i<@gfflist;$i++) { last if ($gfflist[$i]->[4] > $end) }
    splice @gfflist, $i, 0, $gff;
}

sub updategfflist {
    my ($pos) = @_;
    return if defined($lastupdategffpos) && $pos <= $lastupdategffpos;
    $lastupdategffpos = $pos;
    my ($gff,$end);
    while (@gfflist && ($end = $gfflist[0]->[4]) < $pos-1) {
	updatestates($end + 1);
	$gff = shift @gfflist;
	if ($debug{'l'}) { warn "* gfflist_remove: (@$gff)\n" }
    }
    if ($debug{'l'}) { foreach $gff (@gfflist) { warn "* gfflist: (@$gff)\n" } }
}

sub updatestates {
    my ($pos) = @_;
    return if defined($lastupdatestatespos) && $pos <= $lastupdatestatespos;
    $lastupdatestatespos = $pos;
    if ($debug{'u'}) { warn "* updatestates pos=$pos\n" }
    my ($link,$state,$gff,@path);
    for ($state=0;$state<$states;$state++) {
	my $path = {};
	push @path, $path;
	push @{$stateinfo[$state]}, [$pos,$path];
    }
    foreach $link (@links) {
	my ($from,$to,$lencost,$maxlen,$ltied,$rtied,$gap,$stackflag,$gfffilter) = @$link;
	my $topath = $path[$to];
	my $effectivegfflist = \@gfflist;
	if ($gap) { $effectivegfflist = \@dummygfflist }
      GFF: foreach $gff (@$effectivegfflist) {
	  unless ($gap) {
	      if ($stackflag * $gff->[9] < 0) { next GFF }
	      unless (!defined($gfffilter) || eval($gfffilter) || $stackflag < 0) { next GFF }
	      unless ((!$rtied && $gff->[4] >= $pos - 1) || $gff->[4] == $pos - 1) { next GFF }
	  }
	  my $pos_stacks;
	FROMPOS: foreach $pos_stacks (@{$stateinfo[$from]}) {
	    my ($frompos,$fromstacks) = @$pos_stacks;
	    if ($debug{'c'}) { my $text = "* checklink: to=$statename[$to] from=$statename[$from] pos=$pos frompos=$frompos link=(@$link)"; $text .= " gff=(@$gff)" unless $gap; warn "$text\n"; }
	    if ($pos - $frompos > $maxlen) { next FROMPOS }
	    if ((!$gap && $gff->[3] > $frompos) || ($ltied && $gff->[3] != $frompos)) { next FROMPOS }
	    my $gfflen;
	    my $gffscore;
	    if ($gap) {
		$gff = [$lastseqname,"gap","gap",$frompos,$pos-1,0,"+",".","",0];   # pseudo-GFF for gaps
		$gfflen = $pos - $frompos;
		$gffscore = 0;
	    } else {
		$gfflen = $gff->[4] + 1 - $gff->[3];
		$gffscore = $gff->[5] * ($pos-$frompos) / $gfflen;
	    }
	    if (ref $lencost) { $gffscore -= $lencost->[$pos - $frompos] }
	    my $newgff;
	    if ($stackflag>0) {
		$newgff = [@$gff];
		my ($pushstart,$pushend) = split /\s+/, $gff->[8];
		$newgff->[3] = $pushstart + int(($pushend - $pushstart) * ($frompos - $gff->[3]) / $gfflen);
		$newgff->[4] = $pushstart + int(($pushend - $pushstart) * ($pos - $gff->[3]) / $gfflen);
		if ($newgff->[3] > $newgff->[4]) {
		    @{$newgff}[3,4] = @{$newgff}[4,3];
		    $newgff->[6] = "-";
		    $newgff->[8] =~ s/^\S+\s+\S+/$gff->[4] $gff->[3]/;
		} else {
		    $newgff->[6] = "+";
		    $newgff->[8] =~ s/^\S+\s+\S+/$gff->[3] $gff->[4]/;
		}
		if ($newgff->[4] < $gff->[4]) { die "Bad co-ords for stacked GFF" }
		$newgff->[5] = 0;
		addtogfflist($newgff,-(++$gffcount));
		my $i;
		for ($i=0;$i<@stackstart;$i++) { last if ($stackstart[$i] > $newgff->[3]) }
		splice @stackstart, $i, 0, $newgff->[3];
	    }
	    my ($stackstring,$frompath);
	  FROMSTACK: while (($stackstring,$frompath) = each %$fromstacks) {
	      my @stack = split /\s+/, $stackstring;
	      if ($stackflag<0 && (@stack==0 || $gff->[9] != shift(@stack))) { next FROMSTACK }
	      my $score = $frompath->[0] + $gffscore;
	      if ($debug{'f'}) { warn "* fromstack: stackstring=($stackstring) stackflag=$stackflag score=$score\n" }
	      if ($stackflag>0) { unshift @stack, $newgff->[9] }
	      $stackstring = "@stack";
	      if (!defined($topath->{$stackstring}) || $score > $topath->{$stackstring}->[0]) {
		  if ($debug{'s'}) { warn "* stack: to=$statename[$to] from=$statename[$from] stack=(@stack) score=$score gffscore=$gffscore link=(@$link) frompos=$frompos pos=$pos gff=(@$gff) topath=$topath\n" }
		  $topath->{$stackstring} = [$score,$frompath,$gffscore,$link,$frompos,$pos,$gff];
	      }
	  }
	}
      }
    }
    for ($state=0;$state<$states;$state++) {
	while ($stateinfo[$state]->[0]->[0] < $pos - $maxlenfrom[$state]) { shift @{$stateinfo[$state]} }
    }
}

sub updatestack {
    my $pos = shift;
    while (@stackstart && $stackstart[0] <= $pos) {
	$stackstart = shift @stackstart;
	updategfflist($stackstart);
	updatestates($stackstart);
    }
}

sub flushgfflist {
    if ($debug{'f'}) { warn "* flushgfflist\n" }
    while (@gfflist) {
	my $pos = $gfflist[0]->[4] + 1;
	updatestack($pos);
	updatestates($pos);
	updategfflist($pos+1);
    }
    undef $lastupdategffpos;
}

sub flushstates {
    my $s;
    if (@{$stateinfo[$endstate]}) {
	if (exists $stateinfo[$endstate]->[-1]->[1]->{""}) {
	    printpath($stateinfo[$endstate]->[-1]->[1]->{""});
	}
    }
    for ($s=0;$s<$states;$s++) { $stateinfo[$s] = [] }
    push @{$stateinfo[$startstate]}, [0, { "" => [0] }];
    undef $lastupdatestatespos;
}

sub printpath {
    my ($path) = @_;
    my @path;
    while (@$path>1) { push @path, $path; $path = $path->[1] }
    if (defined $model->name) { $modelname = $model->name->[0] }
    else { $modelname = $prog }
    my $i;
    for ($i=$#path;$i>=0;$i--) {
	my ($score,$dummypath,$gffscore,$link,$frompos,$pos,$gff) = @{$path[$i]};
	$fromtag = $statename[$link->[0]];
	$totag = $statename[$link->[1]];
	if (defined($print = $link->[9]->{"print"})) {
	    @$gff[3,4,5,8] = ($frompos,$pos-1,$gffscore,"$gff->[8] $modelname($fromtag->$totag)");
	    $gfftext = join("\t",@$gff[0..8]);
	    foreach $printexpr (@$print) {
		if ($printexpr eq "") { print "$gfftext\n" }
		else { print eval $printexpr }
	    }
	}
    }
}

sub substitute {
    my ($s) = @_;
    my $i;
    for ($i=0;$i<@fields;$i++) {
	$s =~ s/\$$fields[$i]/\$gff->[$i]/g;
    }
    $s;
}
