#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright (c) 2008-2022
#  National Technology and Engineering Solutions of Sandia, LLC
#  Under the terms of Contract DE-NA0003525 with National Technology and
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
#  rights in this software.
#  This software is distributed under the 3-clause BSD License.
#  ___________________________________________________________________________
#
# Unit Tests for Utility Functions
#

import pickle
import os
from os.path import abspath, dirname, join
currdir = dirname(abspath(__file__))+os.sep
import platform

from filecmp import cmp
import pyomo.common.unittest as unittest
from pyomo.environ import AbstractModel, ConcreteModel, Set, Param, Var, Constraint, Objective, Reals, NonNegativeReals, sum_product


using_pypy = platform.python_implementation() == "PyPy"


def obj_rule(model):
    return sum(model.x[a] + model.y[a] for a in model.A)
def constr_rule(model,a):
    return model.x[a] >= model.y[a]
def simple_con_rule(model, i):
    return model.x <= i

class Test(unittest.TestCase):

    def verifyModel(self, ref, new):
        # Verify the block indices
        self.assertEqual(sorted(ref._data.keys()), sorted(new._data.keys()))
        for idx in ref._data.keys():
            self.assertEqual(type(ref._data[idx]),  type(new._data[idx]))
            if idx is not None:
                self.assertNotEqual(id(ref._data[idx]),  id(new._data[idx]))
        self.assertEqual( id(ref.solutions._instance()), id(ref) )
        self.assertEqual( id(new.solutions._instance()), id(new) )

        # Verify the block attributes
        for idx in ref._data.keys():
            # It is important to lookup using ref._data[idx] and not
            # ref[idx], as the model may not be constructed, and
            # accessing __getitem__ on an abstract model generates an
            # exception.
            ref_c = ref._data[idx].component_map()
            new_c = new._data[idx].component_map()
            self.assertEqual( sorted(ref_c.keys()), sorted(new_c.keys()) )
            for a in ref_c.keys():
                self.assertEqual(type(ref_c[a]),  type(new_c[a]))
                self.assertNotEqual(id(ref_c[a]),  id(new_c[a]))

    def test_pickle_empty_abstract_model(self):
        model = AbstractModel()
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_set(self):
        model = AbstractModel()
        model.A = Set(initialize=[1,2,3])
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_virtual_set(self):
        model = AbstractModel()
        model._a = Set(initialize=[1,2,3])
        model.A = model._a * model._a
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_param(self):
        model = AbstractModel()
        model.A = Param(initialize=1)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_indexed_param(self):
        model = AbstractModel()
        model.A = Param([1,2,3], initialize={1:100,2:200,3:300})
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_mutable_param(self):
        model = AbstractModel()
        model.A = Param(initialize=1, mutable=True)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_mutable_indexed_param(self):
        model = AbstractModel()
        model.A = Param([1,2,3], initialize={1:100,3:300}, mutable=True)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_var(self):
        model = AbstractModel()
        model.A = Var(initialize=1)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_indexed_var(self):
        model = AbstractModel()
        model.A = Var([1,2,3], initialize={1:100,2:200,3:300})
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_constant_objective(self):
        model = AbstractModel()
        model.A = Objective(expr=1)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_objective(self):
        model = AbstractModel()
        model.x = Var()
        model.A = Objective(expr=model.x <= 0)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_constraint(self):
        model = AbstractModel()
        model.x = Var()
        model.A = Constraint(expr=model.x <= 0)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_abstract_model_indexed_constraint(self):
        model = AbstractModel()
        model.x = Var()
        model.A = Constraint([1,2,3], rule=simple_con_rule)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    #########

    def test_pickle_empty_concrete_model(self):
        model = ConcreteModel()
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_set(self):
        model = ConcreteModel()
        model.A = Set(initialize=[1,2,3])
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_virtual_set(self):
        model = ConcreteModel()
        model._a = Set(initialize=[1,2,3])
        model.A = model._a * model._a
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_param(self):
        model = ConcreteModel()
        model.A = Param(initialize=1)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_indexed_param(self):
        model = ConcreteModel()
        model.A = Param([1,2,3], initialize={1:100,2:200,3:300})
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_mutable_param(self):
        model = ConcreteModel()
        model.A = Param(initialize=1, mutable=True)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_mutable_indexed_param(self):
        model = ConcreteModel()
        model.A = Param([1,2,3], initialize={1:100,3:300}, mutable=True)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_var(self):
        model = ConcreteModel()
        model.A = Var(initialize=1)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_indexed_var(self):
        model = ConcreteModel()
        model.A = Var([1,2,3], initialize={1:100,2:200,3:300})
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_constant_objective(self):
        model = ConcreteModel()
        model.A = Objective(expr=1)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_objective(self):
        model = ConcreteModel()
        model.x = Var()
        model.A = Objective(expr=model.x**2)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_constraint(self):
        model = ConcreteModel()
        model.x = Var()
        model.A = Constraint(expr=model.x <= 0)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    def test_pickle_concrete_model_indexed_constraint(self):
        model = ConcreteModel()
        model.x = Var()
        model.A = Constraint([1,2,3], rule=simple_con_rule)
        str = pickle.dumps(model)
        tmodel = pickle.loads(str)
        self.verifyModel(model, tmodel)

    ##########


    # tests the ability to pickle an abstract model prior to construction,
    # read it back it, and create an instance from it. validation is relatively
    # weak, in that it only tests the validity of an expression constructed
    # using the resulting model.
    def test_pickle1(self):
        model = AbstractModel()
        model.A = Set(initialize=[1,2,3])
        model.B = Param(model.A,initialize={1:100,2:200,3:300}, mutable=True)
        model.x = Var(model.A)
        model.y = Var(model.A)
        model.obj = Objective(rule=obj_rule)
        model.constr = Constraint(model.A,rule=constr_rule)
        pickle_str = pickle.dumps(model)
        tmodel = pickle.loads(pickle_str)
        instance=tmodel.create_instance()
        expr = sum_product(instance.x,instance.B,instance.y)
        baseline = "B[1]*x[1]*y[1] + B[2]*x[2]*y[2] + B[3]*x[3]*y[3]"
        self.assertEqual( str(expr), baseline )

    # same as above, but pickles the constructed AbstractModel and
    # then operates on the unpickled instance.
    def test_pickle2(self):
        model = AbstractModel()
        model.A = Set(initialize=[1,2,3])
        model.B = Param(model.A,initialize={1:100,2:200,3:300}, mutable=True)
        model.x = Var(model.A)
        model.y = Var(model.A)
        model.obj = Objective(rule=obj_rule)
        model.constr = Constraint(model.A,rule=constr_rule)
        tmp=model.create_instance()
        pickle_str = pickle.dumps(tmp)
        instance = pickle.loads(pickle_str)
        expr = sum_product(instance.x,instance.B,instance.y)
        baseline = "B[1]*x[1]*y[1] + B[2]*x[2]*y[2] + B[3]*x[3]*y[3]"
        self.assertEqual( str(expr), baseline )

    # verifies that the use of lambda expressions as rules yields model instances
    # that are not pickle'able.
    def test_pickle3(self):
        def rule1(model):
            return (1,model.x+model.y[1],2)
        def rule2(model, i):
            return (1,model.x+model.y[1]+i,2)

        model = AbstractModel()
        model.a = Set(initialize=[1,2,3])
        model.A = Param(initialize=1, mutable=True)
        model.B = Param(model.a, mutable=True)
        model.x = Var(initialize=1,within=Reals)
        model.y = Var(model.a, initialize=1,within=Reals)
        model.obj = Objective(rule=lambda model: model.x+model.y[1])
        model.obj2 = Objective(model.a,rule=lambda model,i: i+model.x+model.y[1])
        model.con = Constraint(rule=rule1)
        model.con2 = Constraint(model.a, rule=rule2)
        instance = model.create_instance()
        if using_pypy:
            str_ = pickle.dumps(instance)
            tmp_ = pickle.loads(str_)
        else:
            with self.assertRaises((pickle.PicklingError,
                                    TypeError,
                                    AttributeError)):
                pickle.dumps(instance)

    # verifies that we can print a constructed model and
    # obtain identical results before and after
    # pickling. introduced due to a test case by Gabe that
    # illustrated __getstate__ of various modeling
    # components was incorrectly and unexpectedly modifying
    # object state.
    def test_pickle4(self):

        model = ConcreteModel()
        model.s = Set(initialize=[1,2])
        model.x = Var(within=NonNegativeReals)
        model.x_indexed = Var(model.s, within=NonNegativeReals)
        model.obj = Objective(expr=model.x + model.x_indexed[1] + model.x_indexed[2])
        model.con = Constraint(expr=model.x >= 1)
        model.con2 = Constraint(expr=model.x_indexed[1] + model.x_indexed[2] >= 4)

        OUTPUT=open(join(currdir, "test_pickle4_baseline.out"), "w")
        model.pprint(ostream=OUTPUT)
        OUTPUT.close()
        _out, _txt = join(currdir, "test_pickle4_baseline.out"), join(currdir, "test_pickle4_baseline.txt")
        self.assertTrue(cmp(_out, _txt),
                        msg="Files %s and %s differ" % (_out, _txt))

        str = pickle.dumps(model)

        OUTPUT=open(join(currdir, "test_pickle4_after.out"), "w")
        model.pprint(ostream=OUTPUT)
        OUTPUT.close()
        _out, _txt = join(currdir, "test_pickle4_after.out"), join(currdir, "test_pickle4_baseline.txt")
        self.assertTrue(cmp(_out, _txt),
                        msg="Files %s and %s differ" % (_out, _txt))

if __name__ == "__main__":
    unittest.main()
