#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright (c) 2008-2022
#  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.
#  ___________________________________________________________________________

from pyomo.contrib.pynumero.interfaces.pyomo_nlp import PyomoNLP
from pyomo.contrib.pynumero.interfaces.utils import (build_bounds_mask,
                                                     build_compression_matrix,
                                                     full_to_compressed)
import pyomo.environ as pyo
import numpy as np


def create_basic_model():

    m = pyo.ConcreteModel()
    m.x = pyo.Var([1, 2, 3], domain=pyo.Reals)
    for i in range(1, 4):
        m.x[i].value = i
    m.c1 = pyo.Constraint(expr=m.x[1] ** 2 - m.x[2] - 1 == 0)
    m.c2 = pyo.Constraint(expr=m.x[1] - m.x[3] - 0.5 == 0)
    m.d1 = pyo.Constraint(expr=m.x[1] + m.x[2] <= 100.0)
    m.d2 = pyo.Constraint(expr=m.x[2] + m.x[3] >= -100.0)
    m.d3 = pyo.Constraint(expr=m.x[2] + m.x[3] + m.x[1] >= -500.0)
    m.x[2].setlb(0.0)
    m.x[3].setlb(0.0)
    m.x[2].setub(100.0)
    m.obj = pyo.Objective(expr=m.x[2]**2)
    return m


def main():
    model = create_basic_model()
    solver = pyo.SolverFactory('ipopt')
    solver.solve(model, tee=True)

    # build nlp initialized at the solution
    nlp = PyomoNLP(model)

    # get initial point
    print(nlp.primals_names())
    x0 = nlp.get_primals()

    # vectors of lower and upper bounds
    xl = nlp.primals_lb()
    xu = nlp.primals_ub()

    # demonstrate use of compression from full set of bounds
    # to only finite bounds using masks
    xlb_mask = build_bounds_mask(xl)
    xub_mask = build_bounds_mask(xu)
    # get the compressed vector
    compressed_xl = full_to_compressed(xl, xlb_mask)
    compressed_xu = full_to_compressed(xu, xub_mask)
    # we can also build compression matrices
    Cx_xl = build_compression_matrix(xlb_mask)
    Cx_xu = build_compression_matrix(xub_mask)

    # lower and upper bounds residual
    res_xl = Cx_xl * x0 - compressed_xl
    res_xu = compressed_xu - Cx_xu * x0
    print("Residuals lower bounds x-xl:", res_xl)
    print("Residuals upper bounds xu-x:", res_xu)

    # set the value of the primals (we can skip the duals)
    # here we set them to the initial values, but we could
    # set them to anything
    nlp.set_primals(x0)

    # evaluate residual of equality constraints
    print(nlp.constraint_names())
    res_eq = nlp.evaluate_eq_constraints()
    print("Residuals of equality constraints:", res_eq)

    # evaluate residual of inequality constraints
    res_ineq = nlp.evaluate_ineq_constraints()

    # demonstrate the use of compression from full set of
    # lower and upper bounds on the inequality constraints
    # to only the finite values using masks
    ineqlb_mask = build_bounds_mask(nlp.ineq_lb())
    inequb_mask = build_bounds_mask(nlp.ineq_ub())
    # get the compressed vector
    compressed_ineq_lb = full_to_compressed(nlp.ineq_lb(), ineqlb_mask)
    compressed_ineq_ub = full_to_compressed(nlp.ineq_ub(), inequb_mask)
    # we can also build compression matrices
    Cineq_ineqlb = build_compression_matrix(ineqlb_mask)
    Cineq_inequb = build_compression_matrix(inequb_mask)

    # lower and upper inequalities residual
    res_ineq_lb = Cineq_ineqlb * res_ineq - compressed_ineq_lb
    res_ineq_ub = compressed_ineq_ub - Cineq_inequb*res_ineq
    print("Residuals of inequality constraints lower bounds:", res_ineq_lb)
    print("Residuals of inequality constraints upper bounds:", res_ineq_ub)

    feasible = False
    if np.all(res_xl >= 0) and np.all(res_xu >= 0) \
        and np.all(res_ineq_lb >= 0) and np.all(res_ineq_ub >= 0) and \
        np.allclose(res_eq, np.zeros(nlp.n_eq_constraints()), atol=1e-5):
        feasible = True

    print("Is x0 feasible:", feasible)

    return feasible


if __name__ == '__main__':
    main()
