from sympy import *
init_printing()

A = Matrix([[6,  3,  -2,  4,  1,  0],
            [18,  2,  1,  3,  0,  1]])
c = Matrix([[0, 4,  2,  -1,  0,  0]])

def is_non_negative_col0(m):
    return sum(1 for val in m[:,0] if val < 0) == 0

# поиск единичного базиса
def find_basis(m):
    rows, cols = m.shape
    basis = [None]*rows
    for j in range(cols):
        l_j = list(m[:,j])
        if l_j.count(1) == 1 and l_j.count(0) == rows - 1:
            i = l_j.index(1) # индекс строки с единицей
            basis[i] = j
    return basis

def simplex(m,c):
    if not is_non_negative_col0(m):
        m = m * (-1)

    all_names = 'A0 A1 A2 A3 A4 A5'.split()


    basis = find_basis(m)
    bases_names = Matrix([all_names[i] for i in basis])
    basis_coeffs = Matrix([c[i] for i in basis])
    m01 = bases_names.row_join(basis_coeffs)

    rows, cols = m.shape
    deltas = []
    for j in range(cols):
        x = c[j] - m[:,j].dot(basis_coeffs)
        deltas.append(x)

    # строка заголовок
    m_title = Matrix([["Базис", "C_баз"]+all_names])
    # основные строки
    m_rows = m01.row_join(m)
    # строка оценок
    m_delta = Matrix([["z","z"]+deltas])
    # итоговая таблица
    m_all = m_title.col_join(m_rows).col_join(m_delta)

    rows, cols = A.shape
    while any(x > 0 for x in deltas[1:]): # or deltas[1:].count(0) > rows

        for j in range(1, cols):
            if deltas[j] > 0 and all(x <= 0 for x in A[:,j]):
                print("Целевая функция не ограничена!")
                break
        # ведущий элемент
        max = deltas[1]
        for j in deltas:
            if max <j:
                max = j
            j+=1
        j1 = deltas.index(max)

        print(A[0,0],A[0,j1])
        z = A[0,0]/A[0,j1]
        print(type(z))
        print(z)
        k=1

        i1=0
        while k<rows:
            m = A[k,0]/A[k,j]
            if not comp(z, m):
                z = A[k,0] / A[k,j]
                i1 = k
            k+=1
        j_cur = j1
        i_cur = i1
        A_pivot = A[i_cur,j_cur]
        for i in range(rows):
            if i == i_cur:
                continue
            A.row_op(i, lambda var, j : (var*A_pivot-A[i_cur,j]*A[i,j_cur])/A_pivot)
        A.row_op(i_cur, lambda var, j : var/A_pivot)

        basis = find_basis(A)
        bases_names = Matrix([all_names[i] for i in basis])
        basis_coeffs = Matrix([c[i] for i in basis])
        A01 = bases_names.row_join(basis_coeffs)

        deltas = []
        for j in range(A.shape[1]):
            x = c[0,j] - A[:,j].dot(basis_coeffs)
            deltas.append(x)

        A_title = Matrix([["Базис", "C_баз"]+all_names])
        A_rows = A01.row_join(A)
        A_delta = Matrix([["z","z"]+deltas])
        A_all = A_title.col_join(A_rows).col_join(A_delta)
        if deltas[1:].count(0) > rows:
            if input("Продолжить решение? Y/N ") == "N":
                break

simplex(A,c)