#!/usr/bin/perl

use strict;
use warnings;

#
# Create subset of kmers based on Poisson distribution cutoffs
# Coverage is interpreted as list median
# All counts in resulting whitelist are set to 1
#
# The result is named
# BASENAME_WLEN_whitelist.list
#
# usage:
#   poisson_cutoff.pl LIST|BAM COVERAGE CUTOFF [haploid]
# Use 'median' for coverage to calculate median
#

my $gt4dir = ".";
my $glm = "$gt4dir/glistmaker";
my $glc = "$gt4dir/glistcompare";
my $glq = "$gt4dir/glistquery";
my $b2l = "./bam2list.pl";
my $debug = 1;

if (!defined ($ARGV[2])) {
  printf (STDERR "Usage:\n  poisson_cutoff.pl LIST COVERAGE CUTOFF\n");
  exit (1);
}
my $ncov = 1;
if (defined ($ARGV[3]) && ($ARGV[3] eq "haploid")) {
  $ncov = 0.5;
}

my $list = $ARGV[0];
my $med = $ARGV[1];
my $cutoff = $ARGV[2];
my $wlen = 25;

my @t = split ("/", $list);
@t = split ("\\.", $t[-1]);
my $name = $t[0];
my $suffix = $t[-1];

if ($suffix eq "bam") {
  if ($debug) {
    printf (STDERR "%s seems to be BAM file\n", $name);
  }
  my $cmd = "$b2l $list $wlen 5 tmp_$name";
  if (system ($cmd)) {
    printf (STDERR "Offending command: %s\n", $cmd);
    die ("Cannot create list from BAM");
  }
  $list = "${name}_${wlen}.list";
}

sub poisson {
  my $k = $_[0];
  my $lambda = $_[1];
  my $dk = $k;
  my $dl = $lambda;
  my $p = exp (-$dl);
  while ($k > 0) {
    $p *= $dl;
    $p /= $dk;
    $dk -= 1;
    $k -= 1;
  }
  return $p;
}

if ($med eq "median") {
  my $cmd = "$glq $list -distribution 100";
  if ($debug) {
    printf (STDERR "%s\n", $cmd);
  }
  my $result = qx($cmd);
  if ($debug > 1) {
    printf (STDERR "%s\n", $result);
  }
  my @lines = split ("\n", $result);
  my @v = ();
  foreach (@lines) {
    chomp ($_);
    @t = split ("\t", $_);
    push (@v, $t[1]);
  }
  my $c = 1;
  for (my $i = 1; $i < @v - 1; $i++) {
    printf (STDERR "%s %s %s\n", $v[$i - 1], $v[$i], $v[$i + 1]);
    if (($v[$i - 1] != 0) && ($v[$i - 1] < $v[$i]) && ($v[$i + 1] <= $v[$i])) {
      if ($c == 0) {
        $c += 1;
      } else {
        $med = $i;
        last;
      }
    }
  }
  if ($debug) {
    printf (STDERR "Coverage %s\n", $med);
  }
}
# Find list median
#my $cmd = "$glq $list -median";
#if ($debug) {
#  printf (STDERR "%s\n", $cmd);
#}
#my $result = qx ($cmd);
#@t = split ("\n", $result);
#my @tt = split ("\t", $t[1]);
#$wlen = $tt[1];
#if ($med eq "median") {
#  @tt = split (" ", $t[4]);
#  $med = $tt[5];
#}
#if ($debug) {
#  printf (STDERR "Min %s Max %s Median %s\n", $t[1], $t[3], $med);
#}

# Find cumulative cutoffs
my $sum = poisson (0, $med * $ncov);
my $min = 0;
while ($sum < $cutoff) {
  $min += 1;
  $sum += poisson ($min, $med * $ncov);
}

my $max = $min;
while ($sum < (1 - $cutoff)) {
  $max += 1;
  $sum += poisson ($max, $med * $ncov);
}

if ($debug) {
  printf (STDERR "Cutoffs %s-%s\n", $min, $max);
}

# Create whitelist
my $cmd = "$glc $list $list -i -c $min -o tmp";
if ($debug) {
  printf (STDERR "%s\n", $cmd);
}
if (system ($cmd)) {
  printf (STDERR "Cannot create intersection\n");
  printf (STDERR "Offending command: %s\n", $cmd);
}
$cmd = "$glc tmp_${wlen}_intrsec.list tmp_${wlen}_intrsec.list -i -c $max -o tmp2";
if ($debug) {
  printf (STDERR "%s\n", $cmd);
}
if (system ($cmd)) {
  printf (STDERR "Cannot create intersection\n");
  printf (STDERR "Offending command: %s\n", $cmd);
}
$cmd = "$glc tmp_${wlen}_intrsec.list tmp2_${wlen}_intrsec.list -d -r 1 -o tmp3";
if ($debug) {
  printf (STDERR "%s\n", $cmd);
}
if (system ($cmd)) {
  printf (STDERR "Cannot create intersection\n");
  printf (STDERR "Offending command: %s\n", $cmd);
}
system ("mv tmp3_${wlen}_0_diff1.list ${name}_whitelist_${wlen}.list");
system ("rm tmp_${wlen}_intrsec.list");
system ("rm tmp2_${wlen}_intrsec.list");
