# These are versions of some of our important functions
# that were duplicates or non-generalized forms that I
# am retiring to here in case something goes wrong
import pandas as pd
import numpy as np

def get_probabilities_2state_sys(states):
    num_total = np.shape(states)[0] * np.shape(states)[1]
    
    num_ones = np.shape(np.where(states == 1))[1]
    num_zeros = num_total - num_ones

    prob_one = num_ones/num_total
    prob_zero = num_zeros/num_total
    
    assert round(prob_one+prob_zero, 3) == 1. 
    
    return prob_one, prob_zero

def get_conditional_probability_2parent_2state(Astates, Bstates, Cstates): #I'm in progress of generalizing this in probabilities.ipynb, but feel free to pick up where I left off
    """A and B are parent nodes. C is child node"""
    
    num111, num101, num011, num001, num110, num100, num010, num000 = 0, 0, 0, 0, 0, 0, 0, 0
    num11, num01, num10, num00 = 0,0,0,0
    
    
    #Add up cases where C=1
    indC_one_row = np.where(Cstates == 1)[0] #index of row with a 1
    indC_one_col = np.where(Cstates == 1)[1] #index of corresponding coln with a 1

    for i in range(len(indC_one_col)): # loop through all times C=1
        A = Astates.iloc[indC_one_row[i], indC_one_col[i]]
        B = Bstates.iloc[indC_one_row[i], indC_one_col[i]]
        if A == 1:
            if B == 1:
                num111 += 1  #C=1, B=1, A=1
                num11 += 1  #B=1, A=1
            elif B == 0:
                num101 += 1  #C=1, B=0, A=1
                num01 += 1  #B=0, A=1
        elif A == 0:
            if B == 1:
                num110 += 1  #C=1, B=1, A=0
                num10 += 1  #B=1, A=0
            elif B == 0:
                num100 += 1  #C=1, B=0, A=0
                num00 += 1  #B=0, A=0
                
    # Add up states where C=0
    indC_zero_row = np.where(Cstates == 0)[0] 
    indC_zero_col = np.where(Cstates == 0)[1] 


    for i in range(len(indC_zero_col)): # loop through all times C=0
        A = Astates.iloc[indC_zero_row[i], indC_zero_col[i]]
        B = Bstates.iloc[indC_zero_row[i], indC_zero_col[i]]
        if A == 1:
            if B == 1:
                num011 += 1  #C=0, B=1, A=1
                num11 += 1  #B=1, A=1
            elif B == 0:
                num001 += 1  #C=0, B=0, A=1
                num01 += 1  #B=0, A=1
        elif A == 0:
            if B == 1:
                num010 += 1  #C=0, B=1, A=0
                num10 += 1  #B=1, A=0
            elif B == 0:
                num000 += 1  #C=0, B=0, A=0
                num00 += 1  #B=0, A=0
    
    # Calculate conditional probabilities
    num_total = np.shape(Astates)[0] * np.shape(Astates)[1]
    
    """The variable names are in the order CBA (cases, testing, stay at home). """
    P111 = num111/num11
    P101 = num101/num01
    P110 = num110/num10
    P100 = num100/num00
    P011 = num011/num11
    P001 = num001/num01
    P010 = num010/num10
    P000 = num000/num00
    
    states = np.array(['0|00','1|00','0|01','1|01','0|10','1|10','0|11','1|11']) #I believe this is the order Matt was talking about with top to bottom, left to right in the table from the paper
    conditional_probabilities = np.array([P000,P100,P001,P101,P010,P110,P011,P111]) #in this order, the first two, second two, etc. should each add to 1 as pairs (columns in the paper's table)
    
    # Make graph for circuit input
    ProbA1, ProbA0 = get_probabilities_2state_sys(Astates)
    ProbB1, ProbB0 = get_probabilities_2state_sys(Bstates)
    
    graph = { 'StayAtHome': ([], [ProbA0, ProbA1]), # P(A = 0), P(A = 1)
                'Testing': ([], [ProbB0, ProbB1]),  #P(B = 0), P(B = 1)
                #P(C=0|A=0,B=0), P(C=1|A=0,B=0), P(C=0|A=0,B=1), P(C=1|A=0,B=1), P(C=0|A=1,B=0), P(C=1|A=1,B=0), P(C=0|A=1,B=1), P(C=1|A=1,B=1)
                'Cases': (['StayAtHome','Testing'], [P000, P100, P010, P110, P001, P101, P011, P111])
              }
    
    return states, conditional_probabilities, graph

def Margies_marginal_probabilities(ntwk_results): #a combined version of this with Ella's is kept in the real code
    #ntwk_results: dict, counts resulting from network run (should have 2^n entries)
    #marg_probs: array of length n, marginal probabilities that each qubit is 0,
        #from most significant to least significant
    
    n = len(list(ntwk_results.keys())[0])
    prob = np.zeros(n)
    total = sum(ntwk_results.values())

    for i in range(n):
        for key in ntwk_results:
            if int(key[i]) == 0:
                prob[i] += ntwk_results[key]
        prob[i] = prob[i]/total
    
    return prob

def Ella_3qubit_marginal_probabilities(pstates_dict):
    # Works for 3 qubits only
    C0, B0, A0, numtot = 0, 0, 0, 0

    for key in pstates_dict:
        Cstate=key[0]
        Bstate=key[1]
        Astate=key[2]
        numtot += pstates_dict[key]
        if Cstate == '0':
            C0 += pstates_dict[key]
        if Bstate == '0':
            B0 += pstates_dict[key]
        if Astate == '0':
            A0 += pstates_dict[key]
        
    ProbC0 = C0/numtot
    ProbB0 = B0/numtot
    ProbA0 = A0/numtot

    return ProbC0, ProbB0, ProbA0

def Ella_marginal_probabilities_general(pstates_dict):
    """Works for n qubits"""
    n = len(list(pstates_dict.keys())[0]) #number of qubits in state
    Probs = np.empty(n)
    numZeros, numtot = np.zeros(n), 0

    for key in pstates_dict:
        numtot += pstates_dict[key]
        for i in range(n):
            #print(i)
            state=key[i]
            if state == '0':
                numZeros[i] += pstates_dict[key]
    for i in range(n):
        Probs[i] = numZeros[i]/numtot

    return Probs

# Test the marginal prob functions with some simulated output from our 3 qubit run!

#pick the function to use (comment others out)
funct = Ella_marginal_probabilities_general
#funct = Ella_3qubit_marginal_probabilities
#funct = Margies_marginal_probabilities

result1 = {'000': 2783, '001': 1240, '100': 603, '111': 815, '110': 294, '010': 1712, '101': 485, '011': 260}
print(funct(result1))

def get_lesser_model_states():
    statedataStayHome = {'MarHome' : data['MarHome'], 'AprHome' : data['AprHome'], 'MayHome' : data['MayHome'], 'JunHome' : data['JunHome']}
    statesStayHome = pd.DataFrame(data=statedataStayHome)

    statedataTests = {'MarTest' : data['MarTest'], 'AprTest' : data['AprTest'], 'MayTest' : data['MayTest'], 'JunTest' : data['JunTest']}
    statesTests = pd.DataFrame(data=statedataTests)

    statedataCases = {'MarCases' : data['MarCases'], 'AprCases' : data['AprCases'], 'MayCases' : data['MayCases'], 'JunCases' : data['JunCases']}
    statesCases = pd.DataFrame(data=statedataCases) # 0 = increasing. 1 = flat or decreasing
    
    return statesStayHome, statesTests, statesCases

data = pd.read_csv('data/lesser_model_data.csv')
Astates, Bstates, Cstates = get_lesser_model_states()
print(get_conditional_probability_2parent_2state(Astates, Bstates, Cstates))

""" # ALTERNATE SECTION OF GENERALIZED CONDITIONAL PROBABILITY CODE TO CONSIDER SHORTENING BY LIKE 3 LINES

def f(c, *ps):
	for key in keys:
		num_c, tot_c = 0, 0
		n = len(c)
		for i in range(n):
			all_ps = all([ps[j][i] == key[j] for j in range(len(ps))])
			if all_ps:
				tot_c += 1
				if c[i] == thing before |:
					num_c += 1

"""

def generate_cond_keys_old(s_0, s_i):
    ##############################################
    #THIS FUNCTION WILL GENERATE A LIST OF STRINGS TO USE AS KEYS FOR CONDITIONAL PROBABILITIES
    ### INPUT ###
    # s_0    int    number of states of the child node
    # s_i    list   number of states for each parent node, from most to least significant
    
    ### OUTPUT ###
    # list of strings to use as keys for conditional probabilities (included commas in case there is ever an >11-state node!)
    ##############################################
    
    ranges = [range(0, elem) for elem in list([s_0])+list(s_i)]
    enumed = product(*ranges)
    
    cond_keys = []
    for enum in enumed:
        enum = list(enum)
        parent_str = ",".join(str(x) for x in enum[1:])
        cond_keys.append("%s|%s"%(str(enum[0]), parent_str))
        
    return cond_keys