"""
Description: A Python script to read .mps files which contain MILP problems and convert them to numpy format.
"""

import pulp # library used to read .mps files
from pulp import *
import numpy as np
import time
from pulp.constants import LpConstraintEQ, LpConstraintLE, LpConstraintGE


def read_mps(filename="./B&B/22433.mps"):
    # read .mps file
    info = pulp.LpProblem.fromMPS(filename) # read .mps file
    lp_problem = info[1]

    # Obtain the coefficients of the objective, i.e. c'x
    c = np.array([lp_problem.objective.get(var, 0) for var in lp_problem.variables()])

    # Hold constraints
    A_lt = [] #Ax <= b
    b_lt = []

    A_eq = [] # Ax = b
    b_eq = []


    for constraint in lp_problem.constraints.values():
        row = [constraint.get(var, 0) for var in lp_problem.variables()]  # Coefficients of the constraint
        
        if constraint.sense == LpConstraintEQ: # a'x = b constraint
            A_eq.append(row)
            b_eq.append(-constraint.constant)
            # b_eq.append(constraint.constant)
            # print(f"{constraint.constant}")
        elif constraint.sense == LpConstraintLE: # a'x <= constraint
            A_lt.append(row)
            b_lt.append(-constraint.constant)
            # b_lt.append(constraint.constant)
        elif constraint.sense == LpConstraintGE:
            negative_row = [-element for element in row] # a'x >= constraint
            A_lt.append(negative_row)
            b_lt.append(constraint.constant)
            # b_lt.append(-constraint.constant)

    # convert to numpy
    A_lt = np.array(A_lt) if len(A_lt) != 0 else None
    b_lt = np.array(b_lt) if len(b_lt) != 0 else None
    A_eq = np.array(A_eq) if len(A_eq) != 0 else None
    b_eq = np.array(b_eq) if len(b_eq) != 0 else None

    # Extract the integer and binary constraints
    integer_variables = []
    continuous_variables = []
    bounds = []

    for idx, var in enumerate(lp_problem.variables()):
        bounds.append((var.lowBound, var.upBound))

        if var.cat == 'Integer' or var.cat == 'Binary': # integer constraint
            integer_variables.append(idx)
        else:
            continuous_variables.append(idx)

    integer_variables = np.array(integer_variables)
    continuous_variables = np.array(continuous_variables)

    return c, A_lt, b_lt, A_eq, b_eq, bounds, integer_variables, continuous_variables