#!/usr/bin/perl

use strict;
use Getopt::Std;

my $usage = "perl summarize_transcripts.pl [-b boundary-file] [-a] [-c] [-h] RNAseg-GFF-files\n\n-a\tOutput also alternative transcripts, not only the dominant ones.\n-b\tFile with segment boundaries from rnaseg_parts.pl.\n-c\tMinimum occurrence frequency of a transcript to be output\n-s\tDo not accumulate occurrence frequencies for split transcripts\n\-h\tHelp/Usage\n\n";

my %opts;
getopts('ahsb:c:',\%opts);

if ($opts{'h'} || $#ARGV < 0) { print $usage; exit;}

my %bound_starts;
my %bound_ends;

my $cutoff = 0;

if (defined $opts{'c'}) {
    $cutoff = $opts{'c'};
}

if (defined $opts{'b'}) {
    
    my $boundary_file = $opts{'b'};
    my $first = -1;
    my $last = -1;
    
    open(BOUND, $boundary_file);
    
    while (<BOUND>) {
        chomp;
        my @a = split(/\t/);
        
        if ($first == -1) {$first = $a[0];}
        
        $bound_starts{$a[0]} = 1;
        $bound_ends{$a[1]} = 1;
        $last = $a[1];
        
    }
    ## Don't treat overall ends as boundaries
    delete($bound_starts{$first});
    delete($bound_ends{$last});
    ###
    
    close(BOUND);
}


my %transcripts;
my %chains;
my %entries_per_file;

foreach my $gff_file (@ARGV) {
    
    open(GFF, $gff_file) or die "Could nopt open file $gff_file.\n";
    
    my $entries = 0;
    my $p_start = -1;
    my $p_end = -1;
    my $p_strand = "";
    my @chain = ();
    my %current_trs = ();
    my %current_chns = ();
    
    while(<GFF>) {
        if (/^(\d+) change points/) {$entries++; next;}
        chomp;
        my @a = split(/\t/);
        
        if ($a[2] eq "non_transcript") {next;}
        
        ## Ignore Transcripts starting/ending at a boundary of an analysed part
        if ($bound_starts{$a[3]} == 1 || $bound_ends{$a[4]} == 1) {
            #$transcripts{$a[3]."_".$a[4]."_".$a[6]}{'at_cut'} = 1;
            #print STDERR "Transcript ".$a[3]."_".$a[4]."_".$a[6]." starts/ends at a part boundary! It will be ignored!\n";
            next;
        }
        ########################################################################
        
        if (defined $transcripts{$a[3]."_".$a[4]."_".$a[6]}) {
            $transcripts{$a[3]."_".$a[4]."_".$a[6]}{'count'}+=1;
        } else {
            $transcripts{$a[3]."_".$a[4]."_".$a[6]} = {'count' => 1, 'start' => $a[3], 'end' =>$a[4], 'strand' => $a[6], 'length' => $a[4]-$a[3]+1, 'is_sub' => 0, 'minor' => 0, 'entries' => 0};
        }
        
        $current_trs{$a[3]."_".$a[4]."_".$a[6]} = 1;
        
        if ($p_start == -1) {
            $p_start = $a[3];
            $p_end = $a[4];
            $p_strand = $a[6];
        } else {
            if ( ($a[3] == $p_end + 1 || $a[4] == $p_start - 1) && $p_strand eq $a[6]) {
                
                if ($a[6] eq '+') {
                    
                    push(@chain,$a[3]."_".$a[4]."_".$a[6]);
                    $p_end = $a[4];
                    
                } else {
                    
                    unshift(@chain,$a[3]."_".$a[4]."_".$a[6]);
                    $p_start = $a[3];
                    
                }
            } else {
                $p_start = $a[3];
                $p_end = $a[4];
                $p_strand = $a[6];
                
                for (my $i = 0; $i <= $#chain; $i++) {
                    
                    for (my $j = $i+1; $j <= $#chain; $j++) {
                        
                        my ($start,$ei,$sti)  = split(/_/,$chain[$i]);
                        my ($si,$end,$strand) = split(/_/,$chain[$j]);
                        
                        $current_chns{$start."_".$end."_".$strand} = 1;
                        
                        if (defined $chains{$start."_".$end."_".$strand}) { $chains{$start."_".$end."_".$strand}{'count'}+=1; }
                        else {
                            $chains{$start."_".$end."_".$strand} = {'count' => 1, 'start' => $start, 'end' => $end, 'strand' => $strand, 'entries' =>0};
                        }
                        
                        for (my $k=$i; $k<=$j; $k++) {
                            
                            $chains{$start."_".$end."_".$strand}{'subs'}{$chain[$k]} = 1;
                            
                        }
                        
                    }
                    
                }
                
                @chain = ();
            }
        }
        
    }
    close(GFF);
    
    foreach my $t (keys %current_trs) {
        $transcripts{$t}{'entries'} += $entries;
    }
    foreach my $c (keys %current_chns) {
        $chains{$c}{'entries'} += $entries;
    }
    undef(%current_trs);
    undef(%current_chns);
}

foreach my $t (keys %transcripts) {
    $transcripts{$t}{'count'} = $transcripts{$t}{'count'}/$transcripts{$t}{'entries'};
}

foreach my $c (keys %chains) {
    if (not defined $transcripts{$c}) {delete $chains{$c}; next;}
    if (not defined $opts{'s'}) {
        $chains{$c}{'count'} = $chains{$c}{'count'}/$chains{$c}{'entries'};
    } else {
        $chains{$c}{'count'} = 0;
    }
}

#### Merge transcripts overlapping to >99%

my %to_delete;

foreach my $t (keys %transcripts) {
    
    if (defined $to_delete{$t}) { next; }
    
    foreach my $s (keys %transcripts) {
        
        if ((defined $to_delete{$s}) || (defined $to_delete{$t}) || $s eq $t || $transcripts{$t}{'strand'} ne $transcripts{$s}{'strand'}) { next; }
        
        ### $s must not be part of a chain of $t and vice versa
        if ($chains{$t}{'subs'}{$s} || $chains{$s}{'subs'}{$t} ) {next;}
        ####
    
        my $ol = &olap($transcripts{$t}{'start'},$transcripts{$t}{'end'},$transcripts{$s}{'start'},$transcripts{$s}{'end'});
                
        if ($ol >= 0.99 && $transcripts{$t}{'count'}+$chains{$t}{'count'} > $transcripts{$s}{'count'}+$chains{$s}{'count'}) {
            
            if ($transcripts{$t}{'count'}+$chains{$t}{'count'}+$transcripts{$s}{'count'}+$chains{$s}{'count'} > 1000) {
                print STDERR "Merge: $t and $s ".($transcripts{$t}{'count'}+$chains{$t}{'count'}+$transcripts{$s}{'count'}+$chains{$s}{'count'})."\n";
                print STDERR join("\t",($transcripts{$t}{'count'},$chains{$t}{'count'},$transcripts{$s}{'count'},$chains{$s}{'count'}))."\n";
                
                foreach (keys %{$chains{$s}{'subs'}}) { print STDERR $_."\n"; }
                print STDERR "--\n";
                foreach (keys %{$chains{$t}{'subs'}}) { print STDERR $_."\n"; }
                
            }

            $transcripts{$t}{'count'} = ($transcripts{$t}{'count'} * $transcripts{$t}{'entries'} + $transcripts{$s}{'count'} * $transcripts{$s}{'entries'}) / ($transcripts{$t}{'entries'} + $transcripts{$s}{'entries'});
            $transcripts{$t}{'entries'} = ($transcripts{$t}{'entries'} + $transcripts{$s}{'entries'})/2;
            
            if (not &is_sub($t,$s)) {
                $chains{$t}{'count'} += $chains{$s}{'count'};
            }
            
            foreach my $p (keys %{ $chains{$s}{'subs'} }) {
                $chains{$t}{'subs'}{$p} = 1;
            }
            
            $to_delete{$s} = 1;
            
        }
        
    }
    
}

foreach my $s (keys %to_delete) {
    delete $transcripts{$s};
    delete $chains{$s};
    
    foreach my $c (keys %chains) {
        delete $chains{$c}{'subs'}{$s};
    }
}

##### Identify overlapping transcripts that are mutually exclusive

my %overlaps;
foreach my $t (keys %transcripts) {
    
    foreach my $s  (keys %transcripts) {
        
        if ($s eq $t) { next; }
        
        if ($transcripts{$t}{'start'} <= $transcripts{$s}{'end'} && $transcripts{$t}{'end'} >= $transcripts{$s}{'start'} && $transcripts{$t}{'strand'} eq $transcripts{$s}{'strand'} && ($transcripts{$t}{'count'}+$chains{$t}{'count'} > $transcripts{$s}{'count'}+$chains{$s}{'count'} || ($transcripts{$t}{'count'}+$chains{$t}{'count'} == $transcripts{$s}{'count'}+$chains{$s}{'count'} && $transcripts{$t}{'length'} > $transcripts{$s}{'length'}) ) ) {
            
            $overlaps{$t}{$s} = 1;
        }
        
    }
    
}

#######

foreach my $t ( sort {$transcripts{$b}{'count'}+$chains{$b}{'count'} <=> $transcripts{$a}{'count'}+$chains{$a}{'count'}} keys %overlaps) {
    
    foreach my $s (keys %{ $overlaps{$t} }) {

        delete($overlaps{$s});
        
    }
    
}

foreach my $t (keys %transcripts) {
    
    foreach my $s (keys %{ $overlaps{$t} }) {
        
        $transcripts{$s}{'minor'} = 1;
              
    }
    
}

foreach my $c (keys %chains) {
    
    if (defined $transcripts{$c} && not $transcripts{$c}{'minor'} == 1) {
        
        foreach my $s (keys %{ $chains{$c}{'subs'} }) {
            
            $transcripts{$s}{'is_sub'} = 1;
            
        }
    }
    
}



foreach my $t (sort {$transcripts{$a}{'start'} <=>  $transcripts{$b}{'start'}} keys %transcripts) {
    
    my $red = 255;
    my $green = 255;
    my $blue = 255;
    
    if ($transcripts{$t}{'is_sub'} == 1) {next;}
    if ($transcripts{$t}{'minor'} == 1 && not $opts{'a'}) {next;}
    
    my $score = $transcripts{$t}{'count'}+$chains{$t}{'count'};

    if ($score < $cutoff) {next;}
    
    $green = int(255 - 255*$score);
    $blue = int(255 -255*$score);
    
    if ($transcripts{$t}{'minor'} == 1) {

        print join("\t",("SeqID","RNAseg","alt_transcript",$transcripts{$t}{'start'},$transcripts{$t}{'end'},$score,$transcripts{$t}{'strand'},".","colour=$green 255 $blue"))."\n";
        foreach my $s (sort { $transcripts{$a}{'start'} <=>  $transcripts{$b}{'start'} } keys %{$chains{$t}{'subs'}}) {
            $score = $transcripts{$s}{'count'}+$chains{$s}{'count'};
            $green = int(255 - 255*$score);
            $blue = int(255 -255*$score);
            print join("\t",("SeqID","RNAseg","alt_subtranscript",$transcripts{$s}{'start'},$transcripts{$s}{'end'},$score,$transcripts{$s}{'strand'},".","colour=$green 255 $blue"))."\n";
        }

    } else {

        print join("\t",("SeqID","RNAseg","transcript",$transcripts{$t}{'start'},$transcripts{$t}{'end'},$score,$transcripts{$t}{'strand'},".","colour=$red $green $blue"))."\n";
        foreach my $s (sort { $transcripts{$a}{'start'} <=>  $transcripts{$b}{'start'} } keys %{$chains{$t}{'subs'}}) {
            $score = $transcripts{$s}{'count'}+$chains{$s}{'count'};
            $green = int(255 - 255*$score);
            $blue = int(255 -255*$score);
            print join("\t",("SeqID","RNAseg","subtranscript",$transcripts{$s}{'start'},$transcripts{$s}{'end'},$score,$transcripts{$s}{'strand'},".","colour=$red $green $blue"))."\n";
        }
    
    }
    
}

exit;

sub is_sub {

    my $t = shift;
    my $s = shift;
    
    my %links = ();
    
    foreach my $i (keys %{ $chains{$t}{'subs'} }) {
        if (not defined $chains{$s}{'subs'}{$i}) {
            my ($start,$end,$strand) = split(/_/,$i);
            $links{$start}{$end} = 1;
        }
    }

    foreach my $i (keys %{ $chains{$s}{'subs'} }) {
        if (not defined $chains{$t}{'subs'}{$i}) {
            my ($start,$end,$strand) = split(/_/,$i);
            $links{$start}{$end} = 1;
        }
    }
    
    my ($t_start,$t_end,$t_strand) = split(/_/,$t);
    my ($s_start,$s_end,$s_strand) = split(/_/,$s);
    
    my $start_match = 0;
    
    if ($t_start != $s_start) {
   
        my $i = $s_start;
        my $j = $t_start;
        
        if ($t_start < $s_start) {
            $i = $t_start;
            $j = $s_start;
        } 
        
        my @ends = &get_paths($i,\%links);

        foreach (@ends) {
            if ( $_ == $j-1) {
                $start_match = 1;
            }
        }
    
    } else {
        $start_match = 1;
    }
    
    my $end_match = 0;
    
    if ($t_end != $s_end) {

        my $i = $s_end;
        my $j = $t_end;
        
        if ($t_end < $s_end) {
            $i = $t_end;
            $j = $s_end;
        }
        
        my @ends = &get_paths($i+1,\%links);
        
        foreach (@ends) {
            if ( $_ == $j) { $end_match = 1; }
        }
        
    } else {
        $end_match = 1;
    }
    
    
    return ($start_match && $end_match);
}

sub get_paths {
    my $i = shift;
    my $lnk = shift;
    my %hash = %{ $lnk };
    
    my @res = ();
    
    my @a = keys %{ $hash{$i} };
	
    if ($#a >= 0) {
        foreach (@a) {
            push(@res,$_);
            push(@res,&get_paths(($_ + 1),$lnk));
        }
    } else {
        push(@res,($i-1))
    }
    
    return @res;

}

sub olap {
    my $s1 = shift;
    my $e1 = shift;
    my $s2 = shift;
    my $e2 = shift;
    
    my $ol_start = $s1;
    my $ol_end = $e1;
    
    if ($s2 > $s1) {$ol_start = $s2;}
    if ($e2 < $e1) {$ol_end = $e2;}
    
    my $longer = $e2-$s2+1;
    
    if ($e1-$s1+1 > $longer) {$longer = $e1-$s1+1;}
    
    if ($s1<=$e2 && $e1 >= $s2) {
        return (($ol_end-$ol_start+1)/$longer);
    } else {
        return 0;
    }
    
}
