# working with images ...
#
# (c) A. Dominik, 2020

const UNKNOWN_CLASS = "unknwon_class"

"""
    function mk_image_minibatch(dir, batchsize; split=false, fr=0.2,
                                balanced=false, shuffle=true, train=true,
                                pre_load=false,
                                aug_pipl=nothing, pre_proc=nothing)

Return one or two iterable image-loader-objects that provides
minibatches of images. For training each minibatch is a tupel
`(x,y)` with x: 4-d-array with the minibatch of data and y:
vector of class IDs as Int.

### Arguments:
+ `dir`: base-directory of the image dataset. The first level of
        sub-dirs are used as class names.
+ `batchsize`: size of minibatches

### Keyword arguments:
+ `split`: return two iterators for training and validation
+ `fr`: split fraction
+ `balanced`: return balanced data (i.e. same number of instances
        for all classes). Balancing is achieved via oversampling
+ `shuffle`: if true, shuffle the images everytime the iterator
        restarts
+ `train`: if true, minibatches with (x,y) Tuples are provided,
        if false only x (for prediction)
+ `pre_load`: if `true` all images are loaded in advance;
        otherwise images are loaded on demand durng training.
        (option is *not implemented yet!*)
+ `aug_pipl`: augmentation pipeline for Augmentor.jl. Augmentation
        is performed before the pre_proc-function is applied
+ `pre_proc`: function with preprocessing
        and augmentation algoritms of type x = f(x). In contrast
        to the augmentation that modifies images, is `pre_proc`
        working on Arrays{Float32}.
+ `pre_load=false`: read all images from disk once when populating the
        loader (requires loads of memory, but speeds up training).
"""
function mk_image_minibatch(dir, batchsize; split=false, fr=0.5,
                            balanced=false, shuffle=true, train=true,
                            pre_load=false,
                            aug_pipl=nothing, pre_proc=nothing)

    i_paths = get_files_list(dir)
    i_n = length(i_paths)

    if train
        i_class_names = get_class_names(dir, i_paths)
        classes = unique(i_class_names)
        i_classes = [findall(x->x==c, classes)[1] for c in i_class_names]
    else
        i_class_names = fill(UNKNOWN_CLASS, i_n)
        classes =[UNKNOWN_CLASS]
        i_classes = fill(0, i_n)
    end

    if split                    # return train_loader, valid_loader
        ((xvld,yvld),(xtrn,ytrn)) = do_split(i_paths, i_classes, at=fr)
        if balanced
            (xtrn,ytrn) = do_balance(xtrn, ytrn)
            (xvld,yvld) = do_balance(xvld, yvld)
        end
        trn_loader = ImageLoader(dir, xtrn, ytrn, classes,
                            batchsize, shuffle, train,
                            aug_pipl, pre_proc,
                            pre_load)

        vld_loader = ImageLoader(dir, xvld, yvld, classes,
                            batchsize, shuffle, train,
                            aug_pipl, pre_proc,
                            pre_load)
        return trn_loader, vld_loader
    else
        xtrn, ytrn = i_paths, i_classes
        if balanced
            (xtrn, ytrn) = do_balance(i_paths, i_classes)
        end
        trn_loader = ImageLoader(dir, xtrn, ytrn, classes,
                            batchsize, shuffle, train,
                            aug_pipl, pre_proc,
                            pre_load)
        return trn_loader
    end
end

"""
    function get_class_labels(d::DataLoader)

Extracts a list of class labels from a DataLoader.
"""
function get_class_labels(dl::DataLoader)
    return dl.classes
end

"""
    struct ImageLoader <: DataLoader
        dir
        i_paths
        i_classes
        classes
        batchsize
        shuffle
        train
        aug_pipl
        pre_proc
        pre_load
        i_images
    end

Iterable image loader to provide minibatches of images as
4-d-arrays (x,y,rgb,mb).
"""
mutable struct ImageLoader <: DataLoader
    dir                 # root dir
    i_paths             # list of all file image_paths
    i_classes           # list of classs IDs for each image
    classes             # unique list of class names
    batchsize           #
    shuffle             # if true: shuffle for each start
    i_sequence          # actual sequrence of images to take
    train               # if true: inlcude y in minibatches
    aug_pipl            # Augmentor.jl pipeline
    pre_proc            # function to process one 3d image tensor (RGB)
    pre_load            # load all images on init
    i_images            # list of all images; nothing if not predolad.

    function ImageLoader(dir, i_paths, i_classes, classes,
                         batchsize, shuffle, train,
                         aug_pipl, pre_proc, pre_load)

        i_seq = collect(1:length(i_paths))
        if pre_load
            i_images = pre_load_images(i_paths)
        else
            i_images = nothing
        end
        return new(dir, i_paths, i_classes, classes,
                    batchsize, shuffle, i_seq,
                    train,
                    aug_pipl, pre_proc,
                    pre_load, i_images)
    end
end

# two iterate funs for ImageLoader
#
function Base.iterate(il::ImageLoader)

    if il.shuffle
        Random.shuffle!(il.i_sequence)
    end
    state = 1
    return iterate(il, state)
end

# state is the index of the next image after the currect batch.
#
function Base.iterate(il::ImageLoader, state)

    # check if empty:
    #
    if state > length(il.i_paths)
        return nothing
    end

    # range for next minibatch:
    #
    n = length(il.i_paths)
    mb_start = state
    mb_size = mb_start + il.batchsize > n ? n-mb_start+1 : il.batchsize
    state = mb_start+il.batchsize

    return mk_image_mb(il, mb_start, mb_size), state
end

function mk_image_mb(il, mb_start, mb_size)

    if mb_size < 1
        return nothing
    end

    # nice way:
    image_ns = mb_start:mb_start+mb_size-1

    # make image index form shuffled list entry:
    #
    image_ns = il.i_sequence[image_ns]

    # avoid broadcasting of iterator il:
    #
    mb_i = cat(read_one_image.(image_ns, Ref(il))..., dims=4)
    mb_y = UInt8.(il.i_classes[image_ns])

    if CUDA.functional()
        mb_i = KnetArray(mb_i)
    end

    if il.train
        return mb_i, mb_y
    else
        return mb_i
    end
end

function read_one_image(i, il)

    if il.pre_load
        img = il.i_images[i]
    else
        img = Images.RGB.(Images.load(il.i_paths[i]))
    end

    if il.aug_pipl isa Augmentor.ImmutablePipeline ||
        il.aug_pipl isa Augmentor.ImageOperation
        img = Augmentor.augment(img, il.aug_pipl)
    end

    if isa(img[1,1], Colors.RGB)
        img = Float32.(permutedims(Images.channelview(img), (3,2,1)))
    else
        img = Float32.(Images.channelview(img))
    end

    if il.pre_proc != nothing && il.pre_proc isa Function
        img = il.pre_proc(img)
    end
    return(img)
end

function Base.length(il::ImageLoader)
    n = length(il.i_paths) / il.batchsize
    return ceil(Int, n)
end

# return a list of all files realtive to dir:
#
function get_files_list(dir)

    image_paths = String[]
    for (root, dirs, files) in walkdir(dir)
        if !isempty(files)
            append!(image_paths, joinpath.(root, files))
        end
    end
    # filter out non-images:
    #
    fi = r"(\.png$)|(\.jpg$)|(\.jpeg$)|(\.gif$)"i
    image_paths = filter(i->occursin(fi, i), image_paths)
    return image_paths
end

# extract class names from top-level dir of list of filenames:
#
function get_class_names(dir, image_paths)

    classes = String[]
    for image_path in image_paths
        image_path = replace(image_path, dir=>"")
        dirs = splitpath(image_path)
        while length(dirs) > 0 && !occursin(r"[A-Za-z1-9]", dirs[1])
            popfirst!(dirs)
        end
        if length(dirs) < 2
            class = UNKNOWN_CLASS
        else
            class = dirs[1]
        end
        push!(classes, class)
    end

    return classes
end

# read all images into mem:
#
function pre_load_images(i_paths)

    images = []
    for path in i_paths
        push!(images, Images.RGB.(Images.load(path)))
    end
    return images
end

"""
    function image2array(img)

Take an image and return a 3d-array for RGB and a 2d-array for grayscale
images with the colour channels as last dimension.
"""
function image2array(img)

    ch = Images.channelview(img)

    if length(size(ch)) == 2
     ch = reshape(ch, :,size(ch)...)
    end

    arr = Float32.(permutedims(ch, (3,2,1)))

    if CUDA.functional()
        arr = KnetArray(arr)
    end

    return arr
end

"""
    function array2image(arr)

Take a 3d-array with colour channels as last dimension or a 2d-array
and return an array of RGB or of Gray as Image.
"""
function array2image(arr)

    if length(size(arr)) == 2
        arr = permutedims(arr, (2,1))
        itype = Images.Gray

    elseif length(size(arr)) == 3
        arr = permutedims(arr, (3,2,1))
        if size(arr)[1] == 1
            arr = reshape(arr, size(arr)[2], size(arr)[3])
            itype = Images.Gray
        else
            itype = Images.RGB
        end
    end
    return Images.colorview(itype, arr)
end

"""
    function array2RGB(arr)

Take a 3d-array with colour channels as last dimension or a 2d-array
and return always an array of RGB as Image.
"""
function array2RGB(arr)

    img = array2image(arr)
    return Images.RGB.(img)
end


abstract type Tri end
abstract type Tet end

"""
$(TYPEDEF)

Data structure for holding geometry data.
"""
struct eMesh{T1<:Union{Nothing,Tri},T2<:Union{Nothing,Tet}}
    point::Vector{SVector{3,Float64}}
    tri::Union{Nothing,Vector{SVector{3,Int64}}}
    tet::Union{Nothing,Vector{SVector{4,Int64}}}
    ϵ::Union{Nothing,Vector{Float64}}
    function eMesh( point::Vector{SVector{3,Float64}},
                    tri::Union{Nothing,Vector{SVector{3,Int64}}},
                    tet::Union{Nothing,Vector{SVector{4,Int64}}}=nothing,
                    ϵ::Union{Nothing,Vector{Float64}}=nothing)

        T1_ = ifelse(tri == nothing, Nothing, Tri)
        T2_ = ifelse(tet == nothing, Nothing, Tet)
        if T2_ == Tet
            @assert(isa(ϵ, Vector{Float64}))
            @assert(length(ϵ) == length(point), "length(ϵ) = $(length(ϵ)) but length(point) = $(length(point))")
            (length(ϵ) != 0) && @assert(0.0 < maximum(ϵ), "normalized penetration extent must be non-negative")
            (length(ϵ) != 0) && @assert(minimum(ϵ) == 0.0, "normalized penetration extent must be zero on the surface of the volume mesh")
            for k = 1:length(tet)
                (0.0 < volume(point[tet[k]])) || error("inverted tetrahedron")
            end
        else
            @assert(ϵ == nothing)
        end
        (T1_ == T2_ == Nothing) && error("a whole lot of nothing")
        return new{T1_,T2_}(deepcopy(point), deepcopy(tri), deepcopy(tet), deepcopy(ϵ))
    end
    function eMesh(hm::HomogenousMesh, tet::Union{Nothing,Vector{SVector{4,Int64}}}=nothing,
            ϵ::Union{Nothing,Vector{Float64}}=nothing)

        point = [SVector{3,Float64}(k) for k = hm.vertices]
        tri = [SVector{3,Int64}(k) for k = hm.faces]
        return eMesh(point, tri, tet, ϵ)
    end
    eMesh{Tri,Nothing}() = eMesh(Vector{SVector{3,Float64}}(), Vector{SVector{3,Int64}}(), nothing, nothing)
    eMesh{Nothing,Tet}() = eMesh(Vector{SVector{3,Float64}}(), nothing, Vector{SVector{4,Int64}}(), Vector{Float64}())
    eMesh{Tri,Tet}()     = eMesh(Vector{SVector{3,Float64}}(), Vector{SVector{3,Int64}}(), Vector{SVector{4,Int64}}(), Vector{Float64}())
end

"""
$(SIGNATURES)

Converts an `eMesh` to a mesh that contains only tetrahedrons.
You need to do this before adding an `eMesh` as contact geometry.
"""
as_tet_eMesh(e_mesh::eMesh{Tri,Tet}) = eMesh(e_mesh.point, nothing, e_mesh.tet, e_mesh.ϵ)
as_tet_eMesh(e_mesh::eMesh{Nothing,Tet}) = deepcopy(e_mesh)

"""
$(SIGNATURES)

Converts an `eMesh` to a mesh that contains only triangles.
You need to do this before adding an `eMesh` as contact geometry.
"""
as_tri_eMesh(e_mesh::eMesh{Tri,Tet}, is_repair::Bool=true) = eMesh(e_mesh.point, e_mesh.tri, nothing, nothing)
as_tri_eMesh(e_mesh::eMesh{Tri,Nothing}, is_repair::Bool=true) = deepcopy(e_mesh)
function as_tri_eMesh(eM::eMesh{Nothing,Tet}, is_repair::Bool=true)
	i3 = Vector{SVector{3,Int64}}()
	for k = 1:n_tet(eM)
		i_tet = eM.tet[k]
		ϵ_tet = eM.ϵ[i_tet]
		i_tet_new = sort_so_big_ϵ_last(ϵ_tet, i_tet)
		push!(i3, i_tet_new[1:3])
	end
	eM_tri = eMesh(eM.point, i3)
	if is_repair
		mesh_repair!(eM_tri)
	end
	return eM_tri
end

vertex_pos_for_tri_ind(eM::eMesh{Tri,T2}, k::Int64) where {T2} = eM.point[eM.tri[k]]
vertex_pos_for_tet_ind(eM::eMesh{T1,Tet}, k::Int64) where {T1} = eM.point[eM.tet[k]]

n_point(eM::eMesh) = length(eM.point)
n_tri(eM::eMesh) = length(eM.tri)
n_tet(eM::eMesh) = length(eM.tet)
get_tri(eM::eMesh) = eM.tri
get_tet(eM::eMesh) = eM.tet
get_point(eM::eMesh) = eM.point

function Base.empty!(e_mesh::eMesh{T1,T2}) where {T1,T2}
    empty!(e_mesh.point)
    (T1 == Nothing) || empty!(e_mesh.tri)
    (T2 == Nothing) || empty!(e_mesh.tet)
    (T2 == Nothing) || empty!(e_mesh.ϵ)
    return nothing
end

function Base.isempty(e_mesh::eMesh{T1,T2}) where {T1,T2}
    is_emp = isempty(e_mesh.point)
    if T1 == Tri
        is_emp = is_emp || isempty(e_mesh.tri)
    end
    if T2 == Tet
        is_emp = is_emp || isempty(e_mesh.tet)
        is_emp = is_emp || isempty(e_mesh.ϵ)
    end
    return is_emp
end

function Base.append!(eM_1::eMesh{T1,T2}, eM_2::eMesh{T1,T2}) where {T1,T2}
    n_1 = n_point(eM_1)
    if T1 != Nothing
        for k = eM_2.tri
            push!(eM_1.tri, k .+ n_1)
        end
    end
    if T2 != Nothing
        for k = eM_2.tet
            push!(eM_1.tet, k .+ n_1)
        end
        append!(eM_1.ϵ, eM_2.ϵ)
    end
    append!(eM_1.point, eM_2.point)
    return nothing
end

### VERIFICATION
verify_mesh(eM::eMesh{Tri,Nothing}) = verify_mesh_tri(eM)
verify_mesh(eM::eMesh{Nothing,Tet}) = verify_mesh_tet(eM)
function verify_mesh(eM::eMesh{Tri,Tet})
	eM_tri = as_tri_eMesh(eM)
    eM_tet = as_tet_eMesh(eM)
	verify_mesh_tri(eM_tri)
	verify_mesh_tet(eM_tet)
    eM_tri_2 = as_tri_eMesh(eM_tet)
    (area(eM_tri) ≈ area(eM_tri_2)) || error("surfaces where penetration extent is 0.0 is not the same as the tri surface")
	return nothing
end

function verify_mesh_tet(eM::eMesh{Nothing,Tet})
    length_tet = n_tet(eM)
    length_point = n_point(eM)
    length_ϵ = length(eM.ϵ)
    for k = 1:length_tet
        iΔ = eM.tet[k]
        point_Δ = eM.point[iΔ]
        (0.0 < volume(point_Δ)) || error("tet with negative volume")
    end
    (length_ϵ == length_point) || error("number of points not equal to number or ϵ")
    if !isempty(eM.ϵ)
        (minimum(eM.ϵ) < 0.0) && error("normalized penetration extent must be non-negative")
        (minimum(eM.ϵ) == 0.0) || error("normalized penetration extent must be 0.0 somewhere")
    end
    return nothing
end

function verify_mesh_tri(eM::eMesh{Tri,Nothing})
    for k = 1:n_tri(eM)
        iΔ = eM.tri[k]
        point_Δ = eM.point[iΔ]
    end
    return nothing
end

### MESH MANIPULATION
"""
$(SIGNATURES)

Creates a transformation matrix by creating a `basic_dh` from `extra_arguments` and then transforms the `eMesh`.
See [`basic_dh`](@ref) for how to create common transformation types.
Reflections are not allowed.
"""
function transform!(e_mesh::eMesh{T1,T2}, extra_args...) where {T1,T2}
    dh = basic_dh(extra_args...)
    point = e_mesh.point
    for k = 1:n_point(e_mesh)
        point[k] = dh_vector_mul(dh, point[k])
    end
end

struct PointObj
    p::SVector{3,Float64}
    w::Float64
    b::Bool
    k::Int64
end

function crop_mesh(eM::eMesh{Tri,Nothing}, plane::SMatrix{1,4,Float64,4})
    function add_cropped_triangle!(eM_new::eMesh{Tri,Nothing}, iΔ::SVector{3,Int64})
        function plane_dot_2(eM_new::eMesh{Tri,Nothing}, plane::SMatrix{1,4,Float64,4}, k::Int64)
            p = get_point(eM_new)[k]
            w = dot(plane[:], onePad(p))
            b = -1.0e-12 < w
            return PointObj(p, w, b, k)
        end

        o1 = plane_dot_2(eM_new, plane, iΔ[1])
        o2 = plane_dot_2(eM_new, plane, iΔ[2])
        o3 = plane_dot_2(eM_new, plane, iΔ[3])
        s_bool = o1.b + o2.b + o3.b
        (s_bool == 3) && (push!(eM_new.tri, iΔ))  # all entire triangle
        ((s_bool == 0) || (s_bool == 3)) && (return nothing)  # exit function
        (o1.b == o2.b) && ((o1, o2, o3) = (o2, o3, o1))  # put 1 and 2 on opposite sides
        (o2.b == o3.b) && ((o1, o2, o3) = (o3, o1, o2))  # put 2 and 3 on opposite sides
        (o1.b == o2.b) || (o2.b == o3.b) && error("something is wrong")  # sanity check
        push!(eM_new.point, weightPoly(o1.p, o2.p, o1.w, o2.w))
        i12 = length(get_point(eM_new))
        push!(eM_new.point, weightPoly(o2.p, o3.p, o2.w, o3.w))
        i23 = length(get_point(eM_new))
        if o2.b  # add top of triangle
            push!(eM_new.tri, SVector{3,Int64}(i12, o2.k, i23))
        else  # add bottom quadralateral
            push!(eM_new.tri, SVector{3,Int64}(o1.k, i12, i23))
            push!(eM_new.tri, SVector{3,Int64}(o1.k, i23, o3.k))
        end
    end

    eM_new = eMesh{Tri,Nothing}()
    append!(eM_new.point, eM.point)
    for iΔ = eM.tri
        add_cropped_triangle!(eM_new, iΔ)
    end
    mesh_repair!(eM_new)
    return eM_new
end

function invert!(eM::eMesh{Tri,Nothing})
    for k = 1:n_tri(eM)
        tri_k = eM.tri[k]
        eM.tri[k] = SVector{3,Float64}(tri_k[3], tri_k[2], tri_k[1])
    end
end

### MESH REPAIR
function mesh_repair!(e_mesh::eMesh{T1,T2}) where {T1,T2}
    mesh_remove_unused_points!(e_mesh)
    mesh_inplace_rekey!(e_mesh)
    mesh_remove_unused_points!(e_mesh)
    return delete_triangles!(e_mesh)
end

function remove_degenerate!(eM::eMesh)
	area_or_vol(v::SVector{3,SVector{3,Float64}}) = area(v)
	area_or_vol(v::SVector{4,SVector{3,Float64}}) = volume(v)

	tol = 1.0e-6
	for vec_ind = (eM.tet, eM.tri)
		if vec_ind != nothing
			v = [area_or_vol(eM.point[vec_ind[k]]) for k = 1:length(vec_ind)]
			max_v = maximum(v)
			i_bad = findall(v .< max_v * tol)
			deleteat!(vec_ind, i_bad)
		end
	end
end

rekey!(v::Nothing, i::Vector{Int64}) = nothing
rekey!(v::Vector{SVector{N,Int64}}, i::Vector{Int64}) where {N} = replace!(x -> i[x], v)

function mesh_inplace_rekey(e_mesh::eMesh{T1,T2}) where {T1,T2}
    e_mesh_copy = deepcopy(e_mesh)
    mesh_inplace_rekey!(e_mesh_copy)
    return e_mesh_copy
end

function mesh_inplace_rekey!(e_mesh::eMesh{T1,T2}) where {T1,T2}
    # Expresses the mesh using the minumum number of points

    function shortest_side(e_mesh::eMesh)
        shortest_side(v::Nothing) = Inf
        shortest_side(ind::Vector{SVector{N,Int64}}) where {N} =  minimum([shortest_side(e_mesh.point[k]) for k = ind])
        function shortest_side(v::SVector{N,SVector{3,Float64}}) where {N}
            d_min = Inf
            for k = 1:N
                for kk = 1:(k-1)
                    d_min = min(d_min, norm(v[k] - v[kk]))
                end
            end
            return d_min
        end

        return min(shortest_side(e_mesh.tri), shortest_side(e_mesh.tet))
    end

    function inplace_rekey(point::Vector{SVector{3,Float64}}, min_side_length::Float64)
        balltree = BallTree(point; reorder = false)
        idxs = inrange(balltree, point, min_side_length * 0.499)  # Assume that all points closer to each other than hald a side length are duplicates
        return first.(sort!.(idxs))
    end

    if !isempty(e_mesh)  # reducing over an empty collection is not allowed
        min_side_length = shortest_side(e_mesh)
        new_key = inplace_rekey(e_mesh.point, min_side_length)
        rekey!(e_mesh.tri, new_key)
        rekey!(e_mesh.tet, new_key)
        mesh_remove_unused_points!(e_mesh)
        return new_key
    end
    return nothing
end

function mesh_remove_unused_points!(e_mesh::eMesh{T1,T2}) where {T1,T2}
    truth_vector = falses(n_point(e_mesh))
    (T1 != Nothing) && (for k = e_mesh.tri; truth_vector[k] = true; end)
    (T2 != Nothing) && (for k = e_mesh.tet; truth_vector[k] = true; end)
    new_key = cumsum(truth_vector)
    rekey!(e_mesh.tri, new_key)
    rekey!(e_mesh.tet, new_key)
    ind_truth = findall(truth_vector)
    point_new = e_mesh.point[ind_truth]
    empty!(e_mesh.point)
    append!(e_mesh.point, point_new)
    if T2 != Nothing
        ϵ_new = e_mesh.ϵ[ind_truth]
        empty!(e_mesh.ϵ)
        append!(e_mesh.ϵ, ϵ_new)
    end
    return nothing
end

delete_triangles!(e_mesh::eMesh{Nothing,T2}) where {T2} = -9999
function delete_triangles!(e_mesh::eMesh{Tri,T2}) where {T2}
	# TODO: fix this to only delete triangles for connected meshes

    ### make the first index lowest
    for k = 1:n_tri(e_mesh)
        iΔ = e_mesh.tri[k]
        min_index = minimum(iΔ)
        (min_index == iΔ[1]) || (iΔ = SVector(iΔ[3], iΔ[1], iΔ[2]))
        (min_index == iΔ[1]) || (iΔ = SVector(iΔ[3], iΔ[1], iΔ[2]))
    end

    ### create a dictionary of key repition counts
    key_type = Vector{Tuple{Int64,SVector{3,Int64}}}
    dict_ind = Dict{SVector{3,Int64},key_type}()
    for k = 1:n_tri(e_mesh)
        iΔ = e_mesh.tri[k]
        sort_iΔ = sort(iΔ)
        !haskey(dict_ind, sort_iΔ) && (dict_ind[sort_iΔ] = key_type() )
        push!(dict_ind[sort_iΔ], (k, iΔ))
    end

    ### check for issues is triangle pairs
    i_delete = Vector{Int64}()
    for (key_k, val_k) = dict_ind
        length_val_k = length(val_k)
        if length_val_k == 2
            val_k1 = val_k[1]
            val_k2 = val_k[2]
            (val_k1[2] == val_k2[2]) && error("non-opposing triangles")
            push!(i_delete, val_k1[1], val_k2[1])
        elseif 3 <= length_val_k
            error("something is wrong")
        end
    end

    ### Actually delete the opposing triangles
    i_delete = sort(i_delete)
    deleteat!(e_mesh.tri, i_delete)
	return length(i_delete)
end

function sub_div_mesh(eM_ico::eMesh{Tri,T2}, n_div::Int64) where {T2}
    n_end(n::Int64) = div((n + 1) * n, 2)
    n_start(n::Int64) = 1 + n_end(n - 1)

    function sub_div_triangle(p::SVector{3,SVector{3,Float64}}, n_div::Int64)
        function sub_div_triangle_vert_index(n_div::Int64)
            i_tri = Vector{SVector{3,Int64}}()
            for k = 1:n_div
                for kk = 0:(k - 1)
                    i1 = n_start(k) + kk
                    i2 = i1 + k
                    i3 = i2 + 1
                    push!(i_tri, SVector{3,Int64}(i1, i2, i3))
                end
                for kk = 0:(k - 2)
                    i1 = n_start(k) + kk
                    i2 = i1 + k + 1
                    i3 = i2 - k
                    push!(i_tri, SVector{3,Int64}(i1, i2, i3))
                end
            end
            return i_tri
        end

        function get_new_point(n_vert::Int64, n_div::Int64, p::SVector{3,SVector{3,Float64}})
            function find_layer_1(i_point::Int64)
                i_end_layer = 1
                i_layer = 1
                while i_end_layer < i_point
                    i_layer += 1
                    i_end_layer += i_layer
                end
                norm_extent_1 = ifelse(i_point == 1, 0.0, (i_end_layer - i_point) / (i_layer - 1) )
                return i_layer, norm_extent_1
            end

            i_layer, norm_extent_1 = find_layer_1(n_vert)
            ϕ_1 = (n_div - i_layer + 1) / n_div
            ϕ_2 = (1 - ϕ_1) * norm_extent_1
            ϕ_3 = 1 - ϕ_1 - ϕ_2
            ϕ = SVector{3,Float64}(ϕ_1, ϕ_2, ϕ_3)
            return sum(p .* ϕ)
        end

        point = Vector{SVector{3,Float64}}()
        i_tri_div = sub_div_triangle_vert_index(n_div)
        for k = 1:n_end(n_div + 1)
            push!(point, get_new_point(k, n_div, p))
        end
        return eMesh(point, i_tri_div, nothing, nothing)
    end

    eM_ico_div = eMesh{Tri,Nothing}()
    for k = 1:n_tri(eM_ico)
        p = vertex_pos_for_tri_ind(eM_ico, k)
        append!(eM_ico_div, sub_div_triangle(p, n_div))
    end
    mesh_repair!(eM_ico_div)
    return eM_ico_div
end

### BASIC SHAPES
"""
$(SIGNATURES)

Outputs an `eMesh` for a half-plane.
"""
function eMesh_half_plane(plane_w::Float64=1.0, is_include_vis_sides::Bool=false)
    v1, v2, v3 = [SVector{3,Float64}(cos(theta), sin(theta), 0.0) for theta = (0.0, 2*pi/3, 4*pi/3)]
    v4 = SVector{3,Float64}(0.0, 0.0, -1.0) * plane_w
    point = [v1, v2, v3, v4]
    if is_include_vis_sides
        tri = [SVector{3,Int64}(1,2,3), SVector{3,Int64}(1,3,4),  SVector{3,Int64}(1,4,2), SVector{3,Int64}(2,4,3)]
    else
        tri = [SVector{3,Int64}(1,2,3)]
    end
    tet = [SVector{4,Int64}(4,1,2,3)]
    ϵ = [0.0, 0.0, 0.0, plane_w]
    return eMesh(point, tri, tet, ϵ)
end

"""
$(SIGNATURES)

Outputs an `eMesh` for a sphere. Larger values of `n_div` create finer discretizations.
"""
function eMesh_sphere(rad::Union{Float64,SVector{3,Float64}}=1.0, n_div::Int64=4)
    function make_icosahedron()
        φ = Base.MathConstants.golden

        v = Vector{SVector{3,Float64}}()
        for s_1 = (-1.0, +1.0)
            for s_2 = (-1.0, +1.0)
                push!(v, SVector{3,Float64}(    0.0,     s_1, φ * s_2) )
                push!(v, SVector{3,Float64}(    s_1, φ * s_2,     0.0) )
                push!(v, SVector{3,Float64}(φ * s_2,     0.0,     s_1) )
            end
        end

        n_ico_vert = 12
        d = zeros(n_ico_vert, n_ico_vert)
        for k = 1:n_ico_vert
            for kk = 1:n_ico_vert
                d[k, kk] = norm(v[k] - v[kk])
            end
        end

        v_face_vert = Vector{SVector{3,Int64}}()
        b = d .== 2.0
        for i1 = 1:n_ico_vert
            for i2 = 1:n_ico_vert
                for i3 = 1:n_ico_vert
                    if (i1 < i2 < i3) && b[i1, i2] && b[i2, i3] && b[i1, i3]
                        p1 = v[i1]
                        p2 = v[i2]
                        p3 = v[i3]
                        n = triangleNormal(p1, p2, p3)
                        c = normalize(centroid(p1, p2, p3))
                        i_face = ifelse(n ≈ c, SVector{3,Int64}(i1, i2, i3), SVector{3,Int64}(i1, i3, i2))
                        push!(v_face_vert, i_face)
                    end
                end
            end
        end

        push!(v, SVector{3,Float64}(0,0,0))
        v_tet = Vector{SVector{4,Int64}}()
        for k = 1:length(v_face_vert)
            push!(v_tet, SVector{4,Int64}(13, v_face_vert[k]...))
        end
        ϵ = vcat(zeros(n_ico_vert), 1.0)
        return eMesh(v, v_face_vert, v_tet, ϵ)
    end

    function volumize_about(eM::eMesh{Tri,Nothing})
        i_tet = Vector{SVector{4,Int64}}()
        n_vert = n_point(eM)
        n_center = n_vert + 1
        for k = 1:n_tri(eM)
            push!(i_tet, SVector{4,Int64}(n_center, eM.tri[k]...))
        end
        ϵ = vcat(zeros(n_vert), 1.0)
        point = deepcopy(eM.point)
        push!(point, zeros(SVector{3,Float64}))
        return eMesh(point, eM.tri, i_tet, ϵ)
    end

    function project_to_sphere!(eM::eMesh)
        for k = 1:n_point(eM)
            eM.point[k] = normalize(eM.point[k])
        end
        return nothing
    end

    eM_ico = make_icosahedron()
    eM_ico_div = sub_div_mesh(eM_ico, n_div)
    project_to_sphere!(eM_ico_div)

    rad = ones(SVector{3,Float64}) .* rad
    transform!(eM_ico_div, Diagonal(rad))

    return volumize_about(eM_ico_div)
end

function output_box_ind()
    oriented_box_faces = (
        SVector{4, Int64}(1,3,5,7),  # -x
        SVector{4, Int64}(2,6,4,8),  # +x
        SVector{4, Int64}(1,5,2,6),  # -y
        SVector{4, Int64}(3,4,7,8),  # +y
        SVector{4, Int64}(1,2,3,4),  # -z
        SVector{4, Int64}(5,7,6,8),  # +z
    )
    ϵ = zeros(Float64,8)
    push!(ϵ, 1.0)
    tri = Vector{SVector{3,Int64}}()
    tet = Vector{SVector{4,Int64}}()
    for k = 1:6
        bf_k = oriented_box_faces[k]
        push!(tri, bf_k[SVector{3,Int64}(1,3,4)])
        push!(tri, bf_k[SVector{3,Int64}(1,4,2)])
    end
    for k = 1:length(tri)
        tri_k = tri[k]
        push!(tet, SVector{4,Int64}(9, tri_k[1], tri_k[2], tri_k[3]))
    end
    return tri, tet, ϵ
end

"""
$(SIGNATURES)

Outputs an `eMesh` for a box.
"""
function eMesh_box(r::Union{Float64,SVector{3,Float64}}=1.0, c::SVector{3,Float64}=zeros(SVector{3,Float64}))
    point = [
        SVector{3,Float64}(-1,-1,-1),
        SVector{3,Float64}(+1,-1,-1),
        SVector{3,Float64}(-1,+1,-1),
        SVector{3,Float64}(+1,+1,-1),
        SVector{3,Float64}(-1,-1,+1),
        SVector{3,Float64}(+1,-1,+1),
        SVector{3,Float64}(-1,+1,+1),
        SVector{3,Float64}(+1,+1,+1),
        SVector{3,Float64}( 0, 0, 0),
    ]
    tri, tet, ϵ = output_box_ind()
    e_mesh = eMesh(point, tri, tet, ϵ)
    r = r .* ones(SVector{3,Float64})
    transform!(e_mesh, SMatrix{3,3,Float64,9}(r[1], 0, 0, 0, r[2], 0, 0, 0, r[3]))
    transform!(e_mesh, c)
    return e_mesh
end

const SF3 = SVector{3,Float64}
const SI3 = SVector{3,Int64}
const SI4 = SVector{4,Int64}

function create_2D_circle(rad::Float64=1.0; n::Int64=12)
    points_2D = Vector{SF3}()
    tri_2D = Vector{SI3}()
    i_center = n + 1
    θ_range = LinRange(0.0, 2 * π, n + 1)[2:end]
    for (k, θ) = enumerate(θ_range)
        s, c = sincos(θ)
        push!(points_2D, rad * SF3(c, s, 0.0))
        push!(tri_2D, SI3(k, mod1(k + 1, n), i_center))
    end
    push!(points_2D, SF3(0, 0, 0))
    return eMesh(points_2D, tri_2D)
end

function smallest_first(v::SI4)
	_, i = findmin(v)
	return SI4(v[i], v[mod1(i + 1, 4)], v[mod1(i + 2, 4)], v[mod1(i + 3, 4)])
end

function add_stuff_for_extrude!(tri::Vector{SI3}, tet::Vector{SI4}, k::Int64, i_tri_k::SI3, n_✴_2D::Int64)
	b1, b2, b3 = i_tri_k
	t4, t5, t6 = i_tri_k .+ n_✴_2D
	i_center = k + n_✴_2D * 2
	oriented_slice_faces = (
        SI4(b1,b2,t5,t4),
        SI4(b2,b3,t6,t5),
        SI4(b1,t4,t6,b3)
    )
    tri_add = Vector{SI3}()
    push!(tri_add, SI3(b1,b3,b2))
    push!(tri_add, SI3(t4,t5,t6))
    for bf_k = oriented_slice_faces
		bf_k = smallest_first(bf_k)
		push!(tri_add, bf_k[SI3(1,2,3)])
        push!(tri_add, bf_k[SI3(1,3,4)])
    end
    for tri_k = tri_add
        push!(tri, tri_k)
        push!(tet, SI4(i_center, tri_k[1], tri_k[2], tri_k[3]))
    end
end

function verify_trianges_have_same_normal(surf::eMesh{Tri,Nothing})
	n̂ = [triangleNormal(vertex_pos_for_tri_ind(surf, k)) for k = 1:n_tri(surf)]
	n̂_mean = n̂[1]
	for n̂_k = n̂
		(n̂_k ≈ n̂_mean) || error("All triangles must have the same normal.")
	end
	return n̂_mean
end

"""
$(SIGNATURES)

Extrudes a planar `eMesh` of type `eMesh{Tri,Nothing}` into a type `eMesh{Tri,Tet}` in the direction of the plane
normal. This direction is detected automatically.
The function errors if the points do not lie on a plane.
"""
function extrude_mesh(surf::eMesh{Tri,Nothing}, thick::Float64)
	n̂ = verify_trianges_have_same_normal(surf)

    tri_2D    = surf.tri
    points_2D = surf.point
	n_▲_2D = length(tri_2D )
	n_✴_2D = length(points_2D)
    point_lo = points_2D .- [n̂ * thick / 2]
    point_hi = points_2D .+ [n̂ * thick / 2]
    point_centroid = [centroid(points_2D[k]) for k = tri_2D]
    point = vcat(point_lo, point_hi, point_centroid)
    ϵ = vcat(zeros(2 * n_✴_2D), ones(n_▲_2D))
    tri = Vector{SI3}()
    tet = Vector{SI4}()
    for (k, i_tri_k) = enumerate(tri_2D)
		add_stuff_for_extrude!(tri, tet, k, i_tri_k, n_✴_2D)
    end
    eM_thick = eMesh(point, tri, tet, ϵ)
	mesh_repair!(eM_thick)
    return eM_thick
end

"""
$(SIGNATURES)

Outputs an `eMesh` for a cylinder.
"""
function eMesh_cylinder(rad::Float64=1.0, height::Float64=1.0; n::Int64=6)
	eM_circle = create_2D_circle(rad, n=n)
	return extrude_mesh(eM_circle, height)
end


using CUDA
using KernelAbstractions: @kernel, @index, CUDADevice
using Oceananigans.Architectures: device, GPU
using Oceananigans.Utils: work_layout

function set!(Φ::NamedTuple; kwargs...)
    for (fldname, value) in kwargs
        ϕ = getproperty(Φ, fldname)
        set!(ϕ, value)
    end
    return nothing
end

set!(u::AbstractField, v::Number) = @. u.data.parent = v

set!(u::AbstractField{X, Y, Z, A}, v::AbstractField{X, Y, Z, A}) where {X, Y, Z, A} =
    @. u.data.parent = v.data.parent

# Niceties
const AbstractCPUField =
    AbstractField{X, Y, Z, A, G} where {X, Y, Z, A<:OffsetArray{T, D, <:Array} where {T, D}, G}

const AbstractReducedCPUField =
    AbstractReducedField{X, Y, Z, A, G} where {X, Y, Z, A<:OffsetArray{T, D, <:Array} where {T, D}, G}

"Set the CPU field `u` to the array `v`."
function set!(u::AbstractCPUField, v::Array)

    Sx, Sy, Sz = size(u)
    for k in 1:Sz, j in 1:Sy, i in 1:Sx
        u[i, j, k] = v[i, j, k]
    end

    return nothing
end

""" Returns an AbstractReducedField on the CPU. """
function similar_cpu_field(u::AbstractReducedField)
    FieldType = typeof(u).name.wrapper
    return FieldType(location(u), CPU(), u.grid; dims=u.dims)
end

""" Set the CPU field `u` data to the function `f(x, y, z)`. """
set!(u::AbstractCPUField, f::Function) = interior(u) .= f.(nodes(u; reshape=true)...)

#####
##### set! for fields on the GPU
#####

@hascuda begin
    const AbstractGPUField =
        AbstractField{X, Y, Z, A, G} where {X, Y, Z, A<:OffsetArray{T, D, <:CuArray} where {T, D}, G}

    const AbstractReducedGPUField =
        AbstractReducedField{X, Y, Z, A, G} where {X, Y, Z, A<:OffsetArray{T, D, <:CuArray} where {T, D}, G}

    """ Returns a field on the CPU with `nothing` boundary conditions. """
    function similar_cpu_field(u)
        FieldType = typeof(u).name.wrapper
        return FieldType(location(u), CPU(), u.grid, nothing)
    end

    """ Set the GPU field `u` to the array `v`. """
    function set!(u::AbstractGPUField, v::Array)
        v_field = similar_cpu_field(u)

        set!(v_field, v)
        set!(u, v_field)

        return nothing
    end

    """ Set the GPU field `u` to the CuArray `v`. """
    function set!(u::AbstractGPUField, v::CuArray)

        launch!(GPU(), u.grid, :xyz, _set_gpu!, u.data, v, u.grid,
                include_right_boundaries=true, location=location(u))

        return nothing
    end

    @kernel function _set_gpu!(u, v, grid)
        i, j, k = @index(Global, NTuple)
        @inbounds u[i, j, k] = v[i, j, k]
    end

    """ Set the CPU field `u` data to the GPU field data of `v`. """
    set!(u::AbstractCPUField, v::AbstractGPUField) = u.data.parent .= Array(v.data.parent)

    """ Set the GPU field `u` data to the CPU field data of `v`. """
    set!(u::AbstractGPUField, v::AbstractCPUField) = copyto!(u.data.parent, v.data.parent)

    """ Set the GPU field `u` data to the function `f(x, y, z)`. """
    function set!(u::AbstractGPUField, f::Function)
        # Create a temporary field with bcs = nothing.
        v_field = similar_cpu_field(u)

        set!(v_field, f)
        set!(u, v_field)

        return nothing
    end
end


module Schistoxpkg
include("parameters.jl")
include("functions.jl")

export Environment
export Human
export Parameters
export create_population
export make_age_contact_rate_array
export make_death_rate_array
export birth_of_human
export cercariae_uptake!
export worm_maturity!
export calculate_worm_pairs
export calculate_total_worms
export egg_production!
export egg_production_increasing!
export miracidia_production!
export death_of_human
export birth_of_human
export count_eggs
export human_cercariae_maturity
export update_env
export find_death_rate
export cercariae_death!
export miracidia_death!
export mda
export administer_drug
export administer_vaccine
export update_contact_rate
export update_death_rate
export update_vaccine
export vaccinate
export vac_decay!
export update_mda
export load_population_from_file
export kato_katz
export get_prevalences!
export save_population_to_file
export run_simulation_from_loaded_population
export find_mean_worms_by_age
export run_simulation
export generate_age_distribution
export specified_age_distribution
export create_population_specified_ages
export update_env_to_equilibrium
export create_mda
export vaccine_information
export mda_information
export update_env_no_births_deaths
export update_env_keep_population_same
export collect_prevs
export run_repeated_sims_no_population_change
# export run_repeated_sims_random_births_deaths
export run_repeated_sims_no_population_change_increasing
# export run_repeated_sims_random_births_deaths_increasing
export create_contact_settings
export get_death_age
export generate_ages_and_deaths
export out
export update_env_keep_population_same_save_predisposition
export death_of_human_save_predis
export birth_of_human_specified_predis
export update_env_with_mda_no_births_deaths
export run_repeated_sims_no_births_deaths_increasing
export run_repeated_sims_no_births_deaths
export update_env_constant_population
export update_env_to_equilibrium_increasing
export update_env_constant_population_increasing
export update_env_no_births_deaths_increasing
export cercariae_uptake_with_human_larvae!
export human_larvae_maturity
export update_env_to_equilibrium_human_larvae
export update_env_constant_population_human_larvae
export update_env_no_births_deaths_human_larvae
export run_repeated_sims_no_population_change_human_larvae
export run_repeated_sims_no_births_deaths_human_larvae
export enact_maturity_function_false
export enact_maturity_function_true
# using Distributions
# using Random
# using PoissonRandom

end # module


immutable InvFourierLayer <: Mocha.Layer
    bottoms::Vector{Symbol}
    flags::UInt32
    name::AbstractString
    pathkeys::Vector{PathKey}
    timelimit::Float64
    tops::Vector{Symbol}
end

function InvFourierLayer( ;
        bottoms::Vector{Symbol} = Symbol[],
        flags::UInt32 = FFTW.ESTIMATE,
        name::AbstractString = "invfourier",
        pathkeys::Vector{PathKey} = PathKey[],
        timelimit::Float64 = Inf,
        tops::Vector{Symbol} = Symbol[])
    @assert length(bottoms) == 1
    @assert length(tops) == 1
    InvFourierLayer(bottoms, flags, name, pathkeys, timelimit, tops)
end

Mocha.can_do_bp(::InvFourierLayer) = false
Mocha.has_neuron(::InvFourierLayer) = false
Mocha.has_param(::InvFourierLayer) = false


module MarkovTools

export
    #functions
    count_fixed_points,
    count_cycles,
    readbitstream,

    min,
    min2,
    Spritz_Encrypt,
    Spritz_Decrypt

include("permutations.jl")
include("prng/input.jl")
include("spritz/spritz.jl")
include("monotonicites/monotonicities.jl")

end # module


#
# We use these definitions because when we have other SIMD operations with masks
# LLVM optimizes the masks better.
function truncate_mask!(instrs, input, W, suffix, reverse_load::Bool = false)
  mtyp_input = "i$(max(8,nextpow2(W)))"
  mtyp_trunc = "i$(W)"
  if reverse_load
    bitreverse = "i$(W) @llvm.bitreverse.i$(W)"
    decl = "declare $bitreverse(i$(W))"
    bitrevmask = "bitrevmask.$(suffix)"
    if mtyp_input == mtyp_trunc
      str = """
          %$(bitrevmask) = call $(bitreverse)($(mtyp_trunc) %$(input))
          %mask.$(suffix) = bitcast $mtyp_input %$(bitrevmask) to <$W x i1>
      """
    else
      str = """
          %masktrunc.$(suffix) = trunc $mtyp_input %$input to $mtyp_trunc
          %$(bitrevmask) = call $(bitreverse)($(mtyp_trunc) %masktrunc.$(suffix))
          %mask.$(suffix) = bitcast $mtyp_trunc %$(bitrevmask) to <$W x i1>
      """
    end
  else
    decl = ""
    if mtyp_input == mtyp_trunc
      str = "%mask.$(suffix) = bitcast $mtyp_input %$input to <$W x i1>"
    else
      str = "%masktrunc.$(suffix) = trunc $mtyp_input %$input to $mtyp_trunc\n%mask.$(suffix) = bitcast $mtyp_trunc %masktrunc.$(suffix) to <$W x i1>"
    end
  end
  push!(instrs, str)
  decl
end
function zext_mask!(instrs, input, W, suffix)
  mtyp_input = "i$(max(8,nextpow2(W)))"
  mtyp_trunc = "i$(W)"
  str = if mtyp_input == mtyp_trunc
    "%res.$(suffix) = bitcast <$W x i1> %$input to $mtyp_input"
  else
    "%restrunc.$(suffix) = bitcast <$W x i1> %$input to $mtyp_trunc\n%res.$(suffix) = zext $mtyp_trunc %restrunc.$(suffix) to $mtyp_input"
  end
  push!(instrs, str)
end
function binary_mask_op_instrs(W, op)
  mtyp_input = "i$(max(8,nextpow2(W)))"
  instrs = String[]
  truncate_mask!(instrs, '0', W, 0)
  truncate_mask!(instrs, '1', W, 1)
  push!(instrs, "%combinedmask = $op <$W x i1> %mask.0, %mask.1")
  zext_mask!(instrs, "combinedmask", W, 1)
  push!(instrs, "ret $mtyp_input %res.1")
  join(instrs, "\n")
end
function binary_mask_op(W, U, op, evl::Symbol = Symbol(""))
  instrs = binary_mask_op_instrs(W, op)
  mask = Expr(:curly, evl === Symbol("") ? :Mask : :EVLMask, W)
  gf = GlobalRef(Core, :getfield)
  gf1 = Expr(:call, gf, :m1, 1, false)
  gf2 = Expr(:call, gf, :m2, 1, false)
  llvmc = Expr(:call, GlobalRef(Base, :llvmcall), instrs, U, :(Tuple{$U,$U}), gf1, gf2)
  call = Expr(:call, mask, llvmc)
  evl === Symbol("") ||
    push!(call.args, Expr(:call, evl, :($gf(m1, :evl)), :($gf(m2, :evl))))
  Expr(:block, Expr(:meta, :inline), call)
end

@inline data(m::Mask) = getfield(m, :u)
@inline data(m::EVLMask) = getfield(m, :u)
@inline Base.convert(::Type{Mask{W,U}}, m::EVLMask{W,U}) where {W,U} =
  Mask{W,U}(getfield(m, :u))
for (f, op, evl) ∈ [
  (:vand, "and", :min),
  (:vor, "or", :max),
  (:vxor, "xor", Symbol("")),
  (:veq, "icmp eq", Symbol("")),
  (:vne, "icmp ne", Symbol("")),
]
  @eval begin
    @generated function $f(m1::AbstractMask{W,U}, m2::AbstractMask{W,U}) where {W,U}
      binary_mask_op(
        W,
        U,
        $op,
        ((m1 <: EVLMask) && (m2 <: EVLMask)) ? $(QuoteNode(evl)) : Symbol(""),
      )
    end
  end
end
for f ∈ [:vand, :vor, :vxor] # ignore irrelevant bits, so just bitcast to `Bool`
  @eval @inline $f(a::Vec{W,Bool}, b::Vec{W,Bool}) where {W} =
    vreinterpret(Bool, $f(vreinterpret(UInt8, a), vreinterpret(UInt8, b)))
end
for f ∈ [:vne, :veq] # Here we truncate.
  @eval @inline $f(a::Vec{W,Bool}, b::Vec{W,Bool}) where {W} =
    convert(Bool, $f(convert(Bit, a), convert(Bit, b)))
end

@generated function vconvert(
  ::Type{Vec{W,I}},
  m::AbstractMask{W,U},
) where {W,I<:IntegerTypesHW,U<:Union{UInt8,UInt16,UInt32,UInt64}}
  bits = 8sizeof(I)
  instrs = String[]
  truncate_mask!(instrs, '0', W, 0)
  push!(
    instrs,
    "%res = zext <$W x i1> %mask.0 to <$W x i$(bits)>\nret <$W x i$(bits)> %res",
  )
  gf = Expr(:call, GlobalRef(Core, :getfield), :m, 1, false)
  llvmc = Expr(
    :call,
    GlobalRef(Base, :llvmcall),
    join(instrs, "\n"),
    :(_Vec{$W,$I}),
    :(Tuple{$U}),
    gf,
  )
  Expr(:block, Expr(:meta, :inline), Expr(:call, :Vec, llvmc))
end

@generated function splitint(
  i::S,
  ::Type{T},
) where {S<:Base.BitInteger,T<:Union{Bool,Base.BitInteger}}
  sizeof_S = sizeof(S)
  sizeof_T = sizeof(T)
  if sizeof_T > sizeof_S
    return :(i % T)
  elseif sizeof_T == sizeof_S
    return :i
  end
  W, r = divrem(sizeof_S, sizeof_T)
  @assert iszero(r)
  vtyp = "<$W x i$(8sizeof_T)>"
  instrs = """
      %split = bitcast i$(8sizeof_S) %0 to $vtyp
      ret $vtyp %split
  """
  quote
    $(Expr(:meta, :inline))
    Vec($LLVMCALL($instrs, _Vec{$W,$T}, Tuple{$S}, i))
  end
end
@generated function fuseint(v::Vec{W,I}) where {W,I<:Union{Bool,Base.BitInteger}}
  @assert ispow2(W)
  bytes = W * sizeof(I)
  bits = 8bytes
  @assert bytes ≤ 16
  T = (I <: Signed) ? Symbol(:Int, bits) : Symbol(:UInt, bits)
  vtyp = "<$W x i$(8sizeof(I))>"
  styp = "i$(bits)"
  instrs = """
      %fused = bitcast $vtyp %0 to $styp
      ret $styp %fused
  """
  quote
    $(Expr(:meta, :inline))
    $LLVMCALL($instrs, $T, Tuple{_Vec{$W,$I}}, data(v))
  end
end

function vadd_expr(W, U, instr)
  instrs = String[]
  truncate_mask!(instrs, '0', W, 0)
  truncate_mask!(instrs, '1', W, 1)
  push!(
    instrs,
    """%uv.0 = zext <$W x i1> %mask.0 to <$W x i8>
%uv.1 = zext <$W x i1> %mask.1 to <$W x i8>
%res = $instr <$W x i8> %uv.0, %uv.1
ret <$W x i8> %res""",
  )
  Expr(
    :block,
    Expr(:meta, :inline),
    :(Vec(
      $LLVMCALL(
        $(join(instrs, "\n")),
        _Vec{$W,UInt8},
        Tuple{$U,$U},
        getfield(m1, :u),
        getfield(m2, :u),
      ),
    )),
  )
end
@generated vadd_fast(m1::AbstractMask{W,U}, m2::AbstractMask{W,U}) where {W,U} =
  vadd_expr(W, U, "add")
@generated vsub_fast(m1::AbstractMask{W,U}, m2::AbstractMask{W,U}) where {W,U} =
  vadd_expr(W, U, "sub")
@inline vadd(m1::AbstractMask{W,U}, m2::AbstractMask{W,U}) where {W,U} = vadd_fast(m1, m2)
@inline vsub(m1::AbstractMask{W,U}, m2::AbstractMask{W,U}) where {W,U} = vsub_fast(m1, m2)

@inline Base.:(&)(m::AbstractMask{W,U}, b::Bool) where {W,U} =
  Mask{W,U}(Core.ifelse(b, getfield(m, :u), zero(getfield(m, :u))))
@inline Base.:(&)(b::Bool, m::AbstractMask{W,U}) where {W,U} =
  Mask{W,U}(Core.ifelse(b, getfield(m, :u), zero(getfield(m, :u))))

@inline Base.:(|)(m::AbstractMask{W,U}, b::Bool) where {W,U} =
  Mask{W,U}(Core.ifelse(b, getfield(max_mask(Mask{W,U}), :u), getfield(m, :u)))
@inline Base.:(|)(b::Bool, m::AbstractMask{W,U}) where {W,U} =
  Mask{W,U}(Core.ifelse(b, getfield(max_mask(Mask{W,U}), :u), getfield(m, :u)))

@inline function Base.:(&)(m::EVLMask{W,U}, b::Bool) where {W,U}
  EVLMask{W,U}(
    Core.ifelse(b, getfield(m, :u), zero(getfield(m, :u))),
    Core.ifelse(b, getfield(m, :evl), 0x00000000),
  )
end
@inline function Base.:(&)(b::Bool, m::EVLMask{W,U}) where {W,U}
  EVLMask{W,U}(
    Core.ifelse(b, getfield(m, :u), zero(getfield(m, :u))),
    Core.ifelse(b, getfield(m, :evl), 0x00000000),
  )
end

@inline function Base.:(|)(m::EVLMask{W,U}, b::Bool) where {W,U}
  EVLMask{W,U}(
    Core.ifelse(b, getfield(max_mask(Mask{W,U}), :u), getfield(m, :u)),
    Core.ifelse(b, W % UInt32, getfield(m, :evl)),
  )
end
@inline function Base.:(|)(b::Bool, m::EVLMask{W,U}) where {W,U}
  EVLMask{W,U}(
    Core.ifelse(b, getfield(max_mask(Mask{W,U}), :u), getfield(m, :u)),
    Core.ifelse(b, W % UInt32, getfield(m, :evl)),
  )
end

@inline Base.:(⊻)(m::AbstractMask{W,U}, b::Bool) where {W,U} =
  Mask{W,U}(Core.ifelse(b, ~getfield(m, :u), getfield(m, :u)))
@inline Base.:(⊻)(b::Bool, m::AbstractMask{W,U}) where {W,U} =
  Mask{W,U}(Core.ifelse(b, ~getfield(m, :u), getfield(m, :u)))

@inline vshl(m::AbstractMask{W,U}, i::IntegerTypesHW) where {W,U} =
  Mask{W,U}(shl(getfield(m, :u), i))
@inline vashr(m::AbstractMask{W,U}, i::IntegerTypesHW) where {W,U} =
  Mask{W,U}(shr(getfield(m, :u), i))
@inline vlshr(m::AbstractMask{W,U}, i::IntegerTypesHW) where {W,U} =
  Mask{W,U}(shr(getfield(m, :u), i))

@inline zero_mask(::AbstractSIMDVector{W}) where {W} = Mask(zero_mask(Val(W)))
@inline zero_mask(::VecUnroll{N,W}) where {N,W} = VecUnroll{N}(Mask(zero_mask(Val(W))))
@inline max_mask(::AbstractSIMDVector{W}) where {W} = Mask(max_mask(Val(W)))
@inline max_mask(::VecUnroll{N,W}) where {N,W} = VecUnroll{N}(Mask(max_mask(Val(W))))
@inline zero_mask(::NativeTypes) = false
@inline max_mask(::NativeTypes) = true

for (U, W) in [(UInt8, 8), (UInt16, 16), (UInt32, 32), (UInt64, 64)]
  @eval @inline vany(m::AbstractMask{$W,$U}) = getfield(m, :u) != $(zero(U))
  @eval @inline vall(m::AbstractMask{$W,$U}) = getfield(m, :u) == $(typemax(U))
end
# TODO: use vector reduction intrsincs
@inline function vany(m::AbstractMask{W}) where {W}
  mm = getfield(max_mask(Val{W}()), :u)
  mu = getfield(m, :u)
  (mu & mm) !== zero(mu)
end
@inline function vall(m::AbstractMask{W}) where {W}
  mm = getfield(max_mask(Val{W}()), :u)
  mu = getfield(m, :u)
  mm & mu === mm
end
@inline vany(b::Bool) = b
@inline vall(b::Bool) = b
@inline vsum(m::AbstractMask) = count_ones(getfield(m, :u))
@inline vprod(m::AbstractMask) = vall(m)

@generated function vnot(m::AbstractMask{W,U}) where {W,U}
  mtyp_input = "i$(8sizeof(U))"
  mtyp_trunc = "i$(W)"
  instrs = String[]
  truncate_mask!(instrs, '0', W, 0)
  mask = llvmconst(W, "i1 true")
  push!(instrs, "%resvec.0 = xor <$W x i1> %mask.0, $mask")
  zext_mask!(instrs, "resvec.0", W, 1)
  push!(instrs, "ret $mtyp_input %res.1")
  quote
    $(Expr(:meta, :inline))
    Mask{$W}($LLVMCALL($(join(instrs, "\n")), $U, Tuple{$U}, getfield(m, :u)))
  end
end
@inline vnot(x::Bool) = Base.not_int(x)
# @inline Base.:(~)(m::Mask) = !m

@inline Base.count_ones(m::AbstractMask) = count_ones(getfield(m, :u))
@inline vadd(m::AbstractMask, i::IntegerTypesHW) = i + count_ones(m)
@inline vadd(i::IntegerTypesHW, m::AbstractMask) = i + count_ones(m)

@generated function vzero(::Type{M}) where {W,M<:Mask{W}}
  Expr(
    :block,
    Expr(:meta, :inline),
    Expr(:call, Expr(:curly, :Mask, W), Expr(:call, :zero, mask_type_symbol(W))),
  )
end
@generated function vzero(::Type{M}) where {W,M<:EVLMask{W}}
  Expr(
    :block,
    Expr(:meta, :inline),
    Expr(
      :call,
      Expr(:curly, :EVLMask, W),
      Expr(:call, :zero, mask_type_symbol(W)),
      0x00000000,
    ),
  )
end
@inline vzero(::Mask{W,U}) where {W,U} = Mask{W}(zero(U))
@inline vzero(::EVLMask{W,U}) where {W,U} = EVLMask{W}(zero(U), 0x00000000)
@inline Base.zero(::Type{M}) where {W,M<:AbstractMask{W}} = vzero(M)
@inline zero_mask(::Union{Val{W},StaticInt{W}}) where {W} =
  EVLMask{W}(zero(VectorizationBase.mask_type(Val{W}())), 0x00000000)

@generated function max_mask(::Union{Val{W},StaticInt{W}}) where {W}
  U = mask_type(W)
  :(EVLMask{$W,$U}($(one(U) << W - one(U)), $(UInt32(W))))
end
@inline max_mask(::Type{T}) where {T} = max_mask(pick_vector_width(T))
@generated max_mask(::Type{Mask{W,U}}) where {W,U} =
  EVLMask{W,U}(one(U) << W - one(U), W % UInt32)

@generated function valrem(::Union{Val{W},StaticInt{W}}, l::T) where {W,T<:Union{Integer,StaticInt}}
  ex = ispow2(W) ? :(l & $(T(W - 1))) : Expr(:call, Base.urem_int, :l, T(W))
  Expr(:block, Expr(:meta, :inline), ex)
end

function bzhi_quote(b)
  T = b == 32 ? :UInt32 : :UInt64
  typ = 'i' * string(b)
  instr = "i$b @llvm.x86.bmi.bzhi.$b"
  decl = "declare $instr(i$b, i$b) nounwind readnone"
  instrs = "%res = call $instr(i$b %0, i$b %1)\n ret i$b %res"
  llvmcall_expr(decl, instrs, T, :(Tuple{$T,$T}), typ, [typ, typ], [:a, :b])
end
@generated bzhi(a::UInt32, b::UInt32) = bzhi_quote(32)
@generated bzhi(a::UInt64, b::UInt64) = bzhi_quote(64)

# @generated function _mask(::Union{Val{W},StaticInt{W}}, l::I, ::True) where {W,I<:Union{Integer,StaticInt}}
#   # if `has_opmask_registers()` then we can use bitmasks directly, so we create them via bittwiddling
#   M = mask_type(W)
#   quote # If the arch has opmask registers, we can generate a bitmask and then move it into the opmask register
#     $(Expr(:meta,:inline))
#     evl = valrem(Val{$W}(), (l % $M) - one($M))
#     EVLMask{$W,$M}($(typemax(M)) >>> ($(M(8sizeof(M))-1) - evl), evl + one(evl))
#   end
# end
@generated function _mask_bzhi(::Union{Val{W},StaticInt{W}}, l::I) where {W,I<:Union{Integer,StaticInt}}
  U = mask_type_symbol(W)
  T = W > 32 ? :UInt64 : :UInt32
  quote
    $(Expr(:meta, :inline))
    m = valrem(StaticInt{$W}(), l % $T)
    m = Core.ifelse((m % UInt8) == 0x00, $W % $T, m)
    EVLMask{$W,$U}(bzhi(-1 % $T, m) % $U, m)
  end
end
# @inline function _mask_bzhi(::Union{Val{W},StaticInt{W}}, l::I) where {W,I<:Union{Integer,StaticInt}}
#   U = mask_type(StaticInt(W))
#   # m = ((l) % UInt32) & ((W-1) % UInt32)
#   m = valrem(StaticInt{W}(), l % UInt32)
#   m = Core.ifelse((m % UInt8) == 0x00, W % UInt32, m)
#   # m = Core.ifelse(zero(m) == m, -1 % UInt32, m)
#   EVLMask{W,U}(bzhi(-1 % UInt32, m) % U, m)
# end
# @inline function _mask(::Union{Val{W},StaticInt{W}}, l::I, ::True) where {W,I<:Union{Integer,StaticInt}}
#   U = mask_type(StaticInt(W))
#   m = ((l-one(l)) % UInt32) & ((W-1) % UInt32)
#   m += one(m)
#   EVLMask{W,U}(bzhi(-1 % UInt32, m) % U, m)
# end
# @generated function _mask(::Union{Val{W},StaticInt{W}}, l::I, ::True) where {W,I<:Union{Integer,StaticInt}}
#   M = mask_type_symbol(W)
#   quote
#     $(Expr(:meta,:inline))
#     evl = valrem(Val{$W}(), vsub_nw((l % $M), one($M)))
#     EVLMask{$W}(data(evl ≥ MM{$W}(0)), vadd_nw(evl, one(evl)))
#   end
# end

function mask_shift_quote(W::Int, bmi::Bool)
  if (((Sys.ARCH === :x86_64) || (Sys.ARCH === :i686))) && bmi
    W ≤ 64 && return Expr(:block, Expr(:meta, :inline), :(_mask_bzhi(StaticInt{$W}(), l)))
  end
  MT = mask_type(W)
  quote # If the arch has opmask registers, we can generate a bitmask and then move it into the opmask register
    $(Expr(:meta, :inline))
    evl = valrem(Val{$W}(), (l % $MT) - one($MT))
    EVLMask{$W,$MT}($(typemax(MT)) >>> ($(MT(8sizeof(MT)) - 1) - evl), evl + one(evl))
  end
end
@generated _mask_shift(::StaticInt{W}, l, ::True) where {W} = mask_shift_quote(W, true)
@generated _mask_shift(::StaticInt{W}, l, ::False) where {W} = mask_shift_quote(W, false)
@static if Base.libllvm_version ≥ v"12"
  function active_lane_mask_quote(W::Int)
    quote
      $(Expr(:meta, :inline))
      upper = (l % UInt32) & $((UInt32(W - 1)))
      upper = Core.ifelse(upper == 0x00000000, $(W % UInt32), upper)
      mask(Val{$W}(), 0x00000000, upper)
    end
  end
else
  function active_lane_mask_quote(W::Int)
    quote
      $(Expr(:meta, :inline))
      mask(Val{$W}(), 0x00000000, vsub_nw(l % UInt32, 0x00000001) & $(UInt32(W - 1)))
    end
  end
end
function mask_cmp_quote(W::Int, RS::Int, bmi::Bool)
  M = mask_type_symbol(W)
  bytes = min(RS ÷ W, 8)
  bytes < 4 && return mask_shift_quote(W, bmi)
  T = integer_of_bytes_symbol(bytes, true)
  quote
    $(Expr(:meta, :inline))
    evl = valrem(Val{$W}(), vsub_nw((l % $T), one($T)))
    EVLMask{$W}(data(evl ≥ MM{$W}(zero($T))), vadd_nw(evl, one(evl)))
  end
end
@generated _mask_cmp(
  ::Union{Val{W},StaticInt{W}},
  l::I,
  ::StaticInt{RS},
  ::True,
) where {W,RS,I<:Union{Integer,StaticInt}} = mask_cmp_quote(W, RS, true)
@generated _mask_cmp(
  ::Union{Val{W},StaticInt{W}},
  l::I,
  ::StaticInt{RS},
  ::False,
) where {W,RS,I<:Union{Integer,StaticInt}} = mask_cmp_quote(W, RS, false)
@generated _mask(::Union{Val{W},StaticInt{W}}, l::I, ::True) where {W,I<:Union{Integer,StaticInt}} =
  mask_shift_quote(W, true)
@generated function _mask(::Union{Val{W},StaticInt{W}}, l::I, ::False) where {W,I<:Union{Integer,StaticInt}}
  # Otherwise, it's probably more efficient to use a comparison, as this will probably create some type that can be used directly for masked moves/blends/etc
  if W > 16
    Expr(
      :block,
      Expr(:meta, :inline),
      :(_mask_shift(StaticInt{$W}(), l, has_feature(Val(:x86_64_bmi)))),
    )
    # mask_shift_quote(W)
    # elseif (Base.libllvm_version ≥ v"11") && ispow2(W)
    # elseif ((Sys.ARCH ≢ :x86_64) && (Sys.ARCH ≢ :i686)) && (Base.libllvm_version ≥ v"11") && ispow2(W)
  elseif false
    # cmpval = Base.libllvm_version ≥ v"12" ? -one(I) : zero(I)
    active_lane_mask_quote(W)
  else
    Expr(
      :block,
      Expr(:meta, :inline),
      :(_mask_cmp(
        Val{$W}(),
        l,
        simd_integer_register_size(),
        has_feature(Val(:x86_64_bmi)),
      )),
    )
  end
end
# This `mask` method returns a constant, independent of `has_opmask_registers()`; that only effects method of calculating
# the constant. So it'd be safe to bake in a value.
@inline mask(::Union{Val{W},StaticInt{W}}, L) where {W} =
  _mask(StaticInt(W), L, has_feature(Val(:x86_64_avx512f)) & ge_one_fma(cpu_name()))
@inline mask(::Union{Val{W},StaticInt{W}}, ::StaticInt{L}) where {W,L} =
  _mask(StaticInt(W), L, has_feature(Val(:x86_64_avx512f)) & ge_one_fma(cpu_name()))
@inline mask(::Type{T}, l::Union{Integer,StaticInt}) where {T} =
  _mask(pick_vector_width(T), l, has_feature(Val(:x86_64_avx512f)) & ge_one_fma(cpu_name()))

# @generated function masktable(::Union{Val{W},StaticInt{W}}, rem::Union{Integer,StaticInt}) where {W}
#     masks = Expr(:tuple)
#     for w ∈ 0:W-1
#         push!(masks.args, data(mask(Val(W), w == 0 ? W : w)))
#     end
#     Expr(
#         :block,
#         Expr(:meta,:inline),
#         Expr(:call, Expr(:curly, :Mask, W), Expr(
#             :macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__, Symbol(@__FILE__)),
#             Expr(:call, :getindex, masks, Expr(:call, :+, 1, Expr(:call, :valrem, Expr(:call, Expr(:curly, W)), :rem)))
#         ))
#     )
# end

@inline tomask(m::Unsigned) = Mask{sizeof(m)}(m)
@inline tomask(m::Mask) = m
@generated function tomask(v::Vec{W,Bool}) where {W}
  usize = W > 8 ? nextpow2(W) : 8
  utyp = "i$(usize)"
  U = mask_type_symbol(W)
  instrs = String[]
  push!(instrs, "%bitvec = trunc <$W x i8> %0 to <$W x i1>")
  zext_mask!(instrs, "bitvec", W, 0)
  push!(instrs, "ret i$(usize) %res.0")
  quote
    $(Expr(:meta, :inline))
    Mask{$W}($LLVMCALL($(join(instrs, "\n")), $U, Tuple{_Vec{$W,Bool}}, data(v)))
  end
end
@inline tomask(v::AbstractSIMDVector{W,Bool}) where {W} =
  tomask(vconvert(Vec{W,Bool}, data(v)))
# @inline tounsigned(m::Mask) = getfield(m, :u)
# @inline tounsigned(m::Vec{W,Bool}) where {W} = getfield(tomask(m), :u)
@inline tounsigned(v) = getfield(tomask(v), :u)

@generated function vrem(m::Mask{W,U}, ::Type{I}) where {W,U,I<:IntegerTypesHW}
  bits = 8sizeof(I)
  instrs = String[]
  truncate_mask!(instrs, '0', W, 0)
  push!(
    instrs,
    "%res = zext <$W x i1> %mask.0 to <$W x i$(bits)>\nret <$W x i$(bits)> %res",
  )
  quote
    $(Expr(:meta, :inline))
    Vec($LLVMCALL($(join(instrs, "\n")), _Vec{$W,$I}, Tuple{$U}, data(m)))
  end
end
Vec(m::Mask{W}) where {W} = m % int_type(Val{W}())

# @inline getindexzerobased(m::Mask, i) = (getfield(m, :u) >>> i) % Bool
# @inline function extractelement(m::Mask{W}, i::Union{Integer,StaticInt}) where {W}
#     @boundscheck i > W && throw(BoundsError(m, i))
#     getindexzerobased(m, i)
# end
@generated function extractelement(v::Mask{W,U}, i::I) where {W,U,I}
  instrs = String[]
  truncate_mask!(instrs, '0', W, 0)
  push!(instrs, "%res1 = extractelement <$W x i1> %mask.0, i$(8sizeof(I)) %1")
  push!(instrs, "%res8 = zext i1 %res1 to i8\nret i8 %res8")
  instrs_string = join(instrs, "\n")
  call = :($LLVMCALL($instrs_string, Bool, Tuple{$U,$I}, data(v), i))
  Expr(:block, Expr(:meta, :inline), call)
end
@generated function insertelement(
  v::Mask{W,U},
  x::T,
  i::I,
) where {W,T,U,I<:Union{Bool,IntegerTypesHW}}
  mtyp_input = "i$(max(8,nextpow2(W)))"
  instrs = String["%bit = trunc i$(8sizeof(T)) %1 to i1"]
  truncate_mask!(instrs, '0', W, 0)
  push!(instrs, "%bitvec = insertelement <$W x i1> %mask.0, i1 %bit, i$(8sizeof(I)) %2")
  zext_mask!(instrs, "bitvec", W, 1)
  push!(instrs, "ret $(mtyp_input) %res.1")
  instrs_string = join(instrs, "\n")
  call = :(Mask{$W}($LLVMCALL($instrs_string, $U, Tuple{$U,$T,$I}, data(v), x, i)))
  Expr(:block, Expr(:meta, :inline), call)
end

# @generated function Base.isodd(i::MM{W,1}) where {W}
#     U = mask_type(W)
#     evenfirst = 0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa % U
#     # Expr(:block, Expr(:meta, :inline), :(isodd(getfield(i, :i)) ? Mask{$W}($oddfirst) : Mask{$W}($evenfirst)))
#     Expr(:block, Expr(:meta, :inline), :(Mask{$W}($evenfirst >> (getfield(i, :i) & 0x03))))
# end
# @generated function Base.iseven(i::MM{W,1}) where {W}
#     U = mask_type(W)
#     oddfirst = 0x55555555555555555555555555555555 % U
#     # evenfirst = oddfirst << 1
#     # Expr(:block, Expr(:meta, :inline), :(isodd(getfield(i, :i)) ? Mask{$W}($evenfirst) : Mask{$W}($oddfirst)))
#     Expr(:block, Expr(:meta, :inline), :(Mask{$W}($oddfirst >> (getfield(i, :i) & 0x03))))
# end
@inline Base.isodd(i::MM{W,1}) where {W} = Mask{W}(
  (0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa % mask_type(Val{W}())) >>> (getfield(i, :i) & 0x03),
)
@inline Base.iseven(i::MM{W,1}) where {W} = Mask{W}(
  (0x55555555555555555555555555555555 % mask_type(Val{W}())) >>> (getfield(i, :i) & 0x03),
)

function cmp_quote(W, cond, vtyp, T1, T2 = T1)
  instrs = String["%m = $cond $vtyp %0, %1"]
  zext_mask!(instrs, 'm', W, '0')
  push!(instrs, "ret i$(max(8,nextpow2(W))) %res.0")
  U = mask_type_symbol(W)
  quote
    $(Expr(:meta, :inline))
    Mask{$W}(
      $LLVMCALL(
        $(join(instrs, "\n")),
        $U,
        Tuple{_Vec{$W,$T1},_Vec{$W,$T2}},
        data(v1),
        data(v2),
      ),
    )
  end
end
function icmp_quote(W, cond, bytes, T1, T2 = T1)
  vtyp = vtype(W, "i$(8bytes)")
  cmp_quote(W, "icmp " * cond, vtyp, T1, T2)
end
function fcmp_quote(W, cond, T)
  vtyp = vtype(W, T === Float32 ? "float" : "double")
  cmp_quote(W, "fcmp nsz arcp contract reassoc " * cond, vtyp, T)
end
# @generated function compare(::Val{cond}, v1::Vec{W,I}, v2::Vec{W,I}) where {cond, W, I}
# cmp_quote(W, cond, sizeof(I), I)
# end
# for (f,cond) ∈ [(:(==), :eq), (:(!=), :ne), (:(>), :ugt), (:(≥), :uge), (:(<), :ult), (:(≤), :ule)]
for (f, cond) ∈ [(:veq, "eq"), (:vne, "ne")]
  @eval @generated function $f(
    v1::Vec{W,T1},
    v2::Vec{W,T2},
  ) where {W,T1<:IntegerTypesHW,T2<:IntegerTypesHW}
    if sizeof(T1) != sizeof(T2)
      return Expr(
        :block,
        Expr(:meta, :inline),
        :((v3, v4) = promote(v1, v2)),
        Expr(:call, $f, :v3, :v4),
      )
    end
    icmp_quote(W, $cond, sizeof(T1), T1, T2)
  end
end
for (f, cond) ∈ [(:vgt, "ugt"), (:vge, "uge"), (:vlt, "ult"), (:vle, "ule")]
  @eval @generated function $f(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:Unsigned}
    icmp_quote(W, $cond, sizeof(T), T)
  end
end
for (f, cond) ∈ [(:vgt, "sgt"), (:vge, "sge"), (:vlt, "slt"), (:vle, "sle")]
  @eval @generated function $f(v1::Vec{W,T}, v2::Vec{W,T}) where {W,T<:Signed}
    icmp_quote(W, $cond, sizeof(T), T)
  end
end

# for (f,cond) ∈ [(:veq, "oeq"), (:vgt, "ogt"), (:vge, "oge"), (:vlt, "olt"), (:vle, "ole"), (:vne, "one")]
for (f, cond) ∈ [
  (:veq, "oeq"),
  (:vgt, "ogt"),
  (:vge, "oge"),
  (:vlt, "olt"),
  (:vle, "ole"),
  (:vne, "une"),
]
  # for (f,cond) ∈ [(:veq, "ueq"), (:vgt, "ugt"), (:vge, "uge"), (:vlt, "ult"), (:vle, "ule"), (:vne, "une")]
  @eval @generated function $f(
    v1::Vec{W,T},
    v2::Vec{W,T},
  ) where {W,T<:Union{Float32,Float64}}
    fcmp_quote(W, $cond, T)
  end
end

@inline function vgt(
  v1::AbstractSIMDVector{W,S},
  v2::AbstractSIMDVector{W,U},
) where {W,S<:SignedHW,U<:UnsignedHW}
  (v1 > zero(S)) & (vconvert(U, v1) > v2)
end
@inline function vgt(
  v1::AbstractSIMDVector{W,U},
  v2::AbstractSIMDVector{W,S},
) where {W,S<:SignedHW,U<:UnsignedHW}
  (v2 < zero(S)) | (vconvert(S, v1) > v2)
end

@inline function vge(
  v1::AbstractSIMDVector{W,S},
  v2::AbstractSIMDVector{W,U},
) where {W,S<:SignedHW,U<:UnsignedHW}
  (v1 ≥ zero(S)) & (vconvert(U, v1) ≥ v2)
end
@inline function vge(
  v1::AbstractSIMDVector{W,U},
  v2::AbstractSIMDVector{W,S},
) where {W,S<:SignedHW,U<:UnsignedHW}
  (v2 < zero(S)) | (vconvert(S, v1) ≥ v2)
end

@inline vlt(
  v1::AbstractSIMDVector{W,S},
  v2::AbstractSIMDVector{W,U},
) where {W,S<:SignedHW,U<:UnsignedHW} = vgt(v2, v1)
@inline vlt(
  v1::AbstractSIMDVector{W,U},
  v2::AbstractSIMDVector{W,S},
) where {W,S<:SignedHW,U<:UnsignedHW} = vgt(v2, v1)
@inline vle(
  v1::AbstractSIMDVector{W,S},
  v2::AbstractSIMDVector{W,U},
) where {W,S<:SignedHW,U<:UnsignedHW} = vge(v2, v1)
@inline vle(
  v1::AbstractSIMDVector{W,U},
  v2::AbstractSIMDVector{W,S},
) where {W,S<:SignedHW,U<:UnsignedHW} = vge(v2, v1)
for (op, f) ∈ [(:vgt, :(>)), (:vge, :(≥)), (:vlt, :(<)), (:vle, :(≤))]
  @eval begin
    @inline function $op(
      v1::V1,
      v2::V2,
    ) where {
      V1<:Union{IntegerTypesHW,AbstractSIMDVector{<:Any,<:IntegerTypesHW}},
      V2<:Union{IntegerTypesHW,AbstractSIMDVector{<:Any,<:IntegerTypesHW}},
    }
      V3 = promote_type(V1, V2)
      $op(itosize(v1, V3), itosize(v2, V3))
    end
    @inline function $op(v1, v2)
      v3, v4 = promote(v1, v2)
      $op(v3, v4)
    end
    @inline $op(s1::IntegerTypesHW, s2::IntegerTypesHW) = $f(s1, s2)
    @inline $op(s1::Union{Float32,Float64}, s2::Union{Float32,Float64}) = $f(s1, s2)
  end
end
for (op, f) ∈ [(:veq, :(==)), (:vne, :(≠))]
  @eval begin
    @inline $op(a, b) = ((c, d) = promote(a, b); $op(c, d))
    @inline $op(s1::NativeTypes, s2::NativeTypes) = $f(s1, s2)
  end
end

@generated function vifelse(m::AbstractMask{W,U}, v1::Vec{W,T}, v2::Vec{W,T}) where {W,U,T}
  typ = LLVM_TYPES[T]
  vtyp = vtype(W, typ)
  selty = vtype(W, "i1")
  f = "select"
  if Base.libllvm_version ≥ v"9" && ((T === Float32) || (T === Float64))
    f *= " nsz arcp contract reassoc"
  end
  instrs = String[]
  truncate_mask!(instrs, '0', W, 0)
  push!(instrs, "%res = $f $selty %mask.0, $vtyp %1, $vtyp %2\nret $vtyp %res")
  quote
    $(Expr(:meta, :inline))
    Vec(
      $LLVMCALL(
        $(join(instrs, "\n")),
        _Vec{$W,$T},
        Tuple{$U,_Vec{$W,$T},_Vec{$W,$T}},
        data(m),
        data(v1),
        data(v2),
      ),
    )
  end
end

@inline vifelse(m::Vec{W,Bool}, s1::T, s2::T) where {W,T<:NativeTypes} =
  vifelse(m, Vec{W,T}(s1), Vec{W,T}(s2))
@inline vifelse(m::AbstractMask{W}, s1::T, s2::T) where {W,T<:NativeTypes} =
  vifelse(m, Vec{W,T}(s1), Vec{W,T}(s2))
@inline vifelse(m::AbstractMask{W,U}, s1, s2) where {W,U} =
  ((x1, x2) = promote(s1, s2); vifelse(m, x1, x2))
@inline vifelse(m::AbstractMask{W}, v1::VecUnroll{N,W}, v2::VecUnroll{N,W}) where {N,W} =
  VecUnroll(fmap(vifelse, m, getfield(v1, :data), getfield(v2, :data)))

@inline Base.Bool(m::AbstractMask{1,UInt8}) = (getfield(m, :u) & 0x01) === 0x01
@inline vconvert(::Type{Bool}, m::AbstractMask{1,UInt8}) = (getfield(m, :u) & 0x01) === 0x01
@inline vifelse(m::AbstractMask{1}, s1::T, s2::T) where {T<:NativeTypes} =
  Base.ifelse(Bool(m), s1, s2)
@inline vifelse(
  f::F,
  m::AbstractSIMD{W,B},
  a::Vararg{NativeTypesV,K},
) where {F<:Function,K,W,B<:Union{Bool,Bit}} = vifelse(m, f(a...), a[K])
@inline vifelse(f::F, m::Bool, a::Vararg{NativeTypesV,K}) where {F<:Function,K} =
  ifelse(m, f(a...), a[K])

@inline vconvert(::Type{EVLMask{W,U}}, b::Bool) where {W,U} = b & max_mask(StaticInt{W}())

@inline vifelse(m::AbstractMask{W}, a::AbstractMask{W}, b::AbstractMask{W}) where {W} =
  bitselect(m, b, a)

@inline Base.isnan(v::AbstractSIMD) = v != v
@inline Base.isfinite(x::AbstractSIMD) = iszero(x - x)

@inline Base.flipsign(x::AbstractSIMD, y::AbstractSIMD) = vifelse(y > zero(y), x, -x)
for T ∈ [:Float32, :Float64]
  @eval begin
    @inline Base.flipsign(x::AbstractSIMD, y::$T) = vifelse(y > zero(y), x, -x)
    @inline Base.flipsign(x::$T, y::AbstractSIMD) = vifelse(y > zero(y), x, -x)
  end
end
@inline Base.flipsign(x::AbstractSIMD, y::Real) = ifelse(y > zero(y), x, -x)
@inline Base.flipsign(x::Real, y::AbstractSIMD) = ifelse(y > zero(y), x, -x)
@inline Base.flipsign(x::Signed, y::AbstractSIMD) = ifelse(y > zero(y), x, -x)
@inline Base.isodd(x::AbstractSIMD{W,T}) where {W,T<:Union{Integer,StaticInt}} = (x & one(T)) != zero(T)
@inline Base.iseven(x::AbstractSIMD{W,T}) where {W,T<:Union{Integer,StaticInt}} = (x & one(T)) == zero(T)

@generated function vifelse(m::Vec{W,Bool}, v1::Vec{W,T}, v2::Vec{W,T}) where {W,T}
  typ = LLVM_TYPES[T]
  vtyp = vtype(W, typ)
  selty = vtype(W, "i1")
  f = "select"
  if Base.libllvm_version ≥ v"9" && ((T === Float32) || (T === Float64))
    f *= " nsz arcp contract reassoc"
  end
  instrs = String["%mask.0 = trunc <$W x i8> %0 to <$W x i1>"]
  push!(instrs, "%res = $f $selty %mask.0, $vtyp %1, $vtyp %2\nret $vtyp %res")
  quote
    $(Expr(:meta, :inline))
    Vec(
      $LLVMCALL(
        $(join(instrs, "\n")),
        _Vec{$W,$T},
        Tuple{_Vec{$W,Bool},_Vec{$W,$T},_Vec{$W,$T}},
        data(m),
        data(v1),
        data(v2),
      ),
    )
  end
end
@inline vifelse(b::Bool, w, x) = ((y, z) = promote(w, x); vifelse(b, y, z))
@inline vifelse(b::Bool, w::T, x::T) where {T<:Union{NativeTypes,AbstractSIMDVector}} =
  Core.ifelse(b, w, x)
@inline vifelse(b::Bool, w::T, x::T) where {T<:VecUnroll} =
  VecUnroll(fmap(Core.ifelse, b, getfield(w, :data), getfield(x, :data)))

@generated function vifelse(
  m::AbstractMask{W},
  vu1::VecUnroll{Nm1,Wsplit},
  vu2::VecUnroll{Nm1,Wsplit},
) where {W,Wsplit,Nm1}
  N = Nm1 + 1
  @assert N * Wsplit == W
  U = mask_type_symbol(Wsplit)
  quote
    $(Expr(:meta, :inline))
    vifelse(vconvert(VecUnroll{$Nm1,$Wsplit,Bit,Mask{$Wsplit,$U}}, m), vu1, vu2)
  end
end

@inline vmul(v::AbstractSIMDVector, m::AbstractMask) = vifelse(m, v, zero(v))
@inline vmul(m::AbstractMask, v::AbstractSIMDVector) = vifelse(m, v, zero(v))
@inline vmul(m1::AbstractMask, m2::AbstractMask) = m1 & m2
@inline vmul(v::AbstractSIMDVector, b::Bool) = b ? v : zero(v)
@inline vmul(b::Bool, v::AbstractSIMDVector) = b ? v : zero(v)
@inline vmul(v::VecUnroll{N,W,T}, b::Bool) where {N,W,T} = b ? v : zero(v)
@inline vmul(b::Bool, v::VecUnroll{N,W,T}) where {N,W,T} = b ? v : zero(v)

@static if Base.libllvm_version ≥ v"11"
  """
    mask(::Union{StaticInt{W},Val{W}}, base, N)
    mask(base::MM{W}, N)

  The two arg (`base`, `N`) method takes a base (current index) and last index of a loop.
  Idiomatic use for three-arg version may look like

  ```julia
  using VectorizationBase
  sp = stridedpointer(x);
  for i ∈ 1:8:N
      m = mask(Val(8), (MM{8}(i),), N) # if using an integer base, also needs a `Val` or `StaticInt` to indicate size.
      v = vload(sp, (MM{8}(i),), m)
      # do something with `v`
  end
  ```
  or, a full runnable example:
  ```julia
  using VectorizationBase, SLEEFPirates
  x = randn(117); y = similar(x);
  function vexp!(y, x)
      W = VectorizationBase.pick_vector_width(eltype(x));
      L = length(y);
      spx = stridedpointer(x); spy = stridedpointer(y);
      i = MM(W, 1); # use an `MM` index.
      while (m = mask(i,L); m !== VectorizationBase.zero_mask(W))
          yᵢ = exp(vload(spx, (i,), m))
          vstore!(spy, yᵢ, (i,), m)
          i += W
      end
  end

  vexp!(y, x)
  @assert y ≈ exp.(x)

  # A sum optimized for short vectors (e.g., 10-20 elements)
  function simd_sum(x)
      W = VectorizationBase.pick_vector_width(eltype(x));
      L = length(x);
      spx = stridedpointer(x);
      i = MM(W, 1); # use an `MM` index.
      s = VectorizationBase.vzero(W, eltype(x))
      while (m = mask(i,L); m !== VectorizationBase.zero_mask(W))
          s += vload(spx, (i,), m)
          i += W
      end
      VectorizationBase.vsum(s)
  end
  # or
  function simd_sum(x)
      W = VectorizationBase.pick_vector_width(eltype(x));
      L = length(x);
      spx = stridedpointer(x);
      i = MM(W, 1); # use an `MM` index.
      s = VectorizationBase.vzero(W, eltype(x))
      cond = true
      m = mask(i,L)
      while cond
          s += vload(spx, (i,), m)
          i += W
          m = mask(i,L)
          cond = m !== VectorizationBase.zero_mask(W)
      end
      VectorizationBase.vsum(s)
  end
  ```

  ```julia
  julia> VectorizationBase.mask(Val(8), 1, 6) # starting with `i = 1`, if vector is of length 6, 6 lanes are on
  Mask{8,Bool}<1, 1, 1, 1, 1, 1, 0, 0>

  julia> VectorizationBase.mask(Val(8), 81, 93) # if `i = 81` and the vector is of length 93, we want all lanes on.
  Mask{8,Bool}<1, 1, 1, 1, 1, 1, 1, 1>

  julia> VectorizationBase.mask(Val(8), 89, 93) # But after `i += 8`, we're at `i = 89`, and now want just 5 lanes on.
  Mask{8,Bool}<1, 1, 1, 1, 1, 0, 0, 0>
  ```
  """
  @generated function mask(
    ::Union{Val{W},StaticInt{W}},
    base::T,
    N::T,
  ) where {W,T<:IntegerTypesHW}
    # declare <8 x i1> @llvm.get.active.lane.mask.v8i1.i64(i64 %base, i64 %n)
    bits = 8sizeof(T)
    typ = "i$(bits)"
    decl = "declare <$W x i1> @llvm.get.active.lane.mask.v$(W)i1.$(typ)($(typ), $(typ))"
    instrs = [
      "%m = call <$W x i1> @llvm.get.active.lane.mask.v$(W)i1.$(typ)($(typ) %0, $(typ) %1)",
    ]
    zext_mask!(instrs, 'm', W, 0)
    push!(instrs, "ret i$(max(nextpow2(W),8)) %res.0")
    args = [:base, :N]
    call = llvmcall_expr(
      decl,
      join(instrs, "\n"),
      mask_type_symbol(W),
      :(Tuple{$T,$T}),
      "i$(max(nextpow2(W),8))",
      [typ, typ],
      args,
      true,
    )
    Expr(
      :block,
      Expr(:meta, :inline),
      :(EVLMask{$W}($call, ((N % UInt32) - (base % UInt32)) + 0x00000001)),
    )
  end
  @inline mask(i::MM{W}, N::T) where {W,T<:IntegerTypesHW} =
    mask(Val{W}(), getfield(i, :i), N)
end

@generated function Base.vcat(m1::AbstractMask{_W1}, m2::AbstractMask{_W2}) where {_W1,_W2}
  if _W1 == _W2
    W = _W1
  else
    W = _W1 + _W2
    U = integer_of_bytes_symbol(W, true)
    return Expr(
      :block,
      Expr(:meta, :inline),
      :(Mask{$W}(((data(m1) % $U) << $_W2) | (data(m2) % $U))),
    )
  end
  mtyp_input = "i$(max(8,nextpow2(W)))"
  instrs = String[]
  truncate_mask!(instrs, '0', W, 0)
  truncate_mask!(instrs, '1', W, 1)

  W2 = W + W
  shuffmask = Vector{String}(undef, W2)
  for w ∈ eachindex(shuffmask)
    shuffmask[w] = string(w - 1)
  end
  mask = '<' * join(map(x -> string("i32 ", x), shuffmask), ", ") * '>'

  push!(
    instrs,
    "%combinedmask = shufflevector <$W x i1> %mask.0, <$W x i1> %mask.1, <$(W2) x i32> $mask",
  )

  mtyp_output = "i$(max(8,nextpow2(W2)))"
  zext_mask!(instrs, "combinedmask", W2, 1)
  push!(instrs, "ret $mtyp_output %res.1")
  instrj = join(instrs, "\n")
  U = mask_type_symbol(W)
  U2 = mask_type_symbol(W2)
  Expr(
    :block,
    Expr(:meta, :inline),
    :(Mask{$W2}($LLVMCALL($instrj, $U2, Tuple{$U,$U}, getfield(m1, :u), getfield(m2, :u)))),
  )
end
# @inline function Base.vcat(m1::AbstractMask{W}, m2::AbstractMask{W}) where {W}
#     U = mask_type(Val(W))
#     u1 = data(m1) % U
#     u2 = data(m2) % U
#     (u1 << W) | u2
# end

@inline ifelse(b::Bool, m1::Mask{W}, m2::Mask{W}) where {W} =
  Mask{W}(Core.ifelse(b, getfield(m1, :u), getfield(m2, :u)))
@inline ifelse(b::Bool, m1::Mask{W}, m2::EVLMask{W}) where {W} =
  Mask{W}(Core.ifelse(b, getfield(m1, :u), getfield(m2, :u)))
@inline ifelse(b::Bool, m1::EVLMask{W}, m2::Mask{W}) where {W} =
  Mask{W}(Core.ifelse(b, getfield(m1, :u), getfield(m2, :u)))
@inline ifelse(b::Bool, m1::EVLMask{W}, m2::EVLMask{W}) where {W} = EVLMask{W}(
  Core.ifelse(b, getfield(m1, :u), getfield(m2, :u)),
  Core.ifelse(b, getfield(m1, :evl), getfield(m2, :evl)),
)

@inline vconvert(::Type{<:AbstractMask{W}}, b::Bool) where {W} =
  b ? max_mask(Val(W)) : zero_mask(Val(W))
@inline vconvert(::Type{Mask{W}}, b::Bool) where {W} =
  b ? max_mask(Val(W)) : zero_mask(Val(W))
@inline vconvert(::Type{EVLMask{W}}, b::Bool) where {W} =
  b ? max_mask(Val(W)) : zero_mask(Val(W))

@inline Base.max(x::AbstractMask, y::AbstractMask) = x | y
@inline Base.min(x::AbstractMask, y::AbstractMask) = x & y
@inline Base.FastMath.max_fast(x::AbstractMask, y::AbstractMask) = x | y
@inline Base.FastMath.min_fast(x::AbstractMask, y::AbstractMask) = x & y


ipopt_optimizer =
    JuMP.optimizer_with_attributes(Ipopt.Optimizer, "tol" => 1e-6, "print_level" => 0)
fast_ipopt_optimizer = JuMP.optimizer_with_attributes(
    Ipopt.Optimizer,
    "print_level" => 0,
    "max_cpu_time" => 5.0,
)
# use default print_level = 5 # set to 0 to disable
GLPK_optimizer = JuMP.optimizer_with_attributes(GLPK.Optimizer)
Cbc_optimizer = JuMP.optimizer_with_attributes(Cbc.Optimizer)
OSQP_optimizer =
    JuMP.optimizer_with_attributes(OSQP.Optimizer, "verbose" => false, "max_iter" => 50000)
fast_lp_optimizer = JuMP.optimizer_with_attributes(Cbc.Optimizer, "seconds" => 3.0)
scs_solver = JuMP.optimizer_with_attributes(
    SCS.Optimizer,
    "max_iters" => 100000,
    "eps" => 1e-4,
    "verbose" => 0,
)


# Copyright (C) 2020 Argonne National Laboratory
# Written by Alinson Santos Xavier <axavier@anl.gov>

using RELOG

@testset "compress" begin
    basedir = dirname(@__FILE__)
    instance = RELOG.parsefile("$basedir/../../instances/s1.json")
    compressed = RELOG._compress(instance)

    product_name_to_product = Dict(p.name => p for p in compressed.products)
    location_name_to_facility = Dict()
    for p in compressed.plants
        location_name_to_facility[p.location_name] = p
    end
    for c in compressed.collection_centers
        location_name_to_facility[c.name] = c
    end

    p1 = product_name_to_product["P1"]
    p2 = product_name_to_product["P2"]
    p3 = product_name_to_product["P3"]
    c1 = location_name_to_facility["C1"]
    l1 = location_name_to_facility["L1"]

    @test compressed.time == 1
    @test compressed.building_period == [1]

    @test p1.name == "P1"
    @test p1.transportation_cost ≈ [0.015]
    @test p1.transportation_energy ≈ [0.115]
    @test p1.transportation_emissions["CO2"] ≈ [0.051]
    @test p1.transportation_emissions["CH4"] ≈ [0.0025]

    @test c1.name == "C1"
    @test c1.amount ≈ [1869.12]

    @test l1.plant_name == "F1"
    @test l1.location_name == "L1"
    @test l1.energy ≈ [0.115]
    @test l1.emissions["CO2"] ≈ [0.051]
    @test l1.emissions["CH4"] ≈ [0.0025]
    @test l1.sizes[1].opening_cost ≈ [500]
    @test l1.sizes[2].opening_cost ≈ [1250]
    @test l1.sizes[1].fixed_operating_cost ≈ [60]
    @test l1.sizes[2].fixed_operating_cost ≈ [60]
    @test l1.sizes[1].variable_operating_cost ≈ [30]
    @test l1.sizes[2].variable_operating_cost ≈ [30]
    @test l1.disposal_limit[p2] ≈ [2.0]
    @test l1.disposal_limit[p3] ≈ [2.0]
    @test l1.disposal_cost[p2] ≈ [-10.0]
    @test l1.disposal_cost[p3] ≈ [-10.0]
end


using Oceananigans.Models: ShallowWaterModel
using Oceananigans.Grids: Periodic, Bounded

@testset "Shallow Water Models" begin
    @info "Testing shallow water models..."

    @testset "Model constructor errors" begin
        grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1))
        @test_throws TypeError ShallowWaterModel(architecture=CPU, grid=grid)
        @test_throws TypeError ShallowWaterModel(architecture=GPU, grid=grid)
    end

    topos = (
             (Periodic, Periodic,  Bounded),
             (Periodic,  Bounded,  Bounded),
             (Bounded,   Bounded,  Bounded),
            )

    for topo in topos
        @testset "$topo model construction" begin
            @info "  Testing $topo model construction..."
            for arch in archs, FT in float_types
		        arch isa GPU && topo == (Bounded, Bounded, Bounded) && continue

                grid = RegularCartesianGrid(FT, topology=topo, size=(1, 1, 1), extent=(1, 2, 3))
                model = ShallowWaterModel(grid=grid, architecture=arch, float_type=FT)

                # Just testing that the model was constructed with no errors/crashes.
                @test model isa ShallowWaterModel

                # Test that the grid didn't get mangled (sort of)
                @test size(grid) === size(model.grid)

                too_big_grid = RegularCartesianGrid(FT, topology=topo, size=(1, 1, 2), extent=(1, 2, 3))

                @test_throws ArgumentError ShallowWaterModel(grid=too_big_grid, architecture=arch, float_type=FT)
            end
        end
    end

    @testset "Setting ShallowWaterModel fields" begin
        @info "  Testing setting shallow water model fields..."
        for arch in archs, FT in float_types
            N = (4, 4, 1)
            L = (2π, 3π, 5π)

            grid = RegularCartesianGrid(FT, size=N, extent=L)
            model = ShallowWaterModel(grid=grid, architecture=arch, float_type=FT)

            x, y, z = nodes((Face, Cell, Cell), model.grid, reshape=true)

            uh₀(x, y, z) = x * y^2
            uh_answer = @. x * y^2

            h₀ = rand(size(grid)...)
            h_answer = deepcopy(h₀)

            set!(model, uh=uh₀, h=h₀)

            uh, vh, h = model.solution

            @test all(interior(uh) .≈ uh_answer)
            @test all(interior(h) .≈ h_answer)
        end
    end

    for arch in archs, topo in topos
        @testset "Time-stepping ShallowWaterModels [$arch, $topo]" begin
            @info "  Testing time-stepping ShallowWaterModels [$arch, $topo]..."

            grid = RegularCartesianGrid(size=(1, 1, 1), extent=(2π, 2π, 2π), topology=topo)
            model = ShallowWaterModel(grid=grid, architecture=arch)
            simulation = Simulation(model, Δt=1.0, stop_iteration=1)

            run!(simulation)

            @test model.clock.iteration == 1
        end
    end
end


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

export
    AbstractGraph,
    UndirectedGraph, DirectedGraph, HyperGraph,
    OrderedGraph, OrderedDiGraph, OrderedHyperGraph,
    AbstractNode, AbstractEdge, UndirectedEdge, DirectedEdge,
    hascachestore, hascache, getcache, setcache!, clearcache!, @cachefirst,
    clone,
    neighbors, outneighbors, inneighbors,
    findedgekey, findalledgekeys, getedge, hasedge, nodeattr, edgeattr,
    adjacencies, successors, predecessors,
    incidences, outincidences, in_incidences,
    nodeset, edgeset, edgesiter, neighborsiter, nodeattrs, edgeattrs,
    nodecount, edgecount, neighborcount, degree, indegree, outdegree,
    addnode!, addedge!, setnodeattr!, setedgeattr!,
    unlinknodes, unlinkedges,
    nodeattrtype, edgeattrtype

# Graph types

abstract type AbstractGraph end
abstract type UndirectedGraph <: AbstractGraph end
abstract type DirectedGraph <: AbstractGraph end
abstract type HyperGraph <: AbstractGraph end
abstract type OrderedGraph <: UndirectedGraph end
abstract type OrderedDiGraph <: DirectedGraph end
abstract type OrderedHyperGraph <: HyperGraph end

# Node and edge attributes

abstract type AbstractNode end
abstract type AbstractEdge end
abstract type UndirectedEdge <: AbstractEdge end # TODO: unnecessary?
abstract type DirectedEdge <: AbstractEdge end # TODO: unnecessary?

"""
    hascachestore(graph::AbstractGraph) -> Bool

Return whether the graph has calculated property cache store.
"""
hascachestore(graph::AbstractGraph) = isdefined(graph, :cache)

"""
    hascache(graph::AbstractGraph, key) -> Bool

Return whether the graph has calculated properties with the given key or not.
"""
hascache(graph::AbstractGraph, key) = haskey(graph.cache, key)

"""
    getcache(graph::AbstractGraph, key) -> Any

Get cached calculated properties.
"""
getcache(graph::AbstractGraph, key) = graph.cache[key]

"""
    setcache!(graph::AbstractGraph, key; kwargs...)

Set calculated property caches.
"""
function setcache!(graph::AbstractGraph, key; kwargs...)
    graph.cache[key] = getfield(Graph, key)(graph; kwargs...)
end

"""
    clearcache!(graph::AbstractGraph)

Clear calculated property caches.

Calling `clearcache!` is recommended when the graph nodes/edges are added, removed or reindexed. You can `deepcopy` or [`Graph.clone`](@ref) the graph before destructive operation instead.
"""
clearcache!(graph::AbstractGraph) = empty!(graph.cache)

"""
    @cachefirst expression

A macro for molecular descriptor array cache mechanism.
"""
macro cachefirst(ex)
    func = ex.args[1].args[1]
    dummy = gensym()
    ex.args[1].args[1] = dummy
    return esc(quote
        $ex
        Core.@__doc__ function $func(graph; kwargs...)
            if hascachestore(graph) && isempty(kwargs)
                # Return cache
                symf = nameof($func)
                if hascache(graph, symf)
                    return getcache(graph, symf)
                end
            end
            return $dummy(graph; kwargs...)
        end
    end)
end

"""
    clone(graph::AbstractGraph) -> AbstractGraph

Return deep copy of the graph.
"""
clone = deepcopy  # TODO: use deepcopy

# Lookup

"""
    neighbors(graph, i) -> Dict

Return the mapping of incident edges and adjacent nodes of node `i`.
If the graph is directed graph, both outneighbors and inneighbors are mapped.
"""
neighbors(graph::UndirectedGraph, i) = graph.neighbormap[i]
neighbors(graph::DirectedGraph, i) = merge(outneighbors(graph, i), inneighbors(graph, i))

neighbors(graph::HyperGraph, i) = Dict(inc => graph.edges[inc] for inc in graph.incidences[i])

"""
    outneighbors(graph::DirectedGraph, i) -> Dict

Return the mapping of successor node keys and out edge keys connected to
the given node.
"""
outneighbors(graph::DirectedGraph, i) = graph.outneighbormap[i]

"""
    inneighbors(graph::DirectedGraph, i) -> Dict

Return the mapping of predecessor node keys and in edge keys connected to
the given node.
"""
inneighbors(graph::DirectedGraph, i) = graph.inneighbormap[i]

"""
    findedgekey(graph::UndirectedGraph, u, v) -> Any
    findedgekey(graph::DirectedGraph, u, v) -> Any
    findedgekey(graph::HyperGraph, u, v) -> Any

Return the first edge found that connects nodes u and v.
"""
function findedgekey(graph::UndirectedGraph, u, v)
    for (inc, adj) in neighbors(graph, u)
        adj == v && return inc
    end
end

function findedgekey(graph::DirectedGraph, source, target)
    for (inc, adj) in outneighbors(graph, source)
        adj == target && return inc
    end
end

function findedgekey(graph::HyperGraph, u, v)
    for (inc, adj) in neighbors(graph, u)
        v in adj && return inc
    end
end

"""
    findalledgekeys(graph::UndirectedGraph, u, v) -> Any
    findalledgekeys(graph::DirectedGraph, u, v) -> Any
    findalledgekeys(graph::HyperGraph, u, v) -> Any

Return all edges that connect nodes u and v.
"""
function findalledgekeys(graph::UndirectedGraph, u, v)
    hasmultiedge(graph) || return [findedgekey(graph, u, v)]
    keys = Int[]
    for (inc, adj) in neighbors(graph, u)
        adj == v && push!(keys, inc)
    end
    return keys
end

function findalledgekeys(graph::DirectedGraph, source, target)
    hasmultiedge(graph) || return [findedgekey(graph, source, target)]
    keys = Int[]
    for (inc, adj) in outneighbors(graph, source)
        adj == target && push!(keys, inc)
    end
    return keys
end

function findalledgekeys(graph::HyperGraph, u, v)
    hasmultiedge(graph) || return [findedgekey(graph, u, v)]
    keys = Int[]
    for (inc, adj) in neighbors(graph, u)
        v in adj && push!(keys, inc)
    end
    return keys
end

"""
    getedge(graph::AbstractGraph, i) -> Any

Return an edge.
"""
getedge(graph::AbstractGraph, i) = graph.edges[i]

"""
    hasedge(graph, u, v) -> Bool

Return whether the given two nodes are connected by at least one edge.
"""
hasedge(graph, u, v) = findedgekey(graph, u, v) !== nothing

"""
    nodeattr(graph::AbstractGraph, i) -> AbstractNode

Return the attribute object of node `i`.
"""
nodeattr(graph::AbstractGraph, i) = graph.nodeattrs[i]

"""
    edgeattr(graph::AbstractGraph, i) -> AbstractEdge

Return the attribute object of edge `i`.
"""
edgeattr(graph::AbstractGraph, i) = graph.edgeattrs[i]

"""
    edgeattr(graph::AbstractGraph, u, v) -> Union{AbstractEdge,Nothing}

Return the attribute object of an edge that connects `u` and `v`. If not found,
return nothing.
"""
function edgeattr(graph::AbstractGraph, u, v)
    k = findedgekey(graph, u, v)
    return k === nothing ? nothing : edgeattr(graph, k)
end

# Collection

adjacencies(g::AbstractGraph, i) = Set(values(neighbors(g, i)))
successors(g::DirectedGraph, i) = Set(values(outneighbors(g, i)))
predecessors(g::DirectedGraph, i) = Set(values(inneighbors(g, i)))

incidences(g::AbstractGraph, i) = Set(keys(neighbors(g, i)))
outincidences(g::DirectedGraph, i) = Set(keys(outneighbors(g, i)))
in_incidences(g::DirectedGraph, i) = Set(keys(inneighbors(g, i)))

"""
    nodeset(graph::Union{OrderedGraph,OrderedDiGraph,OrderedHyperGraph}) -> Set{Int}

Return the set of node keys.
"""
nodeset(graph::Union{OrderedGraph,OrderedDiGraph,OrderedHyperGraph}) = Set(1:nodecount(graph))

"""
    edgeset(graph::Union{OrderedGraph,OrderedDiGraph,OrderedHyperGraph}) -> Set{Int}

Return the set of edge keys.
"""
edgeset(graph::Union{OrderedGraph,OrderedDiGraph,OrderedHyperGraph}) = Set(1:edgecount(graph))

edgesiter(graph::Union{OrderedGraph,OrderedDiGraph,OrderedHyperGraph}) = graph.edges

neighborsiter(graph::OrderedGraph) = graph.neighbormap

"""
    nodeattrs(graph::Union{OrderedGraph,OrderedDiGraph,OrderedHyperGraph}) -> Vector{AbstractNode}

Return graph node attributes.
"""
nodeattrs(graph::Union{OrderedGraph,OrderedDiGraph,OrderedHyperGraph}) = graph.nodeattrs

"""
    edgeattrs(graph::Union{OrderedGraph,OrderedDiGraph,OrderedHyperGraph}) -> Vector{AbstractEdge}

Return graph edge attributes.
"""
edgeattrs(graph::Union{OrderedGraph,OrderedDiGraph,OrderedHyperGraph}) = graph.edgeattrs

# Counting

"""
    nodecount(graph::AbstractGraph) -> Int

Return the number of graph nodes.
"""
nodecount(graph::UndirectedGraph) = length(graph.neighbormap)
nodecount(graph::DirectedGraph) = length(graph.outneighbormap)
nodecount(graph::HyperGraph) = length(graph.incidences)

"""
    edgecount(graph::AbstractGraph) -> Int

Return the number of graph edges.
"""
edgecount(graph::AbstractGraph) = length(graph.edges)

"""
    neighborcount(graph::AbstractGraph, n) -> Int
    degree(graph::AbstractGraph, n) -> Int

Return the number of adjacent nodes of the node 'n'.
"""
neighborcount(graph::AbstractGraph, n) = length(neighbors(graph, n))
degree = neighborcount

"""
    outdegree(graph::DirectedGraph, n) -> Int

Return the number of outneighbors of the node 'n'.
"""
outdegree(graph::DirectedGraph, n) = length(outneighbors(graph, n))

"""
    indegree(graph::DirectedGraph, n) -> Int

Return the number of inneighbors of the node 'n'.
"""
indegree(graph::DirectedGraph, n) = length(inneighbors(graph, n))

# Editing

"""
    addnode!(graph) -> Int
    addnode!(graph, attr) -> Int

Add new node and return the node index. If the node attribute type is required,
specify the node attribute object by `node` keyword.
"""
function addnode!(graph::UndirectedGraph, attr)
    push!(graph.neighbormap, Dict())
    push!(graph.nodeattrs, attr)
    return length(graph.neighbormap)
end

function addnode!(graph::UndirectedGraph)
    isdefined(graph, :nodeattrs) && throw(ErrorException("nodeattr required"))
    push!(graph.neighbormap, Dict())
    return length(graph.neighbormap)
end

function addnode!(graph::DirectedGraph, attr)
    push!(graph.outneighbormap, Dict())
    push!(graph.inneighbormap, Dict())
    push!(graph.nodeattrs, attr)
    return length(graph.outneighbormap)
end

function addnode!(graph::DirectedGraph)
    isdefined(graph, :nodeattrs) && throw(ErrorException("nodeattr required"))
    push!(graph.outneighbormap, Dict())
    push!(graph.inneighbormap, Dict())
    return length(graph.outneighbormap)
end

function addnode!(graph::HyperGraph, attr)
    push!(graph.incidences, Set())
    push!(graph.nodeattrs, attr)
    return length(graph.incidences)
end

function addnode!(graph::HyperGraph)
    isdefined(graph, :nodeattrs) && throw(ErrorException("nodeattr required"))
    push!(graph.incidences, Set())
    return length(graph.incidences)
end

"""
    addedge!(graph, u, v) -> Int
    addedge!(graph, u, v, attr) -> Int

Add new edge and return the edge index. If the edge attribute type is required,
specify the edge attribute object by `attr` keyword.
"""
function _addedge!(graph::UndirectedGraph, u, v)
    push!(graph.edges, (u, v))
    i = edgecount(graph)
    graph.neighbormap[u][i] = v
    graph.neighbormap[v][i] = u
    return i
end

function addedge!(graph::UndirectedGraph, u, v, attr)
    push!(graph.edgeattrs, attr)
    return _addedge!(graph, u, v)
end

function addedge!(graph::UndirectedGraph, u, v)
    isdefined(graph, :edgeattrs) && throw(ErrorException("edgeattr required"))
    return _addedge!(graph, u, v)
end

function _addedge!(graph::DirectedGraph, s, t)
    push!(graph.edges, (s, t))
    i = edgecount(graph)
    graph.outneighbormap[s][i] = t
    graph.inneighbormap[t][i] = s
    return i
end

function addedge!(graph::DirectedGraph, s, t, attr)
    push!(graph.edgeattrs, attr)
    return _addedge!(graph, s, t)
end

function addedge!(graph::DirectedGraph, s, t)
    isdefined(graph, :edgeattrs) && throw(ErrorException("edgeattr required"))
    return _addedge!(graph, s, t)
end

function _addedge!(graph::HyperGraph, edge)
    push!(graph.edges, edge)
    i = edgecount(graph)
    for n in edge
        push!(graph.incidences[n], i)
    end
    return i
end

function addedge!(graph::HyperGraph, edge, attr)
    push!(graph.edgeattrs, attr)
    return _addedge!(graph, edge)
end

function addedge!(graph::HyperGraph, edge)
    isdefined(graph, :edgeattrs) && throw(ErrorException("edgeattr required"))
    return _addedge!(graph, edge)
end

"""
    setnodeattr!(graph::AbstractGraph, i, attr::AbstractNode)

Update the node attribute.
"""
function setnodeattr!(graph::AbstractGraph, i, attr)
    graph.nodeattrs[i] = attr
    return
end

"""
    setedgeattr!(graph::AbstractGraph, i, attr::AbstractNode)

Update the edge attribute.
"""
function setedgeattr!(graph::AbstractGraph, i, attr)
    graph.edgeattrs[i] = attr
    return
end

"""
    unlinknodes(graph, nodes)

Delete given nodes and its incident edges from the graph.
"""
function unlinknodes end

"""
    unlinkedges(graph, edges)

Delete given edges from the graph.
"""
function unlinkedges end

# Types

"""
    nodeattrtype(graph) -> Type

Return node attribute type.
"""
nodeattrtype(mol::AbstractGraph) = eltype(mol.nodeattrs)

"""
    edgeattrtype(graph) -> Type

Return edge attribute type.
"""
edgeattrtype(mol::AbstractGraph) = eltype(mol.edgeattrs)


# random functions that dispatch either to rocRAND or GPUArrays' generic RNG

using Random

const GPUARRAY_RNG = Ref{Union{Nothing,GPUArrays.RNG}}(nothing)

function GPUArrays.default_rng(::Type{<:ROCArray})
    if GPUARRAY_RNG[] == nothing
        agent = AMDGPU.default_device().device.agent
        p = Ref{UInt32}()
        GC.@preserve p begin
            AMDGPU.getinfo(agent, HSA.AGENT_INFO_WORKGROUP_MAX_SIZE, p) |> AMDGPU.check
            N = Int(p[])
        end
        state = ROCArray{NTuple{4, UInt32}}(undef, N)
        GPUARRAY_RNG[] = GPUArrays.RNG(state)
        Random.seed!(GPUARRAY_RNG[])
    end
    return GPUARRAY_RNG[]::GPUArrays.RNG
end

gpuarrays_rng() = GPUArrays.default_rng(ROCArray)
const rocrand_rng = librocrand !== nothing ? rocRAND.default_rng : gpuarrays_rng

# the interface is split in two levels:
# - functions that extend the Random standard library, and take an RNG as first argument,
#   will only ever dispatch to rocRAND and as a result are limited in the types they support.
# - functions that take an array will dispatch to either rocRAND or GPUArrays
# - non-exported functions are provided for constructing GPU arrays from only an eltype

# these types are supported by rocrand and we use rocrand whenever possible
const rocRANDUniformType = Union{Type{UInt32}, Type{Cuchar}, Type{UInt16}, Type{Float16}, Type{Float32}, Type{Float64}}
const rocRANDUniformArray = ROCArray{<:Union{UInt32, Cuchar, UInt16, Float16, Float32,Float64}}
const rocRANDNormalType = Union{Type{Float16}, Type{Float32}, Type{Float64}}
const rocRANDNormalArray = ROCArray{<:Union{Float16, Float32,Float64}}

function seed!(seed=Base.rand(UInt64))
    Random.seed!(rocrand_rng(), seed)
    Random.seed!(gpuarrays_rng(), seed)
end

# rocRAND in-place
Random.rand!(A::rocRANDUniformArray) = Random.rand!(rocrand_rng(), A)
Random.randn!(A::rocRANDNormalArray; kwargs...) = Random.randn!(rocrand_rng(), A; kwargs...)

# FIXME: GPUArrays rng does not support Float16
Random.rand!(rng::GPUArrays.RNG, x::ROCArray{Float16}) =
    error("GPUArrays.jl does not support Float16, use rocRAND instead.")
Random.randn!(rng::GPUArrays.RNG, x::ROCArray{Float16}) =
    error("GPUArrays.jl does not support Float16, use rocRAND instead.")

# rocRAND out-of-place
rand(T::rocRANDUniformType, dims::Dims) = Random.rand(rocrand_rng(), T, dims)
randn(T::rocRANDNormalType, dims::Dims; kwargs...) = Random.randn(rocrand_rng(), T, dims; kwargs...)

# support all dimension specifications
rand(T::rocRANDUniformType, dim1::Integer, dims::Integer...) =
    Random.rand(rocrand_rng(), T, Dims((dim1, dims...)))
randn(T::rocRANDNormalType, dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn(rocrand_rng(), T, Dims((dim1, dims...)); kwargs...)

# GPUArrays in-place
Random.rand!(A::ROCArray) = Random.rand!(gpuarrays_rng(), A)
Random.randn!(A::ROCArray; kwargs...) =
    error("AMDGPU.jl does not support generating normally-distributed random numbers of type $(eltype(A))")

# GPUArrays out-of-place
rand(T::Type, dims::Dims) = Random.rand!(ROCArray{T}(undef, dims...))
randn(T::Type, dims::Dims; kwargs...) = Random.randn!(ROCArray{T}(undef, dims...); kwargs...)

# support all dimension specifications
rand(T::Type, dim1::Integer, dims::Integer...) =
    Random.rand!(ROCArray{T}(undef, dim1, dims...))
randn(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn!(ROCArray{T}(undef, dim1, dims...); kwargs...)

# untyped out-of-place
rand(dim1::Integer, dims::Integer...) =
    Random.rand(rocrand_rng(), Dims((dim1, dims...)))
randn(dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn(rocrand_rng(), Dims((dim1, dims...)); kwargs...)

# rand_logn, rand_poisson
const rand_logn = librocrand !== nothing ? rocRAND.rand_logn : (x...;kwargs...) -> error("Not supported without rocRAND.")
const rand_poisson = librocrand !== nothing ? rocRAND.rand_poisson : (x...;kwargs...) -> error("Not supported without rocRAND.")


function _has_lb(model::Optimizer, index::MOI.VariableIndex)
    return _info(model, index).lb !== nothing
end
function _has_ub(model::Optimizer, index::MOI.VariableIndex)
    return _info(model, index).ub !== nothing
end

function MOI.is_valid(
    model::Optimizer,
    c::MOI.ConstraintIndex{MOI.VariableIndex, MOI.LessThan{T}},
) where {T <: Real}
    index = MOI.VariableIndex(c.value)
    return MOI.is_valid(model, index) && _has_ub(model, index)
end

function MOI.is_valid(
    model::Optimizer,
    c::MOI.ConstraintIndex{MOI.VariableIndex, MOI.GreaterThan{T}},
) where {T <: Real}
    index = MOI.VariableIndex(c.value)
    return MOI.is_valid(model, index) && _has_lb(model, index)
end

function MOI.is_valid(
    model::Optimizer,
    c::MOI.ConstraintIndex{MOI.VariableIndex, MOI.Interval{T}},
) where {T <: Real}
    index = MOI.VariableIndex(c.value)
    return MOI.is_valid(model, index) &&
           _has_lb(model, index) &&
           _has_ub(model, index)
end

function MOI.is_valid(
    model::Optimizer,
    c::MOI.ConstraintIndex{MOI.VariableIndex, MOI.EqualTo{T}},
) where {T <: Real}
    index = MOI.VariableIndex(c.value)
    return MOI.is_valid(model, index) &&
           _info(model, index).lb == _info(model, index).ub
end

function _build_constraint(
    model::Optimizer,
    f::MOI.VariableIndex,
    s::MOI.EqualTo{T},
) where {T <: Integer}
    v = _info(model, f).variable
    return XeqC(v, s.value)
end

function MOI.is_valid(
    model::Optimizer,
    c::MOI.ConstraintIndex{MOI.VariableIndex, MOI.ZeroOne},
)
    index = MOI.VariableIndex(c.value)
    if !MOI.is_valid(model, index)
        return false
    end

    return _info(model, index).type == BINARY
end

function MOI.is_valid(
    model::Optimizer,
    c::MOI.ConstraintIndex{MOI.VariableIndex, MOI.Integer},
)
    index = MOI.VariableIndex(c.value)
    if !MOI.is_valid(model, index)
        return false
    end

    return _info(model, index).type == INTEGER
end

function MOI.get(
    model::Optimizer,
    ::MOI.ConstraintFunction,
    c::MOI.ConstraintIndex{MOI.VariableIndex, <:Any},
)
    MOI.throw_if_not_valid(model, c)
    return MOI.VariableIndex(MOI.VariableIndex(c.value))
end

function MOI.set(
    ::Optimizer,
    ::MOI.ConstraintFunction,
    ::MOI.ConstraintIndex{MOI.VariableIndex, S},
    ::MOI.VariableIndex,
) where {S}
    throw(MOI.SettingVariableIndexFunctionNotAllowed())
    return
end

# CP.Domain.
function MOI.supports_constraint(
    ::Optimizer,
    ::MOI.VariableIndex,
    ::CP.Domain{T},
) where {T <: Integer}
    return true
end

function MOI.add_constraint(
    model::Optimizer,
    f::MOI.VariableIndex,
    s::CP.Domain{T},
) where {T <: Integer}
    # Add each possible value separately. Same logic as
    # https://github.com/radsz/jacop/blob/develop/src/main/java/org/jacop/fz/VariablesParameters.java#L238-L267
    v = _info(model, f).variable
    for e in s.values
        jcall(v, "addDom", (jint, jint), e, e)
    end

    index = MOI.ConstraintIndex{F, S}(length(model.constraint_info) + 1)
    model.constraint_info[index] = ConstraintInfo(index, nothing, f, s)
    return index
end