"""
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
from branch_and_bound import BnB

write_flag = False # change to True if you want to write .json and .lp

# read .mps file
info = pulp.LpProblem.fromMPS("./B&B/22433.mps") # read .mps file
# The variable 'info' is a tuple of size 2. 
# - The first entry, info[0] holds a dict of variable names
# - The second entry, info[1] holds a pulp.LpProblem

lp_problem = info[1]

# Returns a list of pulp.LpVariables 
lp_variables = lp_problem.variables()

# Returns the objective of the LP as a pulp.LpAffineExpression  
lp_objective = lp_problem.objective

# Maximize or Minimization problem
type_lookup = {1: 'minimize', -1: 'maximize'}
print(f"Type: {type_lookup[lp_problem.sense]}")

# Print the name of the LP 
print(f"Problem name: {lp_problem.name}")

if write_flag:
    # Convert to .json format (messy and hard to read)
    lp_problem.toJson("./B&B/example_info.json")

    # Write the LP in .lp format (This is easy for humans to read)
    lp_problem.writeLP("./B&B/example_info.lp")

# If you want you can enumerate all the constraints (I wouln't suggest doing this for large LPs)
# for name, constraint in problem.constraints.items():
#     print(f"{name}: {constraint}")

# Convert the pulp.LpProblem to numpy format
print("Converting .mps to numpy format...")
start = time.time()

# 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)
b_lt = np.array(b_lt)
A_eq = np.array(A_eq)
b_eq = np.array(b_eq)

# Extract the integer and binary constraints
binary_variables = []
integer_variables = []
continuous_variables = []
bounds = []
integrality_args = [0]*len(c) # Passes to the linprog if the variable is interger (1) or continuous (0)

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

    if var.cat == 'Integer': # integer constraint
        # bounds.append((0.0, None))
        integer_variables.append(idx)
        # integrality_args[idx] = 1
    elif var.cat == 'Binary': # binary constraint
        # bounds.append((0.0, None))
        binary_variables.append(idx)
    else:
        # bounds.append((var.lowBound, var.upBound))
        # bounds.append((0.0, None))
        continuous_variables.append(idx) # continuous constraint

binary_variables = np.array(binary_variables)
integer_variables = np.array(integer_variables)
continuous_variables = np.array(continuous_variables)
integrality_args = integrality_args if sum(integrality_args) != 0 else 0

end = time.time()
print(f"Finished conversion, took {round(end-start, 5)} sec")

from scipy.optimize import linprog
print("Now starting to solve relaxed MILP...")

if A_lt.shape[0] != 0:
    start = time.time()
    # sol_init = linprog(c=c, A_ub=A_lt, b_ub=b_lt, A_eq=A_eq, b_eq=b_eq, integrality=integrality_args, bounds=bounds)
    sol_init, _ = BnB(c=c, A_ub=A_lt, b_ub=b_lt, A_eq=A_eq, b_eq=b_eq, bounds=bounds, int_indices=integer_variables)
    end = time.time()
else:
    start = time.time()
    # sol_init = linprog(c=c, A_eq=A_eq, b_eq=b_eq, integrality=integrality_args, bounds=bounds)
    # sol_init = linprog(c=c, A_eq=A_eq, b_eq=b_eq, bounds=bounds)
    sol_init, _ = BnB(c=c, A_eq=A_eq, b_eq=b_eq, bounds=bounds, int_indices=integer_variables)
    end = time.time()

print(f"Finished solving, took {round(end-start, 5)} sec")
print(f"solution: f(x*) = {sol_init.obj_val}")
# print(f"Pass: {sol_init.success}")
# print(f"Solution: {sol_init.x}")
# if not sol_init.success:
#     print(f"Solver message: {sol_init.message}")