using Plots
using ForwardDiff
using Roots
include("helper_functions.jl")

begin
    λ1 = 0.5
    λ2 = 1
    α = 1
    μ(α) = 1 / (1.5 * (α))
    ν(α) = 1 / (0.7 * (α))
    parameters = (λ1, λ2, α, μ, ν)
    q = 0.5
    qs = 0:0.001:1
end

# Test 1: Do the derivatives agree?
# function adcost(p, q, λ1, λ2, α, μ, ν)
#     (λ1 * p) / (λ1 * p - ν(α) + λ2 * q)^2 -
#     1/(μ(α) - λ1 + λ1 * p) -
#     1/(λ1 * p - ν(α) + λ2 * q) +
#     (λ1 * (p - 1))/(μ(α) - λ1 + λ1 * p)^2
# end

# dcost = ∂cost_∂p

ps = 0:0.001:1
plot(ps, ∂cost_∂p.(ps, q, parameters...))
plot!(ps, diff_cost.(ps, q, parameters...))
plot!(ylims = (-500, 500))

# @benchmark adcost(0.5, q, parameters...)
# @benchmark dcost(0.5, q, parameters...)

# Test 2: Do the points where derivative is zero agree? (aka pstar)
sol, _ = first.(find_valid_roots(p -> ∂cost_∂p(p, q, parameters...)))
# function apstar(q, λ1, λ2, α, μ, ν)
#     p = (λ1*Complex(μ(α)*(ν(α) - λ2*q))^(1/2) - μ(α)*Complex(μ(α)*(ν(α) - λ2*q))^(1/2) - ν(α)*Complex(μ(α)*(ν(α) - λ2*q))^(1/2) - λ1*ν(α) + 2*μ(α)*ν(α) + λ2*q*Complex(μ(α)*(ν(α) - λ2*q))^(1/2) + λ1*λ2*q - 2*λ2*μ(α)*q)/(λ1*μ(α) - λ1*ν(α) + λ1*λ2*q) |> real
#     return clamp(p, 0, 1)
# end
asol = pstar(q, parameters...)

diff_cost(sol, q, parameters...)
diff_cost(asol, q, parameters...)

# Test 3: Do the pstars agree?
# function pstar(q, λ1, λ2, α, μ, ν)
#     roots = find_valid_roots(p -> ∂cost_∂p(p, q, λ1, λ2, α, μ, ν))
#     if length(roots[1]) > 1
#         root_values, root_is_minimum =
#             classify_roots(p -> ∂cost_∂p(p, q, λ1, λ2, α, μ, ν), roots)
#         return (root_values[root_is_minimum])[1]
#     else
#         return (roots[1])[1]
#     end
# end

qs = 0:0.0001:1
ps = numerical_pstar.(qs, parameters...)
a_ps = pstar.(qs, parameters...)

response_plot = plot(qs, ps, label = "autodiff")
plot!(qs, a_ps, label = "analytic")
plot_assumptions!(response_plot, [ps, a_ps], qs, parameters...)
plot!(response_plot, legend = :outertopleft, size = (690, 400))

costs = cost.(ps, qs, parameters...)
a_costs = cost.(a_ps, qs, parameters...)

cost_plot = plot(qs, costs, label = "autodiff")
plot!(qs, a_costs, label = "analytic")
plot!(ylims = (-50, 50))
plot_assumptions!(cost_plot, [ps, a_ps], qs, parameters..., height = 10)

traffic_plot = plot(
    title = "Traffic at each server for varying q",
    xlabel = "Opponent probability (q)",
    ylabel = "Intensity of arrivals",
    legend = :topleft,
    xlims = (0, 1),
    ylims = ylims = (0, Inf)
    )

plot!(qs, λ1 .* a_ps .+ λ2 .* qs, label = "ν(α)")
hline!([ν.(α)], label = "ν(α) limit", ls = :dash, lw = 1.5)

plot!(qs, λ1 .* (1 .- a_ps), label = "μ(α)")
hline!([μ(α)], label = "μ(α) limit", ls = :dash, lw = 1.5)

intersection_index = last(findmax(λ1 .* a_ps .+ λ2 .* qs .> ν(α)))
x_intersection = qs[intersection_index]
y_intersection_max = maximum([
    (λ1 .* (1 .- a_ps))[intersection_index],
    (λ1 .* a_ps .+ λ2 .* qs)[intersection_index]]
    )
plot!([x_intersection, x_intersection], [0, y_intersection_max], lw = 1, lc = :maroon, label = "")

traffic_plot_anim = @animate for α ∈ vcat(collect(0.5:0.005:2.1), collect(reverse(0.5:0.005:2.1)))
    a_ps = apstar.(qs, (λ1, λ2, α, μ, ν)...)
    traffic_plot = plot(
    title = "Traffic at each server for varying qs at α = $(dp_string(α, 3))",
    # xlabel = "Opponent probability (q)",
    ylabel = "Intensity of arrivals",
    legend = :topleft,
    xlims = (0, 1),
    ylims = ylims = (0, Inf),
    xticks = 0:0.2:1,
    )

    plot!(qs, λ1 .* a_ps .+ λ2 .* qs, label = "ν(α)")
    hline!([ν.(α)], label = "ν(α) limit", ls = :dash, lw = 1.5)

    plot!(qs, λ1 .* (1 .- a_ps), label = "μ(α)")
    hline!([μ(α)], label = "μ(α) limit", ls = :dash, lw = 1.5)

    does_intersect, intersection_index = findmax(λ1 .* a_ps .+ λ2 .* qs .> ν(α))
    x_intersection = qs[intersection_index]
    y_intersection_max = maximum([
        (λ1 .* (1 .- a_ps))[intersection_index],
        (λ1 .* a_ps .+ λ2 .* qs)[intersection_index]]
        )
    plot!([x_intersection, x_intersection], [0, y_intersection_max], ls = :dash, lw = 1, lc = :black, label = "")

    pstar_plot = plot(
        qs,
        a_ps,
        lc = :black,
        label = "Best Response (p⋆)",
        xlabel = "Opponent probability (q)",
        xticks = 0:0.2:1,
        ylabel = "Response probability (p)",
        xlims = (0, 1),
        ylims = (0, 1))
    plot!(
        [x_intersection, x_intersection],
        [0, a_ps[intersection_index]],
        label = "",
        ls = :dash,
        lc = :black)

    if does_intersect
        plot!(
            qs[1:intersection_index],
            zeros(length(qs[1:intersection_index])),
            fillrange = a_ps[1:intersection_index],
            fc = :green,
            fa = 0.35,
            label = "BR Feasibile")
        plot!(
            qs[intersection_index + 1:end],
            zeros(length(qs[intersection_index + 1:end])),
            fillrange = a_ps[intersection_index + 1:end],
            fc = :red,
            fa = 0.35,
            label = "BR Infeasibile")
    else
        plot!(
            qs,
            zeros(length(qs)),
            fillrange = a_ps,
            fc = :green,
            fa = 0.35,
            label = "BR Feasibile")
        # Have to put something here so it doesn't change the legend
        plot!(
            [1, 2],
            zeros(2),
            fillrange = [2, 3],
            fc = :red,
            fa = 0.35,
            label = "BR Infeasibile")
    end

    plot(traffic_plot, pstar_plot, layout = (2, 1), size = (800, 800))
end

gif(traffic_plot_anim, "traffic_plot_anim.gif", fps = 30)


# Operators and mathematical functions overloaded for UnitDuals

# Binary operators for UnitDuals
+(a::UnitDual, b::UnitDual) = UnitDual(value(a) + value(b), add_partials(partials(a), partials(b)))
-(a::UnitDual, b::UnitDual) = UnitDual(value(a) - value(b), sub_partials(partials(a), partials(b)))
-(a::UnitDual) = UnitDual(-value(a), neg_partials(partials(a)))
*(a::UnitDual, b::UnitDual) = UnitDual(value(a)*value(b), mul_partials(partials(a), partials(b), value(b), value(a)))
/(a::UnitDual, b::UnitDual) = UnitDual(value(a)/value(b), mul_partials(partials(a), partials(b), inv(value(b)),-(value(a)/(value(b)*value(b)))))

# Binary operators for UnitDuals & Number
+(a::UnitDual, b::Number) = UnitDual(value(a) + b, partials(a))
+(a::Number, b::UnitDual) = b + a
-(a::UnitDual, b::Number) = UnitDual(value(a) - b, partials(a))
-(a::Number, b::UnitDual) = a + (-b)
*(a::UnitDual, b::Number) = UnitDual(value(a)*b, scale_partials(partials(a), b))
*(a::Number, b::UnitDual) = b*a
/(a::UnitDual, b::Number) = UnitDual(value(a)/b, scale_partials(partials(a), 1/b))
/(a::Number, b::UnitDual) = UnitDual(a/value(b), scale_partials(partials(b), -a/value(b)^2))
inv(a::UnitDual) =  UnitDual(inv(value(a)), scale_partials(partials(a), -1/value(a)^2))
^(a::UnitDual, b::Number) = UnitDual(value(a)^b, scale_partials(partials(a), b*value(a)^(b - 1)))

# Unary mathematical functions using DiffRules
for fun in (:sqrt, :cbrt, :abs, :abs2, :log, :log10, :log2, :exp, :sin, :cos, :tan, :sec, :csc, :cot,
            :acos, :atan, :asec, :acsc, :acot)
    f = :($fun(value(a)))
    fp = DiffRules.diffrule(:Base, fun, :(value(a)))
    @eval function $fun(a::UnitDual)
        UnitDual($f, scale_partials(partials(a), $fp))
    end
end

# Logical operators for UnitDuals
==(a::UnitDual, b::UnitDual) = value(a) == value(b)
==(a::UnitDual, b::Number) = value(a) == b
==(a::Number, b::UnitDual) = a == value(b)
isless(a::UnitDual, b::UnitDual) = isless(value(a), value(b))
isless(a::UnitDual, b::Number) = isless(value(a), b)
isless(a::Number, b::UnitDual) = isless(a, value(b))
isapprox(a::UnitDual, b::UnitDual) = value(a) ≈ value(b) && all(partials(a)[i] ≈ partials(b)[i] for i in 1:length(partials(a)))
isapprox(a::UnitDual, b::Number) = value(a) ≈ b
isapprox(a::Number, b::UnitDual) = a ≈ value(b)


"""
Normal Mixture of `k` components via an online EM algorithm.  `start` is a keyword
argument specifying the initial parameters.

```julia
o = NormalMix(2, LearningRate(); start = MixtureModel(Normal, [(0, 1), (3, 1)]))
mean(o)
var(o)
std(o)
```
"""
type NormalMix{W<:Weight} <: DistributionStat{ScalarInput}
    value::Ds.MixtureModel{Ds.Univariate, Ds.Continuous, Ds.Normal{Float64}}
    s1::VecF
    s2::VecF
    s3::VecF
    w::VecF
    μ::VecF
    σ2::VecF
    weight::W
end
function NormalMix(k::Integer, wgt::Weight = LearningRate();
        start = Ds.MixtureModel(Ds.Normal{Float64}[Ds.Normal(j-1, 10) for j in 1:k])
    )
    NormalMix(
        start,
        ones(k) / k, zeros(k), zeros(k),  # s1, s2, s3
        zeros(k), collect(1.:k), fill(10.0, k),  # w, μ, σ2
        wgt
    )
end
function NormalMix(y::AVec, k::Integer, wgt::Weight = LearningRate(); kw...)
    o = NormalMix(k, wgt; kw...)
    fit!(o, y)
    o
end
NormalMix(k::Integer, y::AVec, wgt::Weight = LearningRate(); kw...) = NormalMix(y, k, wgt; kw...)
Ds.componentwise_pdf(o::NormalMix, y) = Ds.componentwise_pdf(value(o), y)
Ds.ncomponents(o::NormalMix) = Ds.ncomponents(value(o))
Ds.component(o::NormalMix, j) = Ds.component(value(o), j)
Ds.probs(o::NormalMix) = Ds.probs(value(o))
Ds.pdf(o::NormalMix, y) = Ds.pdf(value(o), y)
Ds.cdf(o::NormalMix, y) = Ds.cdf(value(o), y)
Base.mean(o::NormalMix) = mean(value(o))
Base.var(o::NormalMix) = var(value(o))
Base.std(o::NormalMix) = std(value(o))
function Base.show(io::IO, o::NormalMix)
    printheader(io, "NormalMix (k = $(Ds.ncomponents(o)))")
    print_value_and_nobs(io, o)
end
function value(o::NormalMix)
    # try
        vec = map((u,v) -> Ds.Normal{Float64}(u, sqrt(v)), o.μ, o.σ2)
        o.value = Ds.MixtureModel(vec, o.s1)
    # catch
    #     println(sqrt(o.σ2))
    #     println(o.μ)
    #     println(o.s1)
    #     error("Algorithm possibly diverging, nobs = $(nobs(o))")
    # end
end
function _fit!(o::NormalMix, y::Real, γ::Float64)
    k = length(o.μ)
    for j in 1:k
        σinv = 1.0 / sqrt(o.σ2[j])
        o.w[j] = o.s1[j] * σinv * exp(-.5 * σinv * σinv * (y - o.μ[j]) ^ 2)
    end
    sum1 = sum(o.w)
    for j in 1:k
        o.w[j] /= sum1
        o.s1[j] = smooth(o.s1[j], o.w[j], γ)
        o.s2[j] = smooth(o.s2[j], o.w[j] * y, γ)
        o.s3[j] = smooth(o.s3[j], o.w[j] * y * y, γ)
    end
    sum2 = sum(o.s1)
    for j in 1:k
        o.μ[j] = o.s2[j] / o.s1[j]
        o.σ2[j] = (o.s3[j] - o.s2[j] ^ 2 / o.s1[j]) / o.s1[j]
        o.s1[j] /= sum2
        if o.σ2[j] <= _ϵ
            o.σ2 = ones(k)
        end
    end
    o
end
function _fitbatch!{T<:Real}(o::NormalMix, y::AVec{T}, γ::Float64)
    n = length(y)
    k = length(o.μ)
    s1 = copy(o.s1)
    s2 = copy(o.s2)
    s3 = copy(o.s3)

    for yi in y
        for j in 1:k
            σinv = 1.0 / sqrt(o.σ2[j])
            o.w[j] = s1[j] * σinv * exp(-.5 * σinv * σinv * (yi - o.μ[j]) ^ 2)
        end
        sum1 = sum(o.w)
        for j in 1:k
            o.w[j] /= sum1
            if yi == y[1]
                o.s1[j] = smooth(o.s1[j], o.w[j] / n, γ)
                o.s2[j] = smooth(o.s2[j], o.w[j] * yi / n, γ)
                o.s3[j] = smooth(o.s3[j], o.w[j] * yi * yi / n, γ)
            else
                o.s1[j] += γ * o.w[j] / n
                o.s2[j] += γ * o.w[j] * yi / n
                o.s3[j] += γ * o.w[j] * yi * yi / n
            end
        end
    end
    sum2 = sum(o.s1)
    for j in 1:k
        o.μ[j] = o.s2[j] / o.s1[j]
        o.σ2[j] = (o.s3[j] - o.s2[j] ^ 2 / o.s1[j]) / o.s1[j]
        o.s1[j] /= sum2
        if o.σ2[j] <= _ϵ
            o.σ2 = ones(k)
        end
    end
    o
end

# Quantiles.  Starting values based on Normal distribution.
function Base.quantile(o::NormalMix, τ::Real;
        start = quantile(Ds.Normal(mean(o), std(o)), τ), maxit = 20, tol = .001
    )
    @assert 0 < τ < 1
    θ = start
    for i in 1:maxit
        θ += (τ - Ds.cdf(o, θ)) / Ds.pdf(o, θ)
        abs(Ds.cdf(o, θ) - τ) < tol && break
    end
    return θ
end
function Base.quantile{T<:Real}(o::NormalMix, τ::Vector{T};
        start = start = quantile(Ds.Normal(mean(o), std(o)), τ), kw...
    )
    Float64[quantile(o, τ[j]; start = start[j], kw...) for j in 1:length(τ)]
end

# #-----------------------------------------------------------------# Multivariate
# type MvNormalMix{W<:Weight} <: OnlineStat
#     value::Ds.MixtureModel{Ds.Multivariate, Ds.Continuous, Ds.MvNormal}
#     s1::VecF
#     s2::MatF
#     s3::Array{Float64, 3}
#     w::VecF
#     μ::MatF
#     V::Array{Float64, 3}
#     weight::W
#     n::Int
#     nups::Int
# end
# function MvNormalMix(p::Integer, k::Integer, wgt::Weight = LearningRate();
#         start = Ds.MixtureModel(Ds.MvNormal[Ds.MvNormal(ones(p) + j, 10 * eye(p)) for j in 1:k])
#     )
#     V = zeros(p, p, k)
#     # for j in 1:k
#     #     V[:, :, j] = cov(Ds.component(start, j))
#     # end
#
#     MvNormalMix(start,
#         ones(k) / k, zeros(p, k), zeros(p, p, k),  # s1, s2, s3
#         zeros(k), zeros(p, k), V,            # w, μ, V
#         wgt, 0, 0
#     )
# end
# function MvNormalMix{T<:Real}(y::AMat{T}, k::Integer, wgt::Weight = LearningRate(); kw...)
#     o = MvNormalMix(size(y, 2), k, wgt; kw...)
#     fit!(o, y)
#     o
# end
#
# function fit!{T<:Real}(o::MvNormalMix, y::AVec{T})
#     k = length(o.w)
#     γ = weight!(o, 1)
#     for j in 1:k
#         o.w[j] = o.s1[j] * Ds.pdf(Ds.component(o.value, j), y)
#     end
#     sum1 = sum(o.w)
#     for j in 1:k
#         o.w[j] /= sum1
#         o.s1[j] = smooth(o.s1[j], o.w[j], γ)
#         smooth!(o.s2[:, j], o.w[j] * y, γ)
#         smooth!(o.s3[:, :, j], o.w[j] * y * y', γ)
#
#         o.μ[:, j] = o.s2[j] / o.s1[j]
#         o.V[:, :, j] = (o.s3[:, :, j] - o.s2[j] * o.s2[j]' / o.s1[j]) / o.s1[j]
#     end
#
#     sum2 = sum(o.s1)
#     for j in 1:k
#         o.s1[j] /= sum2
#     end
#
#     o.value = Ds.MixtureModel(
#         Ds.MvNormal[
#             Ds.MvNormal(o.μ[:,j], o.V[:,:,j] + .001 * eye(length(o.μ[:,j]))) for j in 1:k
#         ],
#         o.s1 / sum(o.s1)
#     )
# end
#


type SFWParameters{G <: AbstractFloat}
    a::G
    d::G
    e::G
    Ec::G
    rlim::G
    b::G
    q::G
    Jb::G
    ρs::G
    rs::G
    α::G
    β::G
    γ::G

    #secondary parameters
    Φs::G
    xlim::G
    Φ0::G
    Φlim::G

    adjEc::G
    adjJb::G
end

function SFWParameters(a::G, d::G, e::G, Ec::G, rlim::G, b::G, q::G, Jb::G, ρs::G, rs::G, α::G, β::G, γ::G; x0::G = 10.0^-8) where G <: AbstractFloat

    Φs = ρs * rs ^ 2
    xlim = rlim / rs
    adjEc = Φs * Ec
    adjJb = Jb * rs * sqrt(Φs)

    p = SFWParameters(a, d, e, Ec, rlim, b, q, Jb, ρs, rs, α, β, γ, Φs, xlim, 0.0, 0.0, adjEc, adjJb)

    p.Φ0 = gravitational_potential0(x0, p)
    p.Φlim = gravitational_potential(xlim, p)
    return p
end


"""
    InstantiateEvaluateCovariance(cℓ::AbstractCℓ, ConvDens::AbstractConvolvedDensity,
    cosmogrid::CosmologicalGrid, ProbeA::String, ProbeB::String)

This function evaluates and returns the [`aₗₘCovariance`](@ref), according to the following
formula:
```math
\\Sigma_{i j}^{\\mathrm{AB}}(\\ell)=\\sqrt{\\frac{2}{(2 \\ell+1) \\Delta \\ell
f_{\\mathrm{sky}}}}\\left(C_{i j}^{\\mathrm{AB}}(\\ell)+N_{i j}^{\\mathrm{AB}}(\\ell)
\\right),
```
where ``\\mathrm{A}`` and ``\\mathrm{B}`` are the probes, ``\\mathrm{i}`` and
``\\mathrm{j}`` are the tomographic bins, ``\\ell`` and ``\\Delta\\ell`` are respectively the
central value and width of the considered multipole bin, ``f_{\\mathrm{sky}}`` is the sky
fraction covered by the field considered, ``N_{i j}^{\\mathrm{AB}}`` is the noise matrix.
"""
function InstantiateEvaluateCovariance(cℓ::AbstractCℓ, ConvDens::AbstractConvolvedDensity,
    cosmogrid::CosmologicalGrid, ProbeA::String, ProbeB::String)
    Cov = aₗₘCovariance()
    Cov.Cℓ = cℓ
    Cov.Noise      = zeros(size(cℓ.CℓArray))
    Cov.Covariance = zeros(size(cℓ.CℓArray))
    EvaluateNoise!(Cov, ConvDens, ProbeA::String, ProbeB::String)
    EvaluateCovariance!(Cov, cosmogrid)
    InvertCovariance!(Cov)
    return Cov
end

function EvaluateNoise!(Cov::aₗₘCovariance, ConvDens::AbstractConvolvedDensity,
    ProbeA::String, ProbeB::String)
    if ProbeA == ProbeB
        if ProbeA =="Lensing"
            probe_factor = 0.3
        else
            probe_factor = 1.0
        end
        for lidx in 1:length(Cov.Cℓ.CℓArray[:,1,1])
            for iidx in 1:length(Cov.Cℓ.CℓArray[1,:,1])
                Cov.Noise[lidx, iidx, iidx] = probe_factor^2 ./
                (ConvDens.SurfaceDensityArray[iidx]*3437.746771^2)
                #TODO: the multiplication here is to convert from square degrees to
                #steradians. This should be done better...
            end
        end
    end
end

function EvaluateCovariance!(Cov::aₗₘCovariance, cosmogrid::CosmologicalGrid)
    Cov.Covariance = (Cov.Cℓ.CℓArray + Cov.Noise)
    f_sky = 0.363610260832152
    #TODO: this is hardcoded...
    for (ℓidx, ℓvalue) in enumerate(cosmogrid.ℓBinCenters)
        Cov.Covariance[ℓidx,:,:] *= sqrt(2/((2* ℓvalue+1)*f_sky*cosmogrid.ℓBinWidths[ℓidx]))
    end
end

function InvertCovariance!(Cov::AbstractCovariance)
    for ℓidx in 1:length(Cov.Covariance[:,1,1])
        Cov.Covariance⁻¹[ℓidx,:,:] = inv(Cov.Covariance[ℓidx,:,:])
    end
end

"""
    InstantiateEvaluateCovariance(Covaₗₘ::aₗₘCovariance)

This function evaluates and returns the [`CℓCovariance`](@ref), using the following formula:

```math
\\boldsymbol{\\Xi}(\\ell)= \\left(\\boldsymbol{D}_{n}^{T}\\left( \\left(\\boldsymbol{\\Sigma}(\\ell)\\right)^{-1} \\otimes
\\left(\\boldsymbol{\\Sigma}(\\ell)\\right)^{-1}\\right) \\boldsymbol{D}_{n}\\right)^{-1}
```
where ``\\boldsymbol{D}_{n}`` is the [`DuplicationMatrix`](@ref).
"""
function InstantiateEvaluateCovariance(Covaₗₘ::aₗₘCovariance)
    Cov = CℓCovariance()
    ℓnumber = length(Covaₗₘ.Cℓ.CℓArray[:,1,1])
    inumber = length(Covaₗₘ.Cℓ.CℓArray[1,1,:])
    D = DuplicationMatrix(inumber)
    Dᵀ = Transpose(D)
    kronCovaₗₘ = zeros(inumber^2,inumber^2)
    DᵀkronCovaₗₘ = zeros(floor(Int,inumber*0.5*(inumber+1)),inumber^2)
    Cov.Covariance = zeros(ℓnumber, floor(Int,inumber*0.5*(inumber+1)),
    floor(Int,inumber*0.5*(inumber+1)))
    Cov.Covariance = zeros(size(Cov.Covariance))
    Cov.Covariance⁻¹ = zeros(size(Cov.Covariance))
    TempCov = zeros(floor(Int,inumber*0.5*(inumber+1)),
    floor(Int,inumber*0.5*(inumber+1)))
    for ℓ in 1:ℓnumber
        kronCovaₗₘ = kron(Covaₗₘ.Covariance⁻¹[ℓ,:,:], Covaₗₘ.Covariance⁻¹[ℓ,:,:])
        LinearAlgebra.mul!(DᵀkronCovaₗₘ, Dᵀ, kronCovaₗₘ)
        LinearAlgebra.mul!(TempCov, DᵀkronCovaₗₘ, D)
        Cov.Covariance⁻¹[ℓ,:,:] = TempCov
        Cov.Covariance[ℓ,:,:] = inv(Cov.Covariance⁻¹[ℓ,:,:])
        Matrix = (Cov.Covariance[ℓ,:,:] .+ transpose(Cov.Covariance[ℓ,:,:])) ./2
        Cov.Covariance[ℓ,:,:] = Matrix
        Cov.Covariance⁻¹[ℓ,:,:] = inv(Cov.Covariance[ℓ,:,:])
    end
    return Cov
end


## Transformations

# Rotation only
"""
    R2 = rot(θ)

Calculate a two-dimensional rotation transformation `R2`.
`θ` is the angle of rotation around the `z` axis applied to the `xy` plane.
"""
rot(θ::Number) = LinearMap(RotMatrix(θ))

"""
    R3 = rot(qx, qy, qz)

    Calculate a three-dimensional rotation transformation from the supplied parameters.
    `qx`, `qy`, and `qz` specify the rotation via the components of the corresponding
    [unit quaternion](https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation).
    Briefly, `R3` is a rotation around the axis specified by the unit vector `q/|q|`
    (where `q` is the 3-component vector), with an angle `θ` given by `sin(θ/2) = |q|`.
    If `|q| > 1`, `rot` returns `nothing` instead of erroring.
"""
function rot(qx::Number, qy::Number, qz::Number)
    # Unit-quaternion parametrization
    r2 = qx^2 + qy^2 + qz^2
    r2 > 1 && return nothing
    qr = sqrt(1 - r2)
    R = @SMatrix([1 - 2*(qy^2 + qz^2)  2*(qx*qy - qz*qr)    2*(qx*qz + qy*qr);
                  2*(qx*qy + qz*qr)    1 - 2*(qx^2 + qz^2)  2*(qy*qz - qx*qr);
                  2*(qx*qz - qy*qr)    2*(qy*qz + qx*qr)    1 - 2*(qx^2 + qy^2)])
    return LinearMap(RotMatrix(R))
end

#rotation + translation
tfmrigid(dx::Number, dy::Number, θ::Number) = Translation(dx, dy) ∘ rot(θ)

tfmrigid(dx::Number, dy::Number, dz::Number, qx::Number, qy::Number, qz::Number) =
    Translation(dx, dy, dz) ∘ rot(qx, qy, qz)

## Transformations supplied as vectors, for which we pass the array so we can infer the dimensionality

@noinline dimcheck(v, lv) = length(v) == lv || throw(DimensionMismatch("expected $lv parameters, got $(length(v))"))

function rot(theta, img::AbstractArray{T,2}) where {T}
    dimcheck(theta, 1)
    return rot(theta[1])
end

function rot(qs, img::AbstractArray{T,3}) where {T}
    dimcheck(qs, 3)
    qx, qy, qz = qs
    return rot(qx, qy, qz)
end

function tfmrigid(params, img::AbstractArray{T,2}) where {T}
    dimcheck(params, 3)
    dx, dy, θ = params
    return tfmrigid(dx, dy, θ)
end

function tfmrigid(params, img::AbstractArray{T,3}) where {T}
    dimcheck(params, 6)
    dx, dy, dz, qx, qy, qz = params
    return tfmrigid(dx, dy, dz, qx, qy, qz)
end

## Computing mismatch

#rotation + shift, fast because it uses fourier method for shift
function rigid_mm_fast(theta, mxshift, fixed, moving, thresh, SD; initial_tfm=IdentityTransformation())
    # The reason this is `initial_tfm ∘ rot` rather than `rot ∘ initial_tfm`
    # is explained in `linmap`
    tfm = arrayscale(initial_tfm ∘ rot(theta, moving), SD)
    moving, fixed = warp_and_intersect(moving, fixed, tfm)
    bshft, mm = best_shift(fixed, moving, mxshift, thresh; normalization=:intensity)
    return mm
end

#rotation + shift, slow because it warps for every rotation and shift
function rigid_mm_slow(params, fixed, moving, thresh, SD; initial_tfm=IdentityTransformation())
    tfm = arrayscale(initial_tfm ∘ tfmrigid(params, moving), SD)
    moving, fixed = warp_and_intersect(moving, fixed, tfm)
    mm = mismatch0(fixed, moving; normalization=:intensity)
    return ratio(mm, thresh, Inf)
end

function qd_rigid_coarse(fixed, moving, mxshift, mxrot, minwidth_rot;
                         SD=I,
                         initial_tfm=IdentityTransformation(),
                         thresh=0.1*sum(_abs2.(fixed[.!(isnan.(fixed))])),
                         kwargs...)
    #note: if a trial rotation results in image overlap < thresh for all possible shifts then QuadDIRECT throws an error
    f(x) = rigid_mm_fast(x, mxshift, fixed, moving, thresh, SD; initial_tfm=initial_tfm)
    upper = [mxrot...]
    lower = -upper
    root_coarse, x0coarse = _analyze(f, lower, upper;
                                     minwidth=minwidth_rot, print_interval=100, maxevals=5e4, kwargs..., atol=0, rtol=1e-3)
    box_coarse = minimum(root_coarse)
    tfmcoarse0 = initial_tfm ∘ rot(position(box_coarse, x0coarse), moving)
    best_shft, mm = best_shift(fixed, moving, mxshift, thresh; normalization=:intensity, initial_tfm=arrayscale(tfmcoarse0, SD))
    tfmcoarse = tfmcoarse0 ∘ pscale(Translation(best_shft), SD)
    return tfmcoarse, mm
end

function qd_rigid_fine(fixed, moving, mxrot, minwidth_rot;
                       SD=I,
                       initial_tfm=IdentityTransformation(),
                       thresh=0.1*sum(_abs2.(fixed[.!(isnan.(fixed))])),
                       kwargs...)
    f(x) = rigid_mm_slow(x, fixed, moving, thresh, SD; initial_tfm=initial_tfm)
    upper_shft = fill(2.0, ndims(fixed))
    upper_rot = mxrot
    upper = vcat(upper_shft, upper_rot)
    lower = -upper
    minwidth_shfts = fill(0.005, ndims(fixed))
    minwidth = vcat(minwidth_shfts, minwidth_rot)
    root, x0 = _analyze(f, lower, upper; minwidth=minwidth, print_interval=100, maxevals=5e4, kwargs...)
    box = minimum(root)
    tfmfine = initial_tfm ∘ tfmrigid(position(box, x0), moving)
    return tfmfine, value(box)
end

"""
    tform, mm = qd_rigid(fixed, moving, mxshift, mxrot;
                         presmoothed=false,
                         SD=I, minwidth_rot=default_minwidth_rot(fixed, SD),
                         thresh=thresh, initial_tfm=IdentityTransformation(), kwargs...)

Optimize a rigid transformation (rotation + shift) to minimize the mismatch between `fixed` and
`moving` using the QuadDIRECT algorithm.  The algorithm is run twice: the first step finds the optimal rotation,
using a fourier method to speed up the search for the best whole-pixel shift.
The second step performs the rotation + shift in combination to obtain sub-pixel accuracy.
Any `NaN`-valued pixels are not included in the mismatch; you can use this to mask out
any regions of `fixed` that you don't want to align against.

`mxshift` is the maximum-allowed translation, in units of array indices. It can be passed
as a vector or tuple.
`mxrot` is the maximum-allowed rotation, in radians for 2d or
[quaternion-units](https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation) for 3d.
See [`RegisterQD.rot`](@ref) for more information.
`minwidth_rot` optionally specifies the lower limit of resolution for the rotation;
the default is a rotation that moves corner elements by 0.1 pixel.

`kwargs...` can include any keyword argument that can be passed to `QuadDIRECT.analyze`.
It's recommended that you pass your own stopping criteria when possible
(i.e. `rtol`, `atol`, and/or `fvalue`).
If you provide `rtol` and/or `atol` they will apply only to the second (fine) step of the registration;
there is currently no way to adjust these criteria for the coarse step.

The rotation returned will be centered on the origin-of-coordinates, i.e. (0,0) for a 2D image.
Usually it is more natural to consider rotations around the center of the image.
If you would like `mxrot` and the returned rotation to act relative to the center of the image,
then you must move the origin to the center of the image by calling `centered(img)` from the
`ImageTransformations` package.
Call `centered` on both the fixed and moving image to generate the `fixed` and `moving` that you provide as arguments.
If you later want to apply the returned transform to an image you must remember to call `centered`
on that image as well.  Alternatively you can re-encode the transformation in terms of a
different origin by calling `recenter(tform, newctr)` where `newctr` is the displacement of the new center from the old center.

Use `SD` ("spatial displacements") if your axes are not uniformly sampled, for example
`SD = diagm(voxelspacing)` where `voxelspacing` is a vector encoding the spacing along all axes
of the image.
See [`arrayscale`](@ref) for more information about `SD`.

`thresh` enforces a certain amount of sum-of-squared-intensity overlap between
the two images; with non-zero `thresh`, it is not permissible to "align" the images by
shifting one entirely out of the way of the other. The default value for `thresh` is 10%
of the sum-of-squared-intensity of `fixed`.

If you have a good initial guess at the solution, pass it with the `initial_tfm` kwarg to jump-start the search.

Use `presmoothed=true` if you have called [`qsmooth`](@ref) on `fixed` before calling `qd_rigid`.
Do not smooth `moving`.

Both the output `tfm` and any `initial_tfm` are represented in *physical* coordinates;
as long as `initial_tfm` is a rigid transformation, `tfm` will be a pure rotation+translation.
If `SD` is not the identity, use `arrayscale` before applying the result to `moving`.
"""
function qd_rigid(fixed, moving, mxshift::VecLike, mxrot::Union{Number,VecLike};
                  presmoothed=false,
                  SD=I,
                  minwidth_rot=default_minwidth_rot(CartesianIndices(fixed), SD),
                  thresh=0.1*sum(_abs2.(fixed[.!(isnan.(fixed))])),
                  initial_tfm=IdentityTransformation(),
                  print_interval=100,
                  kwargs...)
    #TODO make sure isrotation(initial_tfm.linear) returns true, else throw a warning and an option to terminate? do I still allow the main block to run?
    if initial_tfm == IdentityTransformation() || isrotation(initial_tfm.linear)
    else
        @show "WARNING: initial_tfm is not a rigid transformation"
    end
    fixed, moving = float(fixed), float(moving)
    if presmoothed
        moving = qinterp(eltype(fixed), moving)
    end
    mxrot = [mxrot...]
    print_interval < typemax(Int) && print("Running coarse step\n")
    tfm_coarse, mm_coarse = qd_rigid_coarse(fixed, moving, mxshift, mxrot, minwidth_rot; SD=SD, initial_tfm=initial_tfm, thresh=thresh, print_interval=print_interval, kwargs...)
    print_interval < typemax(Int) && print("Obtained mismatch $mm_coarse from coarse step, running fine step\n")
    final_tfm, mm_fine = qd_rigid_fine(fixed, moving, mxrot./2, minwidth_rot; SD=SD, initial_tfm=tfm_coarse, thresh=thresh, print_interval=print_interval, kwargs...)
    return final_tfm, mm_fine
end


@testset "Test interpolation of outgoing_translation_matrix" begin

    particle_medium = Acoustic(3; ρ=10.0, c=80.0);
    medium = Acoustic(3; ρ=1.0, c=1.0);

    R = 3.0
    r = 1.0

    separation_ratio = 1.02
    ka = 0.1
    k = ka / r
    vol_fraction = 0.3

    basis_order = 2

    ω = k * real(medium.c)

    s1 = Specie(
        particle_medium, Sphere(r);
        number_density = vol_fraction / volume(Sphere(r)),
        exclusion_distance = separation_ratio
    );
    species = [s1]
    material = Material(Sphere(R),species);

    a12 = 2.0 * minimum(outer_radius(s) * s.exclusion_distance for s in material.species)

    tol = 1e-2

    Uinter = outgoing_translation_matrix(ω, medium, material;  basis_order = basis_order, tol = tol)

    rs = LinRange(a12,2R-a12,13);
    φs = LinRange(-pi,pi,13);
    θs = LinRange(0.0,pi,13);

    r2x = radial_to_cartesian_coordinates

    error = [
        begin
            U1 = outgoing_translation_matrix(medium, basis_order, ω, r2x([r,θ,φ]))
            norm(Uinter(r2x([r,θ,φ])) - U1) / norm(U1)
        end
    for r in rs, θ in θs, φ in φs];

    @test maximum(error) < 5 * tol

    @test mean(error) < tol

end


function custom_generator{T<:Union{AbstractFloat,Integer}}(a::Number, b::Number, type_a::Bool, type_b::Bool,
   var_type::Type{T}, p::Float64 = 1.5)
    if a > b
        error("Generotor received lower bound that is greater then upper bound:
        $a > $b")
    end
	if a == -Inf && b == Inf
               if var_type<:AbstractFloat
                   return (size)->(rand(var_type)-.5)*var_type(p)*size
               elseif var_type<:Integer
                   return (size)->rand(Vector{var_type}(-size:size))
               else
                   error("Generator for type $var_type is not specified.")
               end
	elseif b == Inf && type_a == false  # a<<x
		if var_type<:Integer
			return (size) -> var_type(a)+var_type(100)*rand(Vector{var_type}(1:size))
		elseif var_type<:AbstractFloat
			return (size) -> var_type(a)+rand(var_type)*var_type(p)^(size-1)
		else
			error("Generator for type $var_type is not specified.")
		end
	elseif b == Inf && type_a == true  # a<x
		if var_type<:Integer
			return (size) -> var_type(a)+rand(Vector{var_type}(1:size))
		elseif var_type<:AbstractFloat
			return (size) -> var_type(a)+rand(var_type)*var_type(p)*size
		else
			error("Generator for type $var_type is not specified.")
		end
	elseif a == -Inf && type_b == false  # x<<b
		if var_type<:Integer
			return (size) -> var_type(b)-var_type(100)*rand(Vector{var_type}(1:size))
		elseif var_type<:AbstractFloat
			return (size) -> var_type(b)-rand(var_type)*var_type(p)^(size-1)
		else
			error("Generator for type $var_type is not specified.")
		end
	elseif a == -Inf && type_b == true  # x<b
		if var_type<:Integer
			return (size) -> var_type(b)-rand(Vector{var_type}(1:size))
		elseif var_type<:AbstractFloat
			return (size) -> var_type(b)-rand(var_type)*var_type(p)*size
		else
			error("Generator for type $var_type is not specified.")
		end
	elseif  abs(a)!=Inf && abs(b)!=Inf && var_type<:AbstractFloat
		if (type_a,type_b) == (true,true)  # x<a<b
			return (size)->var_type(a)+(var_type(b)-var_type(a))*rand(var_type)  # choose uniform [a,b]
		elseif (type_a,type_b) == (true,false) # a<x<<b
			return (size)->var_type(a)+(var_type(b)-var_type(a))*rand(var_type)/var_type(p)^(size-1)
		elseif (type_a,type_b) == (false,true) # a<<x<b
			return (size)->var_type(b)-(var_type(b)-var_type(a))*rand(var_type)/var_type(p)^(size-1)
		end
	else
		error("Could not specify generator for $a,$b,$var_type")
	end
end


@polly function kernel_trisolv(L, x, b)
    n = size(L,1)
    for i = 1:n
        x[i] = b[i]
        for j = 1:(i-1)
            x[i] -= L[i,j] * x[j]
        end
        x[i] = x[i] / L[i,i]
    end
end

let
    n = 2000

    L = zeros(Float32, n, n)
    x = zeros(Float32, n)
    b = zeros(Float32, n)

    for i = 1:n
        x[i] = -999
        b[i] = i
        for j = 1:i
            L[i,j] = (i+n-j+1)*2/n
        end
    end

    SUITE["trisolv"] = @benchmarkable kernel_trisolv(L, x, b) setup = (L = copy($L); x = copy($x); b = copy($b))
end


# developed with Julia 1.4.2
#
# EMS simulation with a MPC controller
# the online computing of optimal controls is written in LP form
#
# this example can be run with any JuMP compatible LP solver
# e.g replace Clp with CPLEX

using EMSx
using JuMP, Clp

include("arguments.jl")

args = parse_commandline()

mutable struct Mpc <: EMSx.AbstractController
    model::Model
    horizon::Int64
    Mpc() = new()
end

const controller = Mpc()

function EMSx.initialize_site_controller(controller::Mpc, site::EMSx.Site, prices::EMSx.Prices)

    controller = Mpc()

    model = Model(Clp.Optimizer)
    set_optimizer_attribute(model, "LogLevel", 0)

    horizon = 96
    battery = site.battery

    @variable(model, 0 <= u_c[1:horizon])
    @variable(model, 0 <= u_d[1:horizon])
    @variable(model, 0 <= x[1:horizon+1])
    @variable(model, 0 <= z[1:horizon])
    @variable(model, w[1:horizon])
    @variable(model, x0)

    @expression(model, u,  u_c - u_d)

    @constraint(model, u_c .<= battery.power*0.25)
    @constraint(model, u_d .<= battery.power*0.25)
    @constraint(model, x .<= battery.capacity)
    @constraint(model, u.+w .<= z)
    @constraint(model, x[1] == x0)
    @constraint(model, dynamics, diff(x) .== u_c*battery.charge_efficiency .-
        u_d/battery.discharge_efficiency)

    controller.model = model
    controller.horizon = horizon

    return controller

end

function EMSx.compute_control(controller::Mpc, information::EMSx.Information)

    fix(controller.model[:x0], information.soc*information.battery.capacity)
    fix.(controller.model[:w], information.forecast_load - information.forecast_pv)

    # set prices, padding out of test period prices with zero values
    prices = information.prices
    price_window = information.t:min(information.t+controller.horizon-1, size(prices.buy, 1))
    if length(price_window) != controller.horizon
        padding = controller.horizon - length(price_window)
        prices = EMSx.Prices(prices.name, vcat(prices.buy[price_window], zeros(padding)),
            vcat(prices.sell[price_window], zeros(padding)))
    else
        prices = EMSx.Prices(prices.name, prices.buy[price_window], prices.sell[price_window])
    end

    @objective(controller.model, Min,
        sum(prices.buy.*controller.model[:z]-
            prices.sell.*(controller.model[:z]-controller.model[:u]-controller.model[:w])))

    optimize!(controller.model)

    return value(controller.model[:u][1]) / (information.battery.power*0.25)

end

EMSx.simulate_sites(controller,
    joinpath(args["save"], "mpc"),
    args["price"],
    args["metadata"],
    args["test"])


function update_eta!(state::State, data::Data, prior::Prior)
  update_etaC!(state, data, prior)
  update_etaT!(state, data, prior)
end

function update_etaC!(state::State, data::Data, prior::Prior)
  anew = prior.eta.alpha .+ 0

  for n in eachindex(state.lambdaC)
    anew[state.lambdaC[n]] += 1
  end

  if !state.beta
    for n in eachindex(state.lambdaT)
      anew[state.lambdaT[n]] += 1
    end
  end

  state.etaC = rand(Dirichlet(anew))
end

function update_etaT!(state::State, data::Data, prior::Prior)
  anew = prior.eta.alpha .+ 0

  if state.beta
    for n in eachindex(state.lambdaT)
      anew[state.lambdaT[n]] += 1
    end
  end

  state.etaT = rand(Dirichlet(anew))
end


"""
Generates a random value from a truncated normal distribution that represents
a new item. It's used to simulate incoming application to the registry.
"""
function getRegistryCandidate()
    rand(TruncatedNormal(50, 20, 0, 100))
end

"""
Calculates the mean quality of the items in the given registry
"""
function registryMean(registry)
    length(registry) == 0 ? 0 : mean(registry)
end


import WignerSymbols: clebschgordan

"""
    qfunc(a, α)
    qfunc(a, x, y)
    qfunc(a, xvec, yvec)

Husimi Q representation ``⟨α|ρ|α⟩/π`` for the given state or operator `a`. The
function can either be evaluated on one point α or on a grid specified by
the vectors `xvec` and `yvec`. Note that conversion from `x` and `y` to `α` is
done via the relation ``α = \\frac{1}{\\sqrt{2}}(x + i y)``.
"""
function qfunc(rho::AbstractOperator{B,B}, alpha::Number) where B<:FockBasis
    b = basis(rho)
    _qfunc_operator(rho, convert(ComplexF64, alpha), Ket(b), Ket(b))
end

function qfunc(rho::AbstractOperator{B,B}, xvec::Vector{Float64}, yvec::Vector{Float64}) where B<:FockBasis
    b = basis(rho)
    Nx = length(xvec)
    Ny = length(yvec)
    tmp1 = Ket(b)
    tmp2 = Ket(b)
    result = Matrix{ComplexF64}(undef, Nx, Ny)
    @inbounds for j=1:Ny, i=1:Nx
        result[i, j] = _qfunc_operator(rho, complex(xvec[i], yvec[j])/sqrt(2), tmp1, tmp2)
    end
    result
end

function qfunc(psi::Ket{B}, alpha::Number) where B<:FockBasis
    b = basis(psi)
    _alpha = convert(ComplexF64, alpha)
    _conj_alpha = conj(_alpha)
    N = length(psi.basis)
    s = psi.data[N]/sqrt(N-1)
    @inbounds for i=1:N-2
        s = (psi.data[N-i] + s*_conj_alpha)/sqrt(N-i-1)
    end
    s = psi.data[1] + s*_conj_alpha
    return abs2(s)*exp(-abs2(_alpha))/pi
end

function qfunc(psi::Ket{B}, xvec::Vector{Float64}, yvec::Vector{Float64}) where B<:FockBasis
    b = basis(psi)
    Nx = length(xvec)
    Ny = length(yvec)
    points = Nx*Ny
    N = length(b)::Int
    _conj_alpha = [complex(x, -y)/sqrt(2) for x=xvec, y=yvec]
    q = fill(psi.data[N]/sqrt(N-1), size(_conj_alpha))
    @inbounds for n=1:N-2
        f0_ = 1/sqrt(N-n-1)
        x = psi.data[N-n]
        for i=1:points
            q[i] = (x + q[i]*_conj_alpha[i])*f0_
        end
    end
    result = similar(q, Float64)
    x = psi.data[1]
    @inbounds for i=1:points
        result[i] = abs2(x + q[i]*_conj_alpha[i])*exp(-abs2(_conj_alpha[i]))/pi
    end
    result
end

function qfunc(state::Union{Ket{B}, AbstractOperator{B,B}}, x::Number, y::Number) where B<:FockBasis
    qfunc(state, ComplexF64(x, y)/sqrt(2))
end

function _qfunc_operator(rho::AbstractOperator{B,B}, alpha::ComplexF64, tmp1::Ket, tmp2::Ket) where B<:FockBasis
    coherentstate!(tmp1, basis(rho), alpha)
    QuantumOpticsBase.gemv!(complex(1.), rho, tmp1, complex(0.), tmp2)
    a = dot(tmp1.data, tmp2.data)
    return a/pi
end

"""
    wigner(a, α)
    wigner(a, x, y)
    wigner(a, xvec, yvec)

Wigner function for the given state or operator `a`. The
function can either be evaluated on one point α or on a grid specified by
the vectors `xvec` and `yvec`. Note that conversion from `x` and `y` to `α` is
done via the relation ``α = \\frac{1}{\\sqrt{2}}(x + i y)``.
"""
function wigner(rho::DenseOperator{B,B}, x::Number, y::Number) where B<:FockBasis
    b = basis(rho)
    N = b.N::Int
    _2α = complex(convert(Float64, x), convert(Float64, y))*sqrt(2)
    abs2_2α = abs2(_2α)
    w = complex(0.)
    coefficient = complex(0.)
    @inbounds for L=N:-1:1
        coefficient = 2*_clenshaw(L, abs2_2α, rho.data)
        w = coefficient + w*_2α/sqrt(L+1)
    end
    coefficient = _clenshaw(0, abs2_2α, rho.data)
    w = coefficient + w*_2α
    exp(-abs2_2α/2)/pi*real(w)
end

function wigner(rho::DenseOperator{B,B}, xvec::Vector{Float64}, yvec::Vector{Float64}) where B<:FockBasis
    b = basis(rho)
    N = b.N::Int
    _2α = [complex(x, y)*sqrt(2) for x=xvec, y=yvec]
    abs2_2α = abs2.(_2α)
    w = zero(_2α)
    b0 = similar(_2α)
    b1 = similar(_2α)
    b2 = similar(_2α)
    @inbounds for L=N:-1:1
        _clenshaw_grid(L, rho.data, abs2_2α, _2α, w, b0, b1, b2, 2)
    end
    _clenshaw_grid(0, rho.data, abs2_2α, _2α, w, b0, b1, b2, 1)
    @inbounds for i=eachindex(w)
        abs2_2α[i] = exp(-abs2_2α[i]/2)/pi.*real(w[i])
    end
    abs2_2α
end

wigner(psi::Ket, x, y) = wigner(dm(psi), x, y)
wigner(state, alpha::Number) = wigner(state, real(alpha)*sqrt(2), imag(alpha)*sqrt(2))

function _clenshaw_grid(L::Int, ρ::Matrix{ComplexF64},
                abs2_2α::Matrix{Float64}, _2α::Matrix{ComplexF64}, w::Matrix{ComplexF64},
                b0::Matrix{ComplexF64}, b1::Matrix{ComplexF64}, b2::Matrix{ComplexF64}, scale::Int)
    n = size(ρ, 1)-L-1
    points = length(w)
    if n==0
        f = scale*ρ[1, L+1]
        @inbounds for i=1:points
            w[i] = f + w[i]*_2α[i]/sqrt(L+1)
        end
    elseif n==1
        f1 = 1/sqrt(L+1)
        @inbounds for i=1:points
            w[i] = scale*(ρ[1, L+1] - ρ[2, L+2]*(L+1-abs2_2α[i])*f1) + w[i]*_2α[i]*f1
        end
    else
        f0 = sqrt(float((n+L-1)*(n-1)))
        f1 = sqrt(float((n+L)*n))
        f0_ = 1/f0
        f1_ = 1/f1
        fill!(b1, ρ[n+1, L+n+1])
        @inbounds for i=1:points
            b0[i] = ρ[n, L+n] - (2*n-1+L-abs2_2α[i])*f1_*b1[i]
        end
        @inbounds for k=n-2:-1:1
            b1, b2, b0 = b0, b1, b2
            x = ρ[k+1, L+k+1]
            a1 = -(2*k+1+L)
            a2 = -f0*f1_
            @inbounds for i=1:points
                b0[i] = x + (a1+abs2_2α[i])*f0_*b1[i] + a2*b2[i]
            end
            f1 , f1_ = f0, f0_
            f0 = sqrt((k+L)*k)
            f0_ = 1/f0
        end
        @inbounds for i=1:points
            w[i] = scale*(ρ[1, L+1] - (L+1-abs2_2α[i])*f0_*b0[i] - f0*f1_*b1[i]) + w[i]*_2α[i]*f0_
        end
    end
end

function _clenshaw(L::Int, abs2_2α::Float64, ρ::Matrix{ComplexF64})
    n = size(ρ, 1)-L-1
    if n==0
        return ρ[1, L+1]
    elseif n==1
        ϕ1 = -(L+1-abs2_2α)/sqrt(L+1)
        return ρ[1, L+1] + ρ[2, L+2]*ϕ1
    else
        f0 = sqrt(float((n+L-1)*(n-1)))
        f1 = sqrt(float((n+L)*n))
        f0_ = 1/f0
        f1_ = 1/f1
        b2 = complex(0.)
        b1 = ρ[n+1, L+n+1]
        b0 = ρ[n, L+n] - (2*n-1+L-abs2_2α)*f1_*b1
        @inbounds for k=n-2:-1:1
            b1, b2 = b0, b1
            b0 = ρ[k+1, L+k+1] - (2*k+1+L-abs2_2α)*f0_*b1 - f0*f1_*b2
            f1, f1_ = f0, f0_
            f0 = sqrt((k+L)*k)
            f0_ = 1/f0
        end
        return ρ[1, L+1] - (L+1-abs2_2α)*f0_*b0 - f0*f1_*b1
    end
end

"""
    coherentspinstate(b::SpinBasis, θ::Real, ϕ::Real)

A coherent spin state |θ,ϕ⟩ is analogous to the coherent state of the linear harmonic
oscillator. Coherent spin states represent a collection of identical two-level
systems and can be described by two angles θ and ϕ (although this
parametrization is not unique), similarly to a qubit on the
Bloch sphere.
"""
function coherentspinstate(b::SpinBasis, theta::Real, phi::Real,
    result = Ket(b, Vector{ComplexF64}(undef, length(b))))
    data = result.data
    N = BigInt(length(b)-1)
    sinth = sin(0.5theta)
    costh = cos(0.5theta)
    expphi = exp(0.5im*phi)
    expphi_con = conj(expphi)
    @inbounds for n=0:N
        data[n+1] = sqrt(binomial(N, n)) * (sinth*expphi)^n * (costh*expphi_con)^(N-n)
    end
    return result
end

"""
    qfuncsu2(ket,Ntheta;Nphi=2Ntheta)
    qfuncsu2(rho,Ntheta;Nphi=2Ntheta)

Husimi Q SU(2) representation ``⟨θ,ϕ|ρ|θ,ϕ⟩/π`` for the given state.

The function calculates the SU(2) Husimi representation of a state on the
generalised bloch sphere (0 < θ < π and 0 < ϕ < 2 π) with a given
resolution `(Ntheta, Nphi)`.

    qfuncsu2(rho,θ,ϕ)
    qfuncsu2(ket,θ,ϕ)

This version calculates the Husimi Q SU(2) function at a position given by θ and ϕ.
"""
function qfuncsu2(psi::Ket{B}, Ntheta::Int; Nphi::Int=2Ntheta) where B<:SpinBasis
    b = psi.basis
    psi_bra_data = psi.data'
    lb = float(b.spinnumber)
    result = Array{Float64}(undef, Ntheta,Nphi)
    @inbounds  for i = 0:Ntheta-1, j = 0:Nphi-1
        result[i+1,j+1] = (2*lb+1)/(4pi)*abs2(psi_bra_data*coherentspinstate(b,pi-i*pi/(Ntheta-1),j*2pi/(Nphi-1)-pi).data)
    end
    return result
end

function qfuncsu2(rho::DenseOperator{B,B}, Ntheta::Int; Nphi::Int=2Ntheta) where B<:SpinBasis
    b = basis(rho)
    lb = float(b.spinnumber)
    result = Array{Float64}(undef, Ntheta,Nphi)
    @inbounds  for i = 0:Ntheta-1, j = 0:Nphi-1
        c = coherentspinstate(b,pi-i*1pi/(Ntheta-1),j*2pi/(Nphi-1)-pi)
        result[i+1,j+1] = abs((2*lb+1)/(4pi)*c.data'*rho.data*c.data)
    end
    return result
end

function qfuncsu2(psi::Ket{B}, theta::Real, phi::Real) where B<:SpinBasis
    b = basis(psi)
    psi_bra_data = psi.data'
    lb = float(b.spinnumber)
    result = (2*lb+1)/(4pi)*abs2(psi_bra_data*coherentspinstate(b,theta,phi).data)
    return result
end

function qfuncsu2(rho::DenseOperator{B,B}, theta::Real, phi::Real) where B<:SpinBasis
    b = basis(rho)
    lb = float(b.spinnumber)
    c = coherentspinstate(b,theta,phi)
    result = abs((2*lb+1)/(4pi)*c.data'*rho.data*c.data)
    return result
end

"""
    wignersu2(ket,Ntheta;Nphi=2Ntheta)
    wignersu2(rho,Ntheta;Nphi=2Ntheta)

Wigner SU(2) representation for the given state with a resolution `(Ntheta, Nphi)`.

The function calculates the SU(2) Wigner representation of a state on the
generalised bloch sphere (0 < θ < π and 0 < ϕ < 2 π) with a given resolution by
decomposing the state into the basis of spherical harmonics.

    wignersu2(rho,θ,ϕ)
    wignersu2(ket,θ,ϕ)

This version calculates the Wigner SU(2) function at a position given by θ and ϕ
"""
function wignersu2(rho::DenseOperator{B,B}, theta::Real, phi::Real) where B<:SpinBasis

    N = length(basis(rho))-1

    ### Tensor generation ###
    BandT = Array{Vector{Float64}}(undef, N,N+1)
    BandT[1,1] = collect(range(-N/2, stop=N/2, length=N+1))
    BandT[1,2] = -collect(sqrt.(range(1, stop=N, length=N)).*sqrt.(range((N)/2, stop=1/2, length=N)))
    BandT[2,1] = clebschgordan(1,0,1,0,2,0)*BandT[1,1].*BandT[1,1] -
        clebschgordan(1,-1,1,1,2,0)*[zeros(N+1-length(BandT[1,2])); BandT[1,2].*BandT[1,2]] -
        clebschgordan(1,1,1,-1,2,0)*[BandT[1,2].*BandT[1,2]; zeros(N+1-length(BandT[1,2]))]
    BandT[2,2] = clebschgordan(1,0,1,1,2,1)BandT[1,1][1:N].*BandT[1,2]+
        clebschgordan(1,1,1,0,2,1)*BandT[1,2][1:N].*BandT[1,1][2:end]
    BandT[2,3] = BandT[1,2][1:N+1-(2)].*BandT[1,2][2:end]

    @inbounds for S=2:N-1
        BandT[S+1,1] = clebschgordan(1,0,S,0,S+1,0)*BandT[1,1].*BandT[S,1] -
            [zeros(N+1-length(BandT[1,2])); clebschgordan(1,-1,S,1,S+1,0)*BandT[1,2].*BandT[S,2]] -
            clebschgordan(1,1,S,-1,S+1,0)*[BandT[1,2].*BandT[S,2]; zeros(N+1-length(BandT[1,2]))]
        BandT[S+1,S+1] = clebschgordan(1,0,S,S,S+1,S)BandT[1,1][1:N+1-S].*BandT[S,S+1]+
            clebschgordan(1,1,S,S-1,S+1,S)*BandT[1,2][1:N+1-S].*BandT[S,S][2:end]
        BandT[S+1,S+2] = BandT[1,2][1:N+1-(S+1)].*BandT[S,S+1][2:end]
        @inbounds for M=1:S-1
            BandT[S+1,M+1] = clebschgordan(1, 0, S, M, S+1,M)*BandT[1,1][1:N+1-M].*BandT[S,M+1] +
                clebschgordan(1,1,S,M-1,S+1,M)*BandT[1,2][1:N+1-M].*BandT[S,M][2:end] -
                clebschgordan(1,-1,S,M+1,S+1,M)*[zeros(1); BandT[1,2][1:N-M].*BandT[S,M+2][1:N-M]]
        end

    end

    NormT = zeros(N)
    @inbounds for S = 1:N
        NormT[S] = sum(BandT[S,1].^2)
    end

    @inbounds for S = 1:N, M = 0:S
        BandT[S, M + 1] =  BandT[S, M + 1]/sqrt(NormT[S])
    end

    ### State decomposition ###
    c = rho.data
    EVT = Array{ComplexF64}(undef, N,N+1)
    @inbounds for S = 1:N, M = 0:S
        EVT[S,M+1] = conj(sum(BandT[S,M+1].*diag(c,M)))
    end

    wignermap = _wignersu2int(N,theta,phi, EVT)
    return wignermap*sqrt((N+1)/(4pi))
end

function wignersu2(rho::DenseOperator{B,B}, Ntheta::Int; Nphi::Int=2Ntheta) where B<:SpinBasis

    N = length(basis(rho))-1

    ### Tensor generation ###
    BandT = Array{Vector{Float64}}(undef, N,N+1)
    BandT[1,1] = collect(range(-N/2, stop=N/2, length=N+1))
    BandT[1,2] = -collect(sqrt.(range(1, stop=N, length=N)).*sqrt.(range((N)/2, stop=1/2, length=N)))
    BandT[2,1] = clebschgordan(1,0,1,0,2,0)*BandT[1,1].*BandT[1,1] -
        clebschgordan(1,-1,1,1,2,0)*[zeros(N+1-length(BandT[1,2])); BandT[1,2].*BandT[1,2]] -
        clebschgordan(1,1,1,-1,2,0)*[BandT[1,2].*BandT[1,2]; zeros(N+1-length(BandT[1,2]))]
    BandT[2,2] = clebschgordan(1,0,1,1,2,1)BandT[1,1][1:N].*BandT[1,2]+
        clebschgordan(1,1,1,0,2,1)*BandT[1,2][1:N].*BandT[1,1][2:end]
    BandT[2,3] = BandT[1,2][1:N+1-(2)].*BandT[1,2][2:end]

    @inbounds for S=2:N-1
        BandT[S+1,1] = clebschgordan(1,0,S,0,S+1,0)*BandT[1,1].*BandT[S,1] -
            [zeros(N+1-length(BandT[1,2])); clebschgordan(1,-1,S,1,S+1,0)*BandT[1,2].*BandT[S,2]] -
            clebschgordan(1,1,S,-1,S+1,0)*[BandT[1,2].*BandT[S,2]; zeros(N+1-length(BandT[1,2]))]
        BandT[S+1,S+1] = clebschgordan(1,0,S,S,S+1,S)BandT[1,1][1:N+1-S].*BandT[S,S+1]+
            clebschgordan(1,1,S,S-1,S+1,S)*BandT[1,2][1:N+1-S].*BandT[S,S][2:end]
        BandT[S+1,S+2] = BandT[1,2][1:N+1-(S+1)].*BandT[S,S+1][2:end]
        @inbounds for M=1:S-1
            BandT[S+1,M+1] = clebschgordan(1, 0, S, M, S+1,M)*BandT[1,1][1:N+1-M].*BandT[S,M+1] +
                clebschgordan(1,1,S,M-1,S+1,M)*BandT[1,2][1:N+1-M].*BandT[S,M][2:end] -
                clebschgordan(1,-1,S,M+1,S+1,M)*[0.0im; BandT[1,2][1:N-M].*BandT[S,M+2][1:N-M]]
        end
    end

    NormT = zeros(N)
    @inbounds for S = 1:N
        NormT[S] = sum(BandT[S,1].^2)
    end
    @inbounds for S = 1:N, M = 0:S
        BandT[S, M + 1] = BandT[S, M + 1]/sqrt(NormT[S])
    end

    ### State decomposition ###
    c = rho.data
    EVT = Array{ComplexF64}(undef, N,N+1)
    @inbounds for S = 1:N, M = 0:S
        EVT[S,M+1] = conj(sum(BandT[S,M+1].*diag(c,M)))
    end

    wignermap = Array{Float64}(undef, Ntheta,Nphi)
    @inbounds for i = 1:Ntheta, j = 1:Nphi
        wignermap[i,j] = _wignersu2int(N,i*1pi/(Ntheta-1)-1pi/(Ntheta-1),j*2pi/(Nphi-1)-2pi/(Nphi-1)-pi, EVT)
    end
    return wignermap*sqrt((N+1)/(4pi))
end

function _wignersu2int(N::Integer, theta::Real, phi::Real, EVT::Array{ComplexF64, 2})
    UberBand = sqrt(1/(1+N))*ylm(0,0,theta,phi)
    @inbounds for S = 1:N
        @inbounds for M = 1:S
            UberBand += 2*real(EVT[S,M+1]*conj(ylm(S,M,theta,phi)))
        end
        UberBand += EVT[S,1]*ylm(S,0,theta,phi)
    end
    UberBand
end
wignersu2(psi::Ket, args...) = wignersu2(dm(psi), args...)

"""
    ylm(l::Integer,m::Integer,theta::Real,phi::Real)

Spherical harmonics Y(l,m)(θ,ϕ) where l ∈ N,  m = -l,-l+1,...,l-1,l, θ ∈ [0,π],
and ϕ ∈ [0,2π).

This function calculates the value of Y(l,m) spherical harmonic at position θ and ϕ.
"""
function ylm(l::Integer,m::Integer,theta::Real,phi::Real)
    phi_ = mod(phi,2pi)
    theta_ = mod(theta,2pi)
    phase = exp(1.0im*m*phi_)
    if theta_ ≈ 0
        if m == 0
            return @. phase*sqrt((2*l+1)/pi)/2
        else
            return 0
        end
    elseif theta_ ≈ pi
        if m == 0
            return @. phase*(-1)^l*sqrt((2*l+1)/pi)/2
        else
            return 0
        end
    else
        if l == 0
            return 1.0/sqrt(4pi)
        else
            m_ = abs(m)
            norm = _calc_ylm_norm(l, m_)
            sign = m > 0 ? (-1)^m_ : 1
            arg = cos(theta_)
            p_ll = 1.0
            @inbounds for fact = 1.0:l
                p_ll *= @. 1.0/((2*fact))*sqrt(1-arg^2)
            end

            if m_ == l
                return @. p_ll/norm*phase*sign
            elseif l-m_ == 1
                p_llp1 = @. 2*l*arg/sqrt(1-arg^2)*p_ll
                return @. p_llp1/norm*phase*sign
            else
                p_llp1 = @. 2*l*arg/sqrt(1-arg^2)*p_ll
                @inbounds for mr = -l:-m_-2
                    p_llp2 = @. -2*(mr+1)*arg/sqrt(1-arg^2)*p_llp1-(l-mr)*(l+mr+1)*p_ll
                    p_ll = p_llp1
                    p_llp1 = p_llp2
                end
                return @. p_llp1/norm*phase*sign
            end
        end
    end
end

function _calc_ylm_norm(l::Int, m_::Int)
    # TODO: Clean up Int types
    if 0 < l+m_ <= 33
        norm = @. Float64(sqrt(4pi)/sqrt(2*l+1)*sqrt(factorial(Int128(l-m_)))/sqrt(factorial(Int128(l+m_))))
    elseif 0 < l-m_ <= 33 && l+m_ > 33
        norm = @. Float64(sqrt(4pi)/sqrt(2*l+1)*sqrt(factorial(Int128(l-m_)))/sqrt(factorial(BigInt(l+m_))))
    else
        norm = @. Float64(sqrt(4pi)/sqrt(2*l+1)*sqrt(factorial(BigInt(l-m_)))/sqrt(factorial(BigInt(l+m_))))
    end
    norm
end


num_johnson_subsets(simplex_length::Integer) = 2^simplex_length - 1

"""
Compute all subsets of the points in the simplex in a reliable order. The
order is arbitrary, but was chosen to match the implemenation in
S. Cameron, “Enhancing GJK: computing minimum and penetration distances between
convex polyhedra,”. Specifically, subset i contains point j iff the binary
representation of i has a one at the jth bit.
"""
function johnson_subsets(simplex_length::Integer)
    num_subsets = num_johnson_subsets(simplex_length)
    subsets = falses(simplex_length, num_subsets)
    for i in 1:num_subsets
        for j in 1:simplex_length
            subsets[j, i] = (i >> (j - 1)) & 1
        end
    end
    subsets
end

"""
    weights = projection_weights(simplex)

This function implements Johnson's distance subalgorithm, as described in
E. G. Gilbert, D. W. Johnson, and S. S. Keerthi, “A fast procedure for
computing the distance between complex objects in three-dimensional space,”
1988. Given a simplex (a length N+1 vector of points of dimension N), it
returns weights such that dot(weights, simplex) yields the point in the convex
hull of the simplex which is closest to the origin.

This is the critical loop of the GJK algorithm, so it has been heavily optimized
to precompute, inline, and unroll as much of the algorithm as possible. For a
more readable (and much slower) implementation, see
projection_weights_reference()
"""
@generated function projection_weights(simplex::SVector{M, SVector{N, T}}; atol=sqrt(eps(T))) where {M,N,T}
    simplex_length = M
    num_subsets = num_johnson_subsets(M)
    subsets = johnson_subsets(M)
    complements = .!subsets

    expr = quote
        deltas = SVector(tuple($([zero(SVector{simplex_length, T}) for i = 1 : num_subsets]...)))
    end

    # Set the weight of every singleton subset to 1.
    for i in 1:simplex_length
        elements = [:(zero($T)) for j in 1:simplex_length]
        elements[i] = :(one($T))
        arg = :(SVector{$simplex_length, $T}(tuple($(elements...))))

        push!(expr.args, quote
            deltas = setindex(deltas, $(arg), $(2^(i - 1)))
        end)
    end

    for s in 1:(num_subsets - 1)
        k = first((1:M)[subsets[:,s]])
        push!(expr.args, quote
            viable = true
        end)
        for j in (1:M)[complements[:,s]]
            s2 = s + (1 << (j - 1))
            push!(expr.args, quote
                d = deltas[$s2][$j]
            end)
            for i in (1:M)[subsets[:,s]]
                push!(expr.args, quote
                    d += deltas[$s][$i] *
                    (dot(simplex[$i], simplex[$k]) - dot(simplex[$i], simplex[$j]))
                end)
            end
            push!(expr.args, quote
                deltas = setindex(deltas, setindex(deltas[$s2], d, $j), $s2)
            end)
            push!(expr.args, quote
                if d > atol # See #17
                    viable = false
                end
            end)
        end
        push!(expr.args, quote
            viable = viable && all($(Expr(:tuple,
            [:(deltas[$s][$i] >= 0) for i in (1:M)[subsets[:,s]]]...)))
            if viable
                return deltas[$s] ./ sum(deltas[$s])
            end
        end)
    end
    push!(expr.args, quote
        return deltas[$num_subsets] ./ sum(deltas[$num_subsets])
    end)
    return expr
end

function projection_weights_reference(simplex::SVector{M, SVector{N, T}}) where {M,N,T}
    num_subsets = 2^M - 1
    complements = falses(M, num_subsets)
    subsets = falses(M, num_subsets)
    for i in 1:num_subsets
        for j in 1:M
            subsets[j, i] = (i >> (j - 1)) & 1
            complements[j, i] = !subsets[j, i]
        end
    end

    deltas = zeros(MMatrix{M, num_subsets, T})
    viable = ones(MVector{num_subsets, Bool})

    for i in 1:M
        deltas[i, 2^(i - 1)] = 1
    end

    for s in 1:(num_subsets - 1)
        k = first((1:M)[subsets[:,s]])
        for j in (1:M)[complements[:,s]]
            s2 = s + (1 << (j - 1))
            delta_j_s2 = 0
            for i in (1:M)[subsets[:,s]]
                delta_j_s2 += deltas[i, s] * (dot(simplex[i], simplex[k]) - dot(simplex[i], simplex[j]))
            end
            if delta_j_s2 > 0
                viable[s] = false
            elseif delta_j_s2 < 0
                viable[s2] = false
            end
            deltas[j, s2] = delta_j_s2
        end
        if viable[s]
            return deltas[:,s] ./ sum(deltas[:,s])
        end
    end
    return deltas[:,end] ./ sum(deltas[:,end])
end


#!/bin/julia
#
# An implementation of the lattice boltzmann method
#
# @author arbennett
##

# Constant parameters
const T   = 9000          # Number of timesteps to run
const Re  = 20.0           # Reynolds number
const nx  = 150              # Size of the domain in x direction
const ny  = 50             # Size of the domain in y direction
const q   = 9               # number of flow directions
const v   = 0.05            # velocity
const nu  = (v*ny)/(10*Re)  # viscosity
const tau = (3.0*nu+0.5)    # 1/relaxation
const plot_frequency = T/5

## Set up the flow grids, weights, and indices
#
#  Explain here
#
##
dirs =     [  0   1   0  -1   0   1   -1   -1    1 ;
              0   0   1   0  -1   1    1   -1   -1 ]
weights =  [4/9 1/9 1/9 1/9 1/9 1/36 1/36 1/36 1/36]
indices =  [  4   5   6;  1   2   3;   7    8    9 ]
no_slip =  [  1   4   5   2   3   8    9    6    7 ]

## Create the obstacle array
#
#  The obst array is a boolean array that uses true to denote where
#  obstacles to flow are.  When the "fluid" encounters one of these
#  cells it will reverse the flow direction using the `no_slip`
#  condition detailed above
#
##
obst_coords = [nx/2, ny/2, 3]  # [x, y, radius]
obst_coords2 = [nx/3,ny/3,4]
y = collect(1:ny)
x = collect(1:nx)'
obst = (x.-obst_coords[1]).^2 .+ (y.-obst_coords[2]).^2 .<= obst_coords[3].^2
obst = (obst + ((x.-obst_coords2[1]).^2 .+ (y.-obst_coords2[2]).^2 .<= obst_coords2[3].^2) .> 0)
#obst[1,:] = true
#obst[end,:] = true
#obst[:,:] = false

##  Set up initial flow profile
#
#  Allocate the velocity array, u, which has 2 points for each grid
#  point, corresponding to the x and y velocity
#
#  Allocate the distribution array, f, which has 9 points for each grid
#  point, corresponding to the directions in dirs
##
u = zeros(ny,nx,2)
vel = zeros(ny,nx,2)
fEq = zeros(ny,nx,9)
rho = ones(ny,nx)

for j=1:ny, i=1:nx
    vel[j,i,2] = v*(1+0.1*sin((j/ny)/(2*pi)))
end

## Calculate equilibrium for some distribution
#
##
function equilib(u, rho)
    dist = zeros(ny,nx,9)
    for i = 1:9
        for x = 1:nx
            for y = 1:ny
                tmp = u[y,x,1] * dirs[1,i] + u[y,x,2] * dirs[2,i]
                dist[y,x,i] = rho[y,x] * weights[i] * (1 + 3*tmp + 4.5*tmp*tmp - 1.5*(u[y,x,1]*u[y,x,1] + u[y,x,2]*u[y,x,2]))
            end
        end
    end
    return dist
end
fIn = equilib(u, rho)
fLeft = fIn[:,1,:]

## Propogate time
#
##
for t=1:T
    if mod(t, plot_frequency) == 0
        println(t)
        writedlm("out/solution.dat."*lpad(string(t),length(string(int(T))),"0"),sqrt(u[:,:,1].*u[:,:,1] + u[:,:,2].*u[:,:,2]))
    end

    # Right wall conditions
    fIn[:,end,:] = fIn[:,end-1,:]
    rho = fIn[:,:,1] + fIn[:,:,2] + fIn[:,:,3] + fIn[:,:,4] + fIn[:,:,5] + fIn[:,:,6] + fIn[:,:,7] + fIn[:,:,8] + fIn[:,:,9]

    u[:,:,1] = (fIn[:,:,2] - fIn[:,:,4] + fIn[:,:,6] - fIn[:,:,7] - fIn[:,:,8] + fIn[:,:,9]) ./ rho[:,:]
    u[:,:,2] = (fIn[:,:,3] - fIn[:,:,5] + fIn[:,:,6] + fIn[:,:,7] - fIn[:,:,8] - fIn[:,:,9]) ./ rho[:,:]

    # Left wall conditions
    u[:,1,:] = vel[:,1,:]
    #rho[:,1] = 1./(1-u[:,1,2]) .* (fIn[:,1,3] + fIn[:,1,6] + fIn[:,1,7] + 2*(fIn[:,1,5] + fIn[:,1,8] + fIn[:,1,9]))
    #println(rho[:,1])

    fEq = fIn - (1/tau) * (fIn - equilib(u, rho))

    for i=1:q
        for x=1:nx
            for y=1:ny
                if obst[y,x]
                    fEq[y,x,i] = fIn[y, x, no_slip[i]]
                    u[y,x,:] = 0
                end
            end
        end
    end

    for i=1:q
        fIn[:,:,i] = circshift(fEq[:,:,i],dirs[:,i])
    end
end
println("done!")


using KernelRidgeRegression
using StatsBase
using MLKernels

N = 5000
x = rand(1, N) * 4π - 2π
yy = sinc.(x) # vec(sinc.(4 .* x) .+ 0.2 .* sin.(30 .* x))
y = squeeze(yy + 0.1randn(1, N), 1)
xnew = collect(-2.5π:0.01:2.5π)'

@time mykrr = fit(KRR, x, y, 1e-3/5000, GaussianKernel(1.0))
showcompact(mykrr)
show(mykrr)

@time mynystkrr = fit(KernelRidgeRegression.NystromKRR, x, y, 1e-3/5000, 280, MLKernels.GaussianKernel(100.0))
showcompact(mynystkrr)
show(mynystkrr)

@time myfastkrr = fit(KernelRidgeRegression.FastKRR, x, y, 4/5000, 11, MLKernels.GaussianKernel(100.0))
showcompact(myfastkrr)
show(myfastkrr)

@time mytnkrr = fit(KernelRidgeRegression.TruncatedNewtonKRR, x, y, 4/5000, MLKernels.GaussianKernel(100.0), 0.5, 200)
showcompact(mytnkrr)
show(mytnkrr)

@time myrandkrr = fit(KernelRidgeRegression.RandomFourierFeatures, x, y, 1e-3/5000, 200 , 1.0)
showcompact(myrandkrr)
show(myrandkrr)


include("polynomial_feature.jl")
include("gaussian_feature.jl")
include("sigmoidal_feature.jl")


mutable struct Node
	x :: Int
	y :: Int
	value :: Int
	bestdist :: Int
end

struct Edge
	sideA :: Node
	sideB :: Node
end

other(edge :: Edge, node :: Node) = node == edge.sideA ? edge.sideB : node == edge.sideB ? edge.sideA : error("node not on either edge")
getedges(n :: Node, edges :: Vector{Edge}) = [edge for edge in edges if edge.sideA == n || edge.sideB == n]
makenode(value :: Int, x :: Int, y :: Int) = Node(x, y, value, typemax(Int))

function create_edges(grid_to_node :: Dict{Tuple{Int, Int}, Node})
	nodes = values(grid_to_node)
	edges_from_node = Dict{Node, Vector{Edge}}(n => Edge[] for n in nodes)
	for node in nodes
		if haskey(grid_to_node, (node.x + 1, node.y))
			target_node = grid_to_node[(node.x + 1, node.y)]
			new_edge = Edge(node, target_node)
			push!(edges_from_node[node], new_edge)
			push!(edges_from_node[target_node], new_edge)
		end

		if haskey(grid_to_node, (node.x, node.y + 1))
			target_node = grid_to_node[(node.x, node.y + 1)]
			new_edge = Edge(node, target_node)
			push!(edges_from_node[node], new_edge)
			push!(edges_from_node[target_node], new_edge)
		end
	end

	return edges_from_node
end

function dijkstra!(grid_to_node :: Dict{Tuple{Int, Int}, Node}, edges_from_node :: Dict{Node, Vector{Edge}})
	grid_to_node[(1, 1)].bestdist = 0
	eligibles = Set{Node}([grid_to_node[(1, 1)]])
	counter = 0

	while counter < length(grid_to_node)
		current = argmin(n -> n.bestdist, eligibles)
		node_edges = edges_from_node[current]
		for other_edge in node_edges
			target = other(other_edge, current)
			if current.bestdist + target.value < target.bestdist
				target.bestdist = current.bestdist + target.value
				push!(eligibles, target)
			end
		end

		pop!(eligibles, current)
		counter += 1
	end
	# modifies in-place, return not needed
end

function part1()
	lines = readlines("inputD15.txt")

	grid_to_node = Dict{Tuple{Int, Int}, Node}()
	for (i, line) in enumerate(lines)
		for (j, value) in enumerate(line)
			grid_to_node[(i, j)] = makenode(parse(Int, value), i, j)
		end
	end

	edges_from_node = create_edges(grid_to_node)
	dijkstra!(grid_to_node, edges_from_node)

	return grid_to_node[(100, 100)].bestdist
end

function add(n :: Int, amount :: Int)
	for _ in 1:amount
		n == 9 ? n = 1 : n += 1
	end

	return n
end

function part2()
	lines = readlines("inputD15.txt")

	grid_to_node = Dict{Tuple{Int, Int}, Node}()
	for (i, line) in enumerate(lines)
		for (j, value) in enumerate(line)
			value = parse(Int, value)
			for m in 0:4
				for n in 0:4
					grid_to_node[(i+n*100, j+m*100)] = makenode(add(value, n+m), i+n*100, j+m*100)
				end
			end
		end
	end

	edges_from_node = create_edges(grid_to_node)
	dijkstra!(grid_to_node, edges_from_node)

	return grid_to_node[(500, 500)].bestdist
end

println(part1())
println(part2())


using BoltzmannMachines
import BoltzmannMachinesPlots # only needed for plotting capabilities

# ==============================================================================
# Binary data: BernoulliRBM and BasicDBM
# ------------------------------------------------------------------------------

# Generate data and split it in training and test data set
nsamples = 100;
nvariables = 36;
using Random
Random.seed!(1);
x = barsandstripes(nsamples, nvariables);
x, xtest = splitdata(x, 0.3);
datadict = DataDict("Training data" => x, "Test data" => xtest);

# Fit a BernoulliRBM and collect monitoring data during training.
# For small models, the model's loglikelihood can be computed exactly.
monitor = Monitor();
rbm = fitrbm(x; nhidden = 12,
      epochs = 80, learningrate = 0.006,
      monitoring = (rbm, epoch) -> begin
            monitorexactloglikelihood!(monitor, rbm, epoch, datadict)
            monitorreconstructionerror!(monitor, rbm, epoch, datadict)
      end);

# Plot the collected data (requires Gadfly)
BoltzmannMachinesPlots.plotevaluation(monitor, monitorexactloglikelihood)
BoltzmannMachinesPlots.plotevaluation(monitor, monitorreconstructionerror)

# Simplified monitoring:
monitor, rbm = monitored_fitrbm(x; nhidden = 12,
      epochs = 80, learningrate = 0.006,
      monitoringdata = datadict,
      monitoring = [monitorexactloglikelihood!, monitorreconstructionerror!]);

Random.seed!(12);

# DBM-Fitting: Pretraining and Fine-Tuning combined in one function
dbm = fitdbm(x, nhiddens = [6;2], epochs = 20, learningrate = 0.05);

# .. with extensive monitoring
monitor = Monitor(); monitor1 = Monitor(); monitor2 = Monitor();
dbm = fitdbm(x, epochs = 20, learningrate = 0.05,
      monitoringdatapretraining = datadict,
      pretraining = [
            TrainLayer(nhidden = 6, monitoring = (rbm, epoch, datadict) ->
                  monitorreconstructionerror!(monitor1, rbm, epoch, datadict));
            TrainLayer(nhidden = 2, monitoring = (rbm, epoch, datadict) ->
                  monitorreconstructionerror!(monitor2, rbm, epoch, datadict))
            ],
      monitoring = (dbm, epoch) ->
            monitorexactloglikelihood!(monitor, dbm, epoch, datadict));

# Monitoring plots
BoltzmannMachinesPlots.plotevaluation(monitor1)
BoltzmannMachinesPlots.plotevaluation(monitor2)
BoltzmannMachinesPlots.plotevaluation(monitor)

# Evaluate final result with exact computation of likelihood
exactloglikelihood(dbm, xtest)

# Simplified monitoring
monitors, dbm = monitored_fitdbm(x; nhiddens = [6,2],
      epochs = 20, learningrate = 0.05,
      monitoringpretraining = monitorreconstructionerror!,
      monitoring = monitorexactloglikelihood!,
      monitoringdata = datadict);

# If the model has more parameters, in this case more hidden nodes,
# the exact calculation is not feasible any more:
# We need to calculate the likelihood using AIS.

Random.seed!(12);
monitor = Monitor();
rbm = fitrbm(x, nhidden = 36, epochs = 100,
      learningrates = [0.008*ones(80); 0.001*ones(20)],
      monitoring = (rbm, epoch) ->
            monitorloglikelihood!(monitor, rbm, epoch, datadict));

BoltzmannMachinesPlots.plotevaluation(monitor, monitorloglikelihood; sdrange = 3.0)

# Simplified monitoring:
monitor, rbm = monitored_fitrbm(x, nhidden = 36, epochs = 100,
      learningrates = [0.008*ones(80); 0.001*ones(20)],
      monitoring = monitorloglikelihood!,
      monitoringdata = datadict);

# In the DBM, the estimation of the log likelihood is much slower than in the
# RBM because it needs a AIS run for each sample.
# The lower bound of the log probability is faster to estimate:
# (But on the other hand, in some cases it might be too conservative to be
# useful as it is only a lower bound.)

Random.seed!(2);
monitor = Monitor();
dbm = stackrbms(x; nhiddens = [36;10;5], predbm = true, learningrate = 0.05);
traindbm!(dbm, x; epochs = 50, learningrate = 0.008,
      monitoring = (rbm, epoch) ->
            monitorlogproblowerbound!(monitor, rbm, epoch, datadict));

BoltzmannMachinesPlots.plotevaluation(monitor, monitorlogproblowerbound; sdrange = 3.0)

# Simplified monitoring:
monitor, dbm = monitored_traindbm!(dbm, x; epochs = 50, learningrate = 0.008,
      monitoring = monitorlogproblowerbound!,
      monitoringdata = datadict);

# Evaluate final result with AIS-estimated likelihood
loglikelihood(dbm, xtest)

# ==============================================================================
# Categorical data: Softmax0BernoulliRBM
# ------------------------------------------------------------------------------

# Some data set with values in the categories {0, 1, 2}
# Mixing variables with a different number of categories is also possible
# (see also example below for a MultimodalDBM with binary and categorical
# variables).
xcat = map(v -> max(v, 0),
      barsandstripes(100, 4) .+ barsandstripes(100, 4) .- barsandstripes(100, 4));

# Encode values for Softmax0BernoulliRBM:
# Each of the variables, which have three categories, are translated into
# two binary variables.
# This encoding is to similiar to the one-hot encoding with the deviation
# that a zero is encoded as all-zeros.
# (That way, the encoding of binary variables is not changed.)
x01 = oneornone_encode(xcat, 3) # 3 categories
rbm = fitrbm(x01, rbmtype = Softmax0BernoulliRBM, categories = 3,
      nhidden = 4, epochs = 20);

# The monitoring is omitted here for brevity:
# The same monitoring and training options as those for BernoulliRBMs
# are also available for Softmax0BernoulliRBMs.

# A DBM with a Softmax0BernoulliRBM in the first layer and binary hidden layers:
dbm = fitdbm(x01, pretraining = [
         TrainLayer(rbmtype = Softmax0BernoulliRBM, categories = 3, nhidden = 4);
         TrainLayer(nhidden = 4);
         TrainLayer(nhidden = 2)])

# Getting samples with values in the original sample space:
oneornone_decode(samples(dbm, 5), 3)

# ==============================================================================
# Mixing binary and categorical data: MultimodalDBM
# ------------------------------------------------------------------------------

# x1: binary data
x1 = barsandstripes(100, 4);
# x2: values with categories {0, 1, 2}
x2 = map(v -> max(v, 0), x1 .+ barsandstripes(100, 4) .- barsandstripes(100, 4));
x2 = oneornone_encode(x2, 3)
x = hcat(x1, x2);
x, xtest = splitdata(x, 0.3);
datadict = DataDict("Training data" => x, "Test data" => xtest);
monitor1 = Monitor(); monitor2 = Monitor(); monitor3 = Monitor(); monitor = Monitor();
dbm = fitdbm(x; epochspretraining = 50, epochs = 15,
      learningratepretraining = 0.05, learningrate = 0.1,
      monitoringdatapretraining = datadict,
      pretraining = [
            TrainPartitionedLayer([
               TrainLayer(nhidden = 5, nvisible = 4,
                     monitoring = (rbm, epoch, datadict) ->
                           monitorreconstructionerror!(monitor1, rbm, epoch, datadict));
               TrainLayer(nhidden = 7,
                     rbmtype = Softmax0BernoulliRBM, categories = 3,
                     monitoring = (rbm, epoch, datadict) ->
                           monitorreconstructionerror!(monitor2, rbm, epoch, datadict))
            ]),
            TrainLayer(nhidden = 6,
                  monitoring = (rbm, epoch, datadict) ->
                        monitorreconstructionerror!(monitor3, rbm, epoch, datadict))
      ],
      monitoring =
            (dbm, epoch) -> monitorlogproblowerbound!(monitor, dbm, epoch, datadict));

# The same with simplified monitoring:
Random.seed!(1)
monitors, dbm = monitored_fitdbm(x, epochspretraining = 50, epochs = 15,
      pretraining = [
            TrainPartitionedLayer([
                  TrainLayer(nhidden = 5, nvisible = 4);
                  TrainLayer(nhidden = 7,
                        rbmtype = Softmax0BernoulliRBM, categories = 2)]),
            TrainLayer(nhidden = 6)
      ],
      learningratepretraining = 0.05,
      learningrate = 0.1,
      monitoringdata = datadict,
      monitoringpretraining = monitorreconstructionerror!,
      monitoring = monitorlogproblowerbound!);

# Reconstructionerror of BernoulliRBM in first layer:
BoltzmannMachinesPlots.plotevaluation(monitors[1][1])
# Reconstructionerror of Softmax0BernoulliRBM:
BoltzmannMachinesPlots.plotevaluation(monitors[1][2])
# Reconstructionerror of second layer:
BoltzmannMachinesPlots.plotevaluation(monitors[2])
# Lower bound of likelihood during fine-tuning:
BoltzmannMachinesPlots.plotevaluation(monitors[3])

# It is also possible to combine binary and categorical variables in one RBM,
# as a Softmax0BernoulliRBM with variables having two categories is equivalent
# to a BernoulliRBM.
# Here we do not have a partition then. We can use one Softmax0BernoulliRBM for
# the visible layer and have to specify the number of categories for each of the
# variables separately:

categories = [fill(2, 4); fill(3, 4)];
dbm = fitdbm(x;
      pretraining = [
         TrainLayer(nhidden = 7, categories = categories,
               rbmtype = Softmax0BernoulliRBM);
         TrainLayer(nhidden = 6)
      ])

# ==============================================================================
# Real valued data: Intensities, GaussianBernoulliRBM or GaussianBernoulliRBM2
# ------------------------------------------------------------------------------

# Use "iris" dataset as example data for continuous data
using DelimitedFiles
x = convert(Matrix{Float64},
      readdlm(joinpath(dirname(pathof(BoltzmannMachines)), "..",
            "test/data/iris.csv"), ',', header = true)[1][:,1:4]);
Random.seed!(12);
x, xtest = splitdata(x, 0.3);
datadict = DataDict("Training data" => x, "Test data" => xtest);

# Using the intensities transformation and a standard BernoulliRBM
xintensities, xtrafo = intensities_encode(x)
rbm = fitrbm(xintensities)
# Transform samples back into the original sample space:
intensities_decode(samples(rbm, 5), xtrafo)

# Example using a GaussianBernoulliRBM
Random.seed!(1);
monitor = Monitor();
rbm = fitrbm(x, rbmtype = GaussianBernoulliRBM,
      nhidden = 3, epochs = 80, learningrate = 0.0005,
      monitoring = (rbm, epoch) ->
            monitorexactloglikelihood!(monitor, rbm, epoch, datadict));

BoltzmannMachinesPlots.plotevaluation(monitor, monitorexactloglikelihood)

# The alternative variant of an RBM with Gaussian visible nodes and binary hidden
# nodes, as described by Cho et. al. in "Improved learning of Gaussian-Bernoulli
# restricted Boltzmann machines", is also available as "GaussianBernoulliRBM2".
rbm = fitrbm(x, rbmtype = GaussianBernoulliRBM2,
      nhidden = 3, epochs = 70, learningrate = 0.001);

# The same monitoring options as for BernoulliRBMs are available
# for GaussianBernoulliRBMs and GaussianBernoulliRBM2s.

# ==============================================================================
# Data with binomially distributed values in {0,1,2}: Binomial2BernoulliRBM
# ------------------------------------------------------------------------------

# Generate data and split it in training and test data set
nsamples = 100;
nvariables = 25;
Random.seed!(12);
x = barsandstripes(nsamples, nvariables) + barsandstripes(nsamples, nvariables);
x, xtest = splitdata(x, 0.3);
datadict = DataDict("Training data" => x, "Test data" => xtest);

# Fit Binomial2BernoulliRBM.
monitor = Monitor();
rbm = fitrbm(x, rbmtype = Binomial2BernoulliRBM,
      nhidden = 10, epochs = 100, learningrates = [0.003*ones(60); 0.001*ones(40)],
      monitoring = (rbm, epoch) ->
            monitorexactloglikelihood!(monitor, rbm, epoch, datadict));

BoltzmannMachinesPlots.plotevaluation(monitor, monitorexactloglikelihood)

# ==============================================================================
# Examples for conditional sampling
# ------------------------------------------------------------------------------
nsamples = 100;
nvariables = 36;

x = barsandstripes(100, 16);
dbm = fitdbm(x);

# Conditional gibbs sampling of particles
particles = initparticles(dbm, 500)
variablestofixto1 = 1:3
particles[1][:, variablestofixto1] .= 1.0
gibbssamplecond!(particles, dbm, variablestofixto1)
particles[1]

# Short alternative, returns only the visible nodes
samples(dbm, 500, conditions = [i => 1.0 for i in 1:3])

# ==============================================================================
# Examples for cross-validation
# ------------------------------------------------------------------------------

nsamples = 500;
nvariables = 16;
x = barsandstripes(nsamples, nvariables);

# Determine the optimal number of training epochs for a RBM
monitor = crossvalidation(x,
      (x, datadict) ->
            begin
               monitor = BoltzmannMachines.Monitor()
               BoltzmannMachines.fitrbm(x, epochs = 50,
                     learningrate = 0.005,
                     monitoring = (rbm, epoch) ->
                           BoltzmannMachines.monitorexactloglikelihood!(
                                 monitor, rbm, epoch, datadict))
               monitor
            end);
BoltzmannMachinesPlots.crossvalidationcurve(monitor)

# Determine the optimal number of pretraining epochs for a DBM
# given the other parameters
using Distributed
@everywhere using BoltzmannMachines
@everywhere function my_pretraining_monitoring(
      x::Matrix{Float64}, datadict::DataDict, epoch::Int)

   monitor = Monitor()
   monitorlogproblowerbound!(monitor,
         stackrbms(x,
               nhiddens = [16; 12],
               learningrate = 0.001,
               predbm = true,
               epochs = epoch),
         epoch,
         datadict,
         parallelized = false)
end

monitor = crossvalidation(x, my_pretraining_monitoring, 10:10:200);
BoltzmannMachinesPlots.crossvalidationcurve(monitor, monitorlogproblowerbound)

# ==============================================================================
# Examples for custom optimization algorithm
# ------------------------------------------------------------------------------

x = barsandstripes(100, 16);

struct MyRegularizedOptimizer{R<:AbstractRBM} <: AbstractOptimizer{R}
   llopt::LoglikelihoodOptimizer{R}
end

function MyRegularizedOptimizer()
   MyRegularizedOptimizer(loglikelihoodoptimizer(learningrate = 0.01))
end

import BoltzmannMachines: initialized, computegradient!, updateparameters!

function initialized(myoptimizer::MyRegularizedOptimizer,
      rbm::R) where {R<: AbstractRBM}
   MyRegularizedOptimizer(initialized(myoptimizer.llopt, rbm))
end

function computegradient!(myoptimizer::MyRegularizedOptimizer,
      v, vmodel, h, hmodel, rbm)
   computegradient!(myoptimizer.llopt, v, vmodel, h, hmodel, rbm)
   # add L2 regularization term
   myoptimizer.llopt.gradient.weights .-= 0.1 .* rbm.weights
   nothing
end

function updateparameters!(rbm::R, myoptimizer::MyRegularizedOptimizer{R}
      ) where {R<: AbstractRBM}
   updateparameters!(rbm, myoptimizer.llopt)
end

# The optimizer can be used for fitting RBMs ...
rbm = fitrbm(x; optimizer = MyRegularizedOptimizer())

# and also for DBMs
dbm = fitdbm(x; optimizer = MyRegularizedOptimizer())


@testset "Convert" begin

    x = variable()
    for (α, β) in ((1/2, 1/4), (1/4, 1/2), (1,2), (2,1))
        P, Q = Laguerre{α}, Laguerre{β}
        for _ in 1:10
            p = P(rand(1:5,5)); q=P(rand(1:5,5))
            @test convert(P, convert(Q, p))(x) ≈ p(x)
            @test convert(Q, convert(P,q))(x) ≈ q(x)
        end
    end

end


"""
    g = gaussianize(z)

Gaussianize field `z` using a sorted random sequence.
"""
function gaussianize(z)

(m,n)=size(z);
nn=m*n;
zr=z[:];
zs=sort(zr);
i=sortperm(zr);
rrk=zeros(Int16,nn,1);
rrk[i]=collect(1:nn);
gg=randn(nn);
w=sort(gg);
gr=w[rrk];
gs=Statistics.std(gr);
if gs==0
        gs=1;
end
gr=gr/gs;
g=reshape(gr,m,n);
end


"""
    CFLabelPropagationFinder() <: AbstractCommunityFinder

Represents a label propagation search over the hypergraph `h` that finds
a partition into communities (subsets).
"""
struct CFLabelPropagationFinder <: AbstractCommunityFinder
    max_iter::Int
    seed::Int
end

"""
    findcommunities(h::Hypergraph, method::CFLabelPropagationFinder)

Implements the label propagation algorithm over a hypergraph `h`.

NOTE
The algorithm works on a single connected component.
An AssertionError is thrown otherwise.

This algorithm generalizes the one proposed for graphs by Raghavan et al.
(Raghavan, U. N., Albert, R., and Kumara, S. Near linear time algorithm to detect
community structures in large-scale networks. Physical review. E, Statistical,
nonlinear, and soft matter physics 76 (2007).)

The proposed algorithm modifies the propagation rule, splitting it into two phases:
hyperedge labeling and vertex labeling.

For more information see `Section 4` in the paper
Alessia Antelmi, Gennaro Cordasco, Bogumił Kamiński, Paweł Prałat,
Vittorio Scarano, Carmine Spagnuolo, Przemyslaw Szufel
*Analyzing, Exploring, and Visualizing Complex Networks via Hypergraphs Using SimpleHypergraphs.jl.*
Journal Internet Mathematics (2020). https://doi.org/10.24166/im.01.2020
"""
function findcommunities(h::Hypergraph, method::CFLabelPropagationFinder)
    @assert length(get_connected_components(h)) == 1

    rng = MersenneTwister(method.seed)
    vlabels = Dict{Int, Int}()
    helabels = Dict{Int, Int}()

    for v in 1:nhv(h)
        push!(vlabels, v=>v)
    end

    vertices = collect(1:nhv(h))
    edges = collect(1:nhe(h))

    stop = false
    iter = 0

    while !stop && iter < method.max_iter
        stop = true
        shuffle!(rng, edges)

        for e in edges
            length(getvertices(h, e)) == 0 && continue
            l = SimpleHypergraphs.compute_edge_label(h, e, vlabels, helabels, rng)
            push!(helabels, e => l)
        end

        shuffle!(rng, vertices)
        for v in vertices
            l = SimpleHypergraphs.compute_vertex_label(h, v, vlabels, helabels, rng)
            if l != vlabels[v]
                stop = false
                push!(vlabels, v => l)
            end
        end

        iter+=1
    end

    np_vertices = unique(values(vlabels))
    np_edges = unique(values(helabels))

    comms_vertices = Dict{Int, Set{Int}}()
    comms_hyperedges = Dict{Int, Set{Int}}()

    for pv in vlabels
        push!(
            get!(comms_vertices, pv[2], Set{Int}()),
            pv[1]
        )
    end

    for pe in helabels
        push!(
            get!(comms_hyperedges, pe[2], Set{Int}()),
            pe[1]
        )
    end

    labels = Array{Int}(undef, nhv(h))
    for i in 1:nhv(h)
        labels[i] = vlabels[i]
    end

    hlabels = Array{Int}(undef, nhe(h))
    for i in 1:nhe(h)
        if length(getvertices(h, i)) == 0
            hlabels[i] = -1
        else
            hlabels[i] = helabels[i]
        end
    end

    (np=collect(values(comms_vertices)), hep=collect(values(comms_hyperedges)), vlabels=labels, helabels=hlabels, iter=iter)
end

"""
    compute_vertex_label(h::Hypergraph, v::Int, vlabels::Dict{Int,Int}, helabels::Dict{Int,Int}, rng::MersenneTwister)

Vertices labeling phase. Computes the label of each vertex according to the most
frequent label among the hyperedges it belongs to.
"""
function compute_vertex_label(h::Hypergraph, v::Int, vlabels::Dict{Int,Int}, helabels::Dict{Int,Int}, rng::MersenneTwister)
    hesᵥ = gethyperedges(h, v)
    vL = Dict{Int,Int}()

    max = 0
    maxL = Set{Int}()

    for e in shuffle!(rng, collect(keys(hesᵥ)))
        l = helabels[e]

        if !haskey(vL, l)
            push!(vL, l=>0)
        end

        push!(
            vL,
            l => vL[l] + (length(getvertices(h, e)))
        )

        if vL[l] == max
            push!(maxL, l)
        elseif vL[l] > max
            max = vL[l]
            maxL = Set{Int}()
            push!(maxL, l)
        end
    end

    if in(vlabels[v], maxL)
        return vlabels[v]
    end

    return first(maxL)
end

"""
    compute_edge_label(h::Hypergraph, e::Int, vlabels::Dict{Int,Int}, helabels::Dict{Int,Int}, rng::MersenneTwister)

Hyperedges labeling phase. Computes the labels of the hyperedges according  to
the  most frequent label among the vertices contained in that hyperedge.
"""
function compute_edge_label(h::Hypergraph, e::Int, vlabels::Dict{Int,Int}, helabels::Dict{Int,Int}, rng::MersenneTwister)
    vₑ = getvertices(h,e)
    eL = Dict{Int,Int}()

    max = 0
    maxL = Set{Int}()

    for v in shuffle!(rng, collect(keys(vₑ)))
        l = vlabels[v]

        if !haskey(eL, l)
            push!(eL, l=>0)
        end

        push!(
            eL,
            l => eL[l]+1
        )

        if eL[l] == max
            push!(maxL, l)
        elseif eL[l] > max
            max = eL[l]
            maxL = Set{Int}()
            push!(maxL, l)
        end
    end

    if haskey(helabels, e) && in(helabels[e], maxL)
        return helabels[e]
    end

    return first(maxL)
end


# Symmetric triangular distribution
immutable TriangularDist <: ContinuousUnivariateDistribution
    location::Float64
    scale::Float64
    function TriangularDist(l::Real, s::Real)
        s > zero(s) || error("scale must be positive")
        new(float64(l), float64(s))
    end
end

TriangularDist(location::Real) = TriangularDist(location, 1.0)
TriangularDist() = TriangularDist(0.0, 1.0)

## Support
isupperbounded(::Union(TriangularDist, Type{TriangularDist})) = true
islowerbounded(::Union(TriangularDist, Type{TriangularDist})) = true
isbounded(::Union(TriangularDist, Type{TriangularDist})) = true

minimum(d::TriangularDist) = d.location - d.scale
maximum(d::TriangularDist) = d.location + d.scale
insupport(d::TriangularDist, x::Real) = minimum(d) <= x <= maximum(d)

## Properties
mean(d::TriangularDist) = d.location
median(d::TriangularDist) = d.location
mode(d::TriangularDist) = d.location

var(d::TriangularDist) = d.scale^2 / 6.0
skewness(d::TriangularDist) = 0.0
kurtosis(d::TriangularDist) = -0.6

entropy(d::TriangularDist) = 0.5 + log(d.scale)

## Functions
function pdf(d::TriangularDist, x::Real)
    if insupport(d, x)
        (1.0 - abs(x-d.location)/d.scale) / d.scale
    else
        0.0
    end
end

function logpdf(d::TriangularDist, x::Real)
    if insupport(d, x)
        log1p(-abs(x-d.location)/d.scale) -log(d.scale)
    else
        -Inf
    end
end

function cdf(d::TriangularDist, x::Real)
    u = (x - d.location) / d.scale
    if u <= -1.0
        0.0
    elseif u <= 0.0
        0.5*(1.0 + u)^2
    elseif u <= 1.0
        1.0 - 0.5*(1.0 - u)^2
    else
        1.0
    end
end
function ccdf(d::TriangularDist, x::Real)
    u = (x - d.location) / d.scale
    if u <= -1.0
        1.0
    elseif u <= 0.0
        1.0 - 0.5*(1.0 + u)^2
    elseif u <= 1.0
        0.5*(1.0 - u)^2
    else
        0.0
    end
end

function logcdf(d::TriangularDist, x::Real)
    u = (x - d.location) / d.scale
    if u <= -1.0
        -Inf
    elseif u <= 0.0
        loghalf + 2.0*log1p(u)
    elseif u <= 1.0
        log1p(-0.5*(1.0 - u)^2)
    else
        0.0
    end
end
function logccdf(d::TriangularDist, x::Real)
    u = (x - d.location) / d.scale
    if u <= -1.0
        0.0
    elseif u <= 0.0
        log1p(- 0.5*(1.0 + u)^2)
    elseif u <= 1.0
        loghalf + 2.0*log1p(-u)
    else
        -Inf
    end
end

function quantile(d::TriangularDist, p::Real)
    @checkquantile p begin
        if p < 0.5
            d.location - d.scale*(1.0 - sqrt(2.0*p))
        else
            d.location + d.scale*(1.0 - sqrt(2.0*(1.0-p)))
        end
    end
end
function cquantile(d::TriangularDist, p::Real)
    @checkquantile p begin
        if p > 0.5
            d.location - d.scale*(1.0 - sqrt(2.0*(1.0-p)))
        else
            d.location + d.scale*(1.0 - sqrt(2.0*p))
        end
    end
end
function invlogcdf(d::TriangularDist, lp::Real)
    @checkinvlogcdf lp begin
        if lp < loghalf
            d.location + d.scale*expm1(0.5*(lp-loghalf))
        else
            d.location + d.scale*(1.0 - sqrt(-2.0*expm1(lp)))
        end
    end
end
function invlogccdf(d::TriangularDist, lp::Real)
    @checkinvlogcdf lp begin
        if lp > loghalf
            d.location - d.scale*(1.0 - sqrt(-2.0*expm1(lp)))
        else
            d.location - d.scale*expm1(0.5*(lp-loghalf))
        end
    end
end

function mgf(d::TriangularDist, t::Real)
    a = d.scale*t
    a == zero(a) && return one(a)
    4.0*exp(d.location*t)*(sinh(0.5*a)/a)^2
end
function cf(d::TriangularDist, t::Real)
    a = d.scale*t
    a == zero(a) && return complex(one(a))
    4.0*exp(im*d.location*t)*(sin(0.5*a)/a)^2
end

## Sampling
function rand(d::TriangularDist)
    ξ1, ξ2 = rand(), rand()
    return d.location + (ξ1 - ξ2) * d.scale
end


module FESpacesWithLinearConstraintsTests

using Gridap.Algebra
using Gridap.Arrays
using Gridap.Fields
using Gridap.Geometry
using Gridap.FESpaces
using Test
using LinearAlgebra
using Gridap.CellData

domain = (0,1,0,1)
partition = (2,2)
model = CartesianDiscreteModel(domain,partition)
trian = Triangulation(model)

labels = get_face_labeling(model)
add_tag_from_tags!(labels,"dirichlet",[1,2,5])
add_tag_from_tags!(labels,"neumann",[6,7,8])

V = FESpace(
  model=model,valuetype=Float64,reffe=:Lagrangian,order=1,
  conformity=:H1,dirichlet_tags="dirichlet")
test_single_field_fe_space(V)

fdof_to_val = collect(Float64,1:num_free_dofs(V))
ddof_to_val = -collect(Float64,1:num_dirichlet_dofs(V))
vh = FEFunction(V,fdof_to_val,ddof_to_val)

sDOF_to_dof = [1,5,-2]
sDOF_to_dofs = Table([[-1,4],[4,6],[-1,-3]])
sDOF_to_coeffs = Table([[0.5,0.5],[0.5,0.5],[0.5,0.5]])

Vc = FESpaceWithLinearConstraints(
  sDOF_to_dof,
  sDOF_to_dofs,
  sDOF_to_coeffs,
  V)

test_single_field_fe_space(Vc)
@test has_constraints(Vc)

scellids = SkeletonPair([1,2,1,3],[2,4,3,4])
@test isa(get_cell_constraints(Vc,scellids)[1],BlockArrayCoo)

@test Vc.n_fdofs == 6
@test Vc.n_fmdofs == 4

fmdof_to_val = collect(Float64,1:num_free_dofs(Vc))
dmdof_to_val = -collect(Float64,1:num_dirichlet_dofs(Vc))
vch = FEFunction(Vc,fmdof_to_val,dmdof_to_val)
r = [[-1.0, -1.5, 1.0, 1.0], [-1.5, -2.0, 1.0, 2.0], [1.0, 1.0, 3.0, 3.5], [1.0, 2.0, 3.5, 4.0]]
@test get_cell_values(vch) ≈ r

v(x) = sin(4*x[1]+0.4)*cos(5*x[2]+0.7)
vch = interpolate(v,Vc)

#using Gridap.Visualization
#writevtk(trian,"trian",nsubcells=10,cellfields=["vh"=>vh,"vch"=>vch])

u(x) = x[1] + 2*x[2]
f(x) = -Δ(u)(x)
Uc = TrialFESpace(Vc,u)
@test has_constraints(Uc)
uch = interpolate(u,Uc)

btrian = BoundaryTriangulation(model,"neumann")

quad = CellQuadrature(trian,2)
bquad = CellQuadrature(btrian,2)

bn = get_normal_vector(btrian)

strian = SkeletonTriangulation(model)
squad = CellQuadrature(strian,2)

a(u,v) = ∇(v)⋅∇(u)
b1(v) = v*f
b2(v) = v*(bn⋅∇(u))
a3(u,v) = jump(u)*jump(v)
t1 = AffineFETerm(a,b1,trian,quad)
t2 = FESource(b2,btrian,bquad)
t3 = LinearFETerm(a3,strian,squad)
op = AffineFEOperator(Uc,Vc,t1,t2,t3)
uch = solve(op)

#using Gridap.Visualization
#writevtk(trian,"trian",nsubcells=10,cellfields=["uch"=>uch])

e = u - uch

l2(u) = u*u
h1(u) = u*u + a(u,u)

e_l2 = sqrt(sum(integrate(l2(e),trian,quad)))
e_h1 = sqrt(sum(integrate(h1(e),trian,quad)))

tol = 1.e-9
@test e_l2 < tol
@test e_h1 < tol

end # module


## R code 15.3
dlist <- list(
    D_obs = standardize( d$Divorce ),
    D_sd = d$Divorce.SE / sd( d$Divorce ),
    M = standardize( d$Marriage ),
    A = standardize( d$MedianAgeMarriage ),
    N = nrow(d)
)

m15.1 <- ulam(
    alist(
        D_obs ~ dnorm( D_true , D_sd ),
        vector[N]:D_true ~ dnorm( mu , sigma ),
        mu <- a + bA*A + bM*M,
        a ~ dnorm(0,0.2),
        bA ~ dnorm(0,0.5),
        bM ~ dnorm(0,0.5),
        sigma ~ dexp(1)
    ) , data=dlist , chains=4 , cores=4 )

## R code 15.4
precis( m15.1 , depth=2 )

## R code 15.5
dlist <- list(
    D_obs = standardize( d$Divorce ),
    D_sd = d$Divorce.SE / sd( d$Divorce ),
    M_obs = standardize( d$Marriage ),
    M_sd = d$Marriage.SE / sd( d$Marriage ),
    A = standardize( d$MedianAgeMarriage ),
    N = nrow(d)
)

m15.2 <- ulam(
    alist(
        D_obs ~ dnorm( D_est , D_sd ),
        vector[N]:D_est ~ dnorm( mu , sigma ),
        mu <- a + bA*A + bM*M_est[i],
        M_obs ~ dnorm( M_est , M_sd ),
        vector[N]:M_est ~ dnorm( 0 , 1 ),
        a ~ dnorm(0,0.2),
        bA ~ dnorm(0,0.5),
        bM ~ dnorm(0,0.5),
        sigma ~ dexp( 1 )
    ) , data=dlist , chains=4 , cores=4 )

## R code 15.7
N <- 500
A <- rnorm(N)
M <- rnorm(N,-A)
D <- rnorm(N,A)
A_obs <- rnorm(N,A)

## R code 15.8
N <- 100
S <- rnorm( N )
H <- rbinom( N , size=10 , inv_logit(S) )

## R code 15.9
D <- rbern( N ) # dogs completely random
Hm <- H
Hm[D==1] <- NA

## R code 15.10
D <- ifelse( S > 0 , 1 , 0 )
Hm <- H
Hm[D==1] <- NA

## R code 15.11
set.seed(501)
N <- 1000
X <- rnorm(N)
S <- rnorm(N)
H <- rbinom( N , size=10 , inv_logit( 2 + S - 2*X ) )
D <- ifelse( X > 1 , 1 , 0 )
Hm <- H
Hm[D==1] <- NA

## R code 15.12
dat_list <- list(
    H = H,
    S = S )

m15.3 <- ulam(
    alist(
        H ~ binomial( 10 , p ),
        logit(p) <- a + bS*S,
        a ~ normal( 0 , 1 ),
        bS ~ normal( 0 , 0.5 )
    ), data=dat_list , chains=4 )

## R code 15.13
dat_list0 <- list(
    H = H[D==0],
    S = S[D==0] )

m15.4 <- ulam(
    alist(
        H ~ binomial( 10 , p ),
        logit(p) <- a + bS*S,
        a ~ normal( 0 , 1 ),
        bS ~ normal( 0 , 0.5 )
    ), data=dat_list0 , chains=4 )

## R code 15.17
dat_list <- list(
    K = standardize( d$kcal.per.g ),
    B = standardize( d$neocortex.prop ),
    M = standardize( d$logmass )
)

m15.3 <- ulam(
    alist(
        K ~ dnorm( mu , sigma ),
        mu <- a + bB*B + bM*M,
        B ~ dnorm( nu , sigma_B ),
        c(a,nu) ~ dnorm( 0 , 0.5 ),
        c(bB,bM) ~ dnorm( 0, 0.5 ),
        sigma_B ~ dexp( 1 ),
        sigma ~ dexp( 1 )
    ) , data=dat_list , chains=4 , cores=4 )

## R code 15.18
precis( m15.3 , depth=2 )

## R code 15.19
obs_idx <- which( !is.na(d$neocortex.prop) )
dat_list_obs <- list(
    K = dat_list$K[obs_idx],
    B = dat_list$B[obs_idx],
    M = dat_list$M[obs_idx]
)
m15.4 <- ulam(
    alist(
        K ~ dnorm( mu , sigma ),
        mu <- a + bB*B + bM*M,
        B ~ dnorm( nu , sigma_B ),
        c(a,nu) ~ dnorm( 0 , 0.5 ),
        c(bB,bM) ~ dnorm( 0, 0.5 ),
        sigma_B ~ dexp( 1 ),
        sigma ~ dexp( 1 )
    ) , data=dat_list_obs , chains=4 , cores=4 )

## R code 15.22
m15.5 <- ulam(
    alist(
       # K as function of B and M
        K ~ dnorm( mu , sigma ),
        mu <- a + bB*B_merge + bM*M,

       # M and B correlation
        MB ~ multi_normal( c(muM,muB) , Rho_BM , Sigma_BM ),
        matrix[29,2]:MB <<- append_col( M , B_merge ),

       # define B_merge as mix of observed and imputed values
        vector[29]:B_merge <- merge_missing( B , B_impute ),

       # priors
        c(a,muB,muM) ~ dnorm( 0 , 0.5 ),
        c(bB,bM) ~ dnorm( 0, 0.5 ),
        sigma ~ dexp( 1 ),
        Rho_BM ~ lkj_corr(2),
        Sigma_BM ~ exponential(1)
    ) , data=dat_list , chains=4 , cores=4 )

m15.6 <- ulam(
    alist(
        # singing bird model
        ## cat known present/absent:
        notes|RC==0 ~ poisson( lambda ),
        log(lambda) <- a + b*cat,
        ## cat NA:
        notes|RC==1 ~ custom( log_sum_exp(
                log(k) + poisson_lpmf( notes | exp(a + b) ),
                log(1-k) + poisson_lpmf( notes | exp(a) )
            ) ),

        # priors
        a ~ normal(0,1),
        b ~ normal(0,0.5),

        # sneaking cat model
        cat|RC==0 ~ bernoulli(k),
        k ~ beta(2,2)
    ), data=dat , chains=4 , cores=4 )

## R code 15.31
m15.7 <- ulam(
    alist(
        # singing bird model
        notes|RC==0 ~ poisson( lambda ),
        notes|RC==1 ~ custom( log_sum_exp(
                log(k) + poisson_lpmf( notes | exp(a + b) ),
                log(1-k) + poisson_lpmf( notes | exp(a) )
            ) ),
        log(lambda) <- a + b*cat,
        a ~ normal(0,1),
        b ~ normal(0,0.5),

        # sneaking cat model
        cat|RC==0 ~ bernoulli(k),
        k ~ beta(2,2),

        # imputed values
        gq> vector[N]:PrC1 <- exp(lpC1)/(exp(lpC1)+exp(lpC0)),
        gq> vector[N]:lpC1 <- log(k) + poisson_lpmf( notes[i] | exp(a+b) ),
        gq> vector[N]:lpC0 <- log(1-k) + poisson_lpmf( notes[i] | exp(a) )
    ), data=dat , chains=4 , cores=4 )


"""

## spabac!

Function spabac! performs Cholesky forward and back-substitution on
a symmetric skyline global matrix. The loads vector is updated.

###Arguments
```julia
spabac!(kv, loads, kdiag)
```

### Arguments
```julia
* kv::Vector{Float64}       : Skyline vector of global stiffness matrix
* loads::Vector{Float64}    : Load vector (Updated)
* kdiag::Vector{Int}      : Diagonal elemnt index vector
```
"""
function spabac!(kv::Vector{Float64}, loads::Vector{Float64}, kdiag::Vector{Int})
  local x::Float64
  n = size(kdiag, 1)
  loads[1]=loads[1]/kv[1]
  for i in 2:n
    ki = kdiag[i]-i
    l = kdiag[i-1] - ki + 1
    x = loads[i]
    if l !== i
      m = i - 1
      for j in l:m
        x -= kv[ki+j]*loads[j]
      end
    end
    loads[i] = x/kv[ki+i]
  end
  for it in 2:n
    i = n + 2 - it
    ki = kdiag[i] - i
    x = loads[i]/kv[ki+i]
    loads[i] = x
    l = kdiag[i-1] - ki + 1
    if l !== i
      m = i - 1
      for k in l:m
        loads[k] -= x*kv[ki+k]
      end
    end
  end
  loads[1] /= kv[1]
  loads
end


@testset "Find maximum time to nil triangle" begin
    t3 = (
        CausticsEngineering.Vertex3D(5.0, 5.0, 0.0, 0.0, 0.0),
        CausticsEngineering.Vertex3D(6.0, 5.0, 0.0, 0.0, 0.0),
        CausticsEngineering.Vertex3D(5.5, 5.5, 0.0, 0.0, -3.0),
    )
    @test CausticsEngineering.find_maximum_t(t3)[1] ≈ 0.5 / 3.0 atol = 0.01

    t3 = (
        CausticsEngineering.Vertex3D(5.0, 5.0, 0.0, 0.0, 0.0),
        CausticsEngineering.Vertex3D(6.0, 5.0, 0.0, 0.0, 0.0),
        CausticsEngineering.Vertex3D(5.5, 5.5, 0.0, 0.7, -3.0),
    )
    @test CausticsEngineering.find_maximum_t(t3)[1] ≈ 0.5 / 3.0 atol = 0.01

    t3 = (
        CausticsEngineering.Vertex3D(5.0, 5.0, 0.0, 0.0, 0.0),
        CausticsEngineering.Vertex3D(6.0, 5.0, 0.0, 0.0, 0.0),
        CausticsEngineering.Vertex3D(5.5, 5.5, 0.0, -10.0, -3.0),
    )
    @test CausticsEngineering.find_maximum_t(t3)[1] ≈ 0.5 / 3.0 atol = 0.01

    t3 = (
        CausticsEngineering.Vertex3D(5.0, 5.0, 0.01, 0.1, -2.0),
        CausticsEngineering.Vertex3D(6.0, 5.0, 0.01, 0.1, 2.0),
        CausticsEngineering.Vertex3D(5.5, 5.5, 0.01, 0.5, -3.0),
    )
    @test CausticsEngineering.find_maximum_t(t3)[1] ≈ 0.1540 atol = 0.01

    t3 = (
        CausticsEngineering.Vertex3D(5.0, 5.0, 0.01, 0.1, -2.0),
        CausticsEngineering.Vertex3D(6.0, 5.0, 0.01, 0.1, 2.0),
        CausticsEngineering.Vertex3D(5.5, 5.5, 0.01, 0.5, -3.0),
    )
    @test CausticsEngineering.find_maximum_t(t3)[1] ≈ 0.1540 atol = 0.01

end