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

import pyomo.common.unittest as unittest
import pyomo.environ as pyo
import pyomo.contrib.mpc as mpc
from pyomo.contrib.mpc.examples.cstr.run_openloop import run_cstr_openloop

ipopt_available = pyo.SolverFactory("ipopt").available()


@unittest.skipIf(not ipopt_available, "ipopt is not available")
class TestCSTROpenLoop(unittest.TestCase):
    # This data was obtained from a run of this code. The test is
    # intended to make sure that values do not change, not that
    # they are correct in some absolute sense.
    _pred_A_data = [
        1.00000,
        1.50000,
        1.83333,
        2.05556,
        2.20370,
        2.30247,
        2.36831,
        2.41221,
        2.44147,
        2.46098,
        2.47399,
        2.48266,
        2.48844,
        2.36031,
        2.27039,
        2.20729,
        2.16301,
        2.13194,
        2.11013,
        2.09483,
        2.08409,
        2.07656,
        2.07127,
        2.06756,
        2.06495,
        2.16268,
        2.22893,
        2.27385,
        2.30431,
        2.32495,
        2.33895,
        2.34844,
        2.35488,
        2.35924,
        2.36220,
        2.36420,
        2.36556,
        2.36648,
        2.36711,
        2.36753,
        2.36782,
    ]
    _pred_B_data = [
        0.00000,
        0.30200,
        0.61027,
        0.90132,
        1.16380,
        1.39353,
        1.59049,
        1.75683,
        1.89576,
        2.01081,
        2.10544,
        2.18289,
        2.24600,
        2.41517,
        2.54001,
        2.63284,
        2.70242,
        2.75502,
        2.79516,
        2.82605,
        2.85007,
        2.86890,
        2.88380,
        2.89569,
        2.90526,
        2.81484,
        2.75455,
        2.71450,
        2.68802,
        2.67062,
        2.65927,
        2.65195,
        2.64728,
        2.64436,
        2.64258,
        2.64153,
        2.64096,
        2.64067,
        2.64057,
        2.64058,
        2.64064,
    ]

    def _get_input_sequence(self):
        input_sequence = mpc.TimeSeriesData(
            {"flow_in[*]": [0.1, 1.0, 0.7, 0.9]}, [0.0, 3.0, 6.0, 10.0]
        )
        return mpc.data.convert.series_to_interval(input_sequence)

    def test_openloop_simulation(self):
        input_sequence = self._get_input_sequence()
        ntfe = 4
        model_horizon = 1.0
        simulation_steps = 10
        m, sim_data = run_cstr_openloop(
            input_sequence, model_horizon=1.0, ntfe=4, simulation_steps=10
        )
        sim_time_points = [
            model_horizon / ntfe * i for i in range(simulation_steps * ntfe + 1)
        ]

        AB_data = sim_data.extract_variables([m.conc[:, "A"], m.conc[:, "B"]])

        A_cuid = sim_data.get_cuid(m.conc[:, "A"])
        B_cuid = sim_data.get_cuid(m.conc[:, "B"])
        pred_data = {A_cuid: self._pred_A_data, B_cuid: self._pred_B_data}

        self.assertStructuredAlmostEqual(pred_data, AB_data.get_data(), delta=1e-3)
        self.assertStructuredAlmostEqual(
            sim_time_points, AB_data.get_time_points(), delta=1e-7
        )


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