#!/usr/bin/perl
## =========================== NGS tools ==========================================
## NGS tools for metagenomic sequence analysis
## May also be used for other type NGS data analysis
##
##                                      Weizhong Li, UCSD
##                                      liwz@sdsc.edu
## http://weizhongli-lab.org/
## ================================================================================


use Getopt::Std;
getopts("i:o:a:t:r:N:c:P:s:",\%opts);
die usage() unless ($opts{o});


my $sam_file    = $opts{i}; ### input depth file, generated by samtools depth *bam
   $sam_file    = "-" unless ($sam_file);
my $step        = $opts{s};
   $step        = 20 unless ($step);
my $output      = $opts{o};
my $step2       = $step*2 +1;

my ($i, $j, $k, $ll, $cmd);

my $fh;
if ($sam_file eq "-") { $fh = "STDIN";}
else {
  open(TMP, $sam_file) || die "can not open $sam_file";
  $fh = "TMP";
}
open(OUT, "> $output") || die "can not write to $output";


my $last_pos;  #### last pos from last id
my $last_id;
my %node_total = ();
my %node_count = ();
my $num_sample = 0;
my @empty = ();

while($ll=<$fh>) {
  chop($ll);
  my ($id, $pos, @depth) = split(/\t/, $ll);
  if (not $num_sample) {
    $num_sample = ($#depth + 1); 
    for ($i=0; $i<$num_sample; $i++) {
      push(@empty, 0);
    }
  }

  #### a new seq
  if ($id ne $last_id) {
    foreach $t_node (keys %node_total) {
      next unless ($t_node <= $last_pos);
      print OUT "$last_id\t$t_node";
      for ($i=0; $i<$num_sample; $i++) {
        $v = $node_total{$t_node}->[$i] / $node_count{$t_node};
        $v = int($v*10) / 10;
        print OUT "\t$v";
      } 
      print OUT "\n";
      delete( $node_total{$t_node} );
      delete( $node_count{$t_node} );
    } 
    $last_pos = "";
    $last_id = "";
    %node_total = ();
    %node_count = ();
  }

  $i = $pos % $step;
  $last_node = $pos - $i;
  $next_node = $last_node + $step;

  foreach $t_node ($last_node, $next_node) {
    next if ($t_node <=0 );
    if (not defined( $node_total{$t_node} )) {
      $node_total{$t_node} = [@empty];
      $node_count{$t_node} = 0;
    }
    for ($i=0; $i<$num_sample; $i++) {
      $node_total{$t_node}->[$i] += $depth[$i];
    }
    $node_count{$t_node}++;

  }

  foreach $t_node (keys %node_total) {
    next unless ($t_node < $last_node);
    print OUT "$last_id\t$t_node";
    for ($i=0; $i<$num_sample; $i++) {
      $v = $node_total{$t_node}->[$i] / $node_count{$t_node};
      $v = int($v*10) / 10;
      print OUT "\t$v";
    }
    print OUT "\n";
    delete( $node_total{$t_node} );
    delete( $node_count{$t_node} );
  }
  
  $last_id = $id;
  $last_pos = $pos;
}
  if ($last_id) {
    foreach $t_node (keys %node_total) {
      next unless ($t_node <= $last_pos);
      print OUT "$last_id\t$t_node";
      for ($i=0; $i<$num_sample; $i++) {
        $v = $node_total{$t_node}->[$i] / $node_count{$t_node};
        $v = int($v*10) / 10;
        print OUT "\t$v";
      } 
      print OUT "\n";
      delete( $node_total{$t_node} );
      delete( $node_count{$t_node} );
    } 
  }



if ($sam_file ne "-") { close(TMP); }
close(OUT);


sub usage {
<<EOD

Given a coverage file (generated by "samtools depth file1.bam file2.bam ...") and a window size (e.g. 20bp),
this scripts prints the average coverage at each 20th base, averaged over the window.
The average coverage of a base is averaged over 19 bases before and 19 bases after this position
If there are gap in the input BAM file, the average coverage of a base in averaged over
the available bases in the window.


usage:
  $0  -i input_coverage_file -o output

  options
    -i input depth file, generated by samtools depth *bam
    -o output file
    -s window size, default 20
EOD
}
########## END usage

