#!/usr/bin/perl
# Copyright 2012-2017 Selina Vattathil and Paul Scheet.
# hapLOH is freely available "AS IS" with an MIT license. See LICENSE.txt.
#
## rgf: in order to parameterize this and marry it to driver, which is a separate process,
## and try to reduce the tight coupling between the code bodies themselves, I have extracted
## the parameter NAMES into file called names.constants, which is read during startup.
## The result is a kind of crazy reference that looks like:
## $CONSTANT{SOME_NAME} and $PARAMETER{$CONSTANT{SOME_NAME}}
## I hope it's clear that $CONSTANT{SOME_NAME} is a constant string read from names.constants,
## and $PARAMETER{$CONSTANT{SOME_NAME}} is the value of the parameter named by that constant string
## on the command line, so
## bw9.pl total_probability_matrix=fixed
## would produce the following true boolean expression
## $PARAMETER{$CONSTANT{TPM}} == $CONSTANT{TPM_FIXED}

use strict;
use warnings;
use List::Util qw(first max maxstr min minstr reduce shuffle sum);
use File::Basename;
push @INC, dirname($0);
require Constants;
require "haploh_subs.pl";

my %CONSTANT = Constants::load($ARGV[0]);
shift @ARGV;
my %PARAMETER = ();

sub usage() {
    print "
Usage:
    $0 parameter=value parameter=value ...
Parameters are
    $CONSTANT{RANDOM_SEED}
    $CONSTANT{TPM}
    $CONSTANT{TPM_FIXED}
    $CONSTANT{TPM_ESTIMATE}
    $CONSTANT{INITIAL_ALPHAS}
    $CONSTANT{MAX_ITERATIONS}
# For fixed and estimated TPM
    $CONSTANT{MEAN_INFORMATIVE_MARKER_COUNT}
    $CONSTANT{EVENT_PREVALENCE}
# For estimated TPM
    $CONSTANT{GAMMA}
";
}

my $SFX_SWITCH_ENUMERATION = "switch_enumeration";
my $SFX_EM_LOG = "EM_log";
my $SFX_FINALPARAMS = "finalparams";
my $SFX_POSTPROBS = "postprobs";
#my $SFX_DEL_POSTPROBS = "postprobs_state1";
#my $SFX_CNLOH_POSTPROBS = "postprobs_state2";

my $counter = 1;
foreach my $arg (@ARGV) {
    my ($name, $value) = Constants::split_terms($arg, $counter++, 'Argument');
    if ('ERROR' eq $name) {
        print STDERR "Error in $value";
        usage();
        exit(1);
    }
    printf "%s -> %s\n", $name, $value;
    $PARAMETER{$name} = $value;
}

if ($PARAMETER{$CONSTANT{RANDOM_SEED}}) {
    srand($PARAMETER{$CONSTANT{RANDOM_SEED}});
} else {
    srand();
}

# Complain about non-sets
my @complaints = ();
my ($term, $complaint);
# Note, PARAMETERS refers to a parameter file, which is guilding the lily

for $term (qw(PREFIX RANDOM_SEED TPM INITIAL_ALPHAS MAX_ITERATIONS)) {
    #if (! $PARAMETER{$CONSTANT{$term}}) { push @complaints, "$term not provided"; }
}

if (@complaints) {
    for $complaint (@complaints) { printf STDERR "%s\n", $complaint; }
exit(1);
}

if (! grep /$PARAMETER{$CONSTANT{TPM}}/, ($CONSTANT{TPM_ESTIMATE}, $CONSTANT{TPM_FIXED})) {
    push @complaints, sprintf("%s must be either %s or %s", $CONSTANT{TPM}, $CONSTANT{TPM_FIXED}, $CONSTANT{TPM_ESTIMATE});
} else {
    for $term (qw(MEAN_INFORMATIVE_MARKER_COUNT EVENT_PREVALENCE GAMMA)) {
        if (! $PARAMETER{$CONSTANT{$term}}) {
            push @complaints, sprintf("%s not provided for %s", $term, $CONSTANT{TPM_ESTIMATE});
        }
    }
}

if (@complaints) {
    for $complaint (@complaints) { printf STDERR "%s\n", $complaint; }
exit(1);
}

# define filenames
my %filenames = (
     "switchfile" => "$PARAMETER{$CONSTANT{PREFIX}}.$SFX_SWITCH_ENUMERATION",
     "logfile" => "$PARAMETER{$CONSTANT{PREFIX}}.$SFX_EM_LOG",
     "paramsfile" => "$PARAMETER{$CONSTANT{PREFIX}}.$SFX_FINALPARAMS",
     "postprobsfile" => "$PARAMETER{$CONSTANT{PREFIX}}.$SFX_POSTPROBS",
     #"totpostprobsfile" => "$PARAMETER{$CONSTANT{PREFIX}}.$SFX_POSTPROBS",
     #"delpostprobsfile" => "$PARAMETER{$CONSTANT{PREFIX}}.$SFX_DEL_POSTPROBS",
     #"cnpostprobsfile" => "$PARAMETER{$CONSTANT{PREFIX}}.$SFX_CNLOH_POSTPROBS",
    );

# define number of states and options for estimation # 
my $n_starts = $PARAMETER{$CONSTANT{NUM_STARTS}};
my $n_event_states = $PARAMETER{$CONSTANT{NUM_EVENT_STATES}};
my @event_alpha_range = split(/,/, $PARAMETER{$CONSTANT{EVENT_ALPHA_RANGE}}); 
my @states = (0..$n_event_states); # state 0 is normal, others are flexible

## read in data ##
open(IN, $filenames{"switchfile"});
my @y = (-9); 
while (<IN>){
	chomp;
	my @atmp = split;
	push(@y, @atmp);
}
close IN;
my $M = scalar(@y) - 1;  # number of markers

# define initial parameter values, which may or may not be updated #
# from command line input
my $max_iterations = $PARAMETER{$CONSTANT{MAX_ITERATIONS}}; 
my $estimateTPM = ($PARAMETER{$CONSTANT{TPM}} eq $CONSTANT{TPM_ESTIMATE});
my $gamma = $PARAMETER{$CONSTANT{GAMMA}};
my $evprev = $PARAMETER{$CONSTANT{EVENT_PREVALENCE}};
my $genome_mb = $PARAMETER{$CONSTANT{GENOME_MB}};
my $event_mb = $PARAMETER{$CONSTANT{EVENT_MB}};
my $evlength = $PARAMETER{$CONSTANT{MEAN_INFORMATIVE_MARKER_COUNT}};  # for backward compatibility
my @fixed_alphas = split(/,/, $PARAMETER{$CONSTANT{INITIAL_ALPHAS}}); #optional parameter

if($evlength eq "NA") {                                            # if MIMC argument is not provided, 
     $evlength = $M / $genome_mb * $event_mb; # then calculate mean informative marker count using genomic sizes
     $evlength = sprintf("%d", $evlength+0.5); # round to nearest integer
}

# calculated or assembled using above
my @lambdas = (($evprev/((1-$evprev)*$evlength)), (1/$evlength));
my @pstays =  map {1-$_} @lambdas;
my @c; foreach my $i (0..1) { $c[$i] = $gamma * $pstays[$i] } # calculate pseudocounts

## prepare EM log file ##
open(LOG, ">" . $filenames{"logfile"});
#print LOG join("\t", ("p0", "p1", "p2", "p_stay_noloh", "p_stay_loh", "log-likelihood", "convergence"));
print LOG join("\t", (join("\t", map{"alpha" . $_} @states), "p_stay_normal", "p_stay_event", "log-likelihood", "convergence"));
close(LOG);

## set up sets of starting alpha values ##
my @alphastarts; 
if($fixed_alphas[0] eq "NA") {
     my $avgpc = (sum(@y)+9) / $#y; # average phase concordance; compensate for dummy value in first position of vector
     if ($event_alpha_range[0] eq "NA") {  # use default alpha range
	  @event_alpha_range = ($avgpc+0.05, 0.95);      
     } 
     my $step = ($event_alpha_range[1] - $event_alpha_range[0]) / ($n_starts-1);  # distance between values across starts
     foreach my $s (1..$n_starts) {
	  my $this_step = $event_alpha_range[0] + ($step * ($s-1)); # starting value for lowest event state
	  push(@alphastarts, [($avgpc, $this_step, map { $this_step+(0.001*$_) } 1..$n_event_states-1 )]); # problem if n_starts > ((1-alpha_range[1]) / 0.001)
     }
} else {	 # use alpha parameters supplied at command line and shortcut to 0 EM starts
     push(@alphastarts, [@fixed_alphas]);
     #$max_iterations=0;
}

my $final_loglike = "-inf";
my @final_margprobs;
my @final_params;
my @current_params = (@{$alphastarts[0]}, @pstays);

foreach my $astart (0..$#alphastarts) {
     @current_params[0..$n_event_states] = @{$alphastarts[$astart]};
     my $iteration_count = 0;
     my $prev_loglike = "-inf"; # initial loglikelihood value to allow comparisons 
     my $converged = 0;
     &appendLog($filenames{"logfile"}, "\n# start " . $astart);

     my ($loglike, @margprobs) = &calc_beginning_log_likelihood_and_margprobs(($M, [@states], [@current_params], [@y]));
     &appendLog($filenames{"logfile"}, "\n" . join("\t", (@current_params), $loglike, "-")); ## log current parameter estimates
          
     ## perform EM ##
     while (($iteration_count < $max_iterations) & !$converged) {
	  $iteration_count = $iteration_count+1; # increment iteration count
	  print STDOUT "iteration $iteration_count \n";

	  my @funcout = &runEM(($estimateTPM, $gamma, $M, [@states], [@current_params], [@y], [@c])); # had to create funcout to make assigning output easier
	  @current_params = @{$funcout[0]}; 
	  @margprobs = @{$funcout[1]};
	  $loglike = $funcout[2];

	  $converged = (($loglike - $prev_loglike) < 0.0001 ) ? 1:0; # change convergence criterion here
	 
	  $prev_loglike = $loglike;  # save log-likelihood before next  iteration

	  &appendLog($filenames{"logfile"}, "\n" . join("\t", (@current_params, $loglike, ($converged ? "*" : "-")))); ## log current parameter estimates, log-likelihood, and convergence status
     }
     
     ## store results if this start produces max log-likelihood ##
     if($loglike > $final_loglike) {
	  $final_loglike = $loglike;
	  @final_margprobs = @margprobs;
	  @final_params = @current_params;
     }

}

## print results ## 
&printResults(($final_loglike, $M, [@states], [@final_params], [@final_margprobs], \%filenames)); # pass hash by reference
&appendLog($filenames{"logfile"}, "\ninitial values: average number of informative markers per event: $evlength, prevalence: $evprev\n" . "\n");

