text:
### A Pluto.jl notebook ###
# v0.14.8

using Markdown
using InteractiveUtils

# ╔═╡ 4e019dd0-9a70-4034-915b-baf59e3416e5
using Fermi

# ╔═╡ e45bf90a-9cfe-46e8-bb0f-1c42723f1389
using TensorOperations

# ╔═╡ 819f62ec-da7c-11eb-300f-333f486858d2
html"""

  JuliaCon 2021

 Gustavo Aroeira

Center for Computational Quantum Chemistry

University of Georgia, Athens, GA

United States of America

"""

# ╔═╡ e1bd2ce0-a958-4460-b48f-620db0743e45
md"""
## 1. Quantum Chemistry
The main goal in quantum chemistry is to solve the electronic Schrödinger equation for many-electron systems

$\hat{H}|\Psi\rangle = E|\Psi\rangle$

A first approximation to the fermionic wave function, known as *mean-field*, is represented as a Slater determinant

$|\Phi_0\rangle = \mathcal{A}[\phi_1(\mathbf{x_1})\phi_2(\mathbf{x_2})\phi_3(\mathbf{x_3})...\phi_N(\mathbf{x_N})]$

where $\mathcal{A}$ is the antisymmetrizer operator and $\phi_i$ are spin-orbitals, that is, one-electron wave functions used to build the many-electron *ansatz*. Using the *basis set approximation*, we represent the orbitals as an expansion of $M$ pre-selected basis functions

$\phi_i = \sum_\mu^M C_{\mu i} \chi_\mu$

This is known as the **L**inear **C**ombination of **A**tomic **O**rbitals (LCAO) as the basis functions are designed to resemble atomic orbitals. Finally, more precise, i.e. correlated, wave functions are built as combinations of different Slater determinants (they differ by the orbitals within each one).

$|\Psi\rangle = \sum_i c_i|\Phi_i\rangle$

or, in cluster notation

$|\Psi\rangle =  e^{T}|\Phi_0\rangle$

**Fermi** is a quantum chemistry package in Julia specialized in wave function methods. Currently we have implemented the following methods

- Restricted Hartree-Fock
- Møller-Plesset Perturbation Theory
- Coupled Cluster Theory

Moreover, Fermi aims to serve as development platform for molecular quantum mechanics in Julia
"""

# ╔═╡ ff5a766b-716b-47dc-a94f-1f9c3f74ac8b
md"""
## 2. Molecular Integrals

Basic quantities required in all quantum chemistry computations

###### Overlap (S)

$S_{ij} = \int \phi_i^* \phi_j d\tau$

###### Kinetic (T)

$T_{ij} = \int \phi_i^* \left( -\frac{1}{2} \nabla^2 \right) \phi_j d\tau$

###### Nuclear (V)

$V_{ij} = \sum_k \int \phi_i^* \frac{Z_k}{r} \phi_j d\tau$

###### Eletron repulsion integral (ERI)

$G_{ijkl} = \int \phi_i^* \phi_j^* \frac{1}{r_{12}} \phi_k \phi_l d\tau$

Fermi utilizes the library [libcint](https://github.com/sunqm/libcint) to obtain molecular integrals.
"""

# ╔═╡ 67892086-7643-40a0-9e5e-eb56607412de
begin
@set {
	diis false
	oda false
}
nothing
end

# ╔═╡ e8585228-0be6-41dd-8434-045ed69d578c
md"""
The module `Fermi.Libcint` has minimal wraps around libcint functions
"""

# ╔═╡ 3a325265-7699-4f99-a850-207ec00f9c26
names(Fermi.Libcint)

# ╔═╡ 8dced24e-6cc6-46aa-bd50-658907190987
md"""
A more convinient way is to use the `IntegralHelper`.

First, define a molecule and basis set
"""

# ╔═╡ 01228770-9ad9-4634-9aba-afc798ae3e8c
@molecule {
	H 0.0 0.0 0.0
	H 0.76 0.0 0.0
}

# ╔═╡ 597aa086-dd0c-4ef9-b0ef-ab294b89b3a2
@set basis sto-3g

# ╔═╡ bd1dca4c-f126-4353-b09a-e2a641712a09
I = Fermi.Integrals.IntegralHelper()

# ╔═╡ c37cef53-ddfa-4bce-9328-21cde32bf8da
I["S"]

# ╔═╡ 73551bfb-c7c6-4f5b-a931-92c23b7feafd
md"""
## 3. Multiple dispatch in Fermi

`IntegralHelper` guides how amplitudes are updated in a CCSD computation
"""

# ╔═╡ b03d9e3f-8ef5-4e48-b2aa-f6b4295ac228
html"""

"""

# ╔═╡ fb42172a-33df-4b03-88d5-6db7224f132b
md"""
## 4. Tensor Contractions

In methods such as Coupled Cluster, tensor contractions appears in terms like

$t_{ij}^{ab} \leftarrow -2t_{ik}^{ac}\cdot t_{lj}^{bd}\cdot v_{kc}^{ld}$

The package [TensorOperations.jl](https://github.com/Jutho/TensorOperations.jl) is used to handle expressions of this type.

`FermiMDArray` is a thin wrap around a Julia array that allow us to use [TBLIS](https://github.com/devinamatthews/tblis) as a tensor contraction tool.
"""

# ╔═╡ d5942379-2e27-43ca-920e-889a64d32269
begin
	o_size = 10
	v_size = 100
	T2 = FermiMDrand(o_size, o_size, v_size, v_size)
	newT2 = FermiMDrand(o_size, o_size, v_size, v_size)
	Vovov = FermiMDrand(o_size, v_size, o_size, v_size)
	size(Vovov)
end

# ╔═╡ d0c5639e-49f7-45c0-a19e-d1da88b81ce0
t = @elapsed begin
	@tensoropt newT2[i,j,a,b] = -2*T2[i,k,a,c]*T2[l,j,b,d]*Vovov[k,c,l,d]
end

# ╔═╡ e89556da-e726-47bf-ae91-4312b24095b1
@set tblis true

# ╔═╡ Cell order:
# ╟─819f62ec-da7c-11eb-300f-333f486858d2
# ╟─e1bd2ce0-a958-4460-b48f-620db0743e45
# ╟─ff5a766b-716b-47dc-a94f-1f9c3f74ac8b
# ╟─67892086-7643-40a0-9e5e-eb56607412de
# ╠═4e019dd0-9a70-4034-915b-baf59e3416e5
# ╟─e8585228-0be6-41dd-8434-045ed69d578c
# ╠═3a325265-7699-4f99-a850-207ec00f9c26
# ╟─8dced24e-6cc6-46aa-bd50-658907190987
# ╠═01228770-9ad9-4634-9aba-afc798ae3e8c
# ╠═597aa086-dd0c-4ef9-b0ef-ab294b89b3a2
# ╠═bd1dca4c-f126-4353-b09a-e2a641712a09
# ╠═c37cef53-ddfa-4bce-9328-21cde32bf8da
# ╟─73551bfb-c7c6-4f5b-a931-92c23b7feafd
# ╟─b03d9e3f-8ef5-4e48-b2aa-f6b4295ac228
# ╟─fb42172a-33df-4b03-88d5-6db7224f132b
# ╠═e45bf90a-9cfe-46e8-bb0f-1c42723f1389
# ╠═d5942379-2e27-43ca-920e-889a64d32269
# ╠═d0c5639e-49f7-45c0-a19e-d1da88b81ce0
# ╠═e89556da-e726-47bf-ae91-4312b24095b1


text:
module Multinomials

export +, *, -, ^

include("structures.jl")
include("algorithms.jl")

end

ext: jl


text:
@generated function saturated_add(x::I, y::I) where {I <: IntegerTypesHW}
    typ = "i$(8sizeof(I))"
    s = I <: Signed ? 's' : 'u'
    f = "@llvm.$(s)add.sat.$typ"
    decl = "declare $typ $f($typ, $typ)"
    instrs = """
        %res = call $typ $f($typ %0, $typ %1)
        ret $typ %res
    """
    llvmcall_expr(decl, instrs, I, :(Tuple{$I,$I}), typ, [typ,typ], [:x,:y])
end
@generated function saturated_add(x::Vec{W,I}, y::Vec{W,I}) where {W,I}

end

@eval @inline function assume(b::Bool)
    $(llvmcall_expr("declare void @llvm.assume(i1)", "%b = trunc i8 %0 to i1\ncall void @llvm.assume(i1 %b)\nret void", :Cvoid, :(Tuple{Bool}), "void", ["i8"], [:b]))
end

@eval @inline function expect(b::Bool)
    $(llvmcall_expr("declare i1 @llvm.expect.i1(i1, i1)", """
    %b = trunc i8 %0 to i1
    %actual = call i1 @llvm.expect.i1(i1 %b, i1 true)
    %byte = zext i1 %actual to i8
    ret i8 %byte""", :Bool, :(Tuple{Bool}), "i8", ["i8"], [:b]))
end
@generated function expect(i::I, ::Val{N}) where {I <: Integer, N}
    ityp = 'i' * string(8sizeof(I))
    llvmcall_expr("declare i1 @llvm.expect.$ityp($ityp, i1)", """
    %actual = call $ityp @llvm.expect.$ityp($ityp %0, $ityp $N)
    ret $ityp %actual""", I, :(Tuple{$I}), ityp, [ityp], [:i])
end

# for (op,f) ∈ [("abs",:abs)]
# end
if Base.libllvm_version ≥ v"12"
    for (op,f,S) ∈ [("smax",:max,:Signed),("smin",:min,:Signed),("umax",:max,:Unsigned),("umin",:min,:Unsigned)]
        @eval @generated Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W, T <: $S} = llvmcall_expr($op, W, T, (W, W), (T, T))
    end
else
    @inline Base.max(v1::Vec{W,<:Integer}, v2::Vec{W,<:Integer}) where {W} = ifelse(v1 > v2, v1, v2)
    @inline Base.min(v1::Vec{W,<:Integer}, v2::Vec{W,<:Integer}) where {W} = ifelse(v1 < v2, v1, v2)
end
# for T ∈ [Float32, Float64]
#     W = 2
#     while W * sizeof(T) ≤ REGISTER_SIZE
# floating point
for (op,f) ∈ [("sqrt",:sqrt),("fabs",:abs),("floor",:floor),("ceil",:ceil),("trunc",:trunc),("nearbyint",:round)
              ]
    # @eval @generated Base.$f(v1::Vec{W,T}) where {W, T <: Union{Float32,Float64}} = llvmcall_expr($op, W, T, (W,), (T,), "nsz arcp contract afn reassoc")
    @eval @generated Base.$f(v1::Vec{W,T}) where {W, T <: Union{Float32,Float64}} = llvmcall_expr($op, W, T, (W,), (T,), "fast")
end

@generated function Base.round(::Type{Int64}, v1::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
    llvmcall_expr("lrint", W, Int64, (W,), (T,), "nsz arcp contract afn reassoc")
end
@generated function Base.round(::Type{Int32}, v1::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
    llvmcall_expr("lrint", W, Int32, (W,), (T,), "nsz arcp contract afn reassoc")
end
@inline Base.trunc(::Type{I}, v::AbstractSIMD{W,T}) where {W, I<:IntegerTypesHW, T <: NativeTypes} = convert(I, v)

# """
#    setbits(x::Unsigned, y::Unsigned, mask::Unsigned)

# If you have AVX512, setbits of vector-arguments will select bits according to mask `m`, selecting from `y` if 0 and from `x` if `1`.
# For scalar arguments, or vector arguments without AVX512, `setbits` requires the additional restrictions on `y` that all bits for
# which `m` is 1, `y` must be 0.
# That is for scalar arguments or vector arguments without AVX512, it requires the restriction that
# ((y ⊻ m) & m) == m
# """
# @inline setbits(x, y, m) = (x & m) | y

"""
   bitselect(m::Unsigned, x::Unsigned, y::Unsigned)

If you have AVX512, setbits of vector-arguments will select bits according to mask `m`, selecting from `x` if 0 and from `y` if `1`.
For scalar arguments, or vector arguments without AVX512, `setbits` requires the additional restrictions on `y` that all bits for
which `m` is 1, `y` must be 0.
That is for scalar arguments or vector arguments without AVX512, it requires the restriction that
((y ⊻ m) & m) == m
"""
@inline bitselect(m, x, y) = ((~m) & x) | (m & y)
if AVX512F
    # AVX512 lets us use 1 instruction instead of 2 dependent instructions to set bits
    @generated function vpternlog(m::Vec{W,UInt64}, x::Vec{W,UInt64}, y::Vec{W,UInt64}, ::Val{L}) where {W, L}
        @assert W ∈ (2,4,8)
        bits = 64W
        decl64 = "declare <$W x i64> @llvm.x86.avx512.mask.pternlog.q.$(bits)(<$W x i64>, <$W x i64>, <$W x i64>, i32, i8)"
        instr64 = """
            %res = call <$W x i64> @llvm.x86.avx512.mask.pternlog.q.$(bits)(<$W x i64> %0, <$W x i64> %1, <$W x i64> %2, i32 $L, i8 -1)
            ret <$W x i64> %res
        """
        arg_syms = [:(data(m)), :(data(x)), :(data(y))]
        llvmcall_expr(decl64, instr64, :(_Vec{$W,UInt64}), :(Tuple{_Vec{$W,UInt64},_Vec{$W,UInt64},_Vec{$W,UInt64}}), "<$W x i64>", ["<$W x i64>", "<$W x i64>", "<$W x i64>"], arg_syms)
    end
    @generated function vpternlog(m::Vec{W,UInt32}, x::Vec{W,UInt32}, y::Vec{W,UInt32}, ::Val{L}) where {W, L}
        if W ∉ (4,8,16)
            return Expr(:block, Expr(:meta, :inline), :(((~m) & x) | (m & y)))
        end
        bits = 32W
        decl32 = "declare <$W x i32> @llvm.x86.avx512.mask.pternlog.d.$(bits)(<$W x i32>, <$W x i32>, <$W x i32>, i32, i16)"
        instr32 = """
            %res = call <$W x i32> @llvm.x86.avx512.mask.pternlog.d.$(bits)(<$W x i32> %0, <$W x i32> %1, <$W x i32> %2, i32 $L, i16 -1)
            ret <$W x i32> %res
        """
        arg_syms = [:(data(m)), :(data(x)), :(data(y))]
        llvmcall_expr(decl32, instr32, :(_Vec{$W,UInt32}), :(Tuple{_Vec{$W,UInt32},_Vec{$W,UInt32},_Vec{$W,UInt32}}), "<$W x i32>", ["<$W x i32>", "<$W x i32>", "<$W x i32>"], arg_syms)
    end
    # @eval @generated function setbits(x::Vec{W,T}, y::Vec{W,T}, m::Vec{W,T}) where {W,T <: Union{UInt32,UInt64}}
    #     ex = if W*sizeof(T) ∈ (16,32,64)
    #         :(vpternlog(x, y, m, Val{216}()))
    #     else
    #         :((x & m) | y)
    #     end
    #     Expr(:block, Expr(:meta, :inline), ex)
    # end
    @generated function bitselect(m::Vec{W,T}, x::Vec{W,T}, y::Vec{W,T}) where {W,T <: Union{UInt32,UInt64}}
        ex = if W*sizeof(T) ∈ (16,32,64)
            :(vpternlog(m, x, y, Val{172}()))
        else
            :(((~m) & x) | (m & y))
        end
        Expr(:block, Expr(:meta, :inline), ex)
    end
    @inline Base.copysign(v1::Vec{W,Float64}, v2::Vec{W,Float64}) where {W} = reinterpret(Float64, bitselect(Vec{W,UInt64}(0x8000000000000000), reinterpret(UInt64, v1), reinterpret(UInt64, v2)))
    @inline Base.copysign(v1::Vec{W,Float32}, v2::Vec{W,Float32}) where {W} = reinterpret(Float32, bitselect(Vec{W,UInt32}(0x80000000), reinterpret(UInt32, v1), reinterpret(UInt32, v2)))
end

for (op,f) ∈ [("minnum",:min),("maxnum",:max),("copysign",:copysign),
              ]
    @eval @generated function Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
        llvmcall_expr($op, W, T, (W for _ in 1:2), (T for _ in 1:2), "nsz arcp contract afn reassoc")
    end
end
@inline _signbit(v::Vec{W, I}) where {W, I<:Signed} = v & Vec{W,I}(typemin(I))
@inline Base.copysign(v1::Vec{W,I}, v2::Vec{W,I}) where {W, I <: Signed} = ifelse(_signbit(v1) == _signbit(v2), v1, -v1)

@inline Base.copysign(x::Float32, v::Vec{W}) where {W} = copysign(vbroadcast(Val{W}(), x), v)
@inline Base.copysign(x::Float64, v::Vec{W}) where {W} = copysign(vbroadcast(Val{W}(), x), v)
@inline Base.copysign(x::Float32, v::VecUnroll{N,W}) where {N,W} = copysign(vbroadcast(Val{W}(), x), v)
@inline Base.copysign(x::Float64, v::VecUnroll{N,W}) where {N,W} = copysign(vbroadcast(Val{W}(), x), v)
@inline Base.copysign(v::Vec, u::VecUnroll) = VecUnroll(fmap(copysign, v, u.data))
@inline Base.copysign(v::Vec{W,T}, x::NativeTypes) where {W,T} = copysign(v, Vec{W,T}(x))
@inline Base.copysign(v1::Vec{W,T}, v2::Vec{W}) where {W,T} = copysign(v1, convert(Vec{W,T}, v2))

# ternary
for (op,f) ∈ [("fma",:fma),("fmuladd",:muladd)]
    @eval @generated function Base.$f(v1::Vec{W,T}, v2::Vec{W,T}, v3::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
        llvmcall_expr($op, W, T, (W for _ in 1:3), (T for _ in 1:3), $(f === :fma ? nothing : "nsz arcp contract afn reassoc"))
    end
end
# floating vector, integer scalar
# @generated function Base.:(^)(v1::Vec{W,T}, v2::Int32) where {W, T <: Union{Float32,Float64}}
#     llvmcall_expr("powi", W, T, (W, 1), (T, Int32), "nsz arcp contract afn reassoc")
# end
for (op,f) ∈ [
    ("experimental.vector.reduce.v2.fadd",:vsum),
    ("experimental.vector.reduce.v2.fmul",:vprod)
]
    @eval @generated function $f(v1::T, v2::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
        llvmcall_expr($op, -1, T, (1, W), (T, T), "nsz arcp contract afn reassoc")
    end
end
vsum(s::T, v::Vec{W,T}) where {W,T} = Base.FastMath.add_fast(s, vsum(v))
vprod(s::T, v::Vec{W,T}) where {W,T} = Base.FastMath.mul_fast(s, vprod(v))
for (op,f) ∈ [
    ("experimental.vector.reduce.fmax",:vmaximum),
    ("experimental.vector.reduce.fmin",:vminimum)
]
    @eval @generated function $f(v1::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
        llvmcall_expr($op, -1, T, (W,), (T,), "nsz arcp contract afn reassoc")
    end
end
for (op,f,S) ∈ [
    ("experimental.vector.reduce.add",:vsum,:Integer),
    ("experimental.vector.reduce.mul",:vprod,:Integer),
    ("experimental.vector.reduce.and",:vall,:Integer),
    ("experimental.vector.reduce.or",:vany,:Integer),
    ("experimental.vector.reduce.xor",:vxorreduce,:Integer),
    ("experimental.vector.reduce.smax",:vmaximum,:Signed),
    ("experimental.vector.reduce.smin",:vminimum,:Signed),
    ("experimental.vector.reduce.umax",:vmaximum,:Unsigned),
    ("experimental.vector.reduce.umin",:vminimum,:Unsigned)
]
    @eval @generated function $f(v1::Vec{W,T}) where {W, T <: $S}
        llvmcall_expr($op, -1, T, (W,), (T,))
    end
end

#         W += W
#     end
# end
@inline vsum(v::Vec{W,T}) where {W,T<:Union{Float32,Float64}} = vsum(-zero(T), v)
@inline vprod(v::Vec{W,T}) where {W,T<:Union{Float32,Float64}} = vprod(one(T), v)
@inline vsum(x, v::Vec{W,T}) where {W,T<:Union{Float32,Float64}} = vsum(convert(T, x), v)
@inline vprod(x, v::Vec{W,T}) where {W,T<:Union{Float32,Float64}} = vprod(convert(T, x), v)

for (f,f_to,op,reduce,twoarg) ∈ [
    (:reduced_add,:reduce_to_add,:+,:vsum,true),(:reduced_prod,:reduce_to_prod,:*,:vprod,true),
    (:reduced_max,:reduce_to_max,:max,:vmaximum,false),(:reduced_min,:reduce_to_min,:min,:vminimum,false)
]
    @eval begin
        @inline $f_to(x::NativeTypes, y::NativeTypes) = x
        @inline $f_to(x::AbstractSIMD, y::AbstractSIMD) = x
        @inline $f_to(x::AbstractSIMD, y::NativeTypes) = $reduce(x)
        @inline $f(x::NativeTypes, y::NativeTypes) = $op(x,y)
        @inline $f(x::AbstractSIMD, y::AbstractSIMD) = $op(x,y)
    end
    if twoarg
        # @eval @inline $f(y::T, x::AbstractSIMD{W,T}) where {W,T} = $reduce(y, x)
        @eval @inline $f(x::AbstractSIMD, y::NativeTypes) = $reduce(y, x)
        # @eval @inline $f(x::AbstractSIMD, y::NativeTypes) = ((y2,x2,r) = @show (y, x, $reduce(y, x)); r)
    else
        # @eval @inline $f(y::T, x::AbstractSIMD{W,T}) where {W,T} = $op(y, $reduce(x))
        @eval @inline $f(x::AbstractSIMD, y::NativeTypes) = $op(y, $reduce(x))
    end
end

@inline roundint(x::Float32) = round(Int32, x)
if AVX512DQ
    @inline roundint(x::Float64) = round(Int, x)
    @inline roundint(v::Vec{W,Float32}) where {W} = round(Int32, v)
    @inline roundint(v::Vec{W,Float64}) where {W} = round(Int64, v)
else
    @inline roundint(x::Float64) = round(Int32, x)
    @inline roundint(v::Vec{W}) where {W} = round(Int32, v)
end
# binary

function count_zeros_func(W, I, op, tf = 1)
    typ = "i$(8sizeof(I))"
    vtyp = "<$W x $typ>"
    instr = "@llvm.$op.v$(W)$(typ)"
    decl = "declare $vtyp $instr($vtyp, i1)"
    instrs = "%res = call $vtyp $instr($vtyp %0, i1 $tf)\nret $vtyp %res"
    llvmcall_expr(decl, instrs, _Vec{W,I}, Tuple{_Vec{W,I}}, vtyp, (vtyp,), (:(data(v)),))
end
# @generated Base.abs(v::Vec{W,I}) where {W, I <: Integer} = count_zeros_func(W, I, "abs", 0)
@generated Base.leading_zeros(v::Vec{W,I}) where {W, I <: Integer} = count_zeros_func(W, I, "ctlz")
@generated Base.trailing_zeros(v::Vec{W,I}) where {W, I <: Integer} = count_zeros_func(W, I, "cttz")

for (op,f) ∈ [("ctpop", :count_ones)]
    @eval @generated Base.$f(v1::Vec{W,T}) where {W,T} = llvmcall_expr($op, W, T, (W,), (T,))
end

for (op,f) ∈ [("fshl",:funnel_shift_left),("fshr",:funnel_shift_right)
              ]
    @eval @generated function $f(v1::Vec{W,T}, v2::Vec{W,T}, v3::Vec{W,T}) where {W,T}
        llvmcall_expr($op, W, T, (W for _ in 1:3), (T for _ in 1:3))
    end
end
funnel_shift_left(a, b, c) = (a << c) | (b >>> (8sizeof(b) - c))
funnel_shift_right(a, b, c) = (a >>> c) | (b << (8sizeof(b) - c))
@inline rotate_left(a, b) = funnel_shift_left(a, a, b)
@inline rotate_right(a, b) = funnel_shift_right(a, a, b)

# for T ∈ [UInt8,UInt16,UInt32,UInt64]
#     bytes = sizeof(T)
#     W = 2
#     while W * bytes ≤ REGISTER_SIZE

#         for (op,f) ∈ [("ctpop", :count_ones)]
#             @eval @inline Base.$f(v1::Vec{$W,$T}) = $(llvmcall_expr(op, W, T, (W,), (T,)))
#         end

#         for (op,f) ∈ [("fshl",:funnel_shift_left),("fshr",:funnel_shift_right)
#                       ]
#             @eval @inline function $f(v1::Vec{$W,$T}, v2::Vec{$W,$T})
#                 $(llvmcall_expr(op, W, T, (W for _ in 1:3), (T for _ in 1:3)))
#             end
#         end

#         W += W
#     end
# end

@inline vfmadd(a, b, c) = muladd(a, b, c)
@inline vfnmadd(a, b, c) = muladd(-a, b, c)
@inline vfmsub(a, b, c) = muladd(a, b, -c)
@inline vfnmsub(a, b, c) = -muladd(a, b, c)

@inline vfmadd231(a, b, c) = vfmadd(a, b, c)
@inline vfnmadd231(a, b, c) = vfnmadd(a, b, c)
@inline vfmsub231(a, b, c) = vfmsub(a, b, c)
@inline vfnmsub231(a, b, c) = vfnmsub(a, b, c)

if FMA
    @eval begin
        @generated function vfmadd231(a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
            if !(ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE) && (W ≥ (T === Float32 ? 4 : 2)))
                return Expr(:block, Expr(:meta, :inline), :(vfmadd(a, b, c)))
            end
            typ = LLVM_TYPES[T]
            suffix = T == Float32 ? "ps" : "pd"
            vfmadd_str = """%res = call <$W x $(typ)> asm "vfmadd231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
                ret <$W x $(typ)> %res"""
            quote
                $(Expr(:meta, :inline))
                Vec(llvmcall($vfmadd_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
            end
        end
        @generated function vfnmadd231(a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
            if !(ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE) && (W ≥ (T === Float32 ? 4 : 2)))
                return Expr(:block, Expr(:meta, :inline), :(vfnmadd(a, b, c)))
            end
            typ = LLVM_TYPES[T]
            suffix = T == Float32 ? "ps" : "pd"
            vfnmadd_str = """%res = call <$W x $(typ)> asm "vfnmadd231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
                ret <$W x $(typ)> %res"""
            quote
                $(Expr(:meta, :inline))
                Vec(llvmcall($vfnmadd_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
            end
        end
        @generated function vfmsub231(a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
            if !(ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE) && (W ≥ (T === Float32 ? 4 : 2)))
                return Expr(:block, Expr(:meta, :inline), :(vfmsub(a, b, c)))
            end
            typ = LLVM_TYPES[T]
            suffix = T == Float32 ? "ps" : "pd"
            vfmsub_str = """%res = call <$W x $(typ)> asm "vfmsub231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
                ret <$W x $(typ)> %res"""
            quote
                $(Expr(:meta, :inline))
                Vec(llvmcall($vfmsub_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
            end
        end
        @generated function vfnmsub231(a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
            if !(ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE) && (W ≥ (T === Float32 ? 4 : 2)))
                return Expr(:block, Expr(:meta, :inline), :(vfnmsub(a, b, c)))
            end
            typ = LLVM_TYPES[T]
            suffix = T == Float32 ? "ps" : "pd"
            vfnmsub_str = """%res = call <$W x $(typ)> asm "vfnmsub231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
                ret <$W x $(typ)> %res"""
            quote
                $(Expr(:meta, :inline))
                Vec(llvmcall($vfnmsub_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
            end
        end
    end
    if AVX512BW
        @eval begin
            @generated function ifelse(::typeof(vfmadd231), m::Mask{W,U}, a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W,U<:Unsigned,T<:Union{Float32,Float64}}
                if !((W ≥ 8) && ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE))
                    return Expr(:block, Expr(:meta, :inline), :(ifelse(vfmadd, m, a, b, c)))
                end
                typ = LLVM_TYPES[T]
                suffix = T == Float32 ? "ps" : "pd"
                vfmaddmask_str = """%res = call <$W x $(typ)> asm "vfmadd231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
                                ret <$W x $(typ)> %res"""
                quote
                    $(Expr(:meta,:inline))
                    Vec(llvmcall($vfmaddmask_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
                end
            end
            @generated function ifelse(::typeof(vfnmadd231), m::Mask{W,U}, a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W,U<:Unsigned,T<:Union{Float32,Float64}}
                if !((W ≥ 8) && ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE))
                    return Expr(:block, Expr(:meta, :inline), :(ifelse(vfmmadd, m, a, b, c)))
                end
                typ = LLVM_TYPES[T]
                suffix = T == Float32 ? "ps" : "pd"
                vfnmaddmask_str = """%res = call <$W x $(typ)> asm "vfnmadd231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
                            ret <$W x $(typ)> %res"""
                quote
                    $(Expr(:meta,:inline))
                    Vec(llvmcall($vfnmaddmask_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
                end
            end
            @generated function ifelse(::typeof(vfmsub231), m::Mask{W,U}, a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W,U<:Unsigned,T<:Union{Float32,Float64}}
                if !((W ≥ 8) && ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE))
                    return Expr(:block, Expr(:meta, :inline), :(ifelse(vfmsub, m, a, b, c)))
                end
                typ = LLVM_TYPES[T]
                suffix = T == Float32 ? "ps" : "pd"
                vfmsubmask_str = """%res = call <$W x $(typ)> asm "vfmsub231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
                            ret <$W x $(typ)> %res"""
                quote
                    $(Expr(:meta,:inline))
                    Vec(llvmcall($vfmsubmask_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
                end
            end
            @generated function ifelse(::typeof(vfnmsub231), m::Mask{W,U}, a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W,U<:Unsigned,T<:Union{Float32,Float64}}
                if !((W ≥ 8) && ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE))
                    return Expr(:block, Expr(:meta, :inline), :(ifelse(vfnmsub, m, a, b, c)))
                end
                typ = LLVM_TYPES[T]
                suffix = T == Float32 ? "ps" : "pd"
                vfnmsubmask_str = """%res = call <$W x $(typ)> asm "vfnmsub231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
                            ret <$W x $(typ)> %res"""
                quote
                    $(Expr(:meta,:inline))
                    Vec(llvmcall($vfnmsubmask_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
                end
            end
        end
    end
end
    # for T ∈ [Float32,Float64]
    #     W = 16 ÷ sizeof(T)
    #     local suffix = T == Float32 ? "ps" : "pd"
    #     typ = LLVM_TYPES[T]
    #     while W <= VectorizationBase.REGISTER_SIZE ÷ sizeof(T)
    #         vfmadd_str = """%res = call <$W x $(typ)> asm "vfmadd231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
    #             ret <$W x $(typ)> %res"""
    #         vfnmadd_str = """%res = call <$W x $(typ)> asm "vfnmadd231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
    #             ret <$W x $(typ)> %res"""
    #         vfmsub_str = """%res = call <$W x $(typ)> asm "vfmsub231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
    #             ret <$W x $(typ)> %res"""
    #         vfnmsub_str = """%res = call <$W x $(typ)> asm "vfnmsub231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
    #             ret <$W x $(typ)> %res"""
    #         @eval begin
    #             @inline function vfmadd231(a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                 Vec(llvmcall($vfmadd_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
    #             end
    #             @inline function vfnmadd231(a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                 Vec(llvmcall($vfnmadd_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
    #             end
    #             @inline function vfmsub231(a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                 Vec(llvmcall($vfmsub_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
    #             end
    #             @inline function vfnmsub231(a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                 Vec(llvmcall($vfnmsub_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
    #             end
    #         end
    #         if VectorizationBase.AVX512BW && W ≥ 8
    #             vfmaddmask_str = """%res = call <$W x $(typ)> asm "vfmadd231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
    #                 ret <$W x $(typ)> %res"""
    #             vfnmaddmask_str = """%res = call <$W x $(typ)> asm "vfnmadd231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
    #                 ret <$W x $(typ)> %res"""
    #             vfmsubmask_str = """%res = call <$W x $(typ)> asm "vfmsub231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
    #                 ret <$W x $(typ)> %res"""
    #             vfnmsubmask_str = """%res = call <$W x $(typ)> asm "vfnmsub231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
    #                 ret <$W x $(typ)> %res"""
    #             U = VectorizationBase.mask_type(W)
    #             @eval begin
    #                 @inline function ifelse(::typeof(vfmadd231), m::Mask{$W,$U}, a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                     Vec(llvmcall($vfmaddmask_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
    #                 end
    #                 @inline function ifelse(::typeof(vfnmadd231), m::Mask{$W,$U}, a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                     Vec(llvmcall($vfnmaddmask_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
    #                 end
    #                 @inline function ifelse(::typeof(vfmsub231), m::Mask{$W,$U}, a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                     Vec(llvmcall($vfmsubmask_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
    #                 end
    #                 @inline function ifelse(::typeof(vfnmsub231), m::Mask{$W,$U}, a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                     Vec(llvmcall($vfnmsubmask_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
    #                 end
    #             end
    #         end
    #         W += W
    #     end
# end

@inline ifelse(f::F, m::Mask, a::Vararg{Any,K}) where {F<:Function,K} = ifelse(m, f(a...), a[K])

"""
Fast approximate reciprocal.

Guaranteed accurate to at least 2^-14 ≈ 6.103515625e-5.

Useful for special funcion implementations.
"""
@inline inv_approx(x) = Base.FastMath.inv_fast(x)
@inline inv_approx(v::VecUnroll) = VecUnroll(fmap(inv_approx, v.data))
@generated function inv_approx(v::Vec{W,T}) where {W, T <: Union{Float32, Float64}}
    ((Sys.ARCH === :x86_64) || (Sys.ARCH === :i686)) || return Expr(:block, Expr(:meta, :inline), :(inv(v)))
    bits = 8sizeof(T) * W
    if (AVX512F && (bits === 512)) || (AVX512VL && (bits ∈ (128, 256)))
        typ = T === Float64 ? "double" : "float"
        vtyp = "<$W x $(typ)>"
        dors = T === Float64 ? "d" : "s"
        f = "@llvm.x86.avx512.rcp14.p$(dors).$(bits)"
        decl = "declare $(vtyp) $f($(vtyp), $(vtyp), i$(min(8,W))) nounwind readnone"
        instrs = "%res = call $(vtyp) $f($vtyp %0, $vtyp zeroinitializer, i$(min(8,W)) -1)\nret $(vtyp) %res"
        return llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T}}), vtyp, [vtyp], [:(data(v))])
    end
    if (AVX && (W == 8)) && (T === Float32)
        decl = "declare <8 x float> @llvm.x86.avx.rcp.ps.256(<8 x float>) nounwind readnone"
        instrs = "%res = call <8 x float> @llvm.x86.avx.rcp.ps.256(<8 x float> %0)\nret <8 x float> %res"
        return llvmcall_expr(decl, instrs, :(_Vec{8,Float32}), :(Tuple{_Vec{8,Float32}}), "<8 x float>", ["<8 x float>"], [:(data(v))])
    elseif W == 4
        decl = "declare <4 x float> @llvm.x86.sse.rcp.ps(<4 x float>) nounwind readnone"
        instrs = "%res = call <4 x float> @llvm.x86.sse.rcp.ps(<4 x float> %0)\nret <4 x float> %res"
        if T === Float32
            return llvmcall_expr(decl, instrs, :(_Vec{4,Float32}), :(Tuple{_Vec{4,Float32}}), "<4 x float>", ["<4 x float>"], [:(data(v))])
        else#if T === Float64
            argexpr = [:(data(convert(Float32, v)))]
            call = llvmcall_expr(decl, instrs, :(_Vec{4,Float32}), :(Tuple{_Vec{4,Float32}}), "<4 x float>", ["<4 x float>"], argexpr, true)
            return Expr(:block, Expr(:meta, :inline), :(convert(Float64, $call)))
        end
    end
    Expr(:block, Expr(:meta, :inline), :(inv(v)))
end


text:
using Distributions
using Test

function naive_esf(x::AbstractVector{T}) where T <: Real
    n = length(x)
    S = zeros(T, n+1)
    states = hcat(reverse.(digits.(0:2^n-1,base=2,pad=n))...)'

    r_states = vec(mapslices(sum, states, dims=2))

    for r in 0:n
        idx = findall(r_states .== r)
        S[r+1] = sum(mapslices(x->prod(x[x .!= 0]), states[idx, :] .* x', dims=2))
    end
    return S
end

function naive_pb(p::AbstractVector{T}) where T <: Real
    x = p ./ (1 .- p)
    naive_esf(x) * prod(1 ./ (1 .+ x))
end

# test PoissonBinomial PMF algorithms
x = [3.5118, .6219, .2905, .8450, 1.8648]
n = length(x)
p = x ./ (1 .+ x)

naive_sol = naive_pb(p)

@test Distributions.poissonbinomial_pdf_fft(p) ≈ naive_sol
@test Distributions.poissonbinomial_pdf(p) ≈ naive_sol
@test Distributions.poissonbinomial_pdf(Tuple(p)) ≈ naive_sol

@test Distributions.poissonbinomial_pdf_fft(p) ≈ Distributions.poissonbinomial_pdf(p)

# Test the special base where PoissonBinomial distribution reduces
# to Binomial distribution
for (p, n) in [(0.8, 6), (0.5, 10), (0.04, 20)]
    local p

    d = PoissonBinomial(fill(p, n))
    dref = Binomial(n, p)
    println("   testing PoissonBinomial p=$p, n=$n")

    @test isa(d, PoissonBinomial)
    @test minimum(d) == 0
    @test maximum(d) == n
    @test extrema(d) == (0, n)
    @test ntrials(d) == n
    @test @inferred(entropy(d))  ≈ entropy(dref)
    @test @inferred(median(d))   ≈ median(dref)
    @test @inferred(mean(d))     ≈ mean(dref)
    @test @inferred(var(d))      ≈ var(dref)
    @test @inferred(kurtosis(d)) ≈ kurtosis(dref)
    @test @inferred(skewness(d)) ≈ skewness(dref)

    for t=0:5
        @test @inferred(mgf(d, t)) ≈ mgf(dref, t)
        @test @inferred(cf(d, t))  ≈ cf(dref, t)
    end
    for i=0.1:0.1:.9
        @test @inferred(quantile(d, i)) ≈ quantile(dref, i)
    end
    for i=0:n
        @test isapprox(@inferred(cdf(d, i)), cdf(dref, i), atol=1e-15)
        @test isapprox(@inferred(pdf(d, i)), pdf(dref, i), atol=1e-15)
    end

end

# Test against a sum of three Binomial distributions
for (n₁, n₂, n₃, p₁, p₂, p₃) in [(10, 10, 10, 0.1, 0.5, 0.9),
                                 (1, 10, 100, 0.99, 0.1, 0.05),
                                 (5, 1, 3, 0.01, 0.99, 0.999),
                                 (10, 7, 10, 0., 0.9, 0.5)]

    n = n₁ + n₂ + n₃
    p = zeros(n)
    p[1:n₁] .= p₁
    p[n₁+1: n₁ + n₂] .= p₂
    p[n₁ + n₂ + 1:end] .= p₃
    d = PoissonBinomial(p)
    println("   testing PoissonBinomial [$(n₁) × $(p₁), $(n₂) × $(p₂), $(n₃) × $(p₃)]")
    b1 = Binomial(n₁, p₁)
    b2 = Binomial(n₂, p₂)
    b3 = Binomial(n₃, p₃)

    pmf1 = pdf.(b1, support(b1))
    pmf2 = pdf.(b2, support(b2))
    pmf3 = pdf.(b3, support(b3))

    @test @inferred(mean(d)) ≈ (mean(b1) + mean(b2) + mean(b3))
    @test @inferred(var(d))  ≈ (var(b1) + var(b2) + var(b3))
    for t=0:5
        @test @inferred(mgf(d, t)) ≈ (mgf(b1, t) * mgf(b2, t) * mgf(b3, t))
        @test @inferred(cf(d, t))  ≈ (cf(b1, t) * cf(b2, t) * cf(b3, t))
    end

    for k=0:n
        m = 0.
        for i=0:min(n₁, k)
            mc = 0.
            for j=i:min(i+n₂, k)
                mc += (k - j <= n₃) && pmf2[j-i+1] * pmf3[k-j+1]
            end
            m += pmf1[i+1] * mc
        end
        @test isapprox(@inferred(pdf(d, k)), m, atol=1e-15)
    end
end

# Test the _dft helper function
@testset "_dft" begin
    x = Distributions._dft(collect(1:8))
    # Comparator computed from FFTW
    fftw_fft = [36.0 + 0.0im,
                -4.0 + 9.65685424949238im,
                -4.0 + 4.0im,
                -4.0 + 1.6568542494923806im,
                -4.0 + 0.0im,
                -4.0 - 1.6568542494923806im,
                -4.0 - 4.0im,
                -4.0 - 9.65685424949238im]
    @test x ≈ fftw_fft
end

# Test autodiff using ForwardDiff
f = x -> logpdf(PoissonBinomial(x), 0)
at = [0.5, 0.5]
@test isapprox(ForwardDiff.gradient(f, at), fdm(f, at), atol=1e-6)


text:
struct NullData end
const DEFAULT_DATA = Iterators.cycle((NullData(),))
Base.iterate(::NullData, i=1) = nothing
Base.length(::NullData) = 0

get_maxiters(data) = Iterators.IteratorSize(typeof(DEFAULT_DATA)) isa Iterators.IsInfinite ||
                     Iterators.IteratorSize(typeof(DEFAULT_DATA)) isa Iterators.SizeUnknown ?
                     typemax(Int) : length(data)

#=
function update!(x::AbstractArray, x̄::AbstractArray{<:ForwardDiff.Dual})
  x .-= x̄
end

function update!(x::AbstractArray, x̄)
  x .-= getindex.(ForwardDiff.partials.(x̄),1)
end

function update!(opt, x, x̄)
  x .-= Flux.Optimise.apply!(opt, x, x̄)
end

function update!(opt, x, x̄::AbstractArray{<:ForwardDiff.Dual})
  x .-= Flux.Optimise.apply!(opt, x, getindex.(ForwardDiff.partials.(x̄),1))
end

function update!(opt, xs::Flux.Zygote.Params, gs)
    update!(opt, xs[1], gs)
end
=#

maybe_with_logger(f, logger) = logger === nothing ? f() : Logging.with_logger(f, logger)

function default_logger(logger)
    Logging.min_enabled_level(logger) ≤ ProgressLogging.ProgressLevel && return nothing
    if Sys.iswindows() || (isdefined(Main, :IJulia) && Main.IJulia.inited)
          progresslogger = ConsoleProgressMonitor.ProgressLogger()
    else
          progresslogger = TerminalLoggers.TerminalLogger()
    end
    logger1 = LoggingExtras.EarlyFilteredLogger(progresslogger) do log
        log.level == ProgressLogging.ProgressLevel
    end
    logger2 = LoggingExtras.EarlyFilteredLogger(logger) do log
        log.level != ProgressLogging.ProgressLevel
    end
    LoggingExtras.TeeLogger(logger1, logger2)
end

macro withprogress(progress, exprs...)
    quote
      if $progress
          $maybe_with_logger($default_logger($Logging.current_logger())) do
              $ProgressLogging.@withprogress $(exprs...)
            end
      else
        $(exprs[end])
      end
    end |> esc
  end

function __solve(prob::OptimizationProblem, opt, data = DEFAULT_DATA;
                 maxiters::Number = 0, cb = (args...) -> (false),
                 progress = false, save_best = true, kwargs...)

    if data != DEFAULT_DATA
      maxiters = length(data)
    else
      if maxiters <= 0.0
        error("The number of maxiters has to be a non-negative and non-zero number.")
      end
      data = take(data, maxiters)
    end

    # Flux is silly and doesn't have an abstract type on its optimizers, so assume
    # this is a Flux optimizer
    θ = copy(prob.u0)
    ps = Flux.params(θ)

    t0 = time()

    local x, min_err, _loss
    min_err = typemax(eltype(prob.u0)) #dummy variables
    min_opt = 1

    f = instantiate_function(prob.f,prob.u0,prob.f.adtype,prob.p)

    @withprogress progress name="Training" begin
      for (i,d) in enumerate(data)
        gs = Flux.Zygote.gradient(ps) do
            x = prob.f(θ,prob.p, d...)
            first(x)
          end
        x = f.f(θ, prob.p, d...)
        cb_call = cb(θ, x...)
        if !(typeof(cb_call) <: Bool)
          error("The callback should return a boolean `halt` for whether to stop the optimization process. Please see the sciml_train documentation for information.")
        elseif cb_call
          break
        end
        msg = @sprintf("loss: %.3g", x[1])
        progress && ProgressLogging.@logprogress msg i/maxiters
        Flux.update!(opt, ps, gs)

        if save_best
          if first(x) < first(min_err)  #found a better solution
            min_opt = opt
            min_err = x
          end
          if i == maxiters  #Last iteration, revert to best.
            opt = min_opt
            cb(θ,min_err...)
          end
        end
      end
    end

    _time = time()

    SciMLBase.build_solution(prob, opt, θ, x[1])
    # here should be build_solution to create the output message
end

decompose_trace(trace::Optim.OptimizationTrace) = last(trace)
decompose_trace(trace) = trace

function __solve(prob::OptimizationProblem, opt::Optim.AbstractOptimizer,
                 data = DEFAULT_DATA;
                 maxiters = nothing,
                 cb = (args...) -> (false),
                 progress = false,
                 kwargs...)
    local x, cur, state

    if data != DEFAULT_DATA
        maxiters = length(data)
    end

    cur, state = iterate(data)

    function _cb(trace)
        cb_call = opt == NelderMead() ? cb(decompose_trace(trace).metadata["centroid"],x...) : cb(decompose_trace(trace).metadata["x"],x...)
        if !(typeof(cb_call) <: Bool)
            error("The callback should return a boolean `halt` for whether to stop the optimization process.")
        end
        cur, state = iterate(data, state)
        cb_call
    end

    if !(isnothing(maxiters)) && maxiters <= 0.0
        error("The number of maxiters has to be a non-negative and non-zero number.")
    elseif !(isnothing(maxiters))
        maxiters = convert(Int, maxiters)
    end

    f = instantiate_function(prob.f,prob.u0,prob.f.adtype,prob.p)

    !(opt isa Optim.ZerothOrderOptimizer) && f.grad === nothing && error("Use OptimizationFunction to pass the derivatives or automatically generate them with one of the autodiff backends")

    _loss = function(θ)
        x = f.f(θ, prob.p, cur...)
        return first(x)
    end

    fg! = function (G,θ)
        if G !== nothing
            f.grad(G, θ, cur...)
        end
        return _loss(θ)
    end

    if opt isa Optim.KrylovTrustRegion
        optim_f = Optim.TwiceDifferentiableHV(_loss, fg!, (H,θ,v) -> f.hv(H,θ,v,cur...), prob.u0)
    else
        optim_f = TwiceDifferentiable(_loss, (G, θ) -> f.grad(G, θ, cur...), fg!, (H,θ) -> f.hess(H,θ,cur...), prob.u0)
    end

    original = Optim.optimize(optim_f, prob.u0, opt,
                              !(isnothing(maxiters)) ?
                                Optim.Options(;extended_trace = true,
                                               callback = _cb,
                                               iterations = maxiters,
                                               kwargs...) :
                                Optim.Options(;extended_trace = true,
                                               callback = _cb, kwargs...))
    SciMLBase.build_solution(prob, opt, original.minimizer,
                             original.minimum; original=original)
end

function __solve(prob::OptimizationProblem, opt::Union{Optim.Fminbox,Optim.SAMIN},
                 data = DEFAULT_DATA;
                 maxiters = nothing,
                 cb = (args...) -> (false),
                 progress = false,
                 kwargs...)

    local x, cur, state

    if data != DEFAULT_DATA
        maxiters = length(data)
    end

    cur, state = iterate(data)

    function _cb(trace)
        cb_call = !(opt isa Optim.SAMIN) && opt.method == NelderMead() ? cb(decompose_trace(trace).metadata["centroid"],x...) : cb(decompose_trace(trace).metadata["x"],x...)
        if !(typeof(cb_call) <: Bool)
            error("The callback should return a boolean `halt` for whether to stop the optimization process.")
        end
        cur, state = iterate(data, state)
        cb_call
    end

    if !(isnothing(maxiters)) && maxiters <= 0.0
        error("The number of maxiters has to be a non-negative and non-zero number.")
    elseif !(isnothing(maxiters))
        maxiters = convert(Int, maxiters)
    end

    f = instantiate_function(prob.f,prob.u0,prob.f.adtype,prob.p)

    !(opt isa Optim.ZerothOrderOptimizer) && f.grad === nothing && error("Use OptimizationFunction to pass the derivatives or automatically generate them with one of the autodiff backends")

    _loss = function(θ)
        x = f.f(θ, prob.p, cur...)
        return first(x)
    end
    fg! = function (G,θ)
        if G !== nothing
            f.grad(G, θ, cur...)
        end

        return _loss(θ)
    end
    optim_f = OnceDifferentiable(_loss, (G, θ) -> f.grad(G, θ, cur...), fg!, prob.u0)

    original = Optim.optimize(optim_f, prob.lb, prob.ub, prob.u0, opt,
                              !(isnothing(maxiters)) ? Optim.Options(;
                              extended_trace = true, callback = _cb,
                              iterations = maxiters, kwargs...) :
                              Optim.Options(;extended_trace = true,
                              callback = _cb, kwargs...))
    SciMLBase.build_solution(prob, opt, original.minimizer,
                             original.minimum; original=original)
end

function __solve(prob::OptimizationProblem, opt::Optim.ConstrainedOptimizer,
                 data = DEFAULT_DATA;
                 maxiters = nothing,
                 cb = (args...) -> (false),
                 progress = false,
                 kwargs...)

    local x, cur, state

    if data != DEFAULT_DATA
        maxiters = length(data)
    end

    cur, state = iterate(data)

      function _cb(trace)
      cb_call = cb(decompose_trace(trace).metadata["x"],x...)
      if !(typeof(cb_call) <: Bool)
          error("The callback should return a boolean `halt` for whether to stop the optimization process.")
      end
      cur, state = iterate(data, state)
      cb_call
    end

    if !(isnothing(maxiters)) && maxiters <= 0.0
        error("The number of maxiters has to be a non-negative and non-zero number.")
    elseif !(isnothing(maxiters))
        maxiters = convert(Int, maxiters)
    end

    f = instantiate_function(prob.f,prob.u0,prob.f.adtype,prob.p,prob.ucons === nothing ? 0 : length(prob.ucons))

    f.cons_j ===nothing && error("Use OptimizationFunction to pass the derivatives or automatically generate them with one of the autodiff backends")

    _loss = function(θ)
        x = f.f(θ, prob.p, cur...)
        return x[1]
    end
    fg! = function (G,θ)
        if G !== nothing
            f.grad(G, θ, cur...)
        end
        return _loss(θ)
    end
    optim_f = TwiceDifferentiable(_loss, (G, θ) -> f.grad(G, θ, cur...), fg!, (H,θ) -> f.hess(H, θ, cur...), prob.u0)

    cons! = (res, θ) -> res .= f.cons(θ);

    cons_j! = function(J, x)
        f.cons_j(J, x)
    end

    cons_hl! = function (h, θ, λ)
        res = [similar(h) for i in 1:length(λ)]
        f.cons_h(res, θ)
        for i in 1:length(λ)
            h .+= λ[i]*res[i]
        end
    end

    lb = prob.lb === nothing ? [] : prob.lb
    ub = prob.ub === nothing ? [] : prob.ub
    optim_fc = TwiceDifferentiableConstraints(cons!, cons_j!, cons_hl!, lb, ub, prob.lcons, prob.ucons)

    original = Optim.optimize(optim_f, optim_fc, prob.u0, opt,
                              !(isnothing(maxiters)) ? Optim.Options(;
                                extended_trace = true, callback = _cb,
                                iterations = maxiters, kwargs...) :
                                Optim.Options(;extended_trace = true,
                                callback = _cb, kwargs...))
    SciMLBase.build_solution(prob, opt, original.minimizer,
                             original.minimum; original=original)
end

function __init__()
    @require BlackBoxOptim="a134a8b2-14d6-55f6-9291-3336d3ab0209" begin
        decompose_trace(opt::BlackBoxOptim.OptRunController) = BlackBoxOptim.best_candidate(opt)

        struct BBO
            method::Symbol
            BBO(method) = new(method)
        end

        BBO() = BBO(:adaptive_de_rand_1_bin_radiuslimited) # the recommended optimizer as default

        function __solve(prob::OptimizationProblem, opt::BBO, data = DEFAULT_DATA;
                         cb = (args...) -> (false), maxiters = nothing,
                         progress = false, kwargs...)

            local x, cur, state

            if data != DEFAULT_DATA
                maxiters = length(data)
            end

            cur, state = iterate(data)

            function _cb(trace)
              cb_call = cb(decompose_trace(trace),x...)
              if !(typeof(cb_call) <: Bool)
                error("The callback should return a boolean `halt` for whether to stop the optimization process.")
              end
              if cb_call == true
                BlackBoxOptim.shutdown_optimizer!(trace) #doesn't work
              end
              cur, state = iterate(data, state)
              cb_call
            end

            if !(isnothing(maxiters)) && maxiters <= 0.0
                error("The number of maxiters has to be a non-negative and non-zero number.")
            elseif !(isnothing(maxiters))
                maxiters = convert(Int, maxiters)
            end

            _loss = function(θ)
                x = prob.f(θ, prob.p, cur...)
                return first(x)
            end

            bboptre = !(isnothing(maxiters)) ? BlackBoxOptim.bboptimize(_loss;Method = opt.method, SearchRange = [(prob.lb[i], prob.ub[i]) for i in 1:length(prob.lb)], MaxSteps = maxiters, CallbackFunction = _cb, CallbackInterval = 0.0, kwargs...) : BlackBoxOptim.bboptimize(_loss;Method = opt.method, SearchRange = [(prob.lb[i], prob.ub[i]) for i in 1:length(prob.lb)], CallbackFunction = _cb, CallbackInterval = 0.0, kwargs...)

            SciMLBase.build_solution(prob, opt, BlackBoxOptim.best_candidate(bboptre),
                                     BlackBoxOptim.best_fitness(bboptre); original=bboptre)
        end
    end

    @require NLopt="76087f3c-5699-56af-9a33-bf431cd00edd" begin
        function __solve(prob::OptimizationProblem, opt::NLopt.Opt;
                         maxiters = nothing, nstart = 1,
                         local_method = nothing,
                         progress = false, kwargs...)
            local x

            if !(isnothing(maxiters)) && maxiters <= 0.0
                error("The number of maxiters has to be a non-negative and non-zero number.")
            elseif !(isnothing(maxiters))
                maxiters = convert(Int, maxiters)
            end

            f = instantiate_function(prob.f,prob.u0,prob.f.adtype,prob.p)

            _loss = function(θ)
                x = f.f(θ, prob.p)
                return x[1]
            end

            fg! = function (θ,G)
                if length(G) > 0
                    f.grad(G, θ)
                end

                return _loss(θ)
            end

            NLopt.min_objective!(opt, fg!)

            if prob.ub !== nothing
                NLopt.upper_bounds!(opt, prob.ub)
            end
            if prob.lb !== nothing
                NLopt.lower_bounds!(opt, prob.lb)
            end
            if !(isnothing(maxiters))
                NLopt.maxeval!(opt, maxiters)
            end
            if nstart > 1 && local_method !== nothing
                NLopt.local_optimizer!(opt, local_method)
                if !(isnothing(maxiters))
                    NLopt.maxeval!(opt, nstart * maxiters)
                end
            end

            t0 = time()
            (minf,minx,ret) = NLopt.optimize(opt, prob.u0)
            _time = time()

            SciMLBase.build_solution(prob, opt, minx, minf; original=nothing)
        end
    end

    @require MultistartOptimization = "3933049c-43be-478e-a8bb-6e0f7fd53575" begin
        function __solve(prob::OptimizationProblem, opt::MultistartOptimization.TikTak;
                         local_method, local_maxiters = nothing,
                         progress = false, kwargs...)
            local x, _loss

            if !(isnothing(local_maxiters)) && local_maxiters <= 0.0
                error("The number of local_maxiters has to be a non-negative and non-zero number.")
            else !(isnothing(local_maxiters))
                local_maxiters = convert(Int, local_maxiters)
            end

            _loss = function(θ)
                x = prob.f(θ, prob.p)
                return first(x)
            end

            t0 = time()

            P = MultistartOptimization.MinimizationProblem(_loss, prob.lb, prob.ub)
            multistart_method = opt
            if !(isnothing(local_maxiters))
                local_method = MultistartOptimization.NLoptLocalMethod(local_method, maxeval = local_maxiters)
            else
                local_method = MultistartOptimization.NLoptLocalMethod(local_method)
            end
            p = MultistartOptimization.multistart_minimization(multistart_method, local_method, P)

            t1 = time()

            SciMLBase.build_solution(prob, opt, p.location, p.value; original=p)
        end
    end

    @require QuadDIRECT = "dae52e8d-d666-5120-a592-9e15c33b8d7a" begin
        export QuadDirect

        struct QuadDirect
        end

        function __solve(prob::OptimizationProblem, opt::QuadDirect; splits = nothing, maxiters = nothing, kwargs...)

            local x, _loss

            if !(isnothing(maxiters)) && maxiters <= 0.0
                error("The number of maxiters has to be a non-negative and non-zero number.")
            elseif !(isnothing(maxiters))
                maxiters = convert(Int, maxiters)
            end

            if splits === nothing
                error("You must provide the initial locations at which to evaluate the function in `splits` (a list of 3-vectors with values in strictly increasing order and within the specified bounds).")
            end

            _loss = function(θ)
                x = prob.f(θ, prob.p)
                return first(x)
            end

            t0 = time()

            root, x0 = !(isnothing(maxiters)) ? QuadDIRECT.analyze(_loss, splits, prob.lb, prob.ub; maxevals = maxiters, kwargs...) : QuadDIRECT.analyze(_loss, splits, prob.lb, prob.ub; kwargs...)
            box = minimum(root)
            t1 = time()

            SciMLBase.build_solution(prob, opt, QuadDIRECT.position(box, x0), QuadDIRECT.value(box); original=root)
        end
    end

    @require Evolutionary="86b6b26d-c046-49b6-aa0b-5f0f74682bd6" begin
        decompose_trace(trace::Evolutionary.OptimizationTrace) = last(trace)

        function Evolutionary.trace!(record::Dict{String,Any}, objfun, state, population, method::Evolutionary.AbstractOptimizer, options)
            record["x"] = population
        end

        function __solve(prob::OptimizationProblem, opt::Evolutionary.AbstractOptimizer, data = DEFAULT_DATA;
                         cb = (args...) -> (false), maxiters = nothing,
                         progress = false, kwargs...)
            local x, cur, state

            if data != DEFAULT_DATA
                maxiters = length(data)
            end

            cur, state = iterate(data)

            function _cb(trace)
                cb_call = cb(decompose_trace(trace).metadata["x"],trace.value...)
                if !(typeof(cb_call) <: Bool)
                    error("The callback should return a boolean `halt` for whether to stop the optimization process.")
                end
                cur, state = iterate(data, state)
                cb_call
            end

            if !(isnothing(maxiters)) && maxiters <= 0.0
                error("The number of maxiters has to be a non-negative and non-zero number.")
            elseif !(isnothing(maxiters))
                maxiters = convert(Int, maxiters)
            end

            _loss = function(θ)
                x = prob.f(θ, prob.p, cur...)
                return first(x)
            end

            t0 = time()

            result = Evolutionary.optimize(_loss, prob.u0, opt, !isnothing(maxiters) ? Evolutionary.Options(;iterations = maxiters, callback = _cb, kwargs...)
                                                                                : Evolutionary.Options(;callback = _cb, kwargs...))
            t1 = time()

            SciMLBase.build_solution(prob, opt, Evolutionary.minimizer(result), Evolutionary.minimum(result); original=result)
        end
    end
    @require CMAEvolutionStrategy="8d3b24bd-414e-49e0-94fb-163cc3a3e411" begin

        struct CMAEvolutionStrategyOpt end

        function __solve(prob::OptimizationProblem, opt::CMAEvolutionStrategyOpt, data = DEFAULT_DATA;
                         cb = (args...) -> (false), maxiters = nothing,
                         progress = false, kwargs...)
            local x, cur, state

            if data != DEFAULT_DATA
                maxiters = length(data)
            end

            cur, state = iterate(data)

            function _cb(trace)
                cb_call = cb(decompose_trace(trace).metadata["x"],trace.value...)
                if !(typeof(cb_call) <: Bool)
                    error("The callback should return a boolean `halt` for whether to stop the optimization process.")
                end
                cur, state = iterate(data, state)
                cb_call
            end

            _loss = function(θ)
                x = prob.f(θ, prob.p, cur...)
                return first(x)
            end

            if !(isnothing(maxiters)) && maxiters <= 0.0
                error("The number of maxiters has to be a non-negative and non-zero number.")
            elseif !(isnothing(maxiters))
                maxiters = convert(Int, maxiters)
            end

            result = CMAEvolutionStrategy.minimize(_loss, prob.u0, 0.1; lower = prob.lb, upper = prob.ub, maxiter = maxiters, kwargs...)
            CMAEvolutionStrategy.print_header(result)
            CMAEvolutionStrategy.print_result(result)
            println("\n")
            criterion = true

            if (result.stop.reason === :maxtime) #this is an arbitrary choice of convergence (based on the stop.reason values)
                criterion = false
            end

            SciMLBase.build_solution(prob, opt, result.logger.xbest[end], result.logger.fbest[end]; original=result)
        end
    end
end


text:
using DataStructures

"""
Hybrid-λ-PS-PF
--------------

Hybrid model : Prioritized Sweeping (with Particle Filtering) + SARSA-lambda
Prioritized Sweeping with: ParticleFilterJump for the Tsas' estimation
and Perfect Integration for Rsas' estimation.

See learner/prioritizedsweepinglearner.jl, learner/sarsaLlearner.jl and
learner/hybridlearner.jl for the core implementation.

Note 1
------
With particle filtering we need to run multiple times with different seeds.
As final LL we report the MEAN across seeds.
For signals/results we keep track of the ones of all seeds
and in the end report the one with the MAXIMUM LL.
Note 2
------
get_all_puzzles(), get_all_episodes_in_puzzle(), get_data_by_puzzle_episode()
are in rl_data_toolsforalgos.jl
Note 3
------
Temperature of softmax is removed, because this learner in NOT convex and weights
can be in (0, +Inf). The temperature is absorbed in the weights. Otherwise we
it is overparametrized.
"""
function getLL_Hybrid_Prio_Sweep_TParticleFilterJump(data::Array{Int,2},
                                                gamma::Float64,
                                                surpriseprobability::Float64,
                                                stochasticity::Float64,
                                                nr_upd_cycles::Union{Int, Float64},
                                                nparticles::Union{Int, Float64},
                                                alpha::Float64,
                                                lambda::Float64,
                                                w_MB::Float64,
                                                w_MF::Float64;
                                                particlefilterseeds = [1,2,3],
                                                init_R = 0.,
                                                do_debug_print = false
                                                )
    nr_upd_cycles = converttoInt(nr_upd_cycles)
    nparticles = converttoInt(nparticles)
    data_stats, nr_states, nr_actions, nr_samples = get_data_stats_and_dimensions(data)
    if data_stats.nr_subjects != 1
        error("Invalid argument: getLL_Hybrid_Prio_Sweep_TParticleFilterJump supports only one subject")
    end
    all_puzzles = get_all_puzzles(data)

    action_selection_probs_seeds = Array{Array{Float64, 1}, 1}(undef, length(particlefilterseeds))
    [action_selection_probs_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    LL_seeds = zeros(Float64, length(particlefilterseeds))

    Tsas_seeds = Array{Any, 1}(undef, length(particlefilterseeds))
    # signals:
    RPE_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [RPE_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    Surprise_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [Surprise_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    γ_surprise_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [γ_surprise_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    QSA_Hybrid_seeds = Array{Any, 1}(undef, length(particlefilterseeds))
    QSA_MB_seeds = Array{Any, 1}(undef, length(particlefilterseeds))
    QSA_MF_seeds = Array{Any, 1}(undef, length(particlefilterseeds))
    learnerMB = []
    learnerMF = []
    learnerHybrid = []
    modulatingfactor = 0.
    for i_seed in 1:length(particlefilterseeds)
        action_counter = 0
        for i_puzzle in all_puzzles
            # Puzzle = independent block of the experiment.
            # The code was made for general purposes, in this experiment we
            # have only ONE puzzle.
            learnerMB = PrioritizedSweeping(nr_states = nr_states,
                                        nr_actions = nr_actions,
                                        γ = gamma, init_R = init_R,
                                        nr_upd_cycles = nr_upd_cycles,
                                        Testimatetype = TParticleFilterJump,
                                        nparticles = nparticles,
                                        surpriseprobability = surpriseprobability,
                                        stochasticity = stochasticity,
                                        seedlearner = particlefilterseeds[i_seed],
                                        do_debug_print = do_debug_print
                                        )
            learnerMF = SarsaL(nr_states = nr_states,
                                nr_actions = nr_actions,
                                α = alpha,
                                γ = gamma,
                                λ = lambda,
                                initvalue = 0.
                                )
            learnerHybrid = Hybrid(nr_states = nr_states,
                                nr_actions = nr_actions,
                                w_MB = w_MB,
                                w_MF = w_MF
                                )

            all_epi_in_puzzle = get_all_episodes_in_puzzle(data, i_puzzle)
            for i_episode in all_epi_in_puzzle
                episode_data = get_data_by_puzzle_episode(data, i_puzzle, i_episode)
                path_length = size(episode_data)[1]  # = T
                if do_debug_print
                    println("i_puzzle:$i_puzzle | i_episode:$i_episode | path_length:$path_length")
                end
                isendofexperiment = false
                for t in 1:path_length
                    action_counter += 1
                    state = episode_data[t, Int(c_state)]
                    action = episode_data[t, Int(c_action)]
                    reward = episode_data[t, Int(c_reward)]
                    next_state = episode_data[t, Int(c_nextState)]
                    done = reward>0 ? true : false
                    if do_debug_print
                        println("t:$t | s:$state | a:$action | R:$reward | s':$next_state | done: $done")
                    end

                    # Hybrid learner
                    update!(learnerHybrid, state, action, reward, learnerMB.QSA, learnerMF.QSA) #Q_hybrid =  w_MB*learnerMB.QSA + w_MF*learnerMF.QSA

                    p_actions = getSoftmaxProbs(learnerHybrid.QSA[state, :], 1., [0., 0.])
                    action_selection_probs_seeds[i_seed][action_counter] = copy(p_actions[action])
                    LL_seeds[i_seed] += log(p_actions[action])

                    # Prioritized Sweeping
                    update!(learnerMB, state, action, next_state, reward, done)
                    γ_surprise_seeds[i_seed][action_counter] = copy(learnerMB.Testimate.γ_surprise[1])
                    Surprise_seeds[i_seed][action_counter] = copy(learnerMB.Testimate.Sgm[1])
                    Tsas_seeds[i_seed]= deepcopy(learnerMB.Testimate.Ps1a0s0)

                    # Sort out if end of episode or end of experiment
                    next_action = 0
                    if reward==0 #non goal state
                        if nr_samples > action_counter
                            next_action = data[action_counter+1,Int(c_action)]
                        else
                            isendofexperiment = true
                        end
                    end

                    # Sarsa-Lambda with replacing traces
                    update!(learnerMF, state, action, next_state, next_action, reward,
                            isendofexperiment, modulatingfactor)
                    RPE_seeds[i_seed][action_counter] = copy(learnerMF.rpe[1]) # Even though it is the same across seeds, save it anyway...

                    QSA_MB_seeds[i_seed] = copy(learnerMB.QSA)
                    QSA_MF_seeds[i_seed] = copy(learnerMF.QSA)
                    QSA_Hybrid_seeds[i_seed] = copy(learnerHybrid.QSA)

                end # end t
            end # end i_episode
        end # end i_puzzle
    end # end i_seed

    LL = get_meanLL_acrossseeds(LL_seeds)

    action_selection_probs_cat = hcat(action_selection_probs_seeds...)
    action_selection_probs = mean(action_selection_probs_cat, dims=2)[:]
    trial_by_trial_p_sa = (action_selection_probs)

    # For signals: return the ones corresponding to the Max LL across seeds
    LL_max, seed_idx = findmax(LL_seeds)
    RPE_record = copy(RPE_seeds[seed_idx])
    trial_by_trial_p_sa_max = (action_selection_probs_seeds[seed_idx])

    if do_debug_print
        println(QSA_Hybrid_seeds[seed_idx])
        println("action_selection_probs: $action_selection_probs")
        println("LL_mean: $LL_mean")
        println("LL_max: $LL_max")
        println("trial_by_trial_p_sa: $trial_by_trial_p_sa")
    end

    algo = "hybrid_priosweepparticlefilter_sarsaL"
    param_v = [nr_upd_cycles, nparticles, gamma, surpriseprobability,
                stochasticity, alpha, lambda, w_MB, w_MF]
    param_str = "nr_upd_cycles=$nr_upd_cycles, nparticles=$nparticles,
                gamma=$(round(gamma,digits=4)),
                surpriseprobability=$(round(surpriseprobability,digits=4)),
                stochasticity=$(round(stochasticity,digits=4)),
                alpha=$(round(alpha,digits=4)),
                lambda=$(round(lambda,digits=4)),
                w_MB=$(round(w_MB,digits=4)),
                w_MF=$(round(w_MF,digits=4))"

    results_dict = Dict{String,Any}("finalQSA_Hybrid_record" => QSA_Hybrid_seeds[seed_idx],
                                    "finalQSA_MF_record" => QSA_MF_seeds[seed_idx],
                                    "finalQSA_MB_record" => QSA_MB_seeds[seed_idx],
                                    "final_Tsas" => Tsas_seeds[seed_idx],
                                    "Surprise" => Surprise_seeds[seed_idx],
                                    "gamma_surprise" => γ_surprise_seeds[seed_idx],
                                    "nr_upd_cycles" => nr_upd_cycles,
                                    "nparticles" => nparticles,
                                    "particlefilterseeds" => particlefilterseeds,
                                    "trial_by_trial_p_sa_max" => trial_by_trial_p_sa_max)

    signals = Step_by_step_signals(data_stats.all_subject_ids[1],
    RPE_record, #RPE
    [], #SPE
    trial_by_trial_p_sa,
    LL,
    data_stats,
    algo,
    param_v,
    param_str,
    99, #biased_action_idx, dummy value
    results_dict)
    return LL, signals
end


text:
using LiiBRA, UnitSystems, Test, LinearAlgebra

#Construct cell struct
Cell = Construct("LG M50")
Cell.RA.Tlen = 128
Cell.RA.Fs = 1
Cell.RA.H1 = 1:Cell.RA.Tlen
Cell.RA.H2 = 1:Cell.RA.Tlen

function TestDRA(Cell)
    A = B = C = D = Time = x = tuple()
    for Temp in 5.0:10.0:45.0
        #Arrhenius
        Cell.Const.T = 273.15+Temp
        Arr_Factor = (1/Cell.Const.T_ref-1/Cell.Const.T)/R

        #Set Cell Constants
        Cell.Const.κ = Cell.Const.κf(Cell.Const.ce0)*exp(Cell.Const.Ea_κ*Arr_Factor)
        Cell.RA.Nfft = Cell.RA.Nfft!(Cell.RA.Fs, Cell.RA.Tlen)
        Cell.RA.f = Cell.RA.f!(Cell.RA.Nfft)
        Cell.RA.s = Cell.RA.s!(Cell.RA.Fs,Cell.RA.Nfft,Cell.RA.f)
        Cell.Neg.β = Cell.Neg.β!(Cell.RA.s)
        Cell.Pos.β = Cell.Pos.β!(Cell.RA.s)

        #DRA
        A_DRA, B_DRA, C_DRA, D_DRA = DRA(Cell)

        #Flatten output into tuples
        A = flatten_(A,A_DRA)
        B = flatten_(B,B_DRA)
        C = flatten_(C,C_DRA)
        D = flatten_(D,D_DRA)
    end
return A, B, C, D
end

A,B,C,D = TestDRA(Cell)

@test size(A[1],1) == Cell.RA.M+1
@test typeof(A[1]) == Matrix{Float64}
@test typeof(B[1]) == Vector{Float64}
@test typeof(C[1]) == Matrix{Float64}
@test typeof(D[1]) == Vector{Float64}
#@test output eigs and compare to diagonal of A
#@test sim_model output at multiple SOC?

max_stars_repo_name: BradyPlanden/LiBRA

max_forks_repo_name: BradyPlanden/LiBRA


text:
import CLBLAS
using OpenCL: cl

const clblas = CLBLAS
clblas.setup()

device, ctx, queue = clblas.get_next_compute_context()

N = Unsigned(5)
data = Vector{Float32}(5)
data[1] = 1.1
data[2] = 2.22
data[3] = 3.333
data[4] = 4.4444
data[5] = 5.55555

bufX = cl.Buffer(Float32, ctx, (:r, :copy),hostbuf=data)

bufNRM2 = cl.Buffer(Float32, ctx, :w, 1)
scratchBuff = cl.Buffer(Float32, ctx, :rw, 2 * length(data))

offNRM2 = offx = Unsigned(0)
incx = cl.cl_int(1)
ncq = cl.cl_uint(1)
clq = queue.id
newl = cl.cl_uint(0)

event = cl.UserEvent(ctx, retain=true)
ptrEvent = [event.id]

clblas.clblasSnrm2( N, bufNRM2.id, offNRM2, bufX.id, offx, incx, scratchBuff.id, ncq, [clq], newl, C_NULL, ptrEvent);

cl.api.clWaitForEvents(cl.cl_uint(1), ptrEvent)

result = Vector{Float32}(1)
cl.enqueue_read_buffer(queue, bufNRM2, result, Unsigned(0), nothing, true)

expected = BLAS.nrm2(5, data, 1)
println(result)
println("------------")
println(expected)

try
    @assert(isapprox(norm(expected - result[1]), zero(Float32)))
    info("success!")
finally
    clblas.teardown()
end


text:
@testset "Camera" begin

    T = Float64
    p = 3
    model = BicycleGame(p=p)

    roadway_opts = MergingRoadwayOptions(merging_point=1.4, lane_width=0.28)
    roadway = build_roadway(roadway_opts)

    players = Vector{Player{T}}(undef, p)
    players[1] = Player(model, roadway.lane[3])
    players[2] = Player(model, roadway.lane[3])
    players[3] = Player(model, roadway.lane[3])

    sce = Scenario(model, roadway, players)

    # Initialize visualizers
    vis = Visualizer()

    # Visualize trajectories
    set_scenario!(vis, sce)
    set_env!(vis, VehicleState(1.3, 0.0, 0.0, 0.0))
    set_camera!(vis)
    set_camera_birdseye!(vis, height=8.0)

end


text:
import Base.@pure
using LinearAlgebra
using StaticArrays
using SpecialFunctions

const absolute_tolerance = 1e-12
const probability_tolerance = 1/128
const initially_equivariant = false
const σy = 20
const σx = 50
const σz = 5σy
const px = 0.1
const b = SA[0, 6σy, 0]
const s = SA[1, 1, 1]/sqrt(3)
const dist = 1e3
const n_trajectories = 50

@pure function ψG(t, x)
	σXt2 = σx^2 + im * t / 2
	σYt2 = σy^2 + im * t / 2
	σZt2 = σz^2 + im * t / 2
	√(σXt2 * σYt2 * σZt2) / (2*pi*σXt2*σYt2*σZt2)^(3/4) *
	exp( -1/4 * (
		(x[1] - px*t)^2 / σXt2 +
		 x[2]^2         / σYt2 +
		 x[3]^2         / σZt2
		) +
		im * (-px*px/2 + px*x[1])
	)
end

@pure function ∇ψG(t, x)
	σXt2 = σx^2 + im * t / 2
	σYt2 = σy^2 + im * t / 2
	σZt2 = σz^2 + im * t / 2
	SA[-(x[1] - px*t)/(2*σXt2) + px*im,
	   -(x[2]       )/(2*σYt2),
	   -(x[3]       )/(2*σZt2)] * ψG(t, x)
end

@pure function ψ(t, x)
	ψG(t, x - b) + ψG(t, x + b)
end

@pure function ∇ψ(t, x)
	∇ψG(t, x - b) + ∇ψG(t, x + b)
end

@pure function velocity(t, x)
	imag(∇ψ(t,x) / ψ(t,x)) + real(∇ψ(t, x) / ψ(t,x)) × s
end

@pure function vR(t, x)
	velocity(t, x) + s
end

@pure function vL(t,x)
	velocity(t, x) - s
end

@pure function rate(t, x)
	-2s⋅real(∇ψ(t,x) / ψ(t,x))
end

@pure function odestep(t, h, x, v)
	c1 = 0;
	c2 = 1/5;  a21 = 1/5;
	c3 = 3/10; a31 = 3/40;       a32 = 9/40;
	c4 = 3/5;  a41 = 3/10;       a42 = -9/10;   a43 = 6/5;
	c5 = 1;    a51 = -11/54;     a52 = 5/2;     a53 = -70/27;    a54 = 35/27;
	c6 = 7/8;  a61 = 1631/55296; a62 = 175/512; a63 = 575/13824; a64 = 44275/110592; a65 = 253/4096;
	b1 = 37/378; b2 = 0; b3 = 250/621; b4 = 125/594; b5 = 0; b6 = 512/1771;
	d1 = 2825/27648; d2 = 0; d3 = 18575/48384; d4 = 13525/55296; d5 = 277/14336; d6 = 1/4;
	hk1 = h*v(t, x)
	hk2 = h*v(t + c2*h, x + a21 * hk1)
	hk3 = h*v(t + c3*h, x + a31 * hk1 + a32 * hk2)
	hk4 = h*v(t + c4*h, x + a41 * hk1 + a42 * hk2 + a43 * hk3)
	hk5 = h*v(t + c5*h, x + a51 * hk1 + a52 * hk2 + a53 * hk3 + a54 * hk4)
	hk6 = h*v(t + c6*h, x + a61 * hk1 + a62 * hk2 + a63 * hk3 + a64 * hk4 + a65 * hk5)
	xout = x + b1 * hk1 + b2 * hk2 + b3 * hk3 + b4 * hk4 + b5 * hk5 + b6 * hk6
	xerr = (d1-b1) * hk1 + (d2-b2) * hk2 + (d3-b3) * hk3 + (d4-b4) * hk4 + (d5-b5) * hk5 + (d6-b6) * hk6

	err = norm(xerr, Inf)

	return xout, err
end

function trajectory(xi)
	h = 1.0

	steps = 1
	t = told = 0.0
	x = xold = xi
	ts = [t]
	xs = [x]
	χ = rand(Bool)

	while x[1] < dist
		if χ
			r1 = max(+rate(t,x), 0.0)
			xtry, err = odestep(t, h, x, vR)
		else
			r1 = max(-rate(t,x), 0.0)
			xtry, err = odestep(t, h, x, vL)
		end

		if err > absolute_tolerance
			h = 0.9 * h * max((absolute_tolerance / err)^(1/5), 0.1)
			continue
		end

		if χ
			r2 = max(+rate(t + h, xtry), 0.0)
		else
			r2 = max(-rate(t + h, xtry), 0.0)
		end

		r = (r1 + r2) / 2
		if r * h > probability_tolerance
			h = probability_tolerance / r
			continue
		end

		told = t
		xold = x
		t += h
		x = xtry
		steps += 1

		if rand() < r * h
			χ = !χ
			push!(ts, t)
			push!(xs, x)
		end

		h = 0.9 * h * min((absolute_tolerance / err)^(1/5), 10)
		if r * h > probability_tolerance
			h = probability_tolerance / r
		end
	end

	dt = (dist - xold[1]) / ((χ ? vR(told, xold) : vL(told, xold))[1])

	x, _ = odestep(told, dt,  xold, χ ? vR : vL)
	t = told + dt
	push!(ts, t)
	push!(xs, x)

	return ts, xs, steps
end

cd("/tmp/data")
foreach(rm,
	filter(
		endswith(".dat"),
		readdir()
	)
)
println(rand(["Allons-y", "Поехали", "Lad os gå"]))
trajectories = Threads.Atomic{Int}(0)
arrival_coordinates = Vector{SVector{4,Float64}}(undef, n_trajectories)

Threads.@threads for i in 1:n_trajectories

	if initially_equivariant
		δ = [σx,σy,σz].*randn(3)
	else
		# Distribute y-coordinates such that there is the same amount of
		# probability between them
		δuniform = ceil(i/2) / ((n_trajectories/2) + 1)
		δy = -√2 * σy * erfcinv(2 * δuniform)
		δ = SA[0.0, δy, 0.0]
	end

	if iseven(i)
		xi = SVector{3}(-b + δ)
	else
		xi = SVector{3}(+b + δ)
	end
	ts, xs, steps = trajectory(xi)
	open("traj$i.dat", "w") do io
		for i in 1:length(ts)
			println(io, ts[i], '\t', xs[i][1], '\t', xs[i][2], '\t', xs[i][3])
		end
	end
	arrival_coordinates[i] = [ts[end]; xs[end]]

	Threads.atomic_add!(trajectories, 1)
	print("\r$(trajectories[]) / $n_trajectories")
end
open("arrival_coordinates.dat", "w") do io
	for coord in arrival_coordinates
		println(io, coord[1], '\t', coord[2], '\t', coord[3], '\t', coord[4])
	end
end

max_stars_repo_name: kaspermeerts/zigzagelectrons

max_issues_repo_name: kaspermeerts/zigzagelectrons

max_forks_repo_name: kaspermeerts/zigzagelectrons


text:
import Pkg

Pkg.status()
using PyCall
pygui(:tk)
using PyPlot

import PBURF

coefficientlistnumcompila=complex([1.,0.,0.,2.])
coefficientlistdencompila=complex([0.,0.,3.,0.])
result=PBURF.plottingBasinsUnivariateRationalFunctions(coefficientlistnumcompila,
    coefficientlistdencompila)

gcf()

PBURF.plottingBasinsUnivariateRationalFunctions(coefficientlistnumcompila,
    coefficientlistdencompila; colorstrategy="iteration")

gcf()

max_stars_repo_name: luisjavierhernandez/PBURF

max_issues_repo_name: luisjavierhernandez/PBURF

max_forks_repo_name: luisjavierhernandez/PBURF


text:
function getOccupancyString(n::Int64)
    return Tuple([bitstring(i)[end-n+1:end] for i = 0:2^n-1])
end

const ROOK_SHIFTS = (8, -8, -1, 1) #N, S, W, E
const BISHOP_SHIFTS = (9, 7, -7, -9) #NE, NW, SE, SW
function getNumberOfSquares(square::UInt64, rook::Bool, d::Int64)
    rook ? mask = ROOK_MASKS[square] : mask = BISHOP_MASKS[square]
    rook ? shifts = ROOK_SHIFTS : shifts = BISHOP_SHIFTS
    c = 0
    while true
        (square << (shifts[d]*(c+1)) & mask) == EMPTY ? break : c+=1
    end
    return c
end

function getNumberOfBits(square::UInt64, rook::Bool)
    return reduce(+, [getNumberOfSquares(square, rook, d) for d=1:4])
end

function bitsGen(rook::Bool)
    bits = Dict{UInt64,Int64}()
    [push!(bits, sq=>getNumberOfBits(sq, rook)) for sq in values(PGN2UINT)]
    return bits
end
const ROOK_BITS = bitsGen(true)
const BISHOP_BITS = bitsGen(false)

function occupancyGen(rook::Bool)
    occupancies = Dict{UInt64, Tuple{Vararg{UInt64}}}()
    rook ? shifts = ROOK_SHIFTS : shifts = BISHOP_SHIFTS
    for sq in values(PGN2UINT)
        square_occs = []
        occs = [getOccupancyString(getNumberOfSquares(sq, rook, i)) for i=1:4]
        for ds in Iterators.product(occs...)
            occ = EMPTY
            [[if ds[j][i] == '1'; occ |= (sq<<(shifts[j]*i)) end for i=1:length(ds[j])] for j=1:4]
            push!(square_occs, occ)
        end
        push!(occupancies, sq=>tuple(square_occs...))
    end
    return occupancies
end
const ORTHO_OCCS = occupancyGen(true)
const DIAGO_OCCS = occupancyGen(false)

function randomMagic()
    return rand(UInt64) & rand(UInt64) & rand(UInt64)
end

function magicSquareGen(square::UInt64, rook::Bool, bits::Int64)
    rook ? occupancies = ORTHO_OCCS : occupancies = DIAGO_OCCS
    rook ? slider = orthoAttack : slider = diagoAttack
    attacks = [slider(square, var) for var in values(occupancies[square])]
    while true
        magic = randomMagic()
        targets = zeros(UInt64, 2^bits)
        collision = false
        for j=1:length(attacks)
            i = ((occupancies[square][j]*magic) >> (64-bits)) + 1
            if targets[i] == EMPTY
                targets[i] = attacks[j]
            elseif targets[i] == attacks[j]
                continue
            else
                collision = true
                break
            end
        end
        collision ? continue : return magic, tuple(targets...)
    end
end

function magicTableGen(rook::Bool)
    rook ? bits = ROOK_BITS : bits = BISHOP_BITS
    magic_table = Dict{UInt64, UInt64}()
    attacks = Dict{UInt64, Tuple{Vararg{UInt64}}}()
    for s in values(PGN2UINT)
        magic, attack = magicSquareGen(s, rook, bits[s])
        push!(magic_table, s=>magic)
        push!(attacks, s=>attack)
    end
    return magic_table, attacks
end
const ROOK_MAGICS, ROOK_X = magicTableGen(true)
const BISHOP_MAGICS, BISHOP_X = magicTableGen(false)

function getMagicAttack(square::UInt64, blockers::UInt64, rook::Bool)
    rook ? magic = ROOK_MAGICS[square] : magic = BISHOP_MAGICS[square]
    rook ? bits = ROOK_BITS[square] : bits = BISHOP_BITS[square]
    rook ? attacks = ROOK_X[square] : attacks = BISHOP_X[square]
    rook ? mask = ROOK_MASKS[square] : mask = BISHOP_MASKS[square]
    i = ((blockers&mask)*magic) >> (64-bits)
    return attacks[i+1]
end


text:
module Demo

export @demo

function julia_code(expr)
    println("Julia code:")
    println("-----------\n")
    println(expr)
    println("\n...")
end

function outcome()
    println("Output:")
    println("-------\n")
end

macro demo(expr)
    s = string(expr)
    quote
        input = chomp(readline(STDIN))
        julia_code($s)
        input = chomp(readline(STDIN))
        outcome()
        println($(esc(expr)))
        println("\n")
    end
end

end


text:
_kernel_wmma_nwmmas()::Int = 10_000

function kernel_wmma_f16(a_dev, b_dev, c_dev, d_dev)
    conf16 = WMMA.Config{16,16,16,Float16}
    conf32 = WMMA.Config{16,16,16,Float32}

    a_frag = WMMA.load_a(pointer(a_dev), 16, WMMA.RowMajor, conf16)
    b_frag = WMMA.load_b(pointer(b_dev), 16, WMMA.ColMajor, conf16)
    c_frag = WMMA.load_c(pointer(c_dev), 16, WMMA.RowMajor, conf32)

    for _ in 1:_kernel_wmma_nwmmas()
        c_frag = WMMA.mma(a_frag, b_frag, c_frag, conf32)
    end

    WMMA.store_d(pointer(d_dev), c_frag, 16, WMMA.RowMajor, conf32)
    return nothing
end

function kernel_wmma_f16_lowlevel(a_dev, b_dev, c_dev, d_dev)
    a_frag = WMMA.llvm_wmma_load_a_col_m16n16k16_global_stride_f16(pointer(a_dev), 16)
    b_frag = WMMA.llvm_wmma_load_b_col_m16n16k16_global_stride_f16(pointer(b_dev), 16)
    c_frag = WMMA.llvm_wmma_load_c_col_m16n16k16_global_stride_f32(pointer(c_dev), 16)

    for _ in 1:_kernel_wmma_nwmmas()
        c_frag = WMMA.llvm_wmma_mma_col_col_m16n16k16_f32_f32(a_frag, b_frag, c_frag)
    end

    WMMA.llvm_wmma_store_d_col_m16n16k16_global_stride_f32(pointer(d_dev), c_frag, 16)
    return nothing
end

function kernel_wmma_int8_lowlevel(a_dev, b_dev, c_dev, d_dev)
    a_frag = WMMA.llvm_wmma_load_a_col_m16n16k16_global_stride_s8(pointer(a_dev), 16)
    b_frag = WMMA.llvm_wmma_load_b_col_m16n16k16_global_stride_s8(pointer(b_dev), 16)
    c_frag = WMMA.llvm_wmma_load_c_col_m16n16k16_global_stride_s32(pointer(c_dev), 16)

    for _ in 1:_kernel_wmma_nwmmas()
        c_frag = WMMA.llvm_wmma_mma_col_col_m16n16k16_s8(a_frag, b_frag, c_frag)
    end

    WMMA.llvm_wmma_store_d_col_m16n16k16_global_stride_s32(pointer(d_dev), c_frag, 16)
    return nothing
end

function kernel_wmma_tf32_lowlevel(a_dev, b_dev, c_dev, d_dev)
    a_frag = WMMA.llvm_wmma_load_a_col_m16n16k8_global_stride_tf32(pointer(a_dev), 16)
    b_frag = WMMA.llvm_wmma_load_b_col_m16n16k8_global_stride_tf32(pointer(b_dev), 8)
    c_frag = WMMA.llvm_wmma_load_c_col_m16n16k8_global_stride_f32(pointer(c_dev), 16)

    for _ in 1:_kernel_wmma_nwmmas()
        c_frag = WMMA.llvm_wmma_mma_col_col_m16n16k8_tf32(a_frag, b_frag, c_frag)
    end

    WMMA.llvm_wmma_store_d_col_m16n16k8_global_stride_f32(pointer(d_dev), c_frag, 16)
    return nothing
end

function kernel_wmma_bf16_lowlevel(a_dev, b_dev, c_dev, d_dev)
    a_frag = WMMA.llvm_wmma_load_a_col_m16n16k16_global_stride_bf16(pointer(a_dev), 16)
    b_frag = WMMA.llvm_wmma_load_b_col_m16n16k16_global_stride_bf16(pointer(b_dev), 16)
    c_frag = WMMA.llvm_wmma_load_c_col_m16n16k16_global_stride_f32(pointer(c_dev), 16)

    for _ in 1:_kernel_wmma_nwmmas()
        c_frag = WMMA.llvm_wmma_mma_col_col_m16n16k16_bf16(a_frag, b_frag, c_frag)
    end

    WMMA.llvm_wmma_store_d_col_m16n16k16_global_stride_f32(pointer(d_dev), c_frag, 16)
    return nothing
end

"""
    peakflops_gpu_wmmas()
Tries to estimate the peak performance of a GPU in TFLOP/s by measuring the time
it takes to perform `_kernel_wmma_nwmmas()` many WMMAs on Tensor Cores.

**Keyword arguments:**
* `device` (default: `CUDA.device()`): CUDA device to be used.
* `dtype` (default: `Float16`): element type of the matrices. We currently only support `Float16` (`Int8`, `:TensorFloat32`, `:BFloat16`, and `Float64` might or might not work).
* `nkernel` (default: `10`): number of kernel calls that make up one benchmarking sample.
* `nbench` (default: `5`): number of measurements to be performed the best of which is used for the TFLOP/s computation.
* `threads` (default: max. threads per block): how many threads to use per block (part of the kernel launch configuration).
* `blocks` (default: `2048`): how many blocks to use (part of the kernel launch configuration).
* `verbose` (default: `true`): toggle printing.
* `io` (default: `stdout`): set the stream where the results should be printed.
"""
function peakflops_gpu_wmmas(;
    device::CuDevice=CUDA.device(),
    blocks=2048,
    threads=CUDA.attribute(device, CUDA.CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK),
    nbench=5,
    nkernel=10,
    verbose=true,
    dtype=Float16,
    io::IO=stdout,
)
    device!(device) do
        if Symbol(dtype) == :Float16
            m = n = k = 16
            dtype_a = dtype_b = Float16
            dtype_c = dtype_d = Float32
        elseif Symbol(dtype) == :Float16_lowlevel
            m = n = k = 16
            dtype_a = dtype_b = Float16
            dtype_c = dtype_d = Float32
        elseif Symbol(dtype) == :Int8
            if pkgversion(CUDA) < v"3.8.6"
                error("At the time of writing, CUDA#master is required for Int8 WMMA support.")
            end
            m = n = k = 16
            dtype_a = dtype_b = Int8
            dtype_c = dtype_d = Int32
        elseif Symbol(dtype) in (:Float32, :TensorFloat32, :TF32)
            # requires CUDA.jl PR 1419
            if VERSION < v"1.8.0-"
                error("Julia >= 1.8.0 required for `:TensorFloat32`.")
            end
            @warn(
                "Not officially supported by GPUInspector.jl (yet). Expect errors! Requires https://github.com/JuliaGPU/CUDA.jl/pull/1419."
            )
            m = n = 16
            k = 8
            dtype_a = dtype_b = Float32
            dtype_c = dtype_d = Float32
        elseif Symbol(dtype) == :Float64
            # requires CUDA.jl PR 1426
            if VERSION < v"1.8.0-"
                error("Julia >= 1.8.0 required for `Float64`.")
            end
            @warn(
                "Not officially supported by GPUInspector.jl (yet). Expect errors! Requires https://github.com/JuliaGPU/CUDA.jl/pull/1426."
            )
            m = n = 8
            k = 4
            dtype_a = dtype_b = Float64
            dtype_c = dtype_d = Float64
        elseif Symbol(dtype) in (:BFloat16, :BF16)
            # requires CUDA.jl PR 1425
            if VERSION < v"1.8.0-"
                error("Julia >= 1.8.0 required for `:BFloat16`.")
            end
            @warn(
                "Not officially supported by GPUInspector.jl (yet). Expect errors! Requires https://github.com/JuliaGPU/CUDA.jl/pull/1425."
            )
            m = n = k = 16
            dtype_a = dtype_b = BFloat16
            dtype_c = dtype_d = Float32
        else
            throw(ArgumentError("Unsupported dtype."))
        end
        d_a = CUDA.rand(dtype_a, m, k)
        d_b = CUDA.rand(dtype_b, k, n)
        d_c = CUDA.rand(dtype_c, m, n)
        d_d = CUDA.zeros(dtype_d, m, n)

        if Symbol(dtype) == :Float16
            kernel = @cuda launch = false kernel_wmma_f16(d_a, d_b, d_c, d_d)
        elseif Symbol(dtype) == :Float16_lowlevel
            kernel = @cuda launch = false kernel_wmma_f16_lowlevel(d_a, d_b, d_c, d_d)
        elseif Symbol(dtype) == :Int8
            kernel = @cuda launch = false kernel_wmma_int8_lowlevel(d_a, d_b, d_c, d_d)
        elseif Symbol(dtype) in (:Float32, :TensorFloat32, :TF32)
            kernel = @cuda launch = false kernel_wmma_tf32_lowlevel(d_a, d_b, d_c, d_d)
        elseif Symbol(dtype) in (:BFloat16, :BF16)
            kernel = @cuda launch = false kernel_wmma_bf16_lowlevel(d_a, d_b, d_c, d_d)
        else
            throw(ArgumentError("Unsupported dtype."))
        end
        warpsize = CUDA.attribute(device, CUDA.CU_DEVICE_ATTRIBUTE_WARP_SIZE)
        # @show threads
        # @show blocks
        # @show warpsize

        # warm-up
        CUDA.@elapsed kernel(d_a, d_b, d_c, d_d; threads=threads, blocks=blocks)

        t = Inf
        for _ in 1:nbench
            Δt = CUDA.@elapsed begin
                for _ in 1:nkernel
                    # kernel(d_a, d_b, d_c, d_d; threads=threads, blocks=blocks)
                    kernel(d_a, d_b, d_c, d_d; threads=threads, blocks=blocks)
                end
            end
            t = min(t, Δt)
        end
        t /= nkernel

        nwarps = threads / warpsize
        flopcount = 2.0 * m * n * k * blocks * nwarps * _kernel_wmma_nwmmas()
        flops = (flopcount * 1e-12) / t

        if verbose
            printstyled(io,
                "Peakflops ($(Symbol(dtype) == :Int8 ? "TOP" : "TFLOP")/s):\n"; bold=true
            )
            print(io," └ max: ")
            printstyled(io,round(flops; digits=2), "\n"; color=:green, bold=true)
        end
        return flops
    end
end


text:
"""
# Loc

a two-dimensional point (x, y). represents a location on a Quoridor
board

"""
immutable Loc
    x::Int
    y::Int
end

(+)(l1::Loc, l2::Loc) = Loc(l1.x + l2.x, l1.y + l2.y)
(-)(l1::Loc, l2::Loc) = Loc(l1.x - l2.x, l1.y - l2.y)

# TODO bounds check on these somewhere?
left(l::Loc) = l - Loc(1,0)
right(l::Loc) = l + Loc(1,0)
above(l::Loc) = l + Loc(0,1)
below(l::Loc) = l - Loc(0,1)

const DIRS = Vector4{Function}(above, below, left, right)
const VERTICAL = Vector2{Function}(above, below)
const HORIZONTAL = Vector2{Function}(left, right)

const PERPENDICULAR = Dict(left=>VERTICAL, right=>VERTICAL,
                           above=>HORIZONTAL, below=>HORIZONTAL)

isless(l1::Loc, l2::Loc) = isless(l1.x + l1.y, l2.x + l2.y)
mean(l1::Loc, l2::Loc) = Loc((l1.x + l2.x) / 2, (l1.y + l2.y) / 2)

convert(::Type{Loc}, tup::Tuple{Int, Int}) = Loc(tup[1], tup[2])

show(io::IO, loc::Loc) = write(io, "Loc($(loc.x), $(loc.y))")

isoutofbounds(loc::Loc) = (loc.x < 1 || loc.x > 9 ||
                           loc.y < 1 || loc.y > 9)

inbounds(loc::Loc) = !isoutofbounds(loc)

function adjacent(loc::Loc)
    Vector4{Loc}(above(loc), below(loc), right(loc), left(loc))
end


text:
# abstract AbstractState
# type State <: AbstractState
#     s::Vector{Complex}
# end
#
# function call(state::State)

max_stars_repo_name: Roger-luo/QuCmp

max_issues_repo_name: Roger-luo/QuCmp

max_forks_repo_name: Roger-luo/QuCmp


text:
"""[intended for internal use only]
Make sure an Experiment's data are numerical
and do not include missing values, NaNs or infs.
"""
function _assert_clean_data(e::Experiment)
	# Column type cannot be used as for instance a Union type
	# could support missing values but the selected data subset
	# might contain only numbers

	# This excludes strings and missings (and more)
	hasnumbers = getdata(e) |>
					x -> isa.(x, Number) |>
					eachcol |>
					x -> all.(x) |>
					all
	@assert hasnumbers "Selected data include non-numeric values."

	# Exclude NaNs
	hasnonans = getdata(e) |>
					x -> isnan.(x) |>
					eachcol |>
					x -> any.(x) |>
					any |> ~
	@assert hasnonans "Selected data include NaNs."

	# Exclude Inf
	hasnoinf = getdata(e) |>
					x -> isinf.(x) |>
					eachcol |>
					x -> any.(x) |>
					any |> ~
	@assert hasnoinf "Selected data include Inf values."
end

"""[intended for internal use only]
Convert all selected data columns to floats
"""
function _data_to_float!(e::Experiment)
    # Make sure all values are numbers
    @assert all( [x <: Number for x in eltype.(eachcol(getdata(e)))] )
    # Convert each column to floats
    for colname in names(getdata(e))
        e.data[!,colname] = float.(e.data[:,colname])
    end
end


text:
# R2 attributed

using Distributions

n = 10000
x1 = rand(Normal(), n)
x2 = rand(Normal(), n) .+ .2 * x1
x3 = 1000*(rand(Normal(), n) .- .2 * x1 .+ .2 * x2)

x2 = rand(Normal(), n)
x3 = rand(Normal(), n)

u = rand(Normal(), n) .* 2

y = 2 .+ 1.5*x1 .+ 2.7*x2 .- .5* x3 + u

constant = fill(1,n)
X = hcat(constant, x1, x2, x3)

β = inv(X'X) * X'y

ŷ = X*β

s² = ((y-ŷ)' * (y-ŷ)) / (n-4)

seβ₁ = √(s² * inv(x1' * x1))
seβ₂ = √(s² * inv(x2' * x2))
seβ₃ = √(s² * inv(x3' * x3))

√(s² .* inv(X'X))

ccdf(Chisq(4), s²)

SS(x) = sum((x .- mean(x)).^2)

# Sum of Squares Total
SST = SS(y)

# Sum of Squares Explained
SSR = SS(ŷ)

# Sum of Squares Error
SSE = SS(y-ŷ)

# Degrees of freedom model
DFM = 4-1

# Degrees of freedom total
DFT = n-4

# Mean of squares of regression
MSR = SSR/DFM

# Mean of squares of total
MST = SSR/DFT

F = MSR/MST

R2 = SSR/SST

ŷ₁ = x1*β[2]
ŷ₂ = x2*β[3]
ŷ₃ = x3*β[4]

SSR1 = SS(ŷ₁)
SSR2 = SS(ŷ₂)
SSR3 = SS(ŷ₃)

SSR
SSRp = SSR1 + SSR2 + SSR3

SSR1p = SSR1/SSRp
SSR2p = SSR2/SSRp
SSR3p = SSR3/SSRp

##### Normalization Approach
x1 = x1/var(x1)^.5
x2 = x2/var(x2)^.5
x3 = x3/var(x3)^.5

X = hcat(constant, x1, x2, x3)

β = inv(X'X) * X'y

# Asserting Independence - generally vioated
# var(y) = β[2]^2*var(x1) + β[3]^2*var(x2) + β[4]^2*var(x3)

# Since we normalized
denom = β[2]^2 + β[3]^2 + β[4]^2

β[2]^2/denom
β[3]^2/denom
β[4]^2/denom


text:
module Classes
using  ExpressionPatterns.Dispatch
export @class

include("classprep.jl")

@macromethod class(:T{name, Symbol}, block) esc(:(@class ($name) <: ($Object) ($block)))

@macromethod class(:T{name, Symbol} <: superclass, begin *{exprs} end) begin
  parameters, constructor = classinfo(name, exprs)

  esc(
  quote
    mutable struct $name <: $superclass
      $(parameters...)

      $name(::$ClassId) = new()
    end
    $constructor
    $name
  end)
end

function classinfo(name, exprs)
  info = ClassInfo()
  for expr in exprs
    interpret_expr!(name, info, expr)
  end
  parameters  = info.parameters
  constructor = make_constructor(name, info, info.constructor)

  parameters, constructor
end

@metafunction interpret_expr!(name, info, (name(*{args}) = expr)) =
  interpret_expr!(info, :(function constructor($(args...)) $expr end))

@metafunction interpret_expr!(name, info, function name(*{args}) *{statements} end) begin
  if info.constructor.arguments == :undef
    info.constructor.arguments  = args
    info.constructor.statements = statements
  else
    error("This class implementations only accepts one inner constructor!")
  end
end

@metafunction interpret_expr!(name, info, function f(*{xs}) *{statements} end) =
  interpret_expr!(info, :($f($(xs...)) = begin $(statements...) end))

@metafunction interpret_expr!(name, info, f(*{xs}) = expr) begin
  if !(:(($f)::($Function)) in info.parameters)
    push!(info.parameters, :(($f)::($Method)))
    push!(info.constructor.method_defs, :($f($(xs...)) = $expr; this.$f = $Method{$name}($f)))
    push!(info.constructor.new_args, f)
  end
end

@metafunction interpret_expr!(name, info, (:T{x,Symbol}=a)) =
  interpret_expr!(info, :($x::Any=$a))

@metafunction interpret_expr!(name, info, (:T{x,Symbol}::T=a)) begin
  push!(info.parameters, :($x::$T))
  push!(info.constructor.param_defs, :($x = $a; this.$x = $x))
  push!(info.constructor.new_args, x)
end

@metafunction interpret_expr!(name, info, (:T{x,Symbol})) =
  interpret_expr!(info, :($x::Any))

@metafunction interpret_expr!(name, info, (:T{x,Symbol}::T)) begin
  push!(info.parameters, :($x::$T))
  push!(info.undefined,  :($x::$T))
  push!(info.constructor.new_args, x)
end

@metafunction make_def(x::T) = :(this.$x = $x)

# usage

@class Point begin
  x::Number
  y::Number

  add(p) = Point(this.x + p.x, this.y + p.y)
  distanceto(p) = sqrt((this.x-p.x)^2 + (this.y-p.y)^2)
end

Point(1,2).add(Point(1,2)).distanceto(Point(2,4))

@class VectorObject begin
  elements::Vector=[]

  push(a) = (push!(this.elements, a); this)
end

V = VectorObject()
V.push(1).push(2).push(3)

abstract type Shape <: Object end

@class Square <: Shape begin
  a::Point
  b::Point
  c::Point
  d::Point

  function Square(a,b,c,d)
    @assert a.distanceto(b) ==
            b.distanceto(c) ==
            c.distanceto(d)

    this.a = a
    this.b = b
    this.c = c
    this.d = d
  end
end
points(xs) = map(x->Point(x...), xs)
Square(points([(0,0),(2,0),(2,2),(0,2)])...)
end