text:
# Callbacks utilities

"""
Structure specifying the callback context.

Each evaluation callbacks (for objective, gradient or hessian)
is attached to a unique callback context.

"""
mutable struct CallbackContext
    context::Ptr{Cvoid}
    # we have to keep a reference to the optimization model
    model::Model
    # we sometime need some user params
    userparams::Dict

    # oracle's callbacks are context dependent, so we store
    # them inside dedicated CallbackContext
    eval_f::Function
    eval_g::Function
    eval_h::Function
    eval_rsd::Function
    eval_jac_rsd::Function

    function CallbackContext(ptr::Ptr{Cvoid}, model::Model)
        return new(ptr, model, Dict())
    end
end

##################################################
# Utils
##################################################
function KN_set_cb_user_params(m::Model, cb::CallbackContext, userParams=nothing)
    # we have to store the model in the C model to use callbacks
    # function. So, we store the userParams inside the model
    if userParams != nothing
        cb.userparams[:data] = userParams
    end
    # we store the callback context inside KNITRO user data
    c_userdata = cb

    ret = @kn_ccall(set_cb_user_params, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Any),
                    m.env, cb.context, c_userdata)
    _checkraise(ret)
end

"""
Specify which gradient option `gradopt` will be used to evaluate
the first derivatives of the callback functions.  If `gradopt=KN_GRADOPT_EXACT`
then a gradient evaluation callback must be set by `KN_set_cb_grad()`
(or `KN_set_cb_rsd_jac()` for least squares).

"""
function KN_set_cb_gradopt(m::Model, cb::CallbackContext, gradopt::Integer)
    ret = @kn_ccall(set_cb_gradopt, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Cint),
                    m.env, cb.context, gradopt)
    _checkraise(ret)
end

#--------------------
# Set relative step size
#--------------------
"""
Set an array of relative stepsizes to use for the finite-difference
gradient/Jacobian computations when using finite-difference
first derivatives.  Finite-difference step sizes `delta` in Knitro are
computed as:

     delta[i] = relStepSizes[i]*max(abs(x[i]),1);

The default relative step sizes for each component of `x` are sqrt(eps)
for forward finite differences, and eps^(1/3) for central finite
differences.  Use this function to overwrite the default values.

"""
function KN_set_cb_relstepsizes(m::Model, cb::CallbackContext, nindex::Integer, xRelStepSize::Cdouble)
    ret = @kn_ccall(set_cb_relstepsize, Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Cint, Cdouble),
                    m.env, cb.context, nindex, xRelStepSize)
    _checkraise(ret)
end

function KN_set_cb_relstepsizes(m::Model, cb::CallbackContext, xIndex::Vector{Cint}, xRelStepSizes::Vector{Cdouble})
    ncon = length(xIndex)
    @assert length(xRelStepSizes) == ncon
    ret = @kn_ccall(set_cb_relstepsizes, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Cint, Ptr{Cint}, Ptr{Cdouble}),
                    m.env, cb.context, ncon, xIndex, xRelStepSizes)
    _checkraise(ret)
end

function KN_set_cb_relstepsizes(m::Model, cb::CallbackContext, xRelStepSizes::Vector{Cdouble})
    ret = @kn_ccall(set_cb_relstepsizes_all, Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cdouble}),
                    m.env, cb.context, xRelStepSizes)
    _checkraise(ret)
end

#--------------------
# callback context getters
#--------------------
# TODO: dry this code with a macro
function KN_get_cb_number_cons(m::Model, cb::Ptr{Cvoid})
    num = Cint[0]
    ret = @kn_ccall(get_cb_number_cons,
                    Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}),
                    m.env, cb, num)
    _checkraise(ret)
    return num[1]
end

function KN_get_cb_objgrad_nnz(m::Model, cb::Ptr{Cvoid})
    num = Cint[0]
    ret = @kn_ccall(get_cb_objgrad_nnz,
                    Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}),
                    m.env, cb, num)
    _checkraise(ret)
    return num[1]
end

function KN_get_cb_jacobian_nnz(m::Model, cb::Ptr{Cvoid})
    num = Cint[0]
    ret = @kn_ccall(get_cb_jacobian_nnz,
                    Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}),
                    m.env, cb, num)
    _checkraise(ret)
    return num[1]
end

function KN_get_cb_hessian_nnz(m::Model, cb::Ptr{Cvoid})
    num = Cint[0]
    ret = @kn_ccall(get_cb_hessian_nnz,
                    Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}),
                    m.env, cb, num)
    _checkraise(ret)
    return num[1]
end

function KN_get_cb_number_rsds(m::Model, cb::Ptr{Cvoid})
    num = Cint[0]
    ret = @kn_ccall(get_cb_number_rsds,
                    Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}),
                    m.env, cb, num)
    _checkraise(ret)
    return num[1]
end

function KN_get_cb_rsd_jacobian_nnz(m::Model, cb::Ptr{Cvoid})
    num = Cint[0]
    ret = @kn_ccall(get_cb_rsd_jacobian_nnz,
                    Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}),
                    m.env, cb, num)
    _checkraise(ret)
    return num[1]
end

##################################################
# EvalRequest

# low level EvalRequest structure
"""
Low-level structure used to pass back evaluation information for evaluation
callbacks.

# Attributes
 * `evalRequestCode`: indicates the type of evaluation requested
 * `threadID`: the thread ID associated with this evaluation request;
                useful for multi-threaded, concurrent evaluations
 * `x`: values of unknown (primal) variables used for all evaluations
 * `lambda`: values of unknown dual variables/Lagrange multipliers
                     used for the evaluation of the Hessian
 * `sigma`: scalar multiplier for the objective component of the Hessian
 * `vec`: vector array value for Hessian-vector products (only used
          when user option hessopt=KN_HESSOPT_PRODUCT)

"""
mutable struct KN_eval_request
    evalRequestCode::Cint
    threadID::Cint
    x::Ptr{Cdouble}
    lambda::Ptr{Cdouble}
    sigma::Ptr{Cdouble}
    vec::Ptr{Cdouble}
end

# high level EvalRequest structure
"High level structure to wrap `KN_eval_request`."
mutable struct EvalRequest
    evalRequestCode::Cint
    threadID::Cint
    x::Array{Float64}
    lambda::Array{Float64}
    sigma::Float64
    vec::Array{Float64}
end

# Import low level request to Julia object
function EvalRequest(m::Model, evalRequest_::KN_eval_request)
    nx = KN_get_number_vars(m)
    nc = KN_get_number_cons(m)

    # import scaling
    sigma = (evalRequest_.sigma != C_NULL) ? unsafe_wrap(Array, evalRequest_.sigma, 1)[1] : 1.
    # we use only views to C arrays to avoid unnecessary copy
    return EvalRequest(evalRequest_.evalRequestCode,
                       evalRequest_.threadID,
                       unsafe_wrap(Array, evalRequest_.x, nx),
                       unsafe_wrap(Array, evalRequest_.lambda, nx + nc),
                       sigma,
                       unsafe_wrap(Array, evalRequest_.vec, nx))
end

##################################################
# EvalResult

# low level EvalResult structure
"""
Low-level structure used to return results information for evaluation callbacks.
The arrays (and their indices and sizes) returned in this structure are
local to the specific callback structure used for the evaluation.

# Attributes
* `obj`: objective function evaluated at `x` for EVALFC or
                     EVALFCGA request (funcCallback)
* `c`: (length nC) constraint values evaluated at `x` for
                     EVALFC or EVALFCGA request (funcCallback)
* `objGrad`: (length nV) objective gradient evaluated at `x` for
                     EVALGA request (gradCallback) or EVALFCGA request (funcCallback)
* `jac`: (length nnzJ) constraint Jacobian evaluated at `x` for
                     EVALGA request (gradCallback) or EVALFCGA request (funcCallback)
* `hess`: (length nnzH) Hessian evaluated at `x`, `lambda`, `sigma`
                     for EVALH or EVALH_NO_F request (hessCallback)
* `hessVec`: (length n=number variables in the model) Hessian-vector
                     product evaluated at `x`, `lambda`, `sigma`
                     for EVALHV or EVALHV_NO_F request (hessCallback)
* `rsd`: (length nR) residual values evaluated at `x` for EVALR
                     request (rsdCallback)
* `rsdJac`: (length nnzJ) residual Jacobian evaluated at `x` for
                     EVALRJ request (rsdJacCallback)

"""
mutable struct KN_eval_result
    obj::Ptr{Cdouble}
    c::Ptr{Cdouble}
    objGrad::Ptr{Cdouble}
    jac::Ptr{Cdouble}
    hess::Ptr{Cdouble}
    hessVec::Ptr{Cdouble}
    rsd::Ptr{Cdouble}
    rsdJac::Ptr{Cdouble}
end

# high level EvalRequest structure
mutable struct EvalResult
    obj::Array{Float64}
    c::Array{Float64}
    objGrad::Array{Float64}
    jac::Array{Float64}
    hess::Array{Float64}
    hessVec::Array{Float64}
    rsd::Array{Float64}
    rsdJac::Array{Float64}
end

function EvalResult(m::Model, cb::Ptr{Cvoid}, evalResult_::KN_eval_result)
    return EvalResult(
            unsafe_wrap(Array, evalResult_.obj, 1),
            unsafe_wrap(Array, evalResult_.c, KN_get_cb_number_cons(m, cb)),
            unsafe_wrap(Array, evalResult_.objGrad, KN_get_cb_objgrad_nnz(m, cb)),
            unsafe_wrap(Array, evalResult_.jac, KN_get_cb_jacobian_nnz(m, cb)),
            unsafe_wrap(Array, evalResult_.hess, KN_get_cb_hessian_nnz(m, cb)),
            unsafe_wrap(Array, evalResult_.hessVec, KN_get_number_vars(m)),
            unsafe_wrap(Array, evalResult_.rsd, KN_get_cb_number_rsds(m, cb)),
            unsafe_wrap(Array, evalResult_.rsdJac, KN_get_cb_rsd_jacobian_nnz(m, cb))
                     )
end

##################################################
# Callbacks wrappers
# 1/ for eval function
function eval_fc_wrapper(ptr_model::Ptr{Cvoid}, ptr_cb::Ptr{Cvoid},
                         evalRequest_::Ptr{Cvoid},
                         evalResults_::Ptr{Cvoid},
                         userdata_::Ptr{Cvoid})
    # load evalRequest object
    ptr0 = Ptr{KN_eval_request}(evalRequest_)
    evalRequest = unsafe_load(ptr0)::KN_eval_request

    # load evalResult object
    ptr = Ptr{KN_eval_result}(evalResults_)
    evalResult = unsafe_load(ptr)::KN_eval_result

    # and eventually, load callback context
    cb = unsafe_pointer_to_objref(userdata_)
    # we have to ensure that cb is a CallbackContext
    # otherwise, we tell KNITRO that a problem occurs by returning a
    # non-zero status
    if !isa(cb, CallbackContext)
        return Cint(32)
    end

    request = EvalRequest(cb.model, evalRequest)
    result = EvalResult(cb.model, ptr_cb, evalResult)

    cb.eval_f(ptr_model, ptr_cb, request, result, cb.userparams)

    return Cint(0)
end

# 2/ for gradient function
function eval_ga_wrapper(ptr_model::Ptr{Cvoid}, ptr_cb::Ptr{Cvoid},
                         evalRequest_::Ptr{Cvoid},
                         evalResults_::Ptr{Cvoid},
                         userdata_::Ptr{Cvoid})

    # load evalRequest object
    ptr0 = Ptr{KN_eval_request}(evalRequest_)
    evalRequest = unsafe_load(ptr0)::KN_eval_request

    # load evalResult object
    ptr = Ptr{KN_eval_result}(evalResults_)
    evalResult = unsafe_load(ptr)::KN_eval_result

    # and eventually, load callback context
    cb = unsafe_pointer_to_objref(userdata_)
    # we have to ensure that cb is a CallbackContext
    # otherwise, we tell KNITRO that a problem occurs by returning a
    # non-zero status
    if !isa(cb, CallbackContext)
        return Cint(32)
    end

    request = EvalRequest(cb.model, evalRequest)
    result = EvalResult(cb.model, ptr_cb, evalResult)

    cb.eval_g(ptr_model, ptr_cb, request, result, cb.userparams)

    return Cint(0)
end

# 3/ for hessian function
function eval_hess_wrapper(ptr_model::Ptr{Cvoid}, ptr_cb::Ptr{Cvoid},
                           evalRequest_::Ptr{Cvoid},
                           evalResults_::Ptr{Cvoid},
                           userdata_::Ptr{Cvoid})

    # load evalRequest object
    ptr0 = Ptr{KN_eval_request}(evalRequest_)
    evalRequest = unsafe_load(ptr0)::KN_eval_request

    # load evalResult object
    ptr = Ptr{KN_eval_result}(evalResults_)
    evalResult = unsafe_load(ptr)::KN_eval_result

    # and eventually, load callback context
    cb = unsafe_pointer_to_objref(userdata_)
    # we have to ensure that cb is a CallbackContext
    # otherwise, we tell KNITRO that a problem occurs by returning a
    # non-zero status
    if !isa(cb, CallbackContext)
        return Cint(32)
    end

    request = EvalRequest(cb.model, evalRequest)
    result = EvalResult(cb.model, ptr_cb, evalResult)

    cb.eval_h(ptr_model, ptr_cb, request, result, cb.userparams)

    return Cint(0)
end

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

# User callback should be of the form:
# callback(kc, cb, evalrequest, evalresult, usrparams)::Int

"""
    KN_add_eval_callback(m::Model, funccallback::Function)
    KN_add_eval_callback(m::Model, evalObj::Bool, indexCons::Vector{Cint},
                         funccallback::Function)

This is the routine for adding a callback for (nonlinear) evaluations
of objective and constraint functions.  This routine can be called
multiple times to add more than one callback structure (e.g. to create
different callback structures to handle different blocks of constraints).
This routine specifies the minimal information needed for a callback, and
creates the callback structure `cb`, which can then be passed to other
callback functions to set additional information for that callback.

# Parameters
* `evalObj`: boolean indicating whether or not any part of the objective
                  function is evaluated in the callback
* `indexCons`: (length nC) index of constraints evaluated in the callback
                  (set to NULL if nC=0)
* `funcCallback`: a function that evaluates the objective parts
                  (if evalObj=KNTRUE) and any constraint parts (specified by
                  nC and indexCons) involved in this callback; when
                  eval_fcga=KN_EVAL_FCGA_YES, this callback should also evaluate
                  the relevant first derivatives/gradients

# Returns
* `cb`: the callback structure that gets created by
                  calling this function; all the memory for this structure is
                  handled by Knitro

After a callback is created by `KN_add_eval_callback()`, the user can then specify
gradient information and structure through `KN_set_cb_grad()` and Hessian
information and structure through `KN_set_cb_hess()`.  If not set, Knitro will
approximate these.  However, it is highly recommended to provide a callback routine
to specify the gradients if at all possible as this will greatly improve the
performance of Knitro.  Even if a gradient callback is not provided, it is still
helpful to provide the sparse Jacobian structure through `KN_set_cb_grad()` to
improve the efficiency of the finite-difference gradient approximations.
Other optional information can also be set via `KN_set_cb_*()` functions as
detailed below.

"""
function KN_add_eval_callback(m::Model, funccallback::Function)
    # wrap eval_callback_wrapper as C function
    c_f = @cfunction(eval_fc_wrapper, Cint,
                     (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}))

    # define callback context
    rfptr = Ref{Ptr{Cvoid}}()

    # add callback to context
    ret = @kn_ccall(add_eval_callback_all, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}),
                    m.env, c_f, rfptr)
    _checkraise(ret)
    cb = CallbackContext(rfptr.x, m)

    # store function in callback environment:
    cb.eval_f = funccallback

    # store model in user params to access callback in C
    KN_set_cb_user_params(m, cb)

    return cb
end

function KN_add_eval_callback(m::Model, evalObj::Bool, indexCons::Vector{Cint},
                              funccallback::Function)
    nC = length(indexCons)

    # wrap eval_callback_wrapper as C function
    c_f = @cfunction(eval_fc_wrapper, Cint,
                     (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}))

    # define callback context
    rfptr = Ref{Ptr{Cvoid}}()

    # add callback to context
    ret = @kn_ccall(add_eval_callback, Cint,
                    (Ptr{Cvoid}, KNBOOL, Cint, Ptr{Cint}, Ptr{Cvoid}, Ptr{Cvoid}),
                    m.env, KNBOOL(evalObj), nC, indexCons, c_f, rfptr)
    _checkraise(ret)
    cb = CallbackContext(rfptr.x, m)

    # store function in callback environment:
    cb.eval_f = funccallback

    KN_set_cb_user_params(m, cb)

    return cb
end

"""
    KN_set_cb_grad(m::Model, cb::CallbackContext, gradcallback;
                   nV::Integer=KN_DENSE, objGradIndexVars=C_NULL,
                   jacIndexCons=C_NULL, jacIndexVars=C_NULL)

This API function is used to set the objective gradient and constraint Jacobian
structure and also (optionally) a callback function to evaluate the objective
gradient and constraint Jacobian provided through this callback.

"""
function KN_set_cb_grad(m::Model, cb::CallbackContext, gradcallback;
                        nV::Integer=KN_DENSE, objGradIndexVars=C_NULL,
                        jacIndexCons=C_NULL, jacIndexVars=C_NULL)
    # check consistency of arguments
    if (nV == 0 || nV == KN_DENSE )
        (objGradIndexVars != C_NULL) && error("objGradIndexVars must be set to C_NULL when nV = $nV")
    else
        @assert (objGradIndexVars != C_NULL) && (length(objGradIndexVars) == nV)
    end

    nnzJ = KNLONG(0)
    if jacIndexCons != C_NULL && jacIndexVars != C_NULL
        @assert length(jacIndexCons) == length(jacIndexVars)
        nnzJ = KNLONG(length(jacIndexCons))
    end

    if gradcallback != nothing
        # store grad function inside model:
        cb.eval_g = gradcallback

        # wrap gradient wrapper as C function
        c_grad_g = @cfunction(eval_ga_wrapper, Cint,
                              (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}))
    else
        c_grad_g = C_NULL
    end

    ret = @kn_ccall(set_cb_grad, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Cint, Ptr{Cint},
                     KNLONG, Ptr{Cint}, Ptr{Cint}, Ptr{Cvoid}),
                    m.env, cb.context, nV,
                    objGradIndexVars, nnzJ, jacIndexCons, jacIndexVars,
                    c_grad_g)
    _checkraise(ret)

    return nothing
end

"""
    KN_set_cb_hess(m::Model, cb::CallbackContext, nnzH::Integer, hesscallback::Function;
                   hessIndexVars1=C_NULL, hessIndexVars2=C_NULL)

This API function is used to set the structure and a callback function to
evaluate the components of the Hessian of the Lagrangian provided through this
callback.  KN_set_cb_hess() should only be used when defining a user-supplied
Hessian callback function (via the `hessopt=KN_HESSOPT_EXACT` user option).
When Knitro is approximating the Hessian, it cannot make use of the Hessian
sparsity structure.

"""
function KN_set_cb_hess(m::Model, cb::CallbackContext, nnzH::Integer, hesscallback::Function;
                        hessIndexVars1=C_NULL, hessIndexVars2=C_NULL)

    if nnzH == KN_DENSE_ROWMAJOR || nnzH == KN_DENSE_COLMAJOR
        @assert hessIndexVars1 == hessIndexVars2 == C_NULL
    else
        @assert hessIndexVars1 != C_NULL && hessIndexVars2 != C_NULL
        @assert length(hessIndexVars1) == length(hessIndexVars2) == nnzH
    end
    # store hessian function inside model:
    cb.eval_h = hesscallback

    # wrap gradient wrapper as C function
    c_hess = @cfunction(eval_hess_wrapper, Cint,
                        (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}))

    ret = @kn_ccall(set_cb_hess, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, KNLONG, Ptr{Cint}, Ptr{Cint}, Ptr{Cvoid}),
                    m.env, cb.context, nnzH,
                    hessIndexVars1, hessIndexVars2, c_hess)
    _checkraise(ret)

    return nothing
end

##################################################
# Get callbacks info
##################################################
function KN_get_number_FC_evals(m::Model)
    fc_eval = Int32[0]
    ret = @kn_ccall(get_number_FC_evals, Cint, (Ptr{Cvoid}, Ptr{Cint}),
                    m.env, fc_eval)
    _checkraise(ret)
    return fc_eval[1]
end

function KN_get_number_GA_evals(m::Model)
    fc_eval = Int32[0]
    ret = @kn_ccall(get_number_GA_evals, Cint, (Ptr{Cvoid}, Ptr{Cint}),
                    m.env, fc_eval)
    _checkraise(ret)
    return fc_eval[1]
end

function KN_get_number_H_evals(m::Model)
    fc_eval = Int32[0]
    ret = @kn_ccall(get_number_H_evals, Cint, (Ptr{Cvoid}, Ptr{Cint}),
                    m.env, fc_eval)
    _checkraise(ret)
    return fc_eval[1]
end

function KN_get_number_HV_evals(m::Model)
    fc_eval = Int32[0]
    ret = @kn_ccall(get_number_HV_evals, Cint, (Ptr{Cvoid}, Ptr{Cint}),
                    m.env, fc_eval)
    _checkraise(ret)
    return fc_eval[1]
end

##################################################
# Residual callbacks
##################################################
function eval_rsd_wrapper(ptr_model::Ptr{Cvoid}, ptr_cb::Ptr{Cvoid},
                          evalRequest_::Ptr{Cvoid},
                          evalResults_::Ptr{Cvoid},
                          userdata_::Ptr{Cvoid})

    # load evalRequest object
    ptr0 = Ptr{KN_eval_request}(evalRequest_)
    evalRequest = unsafe_load(ptr0)::KN_eval_request

    # load evalResult object
    ptr = Ptr{KN_eval_result}(evalResults_)
    evalResult = unsafe_load(ptr)::KN_eval_result

    # and eventually, load callback context
    cb = unsafe_pointer_to_objref(userdata_)
    # we have to ensure that cb is a CallbackContext
    # otherwise, we tell KNITRO that a problem occurs by returning a
    # non-zero status
    if !isa(cb, CallbackContext)
        return Cint(32)
    end

    request = EvalRequest(cb.model, evalRequest)
    result = EvalResult(cb.model, ptr_cb, evalResult)

    cb.eval_rsd(ptr_model, ptr_cb, request, result, cb.userparams)

    return Cint(0)
end

# TODO: add _one and _* support
"""
    KN_add_lsq_eval_callback(m::Model, rsdCallBack::Function)

Add an evaluation callback for a least-squares models.  Similar to KN_add_eval_callback()
above, but for least-squares models.

* `m`: current KNITRO model
* `rsdCallback`: a function that evaluates any residual parts

After a callback is created by `KN_add_lsq_eval_callback()`, the user can then
specify residual Jacobian information and structure through `KN_set_cb_rsd_jac()`.
If not set, Knitro will approximate the residual Jacobian.  However, it is highly
recommended to provide a callback routine to specify the residual Jacobian if at all
possible as this will greatly improve the performance of Knitro.  Even if a callback
for the residual Jacobian is not provided, it is still helpful to provide the sparse
Jacobian structure for the residuals through `KN_set_cb_rsd_jac()` to improve the
efficiency of the finite-difference Jacobian approximation.

"""
function KN_add_lsq_eval_callback(m::Model, rsdCallBack::Function)

    # wrap eval_callback_wrapper as C function
    c_f = @cfunction(eval_rsd_wrapper, Cint,
                   (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}))

    # define callback context
    rfptr = Ref{Ptr{Cvoid}}()

    # add callback to context
    ret = @kn_ccall(add_lsq_eval_callback_all, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}),
                    m.env, c_f, rfptr)
    _checkraise(ret)
    cb = CallbackContext(rfptr.x, m)

    # store function inside model:
    cb.eval_rsd = rsdCallBack

    KN_set_cb_user_params(m, cb)

    return cb
end

function eval_rj_wrapper(ptr_model::Ptr{Cvoid}, ptr_cb::Ptr{Cvoid},
                         evalRequest_::Ptr{Cvoid},
                         evalResults_::Ptr{Cvoid},
                         userdata_::Ptr{Cvoid})

    # load evalRequest object
    ptr0 = Ptr{KN_eval_request}(evalRequest_)
    evalRequest = unsafe_load(ptr0)::KN_eval_request

    # load evalResult object
    ptr = Ptr{KN_eval_result}(evalResults_)
    evalResult = unsafe_load(ptr)::KN_eval_result

    # and eventually, load callback context
    cb = unsafe_pointer_to_objref(userdata_)
    # we have to ensure that cb is a CallbackContext
    # otherwise, we tell KNITRO that a problem occurs by returning a
    # non-zero status
    if !isa(cb, CallbackContext)
        return Cint(32)
    end

    request = EvalRequest(cb.model, evalRequest)
    result = EvalResult(cb.model, ptr_cb, evalResult)

    cb.eval_jac_rsd(ptr_model, ptr_cb, request, result, cb.userparams)

    return Cint(0)
end

function KN_set_cb_rsd_jac(m::Model, cb::CallbackContext, nnzJ::Integer, evalRJ::Function;
                           jacIndexRsds=C_NULL, jacIndexVars=C_NULL)
    # check consistency of arguments
    if nnzJ == KN_DENSE_ROWMAJOR || nnzJ == KN_DENSE_COLMAJOR || nnzJ == 0
        @assert jacIndexRsds == jacIndexVars == C_NULL
    else
        @assert jacIndexRsds != C_NULL && jacIndexVars != C_NULL
        @assert length(jacIndexRsds) == length(jacIndexVars) == nnzJ
    end

    # store function inside model:
    cb.eval_jac_rsd = evalRJ
    # wrap as C function
    c_eval_rj = @cfunction(eval_rj_wrapper, Cint,
                           (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}))
    ret = @kn_ccall(set_cb_rsd_jac, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, KNLONG, Ptr{Cint}, Ptr{Cint}, Ptr{Cvoid}),
                    m.env, cb.context, nnzJ,
                    jacIndexRsds, jacIndexVars, c_eval_rj)

    _checkraise(ret)

    return cb
end

##################################################
# User callbacks wrapper
##################################################
#--------------------
# New estimate callback
#--------------------
function newpt_wrapper(ptr_model::Ptr{Cvoid},
                       ptr_x::Ptr{Cdouble},
                       ptr_lambda::Ptr{Cdouble},
                       userdata_::Ptr{Cvoid})

    # Load KNITRO's Julia Model
    m = unsafe_pointer_to_objref(userdata_)::Model
    nx = KN_get_number_vars(m)
    nc = KN_get_number_cons(m)

    x = unsafe_wrap(Array, ptr_x, nx)
    lambda = unsafe_wrap(Array, ptr_lambda, nx + nc)
    m.user_callback(ptr_model, x, lambda, m)

    return Cint(0)
end

"""
    KN_set_newpt_callback(m::Model, callback::Function)

Set the callback function that is invoked after Knitro computes a
new estimate of the solution point (i.e., after every iteration).
The function should not modify any Knitro arguments.

Callback is a function with signature:

    callback(kc, x, lambda, userdata)

Argument `kc` passed to the callback from inside Knitro is the
context pointer for the current problem being solved inside Knitro
(either the main single-solve problem, or a subproblem when using
multi-start, Tuner, etc.).
Arguments `x` and `lambda` contain the latest solution estimates.
Other values (such as objective, constraint, jacobian, etc.) can be
queried using the corresonding KN_get_XXX_values methods.

Note: Currently only active for continuous models.

"""
function KN_set_newpt_callback(m::Model, callback::Function)
    # store callback function inside model:
    m.user_callback = callback

    # wrap user callback wrapper as C function
    c_func = @cfunction(newpt_wrapper, Cint,
                        (Ptr{Cvoid}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cvoid}))

    ret = @kn_ccall(set_newpt_callback, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Any),
                    m.env, c_func, m)
    _checkraise(ret)

    return nothing
end

#--------------------
# Multistart callback
#--------------------
function ms_process_wrapper(ptr_model::Ptr{Cvoid},
                            ptr_x::Ptr{Cdouble},
                            ptr_lambda::Ptr{Cdouble},
                            userdata_::Ptr{Cvoid})

    # Load KNITRO's Julia Model
    m = unsafe_pointer_to_objref(userdata_)::Model
    nx = KN_get_number_vars(m)
    nc = KN_get_number_cons(m)

    x = unsafe_wrap(Array, ptr_x, nx)
    lambda = unsafe_wrap(Array, ptr_lambda, nx + nc)
    m.ms_process(ptr_model, x, lambda, m)

    return Cint(0)
end

"""
    KN_set_ms_process_callback(m::Model, callback::Function)

This callback function is for multistart (MS) problems only.
Set the callback function that is invoked after Knitro finishes
processing a multistart solve.

Callback is a function with signature:

    callback(kc, x, lambda, userdata)

Argument `kc` passed to the callback
from inside Knitro is the context pointer for the last multistart
subproblem solved inside Knitro.  The function should not modify any
Knitro arguments.  Arguments `x` and `lambda` contain the solution from
the last solve.

"""
function KN_set_ms_process_callback(m::Model, callback::Function)
    # store callback function inside model:
    m.ms_process = callback

    # wrap user callback wrapper as C function
    c_func = @cfunction(ms_process_wrapper, Cint,
                        (Ptr{Cvoid}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cvoid}))

    ret = @kn_ccall(set_ms_process_callback, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Any),
                    m.env, c_func, m)
    _checkraise(ret)
end

#--------------------
# MIP callback
#--------------------
function mip_node_callback_wrapper(ptr_model::Ptr{Cvoid},
                                   ptr_x::Ptr{Cdouble},
                                   ptr_lambda::Ptr{Cdouble},
                                   userdata_::Ptr{Cvoid})

    # Load KNITRO's Julia Model
    m = unsafe_pointer_to_objref(userdata_)::Model
    nx = KN_get_number_vars(m)
    nc = KN_get_number_cons(m)

    x = unsafe_wrap(Array, ptr_x, nx)
    lambda = unsafe_wrap(Array, ptr_lambda, nx + nc)
    m.mip_callback(ptr_model, x, lambda, m)

    return Cint(0)
end

"""
    KN_set_mip_node_callback(m::Model, callback::Function)

This callback function is for mixed integer (MIP) problems only.
Set the callback function that is invoked after Knitro finishes
processing a node on the branch-and-bound tree (i.e., after a relaxed
subproblem solve in the branch-and-bound procedure).

Callback is a function with signature:

    callback(kc, x, lambda, userdata)

Argument `kc` passed to the callback from inside Knitro is the
context pointer for the last node subproblem solved inside Knitro.
The function should not modify any Knitro arguments.
Arguments `x` and `lambda` contain the solution from the node solve.

"""
function KN_set_mip_node_callback(m::Model, callback::Function)
    # store callback function inside model:
    m.mip_callback = callback

    # wrap user callback wrapper as C function
    c_func = @cfunction(mip_node_callback_wrapper, Cint,
                        (Ptr{Cvoid}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cvoid}))

    ret = @kn_ccall(set_mip_node_callback, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Any),
                    m.env, c_func, m)
    _checkraise(ret)
end

#--------------------
# Multistart procedure callback
#--------------------
function ms_initpt_wrapper(ptr_model::Ptr{Cvoid},
                           nSolveNumber::Cint,
                           ptr_x::Ptr{Cdouble},
                           ptr_lambda::Ptr{Cdouble},
                           userdata_::Ptr{Cvoid})

    # Load KNITRO's Julia Model
    m = unsafe_pointer_to_objref(userdata_)::Model
    nx = KN_get_number_vars(m)
    nc = KN_get_number_cons(m)

    x = unsafe_wrap(Array, ptr_x, nx)
    lambda = unsafe_wrap(Array, ptr_lambda, nx + nc)
    m.ms_initpt_callback(ptr_model, nSolveNumber, x, lambda, m)

    return Cint(0)
end

"""
    KN_set_ms_initpt_callback(m::Model, callback::Function)
Type declaration for the callback that allows applications to
specify an initial point before each local solve in the multistart
procedure.

Callback is a function with signature:

    callback(kc, x, lambda, userdata)

On input, arguments `x` and `lambda` are the randomly
generated initial points determined by Knitro, which can be overwritten
by the user.  The argument `nSolveNumber` is the number of the
multistart solve.

"""
function KN_set_ms_initpt_callback(m::Model, callback::Function)
    # store callback function inside model:
    m.ms_initpt_callback = callback

    # wrap user callback wrapper as C function
    c_func = @cfunction(ms_initpt_wrapper, Cint,
                        (Ptr{Cvoid}, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cvoid}))

    ret = @kn_ccall(set_ms_initpt_callback, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Any),
                    m.env, c_func, m)
    _checkraise(ret)
end

#--------------------
# Puts callback
#--------------------
function puts_callback_wrapper(str::Ptr{Cchar}, userdata_::Ptr{Cvoid})

    # Load KNITRO's Julia Model
    m = unsafe_pointer_to_objref(userdata_)::Model
    res = m.puts_callback(unsafe_string(str), m.userdata)

    return Cint(res)
end

"""
    KN_set_puts_callback(m::Model, callback::Function)

Set the callback that allows applications to handle
output. Applications can set a `put string` callback function to handle
output generated by the Knitro solver.  By default Knitro prints to
stdout or a file named `knitro.log`, as determined by KN_PARAM_OUTMODE.

Callback is a function with signature:

    callback(str::String, userdata)

The KN_puts callback function takes a `userParams` argument which is a pointer
passed directly from KN_solve. The function should return the number of
characters that were printed.

"""
function KN_set_puts_callback(m::Model, callback::Function)
    # store callback function inside model:
    m.puts_callback = callback

    # wrap user callback wrapper as C function
    c_func = @cfunction(puts_callback_wrapper, Cint,
                        (Ptr{Cchar}, Ptr{Cvoid}))

    ret = @kn_ccall(set_puts_callback, Cint,
                    (Ptr{Cvoid}, Ptr{Cvoid}, Any),
                    m.env, c_func, m)
end


text:
import MLJ
import MLJBase
import MLJModels
import MultivariateStats
import ModelSanitizer
import Test

const RidgeRegressor = MLJModels.MultivariateStats_.RidgeRegressor

mutable struct WrappedRidge <: MLJBase.DeterministicNetwork
    ridge_model
end

WrappedRidge(; ridge_model=RidgeRegressor) = WrappedRidge(ridge_model)

function MLJ.fit(model::WrappedRidge, X, y)
    Xs = MLJ.source(X)
    ys = MLJ.source(y)

    stand_model = MLJ.Standardizer()
    stand = MLJ.machine(stand_model, Xs)
    W = MLJ.transform(stand, Xs)

    box_model = MLJ.UnivariateBoxCoxTransformer()
    box = MLJ.machine(box_model, ys)
    z = MLJ.transform(box, ys)

    ridge_model = model.ridge_model
    ridge = MLJ.machine(ridge_model, W, z)
    zhat = MLJ.predict(ridge, W)

    yhat = MLJ.inverse_transform(box, zhat)
    MLJ.fit!(yhat, verbosity=0)

    return yhat
end

boston_task = MLJBase.load_boston()

wrapped_model = WrappedRidge(ridge_model=RidgeRegressor(lambda=0.1))

mach = MLJ.machine(wrapped_model, boston_task)

MLJ.fit!(mach; rows = :)

Test.@test mach.fitresult.nodes[1].data == boston_task.X
Test.@test all(convert(Matrix, mach.fitresult.nodes[1].data) .== convert(Matrix, boston_task.X))
for column in names(boston_task.X)
    Test.@test mach.fitresult.nodes[1].data[column] == boston_task.X[column]
    Test.@test all(mach.fitresult.nodes[1].data[column] .== boston_task.X[column])
end
Test.@test mach.fitresult.nodes[3].data == boston_task.y

ModelSanitizer.sanitize!(ModelSanitizer.Model(mach), ModelSanitizer.Data(boston_task.X), ModelSanitizer.Data(boston_task.y))

Test.@test all(convert(Matrix, mach.fitresult.nodes[1].data) .== 0)
for column in names(boston_task.X)
    Test.@test mach.fitresult.nodes[1].data[column] == zero(boston_task.X[column])
    Test.@test all(mach.fitresult.nodes[1].data[column] .== 0)
end
Test.@test mach.fitresult.nodes[3].data == zero(mach.fitresult.nodes[3].data)
Test.@test all(mach.fitresult.nodes[3].data .== 0)

ext: jl


text:
using PkgBenchmark

result = benchmarkpkg("GAIO")

println(result)


text:
using ScratchUsage

# Touch the scrachspaces and call it good
ScratchUsage.touch_scratch()


text:
include("Lanczos_Ritz_Utilities.jl")
include("../../src/iterative_lanczos.jl")

### ============================================================================
### Lanczos vector project onto the extreme ritz value's eigenvalues
###

function PerformAndProject(n=64; offset=20, exact=false)
    A = Diagonal(LinRange(-1, 1, n).^3)
    il = IterativeLanczos(A, ones(n))
    il.reorthogonalize = exact
    for _ in 1: offset - 1
        il()
    end
    ProjOnRitz = Vector()
    for II in offset:n - 1
        q = il()
        T = il|>GetTMatrix
        Q = il|>GetQMatrix

        s = eigvecs(T)[:, end:-1: (end - 2)]
        push!(ProjOnRitz, q'*Q*s)
    end

return vcat(ProjOnRitz...) end

# ----- Floats --------
ProjOnRitz = PerformAndProject()
ProjOnRitz = ProjOnRitz .|> abs
Iters = 20:63
fig1 = plot(
    Iters,
    ProjOnRitz[:, 1],
    yaxis=:log,
    legend=:bottomleft, size=(750, 500),
    label="\$q^T_kQ_ks_1\$",
    dpi=300, leftmargin=5*Plots.mm
)
plot!(fig1, Iters, ProjOnRitz[:, 2], linestyle=:dashdot, label="\$q^T_kQ_ks_2\$")
plot!(fig1, Iters, ProjOnRitz[:, 3], linestyle=:dash, label="\$q^T_kQ_ks_3\$")
title!("Lanczos vec proj onto Ritz Vectors (Floats)")
yaxis!("\$|q^T_kQS|\$")
xaxis!("iterations")
fig1|>display
savefig(fig1, "$(@__DIR__)/plots/lanczos_proj_on_ritz_float.png")

# ------ Exact --------
ProjOnRitz = PerformAndProject(exact=true)
ProjOnRitz = ProjOnRitz .|> abs
Iters = 20:63
fig1 = plot(
    Iters,
    ProjOnRitz[:, 1],
    yaxis=:log,
    legend=:bottomleft, size=(750, 500),
    label="\$q^T_kQ_ks_1\$",
    dpi=300, leftmargin=5*Plots.mm
)

plot!(fig1, Iters, ProjOnRitz[:, 2], linestyle=:dashdot, label="\$q^T_kQ_ks_2\$")
plot!(fig1, Iters, ProjOnRitz[:, 3], linestyle=:dash, label="\$q^T_kQ_ks_3\$")
title!("Lanczos vec proj on ritz vectors (largest 3) (Exact)")
yaxis!("\$|q^T_kQS|\$")
xaxis!("iterations")
fig1|>display
savefig(fig1, "$(@__DIR__)/plots/lanczos_proj_on_ritz_exact.png")


text:
import NCDatasets

filename = tempname()
#@show filename

dimlen = 10

T = Int32
data = Vector{Vector{T}}(undef,dimlen)
for i = 1:length(data)
    data[i] = T.(collect(1:i) .+ 100 * i)
end

varname = "varname"
vlentypename = "name-vlen"

# write data

ds = NCDatasets.Dataset(filename,"c",format=:netcdf4)
ds.dim["casts"] = dimlen
v = NCDatasets.defVar(ds,varname,Vector{T},("casts",); typename = vlentypename)
@test eltype(v.var) == Vector{T}

#for i = 1:dimlen
#    v.var[i] = data[i]
#end
v.var[:] = data
v.var[1] = data[1]
v.var[1:dimlen] = data[1:dimlen]

close(ds)

# load data

ds = NCDatasets.Dataset(filename)
vv = NCDatasets.variable(ds,"varname")
@test eltype(vv) == Vector{T}
data2 = vv[:]

@test data == data2
@test data[1] == vv[1]
@test data[2] == vv[2]
@test data[1:2] == vv[1:2]

NCDatasets.close(ds)

#@show data
#@show data2


text:
module Vann

using BlackBoxOptim
using Distributions
using DataAssim
using PyPlot
using CSV
using DataFrames

abstract Hydro
abstract Snow

export Hbv, Gr4j
export TinBasic, TinStandard

export run_timestep, run_timestep
export load_data, crop_data, load_operational
export init_states, get_states
export run_model_calib, calib_wrapper
export run_model
export get_param_range, assign_param
export get_input
export plot_sim
export epot_zero, epot_monthly
export particle_filter, enkf_filter
export split_prec, compute_ddf, pot_melt
export kge, nse

include("models/gr4j.jl")
include("models/hbv.jl")
include("models/tinbasic.jl")
include("models/tinstandard.jl")
include("utils_calib.jl")
include("utils_data_assim.jl")
include("utils_data.jl")
include("utils_epot.jl")
include("utils_model.jl")
include("utils_snow.jl")
include("utils_plot.jl")

end

max_stars_repo_name: jmgnve/Vann

max_forks_repo_name: jmgnve/Vann


text:
solver = Solver(model,obj,N=21)
n,m,N = get_sizes(solver)
p,pI,pE = get_num_constraints(solver)
p_N,pI_N,pE_N = get_num_terminal_constraints(solver)
obj.cost.Q .= Diagonal(I,n)
obj.cost.R .= Diagonal(I,m)*2
obj.cost.H .= ones(m,n)

Nx = N*n
Nu = (N-1)*m
Nz = Nx+Nu
X = rand(n,N)
U = rand(m,N-1)
Z = packZ(X,U)

V = [Z; rand(21)]
P = PrimalVars(V,n,m,N)
@test P.Z == Z
@test P.X == X
@test P.U == U
@test get_sizes(P) == (n,m,N)
@test P[3] == [X[:,3]; U[:,3]]

Nx = N*n
Nu = (N-1)*m
Nz = Nx + Nu
Nh = (N-1)*pE + pE_N
Ng = (N-1)*pI + pI_N
NN = 2Nx + Nu + Nh + Ng

names = (:z,:ν,:λ,:μ)
ind1 = TrajectoryOptimization.create_partition([Nz,Nx,Nh,Ng],names)
ind2 = TrajectoryOptimization.create_partition2([Nz,Nx,Nh,Ng],names)

ind_x, ind_u = get_primal_inds(n,m,N)
ind_z = TrajectoryOptimization.create_partition([Nx,Nu],(:x,:u))
ind_h = create_partition([(N-1)*pE,pE_N],(:k,:N))
ind_g = create_partition([(N-1)*pI,pI_N],(:k,:N))
ind_h = (k=ind_h.k, N=ind_h.N, k_mat=reshape(ind_h.k,pE,N-1))
ind_g = (k=ind_g.k, N=ind_g.N, k_mat=reshape(ind_g.k,pI,N-1))

CE = zeros(Nh)
CE[ind_h.k_mat]
CI = zeros(Ng)
CI[ind_g.k_mat]

Z = packZ(X,U)
mycost, c, grad_J, hess_J, jacob_c = gen_usrfun_newton(solver)
mycost(Z)
c.d(Z)
c.I(Z)
c.E(Z)
@test grad_J(Z) == ForwardDiff.gradient(mycost,Z)
@test jacob_c.E(Z) == ForwardDiff.jacobian(c.E,Z)
@testArray(jacob_c.I(Z)) == ForwardDiff.jacobian(c.I,Z)
@test Array(jacob_c.d(Z)) == ForwardDiff.jacobian(c.d,Z)

model,obj = Dynamics.dubinscar_parallelpark
solver = Solver(model,obj,N=21)
n,m,N = get_sizes(solver)
p,pI,pE = get_num_constraints(solver)
p_N,pI_N,pE_N = get_num_terminal_constraints(solver)
c_function!, c_jacobian!, c_labels, cI!, cE! = generate_constraint_functions(solver.obj)
dt = solver.dt
U0 = ones(m,N-1)

X = rand(n,N)
U = rand(m,N-1)
Z = PrimalVars(X,U)
Nx,Nu,Nh,Ng,NN = get_batch_sizes(solver)
Nz = Nx + Nu
part = create_partition([Nz,Nx,Nh,Ng],(:z,:ν,:λ,:μ))
ν = rand(Nx)
λ = rand(Nh)
μ = rand(Ng)
V = NewtonVars(Z,ν,λ,μ)
V1 = copy(V)
V1.V[1] = 100
V1.Z.Z[1] == 100
V1.V[Nz+1] = 200
V1.ν[1] == 200

res,stats = solve(solver,U0)
plot(res.X)

# Convert to PrimalVar
Z = PrimalVars(res)

# Get Constants
n,m,N = get_sizes(solver)
Nz = length(Z.Z)   # Number of decision variables
Nx = N*n           # Number of collocation constraints
nH = 0             # Number of entries in cost hessian

max_issues_repo_name: MehrdadZareian/GMARA

max_forks_repo_name: MehrdadZareian/GMARA


text:
#
# Plots for CR3BP manifolds
#

function Plots.plot(manifold::Manifold; vars=:XYZ, kwargs...)

    indices = vars isa Tuple ? vars : process_vars(vars)

    labels = Dict(
        0 => "Time" * " ($(timeunit(manifold)))",
        1 => "X"    * " ($(lengthunit(manifold)))",
        2 => "Y"    * " ($(lengthunit(manifold)))",
        3 => "Z"    * " ($(lengthunit(manifold)))",
        4 => "Ẋ"    * " ($(velocityunit(manifold)))",
        5 => "Ẏ"    * " ($(velocityunit(manifold)))",
        6 => "Ż"    * " ($(velocityunit(manifold)))"
    )

    defaults = (;
        linewidth = 1,
        linestyle = :dot,
        title     = "Manifold",
        palette   = :blues,
        dpi       = 150,
        label     = :none,
        xguide    = labels[indices[1]],
        yguide    = labels[indices[2]],
        zguide    = length(indices) == 3 ? labels[indices[3]] : :none,
    )
    options = merge(defaults, kwargs)
    fig = plot(; options...)

    for trajectory ∈ manifold.solution.u
        plot!(trajectory; vars=vars, options...)
    end

    return fig
end

function Plots.plot!(manifold::Manifold; vars=:XYZ, kwargs...)

    defaults = (;
        linewidth = 1,
        linestyle = :dot,
        palette   = :blues,
        dpi       = 150,
        label     = :none,
    )
    options = merge(defaults, kwargs)

    for trajectory ∈ manifold.solution.u[1:end-1]
        plot!(trajectory; vars=vars, options...)
    end

    plot!(manifold.solution.u[end]; vars=vars, options...)

end

function Plots.plot!(fig, manifold::Manifold; vars=:XYZ, kwargs...)

    defaults = (;
        linewidth = 1,
        linestyle = :dot,
        palette   = :blues,
        dpi       = 150,
        label     = :none,
    )
    options = merge(defaults, kwargs)

    for trajectory ∈ manifold.solution.u
        plot!(fig, trajectory; vars=vars, options...)
    end

    return fig
end


text:
function check_compatible_sizes(C, A, B)
    n, m = size(C)
    a, k = size(A)
    b, c = size(B)
    @assert (n == a) && (m == c) && (k == b) "matrices of size $(size(C)), $(size(A)), $(size(B)) are incompatible"
    nothing
end

function check_compatible_sizes(C::VecTypes, A, B::VecTypes)
    n    = length(C)
    a, k = size(A)
    b    = length(B)
    @assert (n == a) && (k == b) "matrices of size $(size(C)), $(size(A)), $(size(B)) are incompatible"
    nothing
end

function check_compatible_sizes(C::CoVecTypes, A::CoVecTypes, B::MatTypes)
    m    = length(C)
    n    = length(A)
    a, b = size(B)
    @assert (n == a) && (m == b) "matrices of size $(size(C)), $(size(A)), $(size(B)) are incompatible"
    nothing
end


text:
type GeneDeletion
    r_g
    r_f
    g_f
end

function Base.show(io::IO, m::GeneDeletion)
        println("\nType: GeneDeletion")
        @printf "%30s %7s : %6s %10s\n" "" "Field" "Size" "Fieldtype"
        @printf "%30s %7s : %6s %10s\n" "Reactions => Gene-knockout" fieldnames(m)[1] length(getfield(m,1)) typeof(getfield(m,1))
        @printf "%30s %7s : %6s %10s\n" "Reactions => Flow" fieldnames(m)[2] length(getfield(m,2)) typeof(getfield(m,2))
        @printf "%30s %7s : %6s %10s\n" "Gene-knockout => Flow" fieldnames(m)[3] length(getfield(m,3)) typeof(getfield(m,3))
end
type FBAsolution
	obj
	v
	slack
	rcost
	success
	info
end

type SolverInfo
	solver
	status
end

function Base.show(io::IO, m::FBAsolution)
	@printf "%1s: \n"  "FBAsolution"

	@printf "%12s %9f\n" "obj::" m.obj
	@printf "%12s %3d %2s\n" "v::" length(m.v) "element-array"
	@printf "%12s %3d %2s\n" "slack::" length(m.slack) "element-array"
	@printf "%12s %3d %2s\n" "rcosts::" length(m.rcost) "element-array"
	@printf "%12s %8s\n" "success::" m.success
	@printf "%13s %8s\n" "info:: " m.info
end
# THINK: Specify types wherever appropriate, remove obsolete stuff (sparsestruct?)
type Model
	rxns
	mets
	genes
	S
	lb
	ub
	c
	b
	csense
	rxn_gene_mat
	rxn_name
	rxn_rules
	rxn_subsystem
	rxn_extra
	met_formula
	met_name
	met_extra
	gene_name
	gene_extra
	description
end

function Base.show(io::IO, m::Model)
	for (i,v) in enumerate(fieldnames(m))
		@printf "%15s : %6s %10s\n" v length(getfield(m,i)) typeof(getfield(m,i))
	end
end

Base.length(m::Model) = (length(m.c), length(m.b), length(m.genes))
Base.size(m::Model) = (length(m.c), length(m.b), length(m.genes))

type SLG
	ess
	cond_ess
	non_ess
end

function Base.show(io::IO, m::SLG)
        println("\nType: SLG")
        @printf "%10s  %24s %7s  \n" "Field" "Contains" "Size"
        @printf "%10s  %24s %7d  \n" "ess" "Essential:" length(m.ess)
        @printf "%10s  %24s %7d  \n" "cond_ess" "Conditional Essential:" length(m.cond_ess)
        @printf "%10s  %24s %7d  \n" "non_ess" "Non-Essential:" length(m.non_ess)
end
type Robustness
    result
    reactions
    ranges
end

function Base.show(io::IO, m::Robustness)
    @printf "%1s \n"  "Robustness Analysis"
    @printf "%30s %15s \n"  "result :" string(size(m.result)," Array")

    @printf "%30s \n"  "ranges :"
    @printf "%38s : %15s \n" "reaction" "range:"
    for (i,v) in enumerate(m.ranges)
        @printf "%38s : %15s \n"  m.reactions[i] string("(", round(v[1],5), ",", round(v[end],5), ")")
    end

end

Base.range(a::Robustness, minFlux, maxFlux) = map(y -> ind2sub(a.result, y), find(x -> (x >= minFlux) & (x <= maxFlux), a.result))
Base.getindex(a::Robustness, idx...) = a.result[idx...]
Base.size(a::Robustness) = size(a.result)

max_stars_repo_name: isebarn/CBM

max_issues_repo_name: isebarn/CBM

max_forks_repo_name: isebarn/CBM


text:
using Test

include("../src/GassLaws.jl")
include("../src/Gasses/Chlorine.jl")
include("../src/conversions/TempratureConversion.jl")

@test 288 == absolute_temperature(15,"c")
@test 288 == absolute_temperature(59,"f")
@test 93.25 == charles_law(72,15,100,"c")
@test 93.25 == charles_law(72,59,212,"F")
@test 125 == boyles_law_solve_pressure(500,100,400)
@test 67.0 == round(boyles_law_solve_pressure(500,100,750))
@test 200 == boyles_law_solve_volume(100,2,1)
@test 50 == boyles_law_solve_volume(100,2,4)
@test 11.25503355704698 == gay_lussacs_law(3,25,845,"c")
@test 1.291 == round(grahams_law_velocity_ratio(20,12), digits = 3)
@test 0.775 == round(grahams_law_velocity_ratio(12,20), digits =3)
@test 8.166551644812515 == magnese_dioxide_to_chlorine(10)
@test 2.576489533011272 == magnese_dioxide_to_chlorine(10,"litre")
@test 73.0503795721187 == hydrochloric_acid_needed(43.5)


text:
## Julia program for Box-Muller algorithm
## author: weiya <szcfweiya@gmail.com>
## date: 2018-08-19

function boxmuller()
    x = ones(2)
    u = rand(2)
    logu = sqrt(-2*log(u[1]))
    x[1] = logu * cos(2*pi*u[2])
    x[2] = logu * sin(2*pi*u[2])
    return(x)
end

## example
boxmuller()

max_stars_repo_name: szcf-weiya/MCMC

max_issues_repo_name: szcf-weiya/MCMC

max_forks_repo_name: szcf-weiya/MCMC


text:
import Base: show, showcompact

# Implementation of Classical Laminate Theory

function qmat(m::Dict, mat::Int64=1)
  res = zeros(3, 3)
  if typeof(m[:E1]) == Float64
    @assert mat == 1
    m[:ν21] = m[:ν12]*m[:E2]/m[:E1]
    res = [m[:E1]/(1 - m[:ν12]*m[:ν21]) m[:ν12]*m[:E2]/(1 - m[:ν12]*m[:ν21]) 0;
    m[:ν12]*m[:E2]/(1 - m[:ν12]*m[:ν21]) m[:E2]/(1 - m[:ν12]*m[:ν21]) 0;
    0 0 m[:G12]]
  elseif typeof(m[:E1]) == Array{Float64, 1}
    m[:ν21] = similar(m[:ν12], Float64)
    for i in 1:length(m[:E1])
      m[:ν21][i] = m[:ν12][i]*m[:E2][i]/m[:E1][i]
    end
    res =
      [m[:E1][mat]/(1 - m[:ν12][mat]*m[:ν21][mat]) m[:ν12][mat]*m[:E2][mat]/(1 - m[:ν12][mat]*m[:ν21][mat]) 0;
      m[:ν12][mat]*m[:E2][mat]/(1 - m[:ν12][mat]*m[:ν21][mat]) m[:E2][mat]/(1 - m[:ν12][mat]*m[:ν21][mat]) 0;
      0 0 m[:G12][mat]
    ]
  end
  res
end

function qbarmat(qmat::Array{Float64, 2}, theta::Float64)
  c1 = cos(theta)
  s1 = sin(theta)
  c2 = c1*c1
  s2 = s1*s1
  c3 = c1*c2
  s3 = s1*s2
  c4 = c2*c2
  s4 = s2*s2
  qbar = zeros(Float64, 3, 3)
  qbar[1, 1] = qmat[1, 1]*c4 + 2*(qmat[1, 2] + 2qmat[3, 3])*s2*c2 + qmat[2, 2]*s4
  qbar[1, 2] = (qmat[1, 1] + qmat[2, 2] - 4qmat[3, 3])*s2*c2 + qmat[1, 2]*(s4 + c4)
  qbar[2, 1] = qbar[1, 2]
  qbar[2, 2] = qmat[1, 1]*s4 + 2*(qmat[1, 2] + 2qmat[3, 3])*s2*c2 + qmat[2, 2]*c4
  qbar[1, 3] = (qmat[1, 1] - qmat[1, 2] - 2qmat[3, 3])*s1*c3 +
    (qmat[1, 2] - qmat[2, 2] + 2qmat[3, 3])*s3*c1
  qbar[3, 1] = qbar[1, 3]
  qbar[2, 3] = (qmat[1, 1] - qmat[1, 2] - 2qmat[3, 3])*s3*c1 +
    (qmat[1, 2] - qmat[2, 2] + 2qmat[3, 3])*s1*c3
  qbar[3, 2] = qbar[2, 3]
  qbar[3, 3] = (qmat[1, 1] + qmat[2, 2] - 2qmat[1, 2] - 2qmat[3, 3])*s2*c2 +
    qmat[3, 3]*(s4 + c4)
  qbar
end

function createLaminate!(l::Dict, m::Dict, f::Dict)
  l[:thickness] = 0.0
  if !(:materials in keys(l)) || (:materials in keys(l) && typeof(l[:materials]) == Int64)
    l[:materials] = ones(Int64, l[:nplies])
  end
  if :repeats in keys(l) && l[:repeats] > 1
    l[:nplies] *= l[:repeats]
    l[:orientation] = repeat(l[:orientation], outer=[l[:repeats]])
    l[:materials] = repeat(l[:materials], outer=[l[:repeats]])
  end
  if :symmetric in keys(l) && l[:symmetric] == true
    l[:nplies] *= 2
    l[:orientation] = vcat(l[:orientation], reverse(l[:orientation]))
    l[:materials] = vcat(l[:materials], reverse(l[:materials]))
  end
  if :symmetricrepeats in keys(l) && l[:symmetricrepeats] > 1
    l[:nplies] *= l[:symmetricrepeats]
    l[:orientation] = repeat(l[:orientation], outer=[l[:symmetricrepeats]])
    l[:materials] = repeat(l[:materials], outer=[l[:symmetricrepeats]])
  end
  l[:z] = zeros(Float64, l[:nplies]+1)
  if typeof(m[:thickness]) == Float64
    l[:thickness] = l[:nplies] * m[:thickness]
    l[:z] = linspace(-l[:thickness]/2, l[:thickness]/2, l[:nplies]+1)
  elseif typeof(m[:thickness]) == Array{Float64, 1}
    @assert l[:nplies] == length(l[:materials])
    for i in 1:l[:nplies]
      l[:thickness] += m[:thickness][l[:materials][i]]
    end
    l[:z][1] = -l[:thickness]/2.0
    for i in 1:l[:nplies]
      l[:z][i+1] = l[:z][i] + m[:thickness][l[:materials][i]]
    end
  end
  computeABD!(l, m)
  l[:fTotal] = [f[:Nx],f[:Ny],f[:Nxy],f[:Mx],f[:My],f[:Mxy]] + l[:fT] + l[:fM]
  l[:f] = round(l[:ABD]\l[:fTotal], 15)
end

function computeABD!(l::Dict, m::Dict)
  A = zeros(Float64, 3, 3)
  B = zeros(Float64, 3, 3)
  D = zeros(Float64, 3, 3)
  ntemp = zeros(Float64, 3)
  mtemp = zeros(Float64, 3)
  nmoist = zeros(Float64, 3)
  mmoist = zeros(Float64, 3)
  for i in 1:l[:nplies]
    tk = l[:z][i+1]-l[:z][i]
    tk2 = l[:z][i+1]^2-l[:z][i]^2
    tk3 = l[:z][i+1]^3-l[:z][i]^3

    theta = float(l[:orientation][i])*(pi/180.0)
    A += qbarmat(qmat(m, l[:materials][i]), theta)*tk
    B += (1/2) * qbarmat(qmat(m, l[:materials][i]), theta)*tk2
    D += (1/3) * qbarmat(qmat(m, l[:materials][i]), theta)*tk3

    alpha = [
      m[:alpha1][l[:materials][i]]*cos(theta)^2 + m[:alpha2][l[:materials][i]]*sin(theta)^2,
      m[:alpha1][l[:materials][i]]*sin(theta)^2 + m[:alpha2][l[:materials][i]]*cos(theta)^2,
      2*cos(theta)*sin(theta)*(m[:alpha1][l[:materials][i]]-m[:alpha2][l[:materials][i]])
    ]

    beta = [
      m[:beta1][l[:materials][i]]*cos(theta)^2 + m[:beta2][l[:materials][i]]*sin(theta)^2,
      m[:beta1][l[:materials][i]]*sin(theta)^2 + m[:beta2][l[:materials][i]]*cos(theta)^2,
      2*cos(theta)*sin(theta)*(m[:beta1][l[:materials][i]]-m[:beta2][l[:materials][i]])
    ]

    ntemp += qbarmat(qmat(m, l[:materials][i]), theta)*alpha*l[:deltaTemp]*tk
    mtemp += (1/2)*qbarmat(qmat(m, l[:materials][i]), theta)*alpha*l[:deltaTemp]*tk2
    nmoist += qbarmat(qmat(m, l[:materials][i]), theta)*beta*l[:deltaMoisture]*tk
    mmoist += (1/2)*qbarmat(qmat(m, l[:materials][i]), theta)*beta*l[:deltaMoisture]*tk2
  end

  l[:fT] = vcat(ntemp, mtemp)
  l[:fM] = vcat(nmoist, mmoist)
  l[:ABD] = [ A B; B D]
  l[:abd] = inv(l[:ABD])
  l[:Ex] = 1/(l[:thickness]*l[:abd][1, 1])
  l[:Ey] = 1/(l[:thickness]*l[:abd][2, 2])
  l[:Gxy] = 1/(l[:thickness]*l[:abd][3, 3])
  l[:νxy] = -l[:abd][2, 1] / l[:abd][1, 1]
  l[:Eᶠx] = 12/(l[:thickness]^3*l[:abd][4, 4])
  l[:Eᶠy] = 12/(l[:thickness]^3*l[:abd][5, 5])
  l[:Gᶠxy] = 12/(l[:thickness]^3*l[:abd][6, 6])
  l[:νᶠxy] = -l[:abd][5, 4] / l[:abd][4, 4]
  l[:νᶠyx] = -l[:abd][5, 4] / l[:abd][5, 5]
end

function model_show(io::IO, l::Dict, compact::Bool)
  if compact
    println(m)
  else
    println("Effective material Properties:")
    println("  Ex  (psi) =       $(l[:Ex])")
    println("  Ey  (psi) =       $(l[:Ey])")
    println("  Gxy (psi) =       $(l[:Gxy])")
    println("  νxy       =       $(l[:νxy])")
    #println("  νyx       =       $(l[:νyx])")
    println()
    println("Strain Curvature (EpsilonKappa):")
    println("  εx   =       $(l[:f][1])")
    println("  εy   =       $(l[:f][2])")
    println("  γxy  =       $(l[:f][3])")
    println("  κx   =       $(l[:f][4])")
    println("  κy   =       $(l[:f][5])")
    println("  κxy  =       $(l[:f][6])")
    println()
    println("Effective flexural Properties:")
    println("  Eᶠx  (psi) =       $(l[:Eᶠx])")
    println("  Eᶠy  (psi) =       $(l[:Eᶠy])")
    println("  Gᶠxy (psi) =       $(l[:Gᶠxy])")
    println("  νᶠxy       =       $(l[:νᶠxy])")
    println("  νᶠyx       =       $(l[:νᶠyx])")
    println()
  end
end

show(io::IO, m::Dict) = model_show(io, m, false)
showcompact(io::IO, m::Dict) = model_show(io, m, true)


text:
## Plotting recipe for a setexpansion:
# - determine the plotting grid(depending on the basis)
# - evaluate the expansion in the gridpoints (fast if possible)
# - postprocess the data

@recipe function f(S::Expansion; plot_complex = false, n=200)
    legend --> false
    # title --> "Expansion"
    grid = plotgrid(dictionary(S), n)
    vals = plot_complex ? S.(grid) : real(S.(grid))
    grid, postprocess(dictionary(S), grid, vals)
end

# When a target function is provided, plot the error
# dispatch on dimension to set the logscale
@recipe function f(S::Expansion, target::Function; n=200)
    grid = plotgrid(dictionary(S), n)
    # Determine the return type so we know where to sample
    origvals = sample(grid, target, codomaintype(S))
    vals = abs.(origvals - S(grid))
    dictionary(S), grid, postprocess(dictionary(S), grid, vals)
end
# 1D error plot
@recipe function f(S::Dictionary1d, grid::AbstractGrid, vals)
    title --> "Error"
    legend --> false
    yscale --> :log10
    ylims --> [.1eps(real(codomaintype(S))),Inf]
    vals[vals.==0] .= eps(real(codomaintype(S)))^2
    # yaxis --> (:log10, (,Inf))
    grid, vals
end
# 2D error plot
@recipe function f(S::Dictionary2d, grid::AbstractGrid, vals)
    title --> "Error (log)"
    seriestype --> :heatmap
    grid, log10.(real(vals))
end

# Plot a Dictionary
@recipe function f(F::Dictionary; plot_complex = false, n=200)
    for i in eachindex(F)
        @series begin
            grid = plotgrid(F[i],n)
            z = F[i].(grid)
            vals = plot_complex ? z : real(z)
            grid, postprocess(F[i],grid,vals)
        end
    end
    nothing
end

# Plot a basis function
@recipe function f(φ::AbstractBasisFunction; plot_complex = false, n=200)
    @series begin
        grid = plotgrid(φ,n)
        z = φ.(grid)
        vals = plot_complex ? z : real(z)
        grid, postprocess(φ,grid,vals)
    end
    nothing
end

#
# For regular Expansions, no postprocessing is needed
postprocess(S::Dictionary, grid, vals) = vals

# For function subsets, revert to the underlying Dictionary for postprocessing
postprocess(S::Subdictionary, grid, vals) = postprocess(superdict(S), grid, vals)
postprocess(φ::AbstractBasisFunction, grid, vals) = postprocess(dictionary(φ), grid, vals)

## Plotting grids
# Always plot on equispaced grids for the best plotting resolution
plotgrid(S::Dictionary1d, n) = PeriodicEquispacedGrid(n, support(S))
# look at first element by default
plotgrid(S::MultiDict, n) = plotgrid(component(S,1), n)
plotgrid(S::DerivedDict, n) = plotgrid(superdict(S),n)
# NOTE: This only supoorts multi-dimensional tensor product dicts.
plotgrid(F::TensorProductDict2, n) = plotgrid(component(F,1),n)×plotgrid(component(F,2),n)
plotgrid(F::OperatedDict, n) = plotgrid(superdict(F), n)
plotgrid(F::Subdictionary, n) = plotgrid(superdict(F), n)
plotgrid(φ::AbstractBasisFunction, n) = plotgrid(dictionary(φ), n)
## Split complex plots in real and imaginary parts
# 1D
@recipe function f(A::AbstractArray{S}, B::Array{Complex{T}}) where {S<:Real, T<:Real}
    # Force double layout
    layout := 2
    # Legend is useless here
    legend --> false
    @series begin
        subplot := 1
        title := "Real"
        A,real.(B)
    end
    @series begin
        subplot :=2
        title := "Imaginary"
        A,imag.(B)
    end
    nothing
end

# 2D
@recipe function f(A::LinRange{S},B::LinRange{S},C::Array{T}) where {S<:Real,T<:Complex}
    # Force double layout
    layout := 2
    # Legend is useless here
    legend --> false
    @series begin
        subplot := 1
        title := "Real"
        A,B,real(C)
    end
    @series begin
        subplot :=2
        title := "Imaginary"
        A,B,imag(C)
    end
    nothing
end

@recipe function f(dom::Domain2d; n=300, distance=false, xlim=[-1,1], ylim=[-1,1])
    seriescolor --> :tempo
     seriestype --> :heatmap
    aspect_ratio --> 1
    colorbar --> false
    # xrange = linspace(xlim[1],xlim[2],n)
    xrange = EquispacedGrid(n,xlim[1],xlim[2])
    # yrange = linspace(ylim[1],ylim[2],n)
    yrange = EquispacedGrid(n,ylim[1],ylim[2])
    # grid = [SVector(i,j) for i in xrange , j in yrange]
    grid = xrange×yrange

    plotdata = distance ? dist.(grid, Ref(dom)) : in.(grid, Ref(dom))
    collect(xrange),collect(yrange),plotdata'
end


text:
using Sundials, Compat

## f routine. Compute function f(t,y).

function f(t, y_nv, ydot_nv, user_data)
    y = convert(Vector, y_nv)
    ydot = convert(Vector, ydot_nv)
    ydot[1] = -0.04*y[1] + 1.0e4*y[2]*y[3]
    ydot[3] = 3.0e7*y[2]*y[2]
    ydot[2] = -ydot[1] - ydot[3]
    return Sundials.CV_SUCCESS
end

## g routine. Compute functions g_i(t,y) for i = 0,1.

function g(t, y_nv, gout_ptr, user_data)
    y = convert(Vector, y_nv)
    gout = Sundials.asarray(gout_ptr, (2,))
    gout[1] = y[1] - 0.0001
    gout[2] = y[3] - 0.01
    return Sundials.CV_SUCCESS
end

## Jacobian routine. Compute J(t,y) = df/dy.
# broken -- needs a wrapper from Sundials._DlsMat to Matrix and Jac user function wrapper
function Jac(N, t, ny, fy, Jptr, user_data, tmp1, tmp2, tmp3)
    y = convert(Vector, ny)
    dlsmat = unpack(IOString(@compat unsafe_wrap(convert(Ptr{UInt8}, Jptr),
                                                 (sum(map(sizeof, Sundials._DlsMat))+10,), false)),
                    Sundials._DlsMat)
    J = @compat unsafe_wrap(unsafe_ref(dlsmat.cols), (Int(neq), Int(neq)), false)
    J[1,1] = -0.04
    J[1,2] = 1.0e4*y[3]
    J[1,3] = 1.0e4*y[2]
    J[2,1] = 0.04
    J[2,2] = -1.0e4*y[3] - 6.0e7*y[2]
    J[2,3] = -1.0e4*y[2]
    J[3,2] = 6.0e7*y[2]
    return Sundials.CV_SUCCESS
end

const neq = 3

const t0 = 0.0
const t1 = 0.4
const tmult = 10.0
const nout = 12
const y0 = [1.0, 0.0, 0.0]
const reltol = 1e-4
const abstol = [1e-8, 1e-14, 1e-6]

cvode_mem = Sundials.CVodeCreate(Sundials.CV_BDF, Sundials.CV_NEWTON)
Sundials.@checkflag Sundials.CVodeInit(cvode_mem, f, t0, y0)
Sundials.@checkflag Sundials.CVodeSVtolerances(cvode_mem, reltol, abstol)
Sundials.@checkflag Sundials.CVodeRootInit(cvode_mem, 2, g)
Sundials.@checkflag Sundials.CVDense(cvode_mem, neq)
#Sundials.@checkflag Sundials.CVDlsSetDenseJacFn(cvode_mem, Jac)

iout = 0
tout = t1
const t = [t0]

while iout < nout
    y = similar(y0)
    flag = Sundials.CVode(cvode_mem, tout, y, t, Sundials.CV_NORMAL)
    println("T=", tout, ", Y=", y)
    if flag == Sundials.CV_ROOT_RETURN
        rootsfound = zeros(Cint, 2)
        Sundials.@checkflag Sundials.CVodeGetRootInfo(cvode_mem, rootsfound)
        println("roots=", rootsfound)
    elseif flag == Sundials.CV_SUCCESS
      iout += 1
      tout *= tmult
    end
end


text:
__precompile__(true)
module  MyNGSIM1

# Using packages
using DataFrames
using Reexport
@reexport using AutomotiveDrivingModels
@reexport using AutoViz
@reexport using NGSIM

export
    # functions
    get_scene_id,
    get_neighbor_fore_along_lane_NGSIM,
    get_neighbor_fore_along_left_lane_NGSIM,
    get_neighbor_fore_along_right_lane_NGSIM,
    get_neighbor_rear_along_lane_NGSIM,
    get_neighbor_rear_along_left_lane_NGSIM,
    get_neighbor_rear_along_right_lane_NGSIM,
    # Raw trajectory data in Type NGSIMTrajectorydata
    tdraw_my,
    # Processed trajectory data
    trajdata_my,
    # My road way
    ROADWAY_my,
    # Customized data types
    LaneChangeInfo,
    RefinedLaneChangeInfo,
    LaneKeepInfo,
    RefinedLaneKeepInfo

# =============================================================================
# Data Type
# =============================================================================
# Define a structure for storing lane change information of vehicles.
struct LaneChangeInfo
    changto::AbstractString         # Changed to which direction 'left' or 'right'.
    from_to::Array{Int32, 1}       # [laneID_1, laneID_2] From laneID_1 to laneID_2
    cam_frames::Array{Int, 1}       # The camera frames during which the change happened.
    df_frames::Array{Int,1}         # The dataframe frames during which the change happened.
end

struct RefinedLaneChangeInfo
    changto::AbstractString         # Changed to which direction 'left' or 'right'.
    from_to::Array{Int32, 1}       # [laneID_1, laneID_2] From laneID_1 to laneID_2
    cam_frame::Int                 # The exact camera frame where the vehicle crossed the lane.
    df_frame::Int                  # The exact dataframe frame where the vehicle crossed the lane.
end

# Define a structure for storing lane change information of vehicles.
struct LaneKeepInfo
    cam_frames::Array{Int, 1}       # The camera frames during which the change happened.
    df_frames::Array{Int,1}         # The dataframe frames during which the change happened.
end

struct RefinedLaneKeepInfo
    cam_frame::Int                  # The exact camera frame where the vehicle started braking.
    df_frame::Int                   # The exact dataframe frame where the vehicle started braking.
end

# =============================================================================
# Variables
# =============================================================================
# We define the reference points on the vehicle body. For example, when we need
# to measure the distance between two vehicles, the distance can be from the center point
# of vehcile 1 to the center point of vehicle 2, with reference point VEHICLE_TARGET_POINT_CENTER;
# or the distance can be from the front of vehicle 1 to the back of vehicle 2,
# with reference points VEHICLE_TARGET_POINT_FRONT, VEHICLE_TARGET_POINT_REAR, respectively.
const VEHICLE_TARGET_POINT_CENTER = VehicleTargetPointCenter()
const VEHICLE_TARGET_POINT_FRONT = VehicleTargetPointFront()
const VEHICLE_TARGET_POINT_REAR = VehicleTargetPointRear()

# Load the traffic data of us_101, the units are meters, m/s
# 1. "i101_trajectories-0750am-0805am.txt"
# 2. "i101_trajectories-0805am-0820am.txt"
# 3. "i101_trajectories-0820am-0835am.txt"
# 4. "i80_trajectories-0400-0415.txt",
# 5. "i80_trajectories-0500-0515.txt",
# 6. "i80_trajectories-0515-0530.txt"
# Indicate the proper dataset here.
dataset_num = 1
#
if dataset_num == 1
    trajdata_my = load_trajdata(1) # "i101_trajectories-0750am-0805am.txt"
    road_filename = "i101_trajectories-0750am-0805am.txt"
    ROADWAY_my = ROADWAY_101
elseif dataset_num == 2
    trajdata_my = load_trajdata(2) # "i101_trajectories-0805am-0820am.txt"
    road_filename = "i101_trajectories-0805am-0820am.txt"
    ROADWAY_my = ROADWAY_101
elseif dataset_num == 3
    trajdata_my = load_trajdata(3) # "i101_trajectories-0820am-0835am.txt"
    road_filename = "i101_trajectories-0820am-0835am.txt"
    ROADWAY_my = ROADWAY_101
elseif dataset_num == 4
    trajdata_my = load_trajdata(4) # "i80_trajectories-0400-0415.txt",
    road_filename = "i80_trajectories-0400-0415.txt"
    ROADWAY_my = ROADWAY_80
elseif dataset_num == 5
    trajdata_my = load_trajdata(5) # "i80_trajectories-0500-0515.txt"
    road_filename = "i80_trajectories-0500-0515.txt"
    ROADWAY_my = ROADWAY_80
elseif dataset_num == 6
    trajdata_my = load_trajdata(6) # "i80_trajectories-0515-0530.txt"
    road_filename = "i80_trajectories-0515-0530.txt"
    ROADWAY_my = ROADWAY_80
else
    error("not appropriate dataset number")
end
road_filepath = Pkg.dir("NGSIM", "data", road_filename)
# --Load the un-converted data of us_101. This data is in DataFrammes structure,
# The units are feet, foot/s
tdraw_my = NGSIM.load_ngsim_trajdata(road_filepath)

# =============================================================================
# Functions of MyNGSIM
# =============================================================================
# We define a function to find vehicle scene id from its NGSIM id.
function get_scene_id{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int
    )
    # Use NGSIM_vehicle_index to search for the scene vehicel index.
    # --Get the total number of the vehicles in the current scene.
    total_in_scene = scene.n
    # Note: Be careful, in Julia, if you want to return a value from a for-loop,
    # you should define the corresponding variable before the for-loop. For instance,
    # vehicle_index in the following; otherwise, nothing is returned from the for-loop.
    vehicle_index = 0
    for scene_ind = 1:total_in_scene
        # If the NGSIM_vehicle_index is the same in the current vehicle, this vehicle
        # is what we are looking for.
        if scene[scene_ind].id == NGSIM_vehicle_index
            vehicle_index = scene_ind
            break
        # Otherwise, if the current scene_ind is the last one, a.k.a, we have searched
        # all the vehicles in the scene, we return an error saying no finding.
        elseif scene_ind == total_in_scene
            error("This wanted vehicle is not in the current scene.")
        end
        # If neither case happens, this for-loop goes to the next iteration.
    end
    vehicle_index # return the vehicle' index in the scene
end

# Add a new method to function get_neigbor_fore_along_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_fore_along_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_fore::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_fore_along_lane() from the package AutomotiveDrivingModels
    get_neighbor_fore_along_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_FRONT, VEHICLE_TARGET_POINT_REAR,
        VEHICLE_TARGET_POINT_FRONT, max_distance_fore=max_distance_fore)
end

# Add a new method to function get_neigbor_fore_along_left_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_fore_along_left_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_fore::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_fore_along_left_lane() from the package AutomotiveDrivingModels
    get_neighbor_fore_along_left_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_FRONT, VEHICLE_TARGET_POINT_REAR,
        VEHICLE_TARGET_POINT_FRONT, max_distance_fore=max_distance_fore)
end

# Add a new method to function get_neigbor_fore_along_right_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_fore_along_right_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_fore::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_fore_along_right_lane() from the package AutomotiveDrivingModels
    get_neighbor_fore_along_right_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_FRONT, VEHICLE_TARGET_POINT_REAR,
        VEHICLE_TARGET_POINT_FRONT, max_distance_fore=max_distance_fore)
end

# Add a new method to function get_neigbor_rear_along_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_rear_along_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_rear::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_rear_along_lane() from the package AutomotiveDrivingModels
    get_neighbor_rear_along_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_REAR, VEHICLE_TARGET_POINT_FRONT,
        VEHICLE_TARGET_POINT_REAR, max_distance_rear=max_distance_rear)
end

# Add a new method to function get_neigbor_rear_along_left_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_rear_along_left_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_rear::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_rear_along_left_lane() from the package AutomotiveDrivingModels
    get_neighbor_rear_along_left_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_REAR, VEHICLE_TARGET_POINT_FRONT,
        VEHICLE_TARGET_POINT_REAR, max_distance_rear=max_distance_rear)
end

# Add a new method to function get_neigbor_rear_along_right_lane.
# This function returns the struct NeighborLongitudinalResult{id, distance}.
# The id is the scene_id and the distance is intervehicular distance.
function get_neighbor_rear_along_right_lane_NGSIM{S<:VehicleState,D<:Union{VehicleDef, BicycleModel},I}(
    scene::EntityFrame{S,D,I},
    NGSIM_vehicle_index::Int,
    roadway::Roadway,
    max_distance_rear::Float64 = 50.0 # max distance to search forward [m]
    )
    # Use NGSIM_vehicle_index to search for the scene vehicle index.
    vehicle_index = get_scene_id(scene, NGSIM_vehicle_index)
    # Call the function get_neighbor_rear_along_left_lane() from the package AutomotiveDrivingModels
    get_neighbor_rear_along_right_lane(scene, vehicle_index, roadway,
        VEHICLE_TARGET_POINT_REAR, VEHICLE_TARGET_POINT_FRONT,
        VEHICLE_TARGET_POINT_REAR, max_distance_rear=max_distance_rear)
end

end # end of module

max_stars_repo_name: LongshengJiang/NGSIM-sim

max_issues_repo_name: LongshengJiang/NGSIM-sim

max_forks_repo_name: LongshengJiang/NGSIM-sim


text:
function sortkey(key_values, a)
    indperm=sortperm(key_values)
    return a[indperm]
end


text:
module ADFPEPS

export HamiltonianModel,hamiltonian
export Hubbard,hop_pair,Hubbard_hand,hop_pair_hand,THubbard
export observable

include("hamiltonianmodels.jl")
include("fermion.jl")
include("autodiff.jl")
include("variationalipeps.jl")
include("observable.jl")

end


text:
########################################
# Geolocation functions
########################################
"""
    geolocate(geodata::GeoIP.DB, ip)

Returns geolocation and other information determined in `geodata` by `ip`.
"""
function geolocate(geodata::DB, ip::IPv4)
    ipnet = IPv4Net(ip, 32)

    # TODO: Dict should be changed to a more suitable structure
    res = Dict{String, Any}()

    idx = searchsortedfirst(geodata.index, ipnet) - 1
    if idx == 0 || !(ip in geodata.index[idx])
        return res
    end
    row = geodata.blocks[idx]

    res["v4net"] = row.v4net
    res["geoname_id"] = row.geoname_id
    res["location"] = row.location
    res["registered_country_geoname_id"] = row.registered_country_geoname_id
    res["is_anonymous_proxy"] = row.is_anonymous_proxy
    res["is_satellite_provider"] = row.is_satellite_provider
    res["postal_code"] = row.postal_code
    res["accuracy_radius"] = row.accuracy_radius

    geoname_id = row.geoname_id
    locale = geodata.locs[geodata.localeid]
    idx2 = searchsortedfirst(locale.index, geoname_id)
    if idx2 > length(locale.locs) || idx2 < 1
        return res
    end
    if locale.index[idx2] != geoname_id
        return res
    end

    row2 = locale.locs[idx2]
    for k in keys(row2)
        res[string(k)] = row2[k]
    end

    return res
end

geolocate(geodata::DB, ipstr::AbstractString) = geolocate(geodata, IPv4(ipstr))
geolocate(geodata::DB, ipint::Integer) = geolocate(geodata, IPv4(ipint))
getindex(geodata::DB, ip) = geolocate(geodata, ip)


text:
pushfirst!(LOAD_PATH, joinpath(@__DIR__, "..", ".."))

using Oceananigans
using Oceananigans.Units
using Oceananigans.TimeSteppers: time_step!
using Oceananigans.TurbulenceClosures: VerticallyImplicit, CATKEVerticalDiffusivity
using Oceananigans.Models.HydrostaticFreeSurfaceModels: ColumnEnsembleSize

using CUDA
using BenchmarkTools

Ex, Ey = (400, 20)
sz = ColumnEnsembleSize(Nz=128, ensemble=(Ex, Ey))
halo = ColumnEnsembleSize(Nz=sz.Nz)

grid = RectilinearGrid(size=sz, halo=halo, z=(-128, 0), topology=(Flat, Flat, Bounded))

closure = CuArray([CATKEVerticalDiffusivity() for i=1:Ex, j=1:Ey])

Qᵇ = CuArray([+1e-8 for i=1:Ex, j=1:Ey])
Qᵘ = CuArray([-1e-4 for i=1:Ex, j=1:Ey])
Qᵛ = CuArray([0.0   for i=1:Ex, j=1:Ey])

u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))
v_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵛ))
b_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ))

# Half rotating, half not
f_ij(i, j) = j < Ey/2 ? 1e-4 : 0.0
coriolis_ensemble = CuArray([FPlane(f=f_ij(i, j)) for i=1:Ex, j=1:Ey])

model = HydrostaticFreeSurfaceModel(architecture = GPU(),
                                    grid = grid,
                                    tracers = (:b, :e),
                                    buoyancy = BuoyancyTracer(),
                                    coriolis = coriolis_ensemble,
                                    boundary_conditions = (b=b_bcs, u=u_bcs, v=v_bcs),
                                    closure = closure)

N² = 1e-5
bᵢ(x, y, z) = N² * z
set!(model, b = bᵢ)

function one_step!(model)
    CUDA.@sync time_step!(model, 1e-9)
    return nothing
end

@info "Running one time-step..."
one_step!(model)

@info "Benchmarking..."
@btime one_step!(model)


text:
py"""
class DiffBertsimas(torch.nn.Module):
    '''
    DiffBertsimas

Differentiable Bertsimas model.

Atributes:
- `D::Int`: dimension (i.e. number of assets)
- `w::PyObject`: differentiable weights
- `c::Vector{PyObject}`: convex constraints on weights
- `μ::Array{Float64,1}`: predicted mean vector
- `sqrtΣ::Array{Float64,2}`: predicted upper triangular Cholesky factor of covariance matrix
- ` solver_args::Dict=Dict("use_indirect" => true, "eps" => 1e-5, "max_iters" => 20000}`: solver arguments
    '''
    def __init__(
        self,
        D,
        w,
        c,
        μ,
        sqrtΣ,
        solver_args={'use_indirect': True, 'eps': 1e-5, 'max_iters': 20000}
    ):
        super(DiffBertsimas, self).__init__()
        self.D = D
        self.μ = μ
        Σ = sqrtΣ.T @ sqrtΣ
        self.Σ = Σ
        self.solver_args=solver_args

        λ = cp.Variable()
        π_neg = cp.Variable(D)
        π_pos = cp.Variable(D)
        θ = cp.Variable(D)

        params = cp.Parameter(D+1)

        constraints = [
            λ >= 0,
            π_neg >= 0,
            π_pos >= 0,
            θ >= 0,
            w == π_pos - π_neg,
            cp.multiply(params[0:D], (π_pos + π_neg)) - θ <= λ,
        ]

        objective = cp.Maximize(μ.T @ w - cp.sum(θ) - λ * params[D])
        problem = cp.Problem(objective, constraints + c)

        self.cvxpylayer = CvxpyLayer(
            problem,
            parameters=[params],
            variables=[w, λ, π_neg, π_pos, θ]
        )

    def forward(self, params):
        return self.cvxpylayer(params, solver_args=self.solver_args)[0]

    def get_variables(self, params):
        params_ = torch.tensor(params, dtype=torch.float64, requires_grad=True)
        return self.cvxpylayer(params_, solver_args=self.solver_args)

    def get_w(self, params):
        params_ = torch.tensor(params, dtype=torch.float64, requires_grad=True)
        return self.get_variables(params_)[0]

    def get_lagrangian(self, params):
        params_ = torch.tensor(params, dtype=torch.float64, requires_grad=True)
        D = self.D
        μ_ = torch.tensor(self.μ, dtype=torch.float64)
        Σ_ = torch.tensor(self.Σ, dtype=torch.float64)
        w_, λ_, π_neg_, π_pos_, θ_ = self.cvxpylayer(params_, solver_args=self.solver_args)
        return μ_.T @ w_ - torch.sum(θ_) - λ_ * params_[D]
"""

DiffBertsimas = py"DiffBertsimas"

max_stars_repo_name: molet/DiffDRO

max_issues_repo_name: molet/DiffDRO

max_forks_repo_name: molet/DiffDRO


text:
# Below packages can be used for simulation, data storage and summary
using JLD2, DataFrames, LinearAlgebra, Distributions, Serialization, StatsPlots

# These functions from `GEAS` are to be used in this simulation
import GEAS:breeding, sim_QTL, summarize, create_storage, sim_pt_QTL
import GEAS:rank_QTL, calc_idx

"""
    function gpar(; edit = false, emph = false, cv = 0., g8n = 3)
---
Generate parameters for a simulation scenario.
Note, once the named parameter tuple is set, its value can't be changed.
You can modify the arguments of this function to change elements of `par`.
The return is a named tuple.
That is you can quote a value like `par.nSire`.
To avoid too long tuple names, I used numbers as shortcuts, e.g.,
`generation` → `g8n`.

`cv` is for pleiotropic QTL distribution, which follow a distribution of
`MvNormal(zeros(2), [1 cc; cc 1])`.  `cv` is the covariance of the QTL
effects for traits.
"""
function gpar(;
              binary = false,   # whether to convert challenge trait binary
              edit   = false,
              emph   = false,   # if to emphasize known QTL
              cv     = 0.,      # for pleiotropic trait
              g8n    = 3,       # number of generation to simulate
              u38y   = false)
    par = Dict(
        :nSire=> 100,
        :nDam => 200,   # male:female = 1:2
        :nSib => 40,
        :nC7e => 30,    # number of sib of `nSib` to be challenged
        :nG8n => g8n,   # number of generations to select

        # Note: Using tuples are more efficient than to use vectors
        :h²   => (a = .5, b =.5),
        :w4t  => 1,     # weight on the binary trait EBV
        :cv   => cv,    # covariance of pleiotripic QTL for two traits

        :p17h => .5,    # percentage-of-death after challenge
        :e19e => 1.,    # edit_successsful_rate (e19e), here, 1.0 = 100%
        :nK3n => 10,    # number of known (k3n) QTL on binary trait

        :dd   => 0.01,  # A small value to be added to diagonals
        :edit => edit,  # whether to edit (the biggest) known QTL
        :fix  => emph, # if to fit (the largest) known QTL as fixed effects
        :b4y  => binary, # for debug only, remove the threshold of binary trait
        :u38y => u38y # "use genotypes of current generation only" :40char
    )
    (; par...)
end

"""
    function example_simple_simulation(base, qtl, g8n; dir = ".")
---
# Objectives
- Do a simple breeding simulation with a production and a challenge trait.
- Write intermediate and final results in `dir`, which, by default, is the `pwd`.

# Introduction
This module is designed to run on a large mainframe machine,
where many repeats can be ran simultaneously.
The run is only one repeat of some certain scenario.

This function uses all the set threads for matrix calculation.
Run:
```bash
export JULIA_NUM_THREADS=12
```
before starting Julia.
Or, put above line into your `~/.bashrc`.
You'd better set this number ≤ your available physical CPU cores.
"""
function example_simple_simulation(base, qtl, g8n; dir = ".")
    rst = DataFrame()

    # par = gpar(g8n = g8n, u38y = true)
    par = gpar(g8n = g8n)
    @info "prepare common base + generation 1 of $(par.nG8n)"
    # by default Lapace below.
    ped, snp = create_storage(base, par, qtl)
    serialize("$dir/cmbase.ser", (ped=ped, snp=snp, qtl=qtl))

    @info "Breeding with SNP-BLUP"
    breeding(ped, snp, base, qtl, par)
    df = summarize(ped.prd, snp.prd, qtl[2])
    df.method = repeat(["SNP-BLUP"], nrow(df))
    append!(rst, df)

    @info "Breeding with emphasis on known QTL"
    ped, snp, qtl = deserialize("$dir/cmbase.ser")
    par = gpar(emph = true, g8n = g8n)
    breeding(ped, snp, base, qtl, par)
    df = summarize(ped.prd, snp.prd, qtl[2])
    df.method = repeat(["Fixed"], nrow(df))
    append!(rst, df)

    @info "Breeding with edit on known QTL"
    ped, snp, qtl = deserialize("$dir/cmbase.ser")
    par = gpar(edit = true, g8n = g8n)
    breeding(ped, snp, base, qtl, par)
    df = summarize(ped.prd, snp.prd, qtl[2])
    df.method = repeat(["Edit"], nrow(df))
    append!(rst, df)

    oo = tempname(dir)
    @save "$oo-r.jld2" rst
end

"""
    function one_repeat(nqtl, g8n, jld; dir = ".")
---
Load a base population in `jld`, either real or simulated.
This function simulate two sets of QTL of Laplace distribution for two traits.
The `example_simple_simulation` will do breeding with `SNP-BLUP`, `Fixed`
and `Edit` methods.
Results (`*.jld2`) and intermediate data (`*ser`) are stored in `dir`.

You may find this `example` file by:
```julia
dir = pathof(GEAS)
til = findlast("/src", dir)[1]
example = joinpath(tmp[1:til], "examples/sample-one-repeat.jl")
```
"""
function one_repeat(nqtl, g8n, jld; dir = ".")
    isdir(dir) || mkpath(dir)
    ntd = Threads.nthreads()
    BLAS.set_num_threads(ntd)   # use all specified threads (before REPL)
    # Note `twop` was added to base.jld2, Jan, 2022
    @load jld base  # load population, which can be real, or simulated
    qtl = sim_QTL(base, nqtl, nqtl) # can change this to sim_pt_QTL()
    example_simple_simulation(base, qtl, g8n, dir = dir)
end

"""
    function pt_repeat(nqtl, g8n, jld; dir = ".")
---
Load a base population in `jld`, either real or simulated.
This function simulate one set of QTL of `MvNormal([0, 0], [1 0; 0 1])`
distribution for two traits. QTL are pleiotropic.
The `example_simple_simulation` will do breeding with `SNP-BLUP`, `Fixed`
and `Edit` methods.
Results (`*.jld2`) and intermediate data (`*ser`) are stored in `dir`.
"""
function pt_repeat(nqtl, g8n, jld; dir = ".")
    isdir(dir) || mkpath(dir)
    ntd = Threads.nthreads()
    BLAS.set_num_threads(ntd)
    @load jld base
    qtl = sim_pt_QTL(base, 100, MvNormal([0., 0.], [1 0; 0 1.]))
    example_simple_simulation(base, qtl, g8n, dir = dir)
end

function test_codes(u38y = false)
    ntd = Threads.nthreads()
    BLAS.set_num_threads(ntd)   # use all specified threads (before REPL)
    @load "dat/run/base.jld2" base
    qtl = sim_QTL(base, 100, 100) # can change this to sim_pt_QTL()
    par = gpar(g8n = 3, u38y = u38y)
    @info "prepare common base + generation 1 of $(par.nG8n)"
    # by default Lapace below.
    ped, snp = create_storage(base, par, qtl)

    rkq = rank_QTL(base.hap, qtl[2]) # QTL for binary trait ranking in ↓ order
    Q = par.fix ? abs.(rkq[1:par.nK3n]) : Int[]
    prd, clg = groupby(ped.prd, :g8n), groupby(ped.clg, :g8n) # shorthands

    # note, there are `nG8n+1` generations, including the base.
    @info "Evaluating generation 1"
    calc_idx(ped, snp, prd[2], Q, par) # evaluation of generation one
end


text:
import Oceananigans.OutputWriters: saveproperty!

"""
    AdamsBashforthTimeStepper(float_type, arch, grid, tracers, χ)

Return an AdamsBashforthTimeStepper object with tendency
fields on `arch` and `grid` and AB2 parameter `χ`.
"""
struct AdamsBashforthTimeStepper{T, TG}
      Gⁿ :: TG
      G⁻ :: TG
       χ :: T
end

function AdamsBashforthTimeStepper(float_type, arch, grid, tracers, χ=0.125)
   Gⁿ = Tendencies(arch, grid, tracers)
   G⁻ = Tendencies(arch, grid, tracers)
   return AdamsBashforthTimeStepper{float_type, typeof(Gⁿ)}(Gⁿ, G⁻, χ)
end

# Special savepropety! for AB2 time stepper struct used by the checkpointer so
# it only saves the fields and not the tendency BCs or χ value (as they can be
# constructed by the `Model` constructor).
function saveproperty!(file, location, ts::AdamsBashforthTimeStepper)
    saveproperty!(file, location * "/Gⁿ", ts.Gⁿ)
    saveproperty!(file, location * "/G⁻", ts.G⁻)
end

#####
##### Time steppping
#####

"""
    time_step!(model{<:AdamsBashforthTimeStepper}, Nt, Δt; init_with_euler=true)

Step forward `model` `Nt` time steps with step size `Δt` with an Adams-Bashforth
timestepping method.
"""
function time_step!(model::Model{<:AdamsBashforthTimeStepper}, Nt, Δt; init_with_euler=true)

    if model.clock.iteration == 0
        [ run_diagnostic(model, diag) for diag in values(model.diagnostics) ]
        [ write_output(model, out)    for out  in values(model.output_writers) ]
    end

    for n in 1:Nt
        time_step!(model, Δt; euler=init_with_euler && n==1)

        [ time_to_run(model.clock, diag) && run_diagnostic(model, diag) for diag in values(model.diagnostics) ]
        [ time_to_run(model.clock, out) && write_output(model, out) for out in values(model.output_writers) ]
    end

    return nothing
end

"""
Step forward one time step with a 2nd-order Adams-Bashforth method and pressure-correction
substep.
"""
function time_step!(model::Model{<:AdamsBashforthTimeStepper}, Δt; euler=false)
    χ = ifelse(euler, convert(eltype(model.grid), -0.5), model.timestepper.χ)

    # Convert NamedTuples of Fields to NamedTuples of OffsetArrays
    velocities, tracers, pressures, diffusivities, Gⁿ, G⁻ =
        datatuples(model.velocities, model.tracers, model.pressures, model.diffusivities,
                   model.timestepper.Gⁿ, model.timestepper.G⁻)

    ab2_store_previous_source_terms!(G⁻, model.architecture, model.grid, Gⁿ)

    calculate_explicit_substep!(Gⁿ, velocities, tracers, pressures, diffusivities, model)

    ab2_update_source_terms!(Gⁿ, model.architecture, model.grid, χ, G⁻)

    calculate_pressure_correction!(pressures.pNHS, Δt, Gⁿ, velocities, model)

    complete_pressure_correction_step!(velocities, Δt, tracers, pressures, Gⁿ, model)

    model.clock.time += Δt
    model.clock.iteration += 1

    return nothing
end

#####
##### Adams-Bashforth-specific kernels
#####

""" Store previous source terms for `u`, `v`, and `w` before updating them. """
function ab2_store_previous_velocity_source_terms!(G⁻, grid, Gⁿ)
    @loop_xyz i j k grid begin
        @inbounds G⁻.u[i, j, k] = Gⁿ.u[i, j, k]
        @inbounds G⁻.v[i, j, k] = Gⁿ.v[i, j, k]
        @inbounds G⁻.w[i, j, k] = Gⁿ.w[i, j, k]
    end
    return nothing
end

""" Store previous source terms for a tracer before updating them. """
function ab2_store_previous_tracer_source_term!(Gc⁻, grid, Gcⁿ)
    @loop_xyz i j k grid begin
        @inbounds Gc⁻[i, j, k] = Gcⁿ[i, j, k]
    end
    return nothing
end

""" Store previous source terms before updating them. """
function ab2_store_previous_source_terms!(G⁻, arch, grid, Gⁿ)

    # Velocity fields
    @launch device(arch) config=launch_config(grid, :xyz) ab2_store_previous_velocity_source_terms!(G⁻, grid, Gⁿ)

    # Tracer fields
    for i in 4:length(G⁻)
        @inbounds Gc⁻ = G⁻[i]
        @inbounds Gcⁿ = Gⁿ[i]
        @launch device(arch) config=launch_config(grid, :xyz) ab2_store_previous_tracer_source_term!(Gc⁻, grid, Gcⁿ)
    end

    return nothing
end

"""
Evaluate the right-hand-side terms at time step n+½ using a weighted 2nd-order
Adams-Bashforth method

    `G^{n+½} = (3/2 + χ)G^{n} - (1/2 + χ)G^{n-1}`
"""
function ab2_update_velocity_source_terms!(Gⁿ, grid::AbstractGrid{FT}, χ, G⁻) where FT
    @loop_xyz i j k grid begin
        @inbounds Gⁿ.u[i, j, k] = (FT(1.5) + χ) * Gⁿ.u[i, j, k] - (FT(0.5) + χ) * G⁻.u[i, j, k]
        @inbounds Gⁿ.v[i, j, k] = (FT(1.5) + χ) * Gⁿ.v[i, j, k] - (FT(0.5) + χ) * G⁻.v[i, j, k]
        @inbounds Gⁿ.w[i, j, k] = (FT(1.5) + χ) * Gⁿ.w[i, j, k] - (FT(0.5) + χ) * G⁻.w[i, j, k]
    end

    return nothing
end

"""
Evaluate the right-hand-side terms at time step n+½ using a weighted 2nd-order
Adams-Bashforth method

    `G^{n+½} = (3/2 + χ)G^{n} - (1/2 + χ)G^{n-1}`
"""
function ab2_update_tracer_source_term!(Gcⁿ, grid::AbstractGrid{FT}, χ, Gc⁻) where FT
    @loop_xyz i j k grid begin
        @inbounds Gcⁿ[i, j, k] = (FT(1.5) + χ) * Gcⁿ[i, j, k] - (FT(0.5) + χ) * Gc⁻[i, j, k]
    end
    return nothing
end

"""
Evaluate the right-hand-side terms for velocity fields and tracer fields
at time step n+½ using a weighted 2nd-order Adams-Bashforth method.
"""
function ab2_update_source_terms!(Gⁿ, arch, grid, χ, G⁻)
    # Velocity fields
    @launch device(arch) config=launch_config(grid, :xyz) ab2_update_velocity_source_terms!(Gⁿ, grid, χ, G⁻)

    # Tracer fields
    for i in 4:length(Gⁿ)
        @inbounds Gcⁿ = Gⁿ[i]
        @inbounds Gc⁻ = G⁻[i]
        @launch device(arch) config=launch_config(grid, :xyz) ab2_update_tracer_source_term!(Gcⁿ, grid, χ, Gc⁻)
    end

    return nothing
end