# This program implements a simple 2-node boolean network, computes all possible state transitions (4 of them), and
# plots the resulting STG.
# The Boolean update equations for the 2 nodes (gene_0 and gene_1) are:
#   gene_0 = NOT gene_1
#   gene_1 = gene_0 AND gene_1
# The program then stores the states and the corresponding transitions as nodes and edges in a 'networkx' graph object.
# This graph is nothing but the 'state transition graph' (STG) of the Boolean network that is then plotted.

# IMPORTANT NOTE: The STG is a visualization of the state transitions that shows which states drain into which
# attractors. For a large state space (such as the one for the 10-node BN in the second problem), it may be difficult
# to read off the attractors. Hence, you should write your own logic to identify the attractors in the STG, which would
# basically be a list of states in a particular order.

import networkx as nx
import matplotlib.pyplot as plt

G = nx.DiGraph()  # this only creates the 'graph' object, and does not plot it

# Below is the set of all possible states of 2-gene boolen network
# For your 10-gene network, you have to create this list of all possible states programmatically. DO NOT hard-code them.

possible_states = [(False, False), (False, True), (True, False), (True, True)]
# when represented in Boolean numerals they read as: [(0, 0), (0, 1), (1, 0), (1, 1)]
# corresponding values in Decimals are: [0, 1, 2, 3]

## One way to automatically generate the set of all possible states in as follows:
# possible_states = [(x, y) for x in [False, True] for y in [False, True]]
## Alternatively, you may also use the packages 'itertools': http://docs.python.org/2/library/itertools.html#itertools.product

for (i, j) in possible_states:
    genes = [i, j]
    updated_genes = [False, False] # initially set to some value, which will then get updated in the loop below

    # Below, I have created my own simple boolean expressions. You have to use the expressions given in the paper
    # Gene0 = NOT Gene1
    # Gene1 = Gene0 AND Gene1
    # state of the network = [value of Gene 0, value of Gene 1]

    updated_genes[0] = not genes[1] # Boolean expression for updating gene 1
    updated_genes[1] = genes[0] and genes[1] # Boolean expression for updating gene 2

    # Now, you know that the current state contained in 'genes' goes to the state contained in 'updated_genes' in one update step. 
    # You want to record this transition from the current state to the next state in your STG graph.
    # Both the current state and the next state are nodes in your graph. A node in a graph needs a 'label' to be referred to.
    # So, calculate the decimal number equivalent of the current and next state that are in boolean form and use them as the node labels.
    # For your 10-gene network, compute the decimal value of each state using a loop. DO NOT hard code.

    current_state = int(genes[0]) * pow(2,0) + int(genes[1]) * pow(2,1)  # label for the node representing the current state
    next_state = int(updated_genes[0]) * pow(2,0) + int(updated_genes[1]) * pow(2,1) # label for the next state node
    G.add_node(current_state)
    G.add_node(next_state)
    G.add_edge(current_state, next_state)

print (G.nodes(), G.edges())

plt.figure(figsize=(8,8))
nx.draw_circular(G,node_size=500,with_labels=True) # creates a circular graph
plt.axis('equal')
plt.show()

# you may use other graph plotting options shown in DrawGraph.py