#-----------------------------------------------------------------------
#
# Minimise    f
#
# subject to      -0.5 <= x1 <=  INFTY
#                   -2 <= x2 <=  INFTY
#                    0 <= x3 <=  2
#                   -2 <= x4 <=  2
#                         g1 ==  1
#               -INFTY <= g2 <= -1
#                  2.5 <= g3 <=  5
#
# where         f (x1,x2,x3,x4) = x1^2 + 2 x2^2 - x3
#               g1(x1,x2,x3,x4) = x1^2 + x3^2 + x1x3
#               g2(x1,x2,x3,x4) = x3 - x4
#               g3(x1,x2,x3,x4) = x2 + x4
#
# Optimal solution
#                     x*  = (0, 0.5, 1, 2)
#                   f(x*) = -0.5
#                   g(x*) = (1, -1, 2.5)
#
#-----------------------------------------------------------------------

import worhp

# Objective function
def user_f(opt, wsp, par, cnt):
opt.f = wsp.scale_obj * (opt.x ** 2 + 2.0 * opt.x ** 2 - opt.x)

# Function of constraints
def user_g(opt, wsp, par, cnt):
opt.g = opt.x ** 2 + opt.x ** 2 + opt.x * opt.x
opt.g = opt.x - opt.x
opt.g = opt.x + opt.x

def user_df(opt, wsp, par, cnt):
wsp.df.val = wsp.scale_obj * 2.0 * opt.x
wsp.df.val = wsp.scale_obj * 4.0 * opt.x
wsp.df.val = wsp.scale_obj * -1.0

# Jacobian of constraints
def user_dg(opt, wsp, par, cnt):
wsp.dg.val = 2.0 * opt.x + opt.x
wsp.dg.val = 1.0
wsp.dg.val = opt.x + 2.0 * opt.x
wsp.dg.val = 1.0
wsp.dg.val = -1.0
wsp.dg.val = 1.0

# Hessian of Lagrangian
def user_hm(opt, wsp, par, cnt):
# Only scale the F part of HM
wsp.hm.val = opt.µ
wsp.hm.val = wsp.scale_obj * 2.0 + 2.0 * opt.µ
wsp.hm.val = wsp.scale_obj * 4.0
wsp.hm.val = 2.0 * opt.µ
wsp.hm.val = 0.0

# Check version of library and Python module
if worhp.check_version(worhp.MAJOR, worhp.MINOR, worhp.PATCH):
exit(1)

#
# WORHP data structures
#
# worhp.OptVar contains the variables, constraint values and multipliers.
# worhp.Workspace encapsulates all workspace and working variables for WORHP.
# worhp.Params contains all WORHP parameters.
# worhp.Control contains things for reverse communication flow control.
#
opt = worhp.OptVar()
wsp = worhp.Workspace()
par = worhp.Params()
cnt = worhp.Control()

# Properly zeros everything, or else the following routines could get confused.
worhp.pre_init(opt, wsp, par, cnt)

# worhp.diag(opt, wsp, par, cnt)

# Parameter initialisation routine that must be called when using
worhp.init_params(par)

# We can now set parameters that may be overruled by those in the parameter file.
# This is useful for setting a non-default standard parameter value that may still
# be overwritten.

par.NLPprint = 1  # Let's prefer the slim output format unless the parameter file
# says differently

# Parameter XML import routine that does not reset all parameters to
# default values (init_params does this)
if status == worhp.DATA_ERROR or status == worhp.INIT_ERROR:
exit(1)

#
# WORHP data structure initialisation routine.
# Calling this routine prior to WORHP is mandatory.
# Before calling worhp.init, set the problem and matrix dimensions as
#
# opt.n      = number of variables,
# opt.m      = number of constraints (lin + nonlin, excluding box con's),
# wsp.df.nnz = nonzero entries of the objective function gradient,
# wsp.dg.nnz = nonzero entries of the constraint Jacobian,
# wsp.hm.nnz = nonzero entries of the Lagrange Hessian.
#
# Set nnz to 'worhp.MATRIX_INIT_DENSE' to have worhp.init allocate and
# create a dense matrix structure appropriate for the matrix kind and
# its dimensions. Setting it to its dense dimension achieves the same.
#
opt.n = 4  # This problem has 4 variables
opt.m = 3  # and 3 constraints (excluding box constraints)

# All derivatives for this problem have a sparse structure, so
# set the amount of nonzeros here
wsp.df.nnz = 3
wsp.dg.nnz = 6
wsp.hm.nnz = 1 + opt.n  # 1 entry on strict lower triangle
# plus full diagonal

worhp.init(opt, wsp, par, cnt)
if cnt.status != worhp.FIRST_CALL:
print("Main: Initialisation failed.")
exit(1)

#
#
# opt.x[0:opt.n - 1]  : Optimisation variables
# opt.λ[0:opt.n - 1]  : multipliers for the constraints
#                       on x ("box constraints")
# opt.g[0:opt.m - 1]  : Linear and nonlinear constraints
# opt.µ[0:opt.m - 1]  : multipliers for the constraints on G
#
# Set initial values of x, λ and µ here.
# G need not be initialised.
#
opt.x = 2.0
opt.x = 2.0
opt.x = 1.0
opt.x = 0.0
opt.λ = 0.0  # alternatively you can use opt.Lambda
opt.λ = 0.0
opt.λ = 0.0
opt.λ = 0.0
opt.µ = 0.0  # alternatively you can use opt.Mu
opt.µ = 0.0
opt.µ = 0.0

#
# Set lower and upper bounds on the variables and constraints.
# Use +/-par.infty to signal "unbounded".
#
# xl and xu are lower and upper bounds ("box constraints") on x.
# gl and gu are lower and upper bounds on g.
#
opt.xl = -0.5
opt.xu =  par.infty
opt.xl = -2.0
opt.xu =  par.infty
opt.xl =  0.0
opt.xu =  2.0
opt.xl = -2.0
opt.xu =  2.0

opt.gl =  1.0  # set opt.gl[i] == opt.gu[i]
opt.gu =  1.0  # for equality constraints
opt.gl = -par.infty
opt.gu = -1.0
opt.gl =  2.5
opt.gu =  5.0

#
# Specify matrix structures in CS format, using Fortran indexing,
# i.e. 1...N instead of 0...N-1, to describe the matrix structure.
#

#
# Define df as row vector, column index is ommited
#
if wsp.df.need_structure:
# only set the nonzero entries, so omit the 4th entry,
# which is a structural zero
wsp.df.row = 1
wsp.df.row = 2
wsp.df.row = 3

#
# Define dg as CS-matrix
#
if wsp.dg.need_structure:
# only set the nonzero entries in column-major order */
wsp.dg.row = 1
wsp.dg.col = 1

wsp.dg.row = 3
wsp.dg.col = 2

wsp.dg.row = 1
wsp.dg.col = 3

wsp.dg.row = 2
wsp.dg.col = 3

wsp.dg.row = 2
wsp.dg.col = 4

wsp.dg.row = 3
wsp.dg.col = 4

#
# Define hm as a diagonal LT-CS-matrix, but only if needed by WORHP
#
if wsp.hm.need_structure:
# only set the nonzero entries, with the strictly lower triangle first,
# followed by ALL diagonal entries, so even the entry at (4, 4)
# even though it is a structural zero

# strict lower triangle
wsp.hm.row = 3
wsp.hm.col = 1

# diagonal
for i in range(0, opt.n):
wsp.hm.row[wsp.hm.nnz - opt.n + i] = i + 1
wsp.hm.col[wsp.hm.nnz - opt.n + i] = i + 1

#
# WORHP Reverse Communication loop.
# In every iteration poll get_user_action for the requested action, i.e. one
# of {CALL_WORHP, ITER_OUTPUT, EVAL_F, EVAL_G, EVAL_DF, EVAL_DG, EVAL_HM, FIDIF}.
#
# Make sure to reset the requested user action afterwards by calling
# done_user_action, except for 'CALL_WORHP' and 'FIDIF'.
#
while cnt.status < worhp.TERMINATE_SUCCESS and cnt.status > worhp.TERMINATE_ERROR:
#
# WORHP's main routine.
# Do not manually reset CALL_WORHP, this is only done by the FD routines.
#
if worhp.get_user_action(cnt, worhp.Action.CALL_WORHP):
worhp.worhp(opt, wsp, par, cnt)
# No done_user_action!

#
# Show iteration output.
# The call to iteration_output may be replaced by user-defined code.
#
if worhp.get_user_action(cnt, worhp.Action.ITER_OUTPUT):
worhp.iteration_output(opt, wsp, par, cnt)
worhp.done_user_action(cnt, worhp.Action.ITER_OUTPUT)

#
# Evaluate the objective function.
# The call to user_f may be replaced by user-defined code.
#
if worhp.get_user_action(cnt, worhp.Action.EVAL_F):
user_f(opt, wsp, par, cnt)
worhp.done_user_action(cnt, worhp.Action.EVAL_F)

#
# Evaluate the constraints.
# The call to user_g may be replaced by user-defined code.
#
if worhp.get_user_action(cnt, worhp.Action.EVAL_G):
user_g(opt, wsp, par, cnt)
worhp.done_user_action(cnt, worhp.Action.EVAL_G)

#
# Evaluate the gradient of the objective function.
# The call to user_df may be replaced by user-defined code.
#
if worhp.get_user_action(cnt, worhp.Action.EVAL_DF):
user_df(opt, wsp, par, cnt)
worhp.done_user_action(cnt, worhp.Action.EVAL_DF)

#
# Evaluate the Jacobian of the constraints.
# The call to user_dg may be replaced by user-defined code.
#
if worhp.get_user_action(cnt, worhp.Action.EVAL_DG):
user_dg(opt, wsp, par, cnt)
worhp.done_user_action(cnt, worhp.Action.EVAL_DG)

#
# Evaluate the Hessian matrix of the Lagrange function (L = f + µ*g)
# The call to user_hm may be replaced by user-defined code.
#
if worhp.get_user_action(cnt, worhp.Action.EVAL_HM):
user_hm(opt, wsp, par, cnt)
worhp.done_user_action(cnt, worhp.Action.EVAL_HM)

#
# Use finite differences with RC to determine derivatives
# Do not reset fidif, this is done by the FD routine.
#
if worhp.get_user_action(cnt, worhp.Action.FIDIF):
worhp.fidif(opt, wsp, par, cnt)
# No done_user_action!

#
# Translate the WORHP status flag into a meaningful message.
#
worhp.status_msg(opt, wsp, par, cnt)
#
# Deallocate all data structures.
# Data structures must not be accessed after this call.
#
worhp.free(opt, wsp, par, cnt)