module SIMD

#=

# Various boolean types

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(::Type{Val{$sz}}) = $Boolsz
        inttype(::Type{Val{$sz}}) = $Intsz
        uinttype(::Type{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)})

=#

# 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 FloatingTypes = Union{Float16, Float32, Float64}
const ScalarTypes = Union{IntegerTypes, FloatingTypes}

export Vec
immutable Vec{N,T<:ScalarTypes} <: DenseArray{T,1}   # <: Number
    elts::NTuple{N,T}
    Vec(elts::NTuple{N,T}) = new(elts)
end

function Base.show{N,T}(io::IO, v::Vec{N,T})
    print(io, T, "<")
    for i in 1:N
        i>1 && print(io, ",")
        print(io, v.elts[i])
    end
    print(io, ">")
end

# Base.print_matrix wants to access a second dimension that doesn't exist for
# Vec. (In Julia, every array can be accessed as N-dimensional array, for
# arbitrary N.) Instead of implementing this, output our Vec the usual way.
function Base.print_matrix(io::IO, X::Vec,
                      pre::AbstractString = " ",  # pre-matrix string
                      sep::AbstractString = "  ", # separator between elements
                      post::AbstractString = "",  # post-matrix string
                      hdots::AbstractString = "  \u2026  ",
                      vdots::AbstractString = "\u22ee",
                      ddots::AbstractString = "  \u22f1  ",
                      hmod::Integer = 5, vmod::Integer = 5)
    print(io, X)
end

# Type properties

# eltype and ndims is provided by DenseArray
# Base.eltype{N,T}(::Type{Vec{N,T}}) = T
# Base.ndims{N,T}(::Type{Vec{N,T}}) = 1
Base.length{N,T}(::Type{Vec{N,T}}) = N
Base.size{N,T}(::Type{Vec{N,T}}) = (N,)
Base.size{N,T}(::Type{Vec{N,T}}, n::Integer) = (N,)[n]
# Base.eltype{N,T}(::Vec{N,T}) = T
# Base.ndims{N,T}(::Vec{N,T}) = 1
Base.length{N,T}(::Vec{N,T}) = N
Base.size{N,T}(::Vec{N,T}) = (N,)
Base.size{N,T}(::Vec{N,T}, n::Integer) = (N,)[n]

# Type conversion

@generated function create{N,T<:ScalarTypes}(::Type{Vec{N,T}}, x::T)
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}($(Expr(:tuple, [:x for i in 1:N]...)))
    end
end
Base.convert{N,T}(::Type{Vec{N,T}}, x::T) = create(Vec{N,T}, x)
Base.convert{N,T}(::Type{Vec{N,T}}, x::Number) = create(Vec{N,T}, T(x))
Base.convert{N,T}(::Type{Vec{N,T}}, xs::NTuple{N}) = Vec{N,T}(NTuple{N,T}(xs))

Base.convert{N,T}(::Type{NTuple{N,T}}, v::Vec{N,T}) = v.elts

# Promotion rules

# Type promotion only works for subtypes of Number!
# Base.promote_rule{N,T<:ScalarTypes}(::Type{Vec{N,T}}, ::Type{T}) = 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{T<:FloatingTypes}(::Type{T}) = 8*sizeof(T) - 1 - significand_bits(T)
sign_bits{T<:FloatingTypes}(::Type{T}) = 1

significand_mask{T<:FloatingTypes}(::Type{T}) =
    uint_type(T)(uint_type(T)(1) << significand_bits(T) - 1)
exponent_mask{T<:FloatingTypes}(::Type{T}) =
    uint_type(T)(uint_type(T)(1) << exponent_bits(T) - 1) << significand_bits(T)
sign_mask{T<:FloatingTypes}(::Type{T}) =
    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{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{T<:IntegerTypes}(N::Integer, ::Type{T}) = "v$(N)i$(8*sizeof(T))"
suffix{T<:FloatingTypes}(N::Integer, ::Type{T}) = "v$(N)f$(8*sizeof(T))"

# Type-dependent LLVM constants
function llvmconst{T}(::Type{T}, val)
    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{T}(N::Integer, ::Type{T}, val)
    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

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

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

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

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

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

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

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

llvmins{T<:FloatingTypes}(::Type{Val{:^}}, N, ::Type{T}) =
    "@llvm.pow.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:abs}}, N, ::Type{T}) =
    "@llvm.fabs.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:ceil}}, N, ::Type{T}) =
    "@llvm.ceil.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:copysign}}, N, ::Type{T}) =
    "@llvm.copysign.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:cos}}, N, ::Type{T}) =
    "@llvm.cos.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:exp}}, N, ::Type{T}) =
    "@llvm.exp.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:exp2}}, N, ::Type{T}) =
    "@llvm.exp2.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:floor}}, N, ::Type{T}) =
    "@llvm.floor.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:fma}}, N, ::Type{T}) =
    "@llvm.fma.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:log}}, N, ::Type{T}) =
    "@llvm.log.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:log10}}, N, ::Type{T}) =
    "@llvm.log10.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:log2}}, N, ::Type{T}) =
    "@llvm.log2.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:max}}, N, ::Type{T}) =
    "@llvm.maxnum.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:min}}, N, ::Type{T}) =
    "@llvm.minnum.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:muladd}}, N, ::Type{T}) =
    "@llvm.fmuladd.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:powi}}, N, ::Type{T}) =
    "@llvm.powi.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:round}}, N, ::Type{T}) =
    "@llvm.rint.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:sin}}, N, ::Type{T}) =
    "@llvm.sin.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:sqrt}}, N, ::Type{T}) =
    "@llvm.sqrt.$(suffix(N,T))"
llvmins{T<:FloatingTypes}(::Type{Val{:trunc}}, N, ::Type{T}) =
    "@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 scalar2array(varrec, 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(arr,i)) = " *
                "insertvalue [$siz x $typ] $(accum(arr,i-1)), $typ $sca, $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)
    instrs = []
    accum(nam, i) = i<0 ? "undef" : i==siz-1 ? nam : "$(nam)_iter$i"
    rsiz = siz1 + siz2
    for i in 0:siz-1
        push!(instrs,
            "$(tmp)_elem$i = extractelement <$siz1 x $typ> $vec1, 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

function combinevectors(vec1, siz1, vec2, siz2, typ, rvec, tmp="$(rvec)_cv")
    instrs = []
    accum(nam, i) = i<0 ? "undef" : i==siz-1 ? nam : "$(nam)_iter$i"
    rsiz = siz1 + siz2
    for i in 0:siz1-1
        push!(instrs,
            "$(tmp)_elem$i = extractelement <$siz1 x $typ> $vec1, i32 $i")
        push!(instrs,
            "$(accum(rvec,i)) = "*
                "insertelement <$rsiz x $typ> $(accum(rvec,i-1)), " *
                "$typ $(tmp)_elem$i, $i")
    end
    for i in siz1:siz2-1
        push!(instrs,
            "$(tmp)_elem$i = " *
                "extractelement <$siz2 x $typ> $vec2, i32 $(i-siz1)")
        push!(instrs,
            "$(accum(rvec,i)) = "*
                "insertelement <$rsiz x $typ> $(accum(rvec,i-1)), " *
                "$typ $(tmp)_elem$i, i32 $i")
    end
    instrs
end

# Element-wise access

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

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

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

# Type conversion

@generated function Base.reinterpret{N,R,N1,T1}(::Type{Vec{N,R}},
        v1::Vec{N1,T1})
    @assert N*sizeof(R) == N1*sizeof(T1)
    typ1 = llvmtype(T1)
    atyp1 = "[$N1 x $typ1]"
    vtyp1 = "<$N1 x $typ1>"
    typr = llvmtype(R)
    atypr = "[$N x $typr]"
    vtypr = "<$N x $typr>"
    decls = []
    instrs = []
    append!(instrs, array2vector("%arg1", N1, typ1, "%0", "%arg1arr"))
    push!(instrs, "%res = bitcast $vtyp1 %arg1 to $vtypr")
    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{N1,T1}}, v1.elts))
    end
end

# Generic function wrappers

@generated function llvmwrap{Op,N,T1,R}(::Type{Val{Op}}, v1::Vec{N,T1},
        ::Type{R} = T1)
    @assert isa(Op, Symbol)
    typ1 = llvmtype(T1)
    atyp1 = "[$N x $typ1]"
    vtyp1 = "<$N x $typ1>"
    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)")
        push!(instrs, "%res = call $vtypr $ins($vtyp1 %arg1)")
    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, %arg1")
    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}}, v1.elts))
    end
end

@generated function llvmwrap{Op,N}(::Type{Val{Op}}, v1::Vec{N,Bool},
        ::Type{Bool} = Bool)
    @assert isa(Op, Symbol)
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    abtyp = "[$N x $btyp]"
    ins = llvmins(Val{Op}, N, Bool)
    decls = []
    instrs = []
    append!(instrs, array2vector("%arg1b", N, btyp, "%0", "%arg1arr"))
    push!(instrs, "%arg1 = trunc $vbtyp %arg1b 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")
    append!(instrs, vector2array("%resarr", N, btyp, "%resb"))
    push!(instrs, "ret $abtyp %resarr")
    quote
        $(Expr(:meta, :inline))
        Vec{N,Bool}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,Bool}, Tuple{NTuple{N,Bool}}, v1.elts))
    end
end

@generated function llvmwrap{Op,N,T1,T2,R}(::Type{Val{Op}}, v1::Vec{N,T1},
        v2::Vec{N,T2}, ::Type{R} = T1)
    @assert isa(Op, Symbol)
    typ1 = llvmtype(T1)
    atyp1 = "[$N x $typ1]"
    vtyp1 = "<$N x $typ1>"
    typ2 = llvmtype(T2)
    atyp2 = "[$N x $typ2]"
    vtyp2 = "<$N x $typ2>"
    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"))
    append!(instrs, array2vector("%arg2", N, typ2, "%1", "%arg2arr"))
    if ins[1] == '@'
        push!(decls, "declare $vtypr $ins($vtyp1, $vtyp2)")
        push!(instrs, "%res = call $vtypr $ins($vtyp1 %arg1, $vtyp2 %arg2)")
    else
        push!(instrs, "%res = $ins $vtyp1 %arg1, %arg2")
    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}, NTuple{N,T2}}, v1.elts, v2.elts))
    end
end

@generated function llvmwrap{Op,N,T1,T2}(::Type{Val{Op}}, v1::Vec{N,T1},
        v2::Vec{N,T2}, ::Type{Bool})
    @assert isa(Op, Symbol)
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    abtyp = "[$N x $btyp]"
    typ1 = llvmtype(T1)
    atyp1 = "[$N x $typ1]"
    vtyp1 = "<$N x $typ1>"
    typ2 = llvmtype(T2)
    atyp2 = "[$N x $typ2]"
    vtyp2 = "<$N x $typ2>"
    ins = llvmins(Val{Op}, N, T1)
    decls = []
    instrs = []
    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")
    quote
        $(Expr(:meta, :inline))
        Vec{N,Bool}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,Bool}, Tuple{NTuple{N,T1}, NTuple{N,T2}},
            v1.elts, v2.elts))
    end
end

@generated function llvmwrap{Op,N,T1,T2,R}(::Type{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

@generated function llvmwrap{Op,N}(::Type{Val{Op}}, v1::Vec{N,Bool},
        v2::Vec{N,Bool}, ::Type{Bool} = Bool)
    @assert isa(Op, Symbol)
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    abtyp = "[$N x $btyp]"
    ins = llvmins(Val{Op}, N, Bool)
    decls = []
    instrs = []
    append!(instrs, array2vector("%arg1b", N, btyp, "%0", "%arg1arr"))
    append!(instrs, array2vector("%arg2b", N, btyp, "%1", "%arg2arr"))
    push!(instrs, "%arg1 = trunc $vbtyp %arg1b to <$N x i1>")
    push!(instrs, "%arg2 = trunc $vbtyp %arg2b to <$N x i1>")
    push!(instrs, "%res = $ins <$N x i1> %arg1, %arg2")
    push!(instrs, "%resb = zext <$N x i1> %res to $vbtyp")
    append!(instrs, vector2array("%resarr", N, btyp, "%resb"))
    push!(instrs, "ret $abtyp %resarr")
    quote
        $(Expr(:meta, :inline))
        Vec{N,Bool}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,Bool}, Tuple{NTuple{N,Bool}, NTuple{N,Bool}},
            v1.elts, v2.elts))
    end
end

@generated function llvmwrap{Op,N,T1,T2,T3,R}(::Type{Val{Op}}, v1::Vec{N,T1},
        v2::Vec{N,T2}, v3::Vec{N,T3}, ::Type{R} = T1)
    @assert isa(Op, Symbol)
    typ1 = llvmtype(T1)
    atyp1 = "[$N x $typ1]"
    vtyp1 = "<$N x $typ1>"
    typ2 = llvmtype(T2)
    atyp2 = "[$N x $typ2]"
    vtyp2 = "<$N x $typ2>"
    typ3 = llvmtype(T3)
    atyp3 = "[$N x $typ3]"
    vtyp3 = "<$N x $typ3>"
    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"))
    append!(instrs, array2vector("%arg2", N, typ2, "%1", "%arg2arr"))
    append!(instrs, array2vector("%arg3", N, typ3, "%2", "%arg3arr"))
    if ins[1] == '@'
        push!(decls, "declare $vtypr $ins($vtyp1, $vtyp2, $vtyp3)")
        push!(instrs,
            "%res = call $vtypr $ins($vtyp1 %arg1, $vtyp2 %arg2, $vtyp3 %arg3)")
    else
        push!(instrs, "%res = $ins $vtyp1 %arg1, %arg2, %arg3")
    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}, NTuple{N,T2}, NTuple{N,T3}},
            v1.elts, v2.elts, v3.elts))
    end
end

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

@generated function llvmwrapshift{Op,N,T}(::Type{Val{Op}}, v1::Vec{N,T},
        x2::T)
    @assert isa(Op, Symbol)
    typ = llvmtype(T)
    atyp = "[$N x $typ]"
    vtyp = "<$N x $typ>"
    ins = llvmins(Val{Op}, N, T)
    decls = []
    instrs = []
    append!(instrs, array2vector("%arg1", N, typ, "%0", "%arg1arr"))
    append!(instrs, scalar2vector("%arg2", N, typ, "%1"))
    nbits = 8*sizeof(T)
    push!(instrs, "%tmp = $ins $vtyp %arg1, %arg2")
    push!(instrs, "%inbounds = icmp ult $typ %1, $nbits")
    if Op === :>> && T <: IntTypes
        nbits = llvmconst(N, T, 8*sizeof(T)-1)
        push!(instrs, "%limit = $ins $vtyp %arg1, $nbits")
        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
    append!(instrs, vector2array("%resarr", N, typ, "%res"))
    push!(instrs, "ret $atyp %resarr")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,T}, Tuple{NTuple{N,T}, T}, v1.elts, x2))
    end
end

@generated function llvmwrapshift{Op,N,T}(::Type{Val{Op}}, v1::Vec{N,T},
        v2::Vec{N,T})
    @assert isa(Op, Symbol)
    typ = llvmtype(T)
    atyp = "[$N x $typ]"
    vtyp = "<$N x $typ>"
    ins = llvmins(Val{Op}, N, T)
    decls = []
    instrs = []
    append!(instrs, array2vector("%arg1", N, typ, "%0", "%arg1arr"))
    append!(instrs, array2vector("%arg2", N, typ, "%1", "%arg2arr"))
    push!(instrs, "%tmp = $ins $vtyp %arg1, %arg2")
    nbits = llvmconst(N, T, 8*sizeof(T))
    push!(instrs, "%inbounds = icmp ult $vtyp %arg2, $nbits")
    if Op === :>> && T <: IntTypes
        nbits = llvmconst(N, T, 8*sizeof(T)-1)
        push!(instrs, "%limit = $ins $vtyp %arg1, $nbits")
        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
    append!(instrs, vector2array("%resarr", N, typ, "%res"))
    push!(instrs, "ret $atyp %resarr")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,T}, Tuple{NTuple{N,T}, NTuple{N,T}}, v1.elts, v2.elts))
    end
end

# Conditionals

for op in (:(==), :(!=), :(<), :(<=), :(>), :(>=))
    @eval begin
        @inline Base.$op{N,T}(v1::Vec{N,T}, v2::Vec{N,T}) =
            llvmwrap(Val{$(QuoteNode(op))}, v1, v2, Bool)
    end
end
# @inline Base.isfinite{N,T<:FloatingTypes}(v1::Vec{N,T}) =
#     !(isinf(v1) | isnan(v1))
@inline function Base.isfinite{N,T<:FloatingTypes}(v1::Vec{N,T})
    U = uint_type(T)
    em = Vec{N,U}(exponent_mask(T))
    iv = reinterpret(Vec{N,U}, v1)
    iv & em != em
end
@inline Base.isinf{N,T<:FloatingTypes}(v1::Vec{N,T}) = abs(v1) == Vec{N,T}(Inf)
# @inline function Base.isinf{N,T<:FloatingTypes}(v1::Vec{N,T})
#     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|sm) == em
# end
@inline Base.isnan{N,T<:FloatingTypes}(v1::Vec{N,T}) = v1 != v1
# @inline function Base.isnan{N,T<:FloatingTypes}(v1::Vec{N,T})
#     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 == em)  & (iv & sm != Vec{N,U}(0))
# end
# @inline Base.isnormal{N,T<:FloatingTypes}(v1::Vec{N,T}) = ???
@inline function Base.issubnormal{N,T<:FloatingTypes}(v1::Vec{N,T})
    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{N,T<:FloatingTypes}(v1::Vec{N,T})
    U = uint_type(T)
    sm = Vec{N,U}(sign_mask(T))
    iv = reinterpret(Vec{N,U}, v1)
    iv & sm != Vec{N,U}(0)
end

#=
@generated function Base.ifelse{N,T}(x1::Bool, v2::Vec{N,T}, v3::Vec{N,T})
    typ = llvmtype(T)
    atyp = "[$N x $typ]"
    vtyp = "<$N x $typ>"
    decls = []
    instrs = []
    append!(instrs, array2vector("%arg2", N, typ, "%1", "%arg2arr"))
    append!(instrs, array2vector("%arg3", N, typ, "%2", "%arg3arr"))
    push!(instrs, "%res = select i1 %0, $vtyp %arg2, $vtyp %arg3")
    append!(instrs, vector2array("%resarr", N, typ, "%res"))
    push!(instrs, "ret $atyp %resarr")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,T}, Tuple{Bool, NTuple{N,T}, NTuple{N,T}},
            x1, v2.elts, v3.elts))
    end
end
=#

@generated function Base.ifelse{N,T}(v1::Vec{N,Bool}, v2::Vec{N,T},
        v3::Vec{N,T})
    btyp = llvmtype(Bool)
    vbtyp = "<$N x $btyp>"
    abtyp = "[$N x $btyp]"
    typ = llvmtype(T)
    atyp = "[$N x $typ]"
    vtyp = "<$N x $typ>"
    decls = []
    instrs = []
    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")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,T}, Tuple{NTuple{N,Bool}, NTuple{N,T}, NTuple{N,T}},
            v1.elts, v2.elts, v3.elts))
    end
end

# Integer arithmetic functions

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

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

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

for op in (:<<, :>>, :>>>)
    @eval begin
        @inline Base.$op{N,T<:IntegerTypes,I}(v1::Vec{N,T}, ::Type{Val{I}}) =
            llvmwrapshift(Val{$(QuoteNode(op))}, v1, Val{I})
        @inline Base.$op{N,T<:IntegerTypes}(v1::Vec{N,T}, x2::Int) =
            llvmwrapshift(Val{$(QuoteNode(op))}, v1, T(x2))
        @inline Base.$op{N,T<:IntegerTypes}(v1::Vec{N,T}, x2::Integer) =
            llvmwrapshift(Val{$(QuoteNode(op))}, v1, T(x2))
        @inline Base.$op{N,T<:IntegerTypes}(v1::Vec{N,T}, v2::Vec{N,T}) =
            llvmwrapshift(Val{$(QuoteNode(op))}, v1, v2)
        @inline Base.$op{N,T<:IntegerTypes}(x1::T, v2::Vec{N,T}) =
            $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{N,T<:FloatingTypes}(v1::Vec{N,T}) =
            llvmwrap(Val{$(QuoteNode(op))}, v1)
    end
end
@inline Base.exp10{N,T<:FloatingTypes}(v1::Vec{N,T}) = Vec{N,T}(10)^v1
@inline Base.sign{N,T<:FloatingTypes}(v1::Vec{N,T}) =
    ifelse(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{N,T<:FloatingTypes}(v1::Vec{N,T}, v2::Vec{N,T}) =
            llvmwrap(Val{$(QuoteNode(op))}, v1, v2)
    end
end
@inline Base. ^{N,T<:FloatingTypes}(v1::Vec{N,T},x2::Integer) =
    llvmwrap(Val{:powi}, v1, Int(x2))
@inline Base.flipsign{N,T<:FloatingTypes}(v1::Vec{N,T}, v2::Vec{N,T}) =
    ifelse(signbit(v2), -v1, v1)

for op in (:fma, :muladd)
    @eval begin
        @inline function Base.$op{N,T<:FloatingTypes}(v1::Vec{N,T}, v2::Vec{N,T},
                v3::Vec{N,T})
            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{N}(s1::Bool, v2::Vec{N,Bool}) =
            $op(Vec{N,Bool}(s1), v2)
        @inline Base.$op{N,T<:IntegerTypes}(s1::IntegerTypes, v2::Vec{N,T}) =
            $op(Vec{N,T}(s1), v2)
        @inline Base.$op{N,T<:IntegerTypes}(v1::Vec{N,T}, s2::IntegerTypes) =
            $op(v1, Vec{N,T}(s2))
    end
end
@inline Base.ifelse{N,T<:IntegerTypes}(c::Vec{N,Bool}, s1::IntegerTypes,
        v2::Vec{N,T}) =
    ifelse(c, Vec{N,T}(s1), v2)
@inline Base.ifelse{N,T<:IntegerTypes}(c::Vec{N,Bool}, v1::Vec{N,T},
        s2::IntegerTypes) =
    ifelse(c, v1, Vec{N,T}(s2))

for op in (:muladd,)
    @eval begin
        @inline Base.$op{N,T<:IntegerTypes}(s1::IntegerTypes, v2::Vec{N,T},
                v3::Vec{N,T}) =
            $op(Vec{N,T}(s1), v2, v3)
        @inline Base.$op{N,T<:IntegerTypes}(v1::Vec{N,T}, s2::IntegerTypes,
                v3::Vec{N,T}) =
            $op(v1, Vec{N,T}(s2), v3)
        @inline Base.$op{N,T<:IntegerTypes}(s1::IntegerTypes, s2::IntegerTypes,
                v3::Vec{N,T}) =
            $op(Vec{N,T}(s1), Vec{N,T}(s2), v3)
        @inline Base.$op{N,T<:IntegerTypes}(v1::Vec{N,T}, v2::Vec{N,T},
                s3::IntegerTypes) =
            $op(v1, v2, Vec{N,T}(s3))
        @inline Base.$op{N,T<:IntegerTypes}(s1::IntegerTypes, v2::Vec{N,T},
                s3::IntegerTypes) =
            $op(Vec{N,T}(s1), v2, Vec{N,T}(s3))
        @inline Base.$op{N,T<:IntegerTypes}(v1::Vec{N,T}, s2::IntegerTypes,
                s3::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{N,T<:FloatingTypes}(s1::ScalarTypes, v2::Vec{N,T}) =
            $op(Vec{N,T}(s1), v2)
        @inline Base.$op{N,T<:FloatingTypes}(v1::Vec{N,T}, s2::ScalarTypes) =
            $op(v1, Vec{N,T}(s2))
    end
end
@inline Base.ifelse{N,T<:FloatingTypes}(c::Vec{N,Bool}, s1::ScalarTypes,
        v2::Vec{N,T}) =
    ifelse(c, Vec{N,T}(s1), v2)
@inline Base.ifelse{N,T<:FloatingTypes}(c::Vec{N,Bool}, v1::Vec{N,T},
        s2::ScalarTypes) =
    ifelse(c, v1, Vec{N,T}(s2))

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

# Reduction operations

function getneutral{T}(op::Symbol, ::Type{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

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

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

@generated function Base.reduce{Op,N,T}(::Type{Val{Op}}, v::Vec{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 Base.maximum{N,T<:IntegerTypes}(v::Vec{N,T}) = reduce(Val{:max}, v)
@inline Base.minimum{N,T<:IntegerTypes}(v::Vec{N,T}) = reduce(Val{:min}, v)

# Load and store functions

export vload, vloada
@generated function vload{N,T,Aligned}(::Type{Vec{N,T}}, ptr::Ptr{T},
        ::Type{Val{Aligned}} = Val{false})
    @assert isa(Aligned, Bool)
    typ = llvmtype(T)
    atyp = "[$N x $typ]"
    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 flags *= ", align $align" end
    push!(instrs, "%ptr = bitcast $typ* %0 to $vtyp*")
    push!(instrs, "%res = load $vtyp, $vtyp* %ptr$flags")
    append!(instrs, vector2array("%resarr", N, typ, "%res"))
    push!(instrs, "ret $atyp %resarr")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,T}, Tuple{Ptr{T}}, ptr))
    end
end

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

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

@generated function vload{N,T,Aligned}(::Type{Vec{N,T}}, ptr::Ptr{T},
        mask::Vec{N,Bool}, ::Type{Val{Aligned}} = Val{false})
    @assert isa(Aligned, Bool)
    typ = llvmtype(T)
    atyp = "[$N x $typ]"
    vtyp = "<$N x $typ>"
    btyp = llvmtype(Bool)
    abtyp = "[$N x $btyp]"
    vbtyp = "<$N x $btyp>"
    decls = []
    instrs = []
    if Aligned
        align = N * sizeof(T)
    else
        align = sizeof(T)   # This is overly optimistic
    end
    push!(instrs, "%ptr = bitcast $typ* %0 to $vtyp*")
    append!(instrs, array2vector("%maskb", N, btyp, "%1", "%arg2arr"))
    push!(instrs, "%mask = trunc $vbtyp %maskb 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)))")
    append!(instrs, vector2array("%resarr", N, typ, "%res"))
    push!(instrs, "ret $atyp %resarr")
    quote
        $(Expr(:meta, :inline))
        Vec{N,T}(Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            NTuple{N,T}, Tuple{Ptr{T}, NTuple{N,Bool}}, ptr, mask.elts))
    end
end

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

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

export vstore, vstorea
@generated function vstore{N,T,Aligned}(v::Vec{N,T}, ptr::Ptr{T},
        ::Type{Val{Aligned}} = Val{false})
    @assert isa(Aligned, Bool)
    typ = llvmtype(T)
    atyp = "[$N x $typ]"
    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 flags *= ", align $align" end
    append!(instrs, array2vector("%arg1", N, typ, "%0", "%arg1arr"))
    push!(instrs, "%ptr = bitcast $typ* %1 to $vtyp*")
    push!(instrs, "store $vtyp %arg1, $vtyp* %ptr$flags")
    push!(instrs, "ret void")
    quote
        $(Expr(:meta, :inline))
        Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            Void, Tuple{NTuple{N,T}, Ptr{T}}, v.elts, ptr)
    end
end

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

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

@generated function vstore{N,T,Aligned}(v::Vec{N,T}, ptr::Ptr{T},
        mask::Vec{N,Bool}, ::Type{Val{Aligned}} = Val{false})
    @assert isa(Aligned, Bool)
    typ = llvmtype(T)
    atyp = "[$N x $typ]"
    vtyp = "<$N x $typ>"
    btyp = llvmtype(Bool)
    abtyp = "[$N x $btyp]"
    vbtyp = "<$N x $btyp>"
    decls = []
    instrs = []
    if Aligned
        align = N * sizeof(T)
    else
        align = sizeof(T)   # This is overly optimistic
    end
    append!(instrs, array2vector("%arg1", N, typ, "%0", "%arg1arr"))
    push!(instrs, "%ptr = bitcast $typ* %1 to $vtyp*")
    append!(instrs, array2vector("%maskb", N, btyp, "%2", "%arg3arr"))
    push!(instrs, "%mask = trunc $vbtyp %maskb 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 %arg1, " *
            "$vtyp* %ptr, i32 $align, <$N x i1> %mask)")
    push!(instrs, "ret void")
    quote
        $(Expr(:meta, :inline))
        Base.llvmcall($((join(decls, "\n"), join(instrs, "\n"))),
            Void, Tuple{NTuple{N,T}, Ptr{T}, NTuple{N,Bool}}, v.elts, ptr,
                mask.elts)
    end
end

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

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

end


using YOLO, Test

enable_info() = Base.CoreLogging.disable_logging(Base.CoreLogging.LogLevel(-1)) #Enable printing of @info messages
disable_info() = Base.CoreLogging.disable_logging(Base.CoreLogging.Info) #disable printing of @info messages

disable_info()
datadir = joinpath(dirname(dirname(pathof(YOLO))), "data")
datasetsdir = joinpath(datadir, "datasets")
pretraineddir = joinpath(datadir, "pretrained")

if !isdir(joinpath(datasetsdir, "voc2007", "VOCdevkit", "VOC2007"))
    @info "Downloading VOC dataset for testing..."
    YOLO.download_dataset("voc2007")
end

@testset "Example datasets have been downloaded" begin
    @test isdir(joinpath(datasetsdir, "voc2007", "VOCdevkit", "VOC2007"))
end

@testset "Example pretrained weights files have been downloaded" begin
    @test isfile(joinpath(
        pretraineddir,
        "v2_tiny",
        "voc2007",
        "v2_tiny_voc.weights",
    ))
end

@testset "Loading and running YOLOv2_tiny_voc pretrained model" begin

    num_images = 2

    voc = YOLO.datasets.VOC.populate()
    @test length(voc.image_paths) == 5011
    @test length(voc.label_paths) == 5011
    @test length(voc.objects) == 20

    settings = YOLO.pretrained.v2_tiny_voc.load(minibatch_size = num_images)
    @test settings.num_classes == 20

    vocloaded = YOLO.load(voc, settings, indexes = collect(1:num_images))
    @test size(vocloaded.imstack_mat) == (416, 416, 3, num_images)
    @test length(vocloaded.paddings) == num_images
    @test length(vocloaded.labels) == num_images
    #Checks that the VOC download hasn't changed
    @test vec(sum(vocloaded.imstack_mat, dims = (1, 2, 3))) ≈ [
        140752.47,
        122024.16,
        126477.125,
        114651.555,
        143701.72,
        196821.47,
        179382.5,
        83476.43,
        132382.72,
        166430.16,
    ][1:num_images]

    model = YOLO.v2_tiny.load(settings)
    YOLO.loadWeights!(model, settings)

    res = model(vocloaded.imstack_mat) #run once to deal with compillation overhead
    #accdata = minibatch(inp, out, settings.minibatch_size; xtype = xtype)

    t = @elapsed for i in 1:10
        model(vocloaded.imstack_mat)
    end

    inference_time = (t / 10) / num_images
    inference_rate = 1 / inference_time

    predictions = YOLO.postprocess(res, settings, conf_thresh = 0.3, iou_thresh = 0.3)
    @test length(predictions) == num_images # images
    @test length(predictions[1]) == 0
    @test length(predictions[2]) == 1

    @test isapprox(predictions[2][1].bbox.x, 0.23406872f0, rtol=0.05)
    @test isapprox(predictions[2][1].bbox.y, 0.28794688f0, rtol=0.05)
    @test isapprox(predictions[2][1].bbox.w, 0.8048826f0, rtol=0.05)
    @test isapprox(predictions[2][1].bbox.h, 0.46979588f0, rtol=0.05)
    @test predictions[2][1].class == 7
    @test isapprox(predictions[2][1].conf, 0.5042128f0, rtol=0.05)

    t = @elapsed for i in 1:10
        YOLO.postprocess(res, settings, conf_thresh = 0.3, iou_thresh = 0.3)
    end

    postprocess_time = (t / 10) / num_images
    postprocess_rate = 1 / postprocess_time

    ## Makie Tests
    #disabled because Makie can't be tested on headless CI
    # scene = YOLO.renderResult(vocloaded.imstack_mat[:,:,:,1], predictions[1], settings, save_file = "test.png")
    # @test isfile("test.png")
    # rm("test.png", force=true)

    enable_info()
    @info "YOLO_v2_tiny inference time per image: $(round(inference_time, digits=4)) seconds ($(round(inference_rate, digits=2)) fps)"
    @info "YOLO_v2_tiny postprocess time per image: $(round(postprocess_time, digits=4)) seconds ($(round(postprocess_rate, digits=2)) fps)"
    @info "Total time per image: $(round(inference_time + postprocess_time, digits=4)) seconds ($(round(1/(inference_time + postprocess_time), digits=2)) fps)"

end


# It is about initializing, adding and substracting common mass properties
# consisting of mass m, center of mass rCM and inertia tensor I.
# At initialization common mass properties are computed of each mass super object.
# The result is stored in the root (= element of buffer).
# If the structure changes dynamically, because of gripping, mass properties must be updated
# (e.g. a robot picks something up (compute resulting common mass properties)).

### ---------------------- initialize common mass properties ----------------------
# initializeMassComputation! is called at initialization for computing common mass properties
# of each mass super object. Common mass properties are stored and assigned to the root of
# a super object (root of a super object ^= element of buffer).
# At initialization it needs to differed if a root object has already mass properties,
# otherwise dummy mass properties must be assigned, before computing common properties.
# The initial common mass properties are stored separatelly.
# This information is used for computing common mass for dynamically changing structures.
function initializeMassComputation!(scene::Scene)
    if scene.initMassComp != true
        superObjs = scene.superObjs
        buffer    = scene.buffer
        for i = 1:length(buffer)
            rootSuperObj = buffer[i]
            #=
                if featureHasMass(rootSuperObj) # feature of root object has already mass properties
                    rootSuperObj.objectHasMass = true
                    copyMassPropertiesFromFeatureToObject3D(rootSuperObj.massProperties, rootSuperObj.feature.massProperties.m,
                    rootSuperObj.feature.massProperties.rCM, rootSuperObj.feature.massProperties.I)
                    callInitialInertiaTensor!(rootSuperObj, superObjs[i].superObjMass.superObj)
                else # root obj's feature has no mass properties, but common mass super object do
                    if length(superObjs[i].superObjMass.superObj) > 0
                        rootSuperObj.objectHasMass = true
                        callInitialInertiaTensor!(rootSuperObj, superObjs[i].superObjMass.superObj)
                    end
                end
            =#
            if featureHasMass(rootSuperObj) # feature of root object has already mass properties
                rootSuperObj.hasMass = true
                rootSuperObj.m       = rootSuperObj.feature.massProperties.m
                rootSuperObj.r_CM    = rootSuperObj.feature.massProperties.rCM
                rootSuperObj.I_CM    = rootSuperObj.feature.massProperties.I
                addMassPropertiesOfAllSuperObjChildsToRootSuperObj!(rootSuperObj, superObjs[i].superObjMass.superObj)
            else # root obj's feature has no mass properties, but common mass super object do
                if length(superObjs[i].superObjMass.superObj) > 0
                    rootSuperObj.hasMass = true
                    rootSuperObj.m       = 0.0
                    rootSuperObj.r_CM    = Modia3D.ZeroVector3D
                    rootSuperObj.I_CM    = SMatrix{3,3,Float64,9}(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
                    addMassPropertiesOfAllSuperObjChildsToRootSuperObj!(rootSuperObj, superObjs[i].superObjMass.superObj)
                end
            end
        end
        scene.initMassComp = true
    end
end

# callInitialInertiaTensor! is called at initialization.
# It loops over all objects of an actual mass super object.
# It sums up common mass, inertia tensor and center of mass.
# The results are stored twice in actual root of super object:
#       massProperties: container for actual values
# function callInitialInertiaTensor!(rootSuperObj::Object3D, actualMassSuperObject::Vector{Object3D,1})
function addMassPropertiesOfAllSuperObjChildsToRootSuperObj!(rootSuperObj::Object3D, actualMassSuperObject::Vector{Object3D})
    if length(actualMassSuperObject) > 0
        for i=1:length(actualMassSuperObject)
            obj = actualMassSuperObject[i]
            if rootSuperObj != obj && rootSuperObj == obj.parent
                addOrSubtractMassPropertiesOfChildToRoot!(rootSuperObj, obj; add=true)
        end; end
    else
        rootSuperObj.hasMass = false
    end
#  println(Modia3D.fullName(rootSuperObj), " m ", rootSuperObj.massProperties.m )
#  println(Modia3D.fullName(rootSuperObj), " rCM ", rootSuperObj.massProperties.rCM )
#  println(Modia3D.fullName(rootSuperObj), " I ", rootSuperObj.massProperties.I )
end

### --------------- computing common mass properties for two bodies ---------------
# A movable unit is added or substracted from its parent object.
# Only the parent (root of the super object) and its child (= root of a movable unit)
# hold information about their mass properties.
# First, the childs center of mass must be transformed to parents coordinate system.
# Further, the common center of mass with the help of the common mass can be caluclated.
# Both inertia tensors of child and parent are transformed with Satz von Steiner
#             I_steiner = E*I*E' + m * skew(rCM)' * skew(rCM),
# summed up and again transformed with the common center of mass.
# Adding: add mass properties of a child to its parents mass properties
# Substracting: substract mass properties of a child from its parents mass properties

# function computeInertiaTensorForTwoBodies!(massPropParent, massPropObj, obj; add=true)
function addOrSubtractMassPropertiesOfChildToRoot!(obj_root, obj_child; add=true)::Nothing
    massProperties_child = obj_child.feature.massProperties

    R_child   = obj_child.R_rel
    r_child   = obj_child.r_rel
    m_child   = massProperties_child.m
    rCM_child = massProperties_child.rCM
    I_child   = massProperties_child.I

    m_root    = obj_root.m
    rCM_root  = obj_root.r_CM
    I_root    = obj_root.I_CM

    # transform childs center of mass to parents coordinate system
    rCM_child_new = r_child + R_child' * rCM_child

    # I_child_steiner: I_child needs to be transformed to parents coordinate system
    # I_root_steiner: I_root needs to be transformed to parents coordinate system
    #                 (no need of rotation matrices)
    I_child_steiner = Modia3D.NullRotation * I_child * Modia3D.NullRotation' +
                    m_child * Modia3D.skew(rCM_child_new)' * Modia3D.skew(rCM_child_new)
    I_root_steiner = I_root +
                    m_root * Modia3D.skew(rCM_root)' * Modia3D.skew(rCM_root)

    if add
        # ----------- adding mass properties (parent + child) ----------------------
        # common mass (parent + child)
        m = m_root + m_child

        # common center of mass (parent + child)
        @assert(m > 0.0)
        rCM = (m_root * rCM_root + m_child * rCM_child_new)/m

        # I: substract new common mass multiplied with skew matrices of
        #    center of mass from the sum of I_root_steiner and I_child_steiner
        I = I_root_steiner + I_child_steiner - m * Modia3D.skew(rCM)' * Modia3D.skew(rCM)

        # Assign to obj_root
        obj_root.m    = m
        obj_root.r_CM = rCM
        obj_root.I_CM = I
        return nothing

        #= # other way to compute inertia tensor
        a1 = rCM_new - rCM_child_new
        a2 = rCM_root - rCM_new
        IchildNew =  I_child + m_child * Modia3D.skew(a1)' * Modia3D.skew(a1)
        IrootNew  =  I_root  + m_root  * Modia3D.skew(a2)' * Modia3D.skew(a2)
        InewNew = IchildNew + IrootNew
        =#
    else
        # ----------- substracting mass properties (parent - child) ----------------
        # new mass is the difference between parents and childs mass
        m = m_root - m_child

        # new center of mass is the difference between parents and childs rCM
        if m != 0.0
            rCM = (m_root * rCM_root - m_child * rCM_child_new)/m
        else
            rCM = Modia3D.ZeroVector3D
        end

        # I: substract the sum of I_child_steiner and new mass multiplied with skew matrices of
        #    new center of mass from I_root_steiner
        I = I_root_steiner - I_child_steiner - m * Modia3D.skew(rCM)' * Modia3D.skew(rCM)

        # Assign to obj_root
        obj_root.m    = m
        obj_root.r_CM = rCM
        obj_root.I_CM = I
        return nothing
end; end

# helper function for assigning mass properties
#=
function assignMassProperties(massProperties, m, rCM, I)
    massProperties.m   = m
    massProperties.rCM = rCM
    massProperties.I   = I
end
=#


module SynthControl

using Dates, DataFrames, Optim, LinearAlgebra, RecipesBase, Parameters, FixedEffectModels, Statistics
import CSV

using Reexport
@reexport using TreatmentPanels

export SimpleSCM, SyntheticDiD, fit!, isfitted
export load_brexit, load_germany, load_basque, load_smoking
export load_brexit_panel, load_germany_panel, load_basque_panel, load_smoking_panel

include("SimpleSCM.jl")
include("plotrecipes.jl")
include("load_data.jl")
include("SyntheticDiD.jl")

end # module


export Simulation

struct Simulation{A, C, D}
    x::Array{Float64, 2}
    v::Array{Float64, 2}
    Δ_no::MVector{3, Float64}
    nn_arr::Array{Int64, 2}
    dhp_neutral::Array{Float64, 2}
    dhp_ion::Array{Float64, 2}
    dhp_coup::Array{Float64, 2}
    hp::MVector{3, Float64}
    F::Array{Float64, 2}
    surfp::MVector{A, Int16}
    occnum::MVector{C, Int16}
    surfh::MVector{D, Int16}
    surfpinit::MVector{A, Int16}
    surfpnew::MVector{A, Int16}
    trajzmin::MVector{1, Float64}
    trajtheta::MVector{1, Float64}
    H::Array{Float64, 2}
    ψ::Array{ComplexF64, 2}
    ϕ::Array{ComplexF64, 2}
    Γ::Array{Float64, 2}
    λ::MVector{C, Float64}
    dhdea::Array{Float64, 2}
    dhdv::Array{Float64, 2}
    akl::MVector{1, ComplexF64}
    akk::MVector{1, Float64}
    dm::Array{Float64, 2}
    blk::Array{Float64, 2}
    blk2::Array{Float64, 2}
    Pb::Array{Float64, 1}
    storage_aop::Array{Float64, 1}
    storage_op::Array{Float64, 2}
    storage_e::Array{Float64, 2}
    storage_xno::Array{Float64, 2}
    storage_vno::Array{Float64, 2}
    storage_temp::Array{Float64, 1}
    storage_phop::Array{Float64, 1}
    phipsi::MVector{1, ComplexF64}
    Pbmaxest::MVector{1, Float64}
    vdot::Array{Float64, 2}
    vtemp::Array{Float64, 2}
    vscale::MVector{1, Float64}
    storage_K::Array{Float64, 1}
    storage_P::Array{Float64, 1}
    storage_state::Array{Float64, 1}
    storage_deltaKNO::Array{Float64, 1}
    storage_deltaKAu::Array{Float64, 1}
    storage_hoptimes::Array{Int64, 1}
    exnum::MVector{1, Int64}
    attnum::MVector{1, Int64}
    nf::MVector{1, Int64}
end

#standard type for arrays:

const float_array = Array{Float64, 2}


export plot_avg_longterm_temps, plot_glacier_dataset

function plot_glacier_dataset(gdirs_climate, PDE_refs)
    if plots
        # gdirs_climate = (dates, gdirs, longterm_temps, annual_temps)
        gdirs = gdirs_climate[2]
        theme = Theme(fontsize=18, font="TeX Gyre Heros")
        set_theme!(theme)
        # Ice surface velocities
        figV = Makie.Figure(resolution=(1200, 1800))
        n = 5
        m = floor(Int, length(gdirs_climate[2])/n)
        axsV = [Axis(figV[i, j]) for i in 1:n, j in 1:m]
        hidedecorations!.(axsV)
        tightlimits!.(axsV)
        Vs, hm_Vs = [],[]
        for (i, ax) in enumerate(axsV)
            ax.aspect = DataAspect()
            glacier_gd = xr.open_dataset(gdirs[i].get_filepath("gridded_data"))
            nx = glacier_gd.y.size # glacier extent
            ny = glacier_gd.x.size # really weird, but this is inversed
            Δx = abs(gdirs[i].grid.dx)
            Δy = abs(gdirs[i].grid.dy)
            rgi_id = gdirs[i].rgi_id
            name = gdirs[i].name
            ax.title = "$name - $rgi_id"
            Vx = PDE_refs["V̄x_refs"][i]
            Vy = PDE_refs["V̄y_refs"][i]
            V = (abs.(Vx) .+ abs.(Vy))./2
            push!(Vs,V)
            push!(hm_Vs, Makie.heatmap!(ax, (1:nx).*Δx, (1:ny).*Δy, fillZeros(V), colormap=:speed))
        end
        minV = minimum(minimum.(Vs))
        maxV = maximum(maximum.(Vs))
        Colorbar(figV[2:3,m+1], limits=(minV,maxV), colormap=:speed, label="Ice surface velocity")
        Label(figV[0, :], text = "Glacier dataset", textsize = 30)
        display(figV)
        Makie.save(joinpath(root_plots, "glaciers_V.pdf"), figV, pt_per_unit = 1)

        # Ice thickness
        figH = Makie.Figure(resolution=(1200, 1800))
        axsH = [Axis(figH[i, j]) for i in 1:n, j in 1:m]
        hidedecorations!.(axsH)
        tightlimits!.(axsH)
        Hs, hm_Hs = [],[]
        for (i, ax) in enumerate(axsH)
            ax.aspect = DataAspect()
            glacier_gd = xr.open_dataset(gdirs[i].get_filepath("gridded_data"))
            nx = glacier_gd.y.size # glacier extent
            ny = glacier_gd.x.size # really weird, but this is inversed
            Δx = abs(gdirs[i].grid.dx)
            Δy = abs(gdirs[i].grid.dy)
            rgi_id = gdirs[i].rgi_id
            name = gdirs[i].name
            ax.title = "$name - $rgi_id"
            H = glacier_gd.consensus_ice_thickness.data # initial ice thickness conditions for forward model
            push!(Hs, fillNaN(H))
            push!(hm_Hs, Makie.heatmap!(ax, H, colormap=:ice))
        end
        minH = minimum(minimum.(Hs))
        maxH = maximum(maximum.(Hs))
        Colorbar(figH[2:3,m+1], limits=(minH,maxH), colormap=:ice, label="Ice thickness (m)")
        Label(figH[0, :], text = "Glacier dataset", textsize = 30)
        display(figH)
        Makie.save(joinpath(root_plots, "glaciers_H.pdf"), figH, pt_per_unit = 1)
    end
end

function plot_avg_longterm_temps(climate, gdirs)
    if plots
        mean_longterm_temps, labels = [],[]
        for (climate_glacier, gdir) in zip(climate, gdirs)
            push!(mean_longterm_temps, climate_glacier["longterm_temps"].temp.data)
            push!(labels, gdir.rgi_id)
        end
        display(Plots.plot(mean_longterm_temps, label=permutedims(labels),
                            xlabel="Years", ylabel="Mean longterm air temperature (°C)", legend=:topright;
                            palette=palette(:tab10,15)))
    end
end

# TODO: correctly retrieve the glacier coordinates to plot them in `imshow` as an extent
function plot_monthly_map(climate, variable, year)
    climate = climate[variable].where(climate.time.dt.year == year, drop=true).groupby("time.month")
    fig_clim, ax_clim = pplt.subplots([1:6, 7:12], axheight=2)
    fig_clim.format(
        abc=true, abcloc="ul", suptitle= ("$year - monthly $variable")
    )
    for mon in 1:12
        if variable == "temp"
            m_var = ax_clim[mon].imshow(climate.mean()[mon], cmap="Thermal",
                    vmin=minimum(climate.mean().data), vmax=maximum(climate.mean().data)) # set common min max temp
        else
            m_var = ax_clim[mon].imshow(climate.sum()[mon], cmap="DryWet",
                    vmin=climate.sum().min().data, vmax=climate.sum().max().data) # set common min max precipitation
        end
        ax_clim[mon].set_title(Dates.monthname(mon))
        if(mon == 12)
            if variable == "temp"
                fig_clim.colorbar(m_var, label="Air temperature (°C)")
            else
                fig_clim.colorbar(m_var, label="Accumulated $variable (mm)")
            end
        end
    end
end


export FlowModel, CompiledFlowModel
export compile, nr_params
export getlayers

using TupleTools: flatten

import Base: length, eltype

@doc raw"""
The FlowModel structure is the most generic type of Flow model, in which the layers are not constrained to be of a specific type. The FlowModel structure contains the input dimensionality and a tuple of layers and can be constructed as `FlowModel( dim, (layer1, layer2, ...) )`.

Note: this model can be specialized by constraining the types of layers. This potentially allows for more efficient specialized methods that can deal with specifics of these layers, such as triangular jacobian matrices.
"""
struct FlowModel{N,T<:NTuple{N,AbstractLayer}} <: AbstractFlowModel
    dim     :: Int
    layers  :: T
end

@doc raw"""
Creates a flow model with uninitialized parameters.

Input arguments:
- `dim::Int` - input dimensionality.
- `layers<:NTuple{N,AbstractLayer}` - arbitrarily sized tuple containing abstractlayers.

Return arguments:
- `::Flowmodel` - model containing layers of which are appropriately sized according to the input dimensionality.
"""
function FlowModel(dim::Int, layers::T) where { T <: NTuple{N,AbstractLayerPlaceholder} where { N } }
    return FlowModel(dim, flatten(prepare.(dim, layers)))
end
function FlowModel(layers::T) where { T <: NTuple{N, Union{AbstractLayer, AbstractLayerPlaceholder}} where { N } }
    @assert typeof(first(layers)) <: InputLayer "The FlowModel requires an input dimension to be specified. This can be achieved, either by preceding the layers tuple with an integer as `FlowModel(dim, layers)`, or by starting the tuple of layers with an `InputLayer(dim)` as `FlowModel((InputLayer(dim), layers...))`."
    return FlowModel(getdim(first(layers)), flatten(prepare.(getdim(first(layers)), Base.tail(layers))))
end

# prepare function for setting correct sizes in the layers (without assigning the parameters yet!)
prepare(dim::Int, layers::T) where { T <: NTuple{N,Union{AbstractLayer, AbstractLayerPlaceholder}} where { N }} = _prepare(dim, layers)
Broadcast.broadcasted(::typeof(prepare), dim::Int, layers::T) where { T <: NTuple{N,Union{AbstractLayer, AbstractLayerPlaceholder}} where { N }} = broadcast(_prepare, Ref(dim), layers)

@doc raw"""
The CompiledFlowModel structure is the most generic type of compiled Flow model, in which the layers are not constrained to be of a specific type. The FlowModel structure contains the input dimension and a tuple of compiled layers. Do not manually create a CompiledFlowModel! Instead create a FlowModel first and compile it with `compile(model::FlowModel)`. This will make sure that all layers/mappings are configured with the proper dimensionality and with randomly sampled parameters. Alternatively, if you would like to pass your own parameters, call `compile(model::FlowModel, params::Vector)`.

Note: this model can be specialized by constraining the types of layers. This potentially allows for more efficient specialized methods that can deal with specifics of these layers, such as triangular jacobian matrices.
"""
struct CompiledFlowModel{N,T<:NTuple{N,AbstractLayer}} <: AbstractCompiledFlowModel
    dim     :: Int
    layers  :: T
end

@doc raw"""
`compile()` compiles a model by setting its parameters. It randomly sets parameter values in the layers and flows such that inference in the model can be obtained.

Input arguments
- `model::FlowModel` - a model of which the dimensionality of its layers/flows has been initialized, but its parameters have not been set.

Return arguments
- `::CompiledFlowModel` - a compiled model with set parameters, such that it can be used for processing data.
"""
function compile(model::FlowModel)
    # do not create parameters here for decentralized initializers (TODO)
    return CompiledFlowModel(model.dim, compile.(model.layers))
end

@doc raw"""
`compile(model::FlowModel, params::Vector)` lets you initialize a model `model` with a vector of parameters `params`.

Input arguments
- `model::FlowModel` - a model of which the dimensionality of its layers/flows has been initialized, but its parameters have not been set.
- `params::Vector`   - a vector of parameters with which the model should be compiled.

Return arguments
- `::CompiledFlowModel` - a compiled model with set parameters, such that it can be used for processing data.
"""
function compile(model::FlowModel, params::Vector)

    # assert whether the parameter vector complies with the model
    @assert nr_params(model) == length(params) "The number of parameters in the model does not match the passed number of parameters."

    # compile layers
    return CompiledFlowModel(getdim(model), compile(getlayers(model), params))

end
function compile(x::Tuple, params::Vector)

    # fetch number of parameters in layer
    params_in_layer = nr_params(first(x))

    if params_in_layer > 0
        # if layer has parameters, compile that layer and the other layers with parameters
        return (compile(first(x), params[1:params_in_layer]), compile(Base.tail(x), params[params_in_layer+1:end])...)
    else
        # if layer has no parameters, compile that layer on its own and the other layers with parameters
        return (compile(first(x)), compile(Base.tail(x), params)...)
    end

end
function compile(x::Tuple)
    return (compile(first(x)), compile(Base.tail(x))...)
end
# compile functions to stop the recursion
compile(x::Tuple{}, params::Vector) = ()
compile(x::Tuple{})         = ()

# get-functions for the FlowModel structure
getdim(model::FlowModel)                    = model.dim
getdim(model::CompiledFlowModel)            = model.dim
getlayers(model::FlowModel)                 = model.layers
getlayers(model::CompiledFlowModel)         = model.layers
getforward(model::CompiledFlowModel)        = (x) -> forward(model, x)
getbackward(model::CompiledFlowModel)       = (x) -> backward(model, x)
getjacobian(model::CompiledFlowModel)       = (x) -> jacobian(model, x)
getinv_jacobian(model::CompiledFlowModel)   = (x) -> inv_jacobian(model, x)

# custom Base function for the FlowModel structure
eltype(model::CompiledFlowModel) = promote_type(map(eltype, model.layers)...)
length(model::CompiledFlowModel) = length(model.layers)

# fetch number of parameters of model
nr_params(model::FlowModel)         = nr_params(getlayers(model))
nr_params(model::CompiledFlowModel) = nr_params(getlayers(model))
nr_params(layers::T) where { T <: NTuple{N,AbstractLayer} where N} = nr_params(first(layers)) + nr_params(Base.tail(layers))
nr_params(layers::Tuple{}) = return 0

# forward pass through the Flow model
function _forward(model::CompiledFlowModel, input::AbstractVector{T}) where { T <: Real }

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for result
    output = zeros(Ti, size(input))

    # perform calculations
    forward!(output, model, input)

    # return result
    return output

end

# when calling forward, redirect to _forward
forward(model::CompiledFlowModel, input::AbstractVector{ <: Real }) = _forward(model, input)

# for broadcasting over forward, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(forward), model::CompiledFlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_forward, Ref(model), input)

# inplace forward pass through the Flow model
function forward!(output::AbstractVector{ <: Real }, model::CompiledFlowModel, input::AbstractVector{T}) where { T <: Real }

    # fetch layers
    layers = getlayers(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # decouple changing input from actual input
    input_new = zeros(Ti, size(input))
    input_new .= input

    # perform forward pass over all layers
    forward!(output, layers, input_new)

end

# inplace forward pass through a tuple of layers
function forward!(output::AbstractVector{ <: Real }, layers::T, input::AbstractVector{ <: Real }) where { T <: NTuple{N,AbstractLayer} where N}

    # perform pass over first layer
    forward!(output, first(layers), input)

    # update intermediate input
    input .= output

    # perform pass over all remaining layers
    forward!(output, Base.tail(layers), input)

end

# when no layers are left, stop the inplace recursion
forward!(output::AbstractVector{ <: Real }, layers::Tuple{}, input::AbstractVector{ <: Real }) = return

# backward pass through the Flow model
function _backward(model::CompiledFlowModel, output::AbstractVector{T}) where { T <: Real }

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for result
    input = zeros(Ti, size(output))

    # perform calculations
    backward!(input, model, output)

    # return result
    return input

end

# when calling backward, redirect to _backward
backward(model::CompiledFlowModel, output::AbstractVector{ <: Real }) = _backward(model, output)

# for broadcasting over backward, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(backward), model::CompiledFlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_backward, Ref(model), output)

# inplace backward pass through the Flow model
function backward!(input::AbstractVector{ <: Real }, model::CompiledFlowModel, output::AbstractVector{T}) where { T <: Real }

    # fetch layers
    layers = getlayers(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # decouple changing output from actual output
    output_new = zeros(Ti, size(output))
    output_new .= output

    # perform backward pass over all layers
    backward!(input, layers, output_new)

end

# inplace backward pass through a tuple of layers
function backward!(input::AbstractVector{ <: Real }, layers::T, output::AbstractVector{ <: Real }) where { T <: NTuple{N,AbstractLayer} where N}

    # perform pass over last layer
    backward!(input, last(layers), output)

    # update intermediate output
    output .= input

    # perform pass over all remaining layers
    backward!(input, Base.front(layers), output)

end

# when no layers are left, stop the inplace recursion
backward!(input::AbstractVector{ <: Real }, layers::Tuple{}, output::AbstractVector{ <: Real }) = return

# joint forward-jacobian of the Flow model
function _forward_jacobian(model::CompiledFlowModel, input::AbstractVector{T}) where { T <: Real }

    # fetch layers
    dim    = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for output and jacobian
    output = zeros(Ti, dim)
    J = zeros(Ti, dim, dim)
    for k = 1:dim
        J[k,k] = one(Ti)
    end

    # calculate jacobian
    forward_jacobian!(output, J, model, input)

    # return result
    return output, J

end

# when calling forward_jacobian, redirect to _forward_jacobian
forward_jacobian(model::CompiledFlowModel, input::AbstractVector{ <: Real }) = _forward_jacobian(model, input)

# for broadcasting over forward_jacobian, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(forward_jacobian), model::CompiledFlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_jacobian, Ref(model), input)

# inplace forward_jacobian of the Flow model
function forward_jacobian!(output::AbstractVector{ <: Real }, J::AbstractMatrix{T}, model::CompiledFlowModel, input::AbstractVector{ <: Real }) where { T <: Real }

    # fetch layers
    layers = getlayers(model)
    dim = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for intermediate output and jacobian
    input_new = zeros(Ti, dim)
    input_new .= input
    J_new = zeros(Ti, dim, dim)
    J_old = zeros(Ti, dim, dim)
    for k = 1:dim
        J_old[k,k] = one(Ti)
    end

    # calculate forward_jacobian over all layers
    forward_jacobian!(J, J_new, J_old, output, input_new, layers)

end

# inplace forward_jacobian calculation for a tuple of layers
function forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layers::T) where { T <: NTuple{N,AbstractLayer} where N}

    # perform pass through first layers
    forward_jacobian!(J, J_new, J_old, output, input_new, first(layers))

    # update unless we are on the last layer
    if length(layers) > 1

        # update intermediate input
        input_new .= output

        # update old jacobian
        J_old .= J

    end

    # calculate forward_jacobian of remaining layers
    forward_jacobian!(J, J_new, J_old, output, input_new, Base.tail(layers))

end

# specialized methods
function forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layer::PermutationLayer)
    # calculate new output
    forward!(output, layer, input_new)

    # calculate total jacobian
    mul!(J, jacobian(layer, input_new), J_old)
end

# standard method
function forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layer::AbstractLayer)
    # calculate new jacobian
    forward_jacobian!(output, J_new, layer, input_new)

    # calculate total jacobian
    mul!(J, J_new, J_old)
end

# when no layers are left, stop the inplace recursion
forward_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, output::AbstractVector{ <: Real }, input_new::AbstractVector{ <: Real }, layers::Tuple{}) = return

# joing backward inverse jacobian of the Flow model
function _backward_inv_jacobian(model::CompiledFlowModel, output::AbstractVector{T}) where { T <: Real }

    # fetch layers
    dim    = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for jacobian
    J = zeros(Ti, dim, dim)
    input = zeros(Ti, dim)
    for k = 1:dim
        J[k,k] = one(Ti)
    end

    # calculate jacobian
    backward_inv_jacobian!(input, J, model, output)

    # return result
    return input, J

end

# when calling backward inverse jacobian, redirect to _backward_inv_jacobian
backward_inv_jacobian(model::CompiledFlowModel, output::AbstractVector{ <: Real}) = _backward_inv_jacobian(model, output)

# for broadcasting over backward inverse jacobian, fix the model for multiple inputs
Broadcast.broadcasted(::typeof(backward_inv_jacobian), model::CompiledFlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_backward_inv_jacobian, Ref(model), output)

# inplace inverse backward jacobian of the Flow model
function backward_inv_jacobian!(input::AbstractVector{ <: Real }, J::AbstractMatrix{T}, model::CompiledFlowModel, output::AbstractVector{ <: Real }) where { T <: Real }

    # fetch layers
    layers = getlayers(model)
    dim = getdim(model)

    # promote type for allocating output
    Ti = promote_type(eltype(model), T)

    # allocate space for intermediate output and jacobian
    output_new = zeros(Ti, dim)
    output_new .= output
    J_new = zeros(Ti, dim, dim)
    J_old = zeros(Ti, dim, dim)
    for k = 1:dim
        J_old[k,k] = one(Ti)
    end

    # calculate backward inverse jacobian over all layers
    backward_inv_jacobian!(J, J_new, J_old, input, output_new, layers)

end

# inplace backward inverse jacobian calculation for a tuple of layers
function backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layers::T) where { T <: NTuple{N,AbstractLayer} where N}

    # perform backward pass through layers
    backward_inv_jacobian!(J, J_new, J_old, input, output_new, last(layers))

    # update unless we are on the last/first layer
    if length(layers) > 1

        # update intermediate output
        output_new .= input

        # update old jacobian
        J_old .= J

    end

    # calculate backward inv jacobian of remaining layers
    backward_inv_jacobian!(J, J_new, J_old, input, output_new, Base.front(layers))

end

# specialized methods
function backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layer::PermutationLayer)
    # perform forward pass over last layer
    backward!(input, layer, output_new)

    # calculate total jacobian
    mul!(J, inv_jacobian(layer, output_new), J_old)
end

# standard method
function backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layer::AbstractLayer)
    # perform forward pass over last layer
    backward_inv_jacobian!(input, J_new, layer, output_new)

    # calculate total jacobian
    mul!(J, J_new, J_old)
end

# when no layers are left, stop the inplace recursion
backward_inv_jacobian!(J::AbstractMatrix{ <: Real }, J_new::AbstractMatrix{ <: Real }, J_old::AbstractMatrix{ <: Real }, input::AbstractVector{ <: Real }, output_new::AbstractVector{ <: Real }, layers::Tuple{}) = return

# specify jacobian and inv_jacobian functions based on joint functions
_jacobian(model::CompiledFlowModel, input::AbstractVector{ <: Real })      = forward_jacobian(model, input)[2]
jacobian(model::CompiledFlowModel, input::AbstractVector{ <: Real })       = _jacobian(model, input)
Broadcast.broadcasted(::typeof(jacobian), model::CompiledFlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_jacobian, Ref(model), input)
_inv_jacobian(model::CompiledFlowModel, output::AbstractVector{ <: Real }) = backward_inv_jacobian(model, output)[2]
inv_jacobian(model::CompiledFlowModel, output::AbstractVector{ <: Real })  = _inv_jacobian(model, output)
Broadcast.broadcasted(::typeof(inv_jacobian), model::CompiledFlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = broadcast(_inv_jacobian, Ref(model), output)
function jacobian!(J::AbstractMatrix{T1}, model::CompiledFlowModel, input::AbstractVector{T2}) where { T1 <: Real, T2 <: Real }

    # fetch dimension
    dim    = getdim(model)

    # promote type for allocating output
    T = promote_type(eltype(model), T1, T2)

    # allocate space for output and jacobian
    output = zeros(T, dim)
    J .= zero(T1)
    for k = 1:dim
        J[k,k] = one(T1)
    end

    # calculate jacobian
    forward_jacobian!(output, J, model, input)

end
function inv_jacobian!(J::AbstractMatrix{T1}, model::CompiledFlowModel, output::AbstractVector{T2}) where { T1 <: Real, T2 <: Real }

    # fetch dimension
    dim    = getdim(model)

    # promote type for allocating output
    T = promote_type(eltype(model), T1, T2)

    # allocate space for input and jacobian
    input = zeros(T, dim)
    J .= zero(T1)
    for k = 1:dim
        J[k,k] = one(T1)
    end

    # calculate jacobian
    backward_inv_jacobian!(input, J, model, output)

end

# fallback joint functions over layers
function forward_jacobian!(output::AbstractVector{ <: Real }, J_new::AbstractMatrix{ <: Real }, layer::AbstractLayer, input::AbstractVector{ <: Real })
    forward!(output, layer, input)
    jacobian!(J_new, layer, input)
end
function backward_inv_jacobian!(input::AbstractVector{ <: Real }, J_new::AbstractMatrix{ <: Real }, layer::AbstractLayer, output::AbstractVector{ <: Real })
    backward!(input, layer, output)
    inv_jacobian!(J_new, layer, output)
end

# extra utility functions
det_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})           = det(jacobian(model, input))
absdet_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})        = abs(det_jacobian(model, input))
logdet_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})        = logdet(jacobian(model, input))
logabsdet_jacobian(model::CompiledFlowModel, input::AbstractVector{<: Real})     = logabsdet(jacobian(model, input))

detinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real})       = det(inv_jacobian(model, output))
absdetinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real})    = abs(detinv_jacobian(model, output))
logdetinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real})    = logdet(inv_jacobian(model, output))
logabsdetinv_jacobian(model::CompiledFlowModel, output::AbstractVector{<: Real}) = logabsdet(inv_jacobian(model, output))

# throw an error when the model has not yet been compiled
_forward(model::FlowModel, input::AbstractVector{ <: Real })                                                        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
forward(model::FlowModel, input::AbstractVector{ <: Real })                                                         = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
Broadcast.broadcasted(::typeof(forward), model::FlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } })   = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
forward!(output::AbstractVector{ <: Real }, model::FlowModel, input::AbstractVector{ <: Real })                     = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
_backward(model::FlowModel, output::AbstractVector{ <: Real })                                                      = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
backward(model::FlowModel, output::AbstractVector{ <: Real })                                                       = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
Broadcast.broadcasted(::typeof(backward), model::FlowModel, output::AbstractVector{ <: AbstractVector{ <: Real } }) = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
backward!(input::AbstractVector{ <: Real }, model::FlowModel, output::AbstractVector{ <: Real })                    = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
jacobian!(J_new::AbstractMatrix{ <: Real }, model::FlowModel, input::AbstractVector{ <: Real })                     = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
jacobian(model::FlowModel, input::AbstractVector{ <: Real })                                                        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
Broadcast.broadcasted(::typeof(jacobian), model::FlowModel, input::AbstractVector{ <: AbstractVector{ <: Real } })  = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
_jacobian(model::FlowModel, input::AbstractVector{ <: Real })                                                       = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
_inv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })                                                  = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
inv_jacobian!(J_new::AbstractMatrix{ <: Real }, model::FlowModel, output::AbstractVector{ <: Real })                = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))

det_jacobian(model::FlowModel, input::AbstractVector{ <: Real })           = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
absdet_jacobian(model::FlowModel, input::AbstractVector{ <: Real })        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logdet_jacobian(model::FlowModel, input::AbstractVector{ <: Real })        = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logabsdet_jacobian(model::FlowModel, input::AbstractVector{ <: Real })     = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))

detinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })       = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
absdetinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })    = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logdetinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real })    = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))
logabsdetinv_jacobian(model::FlowModel, output::AbstractVector{ <: Real }) = throw(ArgumentError("Please first compile your model using `compiled_model = compile(model)`."))


module POMDPModelTools

using POMDPs
using Random
using LinearAlgebra
using SparseArrays
using UnicodePlots
import CommonRLInterface
using Tricks: static_hasmethod

import POMDPs: actions, actionindex
import POMDPs: states, stateindex
import POMDPs: observations, obsindex
import POMDPs: initialstate, isterminal, discount
import POMDPs: implemented
import Distributions: pdf, mode, mean, support
import Random: rand, rand!
import Statistics: mean
import Base: ==

import POMDPLinter: @POMDP_require

export
    render
include("visualization.jl")

# info interface
export
    action_info,
    solve_info,
    update_info
include("info.jl")

export
    ordered_states,
    ordered_actions,
    ordered_observations
include("ordered_spaces.jl")

export
    TerminalState,
    terminalstate
include("terminal_state.jl")

export GenerativeBeliefMDP
include("generative_belief_mdp.jl")

export FullyObservablePOMDP
include("fully_observable_pomdp.jl")

export UnderlyingMDP
include("underlying_mdp.jl")

export obs_weight
include("obs_weight.jl")

export
    probability_check,
    obs_prob_consistency_check,
    trans_prob_consistency_check

export
    weighted_iterator
include("distributions/weighted_iteration.jl")

export
    SparseCat
include("distributions/sparse_cat.jl")

export
    BoolDistribution
include("distributions/bool.jl")

export
    Deterministic
include("distributions/deterministic.jl")

export
    Uniform,
    UnsafeUniform
include("distributions/uniform.jl")

export
    ImplicitDistribution
include("distributions/implicit.jl")

export
    StateActionReward,
    FunctionSAR,
    LazyCachedSAR
include("state_action_reward.jl")

# convenient implementations
include("convenient_implementations.jl")

export
    evaluate
include("policy_evaluation.jl")

export
    showdistribution
include("distributions/pretty_printing.jl")

export
    SparseTabularMDP,
    SparseTabularPOMDP,
    transition_matrix,
    reward_vector,
    observation_matrix,
    reward_matrix,
    observation_matrices
include("sparse_tabular.jl")

export
    transition_matrices,
    reward_vectors
include("matrices.jl")

export
    MDPCommonRLEnv,
    POMDPCommonRLEnv
include("common_rl/to_env.jl")

export
    RLEnvMDP,
    RLEnvPOMDP,
    OpaqueRLEnvMDP,
    OpaqueRLEnvPOMDP
include("common_rl/from_env.jl")

export
    add_infonode
include("deprecated.jl")

end # module


# broadcasting

using Base.Broadcast: BroadcastStyle, Broadcasted

struct ROCArrayStyle{N} <: AbstractGPUArrayStyle{N} end
ROCArrayStyle(::Val{N}) where N = ROCArrayStyle{N}()
ROCArrayStyle{M}(::Val{N}) where {N,M} = ROCArrayStyle{N}()

BroadcastStyle(::Type{ROCArray{T,N}}) where {T,N} = ROCArrayStyle{N}()

Base.similar(bc::Broadcasted{ROCArrayStyle{N}}, ::Type{T}) where {N,T} =
    similar(ROCArray{T}, axes(bc))

Base.similar(bc::Broadcasted{ROCArrayStyle{N}}, ::Type{T}, dims) where {N,T} =
    ROCArray{T}(undef, dims)

## replace base functions with libdevice alternatives

rocfunc(f) = f
rocfunc(::Type{T}) where T = (x...) -> T(x...) # broadcasting type ctors isn't GPU compatible

Broadcast.broadcasted(::ROCArrayStyle{N}, f, args...) where {N} =
    Broadcasted{ROCArrayStyle{N}}(rocfunc(f), args, nothing)

const device_intrinsics = :[
    cos, cospi, sin, sinpi, tan, acos, asin, atan,
    cosh, sinh, tanh, acosh, asinh, atanh, angle,
    log, log10, log1p, log2, logb, ilogb,
    exp, exp2, exp10, expm1, ldexp,
    erf, erfinv, erfc, erfcinv, erfcx,
    brev, clz, ffs, byte_perm, popc,
    isfinite, isinf, isnan, nearbyint,
    nextafter, signbit, copysign, abs,
    sqrt, rsqrt, cbrt, rcbrt, pow,
    ceil, floor, saturate,
    lgamma, tgamma,
    j0, j1, jn, y0, y1, yn,
    normcdf, normcdfinv, hypot,
    fma, sad, dim, mul24, mul64hi, hadd, rhadd, scalbn].args

for f in device_intrinsics
    isdefined(Base, f) || continue
    @eval rocfunc(::typeof(Base.$f)) = $f
end

# broadcast ^

rocliteral_pow(::typeof(^), x::T, ::Val{0}) where {T<:Real} = one(x)
rocliteral_pow(::typeof(^), x::T, ::Val{1}) where {T<:Real} = x
rocliteral_pow(::typeof(^), x::T, ::Val{2}) where {T<:Real} = x * x
rocliteral_pow(::typeof(^), x::T, ::Val{3}) where {T<:Real} = x * x * x
rocliteral_pow(::typeof(^), x::T, ::Val{p}) where {T<:Real,p} = pow(x, Int32(p))

rocfunc(::typeof(Base.literal_pow)) = rocliteral_pow
rocfunc(::typeof(Base.:(^))) = pow

using MacroTools

const _rocfuncs = [copy(device_intrinsics); :^]
rocfuncs() = (global _rocfuncs; _rocfuncs)

_rocint(x::Int) = Int32(x)
_rocint(x::Expr) = x.head == :call && x.args[1] == :Int32 && x.args[2] isa Int ? Int32(x.args[2]) : x
_rocint(x) = x

function _rocpowliteral(x::Expr)
    if x.head == :call && x.args[1] == :(AMDGPU.rocfunc(^)) && x.args[3] isa Int32
        num = x.args[3]
        if 0 <= num <= 3
            sym = gensym(:x)
            new_x = Expr(:block, :($sym = $(x.args[2])))

            if iszero(num)
                push!(new_x.args, :(one($sym)))
            else
                unroll = Expr(:call, :*)
                for x = one(num):num
                    push!(unroll.args, sym)
                end
                push!(new_x.args, unroll)
            end

            x = new_x
        end
    end
    x
end
_rocpowliteral(x) = x

function replace_device(ex)
    global _rocfuncs
    MacroTools.postwalk(ex) do x
        x = x in _rocfuncs ? :(AMDGPU.rocfunc($x)) : x
        x = _rocint(x)
        x = _rocpowliteral(x)
        x
    end
end

macro rocfunc(ex)
    global _rocfuncs
    def = MacroTools.splitdef(ex)
    f = def[:name]
    def[:name] = Symbol(:cu, f)
    def[:body] = replace_device(def[:body])
    push!(_rocfuncs, f)
    quote
        $(esc(MacroTools.combinedef(def)))
        AMDGPU.rocfunc(::typeof($(esc(f)))) = $(esc(def[:name]))
    end
end