#!/usr/bin/env perl
# PROJECT: CASAVA
# MODULE:  $RCSfile: filterSmallVariants.pl,v $
# AUTHER: Chris Saunders
#
# Copyright (c) 2006 Solexa, 2007,2008 Illumina
# This software is covered by the "Illumina Genome Analyzer Software
# License Agreement" and the "Illumina Source Code License Agreement",
# and certain third party copyright/licenses, and any user of this
# source file is bound by the terms therein (see accompanying files
# Illumina_Genome_Analyzer_Software_License_Agreement.pdf and
# Illumina_Source_Code_License_Agreement.pdf and third party
# copyright/license notices).
#

use warnings FATAL => 'all';
use strict;

use File::Spec;
use IO::File;
use Sys::Hostname;
use Getopt::Long;

use lib '@CASAVA_FULL_PERL_LIBDIR@';
use Casava::Common::Log;
use Casava::PostAlignment::Sequencing::Config qw(loadConfiguration %CONF_APP %CONF_PROJ %chrEnds readProjectParameters);


my $argvStr     = join ' ', @ARGV;
my $cmdline     = "$0 $argvStr";
my $scriptName = (File::Spec->splitpath($0))[2];

my $projectDir        = "";
my $chrom             = "";
my $sampleRateSetSuffix = '';
my $help              = 0;

my $usage       =<<END;
$scriptName [options]
\t--projectDir   - project directory
\t--chrom        - chromosome
\t-h|--help      - print this message
END


my $result            = GetOptions(
    "projectDir=s" => \$projectDir,
    "chrom=s"      => \$chrom,
    "help|h"       => \$help
);

if ((not $result)  or $help) {
    errorExit "\n\n$usage";
}

errorExit "ERROR: Incorrect number of arguments\n$usage" unless ( $projectDir ne "" );

errorExit "ERROR: Unrecognized arguments: " . join(" ",@ARGV) . "\n" if(@ARGV);

{
    my $hostname = hostname();
    printLog( "Running $hostname:[$cmdline]\n", 0);
}

loadConfiguration($projectDir);

if(not defined $chrEnds{$chrom}){
    errorExit "ERROR: Invalid chromosome label: $chrom\n";
}

my $dirCurrentBuild = $CONF_PROJ{dirBuildParsed};

my %buildChromsBinSizes = ();
readProjectParameters( %buildChromsBinSizes, "BUILD_BIN_SIZES",$projectDir);
my $binCount=$buildChromsBinSizes{$chrom};

my $refSeqLen = $chrEnds{$chrom}->{length};
my $refSeqKnownBases= $chrEnds{$chrom}->{knownBases};

my $chromDir = File::Spec->catdir($dirCurrentBuild, $chrom);
errorExit "ERROR: Can't find chromosome directory: $chromDir\n" unless (-d $chromDir);

my $isCovCutoff = (not ((defined $CONF_PROJ{variantsNoCovCutoff}) and $CONF_PROJ{variantsNoCovCutoff}));

my $isSnpCovCutoff = ((defined $CONF_PROJ{variantsSnpCovCutoff}) and ($CONF_PROJ{variantsSnpCovCutoff} > 0) and $isCovCutoff);
my $isSnpCovCutoffAll = ((defined $CONF_PROJ{variantsSnpsCovCutoffAll}) and $CONF_PROJ{variantsSnpsCovCutoffAll});
my $chromSnpCovCutoff;

my $isNoIndels      = ((defined $CONF_PROJ{variantsNoIndels}) and $CONF_PROJ{variantsNoIndels});

my $isIndelsCovCutoff = ((defined $CONF_PROJ{variantsIndelCovCutoff}) and ($CONF_PROJ{variantsIndelCovCutoff} > 0) and $isCovCutoff);
my $chromIndelsCovCutoff;

my $isKeepUnfiltered = ((defined $CONF_PROJ{variantsKeepUnfiltered}) and $CONF_PROJ{variantsKeepUnfiltered});


#
# get coverage statitistics for coverage summary file and for snp/indel coverage cutoff values:
#

my $totalCoverage = 0;
my $totalUsedCoverage = 0;
my $totalSiteCoverage = 0;
my $totalSiteUsedCoverage = 0;

#if($isSnpCovCutoff or $isIndelsCovCutoff) {
if(1) {

    for ( my $i = 0 ; $i < $binCount ; $i++ ) {
        my $binId = sprintf "%04d", $i;
        my $inputDir = File::Spec->catdir( $chromDir, $binId );
        next if(not -e $inputDir);
        my $varPath = File::Spec->catfile( $inputDir, "variants.info$sampleRateSetSuffix.txt" );

        if(not -e $varPath) {
            # expect no variants info file if sorted file is missing:
            # my $sortedFile = File::Spec->catfile( $inputDir, "sorted.txt");
            # next if(not -e $sortedFile);
            errorExit "ERROR: Can't find variants info file: $varPath\n";
        }

        open(my $varFH, '<', $varPath) or die $!;

        my $is_t_update=0;
        my $is_tu_update=0;

        while(<$varFH>) {
            if (/^NO_REF_N_COVERAGE(|_USED)\s/) {
                my $type = $1;
                my $is_bad = 0;
                if(not /sample_size:\s*(\d+)\s+/) { $is_bad=1; }
                my $ss=$1;
                if(not /nonzero:\s*(\d+)\s+/) { $is_bad=1; }
                my $nz=$1;

                # leave the regex for mean fairly loose to pick up scientific notation, etc..
                if(not /mean:\s*(\d[^\s]*|nan)\s+/) { $is_bad=1; } 
                my $mean=$1;
                errorExit "ERROR: Unexpected format in file: $varPath.\n" if($is_bad);
                my $coverage = ( $ss==0 ? 0 : int($ss*$mean) );
                if($type eq "") {
                    $totalCoverage += $coverage;
                    $totalSiteCoverage += $nz;
                    $is_t_update += 1;
                } else {
                    $totalUsedCoverage += $coverage;
                    $totalSiteUsedCoverage += $nz;
                    $is_tu_update += 1;
                }
            }
        }

        if(not (($is_t_update == 1) and ($is_tu_update == 1))) {
            errorExit "ERROR: unexpected format in variant info file: $varPath\n";
        }

        close($varFH);
    }

    my $meanCoverage = $totalCoverage/$refSeqKnownBases;
    my $meanUsedCoverage = $totalUsedCoverage/$refSeqKnownBases;

    if($isSnpCovCutoff) {
        my $snpCovCutoff = $CONF_PROJ{variantsSnpCovCutoff};
        if($isSnpCovCutoffAll) {
            $chromSnpCovCutoff = $snpCovCutoff*$meanCoverage;
        } else {
            $chromSnpCovCutoff = $snpCovCutoff*$meanUsedCoverage;
        }
    }

    if($isIndelsCovCutoff) {
        $chromIndelsCovCutoff = $CONF_PROJ{variantsIndelCovCutoff}*$meanCoverage;
    }
}


#
# produce chromosome coverage summary file:
#
{
    my $outputCovPath = File::Spec->catfile( $chromDir, "coverage.summary$sampleRateSetSuffix.txt" );
    open(my $outCovFH,'>',$outputCovPath) or die $!;
    print $outCovFH "# ** CASAVA reference sequence coverage summary **\n";
    print $outCovFH "#\$ COLUMNS seq_name sites known_sites known_sites_mapped known_sites_coverage\n";
    print $outCovFH "$chrom\t$refSeqLen\t$refSeqKnownBases\t$totalSiteCoverage\t$totalCoverage\n";
    close($outCovFH);
}



my @unfilteredFiles;


#
# combine snps from all bins and optionally depth-filter:
#
my $outputSnpPath = File::Spec->catfile( $chromDir, "snps$sampleRateSetSuffix.txt" );
open(my $outSnpFH,'>',$outputSnpPath) or die $!;
print $outSnpFH "# ** CASAVA depth-filtered snp calls **\n";
print $outSnpFH "#\$ CMDLINE $cmdline\n" unless $CONF_PROJ{skipVariableMetadata};
print $outSnpFH "#\$ SEQ_MAX_DEPTH $chrom " . ($isSnpCovCutoff ? $chromSnpCovCutoff : "undefined") . "\n";
print $outSnpFH "#\n";


my $columnsDefinition;
my $usedPos;
my $filtPos;

# go through bin snp files:
for ( my $i = 0 ; $i < $binCount ; $i++ ) {
    my $binId = sprintf "%04d", $i;
    my $inputDir = File::Spec->catdir( $chromDir, $binId );
    next if(not -e $inputDir);

    my $snpPath = File::Spec->catfile( $inputDir, "snps.unfiltered$sampleRateSetSuffix.txt" );
    if(not -e $snpPath) {
        # expect no snp file if sorted file is missing:
        my $sortedFile = File::Spec->catfile( $chromDir, "sorted.bam");
        next if(not -e $sortedFile);
        errorExit "ERROR: Can't find snp file: $snpPath\n";
    }

    my $rmSnpFH;
    if($isSnpCovCutoff) {
        my $removedSnpPath = File::Spec->catfile( $inputDir, "snps.removed$sampleRateSetSuffix.txt" );
        open($rmSnpFH,'>',$removedSnpPath) or die $!;
        print $rmSnpFH "# ** CASAVA snp calls removed due to high depth **\n";
        print $rmSnpFH "#\$ CMDLINE $cmdline\n" unless $CONF_PROJ{skipVariableMetadata};
        print $rmSnpFH "#\$ SEQ_MAX_DEPTH $chrom " . ($isSnpCovCutoff ? $chromSnpCovCutoff : "undefined") . "\n";
        print $rmSnpFH "#\n";
        print $rmSnpFH "$columnsDefinition\n" if($columnsDefinition);
    }

    push @unfilteredFiles, $snpPath;

    open(my $snpFH, '<', $snpPath) or die $!;

    while(my $line=<$snpFH>){
        chomp $line;

        if ($line =~ /^#/)
        {
            if ($line =~ /^#\$ COLUMNS\s+(.*)/)
            {
                if(not defined $columnsDefinition)
                {
                    $columnsDefinition = $line;
                    print $outSnpFH "$line\n";
                    print $rmSnpFH "$line\n" if($rmSnpFH);
                    my $pos = 0;
                    for (split (/\s+/, $1))
                    {
                        $usedPos = $pos if ('bcalls_used' eq $_);
                        $filtPos = $pos if ('bcalls_filt' eq $_);
                        ++$pos;
                    }
                }
                elsif ($line ne $columnsDefinition)
                {
                    errorExit "ERROR: Cannot aggregate bin snp files with different column sets\n"
                             ."set1: $columnsDefinition\n"
                             ."set2: $line";
                }
            }
        }
        else
        {
            if(not ((defined $usedPos) and (defined $filtPos))) {
                errorExit "ERROR: Got data line without the position of 'bcalls_used' and 'bcalls_filt' columns defined first\n";
            }
            if($isSnpCovCutoff) {
                my @s = split("\t", $line);
                my $bcalls = $s[$usedPos];
                $bcalls += $s[$filtPos] if($isSnpCovCutoffAll);
                unless ($bcalls < $chromSnpCovCutoff) {
                    print $rmSnpFH "$line\n";
                    next;
                }
            }
            print $outSnpFH "$line\n";
        }
    }
    close($snpFH);
    close($rmSnpFH) if($rmSnpFH);
}
close($outSnpFH);



if(not $isNoIndels) {
    #
    # combine indels from all bins and optionally depth-filter:
    #
    my $outputIndelsPath = File::Spec->catfile( $chromDir, "indels$sampleRateSetSuffix.txt" );
    open(my $outIndelsFH,'>',$outputIndelsPath) or die $!;
    print $outIndelsFH "# ** CASAVA depth-filtered indel calls **\n";
    print $outIndelsFH "#\$ CMDLINE $cmdline\n" unless $CONF_PROJ{skipVariableMetadata};
    print $outIndelsFH "#\$ SEQ_MAX_DEPTH $chrom " . ($isIndelsCovCutoff ? $chromIndelsCovCutoff : "undefined") . "\n";
    print $outIndelsFH "#\n";

    $columnsDefinition = undef;
    my $depthPos;


    # go through bin indel files:
    for ( my $i = 0 ; $i < $binCount ; $i++ ) {
        my $binId = sprintf "%04d", $i;
        my $inputDir = File::Spec->catdir( $chromDir, $binId );
        next if(not -e $inputDir);

        my $indelsPath = File::Spec->catfile( $inputDir, "indels.unfiltered$sampleRateSetSuffix.txt" );

        if(not -e $indelsPath) {
            # expect indel snp file if sorted file is missing or empty:
            my $sortedFile = File::Spec->catfile( $inputDir, "sorted.txt");
            next if((not -e $sortedFile) or (-z $sortedFile));
            errorExit "ERROR: Can't find indel file: $indelsPath\n";
        }

        my $rmIndelsFH;
        if($isIndelsCovCutoff) {
            my $removedIndelsPath = File::Spec->catfile( $inputDir, "indels.removed$sampleRateSetSuffix.txt" );
            open($rmIndelsFH,'>',$removedIndelsPath) or die $!;
            print $rmIndelsFH "# ** CASAVA indel calls removed due to high depth **\n";
            print $rmIndelsFH "#\$ CMDLINE $cmdline\n" unless $CONF_PROJ{skipVariableMetadata};
            print $rmIndelsFH "#\$ SEQ_MAX_DEPTH $chrom " . ($isIndelsCovCutoff ? $chromIndelsCovCutoff : "undefined") . "\n";
            print $rmIndelsFH "#\n";
            print $rmIndelsFH "$columnsDefinition\n" if($columnsDefinition);
        }

        push @unfilteredFiles, $indelsPath;

        open(my $indelsFH, '<', $indelsPath) or die $!;

        while(my $line=<$indelsFH>){
            chomp $line;

            if ($line =~ /^#/)
            {
                if ($line =~ /^#\$ COLUMNS\s+(.*)/)
                {
                    if(not defined $columnsDefinition)
                    {
                        $columnsDefinition = $line;
                        print $outIndelsFH "$line\n";#print the column definition line
                        print $rmIndelsFH "$line\n" if($rmIndelsFH);
                        my $pos = 0;
                        for (split (/\s+/, $1))
                        {
                            $depthPos = $pos if ('depth' eq $_);
                            ++$pos;
                        }
                    }
                    elsif ($line ne $columnsDefinition)
                    {
                        errorExit "ERROR: Cannot aggregate bin indel files with different column sets\n"
                          ."set1: $columnsDefinition\n"
                            ."set2: $line";
                    }
                }
            }
            else
            {
                if(not (defined $depthPos)) {
                    errorExit "ERROR: Got data line without the position of 'depth' column defined first\n";
                }
                if($isIndelsCovCutoff) {
                    my @s = split("\t", $line);
                    my $bcalls = $s[$depthPos];
                    unless ($bcalls < $chromIndelsCovCutoff) {
                        print $rmIndelsFH "$line\n";
                        next;
                    }
                }
                print $outIndelsFH "$line\n";
            }
        }
        close($indelsFH);
        close($rmIndelsFH) if($rmIndelsFH);
    }
    close($outIndelsFH);
}

unlink(@unfilteredFiles);

printLog( "$0 Finished\n", 0);
