#!/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/
## ================================================================================


## copied from  assembly-binning.pl 
## but just bin the possible contigs from host
use Getopt::Std;
getopts("i:j:c:s:o:d:r:n:a:t:T:",\%opts);
die usage() unless ($opts{i} and $opts{j} and $opts{o} and $opts{s});

my $sam_assembly     = $opts{i};
my $sam_reference    = $opts{j};
my $assembly_file    = $opts{s};
my $output           = $opts{o};
my $p_cutoff         = $opts{c}; $p_cutoff = 0.8 unless defined($p_cutoff);
my $host_map_score_cutoff = $opts{T};
   $host_map_score_cutoff = 75 unless ($host_map_score_cutoff);


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

my $fh = "STDOUT";
if ($output) {
  open(OUT, ">$output") || die "Can not write to $OUT";
  $fh = "OUT";
}

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);
  my @lls = split(/\t/, $ll);
  if (($lls[1] =~ /\d+/) or ($lls[3] =~ /\d+/) or ($lls[1] =~ /contaminant/i) ) { #### valid taxid or contaminant
#    print $fh "$ll\n";
    next;
  }
  push(@assembly_ids, $lls[0]);
  $assembly_2_len{$lls[0]} = $lls[7];
}
close(TMP);


open(TMP, $sam_assembly) || die "can not open $sam_assembly";
my %assembly_reads_mapped = ();
my $assembly_mapped_reads = ();
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; }
    next unless ($assembly_2_len{$rid}); #### only the unbinned assembly

    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);
    $assembly_mapped_reads{$id} = 1;
  } #### alignment section
}
close(TMP);


open(TMP, $sam_reference) || die "can not open $sam_reference";
my %read_from_host = ();
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 $score;
    for ($j=11; $j<=$#lls; $j++) { if ($lls[$j] =~ /AS:i:(.+)/) { $score = $1; last; } }
    next unless (defined($score));
    next unless ($score >= $host_map_score_cutoff);

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

    $read_from_host{$id} = 1;
  } #### alignment section
}
close(TMP);


open(TMP, $assembly_file) || die "can not open $assembly_file";
while($ll=<TMP>){
  chop($ll);
  my @lls = split(/\t/, $ll);
  if (($lls[1] =~ /\d+/) or ($lls[3] =~ /\d+/) or ($lls[1] =~ /contaminant/i) ) { #### valid taxid or contaminant
    print $fh "$ll\n";
    next;
  }
  $i = $lls[0];

  my $p = 0;
  my $n_reads = 0;
  if ($assembly_reads_mapped{$i}) {
    my @reads = @{$assembly_reads_mapped{$i}};
    $n_reads = $#reads+1;
    my $host_reads = 0;
    foreach $j (@reads) {
      next unless ($read_from_host{$j});
      $host_reads++;
    }
    $p = $host_reads / $n_reads;
  }

  if ($p >= $p_cutoff) {
    print $fh "$i\tHost\tHost\tHost\tHost\t$p\t$n_reads\t$assembly_2_len{$i}\n";
  }
  else {
    print $fh "$ll\n";
  }
}
close(TMP);
close(OUT) if ($output);


sub usage {
<<EOD
Given a sam file of reads to assembles and 
  another sam file of reads to host 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 host and reads to assembly
mapping information to bin the assembly sequences

assembly-binning-taxon.pl should be run at first and this script
run after that

usage:
  $script_name -i reads-to-assembly-mapping.sam -j reads-to-reference-mapping.sam -o output -s assembly_binning file

  options
    -i reads-to-assembly-mapping.sam
    -j reads-to-reference-mapping.sam
    -o output
    -s assembly-binning file, assembly that are already binned to mapped reference genome
       will be 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
    -T host mapping cutoff score, default 75
       only reads mapped to host at this score are considered to be from host
EOD
}
########## END usage

