# ____________________________________________________________________________________
#
# 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 parametric.mod example from sIPOPT
#
# Original implementation by Hans Pirnay is in pyomo/examples/pyomo/suffixes
#

from pyomo.environ import *
from pyomo.contrib.sensitivity_toolbox.sens import sensitivity_calculation


def create_model():
    '''Create a concrete Pyomo model for this example'''
    m = ConcreteModel()

    m.x1 = Var(initialize=0.15, within=NonNegativeReals)
    m.x2 = Var(initialize=0.15, within=NonNegativeReals)
    m.x3 = Var(initialize=0.0, within=NonNegativeReals)

    m.eta1 = Param(initialize=4.5, mutable=True)
    m.eta2 = Param(initialize=1.0, mutable=True)

    m.const1 = Constraint(expr=6 * m.x1 + 3 * m.x2 + 2 * m.x3 - m.eta1 == 0)
    m.const2 = Constraint(expr=m.eta2 * m.x1 + m.x2 - m.x3 - 1 == 0)
    m.cost = Objective(expr=m.x1**2 + m.x2**2 + m.x3**2)

    return m


def run_example(print_flag=True):
    '''
    Execute the example

    Arguments:
        print_flag: Toggle on/off printing

    Returns
        sln_dict: Dictionary containing solution (used for automated testing)

    '''
    m = create_model()

    m.perturbed_eta1 = Param(initialize=4.0)
    m.perturbed_eta2 = Param(initialize=1.0)

    m_kaug_dsdp = sensitivity_calculation(
        'k_aug', m, [m.eta1, m.eta2], [m.perturbed_eta1, m.perturbed_eta2], tee=True
    )

    if print_flag:
        print("\nOriginal parameter values:")
        print("\teta1 =", m.eta1())
        print("\teta2 =", m.eta2())

        print("Initial point:")
        print("\tObjective =", value(m.cost))
        print("\tx1 =", m.x1())
        print("\tx2 =", m.x2())
        print("\tx3 =", m.x3())

        # Kaug saves only approximated solutions not original solutions
        print("\nNew parameter values:")
        print("\teta1 =", m_kaug_dsdp.perturbed_eta1())
        print("\teta2 =", m_kaug_dsdp.perturbed_eta2())

        print("(Approximate) solution with the new parameter values:")
        print("\tObjective =", m_kaug_dsdp.cost())
        print("\tx1 =", m_kaug_dsdp.x1())
        print("\tx2 =", m_kaug_dsdp.x2())
        print("\tx3 =", m_kaug_dsdp.x3())

    # Save the results in a dictionary.
    # This is optional and makes automated testing convenient.
    # This code is not required for a Minimum Working Example (MWE)
    d = dict()
    d['eta1'] = m.eta1()
    d['eta2'] = m.eta2()
    d['x1_init'] = m.x1()
    d['x2_init'] = m.x2()
    d['x3_init'] = m.x3()
    d['eta1_pert'] = m_kaug_dsdp.perturbed_eta1()
    d['eta2_pert'] = m_kaug_dsdp.perturbed_eta2()
    d['cost_pert'] = m_kaug_dsdp.cost()
    d['x1_pert'] = m_kaug_dsdp.x1()
    d['x2_pert'] = m_kaug_dsdp.x2()
    d['x3_pert'] = m_kaug_dsdp.x3()

    return d


if __name__ == '__main__':
    d = run_example()
