import numpy as np
import math


def notOptimized(table):
    z = table[-1]
    return any(x > 0 for x in z[:-1])


def isBasic(column):
    return sum(column) == 1 and len([c for c in column if c == 0]) == len(column) - 1


def toTable(A, b, c):
    xb = [eq + [x] for eq, x in zip(A, b)]
    z = c + [0]
    return xb + [z]


def vectorOf(table):
    vecs = np.array(table).T
    res = []
    for column in vecs:
        x = 0
        if isBasic(column):
            one_index = column.tolist().index(1)
            x = vecs[-1][one_index]
        res.append(round(x, 3))

    return res


def pivotPosition(table):
    z = table[-1]
    column = next(i for i, x in enumerate(z[:-1]) if x > 0)

    constraints = []
    for eq in table[:-1]:
        el = eq[column]
        constraints.append(math.inf if el <= 0 else eq[-1] / el)

    row = constraints.index(min(constraints))
    return row, column


def pivotStep(table, pivotPosition):
    newTable = [[] for row in range(len(table))]

    i, j = pivotPosition
    pivot_value = table[i][j]
    newTable[i] = np.array(table[i]) / pivot_value

    for eqi, eq in enumerate(table):
        if eqi != i:
            multiplier = np.array(newTable[i]) * table[eqi][j]
            newTable[eqi] = np.array(table[eqi]) - multiplier

    return newTable


def smethod(A, b, c):
    table = toTable(A, b, c)
    while notOptimized(table):
        pivot_position = pivotPosition(table)
        table = pivotStep(table, pivot_position)

    vec = vectorOf(table)[:-1]
    F = 0
    for i in range(len(c)):
        F += c[i] * vec[i]

    print("Optimal plan vector:", vec)
    print("Max F:", F)


# c = np.array([1, 2, 0, 0, 0])
# A = np.array([
#     [5, -2, 1, 0, 0],
#     [-1, 2, 0, 1, 0],
#     [1, 1, 0, 0, -1]
# ])
# b = np.array([4, 4, 4])

c = [1, 2, 0, 0, 0]
A = [
    [5, -2, 1, 0, 0],
    [-1, 2, 0, 1, 0],
    [1, 1, 0, 0, -1]
]
b = [4, 4, 4]

smethod(A, b, c)