#################################################################################
# The Institute for the Design of Advanced Energy Systems Integrated Platform
# Framework (IDAES IP) was produced under the DOE Institute for the
# Design of Advanced Energy Systems (IDAES).
#
# Copyright (c) 2018-2024 by the software owners: The Regents of the
# University of California, through Lawrence Berkeley National Laboratory,
# National Technology & Engineering Solutions of Sandia, LLC, Carnegie Mellon
# University, West Virginia University Research Corporation, et al.
# All rights reserved.  Please see the files COPYRIGHT.md and LICENSE.md
# for full copyright and license information.
#################################################################################
"""
This module contains tests for the variable replace transformation.
"""

import pytest
import pyomo.environ as pyo
import idaes.core.plugins

__author__ = "John Eslick"


@pytest.mark.unit
def test_1():
    # Test scalar variables
    rp = pyo.TransformationFactory("replace_variables")
    m = pyo.ConcreteModel()
    m.x = pyo.Var(initialize=2)
    m.y = pyo.Var(initialize=3)
    m.z = pyo.Var(initialize=0)
    m.c1 = pyo.Constraint(expr=m.z == m.x + m.y)
    m.e1 = pyo.Expression(expr=m.x**m.y)
    m.o1 = pyo.Objective(expr=m.y - m.x)

    assert m.c1.body() == -5  # hope constraint arrangement is deterministic
    assert pyo.value(m.e1) == 8
    assert pyo.value(m.o1) == 1
    rp.apply_to(m, substitute=[(m.y, m.x)])
    assert m.c1.body() == -4
    assert pyo.value(m.e1) == 4
    assert pyo.value(m.o1) == 0


@pytest.mark.unit
def test_2():
    # Test vector variables and sums
    rp = pyo.TransformationFactory("replace_variables")
    m = pyo.ConcreteModel()
    m.x = pyo.Var(["a", "b", "c"], initialize=2)
    m.y = pyo.Var(initialize=3)
    m.z = pyo.Var(initialize=0)
    m.c1 = pyo.Constraint(expr=m.z == m.x["a"] + m.x["b"] + m.x["c"])
    m.e1 = pyo.Expression(expr=sum(m.x[i] for i in m.x))

    assert m.c1.body() == -6  # hope constraint arrangement is deterministic
    assert pyo.value(m.e1) == 6
    rp.apply_to(m, substitute=[(m.x["c"], m.y)])
    assert m.c1.body() == -7
    assert pyo.value(m.e1) == 7


@pytest.mark.unit
def test_3():
    # Test expression in constraint
    rp = pyo.TransformationFactory("replace_variables")
    m = pyo.ConcreteModel()
    m.x = pyo.Var(["a", "b", "c"], initialize=2)
    m.y = pyo.Var(initialize=3)
    m.z = pyo.Var(initialize=0)
    m.e1 = pyo.Expression(expr=sum(m.x[i] for i in m.x))
    m.c1 = pyo.Constraint(expr=m.z == m.e1)

    assert m.c1.body() == -6  # hope constraint arrangement is deterministic
    rp.apply_to(m, substitute=[(m.x["c"], m.y)])
    assert m.c1.body() == -7


@pytest.mark.unit
def test_4():
    # Test expression in objective
    rp = pyo.TransformationFactory("replace_variables")
    m = pyo.ConcreteModel()
    m.x = pyo.Var(["a", "b", "c"], initialize=2)
    m.y = pyo.Var(initialize=3)
    m.z = pyo.Var(initialize=0)
    m.e1 = pyo.Expression(expr=sum(m.x[i] for i in m.x))
    m.c1 = pyo.Constraint(expr=m.z == m.e1)
    m.o1 = pyo.Objective(expr=m.e1)

    assert pyo.value(m.o1) == 6
    rp.apply_to(m, substitute=[(m.x["c"], m.y)])
    assert pyo.value(m.o1) == 7


@pytest.mark.unit
def test_4():
    # Test in a hierarchical model
    rp = pyo.TransformationFactory("replace_variables")
    m = pyo.ConcreteModel()
    m.b1 = pyo.Block()
    m.b1.b2 = pyo.Block()
    x = m.b1.b2.x = pyo.Var(["a", "b", "c"], initialize=2)
    m.y = pyo.Var(initialize=3)
    m.z = pyo.Var(initialize=0)
    m.e1 = pyo.Expression(expr=sum(x[i] for i in x))
    m.b1.c1 = pyo.Constraint(expr=m.z == m.e1)
    m.o1 = pyo.Objective(expr=m.e1)

    assert pyo.value(m.o1) == 6
    assert m.b1.c1.body() == -6
    assert pyo.value(m.e1) == 6
    rp.apply_to(m, substitute=[(x["c"], m.y)])
    assert pyo.value(m.o1) == 7
    assert m.b1.c1.body() == -7
    assert pyo.value(m.e1) == 7


@pytest.mark.unit
def test_5():
    # Test indexed var replace
    rp = pyo.TransformationFactory("replace_variables")
    m = pyo.ConcreteModel()
    m.b1 = pyo.Block()
    m.b1.b2 = pyo.Block()
    x = m.b1.b2.x = pyo.Var(["a", "b", "c"], [1, 2, 3], initialize=2)
    m.y = pyo.Var(["a", "b", "c", "d"], [1, 2, 3], initialize=3)
    m.z = pyo.Var(initialize=0)
    m.e1 = pyo.Expression(expr=sum(x[i] for i in x))
    m.b1.c1 = pyo.Constraint(expr=m.z == m.e1)
    m.o1 = pyo.Objective(expr=m.e1)

    assert pyo.value(m.o1) == 18
    assert m.b1.c1.body() == -18
    assert pyo.value(m.e1) == 18
    rp.apply_to(m, substitute=[(x, m.y)])
    assert pyo.value(m.o1) == 27
    assert m.b1.c1.body() == -27
    assert pyo.value(m.e1) == 27


@pytest.mark.unit
def test_6():
    # Test non-variable exception
    rp = pyo.TransformationFactory("replace_variables")
    m = pyo.ConcreteModel()
    m.b1 = pyo.Block()
    m.b1.b2 = pyo.Block()
    x = m.b1.b2.x = pyo.Var(["a", "b", "c"], [1, 2, 3], initialize=2)
    m.y = pyo.Var(["a", "b", "c", "d"], [1, 2, 3], initialize=3)
    m.z = pyo.Var(initialize=0)

    with pytest.raises(TypeError):
        rp.apply_to(m, substitute=[(x, m.b1)])
    with pytest.raises(TypeError):
        rp.apply_to(m, substitute=[(m.b1, x)])


@pytest.mark.unit
def test_7():
    # Test replace indexed by non-indexed
    rp = pyo.TransformationFactory("replace_variables")
    m = pyo.ConcreteModel()
    m.b1 = pyo.Block()
    m.b1.b2 = pyo.Block()
    x = m.b1.b2.x = pyo.Var(["a", "b", "c"], [1, 2, 3], initialize=2)
    m.y = pyo.Var(["a", "b", "c", "d"], [1, 2, 3], initialize=3)
    m.z = pyo.Var(initialize=0)

    assert x.is_indexed()
    assert not m.z.is_indexed()
    with pytest.raises(TypeError):
        rp.apply_to(m, substitute=[(x, m.z)])


@pytest.mark.unit
def test_8():
    # Test replace indexed by indexed var that doesn't have enough/right indexes
    rp = pyo.TransformationFactory("replace_variables")
    m = pyo.ConcreteModel()
    m.b1 = pyo.Block()
    m.b1.b2 = pyo.Block()
    x = m.b1.b2.x = pyo.Var(["a", "b", "c"], [1, 2, 3], initialize=2)
    m.y = pyo.Var(["a", "b", "d"], [1, 2, 3], initialize=3)

    with pytest.raises(ValueError):
        rp.apply_to(m, substitute=[(x, m.y)])


@pytest.mark.unit
def test_9():
    # test with references the way we handle time indexing a lot in IDAES
    rp = pyo.TransformationFactory("replace_variables")
    block_set = {1, 2, 3}
    m = pyo.ConcreteModel()
    m.b1 = pyo.Block(block_set)
    for i in block_set:
        m.b1[i].x = pyo.Var(initialize=2)
    m.y = pyo.Var([1, 2, 3], initialize=3)
    m.xx = pyo.Reference(m.b1[:].x)
    m.display()

    m.e1 = pyo.Expression(expr=sum(m.xx[i] for i in m.xx))
    m.e2 = pyo.Expression(expr=sum(m.b1[i].x for i in m.b1))

    assert pyo.value(m.e1) == 6
    assert pyo.value(m.e2) == 6
    rp.apply_to(m, substitute=[(m.xx, m.y)])
    assert pyo.value(m.e1) == 9
    assert pyo.value(m.e2) == 9


@pytest.mark.unit
def test_10():
    # test with more than one variable in the list
    rp = pyo.TransformationFactory("replace_variables")
    m = pyo.ConcreteModel()
    m.x = pyo.Var(["a", "b", "c"], initialize=2)
    m.a = pyo.Var(initialize=5)
    m.b = pyo.Var(initialize=6)
    m.c = pyo.Var(initialize=7)

    m.e1 = pyo.Expression(expr=sum(m.x[i] for i in m.x))

    assert pyo.value(m.e1) == 6
    rp.apply_to(m, substitute=[(m.x["a"], m.a), (m.x["b"], m.b), (m.x["c"], m.c)])
    assert pyo.value(m.e1) == 18
