from mip import Model, INTEGER, xsum, OptimizationStatus, minimize
import numpy as np
from itertools import permutations

A = np.array([[0,1,2,3,4,1],
              [2,1,4,1,0,2],
              [1,5,2,3,1,1],
              [0,0,1,2,2,0]])
var_type = INTEGER
max_gap = 0.05
max_seconds=10
c = np.array([-1,2,-2,2,1,1])
m, n = A.shape

def createModelAndCheck(A, b, var_type, max_gap, max_seconds):
    m, n = A.shape
    model = Model()
    x = [model.add_var(var_type=var_type) for j in range(n)]
    
    for i in range(m):
        A_i = A[i]
        b_i = b[i]
        model += xsum(A_i[j]*x[j] for j in range(n)) == b_i
    
    model.objective = minimize(xsum(c[j]*x[j] for j in range(n)))
    
    model.max_gap = max_gap
    status = model.optimize(max_seconds=max_seconds)
    """
    if status != OptimizationStatus.INFEASIBLE:
        print("Found b: ", b)
    if status == OptimizationStatus.OPTIMAL:
        print('optimal solution cost {} found'.format(model.objective_value))
    elif status == OptimizationStatus.FEASIBLE:
        print('sol.cost {} found, best possible: {}'.format(model.objective_value, model.objective_bound))
    elif status == OptimizationStatus.NO_SOLUTION_FOUND:
        print('no feasible solution found, lower bound is: {}'.format(model.objective_bound))
    """
    if status == OptimizationStatus.OPTIMAL:# or status == OptimizationStatus.FEASIBLE:
        print('solution:')
        for v in model.vars:
            if abs(v.x) > 1e-6: # only printing non-zeros
                print('{} : {}'.format(v.name, v.x))
        return model
    else:
        return None

brute_force_b = list(permutations([0,1,2,3,4,5,6,7,8],m))
goodB = []
for b in brute_force_b:
    model = createModelAndCheck(A, np.array(b), var_type, max_gap, max_seconds)
    if  model is not None:
        count = 0
        good = 0
        for v in model.vars:
            if abs(v.x) > 1e-6: 
                count += 1
                if abs(v.x) > 1: good += 1
        if count > 2 and good > 1: 
            goodB +=[b]

print(goodB)
for b in goodB : 
    createModelAndCheck(A, np.array(b), var_type, max_gap, max_seconds)

