#!/usr/bin/env perl
use strict;
use warnings;
use Getopt::Long;

# This code is part of the MOCAT analysis pipeline
# Code is (c) Copyright EMBL and BGI, 2012-2013
# This code is released under GNU GPL v3.
my ( %saved_qual_cutoff, %saved_first_cycle, $length_cutoff, $qual_cutoff, $paired_end_data, $use_5prime_file, $trim_5prime_end, $src_dir, $bin_dir, $solexaqa_or_fastx, $cwd, $temp_dir, $file_formats_array, $ZCAT, $sample, $trim_5_prime_end );
my $TOTmax            = 0;
my $TOTlength         = 0;
my $TOTsum            = 0;
my $TOTinserts        = 0;
my $fastq_trim_filter = "fastq_trim_filter_v5_EMBL"; # v5 is the latest BGI version
my $format_counter    = 0;
GetOptions(
	'length_cutoff=s'      => \$length_cutoff,
	'qual_cutoff=s'        => \$qual_cutoff,
	'paired_end_data=s'    => \$paired_end_data,
	'use_5prime_file=s'    => \$use_5prime_file,
	'trim_5prime_end=s'    => \$trim_5_prime_end,
	'src_dir=s'            => \$src_dir,
	'bin_dir=s'            => \$bin_dir,
	'solexaqa_or_fastx=s'  => \$solexaqa_or_fastx,
	'cwd=s'                => \$cwd,
	'temp_dir=s'           => \$temp_dir,
	'file_formats_array=s' => \$file_formats_array,
	'zcat=s'               => \$ZCAT,
	'sample=s'             => \$sample
);

# If running OSX, change binary file
chomp( my $systemType = `uname -s` );
if ( $systemType =~ m/Darwin/ ) {
	$fastq_trim_filter = $fastq_trim_filter . "_OSX";
}

# Get fq or fq.gz files
#my @allFiles = `ls -1 $cwd/$sample/*fq $cwd/$sample/*fq.gz 2>/dev/null | grep -v 'trimmed.filtered' | grep -v '.single.fq' | grep -v '.pair.'`;
my @allFiles = `ls -1 $cwd/$sample/*fq $cwd/$sample/*fq.gz $cwd/$sample/*bz2 2>/dev/null`;

for my $i ( 0 .. scalar @allFiles - 1 ) {
	chomp( $allFiles[$i] );
	my @J = split /\//, $allFiles[$i];
	$allFiles[$i] = $J[-1];
}

# Delete old files before starting
print "ReadTrimFilter_aux : SAMPLE $sample\n";
print "ReadTrimFilter_aux : SCRIPT VERSION 3\n";
print "ReadTrimFilter_aux : CLEANUP : rm -f $temp_dir/$sample/temp/*.trimmed.filtered.fq*; mkdir -p $cwd/$sample/stats; rm -fr $cwd/$sample/reads.processed.$solexaqa_or_fastx; mkdir -p $cwd/$sample/reads.processed.$solexaqa_or_fastx\n";

system "rm -f $temp_dir/$sample/temp/*.trimmed.filtered.fq*";
system "mkdir -p $cwd/$sample/stats";
system "rm -fr $cwd/$sample/reads.processed.$solexaqa_or_fastx";
system "mkdir -p $cwd/$sample/reads.processed.$solexaqa_or_fastx";

# get quality statistics
foreach my $j (@allFiles) {
	( my $format, my $sanger ) = get_format();
	my $zip = "cat";
	if ( $j =~ /\.gz$/ ) { $zip = "gunzip -c" }
	if ( $j =~ /\.bz2$/ ) { $zip = "bunzip2 -c" }
	print "ReadTrimFilter_aux : GET QUALITY STATS : mode=$solexaqa_or_fastx sanger=$sanger format=$format sample=$j : EXECUTE $zip $cwd/$sample/$j | $bin_dir/fastx_quality_stats -o $temp_dir/$sample/temp/$j.qual_stats.temp $sanger\n";
	system "$zip $cwd/$sample/$j | $bin_dir/fastx_quality_stats -o $temp_dir/$sample/temp/$j.qual_stats.temp $sanger";
	print "ReadTrimFilter_aux : Get 3 prime end\n";
	$saved_qual_cutoff{"$sample$j"} = get_3prime_end( $sample, $j );
	#$saved_qual_cutoff{"$sample$j"} = get_stats( $sample, $j );
}

# Process all files
$format_counter = -1;
foreach my $j (@allFiles) {
	( my $format, my $sanger ) = get_format();
	my $first_cycle;
	my $qual_cutoff = $saved_qual_cutoff{"$sample$j"};
	if ( $use_5prime_file eq 'yes' ) {
		$first_cycle = get_5prime_from_file( $sample, $j );
		$saved_first_cycle{"$sample$j"} = $first_cycle;
	}
	if ( $use_5prime_file eq 'no' ) {
		$first_cycle = get_quality_statistics( $sample, $j );
		$saved_first_cycle{"$sample$j"} = $first_cycle;
		open X, '>>', "$cwd/MOCAT.cutoff5prime_calculated";
		print X "$sample\t$j\t$first_cycle\n";
		close X;
	}
	if ( $trim_5_prime_end eq "no" ) {
		$first_cycle = 1;
		$saved_first_cycle{"$sample$j"} = $first_cycle;
	}

	# Process non pair-end files here
	if ( $paired_end_data eq "no" ) {
		my $basename = $j;
		$basename =~ s/.fq.bz2$//;
		$basename =~ s/.fq.gz$//;
		$basename =~ s/.fq$//;
		
		my $cmd = "$bin_dir/$fastq_trim_filter -m $solexaqa_or_fastx -a $cwd/$sample/$j -f $first_cycle -q $qual_cutoff -l $length_cutoff -p 50 $sanger -o $cwd/$sample/reads.processed.$solexaqa_or_fastx/$basename";
		
		print "ReadTrimFilter_aux : PROCESS SINGLE-END $j : EXECUTE $cmd\n";
		my @out = `$cmd`;
		unless ( scalar @out == 5 ) {
			die "SAMPLE ERROR & EXIT: Sample seem to not have any reads passing read trim filter. Something wrong? Perhaps the cut off value is too high?";
		}
		print "ReadTrimFilter_aux : $fastq_trim_filter OUTPUT: ";
		print $out[0];
		shift @out;
		chomp( my $sum     = $out[0] );
		chomp( my $len     = $out[1] );
		chomp( my $max     = $out[2] );
		chomp( my $inserts = $out[3] );
		$TOTinserts = $TOTinserts + $inserts;
		$TOTsum     = $TOTsum + $sum;
		$TOTlength  = $TOTlength + $len;

		if ( $TOTmax < $max ) {
			$TOTmax = $max;
		}
	}
}

# Post process pair-end files
if ( $paired_end_data eq "yes" ) {
	for ( my $count = 0 ; $count <= scalar @allFiles - 1 ; $count = $count + 2 ) {
		$allFiles[$count] =~ m/^(.+)\.1.fq/;
		my $base_name = $1;
		$format_counter = $count;
		( my $format, my $sanger ) = get_format();
		
		my $cmd = "$bin_dir/$fastq_trim_filter -m $solexaqa_or_fastx -a $cwd/$sample/" . $allFiles[$count] . " -b $cwd/$sample/" . $allFiles[ $count + 1 ] . " -f " . $saved_first_cycle{ $sample . $allFiles[$count] } . " -2 " . $saved_first_cycle{ $sample . $allFiles[ $count + 1 ] } . " -q $qual_cutoff -l $length_cutoff -p 50 $sanger -o $cwd/$sample/reads.processed.$solexaqa_or_fastx/$base_name";
				
		print "ReadTrimFilter_aux : PROCESS PAIR-END $base_name : EXECUTE $cmd\n";
		my @out = `$cmd`;
		unless ( scalar @out == 5 ) {
			die "SAMPLE ERROR & EXIT: Sample seem to not have any reads passing read trim filter. Something wrong? Perhaps the cut off value is too high?";
		}
		shift @out;
		chomp( my $sum     = $out[0] );
		chomp( my $len     = $out[1] );
		chomp( my $max     = $out[2] );
		chomp( my $inserts = $out[3] );
		$TOTinserts = $TOTinserts + $inserts;
		$TOTsum     = $TOTsum + $sum;
		$TOTlength  = $TOTlength + $len;

		if ( $TOTmax < $max ) {
			$TOTmax = $max;
		}
	}
}

# Nice cleanup of files
print "ReadTrimFilter_aux : SUMMARIZE STATS and write $cwd/$sample/stats/$sample.readtrimfilter.$solexaqa_or_fastx.stats\n";
system "rm -f $temp_dir/$sample/temp/*qual_stats.temp; mv $temp_dir/$sample/temp/*qual_stats $cwd/$sample/reads.processed.$solexaqa_or_fastx;";

# Stats section
my $max     = $TOTmax;
my $length  = $TOTlength;
my $inserts = $TOTinserts;
open STAT, '>', "$cwd/$sample/stats/$sample.readtrimfilter.$solexaqa_or_fastx.stats";
if ( $TOTsum == 0 ) {
	die "SAMPLE ERROR & EXIT: Sample $sample seem to not have any reads passing read trim filter. Something wrong?";
}
my $avg = int( ( $TOTlength / $TOTsum ) + 0.5 );
my $kmer;
if ( $avg % 2 == 0 ) {

	#avg_readlen is even
	$kmer = $avg / 2;
	if ( $kmer % 2 == 0 ) {
		$kmer = $kmer + 1;
	}
	else {
		$kmer = $kmer + 2;
	}
}
else {

	#avg_readlen is odd
	$kmer = ( $avg + 1 ) / 2;
	if ( $kmer % 2 == 1 ) {
		$kmer = $kmer + 0;
	}
	else {
		$kmer = $kmer + 1;
	}
}
print STAT "Reads\tBases\tMax\tAvg\tKmer\tInserts\n$TOTsum\t$TOTlength\t$TOTmax\t$avg\t$kmer\t$TOTinserts\n";
close STAT;
print "ReadTrimFilter_aux : COMPLETED READ TRIM FILTER\n";
exit 0;

#-----------
#SUBROUTINES
#-----------
sub get_5prime_from_file {
	my $sample      = $_[0];
	my $j           = $_[1];
	my $first_cycle = "";
	open FPI, '<', "$cwd/MOCAT.cutoff5prime";
	while (<FPI>) {
		chomp;
		my @line = split;
		if (
			$line[0] eq $sample
			&& (         $line[1] eq $j
				|| "$line[1].fq" eq $j
				|| "$line[1].fq.gz" eq $j
				|| "$line[1].fq.bz2" eq $j )
		  )
		{
			$first_cycle = $line[2];
		}
	}
	close FPI;
	if ( $first_cycle eq "" ) {
		die "\nERROR and EXIT: MOCAT.cutoff5prime file incorrect. Last checked sample: $sample. Last checked lane: $j. File should be in format <sample>\\t<lane as complete file name>\\t<base>";
	}
	return $first_cycle;
}

sub get_3prime_end {
	my $sample = $_[0];
	my $j      = $_[1];
	if ( -s "$temp_dir/$sample/temp/$j.qual_stats.temp" ) {
		open IN_TEMP, '<', "$temp_dir/$sample/temp/$j.qual_stats.temp";
	}
	else {
		die "ERROR and EXIT: Cannot open $temp_dir/$sample/temp/$j.qual_stats.temp\n";
	}
	open OUT_TEMP, '>', "$temp_dir/$sample/temp/$j.qual_stats"
	  || die "ERROR and EXIT: Cannot open $temp_dir/$sample/temp/$j.qual_stats\n";
	my ( $raw_read_num, $raw_read_base ) = ( 0, 0 );
	<IN_TEMP>;
	while (<IN_TEMP>) {

		#        unless ($. == 1){
		my @line   = split "\t", $_;
		my $cycle  = $line[0];
		my $mean   = $line[5];
		my $median = $line[7];
		if ( $mean >= $qual_cutoff && $median >= $qual_cutoff ) {
			print OUT_TEMP $_;
		}
		$raw_read_num = $line[1] if ( $raw_read_num < $line[1] );
		$raw_read_base += $line[1];

		#        }
	}
	close IN_TEMP;
	close OUT_TEMP;
	open RAW, "> $cwd/$sample/stats/$j.raw.reads.stats" || die "ERROR and EXIT: Cannot open $cwd/$sample/stats/$j.raw.reads.stats\n";
	print RAW "Reads\tBases\n", $raw_read_num, "\t", $raw_read_base, "\n";
	close RAW;
	return $qual_cutoff;
}

sub get_stats {
	my $sample = $_[0];
	my $j      = $_[1];
	if ( -s "$cwd/$sample/reads.processed.fastx/$j.qual_stats.temp" ) {
		open IN_TEMP, '<', "$cwd/$sample/reads.processed.fastx/$j.qual_stats.temp";
	}
	else {
		die "ERROR & EXIT: Missing $cwd/$sample/reads.processed.fastx/$j.qual_stats.temp";
	}
	my ( $raw_read_num, $raw_read_base ) = ( 0, 0 );
	<IN_TEMP>;
	while (<IN_TEMP>) {
		my @line   = split "\t", $_;
		my $cycle  = $line[0];
		my $mean   = $line[5];
		my $median = $line[7];
		$raw_read_num = $line[1] if ( $raw_read_num < $line[1] );
		$raw_read_base += $line[1];
	}
	close IN_TEMP;
	open RAW, "> $cwd/$sample/stats/$j.raw.reads.stats" || die "ERROR and EXIT: Cannot open $cwd/$sample/stats/$j.raw.reads.stats\n";
	print RAW "Reads\tBases\n", $raw_read_num, "\t", $raw_read_base, "\n";
	close RAW;
	return $qual_cutoff;
}


sub get_quality_statistics {
	my $sample = $_[0];
	my $j      = $_[1];
	my @line   = my @cycle_array = my @cycle_number = ();
	my ( %A, %C, %G, %T );
	open IN, '<', "$temp_dir/$sample/temp/$j.qual_stats"
	  || die "ERROR and EXIT: Cannot open $temp_dir/$sample/temp/$j.qual_stats\n";
	while (<IN>) {
		@line = split "\t", $_;
		$A{ $line[0] } = $line[12];
		$C{ $line[0] } = $line[13];
		$G{ $line[0] } = $line[14];
		$T{ $line[0] } = $line[15];
		push( @cycle_number, $line[0] );
	}
	close IN;
	my $read_length = scalar @cycle_number;
	my ( $max_A, $min_A ) = &get_base_range( \%A, $read_length );
	my ( $max_C, $min_C ) = &get_base_range( \%C, $read_length );
	my ( $max_G, $min_G ) = &get_base_range( \%G, $read_length );
	my ( $max_T, $min_T ) = &get_base_range( \%T, $read_length );
	push( @cycle_array, '0trim' );

	foreach my $m (@cycle_number) {
		if (         $A{$m} > $max_A
			|| $A{$m} < $min_A
			|| $C{$m} > $max_C
			|| $C{$m} < $min_C
			|| $G{$m} > $max_G
			|| $G{$m} < $min_G
			|| $T{$m} > $max_T
			|| $T{$m} < $min_T )
		{
			my $var = $m . "trim";
			push @cycle_array, $var;
		}
		else {
			my $var = $m . "keep";
			push @cycle_array, $var;
		}
	}
	my $first = 1;
	my $tempstring = join( " ", @cycle_array );
	if ( $tempstring =~ m/^(\d+trim )+(\d+)keep/ ) {
		$first = $2;
	}
	return $first;
}

sub get_base_range {
	my ( $base, $read_length ) = @_;
	my ( $sum_base, $avg_base, $sum_of_squares_base, $var_base, $stdev_base );
	foreach my $i ( keys %{$base} ) {
		$sum_base += $$base{$i};
	}
	$avg_base = $sum_base / $read_length;
	foreach my $i ( keys %{$base} ) {
		$sum_of_squares_base += ( $$base{$i} - $avg_base )**2;
	}
	$var_base = $sum_of_squares_base / ( $read_length - 1 );
	$stdev_base = sqrt($var_base);
	return ( $avg_base + ( $stdev_base * 2 ), $avg_base - ( $stdev_base * 2 ) );
}

sub get_format {
	my $format = substr $file_formats_array, $format_counter, 1;
	my $sanger;
	if ( $format eq 's' ) { $format = 'sanger'; $sanger = "-Q 33" }
	if ( $format eq 'x' ) { $format = 'solexa'; $sanger = "" }
	if ( $format eq 'i' ) { $format = 'ia';     $sanger = "" }
	$format_counter++;
	return ( $format, $sanger );
}
