# ____________________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright (c) 2008-2026 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 os
from os.path import abspath, dirname, join

currdir = dirname(abspath(__file__)) + os.sep

from filecmp import cmp
import pyomo.common.unittest as unittest

from pyomo.environ import (
    AbstractModel,
    ConcreteModel,
    ConstraintList,
    Set,
    Param,
    Var,
    Constraint,
    Objective,
    sum_product,
    quicksum,
    sequence,
    prod,
)
from pyomo.core.expr.compare import assertExpressionsEqual


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]


class Test(unittest.TestCase):
    def test_expr0(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.C = Param(model.A, initialize={1: 100, 2: 200, 3: 300}, mutable=False)
        model.x = Var(model.A)
        model.y = Var(model.A)
        instance = model.create_instance()
        expr = sum_product(instance.B, instance.y)
        baseline = "B[1]*y[1] + B[2]*y[2] + B[3]*y[3]"
        self.assertEqual(str(expr), baseline)
        expr = sum_product(instance.C, instance.y)
        self.assertEqual(str(expr), "100*y[1] + 200*y[2] + 300*y[3]")

    def test_expr1(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.C = Param(model.A, initialize={1: 100, 2: 200, 3: 300}, mutable=False)
        model.x = Var(model.A)
        model.y = Var(model.A)
        instance = model.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)
        expr = sum_product(instance.x, instance.C, instance.y)
        self.assertEqual(str(expr), "100*x[1]*y[1] + 200*x[2]*y[2] + 300*x[3]*y[3]")

    def test_expr2(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.C = Param(model.A, initialize={1: 100, 2: 200, 3: 300}, mutable=False)
        model.x = Var(model.A)
        model.y = Var(model.A)
        instance = model.create_instance()
        expr = sum_product(instance.x, instance.B, instance.y, index=[1, 3])
        baseline = "B[1]*x[1]*y[1] + B[3]*x[3]*y[3]"
        self.assertEqual(str(expr), baseline)
        expr = sum_product(instance.x, instance.C, instance.y, index=[1, 3])
        self.assertEqual(str(expr), "100*x[1]*y[1] + 300*x[3]*y[3]")

    def test_expr3(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.C = Param(model.A, initialize={1: 100, 2: 200, 3: 300}, mutable=False)
        model.x = Var(model.A)
        model.y = Var(model.A)
        instance = model.create_instance()
        expr = sum_product(instance.x, instance.B, denom=instance.y, index=[1, 3])
        baseline = "B[1]*x[1]/y[1] + B[3]*x[3]/y[3]"
        self.assertEqual(str(expr), baseline)
        expr = sum_product(instance.x, instance.C, denom=instance.y, index=[1, 3])
        self.assertEqual(str(expr), "100*x[1]/y[1] + 300*x[3]/y[3]")

    def test_expr4(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)
        instance = model.create_instance()
        expr = sum_product(denom=[instance.y, instance.x])
        baseline = "1/(y[1]*x[1]) + 1/(y[2]*x[2]) + 1/(y[3]*x[3])"
        self.assertEqual(str(expr), baseline)

    def test_sum_product_ParamVarVar(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)
        instance = model.create_instance()
        expr = sum_product(instance.B, instance.y, instance.x)
        baseline = "B[1]*y[1]*x[1] + B[2]*y[2]*x[2] + B[3]*y[3]*x[3]"
        self.assertEqual(str(expr), baseline)

    def test_sum_product_ParamParamVar(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 = Param(model.A, mutable=True)
        instance = model.create_instance()
        expr = sum_product(instance.B, instance.y, instance.x)
        baseline = "B[1]*y[1]*x[1] + B[2]*y[2]*x[2] + B[3]*y[3]*x[3]"
        self.assertEqual(str(expr), baseline)

    def test_expr5(self):
        model = ConcreteModel()
        model.A = Set(initialize=[1, 2, 3], doc='set A')
        model.B = Param(
            model.A, initialize={1: 100, 2: 200, 3: 300}, doc='param B', mutable=True
        )
        model.C = Param(initialize=3, doc='param C', mutable=True)
        model.x = Var(model.A, doc='var x')
        model.y = Var(doc='var y')
        model.o = Objective(expr=model.y, doc='obj o')
        model.c1 = Constraint(expr=model.x[1] >= 0, doc='con c1')

        def c2_rule(model, a):
            return model.B[a] * model.x[a] <= 1

        model.c2 = Constraint(model.A, doc='con c2', rule=c2_rule)
        model.c3 = ConstraintList(doc='con c3')
        model.c3.add(model.y <= 0)
        #
        OUTPUT = open(join(currdir, "test_expr5.out"), "w")
        model.pprint(ostream=OUTPUT)
        OUTPUT.close()
        _out, _txt = join(currdir, "test_expr5.out"), join(currdir, "test_expr5.txt")
        self.assertTrue(cmp(_out, _txt), msg="Files %s and %s differ" % (_out, _txt))

    def test_prod1(self):
        self.assertEqual(prod([1, 2, 3, 5]), 30)

    def test_prod2(self):
        model = ConcreteModel()
        model.A = Set(initialize=[1, 2, 3], doc='set A')
        model.x = Var(model.A)
        expr = prod(model.x[i] for i in model.x)
        baseline = "x[1]*x[2]*x[3]"
        self.assertEqual(str(expr), baseline)
        expr = prod(model.x)
        self.assertEqual(expr, 6)

    def test_sum1(self):
        self.assertEqual(quicksum([1, 2, 3, 5]), 11)

    def test_sum2(self):
        model = ConcreteModel()
        model.A = Set(initialize=[1, 2, 3], doc='set A')
        model.x = Var(model.A)
        expr = quicksum(model.x[i] for i in model.x)
        baseline = "x[1] + x[2] + x[3]"
        assertExpressionsEqual(self, expr, model.x[1] + model.x[2] + model.x[3])

    def test_sum3(self):
        model = ConcreteModel()
        model.A = Set(initialize=[1, 2, 3], doc='set A')
        model.x = Var(model.A)
        expr = quicksum(model.x)
        assertExpressionsEqual(self, expr, 6)

    def test_summation_error1(self):
        try:
            sum_product()
            self.fail("Expected ValueError")
        except ValueError:
            pass

    def test_summation_error2(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)
        instance = model.create_instance()
        try:
            expr = sum_product(instance.x, instance.B)
            self.fail("Expected ValueError")
        except ValueError:
            pass

    def test_summation_error3(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)
        instance = model.create_instance()
        try:
            expr = sum_product(denom=(instance.x, instance.B))
            self.fail("Expected ValueError")
        except ValueError:
            pass

    def test_sequence_error1(self):
        try:
            sequence()
            self.fail("Expected ValueError")
        except ValueError:
            pass
        try:
            sequence(1, 2, 3, 4)
            self.fail("Expected ValueError")
        except ValueError:
            pass

    def test_sequence(self):
        self.assertEqual(list(sequence(10)), list(range(1, 11)))
        self.assertEqual(list(sequence(8, 10)), [8, 9, 10])
        self.assertEqual(list(sequence(1, 10, 3)), [1, 4, 7, 10])


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