#!/usr/bin/perl

use strict;
use warnings;
use List::Util (qw(max reduce));

# command line arguments #
my $destdir = $ARGV[0];
my $hapguess = $ARGV[1];
my $bafphased = $ARGV[2];
my $switchpfile = $ARGV[3];
if($#ARGV>4) { srand($ARGV[4])}  # use specified seed, if provided
else{srand()}

# read in data #
my $tp;
my $hbtmp;
my @hb;
my $hgtmp;
my $hgcomptmp;
my @hg;
my @hgcomp;
my @hetlocs_1index;
my @hetlocs;
my $prefix;

($prefix) = $bafphased =~ m=([^/]+)\.[^\./]+$=;
my $hetfile = "$destdir/$prefix.informative";
my $paramfile = "$destdir/$prefix.finalparams";
my $ppfile = "$destdir/$prefix.postprobs";

# baf-phased haplotypes 
open(HB, $bafphased) or die("Could not open $prefix.baf_phased_haplotypes");
<HB>;  # discard 2 header lines
<HB>;
$hbtmp = <HB>;
@hb = split(/\s+/, $hbtmp);  # alleles for 1st baf-phased haplotype
close(HB);

# statistical haplotypes 
open(HG, "$hapguess") or die("Could not open $hapguess");
$hgtmp = <HG>;
$hgcomptmp = <HG>;
@hg = split(/\s+/, $hgtmp);	# alleles for 1st estimated germline haplotype
@hgcomp = split(/\s+/, $hgcomptmp);  # alleles for 2nd estimated germline haplotype
close(HG);

if($#hb != $#hg) {die(sprintf"Error: Lengths of baf-phased and estimated germline haplotypes are not equal ( %d != %d )\n", ($#hb+1), ($#hg+1));}

# informative marker indices 
open(HETS, $hetfile);
my $hettemp = <HETS>;
close(HETS);
@hetlocs = split(/\s+/, $hettemp);

# emission parameter vector 
open(FINALPARAMS, $paramfile);
<FINALPARAMS>;
my @finalparams = split(/\s+/,<FINALPARAMS>);

my @alphas1 = @finalparams[1..$#finalparams];

# postprobs 
# to determine max a posteriori event state 
# and grab emission prob for that state 
# at each marker
my @avector;

open(PP, $ppfile) or die("could not open file $ppfile");
my @pp; 
<PP>; # discard header
while(<PP>) {
     my @ppx = split(/\s+/, $_);  # postprobs for this marker
     my $MAP_state = reduce { $ppx[$a] > $ppx[$b] ? $a : $b } 1..$#ppx;  # exclude state 0; if tie, lastmost state will be returned 
     	  push(@avector, $alphas1[$MAP_state]);
	 
}
close(PP);

## prepare switchprobs vector ##
my @switchprobs;
open(SWITCHACC, "$switchpfile") or die("Error: Couldn't find switchprobs file $switchpfile");
my @raw = split(/\s+/, <SWITCHACC>);
@switchprobs = map(1 - (0.5 + abs($_ - 0.5)), @raw);

if($#switchprobs==0) {   # if average switch accuracy is supplied instead of interval-specific accuracies
     @switchprobs[1..$#hetlocs] = ($switchprobs[0]) x $#hetlocs;	# use average value for every interval
}

## prepare hbhets and hghets ##
my @hbhets;
my @hghets;
my @hgcomphets;
@hbhets = @hb[@hetlocs];
@hghets = @hg[@hetlocs];
@hgcomphets = @hgcomp[@hetlocs];

## check lengths ##
if($#avector+1 != ($#hbhets)) {
     die("Error: Number of informative markers and number of posterior probabilities are not compatible (" 
	 . ($#hbhets+1) . " != " . ($#avector+1) . "+1)"); 
}

if($#switchprobs != $#hbhets) { 
     die("Error: Number of values in switchprobs file does not equal number of informative markers (" 
	 . ($#switchprobs+1) . " != " . ($#hbhets+1) . ")"); 
}

## prepare y vector using hb and hg ##
my @y = (-9);
foreach my $i (0..$#hbhets) {
     my $b = $hbhets[$i];
     my $g = $hghets[$i];
     my $ym;

     if($b eq "A") {
	  $ym = ($g eq "A") ? 0 : 1; 
     }
     if ($b eq "B") {
	  $ym = ($g eq "B") ? 0 : 1;
     }
     push(@y, $ym);
}

# line up emission probs and y vector
push(@avector, $avector[$#avector]);  # duplicate last emission parameter for rightmost marker (results in length increase)
unshift(@avector, -9); # add dummy value to make avector correspond to y vector (results in length increase)

=pod
foreach my $i (0..$#del) { # grab emission probability for state with highest post prob
     if($del[$i] > $cn[$i]) {  
	  push(@avector, $alphas1[0]);
     }
     elsif ($del[$i] < $cn[$i]) {
	  push(@avector, $alphas1[1]);
     }	
     else { push(@avector, $alphas1[int(rand(1))]) }
}
=cut

## HMM ##
sub define_pi {
     my $m = shift(@_);
     my @pi = ();
     
     $pi[0] = [1-$switchprobs[$m], $switchprobs[$m]];
     $pi[1] = [$switchprobs[$m], 1-$switchprobs[$m]];

     return(@pi);
}


sub define_alpha {
     my $m = shift(@_);
     my @alpha = (1-$avector[$m], $avector[$m]);
     return(@alpha);
}

my $M = scalar(@y) - 1;  # number of informative markers

my $sta0;
my $sta1;
$sta0 = 0.5;
$sta1 = 0.5;

# fill forward
my @phi=();
foreach my $m (0..($M)){
	foreach my $a (0,1){
		$phi[$a][$m] = 0;
	}
}

$phi[0][0] =  $sta0;
$phi[1][0] =  $sta1;

foreach my $m (1..($M)){
     my @pi = &define_pi(($m-1));
     my @alpha = &define_alpha($m);

     my $norm = 0;
	foreach my $b (0,1){
		my $sum = 0;
		foreach my $a (0,1){
			$sum += $phi[$a][$m-1] * $pi[$a][$b];
		}
		my $pdat = $alpha[$b]*($y[$m] == 1) + (1 - $alpha[$b])*($y[$m] == 0); # $y[$m] acts like indicator
		$phi[$b][$m] = $sum * $pdat;
		$norm += $phi[$b][$m];
	}

	foreach my $b (0,1){
		$phi[$b][$m] /= $norm ; 
	}
}

# fill backward
my @beta=();
foreach my $m (0..($M+1)){
	foreach my $a (0,1){
		$beta[$a][$m] = 0;
	}
}
$beta[0][$M+1] = 1.;
$beta[1][$M+1] = 1.;

my $m;
for ($m = $M ; $m > 0 ; $m--){
     my @pi = &define_pi($m-1);
     my @alpha = &define_alpha($m);
     my $norm = 0;
     foreach my $b (0,1){
	  my $sum = 0;
	  foreach my $a (0,1){
	       my $pdat;
	       if ($m == $M){ $pdat = 1.; }
	       else {
		    $pdat = $alpha[$b]*($y[$m+1] == 1) + (1 - $alpha[$b])*($y[$m+1] == 0);
	       }

	       $sum += $pdat * $beta[$a][$m+1] * $pi[$b][$a];
	  }
	  $beta[$b][$m] = $sum;
	  $norm += $beta[$b][$m];
     }
     foreach my $b (0,1){
	  $beta[$b][$m] /= $norm;
     }
}

# calculate marginal probabilities
my @proba = ();
foreach my $m (1..($M)){
     my $p1 = ($phi[0][$m] * $beta[0][$m]) ;
     $p1 /= ( $p1 + ($phi[1][$m] * $beta[1][$m]) );
     $p1 = 1. - $p1;    
     $p1 = sprintf "%.3f", $p1;
     push(@proba , $p1);
}

# recover alleles on excess haplotype #
my @excesshap1 = my @excesshap2 = @hg;  # initialize excess haplotype

foreach my $i (0..($#hetlocs)) {
    my $m = $hetlocs[$i];
    my $p = $proba[$i];
    my $a1, my $a2;

    ($a1, $a2) = ($p <= 0.5) ? ($hghets[$i],$hgcomphets[$i]) : ($hgcomphets[$i],$hghets[$i]);
    ($excesshap1[$m], $excesshap2[$m]) = ($a1,$a2);
}

open(ESTOUT, ">$destdir/$prefix.excesshap_haps");
print ESTOUT ("@excesshap1\n@excesshap2\n");
close(ESTOUT);
