#!/usr/bin/perl

use strict;


my @chrs = (1..19, 'X', 'Y');

my ($line, $start, $ch);

my $uname;
my $Nredund;
my $bppos;

my ($oname, $fname, $pname);


foreach $ch ( @chrs ) {
    print("\nchromosome $ch\n");

    my $ifile = join('', 'snps.raw.output.chr', $ch, '.vcf');

    my $cfile = join('', 'snps.raw.SNPclusters.chr', $ch, '.vcf');
    my $ffile = join('', 'snps.raw.failFinal.chr', $ch, '.vcf');
    my $ofile = join('', 'snps.filtered.chr', $ch, '.vcf');


    open(OUT, '>', $ofile) or die "Cannot open file $ofile for writing, reason $!\n";
    open(OUC, '>', $cfile) or die "Cannot open file $cfile for writing, reason $!\n";
    open(OUF, '>', $ffile) or die "Cannot open file $ffile for writing, reason $!\n";

    open(INW, '<', $ifile) or die "Cannot open vcf file $ifile, reason $!\n";

    $line = <INW>;
    $start = substr( $line, 0, 2);

    while ( $start eq '##' ) {
	print(OUT "$line");
	print(OUC "$line");
	print(OUF "$line");

	$line = <INW>;
	$start = substr( $line, 0, 2);
    } # end of headers while loop
 
    my $header_line = $line;
    print(OUT "$header_line");
    print(OUC "$header_line");
    print(OUF "$header_line");


    my ($filt, $info, $i);
    my (@row, @filt_arr, @info_arr, @decom);

    my $Ntot = 0;
    my $Nclus = 0;
    my $Nfinal = 0;
    my $Nlow = 0;
    my $Nmult = 0;
    my $Ninfo_any = 0;
    my $Ninfo_fil = 0;

    while ( defined($line = <INW>) ) {

	chomp($line);

	$Ntot = ( $Ntot + 1 );

	@row = split("\t", $line); 
	$filt = $row[6];
	@filt_arr = split(';', $filt);

	if ( $filt_arr[0] eq 'Final' ) {

	    # write out to fail Final file
	    print(OUF "$line\n");
	    $Nfinal = ( $Nfinal + 1 );

	} else {

	    if ( $filt_arr[0] eq 'SnpCluster' || $filt_arr[1] eq 'SnpCluster' ) {

		# write out to snp cluster file
		print(OUC "$line\n");
		$Nclus = ( $Nclus + 1 );

	    } else {

		# check to make sure doesn't fail filter

		$info = $row[7];
		@info_arr = split(';', $info);

		my %check;

		for ($i=0; $i<=$#info_arr; $i++) {
		    @decom = split('=', $info_arr[$i]);
		    $check{$decom[0]} = $decom[1];
#			print("@decom\n");
		}

		if ( ! exists($check{'DP'}) || ! exists($check{'Dels'}) || ! exists($check{'HRun'}) || ! exists($check{'HaplotypeScore'}) || ! exists($check{'MQ'}) || ! exists($check{'MQ0'}) || ! exists($check{'QD'}) || ! exists($check{'SB'}) ) {

		    $Ninfo_fil = ( $Ninfo_fil + 1 );
		    print(OUF "$line\n");

		} else {

		    if ( $filt_arr[0] eq 'LowQual' ) {
			$Nlow = ( $Nlow + 1 );
		    }

		    # set individual genotypes with low coverage as missing

		    my @new_row = @row[0..8];

		    for ($i=9; $i<29; $i++) {

			my $mouse = $row[$i];
			my @mouse_data = split(':', $mouse);

			my $geno_rd = $mouse_data[2];
			if ( $geno_rd < 5 ) {
			    $mouse_data[0] = './.';
#    print("geno RD $mouse_data[2]\n");
			}

			my $new_mouse = join(':', @mouse_data);
			push(@new_row, $new_mouse);

#$geno_str = $mouse_data[0];
#@geno_arr = split("/", $geno_str);
		    }

		    my $new_line = join("\t", @new_row);
		    print(OUT "$new_line\n");

		} # end of check filter loop
	    } # end of SNP cluster loop
	} # end of Final filter loop
    } # end of while loop
    close(INW);
    close(OUT);

    print("total SNPs $Ntot\nSNPs in clusters $Nclus\nSNPs failing Final filters $Nfinal\nSNPs with missing filter INFO $Ninfo_fil\nSNPs called LowQual $Nlow\n");

} # end of chr loop




sub snphwe {
    my $obs_hets = shift;
    my $obs_hom1 = shift;
    my $obs_hom2 = shift;

    if($obs_hom1 < 0 || $obs_hom2 < 0 || $obs_hets <0) {
	return(-1);
    }

    # rare homozygotes
    my $obs_homr;

    # common homozygotes
    my $obs_homc;
    if($obs_hom1 < $obs_hom2) {
	$obs_homr = $obs_hom1;
	$obs_homc = $obs_hom2;
    } else {
	$obs_homr = $obs_hom2;
	$obs_homc = $obs_hom1;
    }

    # number of rare allele copies
    my $rare_copies = 2 * $obs_homr + $obs_hets;

    # total number of genotypes
    my $genotypes = $obs_homr + $obs_homc + $obs_hets;

    if($genotypes <= 0) {
	return(-1);
    }
    
    # Initialize probability array
    my @het_probs;
    for(my $i=0; $i<=$rare_copies; $i++) {
	$het_probs[$i] = 0.0;
    }

    # start at midpoint
    my $mid = int($rare_copies * (2 * $genotypes - $rare_copies) / (2 * $genotypes));

    # check to ensure that midpoint and rare alleles have same parity
    if(($rare_copies & 1) ^ ($mid & 1)) {
	$mid++;
    }
    
    my $curr_hets = $mid;
    my $curr_homr = ($rare_copies - $mid) / 2;
    my $curr_homc = $genotypes - $curr_hets - $curr_homr;

    $het_probs[$mid] = 1.0;
    my $sum = $het_probs[$mid];
    for($curr_hets = $mid; $curr_hets > 1; $curr_hets -= 2) {
	$het_probs[$curr_hets - 2] = $het_probs[$curr_hets] * $curr_hets * ($curr_hets - 1.0) / (4.0 * ($curr_homr + 1.0) * ($curr_homc + 1.0));
	$sum += $het_probs[$curr_hets - 2];

	# 2 fewer heterozygotes for next iteration -> add one rare, one common homozygote
	$curr_homr++;
	$curr_homc++;
    }

    $curr_hets = $mid;
    $curr_homr = ($rare_copies - $mid) / 2;
    $curr_homc = $genotypes - $curr_hets - $curr_homr;
    for($curr_hets = $mid; $curr_hets <= $rare_copies - 2; $curr_hets += 2) {
	$het_probs[$curr_hets + 2] = $het_probs[$curr_hets] * 4.0 * $curr_homr * $curr_homc / (($curr_hets + 2.0) * ($curr_hets + 1.0));
	$sum += $het_probs[$curr_hets + 2];
	
	# add 2 heterozygotes for next iteration -> subtract one rare, one common homozygote
	$curr_homr--;
	$curr_homc--;
    }

    for(my $i=0; $i<=$rare_copies; $i++) {
	$het_probs[$i] /= $sum;
    }

    # alternate p-value calculation for p_hi/p_lo
#    my $p_hi = $het_probs[$obs_hets];
#    for(my $i=$obs_hets+1; $i<=$rare_copies; $i++) {
#$p_hi += $het_probs[$i];
#    }
#    
#    my $p_lo = $het_probs[$obs_hets];
#    for(my $i=$obs_hets-1; $i>=0; $i--) {
#$p_lo += $het_probs[$i];
#    }
#
#    my $p_hi_lo;
#    if($p_hi < $p_lo) {
#$p_hi_lo = 2 * $p_hi;
#    } else {
#$p_hi_lo = 2 * $p_lo;
#    }

    # Initialise P-value 
    my $p_hwe = 0.0;

    # P-value calculation for p_hwe
    for(my $i = 0; $i <= $rare_copies; $i++) {
	if($het_probs[$i] > $het_probs[$obs_hets]) {
	    next;
	}
	$p_hwe += $het_probs[$i];
    }
    
    if($p_hwe > 1) {
	$p_hwe = 1.0;
    }

    return($p_hwe);
}
