#!/usr/bin/perl

use strict;


my @chrs = (1..19, 'X', 'Y');
#my @chrs = (1);
print("chrs = @chrs\n");

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


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

#    my $pname = join('', 'platypus/vcf/SNPs.filtered.chr', $ch, '.vcf');
#    my $hname = join('', 'platypus/vcf/snps.filtered.hwe.chr', $ch, '.vcf');

    my $pname = join('', 'gatk/vcf/snps.filtered.chr', $ch, '.vcf');
    my $hname = join('', 'gatk/vcf/snps.filtered.hwe.chr', $ch, '.vcf');

    open(OUT, '>', $hname) or die "Cannot open file $hname for writing, reason $!\n";
    open(IN, '<', $pname) or die "Cannot open vcf file $pname, reason $!\n";

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

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

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

    my ($i, $mouse, $pval);
    my (@row, @mouse_data);

    my $Nhwe = 0;
    my $Nmult = 0;
    my $Ntot = 0;
    my $Nnog = 0;

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

	chomp($line);
	@row = split("\t", $line);

	$Ntot++;

	if ( length($row[4] ) > 1 ) {

	    print(OUT "$line\n");
	    $Nmult++;

	} else {

	    my $Nmiss = 0;
	    my $Nref = 0;
	    my $Nhet = 0;
	    my $Nalt = 0;

	    for ($i=9; $i<29; $i++) {
		$mouse = $row[$i];
		@mouse_data = split(':', $mouse);

		my $geno = $mouse_data[0];

		if ( $geno eq './.' ) {
		    $Nmiss++;
		} elsif ( $geno eq '0/0' ) {
		    $Nref++;
		} elsif ( $geno eq '1/1' ) {
		    $Nalt++;
		} elsif ( $geno eq '1/0' || $geno eq '0/1' ) {
		    $Nhet++;
		} else {
		    print("ERROR - unrecognised genotype geno $geno  bp $row[1]  mouse $i  data $mouse\nrow @row\nline $line\n");
		    exit;
		}
	    }

	    # remove SNPs with all missing data

	    if ( $Nmiss == 20 ) {

		$Nnog++;

	    } else {

		# filter on HWE
		$pval = snphwe( $Nhet, $Nref, $Nalt );
#		$pval = snphwe( 0, 19, 1);
#		print("$row[1]  $pval  $Nhet $Nref $Nalt\n");

		if ( $pval < 0.001 ) {

		    if ( $Nhet > ($Nref + $Nalt) ) {
			$Nhwe++;
		    } elsif ( $pval < 0.0001 ) {
			$Nhwe++;
		    } else {
			print(OUT "$line\n");
		    }

		} else {
		    print(OUT "$line\n");
		} # end of HWE
	    } # end of no genotypes loop
	} # end of mult ALT loop
    } # end of while loop
    close(OUT);
    close(IN);

    print("Total SNPs $Ntot\nHW disequilibrium $Nhwe\nMult ALTs $Nmult\nNo genotypes remaining $Nnog\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);
}
