'''
Python re-imagining of haploh_subs.pl
'''
import os
import sys

import random

from scheet import const


def appendLog(logfile, logstring):
    with open(logfile, 'w') as log:
       log.write(logstring + const.NEWLINE)
    pass

def sampleL(states, pL):
    sum = sum(pL)
    pL = [a/sum for a in pL] # normalize

    r = random.random()

    for s in states:
        if r <= sum(pL[0:s]):
            return s
    raise UserWarning("Surprise return from sampleL")


def outputPostprobs(filename, pp):
    with open(filename, 'w') as output:
        output.write(const.TAB.join(pp) + const.NEWLINE)
    pass


def printResults(loglike, M, states, current_params, proba, files):
    # print all final parameter values to file #
    with open(files["paramsfile"], 'w') as final_params:
        final_params.write(const.TAB.join(["p0", "p1", "p2", "p_stay_noloh", "p_stay_loh"]) + const.NEWLINE)
        final_params.write(const.TAB.join(current_params) + const.NEWLINE)
        pass

    probdata_h1 = loglike
    probdata_h0 = -1 * M * math.log(2)
     
    # print prob(non-normal event type) #
    postprobs_for_print = []
    for i in range(len(proba[0])):
        p = 0; 
	for s in range(1,len(states)):
            p = p + proba[s][i];
            pass
        postprobs_for_print.append(p)
        pass
    
    outputPostprobs(files["totpostprobsfile"], postprobs_for_print)
    outputPostprobs(files["delpostprobsfile"], proba[1])     
    if len(states)>1:
        outputPostprobs(files["cnpostprobsfile"], proba[2])
        pass
    pass


def fillForward(M, states, sta, pi, alpha, y, phi):
    for a in states:
        phi[a][0] = sta[a]
        pass
    
    for m in range(1,M):
        for a in states:
            phi[a][m] = 0
            pass
        pass
     
    lognorm = 0;
    for m in range(1,M):
        norm = 0;
        for b in states:
            sum = 0
            for a in states:
                sum += phi[a][m-1] * pi[a][b]
                pass
            pdat = alpha[b]*(y[m] == 1) + (1 - alpha[b])*(y[m] == 0)
            phi[b][m] = sum * pdat
            norm += abs(phi[b][m])
            pass
        lognorm += log(norm)
	 
        for b in states:
            phi[b][m] /= norm; 
            pass
        pass
    return(lognorm, @phi)


def calcForwardProbs(*args):
    loglike, phi = fillForward(args)
    return phi


def calcLogLikelihood(*args):
    loglike, phi = fillForward(args)
    return loglike


def calcBackwardProbs(M, states, pi, alpha, y):
    beta=[]  # backward probabilities
     
    for m in range(M+1): # initialize, since we're starting from the end
        for a in states:
            beta[a][m] = 0
            pass
        pass
     
    for a in states:
        beta[a][M+1] = 1.0
        pass
     
    for m in range(M,0,-1):
        norm = 0;
	for b in states:
            sum = 0
            for a in states:
                if (m == M):
                    pdat = 1.0
                else:
                    pdat = alpha[b]*(y[m+1] == 1) + (1 - alpha[b])*(y[m+1] == 0)
                    pass
                sum += pdat * beta[a][m+1] * pi[b][a];
                pass
            beta[b][m] = sum;
            norm += beta[b][m];
            pass

	for b in states:
            beta[b][m] /= norm
            pass
        pass
    return(beta)


def multiplyForwardBackward(M, states, phi, beta):
    proba = [];
    for a in states:
        for m in range(M):
            proba[a][m] = 0
            pass
        pass

    for m in range(1,M+1):
        sum = 0
        p = []
        for a in states:
            p[a] = phi[a][m] * beta[a][m]
            sum += p[a]
            pass

        if (sum == 0):
            sys.stderr.write("m: Zero sum of phi\n")
            pass
        for a in states:
            p[a] /= sum
            proba[a][(m-1)] = p[a]
            pass
        pass
    return proba


def updateAlpha(M, states, alpha, proba, y, baf, phased, switchprobs, inverse=False):
    for a in states:
        # estimate emission probabilities for all states, including normal states
        phat_est_sum = 0
        for m in range(M):
            phat_est_sum += (proba[a][m] * y[m+1] )
            pass
          
        phat_est = phat_est_sum / sum(proba[a]) if sum(proba[a]) else 0
        alpha[a] = phat_est;
        pass

    if (inverse):
	proba2 = []
	counter1=1
	counter2=2
	counter3=3

	# find max along column dimension of proba
	# develop choosestate
	# randomly sample particular markers 50 times, separating by hg state
	# code abundanceguess with 0 or 1 to indicate which configuration
        # isolate in k0-k2 only those that are of a particular phase guess configuration
        # sort through which samples have a particular phase guess configuration
	# build your own normal function
        # [deleted]
    return alpha


def updateTransitions(M, gamma, states, phi, pi, c):
    p_stay = []

     #sample L at each marker #
     stoL = ["NA" for m in range(M)]
     pL = [phi[a][M-1] for a in states]
     
     stoL[M-1] = sampleL(states, pL)
     
     for m in range(M-2,0,-1):  # sample L at each marker from second-to-last to first
         for a in states:
             pL[a] = phi[a][m] * pi[a][stoL[m+1]] # forward probability * transition probability
             pass
         stoL[m] = sampleL(states, pL)
         pass
     
     # estimate tpm using sEM #
     lengths = [[], []]
     m = 0
     runninglength = 1
     # l is the binary indicator for LOH
     while (m != M-1):  # continue if we're not at end of chain
         m += 1
         if (stoL[m] == stoL[m-1]):  # still in the same state, increment length
             runninglength += 1
         else: # have changed states, store length and reset length to 1
             l = 1*(stoL[m-1] > 0);
             lengths[l].append(runninglength)
             runninglength = 1; 
             pass
         pass

     l = 1*(stoL[m] > 0)  # this is the LOH status indicator at the last marker
     lengths[l].append(runninglength)
     
     for s in [0,1]:
         n_transitions_out = 0
         n_possible_transitions = 0
         n_events = len(lengths[s])+1
         if (n_events != 0):
             if (l != s):
                 n_transitions_out = n_events
                 n_possible_transitions = sum(lengths[s])
                 pass
             if (l == s):
                 n_transitions_out = n_events - 1
                 n_possible_transitions = ( sum(lengths[s]) - 1 )
                 pass
             pass
	            
         p_stay[s] = (( (n_possible_transitions - n_transitions_out) + c[s] ) / 
                      ( n_possible_transitions + gamma) )   # including pseudocounts
         print ("nevents s: n_events " + 
                (n_possible_transitions-n_transitions_out) +
                " n_possible_transitions c[s] gamma p_stay[s]")
         pass
     return p_stay


def runEM(estimateTPM, gamma, M, states, alpha, pstays, y, c, baf, phased, switchprobs, inverse=False):
    # assign parameter values for this iteration #
    sta = [];  # initial distribution
    sta[0] = (1/(1-pstays[0])) / ( (1/(1-pstays[0])) + (1/(1-pstays[1])) )
    for s in range(1,len(states)):
        print("CHECK ME: len(states)+1????")
        sta[s] = (1 - sta[0]) / len(states)
        pass
     
    pi = [] # transition probabilities
    pi[0][0] = pstays[0]
    for s in range(1,len(states)):
        print("CHECK ME: len(states)+1????")

        pi[s][s] = pstays[1];
        pi[0][s] = (1 - pi[0][0]) / len(states);
        pi[s][0] = 1 - pi[s][s];
        for s2 in range(1,len(states)):
            print("CHECK ME: len(states)+1????")
            if s != s2:
                pi[s][s2] = 0  # no transitions between event states
                pass
            pass
        pass
     
    phi = []
    beta = []
    proba = []
    phi = calcForwardProbs(M, states, sta, pi, alpha, y)
    beta = calcBackwardProbs(M, states, pi, alpha, y)
    proba = multiplyForwardBackward(M, states, phi, beta)
    alpha = updateAlpha(M, states, alpha, proba, y, baf, phased, switchprobs, inverse)
     
    if estimateTPM:
        pstays = updateTransitions(M, gamma, states, phi, pi, c)
        pass

    loglike = calcLogLikelihood(M, states, sta, pi, alpha, y)
     
    return([alpha, pstays], proba, loglike)
}
