# ____________________________________________________________________________________
#
# 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.
# ____________________________________________________________________________________

#################################################################################
# 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), and is copyright (c) 2018-2021
# 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.
#################################################################################
import pyomo.common.unittest as unittest

import pyomo.environ as pyo
from pyomo.common.collections import ComponentSet
from pyomo.core.expr.visitor import identify_variables
from pyomo.core.expr.compare import compare_expressions
from pyomo.contrib.mpc.modeling.cost_expressions import (
    get_penalty_from_constant_target,
    get_penalty_from_piecewise_constant_target,
    get_penalty_from_time_varying_target,
    get_penalty_from_target,
)
from pyomo.contrib.mpc.data.scalar_data import ScalarData
from pyomo.contrib.mpc.data.series_data import TimeSeriesData
from pyomo.contrib.mpc.data.interval_data import IntervalData


class TestTrackingCostConstantSetpoint(unittest.TestCase):
    def test_penalty_no_weights(self):
        m = pyo.ConcreteModel()
        m.time = pyo.Set(initialize=[1, 2, 3])
        m.v1 = pyo.Var(m.time, initialize={i: 1 * i for i in m.time})
        m.v2 = pyo.Var(m.time, initialize={i: 2 * i for i in m.time})

        setpoint_data = ScalarData({m.v1[:]: 3.0, m.v2[:]: 4.0})
        variables = [m.v1, m.v2]
        m.var_set, m.tracking_expr = get_penalty_from_constant_target(
            variables, m.time, setpoint_data
        )
        self.assertEqual(len(m.var_set), 2)
        self.assertIn(0, m.var_set)
        self.assertIn(1, m.var_set)

        var_sets = {
            (i, t): ComponentSet(identify_variables(m.tracking_expr[i, t]))
            for i in m.var_set
            for t in m.time
        }
        for i in m.time:
            for j in m.var_set:
                self.assertIn(variables[j][i], var_sets[j, i])
                pred_value = (1 * i - 3) ** 2 if j == 0 else (2 * i - 4) ** 2
                self.assertEqual(pred_value, pyo.value(m.tracking_expr[j, i]))
                pred_expr = (m.v1[i] - 3) ** 2 if j == 0 else (m.v2[i] - 4) ** 2
                self.assertTrue(
                    compare_expressions(pred_expr, m.tracking_expr[j, i].expr)
                )

    def test_penalty_with_weights(self):
        m = pyo.ConcreteModel()
        m.time = pyo.Set(initialize=[1, 2, 3])
        m.v1 = pyo.Var(m.time, initialize={i: 1 * i for i in m.time})
        m.v2 = pyo.Var(m.time, initialize={i: 2 * i for i in m.time})

        setpoint_data = ScalarData({m.v1[:]: 3.0, m.v2[:]: 4.0})
        weight_data = ScalarData({m.v1[:]: 0.1, m.v2[:]: 0.5})
        m.var_set = pyo.Set(initialize=[0, 1])
        variables = [m.v1, m.v2]
        new_set, m.tracking_expr = get_penalty_from_constant_target(
            variables,
            m.time,
            setpoint_data,
            weight_data=weight_data,
            variable_set=m.var_set,
        )
        self.assertIs(new_set, m.var_set)

        var_sets = {
            (i, t): ComponentSet(identify_variables(m.tracking_expr[i, t]))
            for i in m.var_set
            for t in m.time
        }
        for i in m.time:
            for j in m.var_set:
                self.assertIn(variables[j][i], var_sets[j, i])
                pred_value = (
                    0.1 * (1 * i - 3) ** 2 if j == 0 else 0.5 * (2 * i - 4) ** 2
                )
                self.assertAlmostEqual(pred_value, pyo.value(m.tracking_expr[j, i]))
                pred_expr = (
                    0.1 * (m.v1[i] - 3) ** 2 if j == 0 else 0.5 * (m.v2[i] - 4) ** 2
                )
                self.assertTrue(
                    compare_expressions(pred_expr, m.tracking_expr[j, i].expr)
                )

    def test_exceptions(self):
        m = pyo.ConcreteModel()
        m.time = pyo.Set(initialize=[1, 2, 3])
        m.v1 = pyo.Var(m.time, initialize={i: 1 * i for i in m.time})
        m.v2 = pyo.Var(m.time, initialize={i: 2 * i for i in m.time})

        setpoint_data = ScalarData({m.v1[:]: 3.0})
        weight_data = ScalarData({m.v2[:]: 0.1})
        with self.assertRaisesRegex(KeyError, "Setpoint data"):
            _, m.tracking_expr = get_penalty_from_constant_target(
                [m.v1, m.v2], m.time, setpoint_data
            )

        setpoint_data = ScalarData({m.v1[:]: 3.0, m.v2[:]: 4.0})
        with self.assertRaisesRegex(KeyError, "Tracking weight"):
            _, m.tracking_expr = get_penalty_from_constant_target(
                [m.v1, m.v2], m.time, setpoint_data, weight_data=weight_data
            )

    def test_add_set_after_expr(self):
        # A small gotcha that may come up. This is known behavior
        # due to Pyomo's "implicit set" addition.
        m = pyo.ConcreteModel()
        m.time = pyo.Set(initialize=[1, 2, 3])
        m.v1 = pyo.Var(m.time, initialize={i: 1 * i for i in m.time})
        m.v2 = pyo.Var(m.time, initialize={i: 2 * i for i in m.time})

        setpoint_data = ScalarData({m.v1[:]: 3.0, m.v2[:]: 4.0})
        weight_data = ScalarData({m.v1[:]: 0.1, m.v2[:]: 0.5})
        m.var_set = pyo.Set(initialize=[0, 1])
        variables = [m.v1, m.v2]
        new_set, tr_expr = get_penalty_from_constant_target(
            variables,
            m.time,
            setpoint_data,
            weight_data=weight_data,
            variable_set=m.var_set,
        )
        m.tracking_expr = tr_expr  # new_set gets added and assigned a name
        msg = "Attempting to re-assign"
        with self.assertRaisesRegex(RuntimeError, msg):
            # attempting to add the same component twice
            m.variable_set = new_set


class TestTrackingCostPiecewiseSetpoint(unittest.TestCase):
    def _make_model(self, n_time_points=3):
        m = pyo.ConcreteModel()
        m.time = pyo.Set(initialize=list(range(n_time_points)))
        m.comp = pyo.Set(initialize=["A", "B"])
        m.var = pyo.Var(
            m.time, m.comp, initialize={(i, j): 1.1 * i for i, j in m.time * m.comp}
        )
        return m

    def test_piecewise_penalty_no_weights(self):
        m = self._make_model(n_time_points=5)

        variables = [pyo.Reference(m.var[:, "A"]), pyo.Reference(m.var[:, "B"])]
        setpoint_data = IntervalData(
            {m.var[:, "A"]: [2.0, 2.5], m.var[:, "B"]: [3.0, 3.5]}, [(0, 2), (2, 4)]
        )
        m.var_set, m.tracking_cost = get_penalty_from_piecewise_constant_target(
            variables, m.time, setpoint_data
        )
        for i in m.time:
            for j in m.var_set:
                if i <= 2:
                    pred_expr = (
                        (m.var[i, "A"] - 2.0) ** 2
                        if j == 0
                        else (m.var[i, "B"] - 3.0) ** 2
                    )
                else:
                    pred_expr = (
                        (m.var[i, "A"] - 2.5) ** 2
                        if j == 0
                        else (m.var[i, "B"] - 3.5) ** 2
                    )
                pred_value = pyo.value(pred_expr)
                self.assertEqual(pred_value, pyo.value(m.tracking_cost[j, i]))
                self.assertTrue(
                    compare_expressions(pred_expr, m.tracking_cost[j, i].expr)
                )

    def test_piecewise_penalty_with_weights(self):
        m = self._make_model(n_time_points=5)

        variables = [pyo.Reference(m.var[:, "A"]), pyo.Reference(m.var[:, "B"])]
        setpoint_data = IntervalData(
            {m.var[:, "A"]: [2.0, 2.5], m.var[:, "B"]: [3.0, 3.5]}, [(0, 2), (2, 4)]
        )
        weight_data = {
            pyo.ComponentUID(m.var[:, "A"]): 10.0,
            pyo.ComponentUID(m.var[:, "B"]): 0.1,
        }
        m.var_set, m.tracking_cost = get_penalty_from_piecewise_constant_target(
            variables, m.time, setpoint_data, weight_data=weight_data
        )
        for i in m.time:
            for j in m.var_set:
                if i <= 2:
                    pred_expr = (
                        10.0 * (m.var[i, "A"] - 2.0) ** 2
                        if j == 0
                        else 0.1 * (m.var[i, "B"] - 3.0) ** 2
                    )
                else:
                    pred_expr = (
                        10.0 * (m.var[i, "A"] - 2.5) ** 2
                        if j == 0
                        else 0.1 * (m.var[i, "B"] - 3.5) ** 2
                    )
                pred_value = pyo.value(pred_expr)
                self.assertEqual(pred_value, pyo.value(m.tracking_cost[j, i]))
                self.assertTrue(
                    compare_expressions(pred_expr, m.tracking_cost[j, i].expr)
                )

    def test_piecewise_penalty_exceptions(self):
        m = self._make_model(n_time_points=5)

        variables = [pyo.Reference(m.var[:, "A"]), pyo.Reference(m.var[:, "B"])]
        setpoint_data = IntervalData({m.var[:, "A"]: [2.0, 2.5]}, [(0, 2), (2, 4)])
        weight_data = {
            pyo.ComponentUID(m.var[:, "A"]): 10.0,
            pyo.ComponentUID(m.var[:, "B"]): 0.1,
        }
        msg = "Setpoint data does not contain"
        with self.assertRaisesRegex(KeyError, msg):
            tr_cost = get_penalty_from_piecewise_constant_target(
                variables, m.time, setpoint_data, weight_data=weight_data
            )

        setpoint_data = IntervalData(
            {m.var[:, "A"]: [2.0, 2.5], m.var[:, "B"]: [3.0, 3.5]}, [(0, 2), (2, 4)]
        )
        weight_data = {pyo.ComponentUID(m.var[:, "A"]): 10.0}
        msg = "Tracking weight does not contain"
        with self.assertRaisesRegex(KeyError, msg):
            tr_cost = get_penalty_from_piecewise_constant_target(
                variables, m.time, setpoint_data, weight_data=weight_data
            )


class TestTrackingCostVaryingSetpoint(unittest.TestCase):
    def _make_model(self, n_time_points=3):
        m = pyo.ConcreteModel()
        m.time = pyo.Set(initialize=list(range(n_time_points)))
        m.comp = pyo.Set(initialize=["A", "B"])
        m.var = pyo.Var(
            m.time, m.comp, initialize={(i, j): 1.1 * i for i, j in m.time * m.comp}
        )
        return m

    def test_varying_setpoint_no_weights(self):
        m = self._make_model(n_time_points=5)
        variables = [pyo.Reference(m.var[:, "A"]), pyo.Reference(m.var[:, "B"])]
        A_setpoint = [1.0 - 0.1 * i for i in range(len(m.time))]
        B_setpoint = [5.0 + 0.1 * i for i in range(len(m.time))]
        setpoint_data = TimeSeriesData(
            {m.var[:, "A"]: A_setpoint, m.var[:, "B"]: B_setpoint}, m.time
        )
        m.var_set, m.tracking_cost = get_penalty_from_time_varying_target(
            variables, m.time, setpoint_data
        )
        for i, t in enumerate(m.time):
            for j in m.var_set:
                pred_expr = (
                    (m.var[t, "A"] - A_setpoint[i]) ** 2
                    if j == 0
                    else (m.var[t, "B"] - B_setpoint[i]) ** 2
                )
                pred_value = pyo.value(pred_expr)
                self.assertEqual(pred_value, pyo.value(m.tracking_cost[j, t]))
                self.assertTrue(
                    compare_expressions(pred_expr, m.tracking_cost[j, t].expr)
                )

    def test_varying_setpoint_with_weights(self):
        m = self._make_model(n_time_points=5)
        variables = [pyo.Reference(m.var[:, "A"]), pyo.Reference(m.var[:, "B"])]
        A_setpoint = [1.0 - 0.1 * i for i in range(len(m.time))]
        B_setpoint = [5.0 + 0.1 * i for i in range(len(m.time))]
        setpoint_data = TimeSeriesData(
            {m.var[:, "A"]: A_setpoint, m.var[:, "B"]: B_setpoint}, m.time
        )
        weight_data = {
            pyo.ComponentUID(m.var[:, "A"]): 10.0,
            pyo.ComponentUID(m.var[:, "B"]): 0.1,
        }
        m.var_set, m.tracking_cost = get_penalty_from_time_varying_target(
            variables, m.time, setpoint_data, weight_data=weight_data
        )
        for i, t in enumerate(m.time):
            for j in m.var_set:
                pred_expr = (
                    10.0 * (m.var[t, "A"] - A_setpoint[i]) ** 2
                    if j == 0
                    else 0.1 * (m.var[t, "B"] - B_setpoint[i]) ** 2
                )
                pred_value = pyo.value(pred_expr)
                self.assertEqual(pred_value, pyo.value(m.tracking_cost[j, t]))
                self.assertTrue(
                    compare_expressions(pred_expr, m.tracking_cost[j, t].expr)
                )

    def test_varying_setpoint_exceptions(self):
        m = self._make_model(n_time_points=5)
        variables = [pyo.Reference(m.var[:, "A"]), pyo.Reference(m.var[:, "B"])]
        A_setpoint = [1.0 - 0.1 * i for i in range(len(m.time))]
        B_setpoint = [5.0 + 0.1 * i for i in range(len(m.time))]
        setpoint_data = TimeSeriesData(
            {m.var[:, "A"]: A_setpoint, m.var[:, "B"]: B_setpoint},
            [i + 10 for i in m.time],
        )
        weight_data = {
            pyo.ComponentUID(m.var[:, "A"]): 10.0,
            pyo.ComponentUID(m.var[:, "B"]): 0.1,
        }
        msg = "Mismatch in time points"
        with self.assertRaisesRegex(RuntimeError, msg):
            # Time-varying setpoint specifies different time points
            # from our time set.
            var_set, tr_cost = get_penalty_from_time_varying_target(
                variables, m.time, setpoint_data, weight_data=weight_data
            )

        setpoint_data = TimeSeriesData({m.var[:, "A"]: A_setpoint}, m.time)
        msg = "Setpoint data does not contain"
        with self.assertRaisesRegex(KeyError, msg):
            var_set, tr_cost = get_penalty_from_time_varying_target(
                variables, m.time, setpoint_data, weight_data=weight_data
            )

        setpoint_data = TimeSeriesData(
            {m.var[:, "A"]: A_setpoint, m.var[:, "B"]: B_setpoint}, m.time
        )
        weight_data = {pyo.ComponentUID(m.var[:, "A"]): 10.0}
        msg = "Tracking weight does not contain"
        with self.assertRaisesRegex(KeyError, msg):
            tr_cost = get_penalty_from_time_varying_target(
                variables, m.time, setpoint_data, weight_data=weight_data
            )


class TestGetPenaltyFromTarget(unittest.TestCase):
    def _make_model(self, n_time_points=3):
        m = pyo.ConcreteModel()
        m.time = pyo.Set(initialize=list(range(n_time_points)))
        m.comp = pyo.Set(initialize=["A", "B"])
        m.var = pyo.Var(
            m.time, m.comp, initialize={(i, j): 1.1 * i for i, j in m.time * m.comp}
        )
        return m

    def test_constant_setpoint(self):
        m = self._make_model()
        setpoint = {m.var[:, "A"]: 0.3, m.var[:, "B"]: 0.4}
        variables = [pyo.Reference(m.var[:, "B"]), pyo.Reference(m.var[:, "A"])]
        m.var_set, m.penalty = get_penalty_from_target(variables, m.time, setpoint)
        pred_expr = {
            (i, t): (
                (m.var[t, "B"] - 0.4) ** 2 if i == 0 else (m.var[t, "A"] - 0.3) ** 2
            )
            for i, t in m.var_set * m.time
        }
        for t in m.time:
            for i in m.var_set:
                self.assertTrue(
                    compare_expressions(pred_expr[i, t], m.penalty[i, t].expr)
                )
                self.assertEqual(pyo.value(pred_expr[i, t]), pyo.value(m.penalty[i, t]))

    def test_constant_setpoint_with_ScalarData(self):
        m = self._make_model()
        setpoint = ScalarData({m.var[:, "A"]: 0.3, m.var[:, "B"]: 0.4})
        variables = [pyo.Reference(m.var[:, "B"]), pyo.Reference(m.var[:, "A"])]
        m.var_set, m.penalty = get_penalty_from_target(variables, m.time, setpoint)
        pred_expr = {
            (i, t): (
                (m.var[t, "B"] - 0.4) ** 2 if i == 0 else (m.var[t, "A"] - 0.3) ** 2
            )
            for i, t in m.var_set * m.time
        }
        for t in m.time:
            for i in m.var_set:
                self.assertTrue(
                    compare_expressions(pred_expr[i, t], m.penalty[i, t].expr)
                )
                self.assertEqual(pyo.value(pred_expr[i, t]), pyo.value(m.penalty[i, t]))

    def test_varying_setpoint(self):
        m = self._make_model(n_time_points=5)
        A_target = [0.4, 0.6, 0.1, 0.0, 1.1]
        B_target = [0.8, 0.9, 1.3, 1.5, 1.4]
        setpoint = ({m.var[:, "A"]: A_target, m.var[:, "B"]: B_target}, m.time)
        variables = [pyo.Reference(m.var[:, "B"]), pyo.Reference(m.var[:, "A"])]
        m.var_set, m.penalty = get_penalty_from_target(variables, m.time, setpoint)

        target = {
            (i, t): A_target[j] if i == 1 else B_target[t]
            for i in m.var_set
            for (j, t) in enumerate(m.time)
        }
        for i, t in m.var_set * m.time:
            pred_expr = (variables[i][t] - target[i, t]) ** 2
            self.assertTrue(compare_expressions(pred_expr, m.penalty[i, t].expr))
            self.assertEqual(pyo.value(pred_expr), pyo.value(m.penalty[i, t]))

    def test_piecewise_constant_setpoint(self):
        m = self._make_model(n_time_points=5)
        A_target = [0.3, 0.9, 0.7]
        B_target = [1.1, 0.1, 0.5]
        setpoint = (
            {m.var[:, "A"]: A_target, m.var[:, "B"]: B_target},
            [(0.0, 0.0), (0.0, 2.0), (2.0, 4.0)],
        )
        variables = [pyo.Reference(m.var[:, "B"]), pyo.Reference(m.var[:, "A"])]
        m.var_set, m.penalty = get_penalty_from_target(variables, m.time, setpoint)
        target = {
            (i, j): A_target[j] if i == 1 else B_target[j]
            for i in m.var_set
            for j in range(len(A_target))
        }
        for i, t in m.var_set * m.time:
            if t == 0:
                idx = 0
            elif t <= 2.0:
                idx = 1
            elif t <= 4.0:
                idx = 2
            pred_expr = (variables[i][t] - target[i, idx]) ** 2
            self.assertTrue(compare_expressions(pred_expr, m.penalty[i, t].expr))
            self.assertEqual(pyo.value(pred_expr), pyo.value(m.penalty[i, t]))

    def test_bad_argument(self):
        m = self._make_model(n_time_points=3)
        A_target = [0.4, 0.6, 0.1]
        B_target = [0.8, 0.9, 1.3]
        setpoint = ({m.var[:, "A"]: A_target, m.var[:, "B"]: B_target}, m.time)
        variables = [pyo.Reference(m.var[:, "B"]), pyo.Reference(m.var[:, "A"])]
        msg = "tolerance.*can only be used"
        with self.assertRaisesRegex(RuntimeError, msg):
            m.var_set, m.penalty = get_penalty_from_target(
                variables, m.time, setpoint, tolerance=1e-8
            )

    def test_bad_data_tuple(self):
        m = self._make_model(n_time_points=3)
        A_target = [0.4, 0.6, 0.1]
        B_target = [0.8, 0.9, 1.3]
        setpoint = (
            {m.var[:, "A"]: A_target, m.var[:, "B"]: B_target},
            m.time,
            "something else",
        )
        variables = [pyo.Reference(m.var[:, "B"]), pyo.Reference(m.var[:, "A"])]
        msg = "tuple of length two"
        with self.assertRaisesRegex(TypeError, msg):
            m.var_set, m.penalty = get_penalty_from_target(variables, m.time, setpoint)

    def test_bad_data_tuple_entry_0(self):
        m = self._make_model(n_time_points=3)
        A_target = [0.4, 0.6, 0.1]
        B_target = [0.8, 0.9, 1.3]
        setpoint = ([(m.var[:, "A"], A_target), (m.var[:, "B"], B_target)], m.time)
        variables = [pyo.Reference(m.var[:, "B"]), pyo.Reference(m.var[:, "A"])]
        msg = "must be instance of MutableMapping"
        with self.assertRaisesRegex(TypeError, msg):
            m.var_set, m.penalty = get_penalty_from_target(variables, m.time, setpoint)

    def test_empty_time_list(self):
        m = self._make_model(n_time_points=3)
        A_target = []
        B_target = []
        setpoint = ({m.var[:, "A"]: A_target, m.var[:, "B"]: B_target}, [])
        variables = [pyo.Reference(m.var[:, "B"]), pyo.Reference(m.var[:, "A"])]
        msg = "Time sequence.*is empty"
        with self.assertRaisesRegex(ValueError, msg):
            m.var_set, m.penalty = get_penalty_from_target(variables, m.time, setpoint)

    def test_bad_time_list(self):
        m = self._make_model(n_time_points=3)
        A_target = [0.4, 0.6, 0.1]
        B_target = [0.8, 0.9, 1.3]
        setpoint = (
            dict([(m.var[:, "A"], A_target), (m.var[:, "B"], B_target)]),
            [0.0, (0.1, 0.2), 0.3],
        )
        variables = [pyo.Reference(m.var[:, "B"]), pyo.Reference(m.var[:, "A"])]
        msg = "Second entry of data tuple must be"
        with self.assertRaisesRegex(TypeError, msg):
            m.var_set, m.penalty = get_penalty_from_target(variables, m.time, setpoint)


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