break_on(:error)

"""
    teralasso(S_kGram, regtype, regparam, regcoef, tol, niter, ninner)

Implementation of Tensor Graphical Lasso (TeraLasso) precision matrix estimator.
Learns a K-order Kronecker sum covariance model.

# Arguments
- `S_kGram::AbstractVector`: length K array of factorwise
    Gram matrices (S_k in the paper).
    NOTE: use the most general accepted type (UnionAll):
    ```julia-repl
    julia> AbstractVector{<:AbstractMatrix{<:Real}} <: AbstractVector === AbstractArray{T, 1} where T
    true

    julia> AbstractVector{<:AbstractMatrix{<:Real}} <: AbstractVector{Any}
    false
    ```
- `regtype::Symbol=:L1`: type of regularization, :L1, :SCAD, or :MCP.
- `regparam::Real=20`: parameter for SCAD/MCP (a in the paper).
- `regcoef::AbstractVector`: length K vector of regularization
    coefficients for each factor (λ).

# Keywords
- `zeta::Real=0.1`: initial step size.
- `c::Real=0.5`: backtracking constant.
    NOTE: 0.5 too large for active region data; use 0.05.
- `tol::Real=1e-6`: tolerance for convergence criterion (ε).
- `niter::Int=50`: maximum number of iterations allowed.
- `ninner::Int=1`: maximum number of backtracking line search iteration.

# Returns:
- `Psi::AbstractVector{AbstractMatrix{Real}}`: length K array of factorwise
    precision matrices
- `out::AbstractVector{AbstractVector{Any}}`: iteration history.
"""

function teralasso(
    S_kGram::AbstractVector,
    regtype::Symbol=:L1,
    regparam::Real=20,
    regcoef::Union{Nothing, AbstractVector}=nothing,
    ;
    zeta::Real=0.1,
    c::Real=0.1,
    tol::Real=1e-6,
    niter::Int=50,
    ninner::Int=11,
)
    # Set eigenvalue upper bound for nonconvex contraints.
    if regtype == :L1
        mu = 1e-6
    elseif regtype == :SCAD
        mu = 1 / (regparam - 1)
    elseif regtype == :MCP
        mu = 1 / regparam
    end
    kappa = sqrt(2 / mu) * 2

    # Preparation
    K = length(S_kGram)
    ps = [size(S_kGram[k], 1) for k in 1:K]
    p = prod(ps)
    regcoef = isnothing(regcoef) ? sqrt.(log.(ps) ./ (p ./ ps)) : regcoef
    #regcoef .+= 1e-9  # inv(cov) may be ill conditioned if p < n and lambda = 0
    S_tilde = [S - (K - 1) / K * mean(diag(S)) * I for S in S_kGram]

    # Initlalization
    logdet = 0
    logdetN = 0 # updated outside inner loop
    Psi = 1 / K * [Matrix{Float64}(I, ps[k], ps[k]) for k = 1:K]
    PsiN = 1 / K * [Matrix{Float64}(I, ps[k], ps[k]) for k = 1:K]
    M = 1 / K * [Matrix{Float64}(I, ps[k], ps[k]) for k = 1:K]
    out = Array{Any}(undef, niter)

    for iter = 1:niter
        # Backtracking line search
        for cnt = 1:ninner
            for k = 1:K
                PsiN[k] = shrink_regularizer(
                    Psi[k] - zeta * (S_tilde[k] - M[k]),
                    regcoef[k],
                    zeta,
                    regtype,
                    regparam,
                )
            end

            result, logdetN = eval_cond(PsiN, Psi, M, S_kGram, S_tilde, ps, zeta, logdet, kappa)
            #println([iter, cnt, zeta, logdetN])

            if result == 1
                break
            end

            if cnt < ninner - 1
                zeta *= c
            elseif cnt == ninner - 1 # safe step size is not actually safe
                min_eig = sum([minimum(eigvals(Psi[k])) for k = 1:K])
                zeta = min_eig ^ 2 / 2
            else # cnt == ninner
                @warn "Backtracking fails. Use more inner iterations / smaller c."
            end
        end

        # Compute convergence
        dPsi = [PsiN[k] - Psi[k] for k = 1:K]
        nrm = norm([dPsi[k] - I * mean(diag(dPsi[k])) for k = 1:K])^2
        out[iter] = [fun(S_kGram, PsiN), reg(PsiN, regcoef), nrm, PsiN, zeta]
        if iter > 3 && all([out[i][3] <= tol for i in iter-3:iter])
            break
            # Stopping criterion based on fun val doesn't work. f keeps decreasing.
            # The problem is the inner product in f doesn't work properly.
            # Is there a reason why Kristjan uses nrm as the criterion?
            # if abs(out[iter][1] - out[iter-1][1]) < tol
            #     break
            # end
        end

        # Set next initial step
        ~, ~, MsN = project_ksum(PsiN, ps)
        zeta = barzilai_borwein(PsiN, Psi, MsN, M, ps)

        # Update iterate
        M .= MsN  # (copy, deepcopy, vs .=) Use === to check if identical
        Psi .= PsiN
        logdet = logdetN

        #println(iter, ": ", [out[iter][1], out[iter][2]])
        #display(heatmaps(PsiN))
        #println("-----------------------------")
    end
    return Psi, out
end

function barzilai_borwein(PsiN, Psi, MsN, M, ps)
    Num = 0
    Den = 0
    K = size(PsiN, 1)
    dgI = zeros(K)
    for k = 1:K
        dgI[k] = mean(diag((PsiN[k] - Psi[k])))
        Num = Num + prod(ps[[1:k-1; k+1:K]]) * sum((PsiN[k] - Psi[k] - I * dgI[k]) .^ 2)
        dK = PsiN[k] - Psi[k]
        for j = 1:K
            dKK = M[j] - MsN[j]
            if k == j
                Den = Den + prod(ps[[1:k-1; k+1:K]]) * sum(sum(dK .* dKK))
            else
                Den = Den + tr(dK) * tr(dKK) * prod(ps) / (ps[k] * ps[j])
            end
        end
    end
    Num = Num + mean(dgI)^2 * prod(ps)
    zeta = Num / Den
    if Den == 0
        zeta = 0
    elseif abs(zeta) > 1e2
        zeta = 1e2 * sign(zeta)
    elseif zeta < 0
        zeta = 1
    end
    return zeta
end

function eval_cond(PsiN, Psi, M, S_kGram, S_tilde, ps, zeta, logdet, kappa)
    K = length(ps)
    p = prod(ps)
    ms = [p / ps[k] for k = 1:K]
    ergNvec = zeros(p, 1) # eigval of the Kronecker sum of PsiN
    for k = 1:K
        ergNvec = ergNvec + repeat(
            kron(eigvals(PsiN[k]), ones(prod(ps[k+1:end]),1)),
            prod(ps[1:k-1]),
            1,
        ) # repmat is eqv to
    end
    #ergNvec = reduce(⊕, [eigvals(PsiN[k]) for k in 1:K])

    if minimum(ergNvec) <= 0
        logdetN = NaN
        result = 0
        return result, logdetN
    end

    logdetN = -sum(log.(ergNvec))
    if maximum(ergNvec) > kappa
        result = 0
        return result, logdetN
    end

    fNew = logdetN + sum([ms[k] * sum(S_kGram[k] .* PsiN[k]) for k = 1:K])
    # Ψ = PsiN[1] ⊕ PsiN[2]
    # SSSS = S_kGram[1] ⊕ S_kGram[2]
    # my_fNew = -log(det(Ψ)) + sum(Ψ .* SSSS)
    # @bp # sum(Ψ .* SSSS) != sum([ms[k] * sum(S_kGram[k] .* PsiN[k]) for k = 1:K])
    # println([fNew, fun(S_kGram, PsiN), my_fNew])

    dPsi = [PsiN[k] - Psi[k] for k = 1:K]
    tau = [mean(diag(dPsi[k])) for k = 1:K] # dgI, tr(D_k)/d_k
    Num = sum([ms[k] * norm(dPsi[k] - tau[k] * I)^2 for k = 1:K])
    Num += mean(tau)^2 * p  # Norm of change
    Q = logdet# + sum([ms[k] * sum(S_kGram[k].*Psi[k]) for k = 1:K]) + Num*K/(2*zeta)
    for k = 1:K #
        Q += sum(S_kGram[k] .* Psi[k]) * ms[k] + Num / (2 * zeta)
        for j = 1:K
            if j == k
                Q += sum((S_tilde[k] - M[k]) .* dPsi[k]) * prod(ps) / ps[k]
            else
                Q -= tr(M[k] - S_tilde[k]) * tr(dPsi[j]) * prod(ps) / (ps[k]*ps[j])
            end
        end
    end

    result = fNew <= Q

    return result, logdetN
end

function fun(S_kGram, Ψs)
    # - logdet(Ω) + < Ω,S >
    ps = [size(S, 1) for S in S_kGram]
    K = length(ps)
    p = prod(ps)
    T = zeros(ps...)
    for (k, Ψ) in enumerate(Ψs)
        dims = ones(Int, length(ps))
        dims[k] = ps[k]
        T .+= reshape(eigvals(Ψ), dims...)
    end
    # 2x faster than this:
    # T = diag(reduce(⊕, [Diagonal(eigvals(Ψs[k])) for k in 1:K]))
    logdet = sum(log.(T))

    trace = sum([sum(S_kGram[k] .* Ψs[k]) * p / ps[k] for k in 1:K])

    return -logdet + trace
end

function reg(Ψs, ρs)
    ps = [size(Ψ, 1) for Ψ in Ψs]
    p = prod(ps)
    result = 0
    for k in 1:length(Ψs)
        mk = p ÷ ps[k]
        offdiag = Ψs[k] - diagm(diag(Ψs[k]))
        result += mk * ρs[k] * sum(abs.(offdiag))
    end
    return result
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

# asinh, acosh, and atanh are heavily based on FDLIBM code:
# e_sinh.c, e_sinhf, e_cosh.c, e_coshf, s_tanh.c, s_tanhf.c, s_asinh.c,
# s_asinhf.c, e_acosh.c, e_coshf.c, e_atanh.c, and e_atanhf.c
# that are made available under the following licence:

# ====================================================
# Copyright (C) 1993 by Sun Microsystems, Inc. All rights reserved.
#
# Developed at SunSoft, a Sun Microsystems, Inc. business.
# Permission to use, copy, modify, and distribute this
# software is freely granted, provided that this notice
# is preserved.
# ====================================================

@inline function exthorner(x, p::Tuple)
	# polynomial evaluation using compensated summation.
	# much more accurate, especially when lo can be combined with other rounding errors
    hi, lo = p[end], zero(x)
    for i in length(p)-1:-1:1
        pi = p[i]
        prod = hi*x
        err = fma(hi, x, -prod)
        hi = pi+prod
        lo = fma(lo, x, prod - (hi - pi) + err)
    end
    return hi, lo
end

# Hyperbolic functions
# sinh methods
H_SMALL_X(::Type{Float64}) = 2.0^-28
H_MEDIUM_X(::Type{Float64}) = 22.0

H_SMALL_X(::Type{Float32}) = 2f-12
H_MEDIUM_X(::Type{Float32}) = 9f0

H_LARGE_X(::Type{Float64}) = 709.7822265633563 # nextfloat(709.7822265633562)

H_LARGE_X(::Type{Float32}) = 88.72283f0

SINH_SMALL_X(::Type{Float64}) = 2.1
SINH_SMALL_X(::Type{Float32}) = 3.0f0

# For Float64, use DoubleFloat scheme for extra accuracy
function sinh_kernel(x::Float64)
    x2 = x*x
    x2lo = fma(x,x,-x2)
    hi_order = evalpoly(x2, (8.333333333336817e-3, 1.9841269840165435e-4,
                             2.7557319381151335e-6, 2.5052096530035283e-8,
                             1.6059550718903307e-10, 7.634842144412119e-13,
                             2.9696954760355812e-15))
    hi,lo = exthorner(x2, (1.0, 0.16666666666666635, hi_order))
    return muladd(x, hi, muladd(x, lo, x*x2lo*0.16666666666666635))
end
# For Float32, using Float64 is simpler, faster, and doesn't require FMA
function sinh_kernel(x::Float32)
    x=Float64(x)
    res = evalpoly(x*x, (1.0, 0.1666666779967941, 0.008333336726447933,
                         0.00019841001151414065, 2.7555538207080807e-6,
                         2.5143389765825282e-8, 1.6260094552031644e-10))
    return Float32(res*x)
end

function sinh(x::T) where T<:Union{Float32,Float64}
    # Method
    # mathematically sinh(x) is defined to be (exp(x)-exp(-x))/2
    #    1. Sometimes replace x by |x| (sinh(-x) = -sinh(x)).
    #    2. Find the branch and the expression to calculate and return it
    #      a)   0 <= x < SINH_SMALL_X
    #               approximate sinh(x) with a  minimax polynomial
    #      b)   SINH_SMALL_X <= x < H_LARGE_X
    #               return sinh(x) = (exp(x) - exp(-x))/2
    #      d)   H_LARGE_X  <= x
    #               return sinh(x) = exp(x/2)/2 * exp(x/2)
    #               Note that this branch automatically deals with Infs and NaNs

    absx = abs(x)
    if absx <= SINH_SMALL_X(T)
        return sinh_kernel(x)
    elseif absx >= H_LARGE_X(T)
        E = exp(T(.5)*absx)
        return copysign(T(.5)*E*E, x)
    end
    E = exp(absx)
    return copysign(T(.5)*(E - 1/E),x)
end

COSH_SMALL_X(::Type{T}) where T= one(T)

function cosh_kernel(x2::Float32)
    return evalpoly(x2, (1.0f0, 0.49999997f0, 0.041666888f0, 0.0013882756f0, 2.549933f-5))
end

function cosh_kernel(x2::Float64)
    return evalpoly(x2, (1.0, 0.5000000000000002, 0.04166666666666269,
                         1.3888888889206764e-3, 2.4801587176784207e-5,
                         2.7557345825742837e-7, 2.0873617441235094e-9,
                         1.1663435515945578e-11))
end

function cosh(x::T) where T<:Union{Float32,Float64}
    # Method
    # mathematically cosh(x) is defined to be (exp(x)+exp(-x))/2
    #    1. Replace x by |x| (cosh(x) = cosh(-x)).
    #    2. Find the branch and the expression to calculate and return it
    #      a)   x <= COSH_SMALL_X
    #               approximate sinh(x) with a minimax polynomial
    #      b)   COSH_SMALL_X <= x < H_LARGE_X
    #               return cosh(x) = = (exp(x) + exp(-x))/2
    #      e)   H_LARGE_X  <= x
    #               return cosh(x) = exp(x/2)/2 * exp(x/2)
    #      			Note that this branch automatically deals with Infs and NaNs

    absx = abs(x)
    if absx <= COSH_SMALL_X(T)
        return cosh_kernel(x*x)
    elseif absx >= H_LARGE_X(T)
        E = exp(T(.5)*absx)
        return T(.5)*E*E
    end
    E = exp(absx)
    return T(.5)*(E + 1/E)
end

# tanh methods
TANH_LARGE_X(::Type{Float64}) = 44.0
TANH_LARGE_X(::Type{Float32}) = 18.0f0
TANH_SMALL_X(::Type{Float64}) = 1.0
TANH_SMALL_X(::Type{Float32}) = 1.3862944f0       #2*log(2)
@inline function tanh_kernel(x::Float64)
    return evalpoly(x, (1.0, -0.33333333333332904, 0.13333333333267555,
                        -0.05396825393066753, 0.02186948742242217,
                        -0.008863215974794633, 0.003591910693118715,
                        -0.0014542587440487815, 0.0005825521659411748,
                        -0.00021647574085351332, 5.5752458452673005e-5))
end
@inline function tanh_kernel(x::Float32)
    return evalpoly(x, (1.0f0, -0.3333312f0, 0.13328037f0,
                        -0.05350336f0, 0.019975215f0, -0.0050525228f0))
end
function tanh(x::T) where T<:Union{Float32, Float64}
    # Method
    # mathematically tanh(x) is defined to be (exp(x)-exp(-x))/(exp(x)+exp(-x))
    #    1. reduce x to non-negative by tanh(-x) = -tanh(x).
    #    2. Find the branch and the expression to calculate and return it
    #      a) 0 <= x < H_SMALL_X
    #             Use a minimax polynomial over the range
    #      b) H_SMALL_X <= x < TANH_LARGE_X
    #           1 - 2/(exp(2x) + 1)
    #      c) TANH_LARGE_X <= x
    #            return 1
    abs2x = abs(2x)
    abs2x >= TANH_LARGE_X(T) && return copysign(one(T), x)
    abs2x <= TANH_SMALL_X(T) && return x*tanh_kernel(x*x)
    k = exp(abs2x)
    return copysign(1 - 2/(k+1), x)
end

# Inverse hyperbolic functions
AH_LN2(::Type{Float64}) = 6.93147180559945286227e-01
AH_LN2(::Type{Float32}) = 6.9314718246f-01
# asinh methods
function asinh(x::T) where T <: Union{Float32, Float64}
    # Method
    # mathematically asinh(x) = sign(x)*log(|x| + sqrt(x*x + 1))
    # is the principle value of the inverse hyperbolic sine
    # 1. Find the branch and the expression to calculate and return it
    #    a) |x| < 2^-28
    #        return x
    #    b) |x| < 2
    #        return sign(x)*log1p(|x| + x^2/(1 + sqrt(1+x^2)))
    #    c) 2 <= |x| < 2^28
    #        return sign(x)*log(2|x|+1/(|x|+sqrt(x*x+1)))
    #    d) |x| >= 2^28
    #        return sign(x)*(log(x)+ln2))
    if isnan(x) || isinf(x)
        return x
    end
    absx = abs(x)
    if absx < T(2)
        # in a)
        if absx < T(2)^-28
            return x
        end
        # in b)
        t = x*x
        w = log1p(absx + t/(T(1) + sqrt(T(1) + t)))
    elseif absx < T(2)^28
        # in c)
        t = absx
        w = log(T(2)*t + T(1)/(sqrt(x*x + T(1)) + t))
    else
        # in d)
        w = log(absx) + AH_LN2(T)
    end
    return copysign(w, x)
end

# acosh methods
@noinline acosh_domain_error(x) = throw(DomainError(x, "acosh(x) is only defined for x ≥ 1."))
function acosh(x::T) where T <: Union{Float32, Float64}
    # Method
    # mathematically acosh(x) if defined to be log(x + sqrt(x*x-1))
    # 1. Find the branch and the expression to calculate and return it
    #     a) x = 1
    #         return log1p(t+sqrt(2.0*t+t*t)) where t=x-1.
    #     b) 1 < x < 2
    #         return log1p(t+sqrt(2.0*t+t*t)) where t=x-1.
    #     c) 2 <= x <
    #         return log(2x-1/(sqrt(x*x-1)+x))
    #     d) x >= 2^28
    #         return log(x)+ln2
    # Special cases:
    #     if x < 1 throw DomainError

    isnan(x) && return x

    if x < T(1)
        return acosh_domain_error(x)
    elseif x == T(1)
        # in a)
        return T(0)
    elseif x < T(2)
        # in b)
        t = x - T(1)
        return log1p(t + sqrt(T(2)*t + t*t))
    elseif x < T(2)^28
        # in c)
        t = x*x
        return log(T(2)*x - T(1)/(x+sqrt(t - T(1))))
    else
        # in d)
        return log(x) + AH_LN2(T)
    end
end

# atanh methods
@noinline atanh_domain_error(x) = throw(DomainError(x, "atanh(x) is only defined for |x| ≤ 1."))
function atanh(x::T) where T <: Union{Float32, Float64}
    # Method
    # 1.Reduced x to positive by atanh(-x) = -atanh(x)
    # 2. Find the branch and the expression to calculate and return it
    #     a) 0 <= x < 0.5
    #         return 0.5*log1p(2x/(1-x))
    #     b) 0.5 <= x <= 1
    #         return 0.5*log((x+1)/(1-x))
    # Special cases:
    #    if |x| > 1 throw DomainError
    isnan(x) && return x

    absx = abs(x)

    if absx > 1
        atanh_domain_error(x)
    end
    if absx < T(0.5)
        # in a)
        t = log1p(T(2)*absx/(T(1)-absx))
    else
        # in b)
        t = log((T(1)+absx)/(T(1)-absx))
    end
    return T(0.5)*copysign(t, x)
end


using DeterminantalPointProcesses
using Test
using Random, LinearAlgebra, IterTools
import Combinatorics: combinations

rng = MersenneTwister(42)

n = 5
k = 2
A = rand(rng, n, n)
L = Symmetric(A * A')
dpp = DPP(L)

@testset "Ensure correct pmf and logpmf computation" begin
    @testset "For DPP" begin
        # compute the true distribution
        true_pmf = Float64[]
        true_logpmf = Float64[]
        for z in subsets(1:n)
            push!(true_pmf, pmf(dpp, z))
            push!(true_logpmf, logpmf(dpp, z))
        end
        @test sum(true_pmf) ≈ 1.0
        @test all(true_pmf .<= 1.0)
        @test all(true_logpmf .<= 0.0)
    end

    @testset "For k-DPP" begin
        # compute the true distribution
        true_pmf = Float64[]
        true_logpmf = Float64[]
        for z in combinations(1:n, k)
            push!(true_pmf, pmf(dpp, z, k))
            push!(true_logpmf, logpmf(dpp, z, k))
        end
        @test sum(true_pmf) ≈ 1.0
        @test all(true_pmf .<= 1.0)
        @test all(true_logpmf .<= 0.0)
    end
end

@testset "Ensure correct sampling from DPP" begin
    # compute the true distribution
    all_subsets = []
    true_pmf = Float64[]
    for (i, z) in enumerate(subsets(1:n))
        push!(true_pmf, pmf(dpp, z))
        push!(all_subsets, (z, i))
    end
    global subset_to_idx = Dict(all_subsets)

    @testset "Exact sampling" begin
        nb_samples = 1000
        samples = rand(dpp, nb_samples)

        # compute the empirical distribution
        empirical_pmf = zeros(Float64, length(true_pmf))
        for z in samples
            empirical_pmf[subset_to_idx[z]] += 1
        end
        empirical_pmf ./= nb_samples

        # ensure that the empirical pmf is close to the true pmf
        total_variation = maximum(abs.(true_pmf - empirical_pmf))
        @test total_variation ≈ 0.0 atol=1e-1
    end

    @testset "MCMC sampling" begin
        nb_samples = 1000
        samples, state = randmcmc(dpp, nb_samples, return_final_state=true)

        # ensure that L_z_inv makes sense (i.e., noise did not accumulate)
        z, L_z_inv = state
        @test sum(L_z_inv[z, z] * dpp.L[z, z] - I) ≈ 0

        # compute the empirical distribution
        empirical_pmf = zeros(Float64, length(true_pmf))
        for z in samples
            empirical_pmf[subset_to_idx[z]] += 1
        end
        empirical_pmf ./= nb_samples

        # ensure that the empirical pmf is close to the true pmf
        total_variation = maximum(abs.(true_pmf - empirical_pmf))
        @test total_variation ≈ 0.0  atol=1e-1
    end
end

@testset "Ensure correct sampling from k-DPP" begin
    # compute the true distribution
    all_k_subsets = []
    true_pmf = Float64[]
    for (i, z) in enumerate(combinations(1:n, k))
        push!(true_pmf, pmf(dpp, z))
        push!(all_k_subsets, (z, i))
    end
    true_pmf ./= sum(true_pmf)
    k_subset_to_idx = Dict(all_k_subsets)

    @testset "Exact sampling" begin
        nb_samples = 10000
        samples = rand(dpp, nb_samples, k)

        # ensure that samples are of proper cardinality
        @test all(map(length, samples) .== k)

        # compute the empirical distribution
        empirical_pmf = zeros(Float64, length(true_pmf))
        for z in samples
            empirical_pmf[k_subset_to_idx[z]] += 1
        end
        empirical_pmf ./= nb_samples

        # ensure that the empirical pmf is close to the true pmf
        total_variation = maximum(abs.(true_pmf - empirical_pmf))
        @test total_variation ≈ 0.0 atol=1e-1
    end

    @testset "MCMC sampling" begin
        nb_samples = 1000
        samples, state = randmcmc(dpp, nb_samples, k, return_final_state=true)

        # ensure that L_z_inv makes sense (i.e., noise did not accumulate)
        z, L_z_inv = state
        @test sum(L_z_inv[z, z] * dpp.L[z, z] - I) ≈ 0 atol = 1e-1

        # ensure that samples are of proper cardinality
        @test all(map(length, samples) .== k)

        # compute the empirical distribution
        empirical_pmf = zeros(Float64, length(true_pmf))
        for z in samples
            empirical_pmf[k_subset_to_idx[z]] += 1
        end
        empirical_pmf ./= nb_samples

        # ensure that the empirical pmf is close to the true pmf
        total_variation = maximum(abs.(true_pmf - empirical_pmf))
        @test_broken total_variation ≈ 0.0; atol=1e-1
    end
end


module Models
export ModelParams

using FactorizationMachines: FMMatrix, FMFloat

abstract ModelParams

immutable GaussianModelParams <: ModelParams
    k0::Bool
    k1::Bool
    num_factors::Int

    mean::Float64
    stddev::Float64

    GaussianModelParams(; k0 = true, k1 = true, num_factors = 8, mean = .0, stddev = .01) = new(k0, k1, num_factors, mean, stddev)
end
const gauss = GaussianModelParams

function init(params::GaussianModelParams, X::FMMatrix, y::Vector{FMFloat})
    # initialization
    num_attributes, num_samples = size(X)
    # sanity check
    assert(length(y) == num_samples)

    # create initial model
    w0 = .0
    w = zeros(num_attributes)
    V = randn(params.num_factors, num_attributes) .* params.stddev + params.mean

    # new model
    model = FMModel(params.k0, params.k1, w0, w, V, params.num_factors)
end

type FMModel
    k0::Bool
    k1::Bool

    w0::FMFloat
    w::Vector{FMFloat}
    V::Matrix{FMFloat}

    num_factors::Int
end

function predict_instance!(model::FMModel,
                           idx::StridedVector{Int64}, x::StridedVector{FMFloat},
                           f_sum::Vector{FMFloat}, sum_sqr::Vector{FMFloat})
    fill!(f_sum, .0)
    fill!(sum_sqr, .0)
    result = zero(FMFloat)
    if model.k0
        result += model.w0
    end
    if model.k1
        for i in 1:length(idx)
            result += model.w[idx[i]] * x[i]
        end
    end
    for f in 1:model.num_factors
        for i in 1:length(idx)
            d = model.V[f,idx[i]] * x[i]
            f_sum[f] += d
            sum_sqr[f] += d * d
        end
        result += 0.5 * (f_sum[f] * f_sum[f] - sum_sqr[f])
    end
    result
end

end


# based on https://dlmf.nist.gov/11.6.i
# not valid for real z < 0

const order_large = 10

_K_large(ν, z) = (1 / pi) * sum(gamma(k + 0.5) * (0.5z)^(ν - 2k - 1) / gamma(ν + 0.5 - k) for k in 0:order_large)

_K0_large(z) = _K_large(0, z)

_M_large(ν, z) = (1 / pi) * sum((-1)^(k + 1) * gamma(k + 0.5) * (0.5z)^(ν - 2k - 1) / gamma(ν + 0.5 - k) for k in 0:order_large)

_M0_large(z) = _M_large(0, z)


module VStatistic
import GSL:sf_hyperg_2F1
# could have also used scipy.special.hyp2f1 via @pyimports

function vstat(n::Integer, p::Integer, R²)
    R² = R² <= 0 ? .0001 : R²

    r = ((p - 1) * (1 - R²)) / ((n - p) * R²)
    g = min(r, 1)
    g = .4999 < g < .5001 ? .5001 : g

    z = (g - √(g - g^2)) / (2g - 1)
    α = acos((1 - z) / √(1 - 2z * (1 - z)))

    v = begin
        ((2 * cos(α) * gamma((p + 2) / 2)) /
        (√(π) * gamma((p + 1) / 2))) *
        (sf_hyperg_2F1(.5, (1 - p) / 2, 3/2, cos(α)^2) -
        sin(α)^(p - 1))
    end

    return v
end

function sample_size(p::Integer, R²; v=.8)
    n = p
    while vstat(n, p, R²) < v
        n += 1
    end
    return n, round(n / p, 2)
end

end


# NOTE: Add staggered grids
struct Grid
    xmax::Float64
    ymax::Float64

    nx::Int64
    ny::Int64

    Δx::Float64
    Δy::Float64

    xRange::Array{Float64,1}
    yRange::Array{Float64,1}

    function Grid(xmax, ymax, nx, ny)
        Δx = 2*xmax/(nx-1)
        Δy = 2*ymax/(ny-1)

        xRange = (-xmax:Δx:xmax)
        yRange = (-ymax:Δy:ymax)

        return new(xmax, ymax, nx, ny, Δx, Δy, xRange, yRange)
    end
end


using Autotuning, CSV, DataFrames, Distributions

function sample_x_y(samples::Int,
                    distribution_x::Distribution,
                    distribution_y::Distribution)

    configuration = Configuration((x = Parameter(0.5, [0.0, 1.0]),
                                   y = Parameter(0.5, [0.0, 1.0])))

    distribution = (x = distribution_x,
                    y = distribution_y)

    configuration = perturb(configuration, distribution)

    sampled_data = configuration_table(configuration)

    for i = 2:samples
        configuration = perturb(configuration, distribution)
        push!(sampled_data, configuration)
    end

    return sampled_data
end

function sample()
    samples = 1000
    distributions = [
        (
            x = Truncated(Normal(0.5, 0.08), 0.0, 1.0),
            y = Truncated(Normal(0.5, 0.08), 0.0, 1.0),
            name = "mu:0.5,sig:0.08"
        ),
        (
            x = Truncated(Normal(0.5, 0.13), 0.0, 1.0),
            y = Truncated(Normal(0.5, 0.13), 0.0, 1.0),
            name = "mu:0.5,sig:0.13"
        ),
        (
            x = Truncated(Normal(0.1, 0.08), 0.0, 1.0),
            y = Truncated(Normal(0.5, 0.13), 0.0, 1.0),
            name = "mu:(0.1,0.5),sig:(0.08,0.13)"
        ),
        (
            x = MixtureModel(
                [
                    Truncated(Normal(0.0, 0.08), 0.0, 0.4),
                    Truncated(Normal(1.0, 0.08), 0.6, 1.0)
                ]
            ),
            y = MixtureModel(
                [
                    Truncated(Normal(0.0, 0.08), 0.0, 0.4),
                    Truncated(Normal(1.0, 0.08), 0.6, 1.0)
                ]
            ),
            name = "mu:(0.0,1.0)*2,sig:(0.08)*2"
        )
    ]

    sampled_data = sample_x_y(samples, distributions[1][:x], distributions[1][:y])
    sampled_data.name = repeat([distributions[1][:name]], samples)

    for distribution in distributions[2:end]
        new_data = sample_x_y(samples, distribution[:x], distribution[:y])
        new_data.name = repeat([distribution[:name]], samples)

        append!(sampled_data, new_data)
    end

    new_data = sample_x_y(samples, distributions[1][:x], distributions[end][:y])
    new_data.name = repeat([distributions[1][:name] * "x"], samples)

    append!(sampled_data, new_data)

    new_data = sample_x_y(samples, distributions[end][:x], distributions[1][:y])
    new_data.name = repeat([distributions[end][:name] * "x"], samples)

    append!(sampled_data, new_data)

    CSV.write("sampled_data.csv", sampled_data)
end

sample()


import NonlinearEquations
import PyPlot
import SparseArrays

#this block corresponds to the "Parameter values" table
D = 1e-9#m^2/s
M = 4.40098e-2#kg/mol
q = 1.5491e-10#m/s
V_ϕ = 34.8395e-6#m^3/mol
κ = 1e-13#m^2
μ = 6.7549e-5#kg/m/s
ρ_s = 1072.4#kg/m^3
ρ_w = 1066.5#kg/m^3
θ = 0.0844#unitless
g = 9.8067#m/s
Λ = κ / (μ * θ)
Γ = -g * (M - ρ_w * V_ϕ)

#this block corresponds to the "Initial conditions" table
#nx = 1000
#nz = 100
#nx = 50
#nz = 5
#nx = 100
#nz = 10
nx = 100
nz = 10
mins = [0, 0]
maxs = [1000, 100]
deltax = (maxs[1] - mins[1]) / nx
deltaz = (maxs[2] - mins[2]) / nz
tfinal = 3.1536e10#seconds
c0 = [j == nz && i * deltax >= 100 && i * deltax <= 500 ? 750.6 : 4.4457e-2 for j = 1:nz, i = 1:nx]#this is similar to the initial conditions in the document, but matches the dirichlet boundary condition

function plot(args...; filename=false)
	numplots = div(length(args), 2)
	fig, axs = PyPlot.subplots(1, numplots, figsize=(4 * numplots + 1, 4))
	for i = 1:numplots
		x = args[(i - 1) * 2 + 1]
		title = args[(i - 1) * 2 + 2]
		if startswith(title, "log")
			img = axs[i].imshow(x, origin="lower", interpolation="nearest", extent=[mins[1], maxs[1], mins[2], 10 * maxs[2]], vmin=-3)
		else
			img = axs[i].imshow(x, origin="lower", interpolation="nearest", extent=[mins[1], maxs[1], mins[2], 10 * maxs[2]])
		end
		axs[i].set(xlabel="x [meters]", ylabel="z [decimeters]")
		axs[i].set_title(title)
		fig.colorbar(img, ax=axs[i])
	end
	fig.tight_layout()
	display(fig)
	println()
	if  filename != false
		fig.savefig(filename)
	end
	PyPlot.close(fig)
end

#maybe the key is just that it is diffusing in for their example. change the transport boundary conditions so that there’s no CO2 in the reservoir initially, and it only diffuses in
#still, something weird is going on with the pressure -- gravity seems to be absent from the pressure

z(j) = j * deltaz
ij2k(i, j) = j + (i - 1) * nz
@NonlinearEquations.equations function pressure(p, c)
	NonlinearEquations.setnumequations(length(p))
	#do the internal nodes
	for i = 2:nx - 1, j = 1:nz
		k = ij2k(i, j)
		ju = j < nz ? j + 1 : j
		jd = j > 1 ? j - 1 : j
		il = i > 1 ? i - 1 : i
		ir = i < nx ? i + 1 : i
		kl = ij2k(il, j)
		kr = ij2k(ir, j)
		kd = ij2k(i, jd)
		ku = ij2k(i, ju)
		NonlinearEquations.addterm(k, (p[kr] - 2 * p[k] + p[kl]) / deltax ^ 2)
		NonlinearEquations.addterm(k, (p[ku] - 2 * p[k] + p[kd]) / deltaz ^ 2)
		NonlinearEquations.addterm(k, ((ρ_w * g - Γ * c[ku])* (ju - j) - (ρ_w * g - Γ * c[k]) * (j - jd)) * deltaz / deltaz ^ 2)#gravity term
	end
	#do the bottom nodes on the east/west boundary
	k_east = ij2k(nx, nz)
	k_west = ij2k(1, nz)
	#NonlinearEquations.addterm(k_east, p[k_east] - 1.99998975026e7)
	NonlinearEquations.addterm(k_east, p[k_east] - 2e7)
	NonlinearEquations.addterm(k_west, p[k_west] - 2e7)
	#do the east/west boundaries
	for j = 1:nz - 1
		k_east = ij2k(nx, j)
		ku_east = ij2k(nx, j + 1)
		k_west = ij2k(1, j)
		ku_west = ij2k(1, j + 1)
		#NonlinearEquations.addterm(k_east, p[k_east] - p[ku_east] - g * (ρ_w * (1 - c[k_east] * V_ϕ) + M * c[k_east]) * deltaz)
		#NonlinearEquations.addterm(k_west, p[k_west] - p[ku_west] - g * (ρ_w * (1 - c[k_west] * V_ϕ) + M * c[k_west]) * deltaz)
		#NonlinearEquations.addterm(k_west, p[k_west] - g * ρ_w * (nz - 1 + j) * deltaz)
		#NonlinearEquations.addterm(k_east, p[k_east] + g * ρ_w * (nz - 1 + j) * deltaz)
		NonlinearEquations.addterm(k_west, p[k_west] - 2e7 - g * ρ_w * (nz - 1 - j) * deltaz)
		NonlinearEquations.addterm(k_east, p[k_east] - 2e7 - g * ρ_w * (nz - 1 - j) * deltaz)
	end
end
@NonlinearEquations.equations function concentration(p, c)
	NonlinearEquations.setnumequations(length(c))
	cell_volume = deltax * deltaz * 1#assume it is 1m in the y direction
	areax = deltaz * 1
	areaz = deltax * 1
	for i = 1:nx, j = 1:nz
		k = ij2k(i, j)
		if j == nz && i * deltax >= 100 && i * deltax <= 500
			NonlinearEquations.addterm(k, c[k] - 750.6)
		else
			kl = i > 1 ? ij2k(i - 1, j) : k#node to the left
			kr = i < nx ? ij2k(i + 1, j) : k#node to the right
			kd = j > 1 ? ij2k(i, j - 1) : k#node below
			ku = j < nz ? ij2k(i, j + 1) : k#node above
			for ko in [kl, kr]
				k_upstream = p[ko] - p[k] > 0 ? ko : k
				NonlinearEquations.addterm(k, -Λ * (p[ko] - p[k]) / deltax * c[k_upstream] * areax / cell_volume)
			end
			for (ko, gravity_sign) in zip([kd, ku], [1, -1])
				k_viscous_upstream = p[ko] - p[k] > 0 ? ko : k
				NonlinearEquations.addterm(k, -Λ * (p[ko] - p[k]) / deltaz * c[k_viscous_upstream] * areaz / cell_volume)
				if k != ko#don't let gravity move stuff across the boundary
					k_gravity_upstream = gravity_sign == 1 ? k : ko
					NonlinearEquations.addterm(k, gravity_sign * Λ * (g * ρ_w - Γ * c[k_gravity_upstream]) / deltaz * c[k_gravity_upstream] * areaz / cell_volume)
				end
			end
		end
	end
end

#do explicit time stepping
dt = 1e4 / nz^2 * 10^2
c = c0# .+ 1e-3 * rand(size(c0))
p = zeros(size(c0))
t = 0.0
i = 0
frame_counter = 1
#=
	c = zeros(nz, nx)
	A = pressure_p(zeros(nz, nx), c)
	b = -pressure_residuals(zeros(nz, nx), c)
	p = reshape(A \ b, nz, nx)
	@show hash(p)
	plot(log10.(c), "log10(concentration)", p, "pressure")
	error("blah")
	=#
A = pressure_p(zeros(nz, nx), c0)
b = -pressure_residuals(zeros(nz, nx), c0)
p0 = reshape(A \ b, nz, nx)
p = p0
while t < 365 * 10
	global c, t, p, i, frame_counter
	A = pressure_p(zeros(nz, nx), c)
	b = -pressure_residuals(zeros(nz, nx), c)
	lastp = p
	p = reshape(A \ b, nz, nx)
	dc = reshape(-concentration_residuals(p, c), nz, nx)
	lastc = c
	c += dt * dc
	t += dt / 3600 / 24
	if mod(i, 100) == 0
		tstring = round(t; digits=1)
		plot(c, "log10(concentration) at t=$tstring", p, "pressure", p - p0, "p-p0", c - c0, "c-c0", p - lastp, "dp", c - lastc, "dc")
		frame_counter += 1
	end
	i += 1
end
tstring = round(t; digits=1)
plot(log10.(c), "log10(concentration) at t=$tstring days", p, "pressure at t=$tstring days", c - c0, "c-c0")
#plot(log10.(c), "log10(concentration) at t=$tstring days", p, "pressure at t=$tstring days"; filename="figs/frame_$(lpad(frame_counter, 6, '0')).png")
#run(`ffmpeg -y -r 25 -f image2 -s 1920x1080 -i figs/frame_%06d.png -vcodec libx264 -crf 25  -pix_fmt yuv420p movie.mp4`)


"""
    easter_rata(y::Year) → Int

Returns Easter date as a *[Rata Die](https://en.wikipedia.org/wiki/Rata_Die)* number.

Based on *[Algo R](http://www.linuxtopia.org/online_books/programming_books/python_programming/python_ch38.html)*.
"""
function easter_rata(y::Year)

    local c::Int
    local e::Int
    local p::Int

     # Algo R only works after 1582
     if y.value < 1582
          # Are you using this? Send me a postcard!
          error("Year cannot be less than 1582. Provided: $(y.value).")
     end

    # Century
     c = div(y.value , 100) + 1

     # Shifted Epact
     e = mod(14 + 11*(mod(y.value, 19)) - div(3*c, 4) + div(5+8*c, 25), 30)

     # Adjust Epact
     if (e == 0) || ((e == 1) && ( 10 < mod(y.value, 19) ))
        e += 1
     end

     # Paschal Moon
     p = Date(y.value, 4, 19).instant.periods.value - e

     # Easter: locate the Sunday after the Paschal Moon
     return p + 7 - mod(p, 7)
end

"""
    easter_date(y::Year) → Date

Returns result of `easter_rata` as a `Date` instance.
"""
function easter_date(y::Year)
    # Compute the gregorian date for Rata Die number
     return Date(Dates.rata2datetime( easter_rata(y) ))
end

# weekday_target values:
# const Monday,Tuesday,Wednesday,Thursday,Friday,Saturday,Sunday = 1,2,3,4,5,6,7
# See query.jl on Dates module
# See also dayofweek(dt) function.
"""
    findweekday(weekday_target::Int, yy::Int, mm::Int, occurrence::Int, ascending::Bool) → Date

Given a year `yy` and month `mm`, finds a date where a choosen weekday occurs.

`weekday_target` values are declared in module `Base.Dates`:
`Monday,Tuesday,Wednesday,Thursday,Friday,Saturday,Sunday = 1,2,3,4,5,6,7`.

If `ascending` is true, searches from the beggining of the month. If false, searches from the end of the month.

If `occurrence` is `2` and `weekday_target` is `Monday`, searches the 2nd Monday of the given month, and so on.
"""
function findweekday(weekday_target::Int, yy::Int, mm::Int, occurrence::Int, ascending::Bool)

    local dt::Date = Date(yy, mm, 1)
    local dt_dayofweek::Integer
    local offset::Integer

    if occurrence <= 0
        error("occurrence must be >= 1. Provided $(occurrence).")
    end

    if ascending
        dt_dayofweek = dayofweek(dt)
        offset = rem(weekday_target + 7 - dt_dayofweek, 7) # rem = MOD function
    else
        dt = lastdayofmonth(dt)
        dt_dayofweek = dayofweek(dt)
        offset = rem(dt_dayofweek + 7 - weekday_target, 7)
    end

    if occurrence > 1
        offset += 7 * (occurrence - 1)
    end

    if ascending
        return dt + Dates.Day(offset)
    else
        return dt - Dates.Day(offset)
    end
end


module Vikor

import GreyDecision.GreyNumbers.GreyNumber
import GreyDecision.Utility: unitize, apply_columns, weightize, reverseminmax

function vikor(decisionMat::Array{GreyNumber{T}, 2}, weights::Array{Float64,1}, fns::Array{Function,1};  v::Float64=0.5) where {T <: Real}

    w = unitize(weights)

    nalternatives, ncriteria = size(decisionMat)

    col_max = apply_columns(fns, decisionMat)
    col_min = apply_columns(reverseminmax(fns), decisionMat)

    A = similar(decisionMat)

    for i in 1:nalternatives
        for j in 1:ncriteria
            if fns[j] == maximum
                 A[i, j] = abs((col_max[j] - decisionMat[i, j]) / (col_max[j] - col_min[j]))
            elseif fns[j] == minimum
                 A[i, j] = abs((decisionMat[i, j] - col_min[j]) / (col_max[j] - col_min[j]))
            else
                @warn fns[j]
                error("Function must be either maximum or minimum.")
            end
        end
    end

    weightedA = weightize(A, w)

    s = Array{GreyNumber{Float64},1}(undef, nalternatives)
    r = similar(s)
    q = similar(s)

    for i in 1:nalternatives
        s[i] = sum(weightedA[i,:])
        r[i] = maximum(weightedA[i,:])
    end

    smin = minimum(s)
    smax = maximum(s)
    rmin = minimum(r)
    rmax = maximum(r)
    ranges = (smax - smin)
    ranger = (rmax - rmin)
    for i in 1:nalternatives
        #q = v .* ((s .- smin ./ (smax - smin))) .+ (1 - v) .* ((r .- rmin ./ (rmax - rmin)))
        q[i] = v * (s[i] - smin / ranges) + (1.0-v) * (r[i] - rmin / ranger)
    end

    best_index = sortperm(q) |> first

    return Dict(
        :scores => q,
        :bestindex => best_index
    )
end # end of Vikor function

end # End of module Vikor


using MultiscaleGraphSignalTransforms, JLD, Plots, LightGraphs, MultivariateStats
using LaTeXStrings
push!(LOAD_PATH, @__DIR__)
using pSGWT

G = loadgraph("../datasets/simple_tree_graph.lgz")
X = load("../datasets/simple_tree_xy.jld", "xy")
N = nv(G)
L = Matrix(laplacian_matrix(G))
𝛌, 𝚽 = eigen(L)
standardize_eigenvectors!(𝚽)
W = 1.0 * adjacency_matrix(G)
Q = incidence_matrix(G; oriented = true)

ib1 = 36:56
ib2 = 21:35
ib3 = 71:100
ib4 = 57:70
ijc = [3,5,12,16]
ir = setdiff(1:N, ib1, ib2, ib3, ib4, ijc)

##
P = 𝚽.^2
# P = exp.(𝚽) ./ sum(exp.(𝚽), dims = 1)

dist_sROT, Ws, Xs, 𝚯 = eigsROT_Distance(P, W, X; α = 1.0)

##
function find_active_eigenvectors(P, interest_locs; threshold = 0.5)
    N = size(P, 1)
    energy = zeros(N)
    for k=1:N
        energy[k] = norm(P[interest_locs, k], 1) / norm(P[:, k], 1)
    end
    ind = findall(energy .> threshold)
    return ind
end

# index of eigenvectors active at branch k (k = 1,2,3,4)
ieb1 = find_active_eigenvectors(𝚽.^2, ib1)
ieb2 = find_active_eigenvectors(𝚽.^2, ib2)
ieb3 = find_active_eigenvectors(𝚽.^2, ib3)
ieb4 = find_active_eigenvectors(𝚽.^2, ib4)
iejc = find_active_eigenvectors(𝚽.^2, ijc; threshold = 0.1)

function simpletree_mds_plot(E, ieb1, ieb2, ieb3, ieb4, iejc)
    scatter_gplot(E'; c = :grey, ms = 2)
    scatter_gplot!(E[:, ieb1]'; c = :pink, ms = 2)
    scatter_gplot!(E[:, ieb2]'; c = :orange, ms = 2)
    scatter_gplot!(E[:, ieb3]'; c = :green, ms = 2)
    scatter_gplot!(E[:, ieb4]'; c = :yellow, ms = 2)
    scatter_gplot!(E[:, iejc]'; c = :red, ms = 2)
    scatter_gplot!(E[:, 1:1]'; c = :magenta, ms = 4)
    plt = plot!(xaxis = "X₁", yaxis = "X₂", zaxis = "X₃", legend = false,
                cbar = false, grid = true)
    return plt
end

##
f = zeros(N)
ib1 = 36:56; ib2 = 21:35; ib3 = 71:100; ib4 = 57:70; ijc = [3,5,12,16]
f[ib1] = sin.(0.3 * (0:(length(ib1) - 1)))
f[ib2] = cos.(0.4 * (0:(length(ib2) - 1)))
f[ib3] = sin.(0.5 * (0:(length(ib3) - 1)))
f[ib4] = cos.(0.6 * (0:(length(ib4) - 1)))
f[ijc] .= 1;


include("utils.jl")

@testset "Circle" begin
    M = Circle()
    @testset "Real Circle Basics" begin
        @test repr(M) == "Circle(ℝ)"
        @test representation_size(M) == ()
        @test manifold_dimension(M) == 1
        @test !is_manifold_point(M, 9.)
        @test_throws DomainError is_manifold_point(M, 9., true)
        @test !is_tangent_vector(M, 9., 0.)
        @test_throws DomainError is_tangent_vector(M, 9., 0., true)
        @test is_tangent_vector(M, 0., 0.; check_base_point=false)
        @test get_coordinates(M, Ref(0.0), Ref(2.0), DefaultOrthonormalBasis())[] ≈ 2.0
        @test get_coordinates(M, Ref(0.0), Ref(2.0), DiagonalizingOrthonormalBasis(Ref(1.0)))[] ≈ 2.0
        @test get_coordinates(M, Ref(0.0), Ref(-2.0), DiagonalizingOrthonormalBasis(Ref(1.0)))[] ≈ -2.0
        @test get_coordinates(M, Ref(0.0), Ref(2.0), DiagonalizingOrthonormalBasis(Ref(-1.0)))[] ≈ -2.0
        @test get_coordinates(M, Ref(0.0), Ref(-2.0), DiagonalizingOrthonormalBasis(Ref(-1.0)))[] ≈ 2.0
        y = [0.0]
        get_coordinates!(M, y, Ref(0.0), Ref(2.0), DiagonalizingOrthonormalBasis(Ref(1.0)))
        @test y ≈ [2.0]
        @test get_vector(M, Ref(0.0), Ref(2.0), DefaultOrthonormalBasis())[] ≈ 2.0
        @test get_vector(M, Ref(0.0), Ref(2.0), DiagonalizingOrthonormalBasis(Ref(1.0)))[] ≈ 2.0
        @test get_vector(M, Ref(0.0), Ref(-2.0), DiagonalizingOrthonormalBasis(Ref(1.0)))[] ≈ -2.0
        @test get_vector(M, Ref(0.0), Ref(2.0), DiagonalizingOrthonormalBasis(Ref(-1.0)))[] ≈ -2.0
        @test get_vector(M, Ref(0.0), Ref(-2.0), DiagonalizingOrthonormalBasis(Ref(-1.0)))[] ≈ 2.0
        @test flat(M,0.0, FVector(TangentSpace,1.0)) == FVector(CotangentSpace,1.0)
        @test sharp(M,0.0, FVector(CotangentSpace,1.0)) == FVector(TangentSpace,1.0)
        @test vector_transport_to(M,0.0,1.0,1.0, ParallelTransport()) == 1.0
        @test retract(M,0.0,1.0) == exp(M,0.0,1.0)
        @test injectivity_radius(M) ≈ π
        @test injectivity_radius(M, Ref(-2.0)) ≈ π
        @test injectivity_radius(M, Ref(-2.0), ExponentialRetraction()) ≈ π
        @test injectivity_radius(M, ExponentialRetraction()) ≈ π
        @test mean(M, [-π/2,0.,π]) ≈ π/2
        @test mean(M, [-π/2,0.,π], [1., 1., 1.]) == π/2
        z = project(M, 1.5*π)
        z2 = [0.0]
        project!(M,z2,1.5*π)
        @test z2[1]==z
        @test project(M,z) == z
        @test project(M, 1.0, 2.0) == 2.0
    end
    TEST_STATIC_SIZED && @testset "Real Circle and static sized arrays" begin
        v = MVector(0.0)
        x = SVector(0.0)
        log!(M, v, x, SVector(π/4))
        @test norm(M, x, v) ≈ π/4
        @test is_tangent_vector(M, x, v)
        @test is_tangent_vector(M, [], v; check_base_point = false)
        @test project(M, 1.0) == 1.0
        x = MVector(0.0)
        project!(M, x, x)
        @test x == MVector(0.0)
        x .+= 2*π
        project!(M, x, x)
        @test x == MVector(0.0)
        @test project(M, 0.0, 1.) == 1.
    end
    types = [
        Float64,
    ]
    TEST_FLOAT32 && push!(types, Float32)

    basis_types = (DefaultOrthonormalBasis(),)
    basis_types_real = (DefaultOrthonormalBasis(),
        DiagonalizingOrthonormalBasis(Ref(-1.0)),
        DiagonalizingOrthonormalBasis(Ref(1.0))
    )
    for T in types
        @testset "Type $T" begin
            pts = convert.(Ref(T), [-π/4,0.,π/4])
            test_manifold(
                M,
                pts,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vector_spaces = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true,
                test_vee_hat = false,
                is_mutating = false
            )
            ptsS = SVector.(pts)
            test_manifold(
                M,
                ptsS,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true,
                test_vee_hat = true,
                basis_types_vecs = basis_types_real,
                basis_types_to_from = basis_types_real
            )
        end
    end
    Mc = Circle(ℂ)
    @testset "Complex Circle Basics" begin
        @test repr(Mc) == "Circle(ℂ)"
        @test representation_size(Mc) == ()
        @test manifold_dimension(Mc) == 1
        @test is_tangent_vector(Mc, 1im, 0.)
        @test is_manifold_point(Mc, 1im)
        @test !is_manifold_point(Mc, 1+1im)
        @test_throws DomainError is_manifold_point(Mc, 1+1im, true)
        @test !is_tangent_vector(Mc, 1+1im, 0.)
        @test_throws DomainError is_tangent_vector(Mc, 1+1im, 0., true)
        @test !is_tangent_vector(Mc, 1im, 2im)
        @test_throws DomainError is_tangent_vector(Mc, 1im, 2im, true)
        @test flat(Mc,0.0+0.0im, FVector(TangentSpace,1.0im)) == FVector(CotangentSpace,1.0im)
        @test sharp(Mc,0.0+0.0im, FVector(CotangentSpace,1.0im)) == FVector(TangentSpace,1.0im)
        @test norm(Mc,1.0,log(Mc, 1.0,-1.0)) ≈ π
        @test is_tangent_vector(Mc, 1.0, log(Mc, 1.0, -1.0))
        v = MVector(0.0+0.0im)
        x = SVector(1.0+0.0im)
        log!(Mc, v, x, SVector(-1.0+0.0im))
        @test norm(Mc, SVector(1.0), v) ≈ π
        @test is_tangent_vector(Mc, x, v)
        @test project(Mc, 1.0) == 1.0
        project(Mc, 1/sqrt(2.0) + 1/sqrt(2.0) * im) == 1/sqrt(2.0) + 1/sqrt(2.0) * im
        x = MVector(1.0+0.0im)
        project!(Mc, x, x)
        @test x == MVector(1.0+0.0im)
        x .*= 2
        project!(Mc, x, x)
        @test x == MVector(1.0+0.0im)

    end
    types = [Complex{Float64}, ]
    TEST_FLOAT32 && push!(types, Complex{Float32})

    for T in types
        @testset "Type $T" begin
            a = 1/sqrt(2.0)
            pts = convert.(Ref(T), [a-a*im, 1+0im, a+a*im])
            test_manifold(
                Mc,
                pts,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vector_spaces = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true,
                is_mutating = false,
                test_vee_hat = false,
                exp_log_atol_multiplier = 2.0,
                is_tangent_atol_multiplier = 2.0
            )
            ptsS = SVector.(pts)
            test_manifold(
                Mc,
                ptsS,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true,
                test_vee_hat = true,
                exp_log_atol_multiplier = 2.0,
                is_tangent_atol_multiplier = 2.0,
                basis_types_vecs = basis_types,
                basis_types_to_from = basis_types
            )
        end
    end
end


using Base.LinAlg.BLAS

function free_energy(rbm::RBM, vis::Mat{Float64})
    vb = sum(vis .* rbm.vbias, 1)
    Wx_b_log = sum(log(1 + exp(rbm.W * vis .+ rbm.hbias)), 1)
    return - vb - Wx_b_log
end

function score_samples(rbm::RBM, vis::Mat{Float64}; sample_size=10000)
    if issparse(vis)
        # sparse matrices may be infeasible for this operation
        # so using only little sample
        cols = sample(1:size(vis, 2), sample_size)
        vis = full(vis[:, cols])
    end
    n_feat, n_samples = size(vis)
    vis_corrupted = copy(vis)
    idxs = rand(1:n_feat, n_samples)
    for (i, j) in zip(idxs, 1:n_samples)
        vis_corrupted[i, j] = 1 - vis_corrupted[i, j]
    end
    fe = free_energy(rbm, vis)
    fe_corrupted = free_energy(rbm, vis_corrupted)
    return n_feat * log(logsig(fe_corrupted - fe))
end

function recon_error(rbm::RBM, vis::Mat{Float64})
    # Fully forward MF operation to get back to visible samples
    vis_rec = ProbVisCondOnHid(rbm,ProbHidCondOnVis(rbm,vis))
    # Get the total error over the whole tested visible set,
    # here, as MSE
    dif = vis_rec - vis
    mse = mean(dif.*dif)
    return mse
end

function score_samples_TAP(rbm::RBM, vis::Mat{Float64}; n_iter=5)
    _, _, m_vis, m_hid = iter_mag(rbm, vis; n_times=n_iter, approx="tap2")
    eps=1e-6
    m_vis = max(m_vis, eps)
    m_vis = min(m_vis, 1.0-eps)
    m_hid = max(m_hid, eps)
    m_hid = min(m_hid, 1.0-eps)

    m_vis2 = abs2(m_vis)
    m_hid2 = abs2(m_hid)

    S = - sum(m_vis.*log(m_vis)+(1.0-m_vis).*log(1.0-m_vis),1) - sum(m_hid.*log(m_hid)+(1.0-m_hid).*log(1.0-m_hid),1)
    U_naive = - gemv('T',m_vis,rbm.vbias)' - gemv('T',m_hid,rbm.hbias)' - sum(gemm('N','N',rbm.W,m_vis).*m_hid,1)
    Onsager = - 0.5 * sum(gemm('N','N',rbm.W2,m_vis-m_vis2).*(m_hid-m_hid2),1)
    fe_tap = U_naive + Onsager - S
    fe = free_energy(rbm, vis)
    return fe_tap - fe
end


## cauchy

function cauchycircleS(cfs::AbstractVector,z::Number,s::Bool)
    ret=zero(Complex{Float64})

    if s
        zm = one(Complex{Float64})

        #odd coefficients are pos
        @simd for k=1:2:length(cfs)
            @inbounds ret += cfs[k]*zm
            zm *= z
        end
    else
        z=1/z
        zm = z

        #even coefficients are neg
        @simd for k=2:2:length(cfs)
            @inbounds ret -= cfs[k]*zm
            zm *= z
        end
    end

    ret
end

function stieltjes(sp::Laurent{DD,RR},f::AbstractVector,z::Number) where {DD<:Circle,RR}
    d=domain(sp)
    if !d.orientation
        return -stieltjes(reverseorientation(Fun(sp,f)),z)
    end

    z=mappoint(d,Circle(),z)
    -2π*im*cauchycircleS(f,z,abs(z) < 1)
end

function stieltjes(sp::Laurent{DD,RR},f::AbstractVector,z::Directed{s}) where {DD<:Circle,RR,s}
    d=domain(sp)
    if !d.orientation
        return -stieltjes(reverseorientation(Fun(sp,f)),reverseorientation(z))
    end

    z=mappoint(d,Circle(),z)
    -2π*im*cauchycircleS(f,undirected(z),orientation(z))
end

stieltjes(sp::Fourier{DD,RR}, f::AbstractVector, z::Number) where {DD<:Circle,RR} =
    stieltjes(Laurent(domain(sp)),coefficients(f,sp,Laurent(domain(sp))),z)

# we implement cauchy ±1 as canonical
# TODO: reimplement directly
hilbert(sp::Laurent{DD,RR}, f::AbstractVector, z::Number) where {DD<:Circle,RR} =
    (stieltjes(sp,f,(z)⁺)+stieltjes(sp,f,(z)⁻))/(-2π)

## stieltjesintegral and logkernel

function stieltjesintegral(sp::Laurent{DD,RR},f::AbstractVector,z::Number) where {DD<:Circle,RR}
    d=domain(sp)
    @assert d==Circle()  #TODO: radius
    ζ=Fun(d)
    r=stieltjes(integrate(f-f[2]/ζ),z)
    abs(z)<1 ? r : r+2π*im*f[2]*log(z)
end

stieltjesintegral(sp::Fourier{DD,RR},f::AbstractVector,z::Number) where {DD<:Circle,RR}=stieltjesintegral(Fun(Fun(sp,f),Laurent),z)

function logkernel(sp::Fourier{DD,RR},g::AbstractVector,z::Number) where {DD<:Circle,RR}
    d=domain(sp)
    c,r=d.center,d.radius
    z=z-c
    if abs(z) ≤r
        ret=2r*log(r)*g[1]
        for j=2:2:length(g)
            k=div(j,2)
            ret+=-g[j]*sin(k*angle(z))*abs(z)^k/(k*r^(k-1))
        end
        for j=3:2:length(g)
            k=div(j,2)
            ret+=-g[j]*cos(k*angle(z))*abs(z)^k/(k*r^(k-1))
        end
        ret
    else
        ret=2r*logabs(z)*g[1]
        for j=2:2:length(g)
            k=div(j,2)
            ret+=-g[j]*sin(k*angle(z))*r^(k+1)/(k*abs(z)^k)
        end
        for j=3:2:length(g)
            k=div(j,2)
            ret+=-g[j]*cos(k*angle(z))*r^(k+1)/(k*abs(z)^k)
        end
        ret
    end
end

logkernel(sp::Laurent{DD,RR},g::AbstractVector,z) where {DD<:Circle,RR} = logkernel(Fun(Fun(sp,g),Fourier),z)


using AcFDTD
using Base.Test

srand(123)
geo = CuboidRoom(7,8,10,abs(10*randn(6)), FDTDEnv(0.1,IWB()))
f = FDTD(geo)
Nxyz = geo.Nx*geo.Ny*geo.Nz

Nt = 40
xs = [(6, 7, 9), (3, 3, 3), (2, 5, 2)]
xr = [(2, 2, 4), (2, 7, 8)]

B = get_B(geo,Nt)

Exp = get_Exp(geo,Nt,xs)
Sel = get_Sel(geo,Nt,xr)

s = randn(Nt,length(xs))
p = randn(Nt,length(xr))

S = randn(size(geo)...,Nt)
P = randn(size(geo)...,Nt)

println("testing Adjoint with receiver no IC")
SBE = Sel*(B\full(Exp))
p1 =  SBE*s[:]
s1 =  SBE'*p[:]

p2 = fdtd(xr,xs,Nt,f,s)
s2 = fdtdAdj(xr,xs,Nt,f,p)

println("testing calls")
@time s2 = fdtdAdj(xr,xs,Nt, f,p)

s3 = 0.*copy(s2)
p3 = 0.*copy(p2)

fdtd!(p3,xr,xs,Nt,f,s)
fdtdAdj!(s3,xr,xs,Nt,f,p)
fdtd!(p3,xr,xs,Nt,f,s)
println("testing in-place call")
@time fdtdAdj!(s3,xr,xs,Nt,f,p)

@test abs(vecdot(p[:],p1)-vecdot(s1,s[:]))<1e-10
@test abs(vecdot(p,p2)-vecdot(s2,s))<1e-10
@test abs(vecdot(p,p3)-vecdot(s3,s))<1e-10

println("testing Adjoint full no IC")
BE = (B\full(Exp))
P1 =  (BE*s[:])[2*prod(size(geo))+1:end]
s1 =  BE'*[zeros(2*prod(size(geo)));P[:]]

P2 = fdtd(xs,Nt,f,s)
s2 = fdtdAdj(xs,Nt,f,P)

println("testing calls")
@time s2 = fdtdAdj(xs,Nt,f,P)

s3 = 0.*copy(s2)
P3 = 0.*copy(P2)

fdtd!(P3,xs,Nt,f,s)
fdtdAdj!(s3,xs,Nt,f,P)
println("testing in-place call")
@time fdtdAdj!(s3,xs,Nt,f,P)

@test abs(vecdot(P[:],P1)-vecdot(s1,s[:]))<1e-10
@test abs(vecdot(P,P2)-vecdot(s2,s))<1e-10
@test abs(vecdot(P,P3)-vecdot(s3,s))<1e-10

println("testing Adjoint full, source everywhere no IC")

P2 = fdtd(Nt,f,S)
S2 = fdtdAdj(Nt,f,P)

println("testing calls")
@time S2 = fdtdAdj(Nt,f,P)

S3 = 0.*copy(S2)
P3 = 0.*copy(P2)

fdtd!(P3,Nt,f,S)
fdtdAdj!(S3,Nt,f,P)
println("testing in-place call")
@time fdtdAdj!(S3,Nt,f,P)

#@test abs(vecdot(P[:],P1)-vecdot(S1,S[:]))<1e-12
@test abs(vecdot(P,P2)-vecdot(S2,S))<1e-10
@test abs(vecdot(P,P3)-vecdot(S3,S))<1e-10

##TODO make adjoint with ICs


export PermutationMatrix

import Base: *
import LinearAlgebra: transpose , inv, mul!

@doc raw"""
`PermutationMatrix(ind::Array{T})` creates a permutation matrix with ones at coordinates `(k, ind[k]) for k = 1:length(ind)`.

A permutation matrix represents a matrix containing only zeros and ones, which basically permutes the vector or matrix it is multiplied with.
These matrices A are constrained by:
```math
    A_{ij} = \{0, 1\}\\
    ∑_{i} A_{ij} = 1\\
    ∑_{j} A_{ij} = 1
```
An example is the 3-dimensional permutation matrix
```math
A = \begin{bmatrix} 1 & 0 & 0 \\ 0 & 0 & 1 \\ 1 & 0 & 0 \end{bmatrix}
```
"""
struct PermutationMatrix{ T <: Integer } <: AbstractMatrix{T}
    ind::Vector{T}
    function PermutationMatrix(ind::Vector{T}) where { T <: Integer }
        return new{T}(ind)
    end
end

@doc raw"""
`PermutationMatrix(dim::T; switch_first::Bool=true)` Generates a random permutation matrix of size (dim x dim).
The `switch_first` argument specifies whether the first index always has to be permuted.
"""
function PermutationMatrix(dim::T; switch_first::Bool=true) where { T <: Integer }
    ind = shuffle(collect(1:dim))
    if switch_first && ind[1] == 1
        tmp = ind[2]
        ind[2] = ind[1]
        ind[1] = tmp
    end
    return PermutationMatrix(ind)
end

# extensions of base functionality
Base.eltype(::PermutationMatrix{T}) where { T } = T
function Base.size(mat::PermutationMatrix)
    nr_elements = length(mat.ind)
    return (nr_elements, nr_elements)
end
Base.size(mat::PermutationMatrix, d) = d::Integer <= 2 ? length(mat.ind) : 1
Base.length(mat::PermutationMatrix)  = prod(size(mat))

function Base.getindex(mat::PermutationMatrix, i::Int, j::Int)
    if mat.ind[i] == j
        return one(eltype(mat))
    else
        return zero(eltype(mat))
    end
end

LinearAlgebra.inv(mat::PermutationMatrix)       = adjoint(mat)

# get functions
getind(mat::PermutationMatrix)      = mat.ind
getind(mat::Adjoint{T, PermutationMatrix{T}}) where { T }       = sortperm(mat.parent.ind)

# Permutation-vector multiplication
function Base.:*(P::PermutationMatrix, v::AbstractVector)
    y = similar(v)
    mul!(y, P, v)
    return y
end

function LinearAlgebra.mul!(y::AbstractVector, P::PermutationMatrix, v::AbstractVector)
    ind = getind(P)
    @inbounds @simd for k in 1:size(P,1)
        y[k] = v[ind[k]]
    end
end

function Base.:*(P::Adjoint{T, PermutationMatrix{T}}, v::AbstractVector) where { T }
    y = similar(v)
    mul!(y, P, v)
    return y
end

function LinearAlgebra.mul!(y::AbstractVector, P::Adjoint{T, PermutationMatrix{T}}, v::AbstractVector) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        y[ind[k]] = v[k]
    end
end

function Base.:*(P::Transpose{T, PermutationMatrix{T}}, v::AbstractVector) where { T }
    y = similar(v)
    mul!(y, P, v)
    return y
end

function LinearAlgebra.mul!(y::AbstractVector, P::Transpose{T, PermutationMatrix{T}}, v::AbstractVector) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        y[ind[k]] = v[k]
    end
end

# Permutation-matrix multiplication
function Base.:*(P::PermutationMatrix, X::AbstractMatrix)
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, P, X)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, P::PermutationMatrix{T}, X::AbstractMatrix) where { T }
    ind = getind(P)
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[k,ki] = X[ind[k],ki]
        end
    end
end

function Base.:*(P::Adjoint{T, PermutationMatrix{T}}, X::AbstractMatrix) where { T }
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, P, X)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, P::Adjoint{T, PermutationMatrix{T}}, X::AbstractMatrix) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[ind[k],ki] = X[k,ki]
        end
    end
end

function Base.:*(P::Transpose{T, PermutationMatrix{T}}, X::AbstractMatrix) where { T }
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, P, X)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, P::Transpose{T, PermutationMatrix{T}}, X::AbstractMatrix) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[ind[k],ki] = X[k,ki]
        end
    end
end

# transposedvector - matrix product
Base.:*(x::Adjoint{T, S}, P::PermutationMatrix) where { T, S <: AbstractVector } = (P'*x')'
Base.:*(x::Adjoint{T, S}, P::Adjoint{T1, PermutationMatrix{T1}}) where { T, T1, S <: AbstractVector } = (P'*x')'
Base.:*(x::Adjoint{T, S}, P::Transpose{T1, PermutationMatrix{T1}}) where { T, T1, S <: AbstractVector } = (P'*x')'
Base.:*(x::Transpose{T, S}, P::PermutationMatrix) where { T, S <: AbstractVector } = transpose(transpose(P)*transpose(x))
Base.:*(x::Transpose{T, S}, P::Adjoint{T1, PermutationMatrix{T1}}) where { T, T1, S <: AbstractVector } = transpose(transpose(P)*transpose(x))
Base.:*(x::Transpose{T, S}, P::Transpose{T1, PermutationMatrix{T1}}) where { T, T1, S <: AbstractVector } = transpose(transpose(P)*transpose(x))

# matrix-Permutation multiplication
function Base.:*(X::AbstractMatrix, P::PermutationMatrix)
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, X, P)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, X::AbstractMatrix, P::PermutationMatrix)
    ind = getind(P)
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[ki,ind[k]] = X[ki,k]
        end
    end
end

function Base.:*(X::AbstractMatrix, P::Adjoint{T, PermutationMatrix{T}}) where { T }
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, X, P)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, X::AbstractMatrix, P::Adjoint{T, PermutationMatrix{T}}) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[ki,k] = X[ki,ind[k]]
        end
    end
end

function Base.:*(X::AbstractMatrix, P::Transpose{T, PermutationMatrix{T}}, ) where { T }
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, X, P)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, X::AbstractMatrix, P::Transpose{T, PermutationMatrix{T}}, ) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[ki,k] = X[ki,ind[k]]
        end
    end
end

# multiplication of some square matrix with the permutation matrix
function PT_X_P(X::AbstractMatrix, P::PermutationMatrix)

    # allocate output
    Y = copy(X)

    # perform permutation operation
    PT_X_P!(Y, X, P)

    # return output
    return Y

end

function PT_X_P!(Y::AbstractMatrix, X::AbstractMatrix, P::PermutationMatrix)

    # fetch indices of permutation matrix
    ind = getind(P)

    # perform permutation operation
    @inbounds @simd for k1 = 1:size(P,1)
        @inbounds @simd for k2 = 1:size(P,1)
            Y[ind[k1], ind[k2]] = X[k1,k2]
        end
    end

end


@testset "Testing lqcontrol.jl" begin

    rough_kwargs = Dict(:atol => 1e-13, :rtol => 1e-4)

    # set up
    q    = 1.
    r    = 1.
    rf   = 1.
    a    = .95
    b    = -1.
    c    = .05
    β    = .95
    n    = 0.
    capT = 1
    lq_scalar = LQ(q, r, a, b, c, n, β, capT, rf)

    Q  = [0. 0.; 0. 1]
    R  = [1. 0.; 0. 0]
    rf = eye(2) .* 100
    A  = fill(0.95, 2, 2)
    B  = fill(-1.0, 2, 2)
    lq_mat = LQ(Q, R, A, B, bet=β, capT=capT, rf=rf)

    @testset "Test scalar sequences with exact by hand solution" begin
        x0 = 2.0
        x_seq, u_seq, w_seq = compute_sequence(lq_scalar, x0)
        # solve by hand
        u_0 = (-2.*lq_scalar.A*lq_scalar.B*lq_scalar.bet*lq_scalar.rf) /
           (2.*lq_scalar.Q+lq_scalar.bet*lq_scalar.rf*2lq_scalar.B^2)*x0
        x_1 = lq_scalar.A * x0 + lq_scalar.B * u_0 + w_seq[end]

        @test isapprox(u_0[1], u_seq[end]; rough_kwargs...)
        @test isapprox(x_1[1], x_seq[end]; rough_kwargs...)
    end

    @testset "test matrix solutions" begin
        x0 = randn(2) * 25
        x_seq, u_seq, w_seq = compute_sequence(lq_mat, x0)

        @test isapprox(sum(u_seq), 0.95 * sum(x0); rough_kwargs...)
        @test isapprox(x_seq[:,end], zeros(x0); rough_kwargs...)
    end

    @testset "test stationary matrix" begin
        x0 = randn(2) .* 25
        P, F, d = stationary_values(lq_mat)

        f_answer = [-.95 -.95; 0. 0.]
        p_answer = [1. 0; 0. 0.]

        val_func_lq = (x0' * P * x0)[1]
        val_func_answer = x0[1]^2

        @test isapprox(f_answer, F; rough_kwargs...)
        @test isapprox(val_func_lq, val_func_answer; rough_kwargs...)
        @test isapprox(p_answer, P; rough_kwargs...)
    end

    @testset "test runs a (n,k,j) = (2,1,1) model" begin
        # == Model parameters == #
        r = 0.05
        bet = 1 / (1 + r)
        t = 45
        c_bar = 2.0
        sigma = 0.25
        mu = 1.0
        q = 1e6

        # == Formulate as an LQ problem == #
        Q = 1.0
        R = zeros(2, 2)
        Rf = zeros(2, 2); Rf[1, 1] = q
        A = [1.0+r -c_bar+mu;
             0.0 1.0]
        B = [-1.0, 0.0]
        C = [sigma, 0.0]

        # == Compute solutions and simulate == #
        lq = LQ(Q, R, A, B, C; bet=bet, capT=t, rf=Rf)
        x0 = [0.0, 1.0]
        xp, up, wp = compute_sequence(lq, x0)
        @test true == true  # just assert true if we made it to this point
    end

end  # @testset


@testset "test functions var N" begin

    function f1(x)
        return exp(-0.2*sin(3*x))
    end
    function f2(x)
        return exp(-0.3*cos(3*x)) + 20
    end

    c1 = ChebyshevQuantum.Chebyshev.make_Cheb(f1 , a=2, b = 3)
    c2 = ChebyshevQuantum.Chebyshev.make_Cheb(f2 , a=2, b = 3)

    x=2.712
    @test isapprox( (c1+c2)(x), f1(x)+f2(x) , atol=tol_val)
    @test isapprox( (c1-c2)(x), f1(x)-f2(x) , atol=tol_val)
    @test isapprox( (c1^2)(x), f1(x)^2 , atol=tol_val)
    @test isapprox( (c1/c2)(x), f1(x)/f2(x) , atol=tol_val)

    @test isapprox( (-c1)(x), -f1(x) , atol=tol_val)
    @test isapprox( (2*c1)(x), 2*f1(x) , atol=tol_val)
    @test isapprox( (c1*2)(x), 2*f1(x) , atol=tol_val)
    @test isapprox( (2+c1)(x), 2+f1(x) , atol=tol_val)
    @test isapprox( (c1+2)(x), f1(x)+2 , atol=tol_val)

    @test isapprox( (2 / c2)(x), 2 / f2(x) , atol=tol_val)
    @test isapprox( (2 - c2)(x), 2 - f2(x) , atol=tol_val)

end

@testset "test Cheb functions fixed N" begin

    function f1(x)
        return exp(-0.2*sin(3*x))
    end
    function f2(x)
        return exp(-0.3*cos(3*x)) + 20
    end

    c1 = ChebyshevQuantum.Chebyshev.make_Cheb(f1 , N = 20,thr=-1,   a=2, b = 3, edgesearch=false)
    c2 = ChebyshevQuantum.Chebyshev.make_Cheb(f2 , N = 20,thr=-1,  a=2, b = 3, edgesearch=false)

    x=2.712
    @test isapprox( (c1+c2)(x), f1(x)+f2(x) , atol=tol_val)
    @test isapprox( (c1-c2)(x), f1(x)-f2(x) , atol=tol_val)
    @test isapprox( (c1^2)(x), f1(x)^2 , atol=tol_val)
    @test isapprox( (c1/c2)(x), f1(x)/f2(x) , atol=tol_val)

    @test isapprox( (-c1)(x), -f1(x) , atol=tol_val)
    @test isapprox( (2*c1)(x), 2*f1(x) , atol=tol_val)
    @test isapprox( (c1*2)(x), 2*f1(x) , atol=tol_val)
    @test isapprox( (2+c1)(x), 2+f1(x) , atol=tol_val)
    @test isapprox( (c1+2)(x), f1(x)+2 , atol=tol_val)

    @test isapprox( (2 / c2)(x), 2 / f2(x) , atol=tol_val)
    @test isapprox( (2 - c2)(x), 2 - f2(x) , atol=tol_val)

end


@testset "Sparse logistic small ($T)" for T in [Float32, Float64]

    using ProximalOperators
    using ProximalAlgorithms
    using LinearAlgebra
    using Random

    Random.seed!(0)

    A = T[  1.0  -2.0   3.0  -4.0  5.0;
           2.0  -1.0   0.0  -1.0  3.0;
          -1.0   0.0   4.0  -3.0  2.0;
          -1.0  -1.0  -1.0   1.0  3.0]
    b = T[1.0, 2.0, 3.0, 4.0]

    m, n = size(A)

    R = real(T)

    f = Translate(LogisticLoss(ones(R, m), R(1)), -b)
    lam = R(0.1)
    g = NormL1(lam)

    x_star = T[0, 0, 2.114635341704963e-01, 0, 2.845881348733116e+00]

    TOL = R(1e-6)

    # Nonfast/Adaptive

    x0 = zeros(T, n)
    solver = ProximalAlgorithms.ForwardBackward{R}(tol=TOL, adaptive=true, fast=false)
    x, it = solver(x0, f=f, A=A, g=g)
    @test eltype(x) == T
    @test norm(x - x_star, Inf) <= 1e-4
    @test it < 1100

    # Fast/Adaptive

    x0 = zeros(T, n)
    solver = ProximalAlgorithms.ForwardBackward{R}(tol=TOL, adaptive=true, fast=true)
    x, it = solver(x0, f=f, A=A, g=g)
    @test eltype(x) == T
    @test norm(x - x_star, Inf) <= 1e-4
    @test it < 500

    # ZeroFPR/Adaptive

    x0 = zeros(T, n)
    solver = ProximalAlgorithms.ZeroFPR{R}(adaptive=true, tol=TOL)
    x, it = solver(x0, f=f, A=A, g=g)
    @test eltype(x) == T
    @test norm(x - x_star, Inf) <= 1e-4
    @test it < 25

    # PANOC/Adaptive

    x0 = zeros(T, n)
    solver = ProximalAlgorithms.PANOC{R}(adaptive=true, tol=TOL)
    x, it = solver(x0, f=f, A=A, g=g)
    @test eltype(x) == T
    @test norm(x - x_star, Inf) <= 1e-4
    @test it < 50

end


# Multivariate normal setup using non-conjugate prior.
module MVN

module MVNmodel
export Theta, Data, log_likelihood, log_prior, prior_sample!, new_theta, Theta_clear!, Theta_adjoin!, Theta_remove!,
       Hyperparameters, construct_hyperparameters, update_hyperparameters!, update_parameter!, mixrnd, mixture_density

include("Lower.jl")
using .Lower

const NU_SIGMA_PROP = 0.1  # Metropolis proposal std dev for H.R.nu updates

const Data = Array{Float64,1}

# Multivariate normal distribution
type MVN_params
    # read/write
    m::Array{Float64,1} # mean
    # read/restricted-write (must notify after writing)
    L::Array{Float64,2} # lower triangular matrix such that L*L' = R
    # no access (use getter to read when necessary)
    _R::Array{Float64,2} # precision matrix (inverse covariance)
    _R_valid::Bool # flag to indicate whether R has been computed or not
    # read-only
    logdetR::Float64 # log of the determinant of the precision matrix
    d::Int64 # dimension
    n::Int64 # number of data points assigned to this cluster
    sum_x::Array{Float64,1} # sum of the data points x assigned to this cluster
    sum_xx::Array{Float64,2} # sum of x*x' for the data points assigned to this cluster
    function MVN_params(m,R)
        p = new(); p.m = copy(m); p._R = copy(R); p._R_valid = true; p.d = d = length(m)
        p.L = zeros(d,d); Lower.Cholesky!(p.L,p._R,d)
        p.logdetR = Lower.logdetsq(p.L,d)
        p.n = 0
        p.sum_x = zeros(d)
        p.sum_xx = zeros(d,d)
        return p
    end
end
# Interface functions for MVN_params
MVN_logpdf(x,p) = -0.5*Lower.quadratic(x, p.m, p.L, p.d) - 0.5*p.d*log(2*pi) + 0.5*p.logdetR
MVN_sample!(x,p) = Lower.sample_Normal!(x, p.m, p.L, p.d)
MVN_get_R!(p) = (if !p._R_valid; Lower.multiplyMNt!(p._R,p.L,p.L,p.d); p._R_valid=true; end; p._R)
MVN_notify_L!(p) = (p.logdetR = Lower.logdetsq(p.L,p.d); p._R_valid = false) # call this after modifying L
MVN_clear!(p) = (fill!(p.sum_x,0.); fill!(p.sum_xx,0.); p.n=0)
MVN_adjoin!(p,x) = (for i=1:p.d; p.sum_x[i]+=x[i]; for j=1:p.d; p.sum_xx[i,j]+=x[i]*x[j]; end; end; p.n+=1)
MVN_remove!(p,x) = (for i=1:p.d; p.sum_x[i]-=x[i]; for j=1:p.d; p.sum_xx[i,j]-=x[i]*x[j]; end; end; p.n-=1)

# Wishart distribution
type Wishart_params
    # read/restricted-write (must notify after writing)
    W::Array{Float64,2} # inverse of the mean matrix, i.e., W = inv(nu*V)
    # read-only
    nu::Float64 # degrees of freedom (can modify using set_nu)
    M::Array{Float64,2} # lower triangular matrix such that M*M' = inv(nu*W) = V = the scale matrix
    c::Float64 # log normalization constant
    d::Int64 # dimension
    n::Int64 # number of observations
    sum_R::Array{Float64,2} # sum of the R's
    sum_logdetR::Float64 # sum of the log(det(R))'s
    A::Array{Float64,2} # temporary variable
    function Wishart_params(nu,W) # W is the inverse of the mean matrix (d-by-d), nu > d-1 is the degrees of freedom
        p = new(); p.nu = nu; p.W = copy(W); p.d = d = size(p.W,1)
        p.M = zeros(d,d); p.A = zeros(d,d)
        Lower.Cholesky_inverse!(p.M,p.A,p.W,d); Lower.scale!(p.M,sqrt(1/p.nu),d)
        p.c = Wishart_constant(p)
        p.n = 0
        p.sum_R = zeros(d,d)
        p.sum_logdetR = 0.
        return p
    end
end
# Helper functions
matrix_dot(A,B,d) = (r=0.; for i=1:d*d; r += A[i]*B[i]; end; r)
Wishart_constant(p) = (d=p.d; nu=p.nu; r=0.; for j=1:d; r += lgamma(0.5*(nu+1-j)); end;
    0.5*nu*Lower.logdetsq(p.M,d) + 0.5*nu*d*log(2) + 0.25*d*(d-1)*log(pi) + r)
# Interface functions
# Sample L (lower triangular) such that R = L*L' ~ Wishart(Scale=M*M', DOF=nu)
Wishart_sample!(L,p) = Lower.sample!(L, p.M, p.nu, p.d)
# Wishart log pdf of R, using precomputed triangular L such that L*L' = R
Wishart_logpdf(R,L,p) = 0.5*(p.nu-p.d-1)*Lower.logdetsq(L,p.d) - 0.5*p.nu*matrix_dot(R,p.W,p.d) - p.c
Wishart_set_nu!(p,nu) = (Lower.scale!(p.M,sqrt(p.nu/nu),p.d); p.nu=nu; p.c=Wishart_constant(p))
Wishart_notify_W!(p) = (Lower.Cholesky_inverse!(p.M,p.A,p.W,p.d); Lower.scale!(p.M,sqrt(1/p.nu),p.d);
    p.c=Wishart_constant(p))
Wishart_clear!(p) = (fill!(p.sum_R,0.); p.sum_logdetR=0.; p.n=0)
Wishart_adjoin!(p,R,logdetR) = (for i=1:p.d*p.d; p.sum_R[i]+=R[i]; end; p.sum_logdetR+=logdetR; p.n+=1)
# Wishart_remove!(p) = ...

# # LogGamma distribution (X = log(Y) where Y~Gamma(a,b))
type LogGamma_params
    a::Float64
    b::Float64
    c::Float64
    LogGamma_params(a,b) = (p=new(); p.a=a; p.b=b; p.c=lgamma(a)-a*log(b); p)
end
LogGamma_logpdf(x,p) = p.a*x - p.b*exp(x) - p.c

type Hyperparameters
    m::MVN_params # parameters of m's (component means)
    R::Wishart_params # parameters of R's (component precisions)
    mm::MVN_params # parameters of m.m
    mR::Wishart_params # parameters of m.R
    Rn::LogGamma_params # parameters of R.nu
    RW::Wishart_params # parameters of R.W
    mp::MVN_params # temporary variable for the posterior of m's
    Rp::Wishart_params # temporary variable for the posterior of R's
    A::Array{Float64,2} # temporary variable for computations
    x::Array{Float64,1} # temporary variable for computations
    d::Int64 # dimension
end

# likelihood: Normal(x|mean=m,Cov=inv(R)) (Note: R is represented as L*L'.)
const Theta = MVN_params
Theta_clear!,Theta_adjoin!,Theta_remove! = MVN_clear!,MVN_adjoin!,MVN_remove!
log_likelihood(x,p) = MVN_logpdf(x,p)
# prior: Normal(m|mean=H.m.m,Cov=inv(H.m.L*H.m.L')) Wishart(R|Scale=H.R.M*H.R.M',DOF=H.R.nu)
log_prior(p,H) = MVN_logpdf(p.m,H.m) + Wishart_logpdf(MVN_get_R!(p),p.L,H.R)
new_theta(H) = MVN_params(zeros(H.d),eye(H.d))
prior_sample!(p,H) = (MVN_sample!(p.m,H.m); Wishart_sample!(p.L,H.R); MVN_notify_L!(p))

sum_outer(x,m,d,n) = (S=zeros(d,d); for k = 1:n, i=1:d, j=1:d; S[i,j] += (x[k][i]-m[i])*(x[k][j]-m[j]); end; S)

function construct_hyperparameters(options)
    x = options.x
    # x should be an array of n data points of dimension d, e.g., an array of vectors, with n>=d
    n,d = length(x),length(x[1])
    m_hat = mean(x)
    C_hat = sum_outer(x,m_hat,d,n)/n
    R_hat = inv(cholfact(C_hat))

    mm = MVN_params(m_hat,R_hat) # these are fixed
    mR = Wishart_params(d,C_hat) # these are fixed

    m = MVN_params(m_hat,R_hat) # initialize to typical values (E(m.m), E(m.R))

    RW = Wishart_params(d,R_hat) # these are fixed

    # We take log(R.nu-d+1) ~ LogGamma(a,b). (Equivalently, R.nu-d+1 ~ Gamma(a,b).)
    # (This is different from Rasmussen, who takes 1/(R.nu-d+1) ~ Gamma(a,b).)
    Rn = LogGamma_params(2.,2.) # these are fixed

    R = Wishart_params(d,C_hat) # initialize to typical values
    # This choice of initial R.nu,R.W is derived by setting R.nu=E(R.nu)=d and R.W=E(R.W)=C_hat.

    mp = MVN_params(zeros(d),eye(d))
    Rp = Wishart_params(d,eye(d))

    A = zeros(d,d)
    x0 = zeros(d)

    return Hyperparameters(m,R,mm,mR,Rn,RW,mp,Rp,A,x0,d)
end

function NormalWishart_update!(a,b,m_params,R_params,H,active,density)
    d = H.d
    a_R = MVN_get_R!(a)
    m_params_R = MVN_get_R!(m_params)

    # update mean
    for i = 1:d*d; H.A[i] = m_params_R[i] + b.n * a_R[i]; end
    Lower.Cholesky!(H.mp.L, H.A, d); MVN_notify_L!(H.mp)
    for i = 1:d
        r = 0.; for j = 1:d; r += m_params_R[i,j]*m_params.m[j] + a_R[i,j]*b.sum_x[j]; end
        H.x[i] = r
    end
    Lower.solve_Lx_eq_y!(H.mp.L, H.x, H.x, d)
    Lower.solve_Ltx_eq_y!(H.mp.L, H.mp.m, H.x, d)
    if active; MVN_sample!(b.m, H.mp); end

    # update precision matrix
    Wishart_set_nu!(H.Rp, b.n + R_params.nu)
    for i = 1:d, j = 1:d
        H.Rp.W[i,j] = (R_params.nu*R_params.W[i,j] + b.sum_xx[i,j]
            - b.sum_x[i]*b.m[j] - b.m[i]*b.sum_x[j] + b.n*b.m[i]*b.m[j]) / H.Rp.nu
    end; Wishart_notify_W!(H.Rp)
    if active; Wishart_sample!(b.L, H.Rp); MVN_notify_L!(b); end

    # compute density
    return (density? MVN_logpdf(b.m, H.mp) + Wishart_logpdf(MVN_get_R!(b), b.L, H.Rp) : NaN)
end

function Wishart_update!(p,W_params,nu_params,H)
    d = H.d

    # update W (inverse mean matrix)
    Wishart_set_nu!(H.Rp, W_params.nu + p.n * p.nu)
    for i = 1:d*d
        H.Rp.W[i] = (W_params.nu*W_params.W[i] + p.nu*p.sum_R[i]) / H.Rp.nu
    end; Wishart_notify_W!(H.Rp)
    Wishart_sample!(H.A, H.Rp) # H.A*H.A' ~ Wishart(Scale=H.Rp.M*H.Rp.M', DOF=H.Rp.nu)
    Lower.multiplyMNt!(p.W, H.A, H.A, d); Wishart_notify_W!(p) # p.W = H.A*H.A'

    # update nu (degrees of freedom) using a Metropolis move
    nu_0 = p.nu
    x_0 = log(p.nu-d+1)
    log_prior_0 = LogGamma_logpdf(x_0,nu_params)
    log_lik_0 = 0.5*(p.nu-d-1)*p.sum_logdetR - 0.5*p.nu*matrix_dot(p.W,p.sum_R,d) - p.n*p.c

    # make proposal
    x_prop = x_0 + randn()*NU_SIGMA_PROP
    Wishart_set_nu!(p, exp(x_prop)+d-1) # Note: p.c is updated here
    log_prior_prop = LogGamma_logpdf(x_prop,nu_params)
    log_lik_prop = 0.5*(p.nu-d-1)*p.sum_logdetR - 0.5*p.nu*matrix_dot(p.W,p.sum_R,d) - p.n*p.c

    # accept or reject
    p_accept = min(1.0, exp(log_prior_prop + log_lik_prop - log_prior_0 - log_lik_0))
    if rand()>p_accept
        Wishart_set_nu!(p,nu_0) # if we don't accept, set nu back to it's original value
    end
end

function update_parameter!(theta_a,theta_b,H,active,density)
    return NormalWishart_update!(theta_a,theta_b,H.m,H.R,H,active,density)
end

function update_hyperparameters!(H,theta,list,t,x,z)
    # update the hyperparameters of the mean (m.m and m.R)
    MVN_clear!(H.m)
    for k = 1:t; MVN_adjoin!(H.m, theta[list[k]].m); end
    NormalWishart_update!(H.m,H.m,H.mm,H.mR,H,true,false)

    # update the hyperparameters of the precision (R.W and R.nu)
    Wishart_clear!(H.R)
    for k = 1:t; tk=theta[list[k]]; Wishart_adjoin!(H.R, MVN_get_R!(tk), tk.logdetR); end
    Wishart_update!(H.R,H.RW,H.Rn,H)
end

# sample from a discrete distribution with weights proportional to q = [q_1,...,q_n]
randp(q,n) = (cdf=cumsum(q/sum(q)); [findfirst(rand().<cdf)::Int64 for i = 1:n])

# sample from a mixture
function mixrnd(n,probabilities,theta)
    k = length(probabilities)
    z = randp(probabilities[1:k],n) # component assignments
    d = theta[1].d
    x = [MVN_sample!(zeros(d),theta[z[i]])::Array{Float64,1} for i = 1:n]
    return x,z
end

function mixture_density(x,probabilities,theta)
    k = length(probabilities)
    px = 0.0
    for i = 1:k
        px += probabilities[i]*exp(log_likelihood(x,theta[i]))
    end
    return px
end

end # module MVNmodel
using .MVNmodel

# Include generic code
include("generic.jl")

# Include core sampler code
include("coreNonconjugate.jl")

end # module MVN


include("euler/euler.jl")
using .Math: fast_min

function compute(grid::Array{Int,2})::Int
    edge = size(grid)[1]
    for i ∈ edge - 1:-1:1
        grid[edge, i] += grid[edge, i + 1]
        grid[i, edge] += grid[i + 1, edge]
    end
    for i ∈ edge - 1:-1:1 for j ∈ edge - 1:-1:1 grid[i, j] += fast_min(grid[i + 1, j], grid[i, j + 1]) end end
    return grid[1, 1]
end


@with_kw struct IZParameter
    a::SNNFloat = 0.01
    b::SNNFloat = 0.2
    c::SNNFloat = -65
    d::SNNFloat = 2
end

@with_kw mutable struct IZ{VF,VB}
    param::IZParameter = IZParameter()
    N::SNNInt = 100
    v::VF = fill(-65.0, N)
    u::VF = param.b * v
    fire::VB = zeros(eltype(VB), N)
    I::VF = zeros(eltype(VF), N)
    records::Dict = Dict()
end
IZ(x;kwargs...) = IZ{Vector{SNNFloat},Vector{Bool}}(;kwargs...)

function integrate!(p::IZ, param::IZParameter, dt::SNNFloat)
    @unpack N, v, u, fire, I = p
    @unpack a, b, c, d = param
    @inbounds for i = 1:N
        v[i] += 0.5f0dt * (0.04f0v[i]^2 + 5f0v[i] + 140f0 - u[i] + I[i])
        v[i] += 0.5f0dt * (0.04f0v[i]^2 + 5f0v[i] + 140f0 - u[i] + I[i])
        u[i] += dt * (a * (b * v[i] - u[i]))
    end
    @inbounds for i = 1:N
        fire[i] = v[i] > 30f0
        v[i] = ifelse(fire[i], c, v[i])
        u[i] += ifelse(fire[i], d, 0f0)
    end
end


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    BisectFractionPartitioner(normal, fraction=0.5, maxiter=10)

A method for partitioning spatial data into two half spaces
defined by a `normal` direction and a `fraction` of points.
The partition is returned within `maxiter` bisection iterations.
"""
struct BisectFractionPartitioner{T,N} <: AbstractPartitioner
  normal::SVector{N,T}
  fraction::Float64
  maxiter::Int

  function BisectFractionPartitioner{T,N}(normal, fraction, maxiter) where {N,T}
    new(normalize(normal), fraction, maxiter)
  end
end

BisectFractionPartitioner(normal::SVector{N,T}, fraction=0.5, maxiter=10) where {T,N} =
  BisectFractionPartitioner{T,N}(normal, fraction, maxiter)

BisectFractionPartitioner(normal::NTuple{N,T}, fraction=0.5, maxiter=10) where {T,N} =
  BisectFractionPartitioner(SVector(normal), fraction, maxiter)

function partition(object, partitioner::BisectFractionPartitioner)
  bbox = boundbox(object)
  n = partitioner.normal
  f = partitioner.fraction
  c = center(bbox)
  d = diagonal(bbox)

  # maximum number of bisections
  maxiter = partitioner.maxiter

  iter = 0; p = 0
  a = c - d/2 * n
  b = c + d/2 * n
  while iter < maxiter
    m = (a + b) / 2

    p = partition(object, BisectPointPartitioner(n, m))
    g = nelms(p[1]) / nelms(object)

    g ≈ f && break
    g > f && (b = m)
    g < f && (a = m)

    iter += 1
  end

  p
end