# rgf: for the moment, PARAMETER lives on, but in python form
# PERLCONST[TPM] == $CONSTANT{TPM_FIXED}

import os
import sys

import haploh
import bw9_subs

PERLCONST = haploh.PERLCONST

def usage():
    print("""
Usage:
    {0} parameter=value parameter=value ...
Parameters are
    {RANDOM_SEED}
    {TPM}
    {TPM_FIXED}
    {TPM_ESTIMATE}
    {INITIAL_ALPHAS}
    {MAX_ITERATIONS}
# For fixed and estimated TPM
    {MEAN_INFORMATIVE_MARKER_COUNT}
    {EVENT_PREVALENCE}
# For estimated TPM
    {GAMMA}
# For Christina
    {INVERSE}
    {BAF}
    {PHASED}
    {SWITCHPROBS}
""".format(sys.argv[0], **PERLCONST))
    pass

my SFX_SWITCH_ENUMERATION = "switch_enumeration";
my SFX_EM_LOG = "EM_log";
my SFX_PHAT = "phat";
my SFX_FINALPARAMS = "finalparams";
my SFX_LOGLIKELIHOODS = "loglikelihoods";
my SFX_POSTPROBS = "postprobs";
my SFX_DEL_POSTPROBS = "postprobs_state1";
my SFX_CNLOH_POSTPROBS = "postprobs_state2";
my SFX_CHRISTINA = "christina";

if PERLCONST[RANDOM_SEED]:
   pass # already done by haploh.py

filenames = {
     "switchfile" : "PERLCONST[PREFIX].SFX_SWITCH_ENUMERATION",
     "logfile" : "PERLCONST[PREFIX].SFX_EM_LOG",
     "paramsfile" : "PERLCONST[PREFIX].SFX_FINALPARAMS",
     "totpostprobsfile" : "PERLCONST[PREFIX].SFX_POSTPROBS",
     "delpostprobsfile" : "PERLCONST[PREFIX].SFX_DEL_POSTPROBS",
     "cnpostprobsfile" : "PERLCONST[PREFIX].SFX_CNLOH_POSTPROBS",
     "christinafile" : "PERLCONST[PREFIX].SFX_CHRISTINA",
}

def baum_welch(random_seed,
               tpm,
               tpm_fixed,
               tpm_estimate,
               initial_alphas,
               max_iterations,
# For fixed and estimated TPM
               mean_informative_marker_count,
               event_prevalence,
# For estimated TPM
               gamma,
# For Christina
               inverse=False,
               baf=None,
               phased=None,
               switchprobs=None):

    # define states # 
    n_event_states = 2
    states = [n for n in range(n_event_states)]
    # state 0 is normal, others are flexible

    # define initial parameter values, which may or may not be updated #
    alphas = const.COMMA.split(initial_alphas)
    evlength = mean_informative_marker_count
    evprev = event_prevalence
    lambdas = [evprev/((1.0-evprev)*evlength), (1.0/evlength)]
    pstays = [1-L for L in lambdas]

    # calculate pseudocounts
    c = [gamma * pstays[i] for i in range(2)]

    ## read in data ##
    with open(filenames["switchfile"]) as IN:
        y = [-9]; #(0,0,0,1, 1)
        for line in [l.strip() for l in IN.readlines()]:
            atmp = line.split()
            y.append(atmp)
            pass
        pass
    M = len(y) - 1  # number of markers

    if inverse:
        with open(baf) as IN:
            baf = IN.read().strip().split()
            pass

        with open(phased) as IN:
            phased = IN.read().strip().split()
            pass

        with open(switchprobs) as IN:
            switchprobs = IN.read().strip().split()
            pass
               
        pass

    ## prepare EM log file ##
    with open(filenames["logfile"], "w") as LOG:
        print LOG const.TAB.join("p0", "p1", "p2", "p_stay_noloh", "p_stay_loh", "log-likelihood", "convergence")
        pass

    ## set up sets of starting alpha values for 2-state model##
    alphastarts = [[0.5, 0.6], [0.5, 0.7], [0.5, 0.8],
                   [0.5, 0.9], [0.6, 0.7], [0.6, 0.8],
                   [0.6, 0.9], [0.7, 0.8], [0.7, 0.9], [0.8, 0.9]]
    final_loglike = "-inf"

    for alphas in alphastarts:
        iteration_count = 0
        prev_loglike = "-inf" # initial loglikelihood value to allow comparisons 
        margprobs = []

        appendLog(filenames["logfile"], "# start " . const.TAB.join(alphas + pstays))

        ## perform EM for this set of starting parameter values ##
        for iteration_count in range(max_iterations):
            print("iteration %d" % (iteration_count))
     
            funcout = runEM(tpm_estimate, gamma, M, states, alphas, pstays, y, c, baf, phased, switchprobs, inverse)
            alphas, pstays = funcout[0]
            margprobs = funcout[1]
            loglike = funcout[2]
             
            if loglike - prev_loglike < 0.0001:
                # we've converged, break
                break
	 
            prev_loglike = loglike

            appendLog(filenames["logfile"], const.TAB.join(alphas + pstays + [loglike]))
            pass
     
        ## store results if this start produces max log-likelihood ##
        if loglike > final_loglike:
            final_loglike = loglike
            final_margprobs = margprobs
            pass
        pass

    printResults(final_loglike, M, states, alphas, pstays, final_margprobs, filenames)
    appendLog(filenames["logfile"],
              "mean informative marker count is %d, prevalence is %d", (evlength, evprev))


