'''
Indexing of capacity matrix, C, (row = from, column = to),  * = capacity figure
      1   2   3  ....  n
s(=0) *   *   *        *  
1     *   *   *        *
2     *   *   *        *
3     *   *   *        *
:     : 
n-1   *   *   *        *

That is, Pytohn element C[i][jOffset+j] = capacity from i to j, jOffset=-1
x is corresponding matrix for the actual transportation, to be optimized

Decision vector is then:

x_01,x_02,...,x_0n, x_11,x_12,...,x_1n,...,x_(n-1)1,x_(n-2),...,x_(n-1)n

where the indexing, like _01, corrspons to "from" and "to", i.e., not Python indexing

Rows in the constraint matrix, A_i, follow the same structure
'''

import lpLib as LP
import numpy as np

C = [[4,0,3,0,0],\
    [0,2,1,0,0],\
    [0,0,0,0,4],\
    [1,0,0,4,0],\
    [0,2,0,0,2]]
jOffset=-1
s = 0
n = len(C[0]) # Number of the destination node
def findSolution(): # Main routine to maximize flow in the network
    objective = []
    for j in range(1,n+1):
        objective.append(1) # First row in x-matrix, i.e., sum transportation out of the sink s
    for i in range(s+1,n):  # Remaining rows, profit coefficients = 0
        for j in range(1,n+1):
            objective.append(0)
    LP.coefficients(objective, False)  # Add objective function to the LP model
    n_x = len(objective)
    #
    # Since the objective function is a vector, i.e., adding row by row, we need a convenient index function
    #
    def ndx(i,j): # From state s <= i <= n-1, to state 1 <= j <= n
        return jOffset + j + i*n # i,j = "logigcal indexes" ==> indexing to use in A-row
    #
    # Specify slack constraints, x_ij  <= C_ij
    #
    k = 0 # Keep track of corresponding variable in the decision vector
    for i in range(s,n):
        for j in range(1,n+1):
            A_i = np.zeros([n_x]) # Create a new row in the A-matrix
            A_i[k] = 1 # coefficient in front of x_ij
            LP.upperBound(C[i][jOffset + j],A_i)
            k += 1   
    #
    # Specify "mass conservation", i.e., Sum out of s = Sum in to n
    #
    A_i = np.zeros([n_x])
    for j in range(1,n+1): # out of node s to node j
        A_i[ndx(s,j)] = 1
    for i in range(s,n): # from node i to node n 
        A_i[ndx(i,n)] = -1
    LP.equalityConstraint(0,A_i)
    #
    # Specify "mass conservation", i.e., sum out of node k = sum in for node k, all k except s and n
    #
    for k in range(s+1,n):
        A_i = np.zeros([n_x])
        for j in range(1,n+1):
           A_i[ndx(k,j)] = 1 
        for i in range(s,n):
           A_i[ndx(i,k)] = -1 
        LP.equalityConstraint(0,A_i)
    res = LP.lpSolve( True)
    k = 0
    #
    # Format result in a nice manner....
    #
    out = "x_ij|"
    lne = "----+-"
    for j in range(1,n+1):
        out += " " + str(j) + " "
        lne += "---"
    print(out)
    print(lne)
    for i in range(s,n):
        out = " " + str(i) + "  |"
        for j in range(1,n+1):
            out += " " + str(int(res.x[k]+0.00001)) + " "
            k += 1
        print(out+"|")        
    print("Max flow", res.fun)
    return res.fun
#
print("Result for no queue, one direction only")
Z_OK = findSolution()
#
print("\nResult queue, one direction only")
C[3][jOffset+4] = 1 # Queue
C[2][jOffset+n] = 1 # Queue
Z_queue1 = findSolution()
#
print("\nResult queue, two directions")
C[2][jOffset+4] = 2 # 

Z_queue2 = findSolution()
print("\nVerify no queue, both directions")
C[3][jOffset+4] = 4
C[2][jOffset+n] = 4
Z_OK2 = findSolution()
q = 0.1
print("\nAverage, one direction:", (1-q)*Z_OK + q*Z_queue1)
print("Average, both directions:", (1-q)*Z_OK2 + q*Z_queue2)
print('Average "profit", both direction',(1-q)*Z_OK2 + q*Z_queue2-0.15)
print('Measure pays off:',(1-q)*Z_OK2 + q*Z_queue2-0.15>(1-q)*Z_OK + q*Z_queue1)