#! /usr/bin/env python

#
# this python script is a pyomo-centric translation of the AMPL
# script found at: http://www.ampl.com/NEW/LOOP2/stoch2.run
#

# Python imports
import sys
from pyutilib.misc import import_file
from pyomo.opt.base import SolverFactory
from pyomo.opt.parallel import SolverManagerFactory
from pyomo.opt.parallel.manager import solve_all_instances
from pyomo.core import *
import pyomo.environ

# initialize the master instance.
mstr_mdl = import_file("master.py").model
mstr_inst = mstr_mdl.create_instance("master.dat")

# initialize the sub-problem instances.
sb_mdl = import_file("subproblem.py").model
sub_insts = [] # a Python list
sub_insts.append(
    sb_mdl.create_instance(name="Base Sub-Problem", \
                           filename="base_subproblem.dat"))
sub_insts.append(
    sb_mdl.create_instance(name="Low Sub-Problem", \
                           filename="low_subproblem.dat"))
sub_insts.append(
    sb_mdl.create_instance(name="High Sub-Problem", \
                           filename="high_subproblem.dat"))

# initialize the solver / solver manager.
opt = SolverFactory('cplex')
if opt is None or not opt.available():
    print("A CPLEX solver is not available on this machine.")
    sys.exit(1)
solver_manager = SolverManagerFactory("serial")

# miscellaneous initialization.
mstr_inst.Min_Stage2_Profit = float("Inf")

GAP = float("Inf")

max_iterations = 50

# the main benders loop.
for i in sequence(max_iterations):

   print("\nIteration= "+str(i))

   # solve the subproblems.
   solve_all_instances(solver_manager, 'cplex', sub_insts)
   for instance in sub_insts:
      print("Profit for scenario="+instance.name+" is "+
            str(round(instance.Exp_Stage2_Profit(), 4)))
   print("")

   # if not converged, add store the pricing information from the
   # sub-problem solutions in the master.
   mstr_inst.CUTS.add(i)
   for s in sequence(len(sub_insts)):
      inst = sub_insts[s-1]
      dual_suffix = inst.dual
      urc_suffix = inst.urc
      for t in mstr_inst.TWOPLUSWEEKS:
         mstr_inst.time_price[t, s, i] = \
            dual_suffix[inst.Time[t]]
      for p in mstr_inst.PROD:
         mstr_inst.bal2_price[p, s, i] = \
            dual_suffix[inst.Balance2[p]]
      for p in mstr_inst.PROD:
         for t in mstr_inst.TWOPLUSWEEKS:
            mstr_inst.sell_lim_price[p, t, s, i] = \
                urc_suffix[inst.Sell[p, t]]

   # add the master cut.
   cut = sum([mstr_inst.time_price[t, s, i] * mstr_inst.avail[t]
              for t in mstr_inst.TWOPLUSWEEKS
              for s in mstr_inst.SCEN]) + \
         sum([mstr_inst.bal2_price[p, s, i] * (-mstr_inst.Inv1[p])
              for p in mstr_inst.PROD
              for s in mstr_inst.SCEN]) + \
         sum([mstr_inst.sell_lim_price[p, t, s, i] * mstr_inst.market[p, t]
              for p in mstr_inst.PROD
              for t in mstr_inst.TWOPLUSWEEKS
              for s in mstr_inst.SCEN]) - \
         mstr_inst.Min_Stage2_Profit
   mstr_inst.Cut_Defn.add((0.0, cut, None))

   # compute expected second-stage profit
   Exp_Stage2_Profit = sum([sub_insts[s-1].Exp_Stage2_Profit()
                            for s in sequence(mstr_inst.NUMSCEN())])
   print("Expected Stage2 Profit= "+str(round(Exp_Stage2_Profit, 4)))
   print("")

   newGAP = round(mstr_inst.Min_Stage2_Profit.value - \
                  Exp_Stage2_Profit, 6)
   if newGAP == 0:
       # get rid -0.0, which makes this script easier
       # to test against a baseline
       newGAP = 0
   print("New gap= "+str(newGAP)+"\n")

   if newGAP > 0.00001:
      GAP = min(GAP, newGAP)
   else:
      break

   # re-solve the master and update the subproblem inv1 values.
   solve_all_instances(solver_manager, 'cplex', [mstr_inst])

   print("Master expected profit="+str(mstr_inst.Expected_Profit()))

   for instance in sub_insts:
      for p in mstr_inst.PROD:
         # the master inventory values might be slightly
         # less than 0 (within tolerance); threshold here.
         instance.inv1[p] = max(mstr_inst.Inv1[p](),0.0)

print("Benders converged!")

print("\nConverged master solution values:")
for p in sorted(mstr_inst.PROD):
   print("Make1["+p+"]="+str(round(mstr_inst.Make1[p](), 4)))
   print("Sell1["+p+"]="+str(round(mstr_inst.Sell1[p](), 4)))
   print("Inv1["+p+"]="+str(round(mstr_inst.Inv1[p](), 4)))
