from z3 import *

def printMatrix(m, o):
    t1 = list(o[0])
    t2 = list(o[1])

    s = "%3s |" % ""
    for v in t1:
        s += "%5s |" % v
    s += "\n"
    i = 0
    for row in m:
        s += "%3s |" % t2[i]
        for v in row:
            s += "%5.2f |" % v
        s += "\n"
        i += 1
    print s

"""
    checks if given variables with given assignments violate given constraints

    - parameter v: list of variable names
    - parameter a: dictionary of variable names - assigned values
    - parameter c: dictionary of variable names and associated values representing upper bound

    - returns: list of variables violating constraints
"""
def getViolatingVars(v, a, c):
    t = []
    for var in v:
        if var in c:
            if a[var] > c[var]:
                t.append(var)
    return t

"""
    returns a suitable var for pivoting

    - parameter m: matrix (list of lists) in simplified form for simplex algorithm
    - parameter o: variable ordering on matrix (0 is top, 1 is left)
    - parameter p: pivot candidate to search suitable variable for
    - parameter a: variable assignments
    - parameter c: constraints

"""
def getSuitableVar(m, o, p, a, c):
    # possible pivot candidates
    possibleCandidates = o[0]

    # first, check if possibleCandidates are constraint free
    # if so, return one
    for pc in possibleCandidates:
        if not pc in c:
            return pc

    # get row of matrix of pivot candidate
    r = m[o[1].index(p)]

    # if there is need to decrease a var, go for that var
    i = 0
    for v in r:
        if v > 0:
            return possibleCandidates[i]
        i += 1

    return None

"""
    performs pivot operation on given matrix with given variables

    - parameter m: a matrix (list of lists) in simplified form for simplex algorithm
    - parameter o: variable ordering on matrix (0 is top, 1 is left)
    - parameter p1: variable 1
    - parameter p2: variable 2
"""
def performPivotOperation(m, o, p1, p2):
    ip1 = o[1].index(p1)
    ip2 = o[0].index(p2)
    r = m[ip1]

    # swap the values
    r[1-ip2] = -1 * r[1-ip2] / r[ip2]
    r[ip2] = 1 / r[ip2]

    i = 0
    for row in m:
        if i != ip1:
            row[1-ip2] = row[1-ip2] + row[ip2] * r[ip2]
            row[ip2] = row[ip2] * r[ip2]
        i += 1

    # swap in ordering
    o[0][ip2] = p1
    o[1][ip1] = p2


"""
    updates assignment according to matrix

    - parameter m: a matrix (list of lists) in simplified form for simplex algorithm
    - parameter o: variable ordering on matrix (0 is top, 1 is left)
    - parameter c: constraints dictionary
    - parameter v: variable to update
    - parameter a: variable assignments
"""
def updateAssignment(m, o, c, v, a):
    diff = abs(a[v] - c[v])

    # assign new value
    a[v] = a[v] - diff

    i = 0
    for row in m:
        a[o[1][i]] = a[o[0][0]] * row[0] + a[o[0][1]] * row[1]
        i += 1


"""
    performs simplex algorithm

    - parameter m: a matrix (list of lists) in simplified form for simplex algorithm
    - parameter o: variable ordering on matrix (0 is top, 1 is left)
    - parameter c: constraints dictionary of variable names and associated values representing upper bound

    - parameter bv: set of basic vars (for internal use only)
    - parameter a: variable assignments (for internal use only)
    - parameter i: iteration count (for internal use only)

    - returns: sat or unsat
"""
def simplex(m, o, c, bv = None, a = None, i = 1):

    print "---- iteration " + str(i) + " ----"

    if bv == None:
        bv = o[1]

    if a == None:
        a = {var:0.0 for var in bv}
        for var in o[0]:
            a[var] = 0

    if i == 1:
        print "basic vars: " + str(bv)
        print "non-basic vars: " + str(o[0])
        print "assignments: " + str(a)

    violatingVars = getViolatingVars(bv, a, c)

    if violatingVars == []:
        return sat

    varToChoose = None
    for var in [v for t in o for v in t]:
        if var in violatingVars:
            varToChoose = var
            break

    print "variables violating bounds: " + str(violatingVars)
    print "var to choose: " + str(varToChoose)

    suitableVar = getSuitableVar(m, o, varToChoose, a, c)
    if suitableVar == None:
        return unsat

    print "suitable var: " + suitableVar
    print "pivoting " + varToChoose + " with " + suitableVar

    performPivotOperation(m, o, varToChoose, suitableVar)

    updateAssignment(m, o, c, varToChoose, a)

    printMatrix(m, o)
    print "updated assignments: " + str(a)

    return simplex(m, o, c, bv, a, i+1)



matrix = [[-1.0, 1.0],
        [0.0, 1.0],
        [-1.0,-1.0],
        [3.0,-1.0]]

constraints = {
    "s1":1.0,
    "s2":4.0,
    "s3":-6.0,
    "s4":7.0,
}
"""
matrix = [[0.0, 1.0],
        [-3.0, 1.0],
        [-1.0,-1.0],
        [2.0,-1.0]]

constraints = {
    "s1":4.0,
    "s2":-1.0,
    "s3":-5.0,
    "s4":3.0,
}
"""

order = [["x","y"],["s1","s2","s3","s4"]]

print "---- simplex algorithm ----"
print "matrix: "
printMatrix(matrix, order)
print "constraints: " + str(constraints)

print simplex(matrix, order, constraints)