# ____________________________________________________________________________________
#
# 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.
# ____________________________________________________________________________________
#
# Reference: "Optimal Control Theory: An Introduction", Donald E. Kirk,
# 		(1970/1998)
#
# Example 5.2-1
#
# x'(t) = a*x(t)+u(t)
#
# min J(u) = (1/2)*H*x^2(T)+\int_0^T ((1/4)*u^2(t)dt)

from pyomo.environ import (
    ConcreteModel,
    Param,
    Var,
    Objective,
    Constraint,
    Suffix,
    value,
    TransformationFactory,
    SolverFactory,
)
from pyomo.dae import ContinuousSet, DerivativeVar
from pyomo.dae.simulator import Simulator
from pyomo.contrib.sensitivity_toolbox.sens import sensitivity_calculation
from pyomo.common.dependencies.matplotlib import pyplot as plt


def create_model():
    m = ConcreteModel()

    m.a = Param(initialize=-0.2, mutable=True)
    m.H = Param(initialize=0.5, mutable=True)
    m.T = 15

    m.t = ContinuousSet(bounds=(0, m.T))

    m.x = Var(m.t)
    m.F = Var(m.t)
    m.u = Var(m.t, initialize=0, bounds=(-0.2, 0))

    m.dx = DerivativeVar(m.x, wrt=m.t)
    m.df0 = DerivativeVar(m.F, wrt=m.t)

    m.x[0].fix(5)
    m.F[0].fix(0)

    def _x(m, t):
        return m.dx[t] == m.a * m.x[t] + m.u[t]

    m.x_dot = Constraint(m.t, rule=_x)

    def _f0(m, t):
        return m.df0[t] == 0.25 * m.u[t] ** 2

    m.FDiffCon = Constraint(m.t, rule=_f0)

    def _Cost(m):
        return 0.5 * m.H * m.x[m.T] ** 2 + m.F[m.T]

    m.J = Objective(rule=_Cost)

    return m


def initialize_model(m, nfe):
    u_profile = {0: -0.06}

    m.u_input = Suffix(direction=Suffix.LOCAL)
    m.u_input[m.u] = u_profile

    sim = Simulator(m, package='scipy')
    tsim, profiles = sim.simulate(numpoints=100, varying_inputs=m.u_input)

    discretizer = TransformationFactory('dae.collocation')
    discretizer.apply_to(m, nfe=nfe, ncp=1, scheme='LAGRANGE-RADAU')

    sim.initialize_model()


def plot_optimal_solution(m):
    SolverFactory('ipopt').solve(m, tee=True)

    x = []
    u = []
    F = []

    for ii in m.t:
        x.append(value(m.x[ii]))
        u.append(value(m.u[ii]))
        F.append(value(m.F[ii]))

    plt.subplot(131)
    plt.plot(m.t.value, x, 'ro', label='x')
    plt.title('State Soln')
    plt.xlabel('time')

    plt.subplot(132)
    plt.plot(m.t.value, u, 'ro', label='u')
    plt.title('Control Soln')
    plt.xlabel('time')

    plt.subplot(133)
    plt.plot(m.t.value, F, 'ro', label='Cost Integrand')
    plt.title('Anti-derivative of \n Cost Integrand')
    plt.xlabel('time')

    # plt.show()
    return plt


######################################

if __name__ == '__main__':
    m = create_model()
    initialize_model(m, 100)

    #    plt = plot_optimal_solution(m)
    #    plt.show()

    m.perturbed_a = Param(initialize=-0.25)
    m.perturbed_H = Param(initialize=0.55)
    m_sipopt = sensitivity_calculation(
        'sipopt',
        m,
        [m.a, m.H],
        [m.perturbed_a, m.perturbed_H],
        cloneModel=True,
        tee=True,
    )

    for var, val in m_sipopt.sens_sol_state_1.items():
        # To load updated variable values back into the model:
        if var.ctype is not Var:
            continue
        var.set_value(val)

    m_sipopt.a.set_value(value(m_sipopt.perturbed_a))
    m_sipopt.H.set_value(value(m_sipopt.perturbed_H))

    # To solve for the "true solution" (with the full nonlinear
    # model) after perturbing parameters:
    solver = SolverFactory('ipopt')
    solver.solve(m_sipopt)
