module SIMD

# A note on Val{} vs. Val():
#
# For historic reasoons, SIMD's API accepted compile-time constants as
# Val{N} instead of Val(N). The difference is that Val{N} is a type
# (Type{Val{N}}), whereas Val(N) is a value (of type Val{N}). This is
# against the intent of how Val is designed, and is also much slower
# at run time unless functions are either @inline'd or @generated.
#
# The API has now been cleaned up. To preserve backward compatibility,
# passing Val{N} instead of Val(N) is still supported. It might go
# away at the next major release.

#=

# Various boolean types

# Idea (from <Gaunard-simd.pdf>): Use Mask{N,T} instead of booleans
# with different sizes

abstract Boolean <: Integer

for sz in (8, 16, 32, 64, 128)
    Intsz = Symbol(:Int, sz)
    UIntsz = Symbol(:UInt, sz)
    Boolsz = Symbol(:Bool, sz)
    @eval begin
        immutable $Boolsz <: Boolean
            int::$UIntsz
            $Boolsz(b::Bool) =
                new(ifelse(b, typemax($UIntsz), typemin($UIntsz)))
        end
        booltype(::Val($sz)) = $Boolsz
        inttype(::Val($sz)) = $Intsz
        uinttype(::Val($sz)) = $UIntsz

        Base.convert(::Type{Bool}, b::$Boolsz) = b.int != 0

        Base.:~(b::$Boolsz) = $Boolsz(~b.int)
        Base.:!(b::$Boolsz) = ~b
        Base.:&(b1::$Boolsz, b2::$Boolsz) = $Boolsz(b1.int & b2.int)
        Base.:|(b1::$Boolsz, b2::$Boolsz) = $Boolsz(b1.int | b2.int)
        Base.$(:$)(b1::$Boolsz, b2::$Boolsz) = $Boolsz(b1.int $ b2.int)

        Base.:==(b1::$Boolsz, b2::$Boolsz) = $Boolsz(b1.int == b2.int)
        Base.:!=(b1::$Boolsz, b2::$Boolsz) = $Boolsz(b1.int != b2.int)
        Base.:<(b1::$Boolsz, b2::$Boolsz) = $Boolsz(b1.int < b2.int)
        Base.:<=(b1::$Boolsz, b2::$Boolsz) = $Boolsz(b1.int <= b2.int)
        Base.:>(b1::$Boolsz, b2::$Boolsz) = $Boolsz(b1.int > b2.int)
        Base.:>=(b1::$Boolsz, b2::$Boolsz) = $Boolsz(b1.int >= b2.int)
    end
end
Base.convert(::Type{Bool}, b::Boolean) = error("impossible")
Base.convert{I<:Integer}(::Type{I}, b::Boolean) = I(Bool(b))
Base.convert{B<:Boolean}(::Type{B}, b::Boolean) = B(Bool(b))
Base.convert{B<:Boolean}(::Type{B}, i::Integer) = B(i!=0)

booltype{T}(::Type{T}) = booltype(Val(8*sizeof(T)))
inttype{T}(::Type{T}) = inttype(Val(8*sizeof(T)))
uinttype{T}(::Type{T}) = uinttype(Val(8*sizeof(T)))

=#

# Array types for SIMD

using Base: Slice, ScalarIndex

"""
    ContiguousSubArray{T,N,P,I,L}

Like `Base.FastContiguousSubArray` but without requirement for linear
indexing (i.e., type parameter `L` can be `false`).

# Examples
```
julia> A = view(ones(5, 5), :, [1,3]);

julia> A isa Base.FastContiguousSubArray
false

julia> A isa SIMD.ContiguousSubArray
true
```
"""
ContiguousSubArray{T,N,P,
                   I<:Union{Tuple{Union{Slice, AbstractUnitRange}, Vararg{Any}},
                            Tuple{Vararg{ScalarIndex}}},
                   L} = SubArray{T,N,P,I,L}

"""
    ContiguousArray{T,N}

Array types with contiguous first dimension.
"""
ContiguousArray{T,N} = Union{DenseArray{T,N}, ContiguousSubArray{T,N}}

"""
    FastContiguousArray{T,N}

This is the type of arrays that `pointer(A, i)` works.
"""
FastContiguousArray{T,N} = Union{DenseArray{T,N}, Base.FastContiguousSubArray{T,N}}
# https://github.com/eschnett/SIMD.jl/pull/40#discussion_r254131184
# https://github.com/JuliaArrays/MappedArrays.jl/pull/24#issuecomment-460568978

# The Julia SIMD vector type

const BoolTypes = Union{Bool}
const IntTypes = Union{Int8, Int16, Int32, Int64, Int128}
const UIntTypes = Union{UInt8, UInt16, UInt32, UInt64, UInt128}
const IntegerTypes = Union{BoolTypes, IntTypes, UIntTypes}
const IndexTypes = Union{IntegerTypes, Ptr}
const FloatingTypes = Union{Float16, Float32, Float64}
const ScalarTypes = Union{IndexTypes, FloatingTypes}

const VE = Base.VecElement

export Vec
struct Vec{N,T<:ScalarTypes} # <: Number
    elts::NTuple{N,VE{T}}
    @inline Vec{N,T}(elts::NTuple{N, VE{T}}) where {N,T} = new{N,T}(elts)
end

function Base.show(io::IO, v::Vec{N,T}) where {N,T}
    print(io, "<$N x $T>[")
    join(io, [x.value for x in v.elts], ", ")
    print(io, "]")
end

# Type properties
Base.eltype(::Type{Vec{N,T}}) where {N,T} = T
Base.ndims( ::Type{Vec{N,T}}) where {N,T} = 1
Base.length(::Type{Vec{N,T}}) where {N,T} = N
Base.size(  ::Type{Vec{N,T}}) where {N,T} = (N,)
# TODO: This doesn't follow Base, e.g. `size([], 3) == 1`
Base.size(::Type{Vec{N,T}}, n::Integer) where {N,T} = (N,)[n]

Base.eltype(V::Vec) = eltype(typeof(V))
Base.ndims( V::Vec) = ndims(typeof(V))
Base.length(V::Vec) = length(typeof(V))
Base.size(  V::Vec) = size(typeof(V))
Base.size(  V::Vec, n::Integer) = size(typeof(V), n)

# Type conversion

# Create vectors from scalars or tuples
@generated function (::Type{Vec{N,T}})(x::S) where {N,T,S<:ScalarTypes}
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(tuple($([:(VE{T}(T(x))) for i in 1:N]...)))
    end
end
Vec{N,T}(xs::Tuple{}) where {N,T<:ScalarTypes} = error("illegal argument")
@generated function (::Type{Vec{N,T}})(xs::NTuple{N,S}) where {N,T,S<:ScalarTypes}
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(tuple($([:(VE{T}(T(xs[$i]))) for i in 1:N]...)))
    end
end
Vec(xs::NTuple{N,T}) where {N,T<:ScalarTypes} = Vec{N,T}(xs)

# Convert between vectors
@inline Base.convert(::Type{Vec{N,T}}, v::Vec{N,T}) where {N,T} = v

@inline Base.convert(::Type{Vec{N,R}}, v::Vec{N}) where {N,R} =
    Vec{N,R}(NTuple{N, R}(v))

@inline Tuple(v::Vec{N}) where {N} = ntuple(i -> v.elts[i].value, Val(N))
@inline NTuple{N, T}(v::Vec{N}) where{N, T} = ntuple(i -> convert(T, v.elts[i].value), Val(N))

@generated function Base.:%(v::Vec{N,T}, ::Type{Vec{N,R}}) where {N,R,T}
    quote
        $(Expr(:meta, :inline))
        Vec{N,R}(tuple($([:(v.elts[$i].value % R) for i in 1:N]...)))
    end
end

# Convert vectors to tuples
@generated function Base.convert(::Type{NTuple{N,R}}, v::Vec{N,T}) where {N,R,T}
    quote
        $(Expr(:meta, :inline))
        tuple($([:(R(v.elts[$i].value)) for i in 1:N]...))
    end
end
@inline Base.convert(::Type{Tuple}, v::Vec{N,T}) where {N,T} =
    Base.convert(NTuple{N,T}, v)

# Promotion rules

# Note: Type promotion only works for subtypes of Number
# Base.promote_rule{N,T<:ScalarTypes}(::Type{Vec{N,T}}, ::Type{T}) = Vec{N,T}

Base.zero(::Type{Vec{N,T}}) where {N,T} = Vec{N,T}(zero(T))
Base.one(::Type{Vec{N,T}}) where {N,T} = Vec{N,T}(one(T))
Base.zero(::Vec{N,T}) where {N,T} = zero(Vec{N,T})
Base.one(::Vec{N,T}) where {N,T} = one(Vec{N,T})

# Floating point formats

int_type(::Type{Float16}) = Int16
int_type(::Type{Float32}) = Int32
int_type(::Type{Float64}) = Int64
# int_type(::Type{Float128}) = Int128
# int_type(::Type{Float256}) = Int256

uint_type(::Type{Float16}) = UInt16
uint_type(::Type{Float32}) = UInt32
uint_type(::Type{Float64}) = UInt64
# uint_type(::Type{Float128}) = UInt128
# uint_type(::Type{Float256}) = UInt256

significand_bits(::Type{Float16}) = 10
significand_bits(::Type{Float32}) = 23
significand_bits(::Type{Float64}) = 52
# significand_bits(::Type{Float128}) = 112
# significand_bits(::Type{Float256}) = 136

exponent_bits(::Type{T}) where {T<:FloatingTypes} =
    8*sizeof(T) - 1 - significand_bits(T)
sign_bits(::Type{T}) where {T<:FloatingTypes} = 1

significand_mask(::Type{T}) where {T<:FloatingTypes} =
    uint_type(T)(uint_type(T)(1) << significand_bits(T) - 1)
exponent_mask(::Type{T}) where {T<:FloatingTypes} =
    uint_type(T)(uint_type(T)(1) << exponent_bits(T) - 1) << significand_bits(T)
sign_mask(::Type{T}) where {T<:FloatingTypes} =
    uint_type(T)(1) << (significand_bits(T) + exponent_bits(T))

for T in (Float16, Float32, Float64)
    @assert sizeof(int_type(T)) == sizeof(T)
    @assert sizeof(uint_type(T)) == sizeof(T)
    @assert significand_bits(T) + exponent_bits(T) + sign_bits(T) == 8*sizeof(T)
    @assert significand_mask(T) | exponent_mask(T) | sign_mask(T) ==
        typemax(uint_type(T))
    @assert significand_mask(T) ⊻ exponent_mask(T) ⊻ sign_mask(T) ==
        typemax(uint_type(T))
end

# Convert Julia types to LLVM types

llvmtype(::Type{Bool}) = "i8"   # Julia represents Tuple{Bool} as [1 x i8]

# llvmtype(::Type{Bool8}) = "i8"
# llvmtype(::Type{Bool16}) = "i16"
# llvmtype(::Type{Bool32}) = "i32"
# llvmtype(::Type{Bool64}) = "i64"
# llvmtype(::Type{Bool128}) = "i128"

llvmtype(::Type{Int8}) = "i8"
llvmtype(::Type{Int16}) = "i16"
llvmtype(::Type{Int32}) = "i32"
llvmtype(::Type{Int64}) = "i64"
llvmtype(::Type{Int128}) = "i128"
llvmtype(::Type{<:Ptr}) = llvmtype(Int)

llvmtype(::Type{UInt8}) = "i8"
llvmtype(::Type{UInt16}) = "i16"
llvmtype(::Type{UInt32}) = "i32"
llvmtype(::Type{UInt64}) = "i64"
llvmtype(::Type{UInt128}) = "i128"

llvmtype(::Type{Float16}) = "half"
llvmtype(::Type{Float32}) = "float"
llvmtype(::Type{Float64}) = "double"

# Type-dependent optimization flags
# fastflags{T<:IntTypes}(::Type{T}) = "nsw"
# fastflags{T<:UIntTypes}(::Type{T}) = "nuw"
# fastflags{T<:FloatingTypes}(::Type{T}) = "fast"

suffix(N::Integer, ::Type{T}) where {T<:IntegerTypes} = "v$(N)i$(8*sizeof(T))"
suffix(N::Integer, ::Type{T}) where {T<:FloatingTypes} = "v$(N)f$(8*sizeof(T))"

# Type-dependent LLVM constants
function llvmconst(::Type{T}, val) where T
    T(val) === T(0) && return "zeroinitializer"
    typ = llvmtype(T)
    "$typ $val"
end
function llvmconst(::Type{Bool}, val)
    Bool(val) === false && return "zeroinitializer"
    typ = "i1"
    "$typ $(Int(val))"
end
function llvmconst(N::Integer, ::Type{T}, val) where T
    T(val) === T(0) && return "zeroinitializer"
    typ = llvmtype(T)
    "<" * join(["$typ $val" for i in 1:N], ", ") * ">"
end
function llvmconst(N::Integer, ::Type{Bool}, val)
    Bool(val) === false && return "zeroinitializer"
    typ = "i1"
    "<" * join(["$typ $(Int(val))" for i in 1:N], ", ") * ">"
end
function llvmtypedconst(::Type{T}, val) where T
    typ = llvmtype(T)
    T(val) === T(0) && return "$typ zeroinitializer"
    "$typ $val"
end
function llvmtypedconst(::Type{Bool}, val)
    typ = "i1"
    Bool(val) === false && return "$typ zeroinitializer"
    "$typ $(Int(val))"
end

# Type-dependent LLVM intrinsics
llvmins(::Val{:+}, N, ::Type{T}) where {T <: IndexTypes} = "add"
llvmins(::Val{:-}, N, ::Type{T}) where {T <: IndexTypes} = "sub"
llvmins(::Val{:*}, N, ::Type{T}) where {T <: IntegerTypes} = "mul"
llvmins(::Val{:div}, N, ::Type{T}) where {T <: IntTypes} = "sdiv"
llvmins(::Val{:rem}, N, ::Type{T}) where {T <: IntTypes} = "srem"
llvmins(::Val{:div}, N, ::Type{T}) where {T <: UIntTypes} = "udiv"
llvmins(::Val{:rem}, N, ::Type{T}) where {T <: UIntTypes} = "urem"

llvmins(::Val{:~}, N, ::Type{T}) where {T <: IntegerTypes} = "xor"
llvmins(::Val{:&}, N, ::Type{T}) where {T <: IntegerTypes} = "and"
llvmins(::Val{:|}, N, ::Type{T}) where {T <: IntegerTypes} = "or"
llvmins(::Val{:⊻}, N, ::Type{T}) where {T <: IntegerTypes} = "xor"

llvmins(::Val{:<<}, N, ::Type{T}) where {T <: IntegerTypes} = "shl"
llvmins(::Val{:>>>}, N, ::Type{T}) where {T <: IntegerTypes} = "lshr"
llvmins(::Val{:>>}, N, ::Type{T}) where {T <: UIntTypes} = "lshr"
llvmins(::Val{:>>}, N, ::Type{T}) where {T <: IntTypes} = "ashr"

llvmins(::Val{:(==)}, N, ::Type{T}) where {T <: IntegerTypes} = "icmp eq"
llvmins(::Val{:(!=)}, N, ::Type{T}) where {T <: IntegerTypes} = "icmp ne"
llvmins(::Val{:(>)}, N, ::Type{T}) where {T <: IntTypes} = "icmp sgt"
llvmins(::Val{:(>=)}, N, ::Type{T}) where {T <: IntTypes} = "icmp sge"
llvmins(::Val{:(<)}, N, ::Type{T}) where {T <: IntTypes} = "icmp slt"
llvmins(::Val{:(<=)}, N, ::Type{T}) where {T <: IntTypes} = "icmp sle"
llvmins(::Val{:(>)}, N, ::Type{T}) where {T <: UIntTypes} = "icmp ugt"
llvmins(::Val{:(>=)}, N, ::Type{T}) where {T <: UIntTypes} = "icmp uge"
llvmins(::Val{:(<)}, N, ::Type{T}) where {T <: UIntTypes} = "icmp ult"
llvmins(::Val{:(<=)}, N, ::Type{T}) where {T <: UIntTypes} = "icmp ule"

llvmins(::Val{:vifelse}, N, ::Type{T}) where {T} = "select"

llvmins(::Val{:+}, N, ::Type{T}) where {T <: FloatingTypes} = "fadd"
llvmins(::Val{:-}, N, ::Type{T}) where {T <: FloatingTypes} = "fsub"
llvmins(::Val{:*}, N, ::Type{T}) where {T <: FloatingTypes} = "fmul"
llvmins(::Val{:/}, N, ::Type{T}) where {T <: FloatingTypes} = "fdiv"
llvmins(::Val{:inv}, N, ::Type{T}) where {T <: FloatingTypes} = "fdiv"
llvmins(::Val{:rem}, N, ::Type{T}) where {T <: FloatingTypes} = "frem"

llvmins(::Val{:(==)}, N, ::Type{T}) where {T <: FloatingTypes} = "fcmp oeq"
llvmins(::Val{:(!=)}, N, ::Type{T}) where {T <: FloatingTypes} = "fcmp une"
llvmins(::Val{:(>)}, N, ::Type{T}) where {T <: FloatingTypes} = "fcmp ogt"
llvmins(::Val{:(>=)}, N, ::Type{T}) where {T <: FloatingTypes} = "fcmp oge"
llvmins(::Val{:(<)}, N, ::Type{T}) where {T <: FloatingTypes} = "fcmp olt"
llvmins(::Val{:(<=)}, N, ::Type{T}) where {T <: FloatingTypes} = "fcmp ole"

llvmins(::Val{:^}, N, ::Type{T}) where {T <: FloatingTypes} =
    "@llvm.pow.$(suffix(N,T))"
llvmins(::Val{:abs}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.fabs.$(suffix(N,T))"
llvmins(::Val{:ceil}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.ceil.$(suffix(N,T))"
llvmins(::Val{:copysign}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.copysign.$(suffix(N,T))"
llvmins(::Val{:cos}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.cos.$(suffix(N,T))"
llvmins(::Val{:exp}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.exp.$(suffix(N,T))"
llvmins(::Val{:exp2}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.exp2.$(suffix(N,T))"
llvmins(::Val{:floor}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.floor.$(suffix(N,T))"
llvmins(::Val{:fma}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.fma.$(suffix(N,T))"
llvmins(::Val{:log}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.log.$(suffix(N,T))"
llvmins(::Val{:log10}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.log10.$(suffix(N,T))"
llvmins(::Val{:log2}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.log2.$(suffix(N,T))"
llvmins(::Val{:max}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.maxnum.$(suffix(N,T))"
llvmins(::Val{:min}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.minnum.$(suffix(N,T))"
# llvmins(::Val{:max}, N, ::Type{T}) where {T<:FloatingTypes} =
#     "@llvm.maximum.$(suffix(N,T))"
# llvmins(::Val{:min}, N, ::Type{T}) where {T<:FloatingTypes} =
#     "@llvm.minimum.$(suffix(N,T))"
llvmins(::Val{:muladd}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.fmuladd.$(suffix(N,T))"
llvmins(::Val{:powi}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.powi.$(suffix(N,T))"
llvmins(::Val{:round}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.rint.$(suffix(N,T))"
llvmins(::Val{:sin}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.sin.$(suffix(N,T))"
llvmins(::Val{:sqrt}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.sqrt.$(suffix(N,T))"
llvmins(::Val{:trunc}, N, ::Type{T}) where {T<:FloatingTypes} =
    "@llvm.trunc.$(suffix(N,T))"

# Convert between LLVM scalars, vectors, and arrays

function scalar2vector(vec, siz, typ, sca)
    instrs = []
    accum(nam, i) = i<0 ? "undef" : i==siz-1 ? nam : "$(nam)_iter$i"
    for i in 0:siz-1
        push!(instrs,
            "$(accum(vec,i)) = " *
                "insertelement <$siz x $typ> $(accum(vec,i-1)), " *
                "$typ $sca, i32 $i")
    end
    instrs
end

function array2vector(vec, siz, typ, arr, tmp="$(arr)_av")
    instrs = []
    accum(nam, i) = i<0 ? "undef" : i==siz-1 ? nam : "$(nam)_iter$i"
    for i in 0:siz-1
        push!(instrs, "$(tmp)_elem$i = extractvalue [$siz x $typ] $arr, $i")
        push!(instrs,
            "$(accum(vec,i)) = " *
                "insertelement <$siz x $typ> $(accum(vec,i-1)), " *
                "$typ $(tmp)_elem$i, i32 $i")
    end
    instrs
end

function vector2array(arr, siz, typ, vec, tmp="$(vec)_va")
    instrs = []
    accum(nam, i) = i<0 ? "undef" : i==siz-1 ? nam : "$(nam)_iter$i"
    for i in 0:siz-1
        push!(instrs,
            "$(tmp)_elem$i = extractelement <$siz x $typ> $vec, i32 $i")
        push!(instrs,
            "$(accum(arr,i)) = "*
                "insertvalue [$siz x $typ] $(accum(arr,i-1)), " *
                "$typ $(tmp)_elem$i, $i")
    end
    instrs
end

# TODO: change argument order
function subvector(vec, siz, typ, rvec, rsiz, roff, tmp="$(rvec)_sv")
    instrs = []
    accum(nam, i) = i<0 ? "undef" : i==rsiz-1 ? nam : "$(nam)_iter$i"
    @assert 0 <= roff
    @assert roff + rsiz <= siz
    for i in 0:rsiz-1
        push!(instrs,
            "$(tmp)_elem$i = extractelement <$siz x $typ> $vec, i32 $(roff+i)")
        push!(instrs,
            "$(accum(rvec,i)) = " *
                "insertelement <$rsiz x $typ> $(accum(rvec,i-1)), " *
                "$typ $(tmp)_elem$i, i32 $i")
    end
    instrs
end

function extendvector(vec, siz, typ, voff, vsiz, val, rvec, tmp="$(rvec)_ev")
    instrs = []
    accum(nam, i) = i<0 ? "undef" : i==siz+vsiz-1 ? nam : "$(nam)_iter$i"
    rsiz = siz + vsiz
    for i in 0:siz-1
        push!(instrs,
            "$(tmp)_elem$i = extractelement <$siz x $typ> $vec, i32 $i")
        push!(instrs,
            "$(accum(rvec,i)) = " *
                "insertelement <$rsiz x $typ> $(accum(rvec,i-1)), " *
                "$typ $(tmp)_elem$i, i32 $i")
    end
    for i in siz:siz+vsiz-1
        push!(instrs,
            "$(accum(rvec,i)) = " *
                "insertelement <$rsiz x $typ> $(accum(rvec,i-1)), $val, i32 $i")
    end
    instrs
end

# Element-wise access

import Base.setindex
export setindex
@generated function setindex(v::Vec{N,T}, x::Number, ::Val{I}) where {N,T,I}
    @assert isa(I, Integer)
    1 <= I <= N || throw(BoundsError())
    typ = llvmtype(T)
    ityp = llvmtype(Int)
    vtyp = "<$N x $typ>"
    decls = []
    instrs = []
    push!(instrs, "%res = insertelement $vtyp %0, $typ %1, $ityp $(I-1)")
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{T}}, Tuple{NTuple{N,VE{T}}, T}, v.elts, T(x)))
    end
end
@inline function setindex(v::Vec{N,T}, x::Number, ::Type{Val{I}}) where {N,T,I}
    setindex(v, x, Val(I))
end

@generated function setindex(v::Vec{N,T}, x::Number, i::Int) where {N,T}
    typ = llvmtype(T)
    ityp = llvmtype(Int)
    vtyp = "<$N x $typ>"
    decls = []
    instrs = []
    push!(instrs, "%res = insertelement $vtyp %0, $typ %2, $ityp %1")
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        @boundscheck 1 <= i <= N || throw(BoundsError())
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{T}}, Tuple{NTuple{N,VE{T}}, Int, T},
            v.elts, i-1, T(x)))
    end
end
setindex(v::Vec{N,T}, x::Number, i) where {N,T} = setindex(v, x, Int(i))

Base.@propagate_inbounds Base.getindex(v::Vec{N,T}, ::Val{I}) where {N,T,I} = v.elts[I].value
Base.@propagate_inbounds Base.getindex(v::Vec{N,T}, ::Type{Val{I}}) where {N,T,I} = Base.getindex(v, Val(I))
Base.@propagate_inbounds Base.getindex(v::Vec{N,T}, i) where {N,T} = v.elts[i].value

# Type conversion

@generated function Base.reinterpret(::Type{Vec{N,R}},
        v1::Vec{N1,T1}) where {N,R,N1,T1}
    @assert N*sizeof(R) == N1*sizeof(T1)
    typ1 = llvmtype(T1)
    vtyp1 = "<$N1 x $typ1>"
    typr = llvmtype(R)
    vtypr = "<$N x $typr>"
    decls = []
    instrs = []
    push!(instrs, "%res = bitcast $vtyp1 %0 to $vtypr")
    push!(instrs, "ret $vtypr %res")
    quote
        $(Expr(:meta, :inline))
        Vec{N,R}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{R}}, Tuple{NTuple{N1,VE{T1}}}, v1.elts))
    end
end

# Generic function wrappers

# Functions taking one argument
@generated function llvmwrap(::Val{Op}, v1::Vec{N,T1},
                             ::Type{R} = T1) where {Op,N,T1,R}
    @assert isa(Op, Symbol)
    typ1 = llvmtype(T1)
    vtyp1 = "<$N x $typ1>"
    typr = llvmtype(R)
    vtypr = "<$N x $typr>"
    ins = llvmins(Val(Op), N, T1)
    decls = []
    instrs = []
    if ins[1] == '@'
        push!(decls, "declare $vtypr $ins($vtyp1)")
        push!(instrs, "%res = call $vtypr $ins($vtyp1 %0)")
    else
        if Op === :~
            @assert T1 <: IntegerTypes
            otherval = -1
        elseif Op === :inv
            @assert T1 <: FloatingTypes
            otherval = 1.0
        else
            otherval = 0
        end
        otherarg = llvmconst(N, T1, otherval)
        push!(instrs, "%res = $ins $vtyp1 $otherarg, %0")
    end
    push!(instrs, "ret $vtypr %res")
    quote
        $(Expr(:meta, :inline))
        Vec{N,R}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{R}}, Tuple{NTuple{N,VE{T1}}}, v1.elts))
    end
end

# Functions taking one Bool argument
@generated function llvmwrap(::Val{Op}, v1::Vec{N,Bool},
                             ::Type{Bool} = Bool) where {Op,N}
    @assert isa(Op, Symbol)
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    ins = llvmins(Val(Op), N, Bool)
    decls = []
    instrs = []
    push!(instrs, "%arg1 = trunc $vbtyp %0 to <$N x i1>")
    otherarg = llvmconst(N, Bool, true)
    push!(instrs, "%res = $ins <$N x i1> $otherarg, %arg1")
    push!(instrs, "%resb = zext <$N x i1> %res to $vbtyp")
    push!(instrs, "ret $vbtyp %resb")
    quote
        $(Expr(:meta, :inline))
        Vec{N,Bool}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{Bool}}, Tuple{NTuple{N,VE{Bool}}}, v1.elts))
    end
end

# Functions taking two arguments
@generated function llvmwrap(::Val{Op}, v1::Vec{N,T1}, v2::Vec{N,T2},
        ::Type{R} = T1) where {Op,N,T1,T2,R}
    @assert isa(Op, Symbol)
    typ1 = llvmtype(T1)
    vtyp1 = "<$N x $typ1>"
    typ2 = llvmtype(T2)
    vtyp2 = "<$N x $typ2>"
    typr = llvmtype(R)
    vtypr = "<$N x $typr>"
    ins = llvmins(Val(Op), N, T1)
    decls = []
    instrs = []
    if ins[1] == '@'
        push!(decls, "declare $vtypr $ins($vtyp1, $vtyp2)")
        push!(instrs, "%res = call $vtypr $ins($vtyp1 %0, $vtyp2 %1)")
    else
        push!(instrs, "%res = $ins $vtyp1 %0, %1")
    end
    push!(instrs, "ret $vtypr %res")
    quote
        $(Expr(:meta, :inline))
        Vec{N,R}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{R}}, Tuple{NTuple{N,VE{T1}}, NTuple{N,VE{T2}}},
            v1.elts, v2.elts))
    end
end

# Functions taking two arguments, second argument is a scalar
@generated function llvmwrap(::Val{Op}, v1::Vec{N,T1}, s2::ScalarTypes,
        ::Type{R} = T1) where {Op,N,T1,R}
    @assert isa(Op, Symbol)
    typ1 = llvmtype(T1)
    vtyp1 = "<$N x $typ1>"
    typ2 = llvmtype(s2)
    typr = llvmtype(R)
    vtypr = "<$N x $typr>"
    ins = llvmins(Val(Op), N, T1)
    decls = []
    instrs = []
    if ins[1] == '@'
        push!(decls, "declare $vtypr $ins($vtyp1, $typ2)")
        push!(instrs, "%res = call $vtypr $ins($vtyp1 %0, $typ2 %1)")
    else
        push!(instrs, "%res = $ins $vtyp1 %0, %1")
    end
    push!(instrs, "ret $vtypr %res")
    quote
        $(Expr(:meta, :inline))
        Vec{N,R}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{R}}, Tuple{NTuple{N,VE{T1}}, $s2},
            v1.elts, s2))
    end
end

# Functions taking two arguments, returning Bool
@generated function llvmwrap(::Val{Op}, v1::Vec{N,T1}, v2::Vec{N,T2},
        ::Type{Bool}) where {Op,N,T1,T2}
    @assert isa(Op, Symbol)
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    abtyp = "[$N x $btyp]"
    typ1 = llvmtype(T1)
    vtyp1 = "<$N x $typ1>"
    atyp1 = "[$N x $typ1]"
    typ2 = llvmtype(T2)
    vtyp2 = "<$N x $typ2>"
    atyp2 = "[$N x $typ2]"
    ins = llvmins(Val(Op), N, T1)
    decls = []
    instrs = []
    if false && N == 1
        append!(instrs, array2vector("%arg1", N, typ1, "%0", "%arg1arr"))
        append!(instrs, array2vector("%arg2", N, typ2, "%1", "%arg2arr"))
        push!(instrs, "%cond = $ins $vtyp1 %arg1, %arg2")
        push!(instrs, "%res = zext <$N x i1> %cond to $vbtyp")
        append!(instrs, vector2array("%resarr", N, btyp, "%res"))
        push!(instrs, "ret $abtyp %resarr")
    else
        push!(instrs, "%res = $ins $vtyp1 %0, %1")
        push!(instrs, "%resb = zext <$N x i1> %res to $vbtyp")
        push!(instrs, "ret $vbtyp %resb")
    end
    quote
        $(Expr(:meta, :inline))
        Vec{N,Bool}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{Bool}}, Tuple{NTuple{N,VE{T1}}, NTuple{N,VE{T2}}},
            v1.elts, v2.elts))
    end
end

# Functions taking a vector and a scalar argument
# @generated function llvmwrap{Op,N,T1,T2,R}(::Val{Op}, v1::Vec{N,T1},
#         x2::T2, ::Type{R} = T1)
#     @assert isa(Op, Symbol)
#     typ1 = llvmtype(T1)
#     atyp1 = "[$N x $typ1]"
#     vtyp1 = "<$N x $typ1>"
#     typ2 = llvmtype(T2)
#     typr = llvmtype(R)
#     atypr = "[$N x $typr]"
#     vtypr = "<$N x $typr>"
#     ins = llvmins(Val(Op), N, T1)
#     decls = []
#     instrs = []
#     append!(instrs, array2vector("%arg1", N, typ1, "%0", "%arg1arr"))
#     if ins[1] == '@'
#         push!(decls, "declare $vtypr $ins($vtyp1, $typ2)")
#         push!(instrs, "%res = call $vtypr $ins($vtyp1 %arg1, $typ2 %1)")
#     else
#         push!(instrs, "%res = $ins $vtyp1 %arg1, %1")
#     end
#     append!(instrs, vector2array("%resarr", N, typr, "%res"))
#     push!(instrs, "ret $atypr %resarr")
#     quote
#         $(Expr(:meta, :inline))
#         Vec{N,R}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
#             NTuple{N,R}, Tuple{NTuple{N,T1}, T2}, v1.elts, x2))
#     end
# end

# Functions taking two Bool arguments, returning Bool
@generated function llvmwrap(::Val{Op}, v1::Vec{N,Bool}, v2::Vec{N,Bool},
        ::Type{Bool} = Bool) where {Op,N}
    @assert isa(Op, Symbol)
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    ins = llvmins(Val(Op), N, Bool)
    decls = []
    instrs = []
    push!(instrs, "%arg1 = trunc $vbtyp %0 to <$N x i1>")
    push!(instrs, "%arg2 = trunc $vbtyp %1 to <$N x i1>")
    push!(instrs, "%res = $ins <$N x i1> %arg1, %arg2")
    push!(instrs, "%resb = zext <$N x i1> %res to $vbtyp")
    push!(instrs, "ret $vbtyp %resb")
    quote
        $(Expr(:meta, :inline))
        Vec{N,Bool}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{Bool}}, Tuple{NTuple{N,VE{Bool}}, NTuple{N,VE{Bool}}},
            v1.elts, v2.elts))
    end
end

# Functions taking three arguments
@generated function llvmwrap(::Val{Op}, v1::Vec{N,T1}, v2::Vec{N,T2},
        v3::Vec{N,T3}, ::Type{R} = T1) where {Op,N,T1,T2,T3,R}
    @assert isa(Op, Symbol)
    typ1 = llvmtype(T1)
    vtyp1 = "<$N x $typ1>"
    typ2 = llvmtype(T2)
    vtyp2 = "<$N x $typ2>"
    typ3 = llvmtype(T3)
    vtyp3 = "<$N x $typ3>"
    typr = llvmtype(R)
    vtypr = "<$N x $typr>"
    ins = llvmins(Val(Op), N, T1)
    decls = []
    instrs = []
    if ins[1] == '@'
        push!(decls, "declare $vtypr $ins($vtyp1, $vtyp2, $vtyp3)")
        push!(instrs,
            "%res = call $vtypr $ins($vtyp1 %0, $vtyp2 %1, $vtyp3 %2)")
    else
        push!(instrs, "%res = $ins $vtyp1 %0, %1, %2")
    end
    push!(instrs, "ret $vtypr %res")
    quote
        $(Expr(:meta, :inline))
        Vec{N,R}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{R}},
            Tuple{NTuple{N,VE{T1}}, NTuple{N,VE{T2}}, NTuple{N,VE{T3}}},
            v1.elts, v2.elts, v3.elts))
    end
end

@generated function llvmwrapshift(::Val{Op}, v1::Vec{N,T},
        ::Val{I}) where {Op,N,T,I}
    @assert isa(Op, Symbol)
    if I >= 0
        op = Op
        i = I
    else
        if Op === :>> || Op === :>>>
            op = :<<
        else
            @assert Op === :<<
            if T <: Unsigned
                op = :>>>
            else
                op = :>>
            end
        end
        i = -I
    end
    @assert op in (:<<, :>>, :>>>)
    @assert i >= 0
    typ = llvmtype(T)
    vtyp = "<$N x $typ>"
    ins = llvmins(Val(op), N, T)
    decls = []
    instrs = []
    nbits = 8*sizeof(T)
    if (op === :>> && T <: IntTypes) || i < nbits
        count = llvmconst(N, T, min(nbits-1, i))
        push!(instrs, "%res = $ins $vtyp %0, $count")
        push!(instrs, "ret $vtyp %res")
    else
        zero = llvmconst(N, T, 0)
        push!(instrs, "return $vtyp $zero")
    end
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{T}}, Tuple{NTuple{N,VE{T}}}, v1.elts))
    end
end

@generated function llvmwrapshift(::Val{Op}, v1::Vec{N,T},
        x2::Unsigned) where {Op,N,T}
    @assert isa(Op, Symbol)
    typ = llvmtype(T)
    vtyp = "<$N x $typ>"
    ins = llvmins(Val(Op), N, T)
    decls = []
    instrs = []
    append!(instrs, scalar2vector("%count", N, typ, "%1"))
    nbits = 8*sizeof(T)
    push!(instrs, "%tmp = $ins $vtyp %0, %count")
    push!(instrs, "%inbounds = icmp ult $typ %1, $nbits")
    if Op === :>> && T <: IntTypes
        nbits1 = llvmconst(N, T, 8*sizeof(T)-1)
        push!(instrs, "%limit = $ins $vtyp %0, $nbits1")
        push!(instrs, "%res = select i1 %inbounds, $vtyp %tmp, $vtyp %limit")
    else
        zero = llvmconst(N, T, 0)
        push!(instrs, "%res = select i1 %inbounds, $vtyp %tmp, $vtyp $zero")
    end
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        # Note that this function might be called with out-of-bounds
        # values for x2, assuming that the results are then ignored
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{T}}, Tuple{NTuple{N,VE{T}}, T}, v1.elts, x2 % T))
    end
end

@generated function llvmwrapshift(::Val{Op}, v1::Vec{N,T},
        x2::Integer) where {Op,N,T}
    if Op === :>> || Op === :>>>
        NegOp = :<<
    else
        @assert Op === :<<
        if T <: Unsigned
            NegOp = :>>>
        else
            NegOp = :>>
        end
    end
    ValOp = Val(Op)
    ValNegOp = Val(NegOp)
    quote
        $(Expr(:meta, :inline))
        ifelse(x2 >= 0,
               llvmwrapshift($ValOp, v1, unsigned(x2)),
               llvmwrapshift($ValNegOp, v1, unsigned(-x2)))
    end
end

@generated function llvmwrapshift(::Val{Op}, v1::Vec{N,T},
        v2::Vec{N,U}) where {Op,N,T,U<:UIntTypes}
    @assert isa(Op, Symbol)
    typ = llvmtype(T)
    vtyp = "<$N x $typ>"
    ins = llvmins(Val(Op), N, T)
    decls = []
    instrs = []
    push!(instrs, "%tmp = $ins $vtyp %0, %1")
    nbits = llvmconst(N, T, 8*sizeof(T))
    push!(instrs, "%inbounds = icmp ult $vtyp %1, $nbits")
    if Op === :>> && T <: IntTypes
        nbits1 = llvmconst(N, T, 8*sizeof(T)-1)
        push!(instrs, "%limit = $ins $vtyp %0, $nbits1")
        push!(instrs,
            "%res = select <$N x i1> %inbounds, $vtyp %tmp, $vtyp %limit")
    else
        zero = llvmconst(N, T, 0)
        push!(instrs,
            "%res = select <$N x i1> %inbounds, $vtyp %tmp, $vtyp $zero")
    end
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{T}}, Tuple{NTuple{N,VE{T}}, NTuple{N,VE{T}}},
            v1.elts, (v2 % Vec{N,T}).elts))
    end
end

@generated function llvmwrapshift(::Val{Op}, v1::Vec{N,T},
        v2::Vec{N,U}) where {Op,N,T,U<:IntegerTypes}
    if Op === :>> || Op === :>>>
        NegOp = :<<
    else
        @assert Op === :<<
        if T <: Unsigned
            NegOp = :>>>
        else
            NegOp = :>>
        end
    end
    ValOp = Val(Op)
    ValNegOp = Val(NegOp)
    quote
        $(Expr(:meta, :inline))
        vifelse(v2 >= 0,
                llvmwrapshift($ValOp, v1, v2 % Vec{N,unsigned(U)}),
                llvmwrapshift($ValNegOp, v1, -v2 % Vec{N,unsigned(U)}))
    end
end

# Conditionals

for op in (:(==), :(!=), :(<), :(<=), :(>), :(>=))
    @eval begin
        @inline Base.$op(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T} =
            llvmwrap(Val($(QuoteNode(op))), v1, v2, Bool)
    end
end
@inline function Base.cmp(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T}
    I = int_type(T)
    vifelse(isequal(v1, v2), Vec{N,I}(0),
            vifelse(isless(v1, v2), Vec{N,I}(-1), Vec{N,I}(1)))
end
@inline function Base.isfinite(v1::Vec{N,T}) where {N,T<:FloatingTypes}
    U = uint_type(T)
    em = Vec{N,U}(exponent_mask(T))
    iv = reinterpret(Vec{N,U}, v1)
    iv & em != em
end
@inline Base.isinf(v1::Vec{N,T}) where {N,T<:FloatingTypes} = abs(v1) == Vec{N,T}(Inf)
@inline Base.isnan(v1::Vec{N,T}) where {N,T<:FloatingTypes} = v1 != v1
@inline function Base.issubnormal(v1::Vec{N,T}) where {N,T<:FloatingTypes}
    U = uint_type(T)
    em = Vec{N,U}(exponent_mask(T))
    sm = Vec{N,U}(significand_mask(T))
    iv = reinterpret(Vec{N,U}, v1)
    (iv & em == Vec{N,U}(0)) & (iv & sm != Vec{N,U}(0))
end
@inline function Base.signbit(v1::Vec{N,T}) where {N,T<:FloatingTypes}
    U = uint_type(T)
    sm = Vec{N,U}(sign_mask(T))
    iv = reinterpret(Vec{N,U}, v1)
    iv & sm != Vec{N,U}(0)
end

export vifelse
vifelse(c::Bool, x, y) = ifelse(c, x, y)
@generated function vifelse(v1::Vec{N,Bool}, v2::Vec{N,T},
        v3::Vec{N,T}) where {N,T}
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    abtyp = "[$N x $btyp]"
    typ = llvmtype(T)
    vtyp = "<$N x $typ>"
    atyp = "[$N x $typ]"
    decls = []
    instrs = []
    if false && N == 1
        append!(instrs, array2vector("%arg1", N, btyp, "%0", "%arg1arr"))
        append!(instrs, array2vector("%arg2", N, typ, "%1", "%arg2arr"))
        append!(instrs, array2vector("%arg3", N, typ, "%2", "%arg3arr"))
        push!(instrs, "%cond = trunc $vbtyp %arg1 to <$N x i1>")
        push!(instrs, "%res = select <$N x i1> %cond, $vtyp %arg2, $vtyp %arg3")
        append!(instrs, vector2array("%resarr", N, typ, "%res"))
        push!(instrs, "ret $atyp %resarr")
    else
        push!(instrs, "%cond = trunc $vbtyp %0 to <$N x i1>")
        push!(instrs, "%res = select <$N x i1> %cond, $vtyp %1, $vtyp %2")
        push!(instrs, "ret $vtyp %res")
    end
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{T}},
            Tuple{NTuple{N,VE{Bool}}, NTuple{N,VE{T}}, NTuple{N,VE{T}}},
            v1.elts, v2.elts, v3.elts))
    end
end

# Integer arithmetic functions

for op in (:~, :+, :-)
    @eval begin
        @inline Base.$op(v1::Vec{N,T}) where {N,T<:IntegerTypes} =
            llvmwrap(Val($(QuoteNode(op))), v1)
    end
end
@inline Base.:!(v1::Vec{N,Bool}) where {N} = ~v1
@inline function Base.abs(v1::Vec{N,T}) where {N,T<:IntTypes}
    # s = -Vec{N,T}(signbit(v1))
    s = v1 >> Val(8*sizeof(T))
    # Note: -v1 == ~v1 + 1
    (s ⊻ v1) - s
end
@inline Base.abs(v1::Vec{N,T}) where {N,T<:UIntTypes} = v1
# TODO: Try T(v1>0) - T(v1<0)
#       use a shift for v1<0
#       evaluate v1>0 as -v1<0 ?
@inline Base.sign(v1::Vec{N,T}) where {N,T<:IntTypes} =
    vifelse(v1 == Vec{N,T}(0), Vec{N,T}(0),
        vifelse(v1 < Vec{N,T}(0), Vec{N,T}(-1), Vec{N,T}(1)))
@inline Base.sign(v1::Vec{N,T}) where {N,T<:UIntTypes} =
    vifelse(v1 == Vec{N,T}(0), Vec{N,T}(0), Vec{N,T}(1))
@inline Base.signbit(v1::Vec{N,T}) where {N,T<:IntTypes} = v1 < Vec{N,T}(0)
@inline Base.signbit(v1::Vec{N,T}) where {N,T<:UIntTypes} = Vec{N,Bool}(false)

for op in (:&, :|, :⊻, :+, :-, :*, :div, :rem)
    @eval begin
        @inline Base.$op(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:IntegerTypes} =
            llvmwrap(Val($(QuoteNode(op))), v1, v2)
    end
end
@inline Base.copysign(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:IntTypes} =
    vifelse(signbit(v2), -abs(v1), abs(v1))
@inline Base.copysign(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:UIntTypes} = v1
@inline Base.flipsign(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:IntTypes} =
    vifelse(signbit(v2), -v1, v1)
@inline Base.flipsign(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:UIntTypes} = v1
@inline Base.max(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:IntegerTypes} =
    vifelse(v1>=v2, v1, v2)
@inline Base.min(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:IntegerTypes} =
    vifelse(v1>=v2, v2, v1)

@inline function Base.muladd(v1::Vec{N,T}, v2::Vec{N,T},
        v3::Vec{N,T}) where {N,T<:IntegerTypes}
    v1*v2+v3
end

# TODO: Handle negative shift counts
#       use vifelse
#       ensure vifelse is efficient
for op in (:<<, :>>, :>>>)
    @eval begin
        @inline Base.$op(v1::Vec{N,T}, ::Val{I}) where {N,T<:IntegerTypes,I} =
            llvmwrapshift(Val($(QuoteNode(op))), v1, Val(I))
        @inline Base.$op(v1::Vec{N,T}, ::Type{Val{I}}) where {N,T<:IntegerTypes,I} =
            Base.$op(v1, Val(I))
        @inline Base.$op(v1::Vec{N,T}, x2::Unsigned) where {N,T<:IntegerTypes} =
            llvmwrapshift(Val($(QuoteNode(op))), v1, x2)
        @inline Base.$op(v1::Vec{N,T}, x2::Int) where {N,T<:IntegerTypes} =
            llvmwrapshift(Val($(QuoteNode(op))), v1, x2)
        @inline Base.$op(v1::Vec{N,T}, x2::Integer) where {N,T<:IntegerTypes} =
            llvmwrapshift(Val($(QuoteNode(op))), v1, x2)
        @inline Base.$op(v1::Vec{N,T},
                         v2::Vec{N,U}) where {N,T<:IntegerTypes,U<:UIntTypes} =
            llvmwrapshift(Val($(QuoteNode(op))), v1, v2)
        @inline Base.$op(v1::Vec{N,T},
                         v2::Vec{N,U}) where {N,T<:IntegerTypes,U<:IntegerTypes} =
            llvmwrapshift(Val($(QuoteNode(op))), v1, v2)
        @inline Base.$op(x1::T, v2::Vec{N,T}) where {N,T<:IntegerTypes} =
            $op(Vec{N,T}(x1), v2)
    end
end

# Floating point arithmetic functions

for op in (
        :+, :-,
        :abs, :ceil, :cos, :exp, :exp2, :floor, :inv, :log, :log10, :log2,
        :round, :sin, :sqrt, :trunc)
    @eval begin
        @inline Base.$op(v1::Vec{N,T}) where {N,T<:FloatingTypes} =
            llvmwrap(Val($(QuoteNode(op))), v1)
    end
end
@inline Base.exp10(v1::Vec{N,T}) where {N,T<:FloatingTypes} = Vec{N,T}(10)^v1
@inline Base.sign(v1::Vec{N,T}) where {N,T<:FloatingTypes} =
    vifelse(v1 == Vec{N,T}(0.0), Vec{N,T}(0.0), copysign(Vec{N,T}(1.0), v1))

for op in (:+, :-, :*, :/, :^, :copysign, :max, :min, :rem)
    @eval begin
        @inline Base.$op(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:FloatingTypes} =
            llvmwrap(Val($(QuoteNode(op))), v1, v2)
    end
end
# Using `IntegerTypes` here so that this definition "wins" against
# `^(::ScalarTypes, v2::Vec)`.
@inline Base.:^(v1::Vec{N,T}, x2::IntegerTypes) where {N,T<:FloatingTypes} =
    llvmwrap(Val(:powi), v1, Int(x2))
@inline Base.:^(v1::Vec{N,T}, x2::Integer) where {N,T<:FloatingTypes} =
    llvmwrap(Val(:powi), v1, Int(x2))
@inline Base.flipsign(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:FloatingTypes} =
    vifelse(signbit(v2), -v1, v1)

# Do what Base does for HWNumber:
@inline Base.literal_pow(::typeof(^), x::Vec, ::Val{0}) = one(typeof(x))
@inline Base.literal_pow(::typeof(^), x::Vec, ::Val{1}) = x
@inline Base.literal_pow(::typeof(^), x::Vec, ::Val{2}) = x*x
@inline Base.literal_pow(::typeof(^), x::Vec, ::Val{3}) = x*x*x

for op in (:fma, :muladd)
    @eval begin
        @inline function Base.$op(v1::Vec{N,T},
                v2::Vec{N,T}, v3::Vec{N,T}) where {N,T<:FloatingTypes}
            llvmwrap(Val($(QuoteNode(op))), v1, v2, v3)
        end
    end
end

# Type promotion

# Promote scalars of all IntegerTypes to vectors of IntegerTypes, leaving the
# vector type unchanged

for op in (
        :(==), :(!=), :(<), :(<=), :(>), :(>=),
        :&, :|, :⊻, :+, :-, :*, :copysign, :div, :flipsign, :max, :min, :rem)
    @eval begin
        @inline Base.$op(s1::Bool, v2::Vec{N,Bool}) where {N} =
            $op(Vec{N,Bool}(s1), v2)
        @inline Base.$op(s1::IntegerTypes, v2::Vec{N,T}) where {N,T<:IntegerTypes} =
            $op(Vec{N,T}(s1), v2)
        @inline Base.$op(v1::Vec{N,T}, s2::IntegerTypes) where {N,T<:IntegerTypes} =
            $op(v1, Vec{N,T}(s2))
    end
end
@inline vifelse(c::Vec{N,Bool}, s1::IntegerTypes,
        v2::Vec{N,T}) where {N,T<:IntegerTypes} =
    vifelse(c, Vec{N,T}(s1), v2)
@inline vifelse(c::Vec{N,Bool}, v1::Vec{N,T},
        s2::IntegerTypes) where {N,T<:IntegerTypes} =
    vifelse(c, v1, Vec{N,T}(s2))

for op in (:muladd,)
    @eval begin
        @inline Base.$op(s1::IntegerTypes, v2::Vec{N,T},
                v3::Vec{N,T}) where {N,T<:IntegerTypes} =
            $op(Vec{N,T}(s1), v2, v3)
        @inline Base.$op(v1::Vec{N,T}, s2::IntegerTypes,
                v3::Vec{N,T}) where {N,T<:IntegerTypes} =
            $op(v1, Vec{N,T}(s2), v3)
        @inline Base.$op(s1::IntegerTypes, s2::IntegerTypes,
                v3::Vec{N,T}) where {N,T<:IntegerTypes} =
            $op(Vec{N,T}(s1), Vec{N,T}(s2), v3)
        @inline Base.$op(v1::Vec{N,T}, v2::Vec{N,T},
                s3::IntegerTypes) where {N,T<:IntegerTypes} =
            $op(v1, v2, Vec{N,T}(s3))
        @inline Base.$op(s1::IntegerTypes, v2::Vec{N,T},
                s3::IntegerTypes) where {N,T<:IntegerTypes} =
            $op(Vec{N,T}(s1), v2, Vec{N,T}(s3))
        @inline Base.$op(v1::Vec{N,T}, s2::IntegerTypes,
                s3::IntegerTypes) where {N,T<:IntegerTypes} =
            $op(v1, Vec{N,T}(s2), Vec{N,T}(s3))
    end
end

# Promote scalars of all ScalarTypes to vectors of FloatingTypes, leaving the
# vector type unchanged

for op in (
        :(==), :(!=), :(<), :(<=), :(>), :(>=),
        :+, :-, :*, :/, :^, :copysign, :flipsign, :max, :min, :rem)
    @eval begin
        @inline Base.$op(s1::ScalarTypes, v2::Vec{N,T}) where {N,T<:FloatingTypes} =
            $op(Vec{N,T}(s1), v2)
        @inline Base.$op(v1::Vec{N,T}, s2::ScalarTypes) where {N,T<:FloatingTypes} =
            $op(v1, Vec{N,T}(s2))
    end
end
@inline vifelse(c::Vec{N,Bool}, s1::ScalarTypes,
        v2::Vec{N,T}) where {N,T<:FloatingTypes} =
    vifelse(c, Vec{N,T}(s1), v2)
@inline vifelse(c::Vec{N,Bool}, v1::Vec{N,T},
        s2::ScalarTypes) where {N,T<:FloatingTypes} =
    vifelse(c, v1, Vec{N,T}(s2))

for op in (:fma, :muladd)
    @eval begin
        @inline Base.$op(s1::ScalarTypes, v2::Vec{N,T},
                v3::Vec{N,T}) where {N,T<:FloatingTypes} =
            $op(Vec{N,T}(s1), v2, v3)
        @inline Base.$op(v1::Vec{N,T}, s2::ScalarTypes,
                v3::Vec{N,T}) where {N,T<:FloatingTypes} =
            $op(v1, Vec{N,T}(s2), v3)
        @inline Base.$op(s1::ScalarTypes, s2::ScalarTypes,
                v3::Vec{N,T}) where {N,T<:FloatingTypes} =
            $op(Vec{N,T}(s1), Vec{N,T}(s2), v3)
        @inline Base.$op(v1::Vec{N,T}, v2::Vec{N,T},
                s3::ScalarTypes) where {N,T<:FloatingTypes} =
            $op(v1, v2, Vec{N,T}(s3))
        @inline Base.$op(s1::ScalarTypes, v2::Vec{N,T},
                s3::ScalarTypes) where {N,T<:FloatingTypes} =
            $op(Vec{N,T}(s1), v2, Vec{N,T}(s3))
        @inline Base.$op(v1::Vec{N,T}, s2::ScalarTypes,
                s3::ScalarTypes) where {N,T<:FloatingTypes} =
            $op(v1, Vec{N,T}(s2), Vec{N,T}(s3))
    end
end

# Poitner arithmetics between Ptr, IntegerTypes, and vectors of them.

for op in (:+, :-)
    @eval begin
        @inline Base.$op(v1::Vec{N,<:Ptr}, v2::Vec{N,<:IntegerTypes}) where {N} =
            llvmwrap(Val($(QuoteNode(op))), v1, v2)
        @inline Base.$op(v1::Vec{N,<:IntegerTypes}, v2::Vec{N,<:Ptr}) where {N} =
            llvmwrap(Val($(QuoteNode(op))), v1, v2)
        @inline Base.$op(s1::P, v2::Vec{N,<:IntegerTypes}) where {N,P<:Ptr} =
            $op(Vec{N,P}(s1), v2)
        @inline Base.$op(v1::Vec{N,<:IntegerTypes}, s2::P) where {N,P<:Ptr} =
            $op(v1, Vec{N,P}(s2))
    end
end

# Reduction operations

# TODO: map, mapreduce

function getneutral(op::Symbol, ::Type{T}) where T
    zs = Dict{Symbol,T}()
    if T <: IntegerTypes
        zs[:&] = ~T(0)
        zs[:|] = T(0)
    end
    zs[:max] = typemin(T)
    zs[:min] = typemax(T)
    zs[:+] = T(0)
    zs[:*] = T(1)
    zs[op]
end

if VERSION >= v"0.7.0-beta2.195"
    nextpow2(n) = nextpow(2, n)
end

# We cannot pass in the neutral element via Val{}; if we try, Julia refuses to
# inline this function, which is then disastrous for performance
@generated function llvmwrapreduce(::Val{Op}, v::Vec{N,T}) where {Op,N,T}
    @assert isa(Op, Symbol)
    z = getneutral(Op, T)
    typ = llvmtype(T)
    decls = []
    instrs = []
    n = N
    nam = "%0"
    nold,n = n,nextpow2(n)
    if n > nold
        namold,nam = nam,"%vec_$n"
        append!(instrs,
            extendvector(namold, nold, typ, n, n-nold,
                llvmtypedconst(T,z), nam))
    end
    while n > 1
        nold,n = n, div(n, 2)
        namold,nam = nam,"%vec_$n"
        vtyp = "<$n x $typ>"
        ins = llvmins(Val(Op), n, T)
        append!(instrs, subvector(namold, nold, typ, "$(nam)_1", n, 0))
        append!(instrs, subvector(namold, nold, typ, "$(nam)_2", n, n))
        if ins[1] == '@'
            push!(decls, "declare $vtyp $ins($vtyp, $vtyp)")
            push!(instrs,
                "$nam = call $vtyp $ins($vtyp $(nam)_1, $vtyp $(nam)_2)")
        else
            push!(instrs, "$nam = $ins $vtyp $(nam)_1, $(nam)_2")
        end
    end
    push!(instrs, "%res = extractelement <$n x $typ> $nam, i32 0")
    push!(instrs, "ret $typ %res")
    quote
        $(Expr(:meta, :inline))
        Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            T, Tuple{NTuple{N,VE{T}}}, v.elts)
    end
end

@inline Base.all(v::Vec{N,T}) where {N,T<:IntegerTypes} = llvmwrapreduce(Val(:&), v)
@inline Base.any(v::Vec{N,T}) where {N,T<:IntegerTypes} = llvmwrapreduce(Val(:|), v)
@inline Base.maximum(v::Vec{N,T}) where {N,T<:FloatingTypes} =
    llvmwrapreduce(Val(:max), v)
@inline Base.minimum(v::Vec{N,T}) where {N,T<:FloatingTypes} =
    llvmwrapreduce(Val(:min), v)
@inline Base.prod(v::Vec{N,T}) where {N,T} = llvmwrapreduce(Val(:*), v)
@inline Base.sum(v::Vec{N,T}) where {N,T} = llvmwrapreduce(Val(:+), v)

@generated function Base.reduce(::Val{Op}, v::Vec{N,T}) where {Op,N,T}
    @assert isa(Op, Symbol)
    z = getneutral(Op, T)
    stmts = []
    n = N
    push!(stmts, :($(Symbol(:v,n)) = v))
    nold,n = n,nextpow2(n)
    if n > nold
        push!(stmts,
            :($(Symbol(:v,n)) = Vec{$n,T}($(Expr(:tuple,
                [:($(Symbol(:v,nold)).elts[$i]) for i in 1:nold]...,
                [z for i in nold+1:n]...)))))
    end
    while n > 1
        nold,n = n, div(n, 2)
        push!(stmts,
            :($(Symbol(:v,n,"lo")) = Vec{$n,T}($(Expr(:tuple,
                [:($(Symbol(:v,nold)).elts[$i]) for i in 1:n]...,)))))
        push!(stmts,
            :($(Symbol(:v,n,"hi")) = Vec{$n,T}($(Expr(:tuple,
                [:($(Symbol(:v,nold)).elts[$i]) for i in n+1:nold]...)))))
        push!(stmts,
            :($(Symbol(:v,n)) =
                $Op($(Symbol(:v,n,"lo")), $(Symbol(:v,n,"hi")))))
    end
    push!(stmts, :(v1[1]))
    Expr(:block, Expr(:meta, :inline), stmts...)
end
@inline function Base.reduce(::Type{Val{Op}}, v::Vec{N,T}) where {Op,N,T}
    Base.reduce(Val(Op), v)
end

@inline Base.maximum(v::Vec{N,T}) where {N,T<:IntegerTypes} = reduce(Val(:max), v)
@inline Base.minimum(v::Vec{N,T}) where {N,T<:IntegerTypes} = reduce(Val(:min), v)

# Load and store functions

export valloc
function valloc(::Type{T}, N::Int, sz::Int) where T
    @assert N > 0
    @assert sz >= 0
    # We use padding to align the address of the first element, and
    # also to ensure that we can access past the last element up to
    # the next full vector width
    padding = N-1 + mod(-sz, N)
    mem = Vector{T}(undef, sz + padding)
    addr = Int(pointer(mem))
    off = mod(-addr, N * sizeof(T))
    @assert mod(off, sizeof(T)) == 0
    off = fld(off, sizeof(T))
    @assert 0 <= off <= padding
    res = view(mem, off+1 : off+sz)
    addr2 = Int(pointer(res))
    @assert mod(addr2, N * sizeof(T)) == 0
    res
end
function valloc(f, ::Type{T}, N::Int, sz::Int) where T
    mem = valloc(T, N, sz)
    @inbounds for i in 1:sz
        mem[i] = f(i)
    end
    mem
end

export vload, vloada, vloadnt
@generated function vload(::Type{Vec{N,T}}, ptr::Ptr{T},
                          ::Val{Aligned} = Val(false),
                          ::Val{Nontemporal} = Val(false)) where {N,T,Aligned,Nontemporal}
    @assert isa(Aligned, Bool)
    ptyp = llvmtype(Int)
    typ = llvmtype(T)
    vtyp = "<$N x $typ>"
    decls = []
    instrs = []
    if Aligned
        align = N * sizeof(T)
    else
        align = sizeof(T)   # This is overly optimistic
    end
    flags = [""]
    if align > 0
        push!(flags, "align $align")
    end
    if Nontemporal
        push!(flags, "!nontemporal !{i32 1}")
    end
    if VERSION < v"v0.7.0-DEV"
        push!(instrs, "%ptr = bitcast $typ* %0 to $vtyp*")
    else
        push!(instrs, "%ptr = inttoptr $ptyp %0 to $vtyp*")
    end
    push!(instrs, "%res = load $vtyp, $vtyp* %ptr" * join(flags, ", "))
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{T}}, Tuple{Ptr{T}}, ptr))
    end
end
@inline function vload(::Type{Vec{N,T}}, ptr::Ptr{T},
                       ::Type{Val{Aligned}},
                       ::Type{Val{Nontemporal}} = Val{false}) where {N,T,Aligned,Nontemporal}
    vload(Vec{N, T}, ptr, Val(Aligned), Val(Nontemporal))
end

@inline vloada(::Type{Vec{N,T}}, ptr::Ptr{T}) where {N,T} =
    vload(Vec{N,T}, ptr, Val(true))

@inline vloadnt(::Type{Vec{N,T}}, ptr::Ptr{T}) where {N,T} =
    vload(Vec{N,T}, ptr, Val(true), Val(true))

@inline function vload(::Type{Vec{N,T}},
                       arr::FastContiguousArray{T,1},
                       i::Integer,
                       ::Val{Aligned} = Val(false),
                       ::Val{Nontemporal} = Val(false)) where {N,T,Aligned,Nontemporal}
    #TODO @boundscheck 1 <= i <= length(arr) - (N-1) || throw(BoundsError())
    vload(Vec{N,T}, pointer(arr, i), Val(Aligned), Val(Nontemporal))
end
@inline function vload(::Type{Vec{N,T}},
                       arr::FastContiguousArray{T,1},
                       i::Integer,
                       ::Type{Val{Aligned}},k = Val{false},
                       ::Type{Val{Nontemporal}} = Val{false}) where {N,T,Aligned,Nontemporal}
    vload(Vec{N,T}, arr, i, Val(Aligned), Val(Nontemporal))
end
@inline function vloada(::Type{Vec{N,T}},
                        arr::FastContiguousArray{T,1},
                        i::Integer) where {N,T}
    vload(Vec{N,T}, arr, i, Val(true))
end
@inline function vloadnt(::Type{Vec{N,T}},
                        arr::Union{Array{T,1},SubArray{T,1}},
                        i::Integer) where {N,T}
    vload(Vec{N,T}, arr, i, Val(true), Val(true))
end

@inline vload(::Type{Vec{N,T}}, ptr::Ptr{T}, mask::Nothing,
              ::Val{Aligned} = Val(false)) where {N,T,Aligned} =
    vload(Vec{N,T}, ptr, Val(Aligned))
@inline vload(::Type{Vec{N,T}}, ptr::Ptr{T}, mask::Nothing,
              ::Type{Val{Aligned}}) where {N,T,Aligned} =
    vload(Vec{N,T}, ptr, make, Val(Aligned))

@generated function vload(::Type{Vec{N,T}}, ptr::Ptr{T},
                          mask::Vec{N,Bool},
                          ::Val{Aligned} = Val(false)) where {N,T,Aligned}
    @assert isa(Aligned, Bool)
    ptyp = llvmtype(Int)
    typ = llvmtype(T)
    vtyp = "<$N x $typ>"
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    decls = []
    instrs = []
    if Aligned
        align = N * sizeof(T)
    else
        align = sizeof(T)   # This is overly optimistic
    end

    if VERSION < v"v0.7.0-DEV"
        push!(instrs, "%ptr = bitcast $typ* %0 to $vtyp*")
    else
        push!(instrs, "%ptr = inttoptr $ptyp %0 to $vtyp*")
    end
    push!(instrs, "%mask = trunc $vbtyp %1 to <$N x i1>")
    push!(decls,
        "declare $vtyp @llvm.masked.load.$(suffix(N,T))($vtyp*, i32, " *
            "<$N x i1>, $vtyp)")
    push!(instrs,
        "%res = call $vtyp @llvm.masked.load.$(suffix(N,T))($vtyp* %ptr, " *
            "i32 $align, <$N x i1> %mask, $vtyp $(llvmconst(N, T, 0)))")
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{T}}, Tuple{Ptr{T}, NTuple{N,VE{Bool}}}, ptr, mask.elts))
    end
end
@inline function vload(::Type{Vec{N,T}}, ptr::Ptr{T},
                       mask::Vec{N,Bool},
                       ::Type{Val{Aligned}}) where {N,T,Aligned}
    vload(Vec{N,T}, ptr, mask, Val(Aligned))
end

@inline vloada(::Type{Vec{N,T}}, ptr::Ptr{T},
               mask::Union{Vec{N,Bool}, Nothing}) where {N,T} =
    vload(Vec{N,T}, ptr, mask, Val(true))

@inline function vload(::Type{Vec{N,T}},
                       arr::FastContiguousArray{T,1},
                       i::Integer, mask::Union{Vec{N,Bool}, Nothing},
                       ::Val{Aligned} = Val(false)) where {N,T,Aligned}
    #TODO @boundscheck 1 <= i <= length(arr) - (N-1) || throw(BoundsError())
    vload(Vec{N,T}, pointer(arr, i), mask, Val(Aligned))
end
@inline function vload(::Type{Vec{N,T}},
                       arr::FastContiguousArray{T,1},
                       i::Integer, mask::Union{Vec{N,Bool}, Nothing},
                       ::Type{Val{Aligned}}) where {N,T,Aligned}
    vload(Vec{N,T}, arr, i, mask, Val(Aligned))
end
@inline function vloada(::Type{Vec{N,T}},
                        arr::FastContiguousArray{T,1}, i::Integer,
                        mask::Union{Vec{N,Bool}, Nothing}) where {N,T}
    vload(Vec{N,T}, arr, i, mask, Val(true))
end

export vstore, vstorea, vstorent
@generated function vstore(v::Vec{N,T}, ptr::Ptr{T},
                           ::Val{Aligned} = Val(false),
                           ::Val{Nontemporal} = Val(false)) where {N,T,Aligned,Nontemporal}
    @assert isa(Aligned, Bool)
    @assert isa(Nontemporal, Bool)
    ptyp = llvmtype(Int)
    typ = llvmtype(T)
    vtyp = "<$N x $typ>"
    decls = []
    instrs = []
    if Aligned
        align = N * sizeof(T)
    else
        align = sizeof(T)   # This is overly optimistic
    end
    flags = [""]
    if align > 0
        push!(flags, "align $align")
    end
    if Nontemporal
        push!(flags, "!nontemporal !{i32 1}")
    end
    if VERSION < v"v0.7.0-DEV"
        push!(instrs, "%ptr = bitcast $typ* %1 to $vtyp*")
    else
        push!(instrs, "%ptr = inttoptr $ptyp %1 to $vtyp*")
    end
    push!(instrs, "store $vtyp %0, $vtyp* %ptr" * join(flags, ", "))
    push!(instrs, "ret void")
    quote
        $(Expr(:meta, :inline))
        Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
                      Cvoid, Tuple{NTuple{N,VE{T}}, Ptr{T}}, v.elts, ptr)
    end
end
@inline function vstore(v::Vec{N,T}, ptr::Ptr{T},
                        ::Type{Val{Aligned}},
                        ::Type{Val{Nontemporal}} = Val{false}) where {N,T,Aligned,Nontemporal}
    vstore(v, ptr, Val(Aligned), Val(Nontemporal))
end

@inline vstorea(v::Vec{N,T}, ptr::Ptr{T}) where {N,T} = vstore(v, ptr, Val{true})

@inline vstorent(v::Vec{N,T}, ptr::Ptr{T}) where {N,T} = vstore(v, ptr, Val{true}, Val{true})

@inline function vstore(v::Vec{N,T},
                        arr::FastContiguousArray{T,1},
                        i::Integer,
                        ::Val{Aligned} = Val(false),
                        ::Val{Nontemporal} = Val(false)) where {N,T,Aligned,Nontemporal}
    @boundscheck 1 <= i <= length(arr) - (N-1) || throw(BoundsError())
    vstore(v, pointer(arr, i), Val{Aligned}, Val{Nontemporal})
end
@inline function vstore(v::Vec{N,T},
                        arr::FastContiguousArray{T,1},
                        i::Integer,
                        ::Type{Val{Aligned}},
                        ::Type{Val{Nontemporal}} = Val{false}) where {N,T,Aligned,Nontemporal}
    vstore(v, arr, i, Val(Aligned), Val(Nontemporal))
end
@inline function vstorea(v::Vec{N,T}, arr::FastContiguousArray{T,1},
                         i::Integer) where {N,T}
    vstore(v, arr, i, Val{true})
end
@inline function vstorent(v::Vec{N,T}, arr::FastContiguousArray{T,1},
                         i::Integer) where {N,T}
    vstore(v, arr, i, Val{true}, Val{true})
end

@inline vstore(v::Vec{N,T}, ptr::Ptr{T}, mask::Nothing,
               ::Val{Aligned} = Val(false)) where {N,T,Aligned} =
    vstore(v, ptr, Val{Aligned})
@inline vstore(v::Vec{N,T}, ptr::Ptr{T}, mask::Nothing,
               ::Type{Val{Aligned}}) where {N,T,Aligned} =
    vstore(v, ptr, mask, Val(Aligned))

@generated function vstore(v::Vec{N,T}, ptr::Ptr{T},
                           mask::Vec{N,Bool},
                           ::Val{Aligned} = Val(false)) where {N,T,Aligned}
    @assert isa(Aligned, Bool)
    ptyp = llvmtype(Int)
    typ = llvmtype(T)
    vtyp = "<$N x $typ>"
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    decls = []
    instrs = []
    if Aligned
        align = N * sizeof(T)
    else
        align = sizeof(T)   # This is overly optimistic
    end
    if VERSION < v"v0.7.0-DEV"
        push!(instrs, "%ptr = bitcast $typ* %1 to $vtyp*")
    else
        push!(instrs, "%ptr = inttoptr $ptyp %1 to $vtyp*")
    end
    push!(instrs, "%mask = trunc $vbtyp %2 to <$N x i1>")
    push!(decls,
        "declare void @llvm.masked.store.$(suffix(N,T))($vtyp, $vtyp*, i32, " *
            "<$N x i1>)")
    push!(instrs,
        "call void @llvm.masked.store.$(suffix(N,T))($vtyp %0, $vtyp* %ptr, " *
            "i32 $align, <$N x i1> %mask)")
    push!(instrs, "ret void")
    quote
        $(Expr(:meta, :inline))
        Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            Cvoid, Tuple{NTuple{N,VE{T}}, Ptr{T}, NTuple{N,VE{Bool}}},
            v.elts, ptr, mask.elts)
    end
end
@inline function vstore(v::Vec{N,T}, ptr::Ptr{T},
                        mask::Vec{N,Bool},
                        ::Type{Val{Aligned}}) where {N,T,Aligned}
    vstore(v, ptr, mask, Val(Aligned))
end

@inline vstorea(v::Vec{N,T}, ptr::Ptr{T},
                mask::Union{Vec{N,Bool}, Nothing}) where {N,T} =
    vstore(v, ptr, mask, Val{true})

@inline function vstore(v::Vec{N,T},
                        arr::FastContiguousArray{T,1},
                        i::Integer,
                        mask::Union{Vec{N,Bool}, Nothing},
                        ::Val{Aligned} = Val(false),
                        ::Val{Nontemporal} = Val(false)) where {N,T,Aligned,Nontemporal}
    #TODO @boundscheck 1 <= i <= length(arr) - (N-1) || throw(BoundsError())
    vstore(v, pointer(arr, i), mask, Val{Aligned}, Val{Nontemporal})
end
@inline function vstore(v::Vec{N,T},
                        arr::FastContiguousArray{T,1},
                        i::Integer,
                        mask::Union{Vec{N,Bool}, Nothing},
                        ::Type{Val{Aligned}},
                        ::Type{Val{Nontemporal}} = Val{false}) where {N,T,Aligned,Nontemporal}
    vstore(v, arr, i, mask, Val(Aligned), Val(Nontemporal))
end
@inline function vstorea(v::Vec{N,T},
                         arr::FastContiguousArray{T,1},
                         i::Integer,
                         mask::Union{Vec{N,Bool}, Nothing}) where {N,T}
    vstore(v, arr, i, mask, Val{true})
end

export vgather, vgathera

@inline vgather(
        ::Type{Vec{N,T}}, ptrs::Vec{N,Ptr{T}}, mask::Nothing,
        ::Val{Aligned} = Val(false)) where {N,T,Aligned} =
    vgather(Vec{N,T}, ptrs, Vec(ntuple(_ -> true, N)), Val(Aligned))
@inline vgather(
        ::Type{Vec{N,T}}, ptrs::Vec{N,Ptr{T}}, mask::Nothing,
        ::Type{Val{Aligned}}) where {N,T,Aligned} =
    vgather(Vec{N,T}, ptrs, mask, Val(Aligned))

@generated function vgather(
        ::Type{Vec{N,T}}, ptrs::Vec{N,Ptr{T}}, mask::Vec{N,Bool},
        ::Val{Aligned} = Val(false)) where {N,T,Aligned}
    @assert isa(Aligned, Bool)
    ptyp = llvmtype(Int)
    typ = llvmtype(T)
    vptyp = "<$N x $typ*>"
    vtyp = "<$N x $typ>"
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    decls = []
    instrs = []
    if Aligned
        align = N * sizeof(T)
    else
        align = sizeof(T)   # This is overly optimistic
    end

    if VERSION < v"v0.7.0-DEV"
        push!(instrs, "%ptrs = bitcast <$N x $typ*> %0 to $vptyp")
    else
        push!(instrs, "%ptrs = inttoptr <$N x $ptyp> %0 to $vptyp")
    end
    push!(instrs, "%mask = trunc $vbtyp %1 to <$N x i1>")
    push!(decls,
        "declare $vtyp @llvm.masked.gather.$(suffix(N,T))($vptyp, i32, " *
            "<$N x i1>, $vtyp)")
    push!(instrs,
        "%res = call $vtyp @llvm.masked.gather.$(suffix(N,T))($vptyp %ptrs, " *
            "i32 $align, <$N x i1> %mask, $vtyp $(llvmconst(N, T, 0)))")
    push!(instrs, "ret $vtyp %res")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,VE{T}}, Tuple{NTuple{N,VE{Ptr{T}}}, NTuple{N,VE{Bool}}},
            ptrs.elts, mask.elts))
    end
end
@inline function vgather(
        ::Type{Vec{N,T}}, ptrs::Vec{N,Ptr{T}}, mask::Vec{N,Bool},
        ::Type{Val{Aligned}}) where {N,T,Aligned}
    vgather(Vec{N,T}, ptrs, mask, Val(Aligned))
end

@inline vgathera(::Type{Vec{N,T}}, ptrs::Vec{N,Ptr{T}},
                 mask::Union{Vec{N,Bool}, Nothing}) where {N,T} =
    vgather(Vec{N,T}, ptrs, mask, Val{true})

@inline vgather(arr::FastContiguousArray{T,1},
                idx::Vec{N,<:Integer},
                mask::Union{Vec{N,Bool}, Nothing} = nothing,
                ::Val{Aligned} = Val(false)) where {N,T,Aligned} =
    vgather(Vec{N,T},
            pointer(arr) + sizeof(T) * (idx - 1),
            mask, Val{Aligned})
@inline vgather(arr::FastContiguousArray{T,1},
                idx::Vec{N,<:Integer},
                mask::Union{Vec{N,Bool}, Nothing},
                ::Type{Val{Aligned}}) where {N,T,Aligned} =
    vgather(arr, idx, mask, Val(Aligned))

@inline vgathera(arr::FastContiguousArray{T,1},
                 idx::Vec{N,<:Integer},
                 mask::Union{Vec{N,Bool}, Nothing} = nothing) where {N,T} =
    vgather(arr, idx, mask, Val{true})

export vscatter, vscattera

@inline vscatter(
        v::Vec{N,T}, ptrs::Vec{N,Ptr{T}}, mask::Nothing,
        ::Val{Aligned} = Val(false)) where {N,T,Aligned} =
    vscatter(v, ptrs, Vec(ntuple(_ -> true, N)), Val(Aligned))
@inline vscatter(
        v::Vec{N,T}, ptrs::Vec{N,Ptr{T}}, mask::Nothing,
        ::Type{Val{Aligned}}) where {N,T,Aligned} =
    vscatter(v, ptrs, mask, Val(Aligned))

@generated function vscatter(
        v::Vec{N,T}, ptrs::Vec{N,Ptr{T}}, mask::Vec{N,Bool},
        ::Val{Aligned} = Val(false)) where {N,T,Aligned}
    @assert isa(Aligned, Bool)
    ptyp = llvmtype(Int)
    typ = llvmtype(T)
    vptyp = "<$N x $typ*>"
    vtyp = "<$N x $typ>"
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    decls = []
    instrs = []
    if Aligned
        align = N * sizeof(T)
    else
        align = sizeof(T)   # This is overly optimistic
    end
    if VERSION < v"v0.7.0-DEV"
        push!(instrs, "%ptrs = bitcast <$N x $typ*> %1 to $vptyp")
    else
        push!(instrs, "%ptrs = inttoptr <$N x $ptyp> %1 to $vptyp")
    end
    push!(instrs, "%mask = trunc $vbtyp %2 to <$N x i1>")
    push!(decls,
        "declare void @llvm.masked.scatter.$(suffix(N,T))" *
            "($vtyp, $vptyp, i32, <$N x i1>)")
    push!(instrs,
        "call void @llvm.masked.scatter.$(suffix(N,T))" *
            "($vtyp %0, $vptyp %ptrs, i32 $align, <$N x i1> %mask)")
    push!(instrs, "ret void")
    quote
        $(Expr(:meta, :inline))
        Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            Cvoid,
            Tuple{NTuple{N,VE{T}}, NTuple{N,VE{Ptr{T}}}, NTuple{N,VE{Bool}}},
            v.elts, ptrs.elts, mask.elts)
    end
end
@inline function vscatter(
        v::Vec{N,T}, ptrs::Vec{N,Ptr{T}}, mask::Vec{N,Bool},
        ::Type{Val{Aligned}}) where {N,T,Aligned}
    vscatter(v, ptrs, mask, Val(Aligned))
end

@inline vscattera(v::Vec{N,T}, ptrs::Vec{N,Ptr{T}},
                  mask::Union{Vec{N,Bool}, Nothing}) where {N,T} =
    vscatter(v, ptrs, mask, Val{true})

@inline vscatter(v::Vec{N,T}, arr::FastContiguousArray{T,1},
                 idx::Vec{N,<:Integer},
                 mask::Union{Vec{N,Bool}, Nothing} = nothing,
                 ::Val{Aligned} = Val(false)) where {N,T,Aligned} =
    vscatter(v, pointer(arr) + sizeof(T) * (idx - 1), mask, Val(Aligned))
@inline vscatter(v::Vec{N,T}, arr::FastContiguousArray{T,1},
                 idx::Vec{N,<:Integer},
                 mask::Union{Vec{N,Bool}, Nothing},
                 ::Type{Val{Aligned}}) where {N,T,Aligned} =
    vscatter(v, arr, idx, mask, Val(Aligned))

@inline vscattera(v::Vec{N,T}, arr::FastContiguousArray{T,1},
                  idx::Vec{N,<:Integer},
                  mask::Union{Vec{N,Bool}, Nothing} = nothing) where {N,T} =
    vscatter(v, arr, idx, mask, Val{true})

# Vector shuffles

function shufflevector_instrs(N, T, I, two_operands)
    typ = llvmtype(T)
    vtyp2 = vtyp1 = "<$N x $typ>"
    M = length(I)
    vtyp3 = "<$M x i32>"
    vtypr = "<$M x $typ>"
    mask = "<" * join(map(x->string("i32 ", x), I), ", ") * ">"
    instrs = []
    v2 = two_operands ? "%1" : "undef"
    push!(instrs, "%res = shufflevector $vtyp1 %0, $vtyp2 $v2, $vtyp3 $mask")
    push!(instrs, "ret $vtypr %res")
    return M, [], instrs
end

export shufflevector
@generated function shufflevector(v1::Vec{N,T}, v2::Vec{N,T},
                                  ::Val{I}) where {N,T,I}
    M, decls, instrs = shufflevector_instrs(N, T, I, true)
    quote
        $(Expr(:meta, :inline))
        Vec{$M,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{$M,VE{T}},
            Tuple{NTuple{N,VE{T}}, NTuple{N,VE{T}}},
            v1.elts, v2.elts))
    end
end
@inline function shufflevector(v1::Vec{N,T}, v2::Vec{N,T},
                               ::Type{Val{I}}) where {N,T,I}
    shufflevector(v1, v2, Val(I))
end

@generated function shufflevector(v1::Vec{N,T}, ::Val{I}) where {N,T,I}
    M, decls, instrs = shufflevector_instrs(N, T, I, false)
    quote
        $(Expr(:meta, :inline))
        Vec{$M,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{$M,VE{T}},
            Tuple{NTuple{N,VE{T}}},
            v1.elts))
    end
end
@inline function shufflevector(v1::Vec{N,T}, ::Type{Val{I}}) where {N,T,I}
    shufflevector(v1, Val(I))
end

export VecRange

"""
    VecRange{N}(i::Int)

Analogous to `UnitRange` but for loading SIMD vector of width `N` at
index `i`.

# Examples
```jldoctest
julia> xs = ones(4);

julia> xs[VecRange{4}(1)]  # calls `vload(Vec{4,Float64}, xs, 1)`
<4 x Float64>[1.0, 1.0, 1.0, 1.0]
```
"""
struct VecRange{N}
    i::Int
end

@inline Base.length(idx::VecRange{N}) where {N} = N
@inline Base.first(idx::VecRange) = idx.i
@inline Base.last(idx::VecRange) = idx.i + length(idx) - 1

@inline Base.:+(idx::VecRange{N}, j::Integer) where N = VecRange{N}(idx.i + j)
@inline Base.:+(j::Integer, idx::VecRange{N}) where N = VecRange{N}(idx.i + j)
@inline Base.:-(idx::VecRange{N}, j::Integer) where N = VecRange{N}(idx.i - j)

Base.checkindex(::Type{Bool}, inds::AbstractUnitRange, idx::VecRange) =
    (first(inds) <= first(idx)) && (last(idx) <= last(inds))

Base.checkindex(::Type{Bool}, inds::AbstractUnitRange, idx::Vec) =
    all(first(inds) <= idx) && all(idx <= last(inds))

@inline _checkarity(::AbstractArray{<:Any,N}, ::Vararg{<:Any,N}) where {N} =
    nothing

@inline _checkarity(::T, ::Any) where {T <: AbstractArray} =
    if IndexStyle(T) isa IndexLinear
        nothing
    else
        throw(ArgumentError("""
        Array type $T does not support indexing with a single index.
        Exactly $(ndims(T)) (non-mask) indices have to be specified.
        """))
    end

_checkarity(::AbstractArray{<:Any,N}, ::Vararg{<:Any,M}) where {N,M} =
    throw(ArgumentError("""
    $M indices are given to $N-dimensional array.
    Exactly $N (non-mask) indices have to be specified when using SIMD.
    """))

# Combined with `_preprocessindices`, helper function `_extractmask`
# extracts `mask` in the tail position.  As slicing tuple is not
# type-stable, we use reverse-of-tail-of-reverse hack to extract
# `mask` at the end of `args`.
@inline _extractmask(mask::Vec{N,Bool}, R::Vararg{Integer}) where N =
    (reverse(R), mask)
@inline _extractmask(R::Vararg{Integer}) = (reverse(R), nothing)
@inline _extractmask(mask::Vec{N,Bool}) where {N} = ((), mask)
@inline _extractmask() = ((), nothing)

@noinline _extractmask(rargs...) =
    throw(ArgumentError("""
    Using SIMD indexing `array[idx, i2, ..., iN, mask]` for `N`-dimensional
    array requires `i2` to `iN` to be all integers and `mask` to be optionally
    a SIMD vector `Vec` of `Bool`s.  Given `(i2, ..., iN, mask)` is
    $(summary(reverse(rargs)))
    """))

_maskedidx(idx, ::Nothing, ::Any) = idx
_maskedidx(idx::Vec, mask::Vec, fst) = vifelse(mask, idx, fst)
_maskedidx(idx::VecRange, mask::Vec, fst) =
    _maskedidx(Vec(ntuple(i -> i - 1 + idx.i, length(mask))), mask, fst)

Base.@propagate_inbounds function _preprocessindices(arr, idx, args)
    I, mask = _extractmask(reverse(args)...)
    _checkarity(arr, idx, I...)
    @boundscheck checkbounds(arr,
                             _maskedidx(idx, mask, first(axes(arr, 1))),
                             I...)
    return I, mask
end

"""
    _pointer(arr, i, I)

Pointer to the element `arr[i, I...]`.
"""
Base.@propagate_inbounds _pointer(arr::Array, i, I) =
    pointer(arr, LinearIndices(arr)[i, I...])
Base.@propagate_inbounds _pointer(arr::Base.FastContiguousSubArray, i, I) =
    pointer(arr, (i, I...))
Base.@propagate_inbounds _pointer(arr::SubArray, i, I) =
    pointer(Base.unsafe_view(arr, 1, I...), i)

Base.@propagate_inbounds function Base.getindex(
        arr::ContiguousArray{T}, idx::VecRange{N},
        args::Vararg{Union{Integer,Vec{N,Bool}}}) where {N,T}
    I, mask = _preprocessindices(arr, idx, args)
    return vload(Vec{N,T}, _pointer(arr, idx.i, I), mask)
end

Base.@propagate_inbounds function Base.setindex!(
        arr::ContiguousArray{T}, v::Vec{N,T}, idx::VecRange{N},
        args::Vararg{Union{Integer,Vec{N,Bool}}}) where {N,T}
    I, mask = _preprocessindices(arr, idx, args)
    vstore(v, _pointer(arr, idx.i, I), mask)
    return arr
end

Base.@propagate_inbounds function Base.getindex(
        arr::ContiguousArray{T}, idx::Vec{N,<:Integer},
        args::Vararg{Union{Integer,Vec{N,Bool}}}) where {N,T}
    I, mask = _preprocessindices(arr, idx, args)
    ptrs = _pointer(arr, 1, I) - sizeof(T) + sizeof(T) * idx
    return vgather(Vec{N,T}, ptrs, mask)
end

Base.@propagate_inbounds function Base.setindex!(
        arr::ContiguousArray{T}, v::Vec{N,T}, idx::Vec{N,<:Integer},
        args::Vararg{Union{Integer,Vec{N,Bool}}}) where {N,T}
    I, mask = _preprocessindices(arr, idx, args)
    ptrs = _pointer(arr, 1, I) - sizeof(T) + sizeof(T) * idx
    vscatter(v, ptrs, mask)
    return arr
end

end


using CairoMakie
using Images, FileIO
using ImageTransformations
using WindMixing
using JLD2
using WindMixing: animate_profiles_fluxes_final

# convective_adjustment = parse(Bool, ARGS[1])
convective_adjustment = false
# params_type = ARGS[2]
params_type = "18simBFGST0.8nograd"
# NN_type = ARGS[3]
NN_type = "leakyrelu"
# num = parse(Int, ARGS[1])
# NN_version = ARGS[2]

num=1
NN_version = "old"

if convective_adjustment
    CA_str = "CA"
else
    CA_str = "noCA"
end

u_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "u.png"))))
v_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "v.png"))))
T_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "T.png"))))
uw_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "uw.png"))))
vw_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "vw.png"))))
wT_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "wT.png"))))
Ri_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "Ri.png"))))
z_img = load(assetpath(joinpath(pwd(), "images_plot", "z.png")))

axis_images = (
    u = u_img,
    v = v_img,
    T = T_img,
    z = z_img,
    uw = uw_img,
    vw = vw_img,
    wT = wT_img,
    Ri = Ri_img,
)

EXTRACTED_FILE_DIR = joinpath(pwd(), "extracted_training_output")
if NN_version == "old"
    EXTRACTED_FILE_PATH = joinpath(EXTRACTED_FILE_DIR, "NDE_18sim_windcooling_windheating_$(params_type)_divide1f5_gradient_smallNN_$(NN_type)_rate_2e-4_T0.8_extracted.jld2")
    FILE_DIR = joinpath(pwd(), "final_results", "18sim_old")
else
    EXTRACTED_FILE_PATH = joinpath(EXTRACTED_FILE_DIR, "NDE_18sim_windcooling_windheating_$(params_type)_divide1f5_gradient_smallNN_$(NN_type)_rate_2e-4_T0.8_1e-4_extracted.jld2")
    FILE_DIR = joinpath(pwd(), "final_data_2", "NDE_18sim_windcooling_windheating_$(params_type)_divide1f5_gradient_smallNN_$(NN_type)_rate_2e-4_T0.8_1e-4_$(CA_str)")
end

if !ispath(FILE_DIR)
    mkdir(FILE_DIR)
end

file = jldopen(EXTRACTED_FILE_PATH)
train_files = file["training_info/train_files"]
loss_scalings = file["training_info/loss_scalings"]
train_parameters = file["training_info/parameters"]
close(file)

test_files_all = [
    # "wind_-5e-4_cooling_3e-8_new",
    # "wind_-3.5e-4_cooling_3e-8_new",
    # "wind_-2e-4_cooling_3e-8_new",

    # "wind_-5e-4_heating_-3e-8_new",
    # "wind_-3.5e-4_heating_-3e-8_new",
    # "wind_-2e-4_heating_-3e-8_new",

    # "wind_-5e-4_cooling_2e-8_new",
    # "wind_-5e-4_cooling_1e-8_new",
    # "wind_-3.5e-4_cooling_2e-8_new",
    # "wind_-3.5e-4_cooling_1e-8_new",
    # "wind_-2e-4_cooling_2e-8_new",
    # "wind_-2e-4_cooling_1e-8_new",
    # "wind_-5e-4_heating_-2e-8_new",
    # "wind_-5e-4_heating_-1e-8_new",
    # "wind_-3.5e-4_heating_-2e-8_new",
    # "wind_-3.5e-4_heating_-1e-8_new",
    # "wind_-2e-4_heating_-2e-8_new",
    # "wind_-2e-4_heating_-1e-8_new",

    # "wind_-4.5e-4_cooling_2.5e-8",
    # "wind_-2.5e-4_cooling_1.5e-8",
    # "wind_-4.5e-4_cooling_1.5e-8",
    # "wind_-2.5e-4_cooling_2.5e-8",

    # "wind_-4.5e-4_heating_-2.5e-8",
    # "wind_-2.5e-4_heating_-1.5e-8",
    # "wind_-4.5e-4_heating_-1.5e-8",
    # "wind_-2.5e-4_heating_-2.5e-8",

    # "wind_-5.5e-4_diurnal_5.5e-8",
    # "wind_-1.5e-4_diurnal_5.5e-8",

    # "wind_-5.5e-4_new",

    # "wind_-5.5e-4_heating_-3.5e-8",
    # "wind_-1.5e-4_heating_-3.5e-8",
    # "wind_-5.5e-4_cooling_3.5e-8",
    "wind_-1.5e-4_cooling_3.5e-8",
]

test_files = [test_files_all[num]]

ν₀ = train_parameters["ν₀"]
ν₋ = train_parameters["ν₋"]
ΔRi = train_parameters["ΔRi"]
Riᶜ = train_parameters["Riᶜ"]
Pr = train_parameters["Pr"]

# solve_oceananigans_modified_pacanowski_philander_nn(test_files, EXTRACTED_FILE_PATH, FILE_DIR,
#                                                         timestep=1, convective_adjustment=convective_adjustment)

diurnal = occursin("diurnal", test_files[1])

if num <= 18
    DIR_NAME = "train_$(test_files[1])"
    animation_type = "Training"
elseif num >= 19 && num <= 26
    DIR_NAME = "test_$(test_files[1])"
    animation_type = "Interpolating"
else
    DIR_NAME = "test_$(test_files[1])"
    animation_type = "Extrapolating"
end

DIR_NAME = "test_$(test_files[1])"
animation_type = "Extrapolating"

n_trainings = length(train_files)

plot_data = NDE_profile_oceananigans(joinpath(FILE_DIR, DIR_NAME), train_files, test_files,
                                  ν₀=ν₀, ν₋=ν₋, ΔRi=ΔRi, Riᶜ=Riᶜ, Pr=Pr,
                                  loss_scalings=loss_scalings,
                                  OUTPUT_PATH=joinpath("C:\\Users\\xinle\\Documents\\OceanParameterizations.jl\\final_results\\18sim_old", DIR_NAME, "profiles_fluxes_oceananigans.jld2"))

# animate_profiles_fluxes_final(plot_data, axis_images,
#                     joinpath(FILE_DIR, DIR_NAME, "$(test_files[1])_60fps"),
#                     animation_type=animation_type, n_trainings=n_trainings, training_types="Wind + Cooling, Wind + Heating", fps=60)


using Test
using MPI
using CLIMA.Topologies

function main()
  MPI.Init()
  comm = MPI.COMM_WORLD
  crank = MPI.Comm_rank(comm)
  csize = MPI.Comm_size(comm)

  @assert csize == 2

  DFloat = Float64
  Nx = 3
  Ny = 2
  x = range(DFloat(0); length=Nx+1, stop=1)
  y = range(DFloat(0); length=Ny+1, stop=1)

  topology = BrickTopology(comm, (x,y); boundary=[1 3; 2 4],
                           periodicity=(true,true))

  elems       = topology.elems
  realelems   = topology.realelems
  ghostelems  = topology.ghostelems
  sendelems   = topology.sendelems
  elemtocoord = topology.elemtocoord
  elemtoelem  = topology.elemtoelem
  elemtoface  = topology.elemtoface
  elemtoordr  = topology.elemtoordr
  elemtobndy  = topology.elemtobndy
  nabrtorank  = topology.nabrtorank
  nabrtorecv  = topology.nabrtorecv
  nabrtosend  = topology.nabrtosend

  globalelemtoface = [2  2  2  2  2  2
                      1  1  1  1  1  1
                      4  4  4  4  4  4
                      3  3  3  3  3  3]

  globalelemtoordr = ones(Int, size(globalelemtoface))
  globalelemtobndy = zeros(Int, size(globalelemtoface))

  if crank == 0
    nrealelem = 3
    globalelems = [1,2,3,4,5,6]
    elemtoelem_expect = [6 4 2 4 5 6
                         5 3 4 4 5 6
                         2 1 5 4 5 6
                         2 1 5 4 5 6]
    nabrtorank_expect = [1]
    nabrtorecv_expect = UnitRange{Int}[1:3]
    nabrtosend_expect = UnitRange{Int}[1:3]
  elseif crank == 1
    nrealelem = 3
    globalelems = [4,5,6,1,2,3]
    elemtoelem_expect = [6 4 2 4 5 6
                         5 3 4 4 5 6
                         3 6 1 4 5 6
                         3 6 1 4 5 6]
    nabrtorank_expect = [0]
    nabrtorecv_expect = UnitRange{Int}[1:3]
    nabrtosend_expect = UnitRange{Int}[1:3]
  end

  @test elems == 1:length(globalelems)
  @test realelems == 1:nrealelem
  @test ghostelems == nrealelem+1:length(globalelems)
  @test elemtoface[:,realelems] == globalelemtoface[:,globalelems[realelems]]
  @test elemtoelem == elemtoelem_expect
  @test elemtobndy == globalelemtobndy[:,globalelems]
  @test elemtoordr == ones(eltype(elemtoordr), size(elemtoordr))
  @test nabrtorank == nabrtorank_expect
  @test nabrtorecv == nabrtorecv_expect
  @test nabrtosend == nabrtosend_expect

  MPI.Finalize()
end

main()


# __BEGIN_LICENSE__
#
# ThreeDeconv.jl
#
# Copyright (c) 2018, Stanford University
#
# All rights reserved.
#
# Redistribution and use in source and binary forms for academic and other
# non-commercial purposes with or without modification, are permitted provided
# that the following conditions are met:
#
# * Redistributions of source code, including modified source code, must retain
#   the above copyright notice, this list of conditions and the following
#   disclaimer.
#
# * Redistributions in binary form or a modified form of the source code must
#   reproduce the above copyright notice, this list of conditions and the
#   following disclaimer in the documentation and/or other materials provided with
#   the distribution.
#
# * Neither the name of The Leland Stanford Junior University, any of its
#   trademarks, the names of its employees, nor contributors to the source code
#   may be used to endorse or promote products derived from this software without
#   specific prior written permission.
#
# * Where a modified version of the source code is redistributed publicly in
#   source or binary forms, the modified source code must be published in a freely
#   accessible manner, or otherwise redistributed at no charge to anyone
#   requesting a copy of the modified source code, subject to the same terms as
#   this agreement.
#
# THIS SOFTWARE IS PROVIDED BY THE TRUSTEES OF THE LELAND STANFORD JUNIOR
# UNIVERSITY "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE LELAND STANFORD JUNIOR
# UNIVERSITY OR ITS TRUSTEES BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
# IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
# OF SUCH DAMAGE.
#
# __END_LICENSE__

using QuadGK
import Optim

const ϵ_reg = 1e-6

struct ParametricNoiseModel
    m_hat::Vector{Float64}
    σ_hat::Vector{Float64}
    model::Function
    model_init::Function
    params
    params_init
end

struct localEstimation{T<:AbstractFloat}
    y::T
    σ::T
    κ::T
    n::Int
end

function foi_noiseestimation(z::AbstractArray{Float32}; τ=0.2, maxnum_pairs::Int, verbose::Bool=false)
    # τ = 0.2 is good to reject a lot of outliers.

    @assert maximum(z) <= 1.0
    @assert minimum(z) >= 0.0

    # Compute the local sample mean and standard deviation over a smooth region
    println("Computing local noise variance.")
    num_imgs = size(z,3)
    z_stack = (z[:,:,i] for i in 1:num_imgs)
    est_pairs_array = map(x->localnoiseestimation(x, τ)[1], z_stack)
    est_pairs = vcat(est_pairs_array...)

    println("Initializing parameters by least-squares.")
    ab_init, y_hat, σ_hat = initialize_parameters(est_pairs)

    num_pairs = length(est_pairs)
    if num_pairs > maxnum_pairs
        idx = rand(1:num_pairs, maxnum_pairs)
        est_pairs = est_pairs[idx]
    end

    println("Initialization done.")
    println("Starting likelihood maximization.")

    likelihood0(x::Vector{Float64}) = nonclipped_negloglikelihood(x, est_pairs)
    result = Optim.optimize(likelihood0, ab_init, Optim.NelderMead(), Optim.Options(show_trace = verbose))

    println("Finished the maximization.")

    ab_hat = Optim.minimizer(result)
    @assert Optim.converged(result)

    return ab_hat
end

function localnoiseestimation(z::AbstractArray{Float32,2}, τ)
    # Wavelet and scaling functions used in the original paper are the followings:
    # ψ = Array{Float32}([0.035, 0.085, -0.135, -0.460, 0.807, -0.333])
    # ϕ = Array{Float32}([0.025, -0.060, -0.095, 0.325, 0.571, 0.235])
    ϕ = [0.035226291882100656f0, -0.08544127388224149f0, -0.13501102001039084f0, 0.4598775021193313f0, 0.8068915093133388f0, 0.3326705529509569f0]
    ϕ ./= sum(ϕ)

    ψ = [-0.3326705529509569f0, 0.8068915093133388f0, -0.4598775021193313f0, -0.13501102001039084f0, 0.08544127388224149f0, 0.035226291882100656f0]
    ψ ./= norm(ψ)

    σ_gauss = 1.2f0
    gauss = [exp(-x^2 / (2.0f0 * σ_gauss^2)) for x in -10.f0:10.f0]
    gauss ./= sum(gauss)

    z_wdet = circconv(z, ψ)[1:2:end, 1:2:end]
    z_wapp = circconv(z, ϕ)[1:2:end, 1:2:end]

    ω = ones(Float32, 7) ./ 7.0f0
    z_smo = circconv(z_wapp, ω, ω)
    s = sqrt(0.5 * π) .* circconv(abs.(z_wdet), ω, ω)

    g = [-0.5f0, 0.0f0, 0.5f0]
    smoothed_zwapp = circconv(z_wapp, gauss, gauss)
    dx_wapp = circconv(smoothed_zwapp, [1.0f0], g)
    dy_wapp = circconv(smoothed_zwapp, g, [1.0f0])

    x_smo = sqrt.(dx_wapp.^2 .+ dy_wapp.^2) .< τ .* s
    N = length(z_smo)
    num_bins = 300
    histogram_zwapp = [Vector{Float32}() for _ in 1:num_bins]
    histogram_zwdet = [Vector{Float32}() for _ in 1:num_bins]
    min_wapp = sum(ϕ[ϕ .<  0.f0])
    max_wapp = sum(ϕ[ϕ .>= 0.f0])
    Δ = (max_wapp - min_wapp) / num_bins

    for i in 1:N
        if x_smo[i]
            idx = floor(Int, (z_smo[i] - min_wapp) / Δ)
            push!(histogram_zwapp[idx], z_wapp[i])
            push!(histogram_zwdet[idx], z_wdet[i])
        end
    end

    est_pairs = Vector{localEstimation{Float64}}()
    for i in 1:num_bins
        if histogram_zwdet[i] != []
            n = length(histogram_zwdet[i])
            κ = madBiasFactor(n)
            σ = sqrt(max(Float64(mad(histogram_zwdet[i], κ))^2, .0))
            y = Float64(mean(histogram_zwapp[i]))
            push!(est_pairs, localEstimation(y, σ, κ, n))
        end
    end

    return est_pairs, x_smo
end

function initialize_parameters(est_pairs::Vector{localEstimation{T}}) where T
    num_pairs = length(est_pairs)

    y_hat = zeros(num_pairs)
    σ_hat = zeros(num_pairs)

    Φ_tmp = Vector{Float64}()
    v = Vector{Float64}()

    for i = 1:num_pairs
        tmp = est_pairs[i]
        y_hat[i] = tmp.y
        σ_hat[i] = tmp.σ
        push!(Φ_tmp, tmp.y)
        push!(v, tmp.σ^2)
    end
    Φ = hcat(Φ_tmp, ones(length(Φ_tmp)))
    ab0 =  Φ \ v
    return ab0, y_hat, σ_hat
end

function nonclipped_negloglikelihood(ab::Vector{Float64},  est_pairs::Vector{localEstimation{T}}) where T

    Δ = .1
    total_val = .0

    for l in est_pairs
        c_i = 1.0 / l.n
        d_i = 1.35 / (l.n + 1.5)
        y_i = l.y
        σ_i = l.σ

        # This integrand is sometimes a very sharp peak-like function like a Dirac funciton.
        # Therefore, the direct numerical integration over [0.0, 1.0] is realatively difficult
        # because the algorithm may not evaluate the integrand around the peak.
        # To avoid this issue, the integration interval is separated into multiple intervals.
        # Since the peak is known to be close to y_i, the function value at y_i should be enough large than zero.
        integrand(y::Float64)::Float64 = 1.0 / σsq_reg(y, ab) *
            exp(-1.0 / (2.0 * σsq_reg(y, ab)) * ( (y_i - y)^2 / c_i + (σ_i - sqrt(σsq_reg(y, ab)) )^2 / d_i))

        val = .0
        if y_i - Δ > .0
            val += quadgk(integrand, .0, y_i - Δ)[1]
        end
        val += quadgk(integrand, max(.0, y_i - Δ), y_i)[1]
        val += quadgk(integrand, y_i, min(1.0, y_i + Δ))[1]
        if y_i + Δ < 1.0
            val += quadgk(integrand, y_i + Δ, 1.0)[1]
        end

        total_val -= log(1 / (2π * sqrt(c_i * d_i) ) * abs(val) + ϵ_reg)
    end

    return total_val
end

σsq_reg(y::Float64, ab::Vector{Float64}) = max(ϵ_reg^2, σsq(y, ab))
σsq(y::Float64, ab::Vector{Float64}) = ab[1] * y + ab[2]


@testset "GravityEffect/layerResponse write" begin
	sensor=Dict(:x => 100., :y => 150., :z => 100., :sensHeight => 1.0)
	layers=Dict(:start => [0.0, 1.0],:stop  => [1.0, 2.0])
	dem_in = joinpath(dirname(@__DIR__),"test","input","dem_data.asc");
	zones =Dict(:dem   => [dem_in,dem_in], :radius=> [50.,200.],
				:resolution => [0.1,0.2],:interpAltitude => [true,false]);
	exclude = Dict(:cylinder=>Dict(:radius=>1.,:start=>0.5,:stop=>2.),
			   :prism=>Dict(:x=>[100.,106.],:y=>[155.,156.],
							:dx=>[2.,3.],:dy=>[1.5,1.8],
							:start=>[0.,1.],:stop=>[1.,2.]),
			   :polygon=>Dict(:file=>joinpath(dirname(@__DIR__),"test","input","exclusion_polygon.txt"),
							  :start=>0.,:stop=>2.))
	outputfile = joinpath(dirname(@__DIR__),"test","output","write_layerResponse.txt");
	#outdata = GravityEffect.layerResponse(sensor,layers,zones,
	#					exclude=exclude,nanheight=true,
	#					outfile=outputfile,def_density=10.);
	outdata = DataFrame(layer=[1,2],start=layers[:start],stop=layers[:stop],
	 		   	   total=[3.82615,3.76013].*1e-9,zone1=[3.76682,3.66139].*1e-9,
				   zone2=[0.05933,0.09874].*1e-9);
    FileTools.write_layerResponse(sensor,layers,zones,exclude,false,
			   		    outputfile,10.,outdata);
	t = readdlm(outputfile,comments=true,comment_char='%');
	@test size(t) == (2,6)
	@test all(isapprox.(t[:,4],outdata[:total].*1e+9))
	@test all(isapprox.(t[:,5],outdata[:zone1].*1e+9))
	@test all(isapprox.(t[:,6],outdata[:zone2].*1e+9))
	@test t[:,1]==[1.,2.]
	@test t[:,2]==layers[:start]
	@test t[:,3]==layers[:stop]
end

@testset "GravityEffect/layerResponse read" begin
	filein = joinpath(dirname(@__DIR__),"test","input","read_layerResponse.txt");
	t = read_layerResponse(filein,"results")
	@test t[:layer]==[1.,2.]
	@test t[:start]==[0.,1.]
	@test t[:stop]==[1.,2.]
	@test t[:total]==[4.,3.]
	@test t[:zone1]==[3.1,2.2]
	@test t[:zone2]==[0.9,0.8]

	t = read_layerResponse(filein,"layers")
	@test t == Dict(:start => [0.0, 1.0],:stop  => [1.0, 2.0])
	t = read_layerResponse(filein,"sensor")
	@test t == Dict(:x => 100., :y => 150., :z => 100., :sensHeight => 1.0)
	t = read_layerResponse(filein,"def_density")
	@test t == 10.
	t = read_layerResponse(filein,"nanheight")
	@test !t

	dem_in = joinpath(dirname(@__DIR__),"test","input","dem_data.asc");
	zones =Dict(:dem   => [dem_in,dem_in], :radius=> [50.,200.],
				:resolution => [0.1,0.2],:interpAltitude => [true,false]);
	t = read_layerResponse(filein,"zones")
	#@test t == zones

	exclude = Dict(:cylinder=>Dict(:radius=>1.,:start=>0.5,:stop=>2.),
			   :prism=>Dict(:x=>[100.,106.],:y=>[155.,156.],
							:dx=>[2.,3.],:dy=>[1.5,1.8],
							:start=>[0.,1.],:stop=>[1.,2.]),
			   :polygon=>Dict(:file=>joinpath(dirname(@__DIR__),"test","input","exclusion_polygon.txt"),
							  :start=>0.,:stop=>2.))
	t = read_layerResponse(filein,"exclude")
	#@test t == exclude
end


export get_fields, get_records, build_result

function get_fields
end

function get_records
end

function build_result
end

try
    # interface for `DataFrame`s
    get_fields(df :: DataFrame) = names(df)
    get_records(df :: DataFrame) = zip(DataFrames.columns(df)...)
    function build_result(::Type{DataFrame}, fields, typs, source :: Union{Base.Generator, Vector{Tp}}) where Tp <: Tuple
        res = Tuple(typ[] for typ in typs)
        for each in source
            push!.(res, each)
        end
        DataFrame(collect(res), fields)
    end
catch
    @warn """
    DataFrames.jl not installed. If you want to peform queries on DataFrames with Wanderer.jl,
    build Wanderer.jl once again after installed DataFrames.jl.
    """
end

# interface for `Dict`s
get_fields(df :: D) where D <: Dict{Symbol} = collect(keys(df))
get_records(df :: D) where D <: Dict{Symbol} = zip((df[key] for key in keys(df))...)
function build_result(::Type{D}, fields, typs, source :: Union{Base.Generator, Vector{Tp}}) where {D <: Dict{Symbol}, Tp <: Tuple}
    res = Tuple(typ[] for typ in typs)
    for each in source
        push!.(res, each)
    end
    Dict(zip(fields, res))
end