"""
The function simplex takes initial feasible dictionary of a LP problem with
m basic variables and n non-basic variables, represented by (m+1)x(n+1) matrix
wrapped in numpy.array, as its only input.

The format of the input matrix input is as follows:
In each row, the first column is reserved for additive constant
while the rest of the columns stores coefficients of non-basic variable.
The first row describes the objective function while the rest of the rows
list constraints. You can check some examples near the end of the file.
"""

import numpy as np

def simplex(tableau):
    m = tableau.shape[0] - 1
    n = tableau.shape[1] - 1

    nonbasic_vars = [i for i in range(1, n + 1)]
    basic_vars = [n + i for i in range(1, m + 1)]

    while True:
        print(f'NONBASIC VARS: {nonbasic_vars}')
        print(f'BASIC VARS: {basic_vars}')
        print(tableau)

        entering_idx = -1
        leaving_idx = -1
        largest_coeff = -348573248583457
        min_inc = 8943738945723545

        # select entering and leaving variables

        for j in range(n, 0, -1):
            if tableau[0][j] > 0 and tableau[0][j] >= largest_coeff:
                entering_idx = j
                largest_coeff = tableau[0][j]

        if entering_idx == -1:
            print(f'OPTIMAL SOLUTION FOUND!! z={tableau[0][0]}\n')
            break

        for i in range(m, 0, -1):
            if tableau[i][entering_idx] < 0:
                x_inc = tableau[i][0] / -tableau[i][entering_idx]
                if x_inc <= min_inc:
                    leaving_idx = i
                    min_inc = x_inc

        if leaving_idx == -1:
            print('UNBOUNDED LP PROBLEM!!\n')
            break

        print(f'ENTERING: x_{nonbasic_vars[entering_idx - 1]}')
        print(f'LEAVING: x_{basic_vars[leaving_idx - 1]}')

        tmp = basic_vars[leaving_idx - 1]
        basic_vars[leaving_idx - 1] = nonbasic_vars[entering_idx - 1]
        nonbasic_vars[entering_idx - 1] = tmp

        # update dictionary

        c = tableau[leaving_idx][entering_idx]
        tableau[leaving_idx][entering_idx] = -1
        tableau[leaving_idx][:] *= -1 / c

        for i in range(0, m + 1):
            if i == leaving_idx:
                continue
            c = tableau[i][entering_idx]
            tableau[i][entering_idx] = 0
            tableau[i][:] += c * tableau[leaving_idx][:]

        # perform bubble sort if needed

        run = n
        swapped = True

        while swapped:
            swapped = False
            for j in range(1, run):
                if nonbasic_vars[j - 1] > nonbasic_vars[j]:
                    tmp = nonbasic_vars[j - 1]
                    nonbasic_vars[j - 1] = nonbasic_vars[j]
                    nonbasic_vars[j] = tmp
                    tableau[:, [j, j + 1]] = tableau[:, [j + 1, j]]
                    swapped = True
            run -= 1

        run = m
        swapped = True

        while swapped:
            swapped = False
            for i in range(1, run):
                if basic_vars[i - 1] > basic_vars[i]:
                    tmp = basic_vars[i - 1]
                    basic_vars[i - 1] = basic_vars[i]
                    basic_vars[i] = tmp
                    tableau[[i, i + 1]] = tableau[[i + 1, i]]
                    swapped = True
            run -= 1

        print()

# hw2 5a
test1 = np.array([[0, 2, 3, 3],
                  [60, -3, -1, 0],
                  [10, 1, -1, -4],
                  [15, -2, 2, -5]]).astype(float)

# hw2 5b
test2 = np.array([[0, 3, 2, 4],
                  [4, -1, -1, -2],
                  [5, -2, 0, -3],
                  [7, -2, -1, -3]]).astype(float)

# lec6 example
test3 = np.array([[0, 4, 3],
                  [40, -2, -1],
                  [30, -1, -1],
                  [15, -1, 0]]).astype(float)

# lec7 unbounded lp prob. example
test4 = np.array([[10, 1, -1, -2],
                  [1, 0, -1, -1],
                  [1, 2, 0, -1],
                  [1, 1, 1, -1]]).astype(float)

# lec7 another unbounded lp prob. example
test5 = np.array([[0, 0, 2, 1],
                  [5, -1, 1, -1],
                  [3, 2, -1, 0],
                  [5, 0, -1, 2]]).astype(float)

simplex(test1)
simplex(test2)
simplex(test3)
simplex(test4)
simplex(test5)