import gurobipy as gp
from gurobipy import GRB

# Setting up model and variables
# Introducing auxilliary 2 variables 'Deficit_stock' and 'Excess_stock' to replance non-linear Max() logic in objective function
model = gp.Model('Network Balancing - Pilot')
X = model.addVars(s_nodes, r_nodes, parts, name='x') # x_srp variable
Deficit_stock = model.addVars(nodes, parts, name='d_stock') # deficit stock
Excess_stock = model.addVars(nodes, parts, name='e_stock') # excess stock
IE = model.addVars(nodes, parts, name='IE') #IE_srp variable (ending node)


# Setting up constraints
x_nonneg_const = model.addConstrs((X[(s, r, p)] >= 0 for s in s_nodes for r in r_nodes for p in parts), 'X_Non_Negativity')
ie_nonneg_const = model.addConstrs((IE[(n, p)] >= 0 for n in nodes for p in parts), 'IE_Non_Negativity')
for n in nodes:
    for p in parts:
        model.addConstr((IE[(n, p)] == IS_dict[(n, p)] + gp.quicksum(X[(s, n, p)] for s in s_nodes) - gp.quicksum(X[(n, r, p)] for r in r_nodes)), "balance")
deficit_stock_const = model.addConstrs((Deficit_stock[(n, p)] >= 0 for n in nodes for p in parts), 'deficit_stock_const #1')
deficit_stock_const = model.addConstrs((Deficit_stock[(n, p)] >= DS_dict[(n,p)] - IE[(n,p)] for n in nodes for p in parts), 'deficit_stock_const #2')
excess_stock_const = model.addConstrs((Excess_stock[(n, p)] >= 0 for n in nodes for p in parts), 'excess_stock_const #1')
excess_stock_const = model.addConstrs((Excess_stock[(n, p)] >= IE[(n,p)] - ES_dict[(n,p)] for n in nodes for p in parts), 'excess_stock_const #2')

# Setting up objective function
# This is the point where memory error occurs
obj = gp.quicksum(k*routes_dict[(s, r)] * gp.quicksum(X[(s,r,p)] for p in parts) for s in s_nodes for r in r_nodes) \
+ gp.quicksum(unit_value_dict[p] * gp.quicksum(Excess_stock[(n, p)] for n in nodes) for p in parts) \
+ gp.quicksum(unit_value_dict[p] * gp.quicksum(Deficit_stock[(n, p)] for n in nodes) for p in parts)
model.setObjective(obj, GRB.MINIMIZE)

model.optimize()