#!/usr/local/bin/perl -w

# Author: lh3

use strict;
use warnings;
use Getopt::Std;
use Cwd qw/getcwd/;

&usage if (@ARGV < 1);
# global variables
my @nucl_type = (0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4);
my @conv = split('', 'XACMGRSVTWYHKDBN');

my $command = shift(@ARGV);
my %func = (easyrun=>\&easyrun, chrpt2snp=>\&chrpt2snp, fastq2bfq=>\&fastq2bfq, cat2pair=>\&cat2pair);
die("Unknown command \"$command\".\n") if (!defined($func{$command}));
&{$func{$command}}();
exit(0);

sub cat2pair
{
	die qq(Usage: mapass2.pl cat2pair <read1_len> <input.fastq> [0]\n) if (@ARGV < 2);
	my $size1 = $ARGV[0];
	my $trim_len = $ARGV[2]? $ARGV[2] : 0;
	my $fn = $ARGV[1];
	my $fn_out = $ARGV[1];
	$fn_out =~ s/^.*\/([^\/\s]+)$/$1/ if ($fn_out =~ /\//);
	mkdir("read1"); mkdir("read2");
	my ($fh1, $fh2, $fhin);
	open($fhin, $fn) || die;
	open($fh1, ">read1/$fn_out") || die;
	open($fh2, ">read2/$fn_out") || die;
	while (<$fhin>) {
		if (/^@/) {
			print $fh1 $_; print $fh2 $_;
			$_ = <$fhin>; chomp; # sequence
			print $fh1 substr($_, $trim_len, $size1-$trim_len), "\n+\n";
			print $fh2 substr($_, $size1+$trim_len), "\n+\n";
			<$fhin>; $_ = <$fhin>; chomp; # qualities
			print $fh1 substr($_, $trim_len, $size1-$trim_len), "\n";
			print $fh2 substr($_, $size1+$trim_len), "\n";
		}
	}
	close($fhin); close($fh1); close($fh2);
}

sub fastq2bfq
{
	my %opts;
	getopts('s:e:', \%opts);
	die qq(Usage: mapass2.pl fastq2bfq [-s <size>=500000] [-e <ext>=fastq] <dst_dir> <src_dir>\n) if (@ARGV < 2);
	my $size = ($opts{s})? $opts{s} : 500000; # 500k by default
	my $ext = ($opts{e})? $opts{e} : 'fastq';
	my ($dst, $src) = ($ARGV[0], $ARGV[1]);
	my $prefix = $ARGV[0];
	$size *= 4;
	die("FATAL ERROR: source directory $src does not exist!\n") unless (-d $src);
	print qq(mkdir $dst;\n) unless (-d $dst);
	my $cwd = getcwd();
	chdir($src);
	my @list = `(find . -name "*.$ext"; find . -name "*.$ext.gz")`;
	chdir($cwd);
	foreach (@list) {
		chomp;
		s/^\.\///;
		my $ori = $_;
		s/\//-/g; s/\.gz$//; s/\.$ext$//;
		my ($cur_size, $name);
		my $is_paired = (/read(s?)[12]/)? 1 : 0;
		if ($is_paired) {
			$cur_size = int($size/8)*4;
			$name = (/read(s?)1/)? "reads1.bfq" : "reads2.bfq";
			s/read(s?)[12]//;
			s/--/-/g;
		} else {
			$cur_size = $size;
			$name = "reads.bfq";
		}
		my $prog = ($ori =~ /\.gz$/)? 'zcat' : 'cat';
		my $lineno = `$prog $src/$ori | wc -l`;
		$lineno =~ s/^\s*(\d+).*$/$1/;
		for (my $i = 0; $i < $lineno; $i += $cur_size) {
			my $dst_dir = "$dst/$_-$i";
			my ($begin, $end) = ($i+1, $i+$cur_size);
			print qq(mkdir -p $dst_dir; $prog $src/$ori | awk 'NR>=$begin&&NR<=$end' | mapass2 fastq2bfq - $dst_dir/$name;\n);
		}
	}
	print qq(find $dst -type d | tail +2 > $prefix.list\n);
	my $n_jobs = `wc -l $prefix.list`;
	$n_jobs =~ /^\s+(\d+)/;
	$n_jobs = $1;
	print qq(cat >$prefix.pl <<\@\@CONTENT\@\@
%fr_config =
(
 run_list=>'$prefix.list',
 n_jobs=>$n_jobs,
 LSF_queue=>'normal',
 LSF_resource=>'select[mem>1000 && type==LINUX64] rusage[mem=1000]',
 binary_path=>'/lustre/work1/sanger/lh3/lh3/lsf-prog',
 script_path=>'/lustre/work1/sanger/lh3/lh3/lsf-prog/scripts',
 action=>\\&func
);

sub func
{
	if (-f "reads.bfq") {
		system("mapass2 match -n 3 $prefix.map -d ../../adaptor ../../b36.bfa reads.bfq");
	} elsif (-f "reads1.bfq" && -f "reads2.bfq") {
		system("mapass2 match -i 60 -a 250 -n 3 $prefix.map -d ../../adaptor ../../b36.bfa reads1.bfq reads2.bfq");
	}
}
\@\@CONTENT\@\@\n);
}

sub easyrun
{
	my %opts;
	getopts('d:', \%opts);
	die qq(Usage: mapass2.pl easyrun [-d dstdir] ref.fa reads1.fastq [reads2.fastq [...]]\n) if (@ARGV < 2);
	my $d = ($opts{d})? $opts{d} : "easyrun-$$";
	mkdir($d) unless (-d $d);
	my $exe = gwhich('mapass2') || die;
	# run fasta2bfa
	warn("[easyrun] Converting the reference sequence '$ARGV[0]' to binary fasta format (.bfa)...\n");
	system("$exe fasta2bfa $ARGV[0] $d/ref.bfa") && die;
	# run fastq2bfq && match
	shift(@ARGV);
	my $i = 0;
	my (@map_files, $map_file);
	foreach my $file (@ARGV) {
		++$i;
		warn("[easyrun] Converting fastq file '$file' to binary fastq format (.bfq)...\n");
		system("$exe fastq2bfq $file $d/reads-$i.bfq") && die;
		warn("[easyrun] Aligning reads-$i.bfq to the reference...\n");
		system("(cd $d; $exe match reads-$i.map ref.bfa reads-$i.bfq)") && die;
		push(@map_files, "reads-$i.map");
	}
	if ($i > 1) { # we should merge them
		system("cd $d; $exe mapmerge merged.map @map_files") && die;
		$map_file = "merged.map";
	} else { $map_file = "reads-1.map"; }
	warn("[easyrun] Do mapcheck...\n");
	system("(cd $d; $exe mapcheck ref.bfa $map_file > mapcheck.txt)") && die;
	warn("[easyrun] Consensus base-calling...\n");
	system("(cd $d; $exe assemble consensus.cns ref.bfa $map_file 2> assemble.log)") && die;
	# do some statistics
	warn("[easyrun] Extracting consensus in the fastq format...\n");
	system("$exe cns2fq $d/consensus.cns > $d/cns.fq") && die;
	warn("[easyrun] Extracting SNPs...\n");
	system("$exe cns2snp $d/consensus.cns > $d/cns.snp") && die;
	exit(0);
}

sub chrpt2snp
{
	# calculate @mm
	my @mm;
	for (my $i = 0; $i != 0x10000; ++$i) {
		my ($x, $k) = (1, 0);
		for (my $l = 0; $l != 16; ++$l, $x <<= 1) {
			++$k if ($i & $x);
		}
		$mm[$i] = $k;
	}
	my $fh;
	open($fh, "| sort -k1,1 -k2,2n");
	while (<>) {
		next unless (/^\d+/ && /reference$/);
		my @t = split("\t", $_);
		next unless ($t[11] =~ /^\d+/); # no position
		my $vs = $mm[$t[16]&0xffff] + $mm[$t[16]>>16&0xffff];
		printf $fh ("chr$t[6]\t$t[11]\tN\tN\t%d\t%.3f\n", $vs*10, ($t[13] ne ' ')? $t[13] : 0.0);
	}
	close($fh);
}

sub usage
{
	die qq(
Usage:   mapass2.pl <command> [options] <input> [...]

Command: easyrun      simple pipeline for small dataset
         chrpt2snp    convert dbSNP's chr_rpt file to .snp file
         fastq2bfq    convert fastq in batch
         cat2pair     convert paired end reads format
\n);
}

# the following codes are copied from treefam::generic

sub dirname
{
	my $prog = shift;
	my $cwd = getcwd;
	return $cwd if (!($prog =~ /\//));
	$prog =~ s/\/[^\s\/]+$//g;
	return $prog;
}
sub which
{
	my $file = shift;
	my $path = (@_)? shift : $ENV{PATH};
	return if (!defined($path));
	foreach my $x (split(":", $path)) {
		$x =~ s/\/$//;
		return "$x/$file" if (-x "$x/$file");
	}
	return;
}
sub gwhich
{
	my $progname = shift;
	my $addtional_path = shift if (@_);
	my $dirname = &dirname($0);
	my $tmp;
	my $cwd = getcwd;

	chomp($dirname);
	if ($progname =~ /^\//) {
		return $progname;
	} elsif (-x $progname) {
		return "$cwd/$progname";
	} elsif (defined($addtional_path) && ($tmp = &which($progname, $addtional_path))) {
		return $tmp;
	} elsif (-x "./$progname") {
		return "$cwd/$progname";
	} elsif (defined($dirname) && (-x "$dirname/$progname")) {
		return "$dirname/$progname";
	} elsif (($tmp = &which($progname))) {
		return $tmp;
	} else {
		warn("[generic::gwhich] fail to find executable $progname anywhere.");
		return;
	}
}
