#hash stuff
def create_hash_table(m):
    """ initial hash table, m empty entries """
    return [[] for i in range(m)]

def find(htable, key):
    """ return True if key is in htable, else False """
    ht_size = len(htable)
    index = hash(key) % ht_size
    if key in htable[index]:
        return index
    return None

def insert(htable, key):
    """ insert key into htable """
    if find(htable, key) == None:
        ht_size = len(htable)
        index = hash(key) % ht_size
        htable[index].append(key)

def repeat_hash(lst):
    size = len(lst)
    htable = create_hash_table(size)
    for s in lst:
        if find(htable,s) != None:
            return s
        insert(htable, s)
    return None

#########################################
## Boolean model for regulatory networks
#########################################

def update_node(states, i):
    ''' return new state of node i given states vector '''
    n = len(states)
    s = 0
    
    for j in range(n):
        s += edges[j][i]*states[j]
    if s>0:
        new = 1
    elif s<0:
        new = 0
    else: #s==0
        new = states[i]
    return new

def update(states):
    ''' return the next states vector in the next time step - apply transition function to all nodes '''    
    n = len(states)
    next_states = [-1]*n
    for i in range(n):
        next_states[i] = update_node(states, i)
    return next_states


def run(init, track = False):
    ''' start with init vector, and run until steady state '''

    states = init
    next_states = update(states)
    if track:
        print(states)
    insert(ht,tuple(states))
    #print(ht)
    c = 1 #if you want to count steps
    
    while states != next_states \
          and find(ht,tuple(next_states))==None:
        states = next_states
        next_states = update(states)
        if track:
            print(states)

        c += 1
        insert(ht,tuple(states))
        #print(ht)

    if states == next_states:
        return states
    else:
        print("LOOP!!!")



import itertools

def basin_size(attractor):
    ''' count how many initial vectors end in the attractor vector '''
    cnt=0

    for states in ["".join(seq) for seq in itertools.product("01", repeat=len(nodes))]:
        states = [int(s) for s in states]
        final = run(states)
        if final == attractor:
            cnt+=1
    return cnt

#############################################
### check input correctness
def check():
    n = len(nodes)
    assert len(nodes) == len(init)
    assert all([len(edges[i]) == n for i in range(n)])


#############################################
### Input network - example for infinite loop

nodes = ['A','B'] 

init = [1, 0]

edges = [
    [-1,1],
    [1,-1]
    ]

check()
ht = create_hash_table(len(nodes))
run(init, True)

"""
#############################################
### Input network - toy example from lecture

nodes = ['A','B','C','D','E','F'] 

init = [0, 0, 0, 1 , 0, 0]

edges = [
    [0,0,0,1,0,0],
    [-1,0,0,-1,-1,0],
    [-1,0,0,0,0,0],
    [1,0,0,0,1,0],
    [0,0,0,0,0,1],
    [0,0,0,-1,1,-1]
    ]

check()
run(init, True)

#############################################
### Input network - yeast cell-cycle
### "The Yeast Cell-Cycle Network is Robustly Designed", Li et.al, PNAS, 2004

nodes = ['Cln3','MBF','SBF','Cln1,2','Cdh1','Swi5','Cdc20/Cdc14','Clb5,6','Sic1','Clb1,2','Mcm1/SFF'] 

init = [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0] #simulation of this is in the paper

edges = [
    [-1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
    [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, -1, -1, 0, 0, 0, -1, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, -1, 0],
    [0, 0, 0, 0, 0, -1, 0, 0, 1, 0, 0],
    [0, 0, 0, 0, 1, 1, -1, -1, 1, -1, 0],
    [0, 0, 0, 0, -1, 0, 0, 0, -1, 1, 1],
    [0, 0, 0, 0, 0, 0, 0, -1, 0, -1, 0],
    [0, -1, -1, 0, -1, -1, 1, 0, -1, 0, 1],
    [0, 0, 0, 0, 0, 1, 1, 0, 0, 1, -1]
    ]

check()
run(init, True)

final = [0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0] #according to paper, 1764 initial vectors end here
print(basin_size(final), "initial vectors end at", final)
"""