#!/usr/bin/env perl
# SuperRead pipeline
# Copyright (C) 2012  Genome group at University of Maryland.
# 
# This program is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# This script reads the the config file and runs the full super reads
# pipeline.  An example of configuration file is in MasurcaConf.pm

use strict;
use warnings;
use Env qw(@PATH @LD_LIBRARY_PATH);
use File::Spec;
use Getopt::Long;

use FindBin qw($RealBin);
use lib $RealBin;

use MasurcaConf qw(fail $config_file);
use MasurcaCommon qw($rerun_pe $rerun_sj $list_frg_files);
use MasurcaSuperReads;
use MasurcaPaths;

#global variables
my (@pe_files, @sj_files);
my ($list_pe_files, $list_jump_files);

# Those seem to be constants
my $WINDOW                  = 10;
my $MAX_ERR_PER_WINDOW      = 3;

# Options to masurca script
my $assembly_script = "assemble.sh";
my $generate_default;
my $skip_checking;
my @prep_path;
my @prep_ld;
my $help;

my $usage = "USAGE: $0 <config_file>";
my $help_str = <<"EOS" ;
Create the assembly script from a MaSuRCA configuration file. An
sample configuration file can be generated with the -g switch. The
assembly script will run the assembly proper.

Options:
 -o, --output              Assembly script ($assembly_script)
 -g, --generate            Generate example configuration file
 -p, --path                Prepend to PATH in assembly script
 -l, --ld-library-path     Prepend to LD_LIBRARY_PATH in assembly script
     --skip-checking       Skip checking availability of other executables
 -h, --help                This message
EOS

GetOptions("o|output=s"          => \$assembly_script,
           "g|generate"          => \$generate_default,
           "p|path=s"            => \@prep_path,
           "l|ld-library-path=s" => \@prep_ld,
           "skip-checking"       => \$skip_checking,
           "h|help"              => \$help) or fail $usage;

if($help) {
  print "$help_str";
  exit(0);
}
fail $usage if @ARGV != 1;
$config_file = $ARGV[0];

if($generate_default) {
  fail "Refusing to overwrite existing configuration file '$config_file'" if(-e $config_file);
  open(FILE, ">", $config_file) or fail "Error opening configuration file '$config_file': $!";
  print(FILE MasurcaConf::default_config);
  close(FILE);
  exit(0);
}

sub script_done {
  close(FILE);
  chmod 0755, $assembly_script;
  print("done.\n",
        "execute $assembly_script to run assembly\n");
  exit(0);
}

sub check_exec {
  for my $e (@_) {
    system("$e --help >/dev/null 2>&1") == 0 or fail "$e not found or failed to run";
    print "$e OK\n";
  }
}

# Parse configuration file
my %config = MasurcaConf::parse($config_file);
fail("no read data files specified") if(@{$config{PE_INFO}} + @{$config{JUMP_INFO}} + @{$config{OTHER_INFO}} == 0);

if(@prep_path) {
  unshift(@PATH, @prep_path);
} else {
  unshift(@PATH, $RealBin);
}
unshift(@LD_LIBRARY_PATH, @prep_ld) if @prep_ld;

unless($skip_checking) {
  print "Verifying PATHS...\n";
  my @progs=("jellyfish","createSuperReadsForDirectory.perl");
  &check_exec(@progs);
}

fail ("no PE data specified") if(@{$config{PE_INFO}} == 0);

print("creating script file for the actions...");

my $config_abs_path = File::Spec->rel2abs($config_file);
my $cmd_abs_path = File::Spec->rel2abs(__FILE__);

open(FILE,">", $assembly_script) or fail "Can't open output file '$assembly_script': $!";
print FILE <<"EOS";
#!/bin/bash

# $assembly_script generated by masurca
CONFIG_PATH="$config_abs_path"
CMD_PATH="$cmd_abs_path"

# Test that we support <() redirection
(eval "cat <(echo test) >/dev/null" 2>/dev/null) || {
  echo >&2 "ERROR: The shell used is missing important features."
  echo >&2 "       Run the assembly script directly as './\$0'"
  exit 1
}

# Parse command line switches
while getopts ":rc" o; do
  case "\${o}" in
    c)
    echo "configuration file is '\$CONFIG_PATH'"
    exit 0
    ;;
    r)
    echo "Rerunning configuration"
    exec perl "\$CMD_PATH" "\$CONFIG_PATH"
    echo "Failed to rerun configuration"
    exit 1
    ;;
    *)
    echo "Usage: \$0 [-r] [-c]"
    exit 1
    ;;
  esac
done
EOS

print FILE <<'EOS';
set +e
# Set some paths and prime system to save environment variables
save () {
  (echo -n "$1=\""; eval "echo -n \"\$$1\""; echo '"') >> environment.sh
}
GC=
RC=
NC=
if tty -s < /dev/fd/1 2> /dev/null; then
  GC='\e[0;32m'
  RC='\e[0;31m'
  NC='\e[0m'
fi
log () {
  d=$(date)
  echo -e "${GC}[$d]${NC} $@"
}
fail () {
  d=$(date)
  echo -e "${RC}[$d]${NC} $@"
  exit 1
}
signaled () {
  fail Interrupted
}
trap signaled TERM QUIT INT
rm -f environment.sh; touch environment.sh

# To run tasks in parallel
run_bg () {
  semaphore -j $NUM_THREADS --id masurca_$$ -- "$@"
}
run_wait () {
  semaphore -j $NUM_THREADS --id masurca_$$ --wait
}
EOS
    ;

if(@prep_path) {
  print FILE "export PATH=\"", join(":", @prep_path), ":\$PATH\"\n";
} else {
  print FILE "export PATH=\"$RealBin:$RealBin/../CA/Linux-amd64/bin:\$PATH\"\n";
}
print FILE "save PATH\n";

if(@prep_ld) {
  print FILE "export LD_LIBRARY_PATH=\"", join(":", @prep_ld), ":\$LD_LIBRARY_PATH\"\n";
  print FILE "save LD_LIBRARY_PATH\n";
}
print FILE "export PERL5LIB=", $RealBin . "/" . MasurcaPaths::rel_ext_bin, "\${PERL5LIB:+:\$PERL5LIB}\n";
print FILE "save PERL5LIB\n";
print FILE "NUM_THREADS=$config{NUM_THREADS}\n";
print FILE "save NUM_THREADS\n";

###renaming reads###

@pe_files        = MasurcaCommon::rename_reads(*FILE, "pe", @{$config{PE_INFO}});
print FILE "run_wait\n";
$list_pe_files   = join(" ", @pe_files);

###done renaming reads###
print FILE "\n";

###compute minimum and average PE read length and gc content, and kmer size###
print FILE "head -q -n 40000  $list_pe_files | grep --text -v '^\+' | grep --text -v '^\@' > pe_data.tmp\n";
print FILE "export PE_AVG_READ_LENGTH=`awk '{if(length(\$1)>31){n+=length(\$1);m++;}}END{print int(n/m)}' pe_data.tmp`\n";
print FILE "save PE_AVG_READ_LENGTH\n";
print FILE "echo \"Average PE read length \$PE_AVG_READ_LENGTH\"\n";

if(uc($config{KMER}) eq "AUTO"){
#here we have to estimate gc content and recompute kmer length for the graph
     MasurcaCommon::estimate_optimal_kmer(*FILE, $list_pe_files,"KMER");
     print FILE "save KMER\n";
     print FILE "echo \"choosing kmer size of \$KMER for the graph\"\n";
}else{
     print FILE "KMER=$config{KMER}\n";
}

###done###

###Jellyfish###

MasurcaCommon::get_MIN_Q_CHAR(*FILE, $list_pe_files);

print FILE "JF_SIZE=`ls -l *.fastq | awk '{n+=\$5}END{s=int(n/50); if(s>$config{JF_SIZE})print s;else print \"$config{JF_SIZE}\";}'`\n";
print FILE "save JF_SIZE\n";
print FILE "perl -e '{if(int('\$JF_SIZE')>$config{JF_SIZE}){print \"WARNING: JF_SIZE set too low, increasing JF_SIZE to at least '\$JF_SIZE', this automatic increase may be not enough!\\n\"}}'\n"; 
if(not(-e "quorum_mer_db.jf") || $rerun_pe==1){
    print FILE "log Creating mer database for Quorum.\n";
    MasurcaCommon::count_kmers_for_quorum(*FILE, $list_pe_files, $config{NUM_THREADS}, 5);
    $rerun_pe=1;
    $rerun_sj=1;
}

###done Jellyfish###
print FILE "\n";
###error correct PE###

my $homo_polymer_flag = $config{DO_HOMOPOLYMER_TRIM} ? "--homo-trim $config{TRIM_PARAM}" : "";
my $mmap_flag = $config{NO_MMAP} ? "-M" : "";

if(not(-e "pe.cor.fa")||$rerun_pe==1){
  print FILE <<"EOS";
log Error correct PE.

quorum_error_correct_reads  -q \$((MIN_Q_CHAR + $config{RELIABLE_Q_PARAM})) --contaminant=$RealBin/../share/adapter.jf -m $config{KMER_COUNT_THRESHOLD} -s 1 -g 1 -a 1 -t $config{NUM_THREADS} -w $WINDOW -e $MAX_ERR_PER_WINDOW $mmap_flag $homo_polymer_flag quorum_mer_db.jf $list_pe_files --no-discard -o pe.cor --verbose 1>quorum.err 2>&1 || {
  mv pe.cor.fa pe.cor.fa.failed && fail Error correction of PE reads failed. Check pe.cor.log.
}
EOS
    $rerun_pe=1;
}

###done error correct PE###
print FILE "\n";
###estimate genome size###
my $k_u_arg="pe.cor.fa";

print FILE "log Estimating data size.\n";
if(not(-e "k_u_hash_0")||$rerun_pe==1||$rerun_sj==1){
    print FILE "jellyfish count -m 31 -t $config{NUM_THREADS} -C -s \$JF_SIZE -o k_u_hash_0 $k_u_arg\n";
}

print FILE "export ESTIMATED_GENOME_SIZE=`jellyfish histo -t $config{NUM_THREADS} -h 1 k_u_hash_0 | tail -n 1 |awk '{print \$2}'`\n";
print FILE "save ESTIMATED_GENOME_SIZE\n";
print FILE "echo \"Estimated data size: \$ESTIMATED_GENOME_SIZE\"\n";

####done estimate genome size###
print FILE "\n";
###build k-unitigs###

if(not(-e "guillaumeKUnitigsAtLeast32bases_all.fasta")||$rerun_pe==1||$rerun_sj==1){
    print FILE "log Creating k-unitigs with k=\$KMER\n";
    print FILE "create_k_unitigs_large_k2 -q 1 -c \$((\$KMER-1)) -t $config{NUM_THREADS} -m \$KMER -n \$((\$ESTIMATED_GENOME_SIZE*2)) -l \$KMER $k_u_arg  | grep --text -v '^>' | perl -ane '{\$seq=\$F[0]; \$F[0]=~tr/ACTGactg/TGACtgac/;\$revseq=reverse(\$F[0]); \$h{(\$seq ge \$revseq)?\$seq:\$revseq}=1;}END{\$n=0;foreach \$k(keys \%h){print \">\",\$n++,\" length:\",length(\$k),\"\\n\$k\\n\"}}' > guillaumeKUnitigsAtLeast32bases_all.fasta.tmp && mv guillaumeKUnitigsAtLeast32bases_all.fasta.tmp guillaumeKUnitigsAtLeast32bases_all.fasta\n";
    $rerun_pe=1;
}

###done build k-unitigs###
print FILE "\n";
###super-reads###
  print FILE "createSuperReadsForDirectory.perl -low-memory -l \$KMER -mean-and-stdev-by-prefix-file meanAndStdevByPrefix.pe.txt -kunitigsfile guillaumeKUnitigsAtLeast32bases_all.fasta -t $config{NUM_THREADS} -mikedebug work1 pe.cor.fa 1> super1.err 2>&1\n";
  print FILE "if [[ ! -e work1/superReads.success ]];then\n";
  print FILE "fail Super reads failed, check super1.err and files in ./work1/\n";
  print FILE "fi\n";
  print FILE "log 'All done, the super reads are in work1/superReadSequences.fasta'\n";

script_done


