#!/usr/bin/perl -w 
#Diploid Mode Output
#This will take as input HAPMIX 16 column file using 38:1 option and create 
#a) 4 phased EIGENSTRAT files: 2 for local ancestry at each chromosome, and 2 for genotypes
use Switch;

$DIR = $ARGV[0];
$NADMIX = $ARGV[1];
$CHR = $ARGV[2];
$thresh = $ARGV[3];
$admixpop = $ARGV[4];

if($thresh == -1)
  {
    $thresh = 0.0;
  }

for($n=0;$n<$NADMIX;$n++)
  {
    open(ANCFILE,"$DIR/$admixpop.DIPLOID.$n.$CHR") ||die("COF");

    @row = <ANCFILE>;
    $nsnps = scalar @row;
    for($i = 0; $i < scalar @row; $i++)
      {
	chomp($row[$i]);
	$key=$n.":".$i;
	@data = split(' ',$row[$i]);
	if((scalar @data == 16) || (scalar @data == 17))
	  {
	    #Need to convert the 16 column file into a 10 column file 
	    #by adding probabilities of certain interchangebale events
	    #and then we will use a threshold
	    
	    #Check if any of them is greater than 0.9
	    if(scalar @data == 16)
	      {
		$data[0] = $data[0];
		$data[1] = $data[1]+$data[4];
		$data[2] = $data[2]+$data[8];
		$data[3] = $data[3]+$data[12];
		$data[4] = 0.0;
		$data[5] = $data[5];
		$data[6] = $data[6]+$data[9];
		$data[7] = $data[7]+$data[13];
		$data[8] = 0.0;
		$data[9] = 0.0;
		$data[10] = $data[10];
		$data[11]=$data[11]+$data[14];
		$data[12] = 0.0;
		$data[13] = 0.0;
		$data[14] = 0.0;
		$data[15] = $data[15];
		$l = 0;
	      }
	    elsif(scalar @data == 17)
	      {
		$data[1]= $data[1];
		$data[2] = $data[2]+$data[5];
		$data[3] = $data[3]+$data[9];
		$data[4] = $data[4]+$data[13];
		$data[5]= 0.0;
		$data[6] = $data[6];
		$data[7] = $data[7]+$data[10];
		$data[8] = $data[8]+$data[14];
		$data[9]=0.0;
		$data[10]=0.0;
		$data[11]=$data[11];
		$data[12]=$data[12]+$data[15];
		$data[13]=0.0;
		$data[14]=0.0;
		$data[15]=0.0;
		$data[16]=$data[16];
		$l = 1;
	      }
	    
	    
#	    $col = -1;
	    $maxind = $l;
	    while($l < scalar @data)
	      {
		if($data[$l] > $data[$maxind]){$maxind = $l;}
		$l++;
	      }
		
	    if($data[$maxind] > $thresh)
	      {
		$col = $maxind;
	      }
	    else
	      {
		$col = -1;
	      }
	      
	    if(scalar @data == 17)
	      {
		$col--;
	      }
	    
	    switch($col)
	      {
		case 0 {$anc1{$key} = 0;$geno1{$key} =0;$anc2{$key}=0;$geno2{$key}=0;}
		case 1 { $anc1{$key} = 0;$geno1{$key} =1;$anc2{$key}=0;$geno2{$key}=0;}
		case 2 { $anc1{$key} = 1;$geno1{$key} =0;$anc2{$key}=0;$geno2{$key}=0;}
		case 3 { $anc1{$key} = 1;$geno1{$key} =1;$anc2{$key}=0;$geno2{$key}=0;}
#		case 4 { $anc1{$key} = 0;$geno1{$key} =1;$anc2{$key}=0;$geno2{$key}=0;}
		case 5 { $anc1{$key} = 0;$geno1{$key} =1;$anc2{$key}=0;$geno2{$key}=1;}
		case 6 { $anc1{$key} = 1;$geno1{$key} =0;$anc2{$key}=0;$geno2{$key}=1;}
		case 7 { $anc1{$key} = 1;$geno1{$key} =1;$anc2{$key}=0;$geno2{$key}=1;}
#		case 8 { $anc1{$key} = 1;$geno1{$key} =0;$anc2{$key}=0;$geno2{$key}=0;}
#		case 9 { $anc1{$key} = 1;$geno1{$key} =0;$anc2{$key}=0;$geno2{$key}=1;}
		case 10 { $anc1{$key} = 1;$geno1{$key} =0;$anc2{$key}=1;$geno2{$key}=0;}
		case 11 { $anc1{$key} = 1;$geno1{$key} =1;$anc2{$key}=1;$geno2{$key}=0;}
#		case 12 { $anc1{$key} = 1;$geno1{$key} =1;$anc2{$key}=0;$geno2{$key}=0;}
#		case 13 { $anc1{$key} = 1;$geno1{$key} =1;$anc2{$key}=0;$geno2{$key}=1;}
#		case 14 { $anc1{$key} = 1;$geno1{$key} =1;$anc2{$key}=1;$geno2{$key}=0;}
		case 15 { $anc1{$key} = 1;$geno1{$key} =1;$anc2{$key}=1;$geno2{$key}=1;}
		else {$anc1{$key} = 9;$geno1{$key} =9;$anc2{$key}=9;$geno2{$key}=9;}
	      }
	  }
      }
    close ANCFILE;
  }


open(EIGENANCH1FILE,">$DIR/$admixpop.DIPLOID.ANCHAP1.$CHR") || die("COF");
open(EIGENANCH2FILE,">$DIR/$admixpop.DIPLOID.ANCHAP2.$CHR") || die("COF");
open(EIGENGENOH1FILE,">$DIR/$admixpop.DIPLOID.GENOHAP1.$CHR") || die("COF");
open(EIGENGENOH2FILE,">$DIR/$admixpop.DIPLOID.GENOHAP2.$CHR") || die("COF");

for($j = 0; $j < $nsnps;$j++)
  {
    for($k=0;$k<$NADMIX;$k++)
      {
	$key = $k.":".$j;
	print EIGENANCH1FILE (1.0-$anc1{$key});
	print EIGENANCH2FILE (1.0-$anc2{$key});
	print EIGENGENOH1FILE $geno1{$key};
	print EIGENGENOH2FILE $geno2{$key};
      }
    print EIGENANCH1FILE "\n";
    print EIGENANCH2FILE "\n";
    print EIGENGENOH1FILE "\n";
    print EIGENGENOH2FILE "\n";
  }