#!/usr/bin/perl -w
use strict;
use warnings;
use Statistics::Distributions qw< chisqrprob >;


my $hapmapfile = $ARGV[0];
my $pvalue_cutoff= 0.001;

open IN, "$hapmapfile";
open OUT, ">filtered.hmp";
open JM, ">filtered.loc";
open COUNT, ">count.txt";
my $line = <IN>;
chomp $line;
print OUT $line, "\n";
my @inds = split "\t", $line;
@inds = @inds[11..$#inds];
my $inds_count = @inds +0;
my $inds_line = join "\n", @inds;
my $num_loci =0;

my %iupac = ("M"=>["M", "A", "C"], 
			"R"=>["R", "A", "G"],
			"W"=>["W", "A", "T"],
			"S"=>["S","C", "G"],
			"Y"=>["Y", "C", "T"],
			"K"=>["K", "G", "T"]
);

my %iupac2 = ("M"=>"AC", 
			"R"=>"AG",
			"W"=>"AT",
			"S"=>"CG",
			"Y"=>"CT",
			"K"=>"GT"
);
my $JMlines = "";
LOOP:while ($line=<IN>) 
{
	chomp $line;
	my @data = split "\t", $line;
	my $site = $data[0];
	my @gdata = @data[11..$#data];
	my $p1 = pop @gdata;

	#my @genotypes = @{$iupac{$p1}};
	my %tophomoalleles = ();
	foreach  (@gdata) 
	{
		if (/[ACGTMRWSKY]/) 
		{
			$tophomoalleles{$_} ++;
		}
	}
	my @tophomoalleles =reverse sort {$tophomoalleles{$a} <=> $tophomoalleles{$b}} (keys %tophomoalleles);
	my $hetgenotype = "";
	#print @tophomoalleles, "\n"; 
	LOOPH:foreach  (@tophomoalleles) 
	{
		if (/[MRWSYK]/) 
		{
			$hetgenotype = $_;
			last LOOPH;
		}
	}
	
	unless ($hetgenotype=~/[MRWSYK]/) 
	{
		next LOOP;
	}
	

	my @genotypes = ($hetgenotype, (split "", $iupac2{$hetgenotype}));
	if (@genotypes <3) 
	{
		next LOOP;
	}


	my %ghash = ();
	foreach  (@genotypes) {$ghash{$_} =0;}
	foreach  (@gdata) 
	{
		$ghash{$_} ++;
	}
	my @observed = ($ghash{$genotypes[0]}, $ghash{$genotypes[1]}, $ghash{$genotypes[2]});
	my $sum = $observed[0] + $observed[1] + $observed[2];
	my @expected = ($sum * 0.5, $sum * 0.25, $sum*0.25);
	my $p = chi_squared_test(\@observed, \@expected);
	
	foreach  (@genotypes) 
	{
		print COUNT $_, "\t", $ghash{$_}, "\t";
	}
	print COUNT "\n";
	my $pt_h = substr($iupac2{$hetgenotype}, 0, 1);
	my $pt_k = substr($iupac2{$hetgenotype}, 1, 1);
	if ($p>$pvalue_cutoff) 
	{
		$num_loci ++;
		print OUT $line, "\n";
		$JMlines .= "$site\t<hkxhk>\n";
		foreach my $hmpgt (@gdata) 
		{
			if ($hmpgt=~/^[RYSWKM]$/) 
			{
				$JMlines .= "hk\t";
			}
			elsif ($hmpgt eq $pt_h) 
			{
				$JMlines .= "hh\t";	
			}
			elsif ($hmpgt eq $pt_k)
			{
				$JMlines .= "kk\t";	
			}
			else
			{
				$JMlines .="--\t"
			}
		}
		$JMlines=~s/\t$//;
		$JMlines .= "\n";
	}
}

print JM <<"EOT";
name = NAME
popt = F2
nloc = $num_loci
nind = $inds_count

$JMlines

individual names:
$inds_line
EOT

sub chi_squared_test {
  my $observed = shift ;
  my $expected = shift ;
  my $chi_squared=0;
  foreach my $o (@$observed) 
  {
	  my $e = shift @$expected;
	  if ($e==0) 
	  {
		  return 0;
	  }
	  $chi_squared += ($o - $e)**2 / $e;
  }
  my $degrees_of_freedom = @$observed - 1;
  my $probability = chisqrprob($degrees_of_freedom, $chi_squared);
  return $probability;
}
