# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

"""
Parameters
----------
dimension
    dimension of surface, 1 for 2d problems (plane strain, plane stress,
    axisymmetric) and 2 for 3d problems. It not given, try to determine problem
    dimension from first element
rotate_normals
    if all surface elements are in cw order instead of ccw, this can be used to
    swap normal directions so that normals point to outward of body
adjust
    for elasticity problems only; closes any gaps between surfaces if found
dual_basis
    use bi-orthogonal basis when interpolating Lagrange multiplier space
use_forwarddiff
    use forwarddiff to linearize contact constraints directly from weighted
    gap function
distval
    charasteristic measure, contact pairs with distance over this value are
    skipped from contact segmentation algorithm
linear_surface_elements
    convert quadratic surface elements to linear elements on the fly, notice
    that middle nodes are missing Lagrange multipliers
split_quadratic_slave_elements
    split quadratic surface elements to several linear sub-elements to get
    Lagrange multiplier to middle nodes also
split_quadratic_master_elements
    split quadratic master elements to several linear sub-elements
store_fields
    not used
"""
mutable struct Mortar <: BoundaryProblem
    dimension :: Int
    rotate_normals :: Bool
    adjust :: Bool
    dual_basis :: Bool
    use_forwarddiff :: Bool
    distval :: Float64
    linear_surface_elements :: Bool
    split_quadratic_slave_elements :: Bool
    split_quadratic_master_elements :: Bool
    alpha :: Float64
    drop_tolerance :: Float64
    store_fields :: Vector{Symbol}
end

function Mortar()
    default_fields = []
    return Mortar(-1, false, false, false, false, Inf, true, true, true, 0.0, 1.0e-9, default_fields)
end

function assemble!(problem::Problem{Mortar}, time::Float64)
    if length(problem.elements) == 0
        @warn("No elements defined in interface $(problem.name), this will result empty assembly!")
        return
    end
    if problem.properties.dimension == -1
        problem.properties.dimension = dim = size(first(problem.elements), 1)
        @info("Assuming dimension of mesh tie surface is $dim. If this is wrong set is manually using problem.properties.dimension")
    end
    dimension = Val{problem.properties.dimension}
    use_forwarddiff = Val{problem.properties.use_forwarddiff}
    assemble!(problem, time, dimension, use_forwarddiff)
end

function get_slave_elements(problem::Problem)
    cond(el) = haskey(el, "master elements") || haskey(el, "potential master elements")
    return filter(cond, get_elements(problem))
end

""" Given a CCW ordered set of vertices, calculate area of polygon.

Examples
--------

julia> P = Vector[[1/3, 5/12, 1/2], [1/3, 1/2, 1/2], [1/2, 1/2, 1/2], [1/2, 1/3, 1/2], [5/12, 1/3, 1/2]]
5-element Array{Array{T,1},1}:
 [0.333333,0.416667,0.5]
 [0.333333,0.5,0.5]
 [0.5,0.5,0.5]
 [0.5,0.333333,0.5]
 [0.416667,0.333333,0.5]

julia> A = calculate_polygon_area(P)
0.02430555555555556

julia> isapprox(A, 7/288)
true

"""
function calculate_polygon_area(P)
    N_P = length(P)
    A = sum([norm(1/2*cross(P[i]-P[1], P[mod(i,N_P)+1]-P[1])) for i=2:N_P])
    return A
end

""" Function to print useful debug information from interface to find bugs. """
function diagnose_interface(problem::Problem{Mortar}, time::Float64)
    @info("Diagnosing Mortar interface...")
    props = problem.properties
    field_dim = get_unknown_field_dimension(problem)
    field_name = get_parent_field_name(problem)
    slave_elements = get_slave_elements(problem)

    I_area = 0.0

    if props.split_quadratic_slave_elements
        @info("props.split_quadratic_slave_elements = true")
        if !props.linear_surface_elements
            @warn("Mortar3D: split_quadratic_surfaces = true and linear_surface_elements = false maybe have unexpected behavior")
        end
        slave_elements = split_quadratic_elements(slave_elements, time)
    end
    @info("Number of slave elements in interface: $(length(slave_elements))")

    # 1. calculate nodal normals and tangents for slave element nodes j ∈ S
    normals = calculate_normals(slave_elements, time, Val{2};
                                rotate_normals=props.rotate_normals)
    update!(slave_elements, "normal", time => normals)

    S_areas = []
    C_areas = []
    P_areas = []

    for slave_element in slave_elements

        @info(repeat("-", 80))
        @info("Processing slave element $(slave_element.id), type = $(get_element_type(slave_element))")
        @info(repeat("-", 80))

        S_area = 0.0
        S_area_in_contact = 0.0
        for ip in get_integration_points(slave_element)
            S_area += ip.weight*slave_element(ip, time, Val{:detJ})
        end
        @info("Total area of slave element = $S_area")

        if props.linear_surface_elements
            @info("Converting slave element to linear surface element")
            slave_element = convert_to_linear_element(slave_element)
        end

        slave_element_nodes = get_connectivity(slave_element)
        @info("Slave element connectivity = $slave_element_nodes")
        nsl = length(slave_element)
        X1 = slave_element("geometry", time)
        n1 = tuple(collect(normals[j] for j in slave_element_nodes)...)

        # project slave nodes to auxiliary plane (x0, Q)
        xi = get_mean_xi(slave_element)
        N = vec(get_basis(slave_element, xi, time))
        x0 = interpolate(N,X1)
        n0 = interpolate(N,n1)
        @info("Auxiliary plane x0 = $x0, n0 = $n0")
        S = Vector[project_vertex_to_auxiliary_plane(X1[i], x0, n0) for i=1:nsl]
        check_orientation!(S, n0)
        @info("Slave element $(slave_element.id) vertices in auxiliary plane: $S")

        # 3. loop all master elements
        master_elements = slave_element("master elements", time)
        if props.split_quadratic_master_elements
            master_elements = split_quadratic_elements(master_elements, time)
        end

        for master_element in master_elements

            if props.linear_surface_elements
                master_element = convert_to_linear_element(master_element)
            end

            master_element_nodes = get_connectivity(master_element)
            nm = length(master_element)
            X2 = master_element("geometry", time)

            if norm(mean(X1) - mean(X2)) > problem.properties.distval
                # elements are "far enough"
                continue
            end

            # 3.1 project master nodes to auxiliary plane and create polygon clipping
            M = Vector[project_vertex_to_auxiliary_plane(X2[i], x0, n0) for i=1:nm]
            check_orientation!(M, n0)
            P = get_polygon_clip(S, M, n0)
            if length(P) < 3
                if length(P) == 0
                    continue
                end
                if length(P) == 1
                    @info("length(P) == 1, shared vertex")
                end
                if length(P) == 2
                    @info("length(P) == 2, shared edge")
                end
                continue
            end
            @info("Master element $(master_element.id) vertices in auxiliary plane = $M")
            check_orientation!(P, n0)
            P_area_ = calculate_polygon_area(P)
            @info("Polygon clip found, P=$P, N_P = $(length(P)), area of polygon = $P_area_")
            if isapprox(P_area_, 0.0)
                error("Polygon P has zero area: $P_area_")
            end

            P_area = 0.0

            C0 = calculate_centroid(P)
            @info("Centroid of polygon = $C0")

            # 4. loop integration cells
            all_cells = get_cells(P, C0)
            @info("Polygon is splitted to $(length(all_cells)) integration cells.")
            for (cell_id, cell) in enumerate(all_cells)
                C_area = 0.0
                virtual_element = Element(Tri3, Int[])
                update!(virtual_element, "geometry", tuple(cell...))

                # 5. loop integration point of integration cell
                for ip in get_integration_points(virtual_element, 3)
                    N = vec(get_basis(virtual_element, ip, time))
                    detJ = virtual_element(ip, time, Val{:detJ})
                    w = ip.weight*detJ
                    # project gauss point from auxiliary plane to master and slave element
                    x_gauss = virtual_element("geometry", ip, time)
                    xi_s, alpha = project_vertex_to_surface(x_gauss, x0, n0, slave_element, X1, time)
                    xi_m, alpha = project_vertex_to_surface(x_gauss, x0, n0, master_element, X2, time)
                    C_area += w
                end # integration points done
                @info("Cell $cell_id has area of $C_area")
                P_area += C_area
                push!(C_areas, C_area)
            end # integration cells done

            if !isapprox(P_area, P_area_)
                error("P_area = $P_area, should be $P_area_")
            end

            S_area_in_contact += P_area
            push!(P_areas, P_area)

        end # master elements done

        S_perc = S_area_in_contact / S_area * 100.0
        push!(S_areas, S_area_in_contact)
        @info("Area of slave element in contact: $S_area_in_contact, it's $S_perc % of total element area")

        I_area += S_area_in_contact

    end # slave elements done, contact virtual work ready

    @info("Area of interface: $I_area")
    @info("Smallest cell area: $(minimum(C_areas))")
    @info("Smallest polygon area: $(minimum(P_areas))")
    @info("Smallest slave element area in contact: $(minimum(S_areas))")

end


# ## Setup
using Pkg
Pkg.activate(joinpath(@__DIR__, "..", "..", ".."))
Pkg.instantiate()
using Dojo

# Examples
Pkg.activate(joinpath(@__DIR__, "..", "..", "..", "examples"))
Pkg.instantiate()
using StaticArrays
using DojoEnvironments
using RobotVisualizer
using Quaternions

# include("methods/initialize.jl")
# ##

# ## visualizer
vis = Visualizer()
open(vis)

mech = DojoEnvironments.get_mechanism(:tugbot; gravity=-1.81, timestep=0.10, radius=0.32)

function ctrl!(mechanism::Mechanism, k::Int; x_target=SVector{3}(0,3,3.0), kp=1.5, kd=2.0)
    drone_body = get_body(mechanism, :drone)
    drone_joint = get_joint(mechanism, :drone_joint)
    x = Dojo.current_position(drone_body.state)
    v, ω = Dojo.current_velocity(drone_body.state)
    θ = mrp(current_orientation(mech.bodies[1].state))
    u_gravity = -drone_body.mass * mechanism.gravity
    u_tra = u_gravity + kp*(x_target - x) - kd .* v
    u_rot = -kp/5 * θ - kd/10 .* ω
    set_input!(drone_joint, [u_tra; u_rot])
    return nothing
end

initialize!(mech, :tugbot, drone_position=[1,0,0])

storage = simulate!(mech, 5.2, ctrl!, record=true, verbose=true)
vis, anim = visualize(mech, storage, vis=vis, show_contact=false)

# add drone mesh
tugbot_visualizer!(vis, mech)

# find the trajectories of the two attach points
contact = get_contact(mech, :body_body)
parent_traj = []
child_traj = []
for i = 1:52
    z = get_maximal_state(storage, i)
    set_maximal_state!(mech, z)
    x_parent, x_child = string_location(mech, contact.model.collision)
    push!(parent_traj, Vector(x_parent))
    push!(child_traj, Vector(x_child))
end

# add rope to visualizer
N = 50
build_rope(vis, N=N, name=:rope2, rope_radius=0.01)
set_loose_rope(vis, [1,-1,1.0], [1,1,1.0], rope_length=10, N=N, min_altitude=-0.0, name=:rope2)

vis, anim = animate_loose_rope(vis, parent_traj, child_traj, rope_length=2.0, N=N,
    anim=anim, name=:rope2, min_altitude=0.02)

convert_frames_to_video_and_gif("ballincup")


# * BitIntegers

module BitIntegers

import Base: &, *, +, -, <, <<, <=, ==, >>, >>>, |, ~, AbstractFloat, add_with_overflow,
             bswap, checked_abs, count_ones, div, flipsign, leading_zeros, mod,
             mul_with_overflow, ndigits0zpb, promote_rule, rem, sub_with_overflow,
             trailing_zeros, typemax, typemin, unsigned, xor

using Base: add_int, and_int, ashr_int, bswap_int, checked_sadd_int, checked_sdiv_int,
            checked_smul_int, checked_srem_int, checked_ssub_int, checked_uadd_int,
            checked_udiv_int, checked_umul_int, checked_urem_int, checked_usub_int, ctlz_int,
            ctpop_int, cttz_int, flipsign_int, lshr_int, mul_int, ndigits0z, ndigits0znb,
            neg_int, not_int, or_int, shl_int, sitofp, sle_int, slt_int, sub_int, uinttype,
            uitofp, ule_int, ult_int, xor_int

using Base.GMP: ispos, Limb

using Core: bitcast, check_top_bit, checked_trunc_sint, checked_trunc_uint, sext_int,
            trunc_int, zext_int

import Random: rand, Sampler
using Random: AbstractRNG, Repetition, SamplerType, LessThan, Masked

# * types definition & aliases

abstract type AbstractBitUnsigned <: Unsigned end
abstract type AbstractBitSigned   <: Signed   end

# ** @define_integers

macro define_integers(n::Int, SI=nothing, UI=nothing)
    if SI === UI === nothing
        SI = Symbol(:Int, n)
        UI = Symbol(:UInt, n)
    elseif SI === nothing || UI === nothing
        throw(ArgumentError(
            "@define_integers requires exactly 1 or 3 arguments, got 2"))
    end

    WSI = Symbol(:Int, 2n) # Wide
    WUI = Symbol(:UInt, 2n)

    WSI = isdefined(__module__, WSI) ? WSI : BigInt
    WUI = isdefined(__module__, WUI) ? WUI : BigInt

    quote
        # `esc` is necessary only on versions < 1.1
        primitive type $(esc(SI)) <: AbstractBitSigned   $n end
        primitive type $(esc(UI)) <: AbstractBitUnsigned $n end

        Base.Signed(x::$(esc(UI)))   = $(esc(SI))(x)
        Base.Unsigned(x::$(esc(SI))) = $(esc(UI))(x)
        Base.uinttype(::Type{$(esc(SI))}) = $(esc(UI))
        Base.uinttype(::Type{$(esc(UI))}) = $(esc(UI))

        Base.widen(::Type{$(esc(SI))}) = $(esc(WSI))
        Base.widen(::Type{$(esc(UI))}) = $(esc(WUI))
    end
end

# ** instanciation of types

const _DEFINED_SIZES = (256, 512, 1024)

# reverse so that widen knows about bigger types first:
for n = reverse(_DEFINED_SIZES)
    SI = Symbol(:Int, n)
    UI = Symbol(:UInt, n)
    sistr = Symbol(:int, n, :_str)
    uistr = Symbol(:uint, n, :_str)
    @eval begin
        @define_integers $n

        macro $sistr(s)
            return parse($SI, s)
        end
        macro $uistr(s)
            return parse($UI, s)
        end

        export $SI, $UI, $(Symbol("@", sistr)), $(Symbol("@", uistr))
    end
end

# ** type lists & aliases

const BitUnsigned_types = # (UInt256, ...)
    tuple((getfield(@__MODULE__, (Symbol(:UInt, n))) for n in _DEFINED_SIZES)...)

const BitSigned_types = # (Int256, ...)
    tuple((getfield(@__MODULE__, (Symbol(:Int, n))) for n in _DEFINED_SIZES)...)

const BitInteger_types = (BitSigned_types..., BitUnsigned_types...)

const BitUnsigned = Union{BitUnsigned_types...}
const BitSigned   = Union{BitSigned_types...}
const BitInteger  = Union{BitInteger_types...}

## convenient abbreviations, only for internal use

const BBS = Base.BitSigned
const BBU = Base.BitUnsigned
const BBI = Base.BitInteger

# eXtended
const XBS = AbstractBitSigned
const XBU = AbstractBitUnsigned
const XBI = Union{XBS,XBU}

# Union, Unified
const UBS = Union{BBS,XBS}
const UBU = Union{BBU,XBU}
const UBI = Union{BBI,XBI}

# ** typemin, typemax

typemin(::Type{T}) where {T<:XBU} = convert(T, 0)
typemax(::Type{T}) where {T<:XBU} = ~convert(T, 0)

typemin(::Type{T}) where {T<:XBS} = convert(T, 1) << (sizeof(T)*8-1)
typemax(::Type{T}) where {T<:XBS} = bitcast(T, typemax(uinttype(T)) >> 1)

# * conversions, promotions

# ** unsigned

unsigned(x::XBS) = reinterpret(typeof(convert(Unsigned, zero(x))), x)

# ** integers

# U -> X
(::Type{T})(x::Union{UBI,Bool}) where {T<:XBI} = convertto(T, x)::T

# X -> B
for T in Base.BitInteger_types
    @eval Core.$(Symbol(:to, T))(x::XBI) = convertto($T, x)
    @eval Base.$(Symbol(T))(x::XBI) = convertto($T, x)::$T
end

@generated function convertto(::Type{T}, x::Union{UBI,Bool}) where {T<:UBI}
    x === Bool && return :(and_int(zext_int(T, x), T(1)))
    if T <: Unsigned
        if x <: Signed
            if sizeof(x) < sizeof(T)
                :(sext_int(T, check_top_bit(x)))
            elseif sizeof(x) == sizeof(T)
                :(bitcast(T, check_top_bit(x)))
            else
                :(checked_trunc_uint(T, x))
            end
        else # x <: Unsigned
            if sizeof(x) < sizeof(T)
                :(zext_int(T, x))
            elseif sizeof(x) == sizeof(T)
                x === T ?
                    :x  :
                    :(reinterpret(T, x))
            else
                :(checked_trunc_uint(T, x))
            end
        end
    else # T <: Signed
        if x <: Signed
            if sizeof(x) < sizeof(T)
                :(sext_int(T, x))
            elseif sizeof(x) == sizeof(T)
                x === T ?
                    :x  :
                    :(reinterpret(T, x))
            else
                :(checked_trunc_sint(T, x))
            end
        else # x <: Unsigned
            if sizeof(x) < sizeof(T)
                :(zext_int(T, x))
            elseif sizeof(x) == sizeof(T)
                :(bitcast(T, check_top_bit(x)))
            else
                :(checked_trunc_sint(T, check_top_bit(x)))
            end
        end
    end
end

@generated function _rem(x::Union{UBI,Bool}, ::Type{to}) where {to<:UBI}
    from = x
    to === from && return :x # this replaces Base's method for BBI
    if to.size < from.size
        :(trunc_int(to, x))
    elseif from === Bool
        :(convert(to, x))
    elseif from.size < to.size
        if from <: Signed
            :(sext_int(to, x))
        else
            :(convert(to, x))
        end
    else
        :(bitcast(to, x))
    end
end

rem(x::Union{UBI,Bool}, ::Type{to}) where {to<:XBI} = _rem(x, to)
rem(x::XBI, ::Type{to}) where {to<:UBI} = _rem(x, to)
# to disambiguate
rem(x::XBI, ::Type{to}) where {to<:XBI} = _rem(x, to)

@generated function promote_rule(::Type{X}, ::Type{Y}) where {X<:XBI,Y<:UBI}
    if X.size > Y.size
        X
    elseif X.size == Y.size
        X <: Unsigned ?
            X :
            Y
    else
        Y
    end
end

# ** BigInt

function rem(x::BigInt, ::Type{T}) where T<:XBI
    if sizeof(T) <= sizeof(Limb)
        iszero(x) ? zero(T) : flipsign(unsafe_load(x.d) % T, x.size)
    else
        u = zero(T)
        for l = 1:min(abs(x.size), cld(sizeof(T), sizeof(Limb)))
            u += (unsafe_load(x.d, l) % T) << ((sizeof(Limb)<<3)*(l-1))
        end
        flipsign(u, x.size)
    end
end

function (::Type{T})(x::BigInt) where T<:XBU
    if sizeof(T) < sizeof(Limb)
        convert(T, convert(Limb,x))
    else
        0 <= x.size <= cld(sizeof(T),sizeof(Limb)) || throw(InexactError(Symbol(string(T)), T, x))
        x % T
    end
end

function (::Type{T})(x::BigInt) where T<:XBS
    n = abs(x.size)
    if sizeof(T) < sizeof(Limb)
        SLimb = typeof(Signed(one(Limb)))
        convert(T, convert(SLimb, x))
    else
        0 <= n <= cld(sizeof(T),sizeof(Limb)) || throw(InexactError(Symbol(string(T)), T, x))
        y = x % T
        ispos(x) ⊻ (y > 0) && throw(InexactError(Symbol(string(T)), T, x)) # catch overflow
        y
    end
end

# ** floats

AbstractFloat(x::XBI) = Float64(x)

for T in (Float32, Float64)
    @eval begin
        (::Type{$T})(x::XBS) = sizeof(x) > 16 ? $T(big(x)) : sitofp($T, x)
        (::Type{$T})(x::XBU) = sizeof(x) > 16 ? $T(big(x)) : uitofp($T, x)
        promote_rule(::Type{$T}, ::Type{<:XBI}) = $T
    end
end

promote_rule(::Type{Float16}, ::Type{<:XBI}) = Float16

# TODO: avoid conversion to BigInt
(::Type{T})(x::AbstractFloat) where {T<:XBI} = T(BigInt(x))::T
# to disambiguate:
(::Type{T})(x::Float16) where {T<:XBI} = T(BigInt(x))::T

# * comparisons

(<)(x::T, y::T) where {T<:XBU} = ult_int(x, y)
(<)(x::T, y::T) where {T<:XBS} = slt_int(x, y)

(<=)(x::T, y::T) where {T<:XBU} = ule_int(x, y)
(<=)(x::T, y::T) where {T<:XBS} = sle_int(x, y)

==(x::UBS, y::UBU) = (x >= 0) & (unsigned(x) == y)
==(x::UBU, y::UBS) = (y >= 0) & (x == unsigned(y))
<( x::UBS, y::UBU) = (x <  0) | (unsigned(x) <  y)
<( x::UBU, y::UBS) = (y >= 0) & (x <  unsigned(y))
<=(x::UBS, y::UBU) = (x <  0) | (unsigned(x) <= y)
<=(x::UBU, y::UBS) = (y >= 0) & (x <= unsigned(y))

# * bit operations

(~)(x::XBI) = not_int(x)
(&)(x::T, y::T) where {T<:XBI} = and_int(x, y)
(|)(x::T, y::T) where {T<:XBI} = or_int(x, y)
xor(x::T, y::T) where {T<:XBI} = xor_int(x, y)

>>( x::UBS, y::UBU) = ashr_int(x, y)
>>( x::UBU, y::UBU) = lshr_int(x, y)
>>>(x::UBI, y::UBU) = lshr_int(x, y)
<<( x::UBI, y::UBU) = shl_int(x, y)

>>( x::UBI, y::Int) = ifelse(0 <= y, x >> unsigned(y),  x << unsigned(-y))
<<( x::UBI, y::Int) = ifelse(0 <= y, x << unsigned(y),  x >> unsigned(-y))
>>>(x::UBI, y::Int) = ifelse(0 <= y, x >>> unsigned(y), x << unsigned(-y))

count_ones(    x::XBI) = Int(ctpop_int(x))
leading_zeros( x::XBI) = Int(ctlz_int(x))
trailing_zeros(x::XBI) = Int(cttz_int(x))

function bswap(x::XBI)
    if sizeof(x) % 2 != 0
        # llvm instruction is invalid
        error("unimplemented")
    else
        bswap_int(x)
    end
end

flipsign(x::T, y::T) where {T<:XBS} = flipsign_int(x, y)

# this doesn't catch flipsign(x::BBS, y::BBS), which is more specific in Base
flipsign(x::UBS, y::UBS) = flipsign_int(promote(x, y)...) % typeof(x)

# * arithmetic operations

(-)(x::XBI)                    = neg_int(x)
(-)(x::T, y::T) where {T<:XBI} = sub_int(x, y)
(+)(x::T, y::T) where {T<:XBI} = add_int(x, y)
(*)(x::T, y::T) where {T<:XBI} = mul_int(x, y)

div(x::XBS, y::Unsigned) = flipsign(signed(div(unsigned(abs(x)), y)), x)
div(x::Unsigned, y::XBS) = unsigned(flipsign(signed(div(x, unsigned(abs(y)))), y))

rem(x::XBS, y::Unsigned) = flipsign(signed(rem(unsigned(abs(x)), y)), x)
rem(x::Unsigned, y::XBS) = rem(x, unsigned(abs(y)))

mod(x::XBS, y::Unsigned) = rem(y + unsigned(rem(x, y)), y)

# these operations fail LLVM for bigger types than UInt128
div(x::T, y::T) where {T<:XBS} = sizeof(T) > 16 ? T(div(big(x), big(y))) : checked_sdiv_int(x, y)
rem(x::T, y::T) where {T<:XBS} = sizeof(T) > 16 ? T(rem(big(x), big(y))) : checked_srem_int(x, y)
div(x::T, y::T) where {T<:XBU} = sizeof(T) > 16 ? T(div(big(x), big(y))) : checked_udiv_int(x, y)
rem(x::T, y::T) where {T<:XBU} = sizeof(T) > 16 ? T(rem(big(x), big(y))) : checked_urem_int(x, y)

# ** checked operations

add_with_overflow(x::T, y::T) where {T<:XBS} = checked_sadd_int(x, y)
add_with_overflow(x::T, y::T) where {T<:XBU} = checked_uadd_int(x, y)

sub_with_overflow(x::T, y::T) where {T<:XBS} = checked_ssub_int(x, y)
sub_with_overflow(x::T, y::T) where {T<:XBU} = checked_usub_int(x, y)

mul_with_overflow(x::T, y::T) where {T<:XBS} = sizeof(T) >= 16 ? broken_mul_with_overflow(x, y) : checked_smul_int(x, y)
mul_with_overflow(x::T, y::T) where {T<:XBU} = sizeof(T) >= 16 ? broken_mul_with_overflow(x, y) : checked_umul_int(x, y)

# cf. base/checked.jl
# TODO: check whether the specific implementation for [U]Int128 is better suited here
function broken_mul_with_overflow(x::T, y::T) where T<:XBS
    r = widemul(x, y)
    f = r % T != r
    r % T, f
end

#= broken
function broken_mul_with_overflow(x::T, y::T) where T<:XBU
    r = widemul(x, y)
    f = r % T != r
    r % T, f
end
=#

function checked_abs(x::XBS)
    r = ifelse(x<0, -x, x)
    r<0 && throw(OverflowError(string("checked arithmetic: cannot compute |x| for x = ", x, "::", typeof(x))))
    r
end

checked_abs(x::XBU) = x

# * misc

function ndigits0zpb(x::XBU, b::Int)
    # precondition: b > 1
    x == 0 && return 0
    b < 0   && return ndigits0znb(signed(x), b)
    b == 2  && return sizeof(x)<<3 - leading_zeros(x)
    b == 8  && return (sizeof(x)<<3 - leading_zeros(x) + 2) ÷ 3
    b == 16 && return sizeof(x)<<1 - leading_zeros(x)>>2
    # b == 10 && return ndigits0z(x) # TODO: implement ndigits0z(x)

    d = 0
    while x > typemax(Int)
        x = div(x,b)
        d += 1
    end
    x = div(x,b)
    d += 1

    m = 1
    while m <= x
        m *= b
        d += 1
    end
    return d
end

ndigits0zpb(x::XBS, b::Integer) = ndigits0zpb(unsigned(abs(x)), Int(b))
ndigits0zpb(x::XBU, b::Integer) = ndigits0zpb(x, Int(b))

# * rand

# ** scalar

rand(rng::AbstractRNG, ::SamplerType{T}) where {T<:XBI} = rand(rng, NBits{T}(sizeof(T) << 3))

# sampler produce a T with at least n random bits (from low to high bits)
struct NBits{T<:XBI} <: Sampler{T}
    n::Int
end

function rand(rng::AbstractRNG, sp::NBits{T}) where {T<:XBI}
    n = (sp.n + 7) >> 3 # bytes
    if n <= 16 # 128 bits
        n > 8 ? rand(rng, UInt128) % T :
        n > 4 ? rand(rng, UInt64)  % T :
        n > 2 ? rand(rng, UInt32)  % T :
        n > 1 ? rand(rng, UInt16)  % T :
                rand(rng, UInt8)   % T
    else
        u = rand(rng, UInt128) % T
        while n > 16
            u = (u << 128) | rand(rng, UInt128)
            n -= 16
        end
        u
    end
end

# ** ranges

# we use SamplerRangeFast as divisions are not yet efficient (for SamplerRangeInt)

Sampler(::Type{<:AbstractRNG}, r::AbstractUnitRange{T}, ::Repetition) where {T<:XBI} =
    SamplerRangeFast(r)

# have to redefine SamplerRangeFast, as it's defined only for Base types in Random
struct SamplerRangeFast{U<:XBU,T<:XBI} <: Sampler{T}
    a::T      # first element of the range
    bw::UInt  # bit width
    m::U      # range length - 1
    mask::U   # mask generated values before threshold rejection
end

SamplerRangeFast(r::AbstractUnitRange{T}) where T<:BitInteger =
    SamplerRangeFast(r, uinttype(T))

function SamplerRangeFast(r::AbstractUnitRange{T}, ::Type{U}) where {T,U}
    isempty(r) && throw(ArgumentError("range must be non-empty"))
    m = (last(r) - first(r)) % uinttype(T) % U # % uinttype(T) to not propagate sign bit
    bw = (sizeof(U) << 3 - leading_zeros(m)) % UInt # bit-width
    mask = ((1 % U) << bw) - (1 % U)
    SamplerRangeFast{U,T}(first(r), bw, m, mask)
end

rand(rng::AbstractRNG, sp::SamplerRangeFast{<:XBI,T}) where {T} =
    rand(rng, LessThan(sp.m, Masked(sp.mask, NBits{T}(sp.bw)))) % T + sp.a

end # module


# from Float64 to Posit8,16,32
Posit8(x::Float64) = ccall((:convertDoubleToP8, SoftPositPath), Posit8, (Float64,),x)
Posit16(x::Float64) = ccall((:convertDoubleToP16, SoftPositPath), Posit16, (Float64,),x)
Posit32(x::Float64) = ccall((:convertDoubleToP32, SoftPositPath), Posit32, (Float64,),x)

# from Float32/16 to Posit8,16,32
Posit8(x::T where {T <: Float16or32}) = Posit8(Float64(x))
Posit16(x::T where {T <: Float16or32}) = Posit16(Float64(x))
Posit32(x::T where {T <: Float16or32}) = Posit32(Float64(x))

# from Float64 to Posit_1 (convertDoubleToPX1 currently not available)
# Posit8_1(x::Float64) = ccall((:convertDoubleToPX1, SoftPositPath), Posit8_1, (Float64,Int64),x,8)
# Posit16_1(x::Float64) = ccall((:convertDoubleToPX1, SoftPositPath), Posit16_1, (Float64,Int64),x,16)
# Posit24_1(x::Float64) = ccall((:convertDoubleToPX1, SoftPositPath), Posit24_1, (Float64,Int64),x,24)

# use detour
Posit8_1(x::Float64) = Posit8_1(Posit32(x))
Posit16_1(x::Float64) = Posit16_1(Posit32(x))
Posit24_1(x::Float64) = Posit24_1(Posit32(x))

# from Float32/16 to Posit_1
Posit8_1(x::T where {T <: Float16or32}) = Posit8_1(Float64(x))
Posit16_1(x::T where {T <: Float16or32}) = Posit16_1(Float64(x))
Posit24_1(x::T where {T <: Float16or32}) = Posit24_1(Float64(x))

# from Float64 to Posit_2
Posit8_2(x::Float64) = ccall((:convertDoubleToPX2, SoftPositPath), Posit8_2, (Float64,Int64),x,8)
Posit16_2(x::Float64) = ccall((:convertDoubleToPX2, SoftPositPath), Posit16_2, (Float64,Int64),x,16)
Posit24_2(x::Float64) = ccall((:convertDoubleToPX2, SoftPositPath), Posit24_2, (Float64,Int64),x,24)

# from Float32/16 to Posit_2
Posit8_2(x::T where {T <: Float16or32}) = Posit8_2(Float64(x))
Posit16_2(x::T where {T <: Float16or32}) = Posit16_2(Float64(x))
Posit24_2(x::T where {T <: Float16or32}) = Posit24_2(Float64(x))


#NOTE: These are helper functions to get data from JuMP models into the PIPS-NLP Interface.
#The implementation here was heavily framed around the code in Ipopt.jl

mutable struct ConstraintData
    linear_le_constraints::Vector{JuMP.ScalarConstraint{JuMP.GenericAffExpr{Float64,JuMP.VariableRef},MathOptInterface.LessThan{Float64}}}
    linear_ge_constraints::Vector{JuMP.ScalarConstraint{JuMP.GenericAffExpr{Float64,JuMP.VariableRef},MathOptInterface.GreaterThan{Float64}}}
	linear_interval_constraints::Vector{JuMP.ScalarConstraint{JuMP.GenericAffExpr{Float64,JuMP.VariableRef},MathOptInterface.Interval{Float64}}}
    linear_eq_constraints::Vector{JuMP.ScalarConstraint{JuMP.GenericAffExpr{Float64,JuMP.VariableRef},MathOptInterface.EqualTo{Float64}}}
    quadratic_le_constraints::Vector{JuMP.ScalarConstraint{JuMP.GenericQuadExpr{Float64,JuMP.VariableRef},MathOptInterface.LessThan{Float64}}}
    quadratic_ge_constraints::Vector{JuMP.ScalarConstraint{JuMP.GenericQuadExpr{Float64,JuMP.VariableRef},MathOptInterface.GreaterThan{Float64}}}
	quadratic_interval_constraints::Vector{JuMP.ScalarConstraint{JuMP.GenericQuadExpr{Float64,JuMP.VariableRef},MathOptInterface.Interval{Float64}}}
    quadratic_eq_constraints::Vector{JuMP.ScalarConstraint{JuMP.GenericQuadExpr{Float64,JuMP.VariableRef},MathOptInterface.EqualTo{Float64}}}
	nonlinear_constraints::Vector{JuMP._NonlinearConstraint}
end

ConstraintData() = ConstraintData([], [], [], [], [], [],[],[],[])

function get_constraint_data(node::OptiNode)
	con_data = ConstraintData()
	constraint_types = JuMP.list_of_constraint_types(node)

    for (func,set) in constraint_types
        if func == JuMP.VariableRef 	#This is a variable bound, not a PIPS-NLP constraint
			continue
		else
	        constraint_refs = JuMP.all_constraints(node, func, set)
	        for constraint_ref in constraint_refs
	            constraint = JuMP.constraint_object(constraint_ref)

				func_type = typeof(constraint.func)
				con_type = typeof(constraint.set)

				if func_type == JuMP.GenericAffExpr{Float64,JuMP.VariableRef}
					if con_type == MOI.LessThan{Float64}
						push!(con_data.linear_le_constraints,constraint)
					elseif con_type == MOI.GreaterThan{Float64}
						push!(con_data.linear_ge_constraints,constraint)
					elseif con_type == MOI.Interval{Float64}
						push!(con_data.linear_interval_constraints,constraint)
					elseif con_type == MOI.EqualTo{Float64}
						push!(con_data.linear_eq_constraints,constraint)
					end
				elseif func_type == JuMP.GenericQuadExpr{Float64,JuMP.VariableRef}
					if con_type == MOI.LessThan{Float64}
						push!(con_data.quadratic_le_constraints,constraint)
					elseif con_type == MOI.GreaterThan{Float64}
						push!(con_data.quadratic_ge_constraints,constraint)
					elseif con_type == MOI.Interval{Float64}
						push!(con_data.quadratic_interval_constraints,constraint)
					elseif con_type == MOI.EqualTo{Float64}
						push!(con_data.quadtratic_eq_constraints,constraint)
					end
	            else
	                error("Could not figure out constraint type for $(constraint_ref)")
	            end
	        end
        end
    end
	if node.nlp_data != nothing
		con_data.nonlinear_constraints = node.nlp_data.nlconstr
	end
	return con_data
end

linear_le_offset(constraints::ConstraintData) = 0
linear_ge_offset(constraints::ConstraintData) = length(constraints.linear_le_constraints)
linear_interval_offset(constraints::ConstraintData) = linear_ge_offset(constraints) + length(constraints.linear_ge_constraints)
linear_eq_offset(constraints::ConstraintData) = linear_interval_offset(constraints) + length(constraints.linear_interval_constraints)
quadratic_le_offset(constraints::ConstraintData) = linear_eq_offset(constraints) + length(constraints.linear_eq_constraints)
quadratic_ge_offset(constraints::ConstraintData) = quadratic_le_offset(constraints) + length(constraints.quadratic_le_constraints)
quadratic_interval_offset(constraints::ConstraintData) =  quadratic_ge_offset(constraints) + length(constraints.quadratic_ge_constraints)
quadratic_eq_offset(constraints::ConstraintData) = quadratic_interval_offset(constraints) + length(constraints.quadratic_interval_constraints)
nlp_constraint_offset(constraints::ConstraintData) = quadratic_eq_offset(constraints) + length(constraints.quadratic_eq_constraints)

function numconstraints(con_data::ConstraintData)
	fields = fieldnames(ConstraintData)
	return sum([length(getfield(con_data,field)) for field in fields])
end

function constraintbounds(con_data::ConstraintData)
	num_cons = numconstraints(con_data)
	# Setup indices.  Need to get number of constraints
	constraint_lower = ones(num_cons)*-Inf
	constraint_upper = ones(num_cons)*Inf

	#linear_le_constraints
	constraints = con_data.linear_le_constraints
	offset = linear_le_offset(con_data)
	constraint_upper[offset + 1:offset + length(constraints)] .= [constraints[i].set.upper for i in 1:length(constraints)]

	#linear_ge_constraints
	constraints = con_data.linear_ge_constraints
	offset = linear_ge_offset(con_data)
	constraint_lower[offset + 1:offset + length(constraints)] .= [constraints[i].set.lower for i in 1:length(constraints)]

	#linear_interval_constraints
	constraints = con_data.linear_interval_constraints
	offset = linear_interval_offset(con_data)
	constraint_lower[offset + 1:offset + length(constraints)] .= [constraints[i].set.lower for i in 1:length(constraints)]
	constraint_upper[offset + 1:offset + length(constraints)] .= [constraints[i].set.upper for i in 1:length(constraints)]

	#linear_eq_constraints
	constraints = con_data.linear_eq_constraints
	offset = linear_eq_offset(con_data)
	constraint_lower[offset + 1:offset + length(constraints)] .= [constraints[i].set.value for i in 1:length(constraints)]
	constraint_upper[offset + 1:offset + length(constraints)] .= [constraints[i].set.value for i in 1:length(constraints)]

	#quadratic_le_constraints
	constraints = con_data.quadratic_le_constraints
	offset = quadratic_le_offset(con_data)
	constraint_upper[offset + 1:offset + length(constraints)] .= [constraints[i].set.upper for i in 1:length(constraints)]

	#quadratic_ge_constraints
	constraints = con_data.quadratic_ge_constraints
	offset = quadratic_ge_offset(con_data)
	constraint_lower[offset + 1:offset + length(constraints)] .= [constraints[i].set.lower for i in 1:length(constraints)]

	#quadratic_interval_constraints
	constraints = con_data.quadratic_interval_constraints
	offset = quadratic_interval_offset(con_data)
	constraint_lower[offset + 1:offset + length(constraints)] .= [constraints[i].set.lower for i in 1:length(constraints)]
	constraint_upper[offset + 1:offset + length(constraints)] .= [constraints[i].set.upper for i in 1:length(constraints)]

	#quadratic_eq_constraints
	constraints = con_data.quadratic_eq_constraints
	offset = quadratic_eq_offset(con_data)
	constraint_lower[offset + 1:offset + length(constraints)] .= [constraints[i].set.value for i in 1:length(constraints)]
	constraint_upper[offset + 1:offset + length(constraints)] .= [constraints[i].set.value for i in 1:length(constraints)]

	#nonlinear constraints
	constraints = con_data.nonlinear_constraints
	offset = nlp_constraint_offset(con_data)
	constraint_lower[offset + 1:offset + length(constraints)] .= [constraints[i].lb for i in 1:length(constraints)]
	constraint_upper[offset + 1:offset + length(constraints)] .= [constraints[i].ub for i in 1:length(constraints)]

	return constraint_lower,constraint_upper
end

function numconstraints(node::OptiNode)
    num_cons = 0
    constraint_types = JuMP.list_of_constraint_types(node)
    for (func,set) in constraint_types
        if func != JuMP.VariableRef #This is a variable bound, not a PIPS-NLP constraint
            num_cons += JuMP.num_constraints(node,func,set)
        end
    end
    num_cons += JuMP.num_nl_constraints(node)
    return num_cons
end

function variableupperbounds(node::OptiNode)
    #Get upper bound variable constraints
    upper_bounds = ones(JuMP.num_variables(node))*Inf #Assume no upper bound by default

    var_bound_constraints = JuMP.all_constraints(node,JuMP.VariableRef,MOI.LessThan{Float64})
    for var_bound_ref in var_bound_constraints
        var_constraint = JuMP.constraint_object(var_bound_ref)
        var = var_constraint.func
        index = var.index
        upper_bound = var_constraint.set.upper
        upper_bounds[index.value] = upper_bound
    end
    #Get fixed variable constraints
    var_equal_constraints = JuMP.all_constraints(node,JuMP.VariableRef,MOI.EqualTo{Float64})
    for var_equal_ref in var_equal_constraints
        var_constraint = JuMP.constraint_object(var_equal_ref)
        var = var_constraint.func
        index = var.index
        upper_bound = var_constraint.set.value
        upper_bounds[index.value] = upper_bound
    end

    #TODO sort?
    return upper_bounds
end

function variablelowerbounds(node::OptiNode)
    lower_bounds = ones(JuMP.num_variables(node))*-Inf #Assume no upper bound by default
    var_bound_constraints = JuMP.all_constraints(node,JuMP.VariableRef,MOI.GreaterThan{Float64})
    for var_bound_ref in var_bound_constraints
        var_constraint = JuMP.constraint_object(var_bound_ref)
        var = var_constraint.func
        index = var.index
        lower_bound = var_constraint.set.lower
        lower_bounds[index.value] = lower_bound
    end

    var_equal_constraints = JuMP.all_constraints(node,JuMP.VariableRef,MOI.EqualTo{Float64})
    for var_equal_ref in var_equal_constraints
        var_constraint = JuMP.constraint_object(var_equal_ref)
        var = var_constraint.func
        index = var.index
        lower_bound = var_constraint.set.value
        lower_bounds[index.value] = lower_bound
    end

    #TODO sort?
    return lower_bounds
end

#####################################
# OBJECTIVE FUNCTION EVALUATION
#####################################
#x is an array of variable values returned from the solver
function eval_function(var::JuMP.VariableRef, x)
    return x[var.index.value]
end

function eval_function(aff::JuMP.GenericAffExpr{Float64,JuMP.VariableRef}, x)
    function_value = aff.constant
    for (var,coeff) in aff.terms
        # NOTE the implicit assumtion that VariableIndex values match up with
        # x indices. This is valid because in this wrapper ListOfVariableIndices
        # is always [1, ..., NumberOfVariables].
        function_value += coeff*x[var.index.value]
    end
    return function_value
end

function eval_function(quad::JuMP.GenericQuadExpr{Float64,JuMP.VariableRef}, x)
    function_value = quad.aff.constant
    for (var,coeff) in quad.aff.terms
        function_value += coeff*x[var.index.value]
    end
    for (terms,coeff) in quad.terms
        row_idx = terms.a.index
        col_idx = terms.b.index
		function_value += coeff*x[row_idx.value]*x[col_idx.value]
    end
    return function_value
end

function pips_eval_objective(d::JuMP.NLPEvaluator, x)
    # The order of the conditions is important. NLP objectives override regular objectives.
	m = d.m
    if has_nl_objective(m)
        return MOI.eval_objective(d, x)
    elseif JuMP.objective_function(m) !== nothing
        return eval_function(JuMP.objective_function(m), x)
    else
        return 0.0 # No objective function set. This could happen with FEASIBILITY_SENSE.
    end
end

#####################################
# OBJECTIVE GRADIENT EVALUATION
#####################################
function fill_gradient!(grad, x, var::JuMP.VariableRef)
    fill!(grad, 0.0)
    grad[var.index.value] = 1.0
	return
end

function fill_gradient!(grad, x, aff::JuMP.GenericAffExpr{Float64,JuMP.VariableRef})
    fill!(grad, 0.0)
	for	(var,coeff) in aff.terms
        grad[var.index.value] += coeff
    end
	return
end

function fill_gradient!(grad, x, quad::JuMP.GenericQuadExpr{Float64,JuMP.VariableRef})
    fill!(grad, 0.0)
    for	(var,coeff) in quad.aff.terms
        grad[var.index.value] += coeff
    end
    #for term in quad.quadratic_terms
	for (terms,coeff) in quad.terms
        row_idx = terms.a.index
        col_idx = terms.b.index
        if row_idx == col_idx
            grad[row_idx.value] += 2*coeff*x[row_idx.value]
        else
            grad[row_idx.value] += coeff*x[col_idx.value]
            grad[col_idx.value] += coeff*x[row_idx.value]
        end
    end
	return
end

function pips_eval_objective_gradient(d::JuMP.NLPEvaluator, grad, x)
	m = d.m
    if has_nl_objective(m)
        MOI.eval_objective_gradient(d, grad, x)
    elseif JuMP.objective_function(m) !== nothing
        fill_gradient!(grad, x, JuMP.objective_function(m))
    else
        fill!(grad, 0.0)
    end
    return
end

#####################################
# CONSTRAINT EVALUATION
#####################################
# Refers to local variables in eval_constraint() below.
macro eval_function(array_name)
    escrow = esc(:row)
    quote
        for constraint in $(esc(array_name))
			func = constraint.func
            $(esc(:g))[$escrow] = eval_function(func, $(esc(:x)))
            $escrow += 1
        end
    end
end

function pips_eval_constraint(d::JuMP.NLPEvaluator, g, x)
    row = 1
	con_data = getnode(d.m).ext[:constraint_data]
    @eval_function con_data.linear_le_constraints
    @eval_function con_data.linear_ge_constraints
	@eval_function con_data.linear_interval_constraints
    @eval_function con_data.linear_eq_constraints
    @eval_function con_data.quadratic_le_constraints
    @eval_function con_data.quadratic_ge_constraints
	@eval_function con_data.quadratic_interval_constraints
    @eval_function con_data.quadratic_eq_constraints
    nlp_g = view(g, row:length(g))
    MOI.eval_constraint(d, nlp_g, x)
    return
end

#####################
# JACOBIAN STRUCTURE
#####################
function append_to_jacobian_sparsity!(jacobian_sparsity, func::JuMP.GenericAffExpr{Float64,JuMP.VariableRef}, row)
	aff = func
    for term in keys(aff.terms)
        push!(jacobian_sparsity, (row, term.index.value))
    end
end

function append_to_jacobian_sparsity!(jacobian_sparsity, func::JuMP.GenericQuadExpr{Float64,JuMP.VariableRef}, row)
	quad = func
    for term in keys(quad.aff.terms)
        push!(jacobian_sparsity, (row, term.index.value))
    end
    for term in keys(quad.terms)
        row_idx = term.a.index
        col_idx = term.b.index
        if row_idx == col_idx
            push!(jacobian_sparsity, (row, row_idx.value))
        else
            push!(jacobian_sparsity, (row, row_idx.value))
            push!(jacobian_sparsity, (row, col_idx.value))
        end
    end
end

macro append_to_jacobian_sparsity(array_name)
    escrow = esc(:row)
    quote
        for constraint in $(esc(array_name))
			func = constraint.func
            append_to_jacobian_sparsity!($(esc(:jacobian_sparsity)), func, $escrow)
            $escrow += 1
        end
    end
end

function pips_jacobian_structure(d::JuMP.NLPEvaluator)
    num_nlp_constraints = length(d.m.nlp_data.nlconstr)
    if num_nlp_constraints > 0
		MOI.initialize(d,[:Grad,:Jac,:Hess])
        nlp_jacobian_sparsity = MOI.jacobian_structure(d)
    else
        nlp_jacobian_sparsity = []
    end

    jacobian_sparsity = Tuple{Int64,Int64}[]
    row = 1
	con_data = get_constraint_data(getnode(d.m))

    @append_to_jacobian_sparsity con_data.linear_le_constraints
    @append_to_jacobian_sparsity con_data.linear_ge_constraints
	@append_to_jacobian_sparsity con_data.linear_interval_constraints
    @append_to_jacobian_sparsity con_data.linear_eq_constraints
    @append_to_jacobian_sparsity con_data.quadratic_le_constraints
    @append_to_jacobian_sparsity con_data.quadratic_ge_constraints
	@append_to_jacobian_sparsity con_data.quadratic_interval_constraints
    @append_to_jacobian_sparsity con_data.quadratic_eq_constraints
    for (nlp_row, column) in nlp_jacobian_sparsity
        push!(jacobian_sparsity, (nlp_row + row - 1, column))
    end
    return jacobian_sparsity
end

#####################
# HESSIAN STRUCTURE
#####################
function has_nl_objective(m::JuMP.Model)
	if m.nlp_data == nothing
		return false
	elseif m.nlp_data.nlobj != nothing
		return true
	else
		return false
	end
end

append_to_hessian_sparsity!(hessian_sparsity, ::Union{JuMP.VariableRef,JuMP.GenericAffExpr}) = nothing

function append_to_hessian_sparsity!(hessian_sparsity, quad::JuMP.GenericQuadExpr{Float64,JuMP.VariableRef})
    for term in keys(quad.terms)
        push!(hessian_sparsity, (term.a.index.value,term.b.index.value))
    end
end

#NOTE: there can be duplicate entries
function pips_hessian_lagrangian_structure(d::JuMP.NLPEvaluator)
    hessian_sparsity = Tuple{Int64,Int64}[]
	m = d.m
	con_data = get_constraint_data(getnode(m))
	if m.nlp_data == nothing
	    append_to_hessian_sparsity!(hessian_sparsity, JuMP.objective_function(m))
	elseif m.nlp_data.nlobj == nothing
		append_to_hessian_sparsity!(hessian_sparsity, JuMP.objective_function(m))
	end
    for constraint in con_data.quadratic_le_constraints
		quad = constraint.func
        append_to_hessian_sparsity!(hessian_sparsity, quad)
    end
    for constraint in con_data.quadratic_ge_constraints
		quad = constraint.func
        append_to_hessian_sparsity!(hessian_sparsity, quad)
    end
	for constraint in con_data.quadratic_interval_constraints
		quad = constraint.func
		append_to_hessian_sparsity!(hessian_sparsity, quad)
	end
    for constraint in con_data.quadratic_eq_constraints
		quad = constraint.func
        append_to_hessian_sparsity!(hessian_sparsity, quad)
    end
    nlp_hessian_sparsity = MOI.hessian_lagrangian_structure(d)
    append!(hessian_sparsity, nlp_hessian_sparsity)
    return hessian_sparsity
end

#####################################
# JACOBIAN EVALUATION
#####################################
function fill_constraint_jacobian!(jac_values, start_offset, x, aff::JuMP.GenericAffExpr{Float64,JuMP.VariableRef})
    num_coefficients = length(aff.terms)
    #for i in 1:num_coefficients
	i = 1
	for coeff in values(aff.terms) #ordered dictionary
        jac_values[start_offset+i] = coeff
		i += 1
    end
    return num_coefficients
end

function fill_constraint_jacobian!(jac_values, start_offset, x, quad::JuMP.GenericQuadExpr{Float64,JuMP.VariableRef})
	aff = quad.aff
    num_affine_coefficients = length(aff.terms)
	i = 1
	for coeff in values(aff.terms)
        jac_values[start_offset+i] = coeff #quad.affine_terms[i].coefficient
		i += 1
    end
    num_quadratic_coefficients = 0

	for (terms,coeff) in quad.terms
        row_idx = terms.a.index
        col_idx = terms.b.index
        if row_idx == col_idx
            jac_values[start_offset+num_affine_coefficients+num_quadratic_coefficients+1] = 2*coeff*x[row_idx.value]
            num_quadratic_coefficients += 1
        else
            # Note that the order matches the Jacobian sparsity pattern.
            jac_values[start_offset+num_affine_coefficients+num_quadratic_coefficients+1] = coeff*x[col_idx.value]
            jac_values[start_offset+num_affine_coefficients+num_quadratic_coefficients+2] = coeff*x[row_idx.value]
            num_quadratic_coefficients += 2
        end
    end
    return num_affine_coefficients + num_quadratic_coefficients
end

# Refers to local variables in eval_constraint_jacobian() below.
macro fill_constraint_jacobian(array_name)
    esc_offset = esc(:offset)
    quote
        for constraint in $(esc(array_name))
			func = constraint.func
            $esc_offset += fill_constraint_jacobian!($(esc(:jac_values)),
                                                     $esc_offset, $(esc(:x)),
                                                     func)
        end
    end
end

function pips_eval_constraint_jacobian(d::JuMP.NLPEvaluator, jac_values, x)
    offset = 0
	m = d.m
	con_data = getnode(m).ext[:constraint_data]
    @fill_constraint_jacobian con_data.linear_le_constraints
    @fill_constraint_jacobian con_data.linear_ge_constraints
	@fill_constraint_jacobian con_data.linear_interval_constraints
    @fill_constraint_jacobian con_data.linear_eq_constraints
    @fill_constraint_jacobian con_data.quadratic_le_constraints
    @fill_constraint_jacobian con_data.quadratic_ge_constraints
	@fill_constraint_jacobian con_data.quadratic_interval_constraints
    @fill_constraint_jacobian con_data.quadratic_eq_constraints

    nlp_values = view(jac_values, 1+offset:length(jac_values))
    MOI.eval_constraint_jacobian(d, nlp_values, x)
    return
end

##########################################
# HESSIAN OF THE LAGRANGIAN EVALUATION
#########################################
function fill_hessian_lagrangian!(hess_values, start_offset, scale_factor,::Union{JuMP.VariableRef,JuMP.GenericAffExpr{Float64,JuMP.VariableRef},Nothing})
    return 0
end

function fill_hessian_lagrangian!(hess_values, start_offset, scale_factor, quad::JuMP.GenericQuadExpr{Float64,JuMP.VariableRef})
	i = 1
	for (terms,coeff) in quad.terms
		row_idx = terms.a.index
		col_idx = terms.b.index
		if row_idx == col_idx
			hess_values[start_offset + i] = 2*scale_factor*coeff
		else
			hess_values[start_offset + i] = scale_factor*coeff
		end
		i += 1
	end
    return length(quad.terms)
end

function pips_eval_hessian_lagrangian(d::JuMP.NLPEvaluator, hess_values, x, obj_factor, lambda)
    offset = 0
	m = d.m
	con_data = getnode(m).ext[:constraint_data]
	if !(has_nl_objective(m))
        offset += fill_hessian_lagrangian!(hess_values, 0, obj_factor, JuMP.objective_function(m))
    end
    for (i, constraint) in enumerate(con_data.quadratic_le_constraints)
		quad = constraint.func
        offset += fill_hessian_lagrangian!(hess_values, offset, lambda[i+quadratic_le_offset(con_data)], quad)
    end
    for (i, constraint) in enumerate(con_data.quadratic_ge_constraints)
		quad = constraint.func
        offset += fill_hessian_lagrangian!(hess_values, offset, lambda[i+quadratic_ge_offset(con_data)], quad)
    end
	for (i, constraint) in enumerate(con_data.quadratic_interval_constraints)
		quad = constraint.func
        offset += fill_hessian_lagrangian!(hess_values, offset, lambda[i+quadratic_interval_offset(con_data)], quad)
    end
    for (i, constraint) in enumerate(con_data.quadratic_eq_constraints)
		quad = constraint.func
        offset += fill_hessian_lagrangian!(hess_values, offset, lambda[i+quadratic_eq_offset(con_data)], quad)
    end
    nlp_values = view(hess_values, 1 + offset : length(hess_values))
    nlp_lambda = view(lambda, 1 + nlp_constraint_offset(con_data) : length(lambda))
    MOI.eval_hessian_lagrangian(d, nlp_values, x, obj_factor, nlp_lambda)
end


using OffsetArrays: OffsetArray
using Oceananigans.Architectures: device_event

#####
##### General halo filling functions
#####

fill_halo_regions!(::Nothing, args...) = nothing

"""
    fill_halo_regions!(fields::Union{Tuple, NamedTuple}, arch, args...)

Fill halo regions for each field in the tuple `fields` according to their boundary
conditions, possibly recursing into `fields` if it is a nested tuple-of-tuples.
"""
function fill_halo_regions!(fields::Union{Tuple, NamedTuple}, arch, args...)

    for field in fields
        fill_halo_regions!(field, arch, args...)
    end

    return nothing
end

# Some fields have `nothing` boundary conditions, such as `FunctionField` and `ZeroField`.
fill_halo_regions!(c::OffsetArray, ::Nothing, args...; kwargs...) = nothing

"Fill halo regions in x, y, and z for a given field's data."
function fill_halo_regions!(c::OffsetArray, field_bcs, arch, grid, args...; kwargs...)

    fill_halos! = [
        fill_west_and_east_halo!,
        fill_south_and_north_halo!,
        fill_bottom_and_top_halo!,
    ]

    field_bcs_array_left = [
        field_bcs.west,
        field_bcs.south,
        field_bcs.bottom,
    ]

    field_bcs_array_right = [
        field_bcs.east,
        field_bcs.north,
        field_bcs.top,
    ]

    perm = sortperm(field_bcs_array_left, lt=fill_first)
    fill_halos! = fill_halos![perm]
    field_bcs_array_left  = field_bcs_array_left[perm]
    field_bcs_array_right = field_bcs_array_right[perm]

    for task = 1:3

        barrier = device_event(arch)

        fill_halo!  = fill_halos![task]
        bc_left     = field_bcs_array_left[task]
        bc_right    = field_bcs_array_right[task]

        events      = fill_halo!(c, bc_left, bc_right, arch, barrier, grid, args...; kwargs...)

        wait(device(arch), events)

    end

    return nothing
end

@inline validate_event(::Nothing) = NoneEvent()
@inline validate_event(event) = event

# Fallbacks split into two calls
function fill_west_and_east_halo!(c, west_bc, east_bc, args...; kwargs...)
     west_event = validate_event(fill_west_halo!(c, west_bc, args...; kwargs...))
     east_event = validate_event(fill_east_halo!(c, east_bc, args...; kwargs...))
    multi_event = MultiEvent((west_event, east_event))
    return multi_event
end

function fill_south_and_north_halo!(c, south_bc, north_bc, args...; kwargs...)
    south_event = validate_event(fill_south_halo!(c, south_bc, args...; kwargs...))
    north_event = validate_event(fill_north_halo!(c, north_bc, args...; kwargs...))
    multi_event = MultiEvent((south_event, north_event))
    return multi_event
end

function fill_bottom_and_top_halo!(c, bottom_bc, top_bc, args...; kwargs...)
    bottom_event = validate_event(fill_bottom_halo!(c, bottom_bc, args...; kwargs...))
       top_event = validate_event(fill_top_halo!(c, top_bc, args...; kwargs...))
     multi_event = MultiEvent((bottom_event, top_event))
     return multi_event
end

#####
##### Halo filling order
#####

fill_first(bc1::PBC, bc2)      = false
fill_first(bc1, bc2::PBC)      = true
fill_first(bc1::PBC, bc2::PBC) = true
fill_first(bc1, bc2)           = true


module stuff
  using GaussianProcesses
end