#!/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:j:c:s:o:d:r:n:a:",\%opts);
die usage() unless ($opts{i} and $opts{j} and $opts{o} and $opts{s} and $opts{a});

my $sam_assembly     = $opts{i};
my $sam_reference    = $opts{j};
my $assembly_file    = $opts{s};
my $output           = $opts{o};
my $n_cutoff         = $opts{n}; $n_cutoff = 10  unless defined($n_cutoff);
my $p_cutoff         = $opts{c}; $p_cutoff = 0.5 unless defined($p_cutoff);
my $tax_str          = $opts{a};

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

open(TMP, $assembly_file) || die "can not open $assembly_file";
my @assembly_ids = ();
my %assembly_2_len = ();
my $seq_id;
while($ll=<TMP>){
  chop($ll);
  if ($ll =~ /^>/) {
    $seq_id = substr($ll,1);
    $seq_id =~ s/\s.+$//;
    $assembly_2_len{$seq_id} = 0;
    push(@assembly_ids, $seq_id);
  }
  else {
    $ll =~ s/\s+$//;
    $assembly_2_len{$seq_id}  += length($ll);
  }
}
close(TMP);

open(TMP, $sam_assembly) || die "can not open $sam_assembly";
my %assembly_reads_mapped = ();
while($ll=<TMP>){
  if ($ll =~ /^\@/) { #### headers
    next;
  }
  else { #### alignment
    chop($ll);
    my @lls = split(/\t/,$ll);
    my $id  = $lls[0];
    my $rid = $lls[2];    if ($rid eq "*") {  next; }
    my $FLAG = $lls[1];
    next unless ( $FLAG & 0x0040 ); #### count R1 only 

    if (not defined($assembly_reads_mapped{$rid})) {
      $assembly_reads_mapped{$rid} = [];
    }
    push(@{$assembly_reads_mapped{$rid}}, $id);
  } #### alignment section
}
close(TMP);

open(TMP, $sam_reference) || die "can not open $sam_reference";
my %t_taxids = ();
my $last_id = "";
my %read_2_taxids = ();
while($ll=<TMP>){
  if ($ll =~ /^\@/) { #### headers
    next;
  }
  else { #### alignment
    chop($ll);
    my @lls = split(/\t/,$ll);
    my $id  = $lls[0];
    my $rid = $lls[2];    if ($rid eq "*") {  next; }
    my $tid = "";
    if ($rid =~ /$tax_str\|(\d+)/) { $tid = $1; } else { next; }

    my $FLAG = $lls[1];
    next unless ( $FLAG & 0x0040 ); #### count R1 only 

    if    (($id  ne $last_id) and $last_id ) {
      my @t_taxids = keys %t_taxids;
      if ($#t_taxids+1 <= $n_cutoff) {
        $read_2_taxids{$last_id} = [@t_taxids];
      }
      %t_taxids = ();
    }

    $t_taxids{$tid} = 1;
    $last_id = $id;
  } #### alignment section
}
close(TMP);
    if    (($id  ne $last_id) and $last_id ) {
      my @t_taxids = keys %t_taxids;
      if ($#t_taxids+1 <= $n_cutoff) {
        $read_2_taxids{$last_id} = [@t_taxids];
      }
    }


my $fh = "STDOUT";
if ($output) {
  open(OUT, ">$output") || die "Can not write to $OUT";
  $fh = "OUT";
}
foreach $i (@assembly_ids) {
  my $taxid = "Unknown";
  my $p = 1;
  my $n_reads = 0;
  if ($assembly_reads_mapped{$i}) {
    my @reads = @{$assembly_reads_mapped{$i}};
    $n_reads = $#reads+1;
    my %mapped_tids = ();

    foreach $j (@reads) {
      next unless ($read_2_taxids{$j});
      my @tids = @{$read_2_taxids{$j}};
      my $n_tids = $#tids + 1;
      foreach $k (@tids) {
        $mapped_tids{$k} += 1.0 / $n_tids / $n_reads;
      }
    }
    my @mapped_tids = keys %mapped_tids;
       @mapped_tids = sort { $mapped_tids{$b} <=> $mapped_tids{$a} } @mapped_tids;

    if ($mapped_tids{ $mapped_tids[0] } >= $p_cutoff) {
      $taxid = $mapped_tids[0];
      $p = $mapped_tids{ $mapped_tids[0] };
    }
  }

  print OUT "$i\t$taxid\t$p\t$n_reads\t$assembly_2_len{$i}\n";
}
close(OUT) if ($output);


sub usage {
<<EOD
Given a sam file of reads to assembles and 
  another sam file of reads to reference genome generated by 
  bwa mem reference R1.fa_or_fq R2.fa_or_fq 
and (better) filtered by sam-filter-top-pair.pl

This script uses the reads to reference and reads to assembly
mapping information to bin the assembly sequences

usage:
  $script_name -i reads-to-assembly-mapping.sam -j reads-to-reference-mapping.sam -o output -s assembly_fasta_file -a tid

  options
    -i reads-to-assembly-mapping.sam
    -j reads-to-reference-mapping.sam
    -o output
    -s fasta file of assembly
    -n number alignment cutoff, default 10, 
       if one read is mapped to multiple taxids (> the cutoff), this read is skipped.
    -c probability cutoff, default 0.5
       each read (R) has a probality to a genome (G), if a reads mapped to
       N multiple genomes, the probability is 1/N
       for a assembly (S) , the probability for it belong to the same species of 
       a genome (G) is sum of P(R,G) / number of reads mapped to S
    -a taxid string, the name of reference genome can be ...sequence_name|tid|12345,
       here 12345 is the taxon ID. in this case, use -a tid so that the script can 
       parse the taxon ID 
EOD
}
########## END usage

