# ____________________________________________________________________________________
#
# 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
from pyomo.common.tempfiles import TempfileManager

from pyomo.contrib.pynumero.dependencies import (
    numpy as np,
    numpy_available,
    scipy,
    scipy_available,
)
from pyomo.common.dependencies.scipy import sparse as spa

if not (numpy_available and scipy_available):
    raise unittest.SkipTest("Pynumero needs scipy and numpy to run NLP tests")

from pyomo.contrib.pynumero.asl import AmplInterface

if not AmplInterface.available():
    raise unittest.SkipTest(
        "Pynumero needs the ASL extension to run CyIpoptSolver tests"
    )

from pyomo.contrib.pynumero.algorithms.solvers.cyipopt_solver import cyipopt_available

if not cyipopt_available:
    raise unittest.SkipTest("Pynumero needs cyipopt to run CyIpoptSolver tests")

from pyomo.contrib.pynumero.algorithms.solvers.pyomo_ext_cyipopt import (
    ExternalInputOutputModel,
    PyomoExternalCyIpoptProblem,
)
from pyomo.contrib.pynumero.algorithms.solvers.cyipopt_solver import CyIpoptSolver


class PressureDropModel(ExternalInputOutputModel):
    def __init__(self):
        self._Pin = None
        self._c1 = None
        self._c2 = None
        self._F = None

    def set_inputs(self, input_values):
        assert len(input_values) == 4
        self._Pin = input_values[0]
        self._c1 = input_values[1]
        self._c2 = input_values[2]
        self._F = input_values[3]

    def evaluate_outputs(self):
        P1 = self._Pin - self._c1 * self._F**2
        P2 = P1 - self._c2 * self._F**2
        return np.asarray([P1, P2], dtype=np.float64)

    def evaluate_derivatives(self):
        jac = [
            [1, -self._F**2, 0, -2 * self._c1 * self._F],
            [1, -self._F**2, -self._F**2, -2 * self._F * (self._c1 + self._c2)],
        ]
        jac = np.asarray(jac, dtype=np.float64)
        return spa.coo_matrix(jac)


class TestExternalInputOutputModel(unittest.TestCase):
    def test_interface(self):
        # weird, this is really a test of the test class above
        # but we could add code later, so...
        iom = PressureDropModel()
        iom.set_inputs(np.ones(4))
        o = iom.evaluate_outputs()
        expected_o = np.asarray([0.0, -1.0], dtype=np.float64)
        self.assertTrue(np.array_equal(o, expected_o))

        jac = iom.evaluate_derivatives()
        expected_jac = np.asarray([[1, -1, 0, -2], [1, -1, -1, -4]], dtype=np.float64)
        self.assertTrue(np.array_equal(jac.todense(), expected_jac))

    def test_pyomo_external_model(self):
        m = pyo.ConcreteModel()
        m.Pin = pyo.Var(initialize=100, bounds=(0, None))
        m.c1 = pyo.Var(initialize=1.0, bounds=(0, None))
        m.c2 = pyo.Var(initialize=1.0, bounds=(0, None))
        m.F = pyo.Var(initialize=10, bounds=(0, None))

        m.P1 = pyo.Var()
        m.P2 = pyo.Var()

        m.F_con = pyo.Constraint(expr=m.F == 10)
        m.Pin_con = pyo.Constraint(expr=m.Pin == 100)

        # simple parameter estimation test
        m.obj = pyo.Objective(expr=(m.P1 - 90) ** 2 + (m.P2 - 40) ** 2)

        cyipopt_problem = PyomoExternalCyIpoptProblem(
            m, PressureDropModel(), [m.Pin, m.c1, m.c2, m.F], [m.P1, m.P2]
        )

        # check that the dummy variable is initialized
        # + pyo.value(m.P1) + pyo.value(m.P2)
        # not initialized - therefore should use zero
        expected_dummy_var_value = (
            pyo.value(m.Pin)
            + pyo.value(m.c1)
            + pyo.value(m.c2)
            + pyo.value(m.F)
            + 0
            + 0
        )
        self.assertAlmostEqual(
            pyo.value(m._dummy_variable_CyIpoptPyomoExNLP), expected_dummy_var_value
        )

        # solve the problem
        solver = CyIpoptSolver(
            cyipopt_problem, {'hessian_approximation': 'limited-memory'}
        )
        x, info = solver.solve(tee=False)
        cyipopt_problem.load_x_into_pyomo(x)
        self.assertAlmostEqual(pyo.value(m.c1), 0.1, places=5)
        self.assertAlmostEqual(pyo.value(m.c2), 0.5, places=5)

    def test_pyomo_external_model_scaling(self):
        m = pyo.ConcreteModel()
        m.Pin = pyo.Var(initialize=100, bounds=(0, None))
        m.c1 = pyo.Var(initialize=1.0, bounds=(0, None))
        m.c2 = pyo.Var(initialize=1.0, bounds=(0, None))
        m.F = pyo.Var(initialize=10, bounds=(0, None))

        m.P1 = pyo.Var()
        m.P2 = pyo.Var()

        m.F_con = pyo.Constraint(expr=m.F == 10)
        m.Pin_con = pyo.Constraint(expr=m.Pin == 100)

        # simple parameter estimation test
        m.obj = pyo.Objective(expr=(m.P1 - 90) ** 2 + (m.P2 - 40) ** 2)

        # set scaling parameters for the pyomo variables and constraints
        m.scaling_factor = pyo.Suffix(direction=pyo.Suffix.EXPORT)
        m.scaling_factor[m.obj] = 0.1  # scale the objective
        m.scaling_factor[m.Pin] = 2.0  # scale the variable
        m.scaling_factor[m.c1] = 3.0  # scale the variable
        m.scaling_factor[m.c2] = 4.0  # scale the variable
        m.scaling_factor[m.F] = 5.0  # scale the variable
        m.scaling_factor[m.P1] = 6.0  # scale the variable
        m.scaling_factor[m.P2] = 7.0  # scale the variable
        m.scaling_factor[m.F_con] = 8.0  # scale the pyomo constraint
        m.scaling_factor[m.Pin_con] = 9.0  # scale the pyomo constraint

        with TempfileManager.new_context() as temp:
            cyipopt_problem = PyomoExternalCyIpoptProblem(
                pyomo_model=m,
                ex_input_output_model=PressureDropModel(),
                inputs=[m.Pin, m.c1, m.c2, m.F],
                outputs=[m.P1, m.P2],
                outputs_eqn_scaling=[10.0, 11.0],
                nl_file_options={'file_determinism': 2},
            )
            logfile = temp.create_tempfile('_cyipopt-pyomo-ext-scaling.log')
            # solve the problem
            options = {
                'hessian_approximation': 'limited-memory',
                'nlp_scaling_method': 'user-scaling',
                'output_file': logfile,
                'file_print_level': 10,
                'max_iter': 0,
            }
            solver = CyIpoptSolver(cyipopt_problem, options=options)
            x, info = solver.solve(tee=False)
            cyipopt_problem.close()

            with open(logfile, 'r') as fd:
                solver_trace = fd.read()

        self.assertIn('nlp_scaling_method = user-scaling', solver_trace)
        self.assertIn(f"output_file = {logfile}", solver_trace)
        self.assertIn('objective scaling factor = 0.1', solver_trace)
        self.assertIn('x scaling provided', solver_trace)
        self.assertIn('c scaling provided', solver_trace)
        self.assertIn('d scaling provided', solver_trace)
        self.assertIn('DenseVector "x scaling vector" with 7 elements:', solver_trace)
        self.assertIn('x scaling vector[    1]= 6.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    2]= 7.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    3]= 5.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    4]= 2.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    5]= 1.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    6]= 3.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    7]= 4.0000000000000000e+00', solver_trace)
        self.assertIn('DenseVector "c scaling vector" with 5 elements:', solver_trace)
        self.assertIn('c scaling vector[    1]= 8.0000000000000000e+00', solver_trace)
        self.assertIn('c scaling vector[    2]= 9.0000000000000000e+00', solver_trace)
        self.assertIn('c scaling vector[    3]= 1.0000000000000000e+00', solver_trace)
        self.assertIn('c scaling vector[    4]= 1.0000000000000000e+01', solver_trace)
        self.assertIn('c scaling vector[    5]= 1.1000000000000000e+01', solver_trace)

    def test_pyomo_external_model_ndarray_scaling(self):
        m = pyo.ConcreteModel()
        m.Pin = pyo.Var(initialize=100, bounds=(0, None))
        m.c1 = pyo.Var(initialize=1.0, bounds=(0, None))
        m.c2 = pyo.Var(initialize=1.0, bounds=(0, None))
        m.F = pyo.Var(initialize=10, bounds=(0, None))

        m.P1 = pyo.Var()
        m.P2 = pyo.Var()

        m.F_con = pyo.Constraint(expr=m.F == 10)
        m.Pin_con = pyo.Constraint(expr=m.Pin == 100)

        # simple parameter estimation test
        m.obj = pyo.Objective(expr=(m.P1 - 90) ** 2 + (m.P2 - 40) ** 2)

        # set scaling parameters for the pyomo variables and constraints
        m.scaling_factor = pyo.Suffix(direction=pyo.Suffix.EXPORT)
        m.scaling_factor[m.obj] = 0.1  # scale the objective
        m.scaling_factor[m.Pin] = 2.0  # scale the variable
        m.scaling_factor[m.c1] = 3.0  # scale the variable
        m.scaling_factor[m.c2] = 4.0  # scale the variable
        m.scaling_factor[m.F] = 5.0  # scale the variable
        m.scaling_factor[m.P1] = 6.0  # scale the variable
        m.scaling_factor[m.P2] = 7.0  # scale the variable
        m.scaling_factor[m.F_con] = 8.0  # scale the pyomo constraint
        m.scaling_factor[m.Pin_con] = 9.0  # scale the pyomo constraint

        # test that this all works with ndarray input as well
        with TempfileManager.new_context() as temp:
            cyipopt_problem = PyomoExternalCyIpoptProblem(
                pyomo_model=m,
                ex_input_output_model=PressureDropModel(),
                inputs=[m.Pin, m.c1, m.c2, m.F],
                outputs=[m.P1, m.P2],
                outputs_eqn_scaling=np.asarray([10.0, 11.0], dtype=np.float64),
                nl_file_options={'file_determinism': 2},
            )
            logfile = temp.create_tempfile('_cyipopt-pyomo-ext-scaling-ndarray.log')
            # solve the problem
            options = {
                'hessian_approximation': 'limited-memory',
                'nlp_scaling_method': 'user-scaling',
                'output_file': logfile,
                'file_print_level': 10,
                'max_iter': 0,
            }
            solver = CyIpoptSolver(cyipopt_problem, options=options)
            x, info = solver.solve(tee=False)
            cyipopt_problem.close()

            with open(logfile, 'r') as fd:
                solver_trace = fd.read()

        self.assertIn('nlp_scaling_method = user-scaling', solver_trace)
        self.assertIn(f'output_file = {logfile}', solver_trace)
        self.assertIn('objective scaling factor = 0.1', solver_trace)
        self.assertIn('x scaling provided', solver_trace)
        self.assertIn('c scaling provided', solver_trace)
        self.assertIn('d scaling provided', solver_trace)
        self.assertIn('DenseVector "x scaling vector" with 7 elements:', solver_trace)
        self.assertIn('x scaling vector[    1]= 6.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    2]= 7.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    3]= 5.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    4]= 2.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    5]= 1.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    6]= 3.0000000000000000e+00', solver_trace)
        self.assertIn('x scaling vector[    7]= 4.0000000000000000e+00', solver_trace)
        self.assertIn('DenseVector "c scaling vector" with 5 elements:', solver_trace)
        self.assertIn('c scaling vector[    1]= 8.0000000000000000e+00', solver_trace)
        self.assertIn('c scaling vector[    2]= 9.0000000000000000e+00', solver_trace)
        self.assertIn('c scaling vector[    3]= 1.0000000000000000e+00', solver_trace)
        self.assertIn('c scaling vector[    4]= 1.0000000000000000e+01', solver_trace)
        self.assertIn('c scaling vector[    5]= 1.1000000000000000e+01', solver_trace)

    def test_pyomo_external_model_dummy_var_initialization(self):
        m = pyo.ConcreteModel()
        m.Pin = pyo.Var(initialize=100, bounds=(0, None))
        m.c1 = pyo.Var(initialize=1.0, bounds=(0, None))
        m.c2 = pyo.Var(initialize=1.0, bounds=(0, None))
        m.F = pyo.Var(initialize=10, bounds=(0, None))

        m.P1 = pyo.Var(initialize=75.0)
        m.P2 = pyo.Var(initialize=50.0)

        m.F_con = pyo.Constraint(expr=m.F == 10)
        m.Pin_con = pyo.Constraint(expr=m.Pin == 100)

        # simple parameter estimation test
        m.obj = pyo.Objective(expr=(m.P1 - 90) ** 2 + (m.P2 - 40) ** 2)

        cyipopt_problem = PyomoExternalCyIpoptProblem(
            m, PressureDropModel(), [m.Pin, m.c1, m.c2, m.F], [m.P1, m.P2]
        )

        # check that the dummy variable is initialized
        expected_dummy_var_value = (
            pyo.value(m.Pin)
            + pyo.value(m.c1)
            + pyo.value(m.c2)
            + pyo.value(m.F)
            + pyo.value(m.P1)
            + pyo.value(m.P2)
        )
        self.assertAlmostEqual(
            pyo.value(m._dummy_variable_CyIpoptPyomoExNLP), expected_dummy_var_value
        )
        # check that the dummy constraint is satisfied
        self.assertAlmostEqual(
            pyo.value(m._dummy_constraint_CyIpoptPyomoExNLP.body),
            pyo.value(m._dummy_constraint_CyIpoptPyomoExNLP.lower),
        )
        self.assertAlmostEqual(
            pyo.value(m._dummy_constraint_CyIpoptPyomoExNLP.body),
            pyo.value(m._dummy_constraint_CyIpoptPyomoExNLP.upper),
        )

        # solve the problem
        solver = CyIpoptSolver(
            cyipopt_problem, {'hessian_approximation': 'limited-memory'}
        )
        x, info = solver.solve(tee=False)
        cyipopt_problem.load_x_into_pyomo(x)
        self.assertAlmostEqual(pyo.value(m.c1), 0.1, places=5)
        self.assertAlmostEqual(pyo.value(m.c2), 0.5, places=5)
