# Bogumił Kamiński, 2019-03-25

# Reworked example from
# https://github.com/FluxML/model-zoo/blob/master/vision/mnist/mlp.jl

using Statistics
using Flux

const TRAIN = (class=UInt8[], image=Vector{Float64}[])
const TEST = (class=UInt8[], image=Vector{Float64}[])

for (filename, data) in [("mnist_train.int", TRAIN),
                         ("mnist_test.int", TEST)]
    open(filename) do f
        while !eof(f)
            c = read(f, UInt8)
            v = reshape(read(f, 28^2), 28, 28)
            push!(data.class, c)
            push!(data.image, vec(v[4:end-3, 4:end-3]) ./ 255)
        end
    end
end

Y = Flux.onehotbatch(TRAIN.class, 0:9);
X = reduce(hcat, TRAIN.image);

m = Chain(Dense(484, 64, relu),
          Dense(64, 64, relu),
          Dense(64, 64, relu),
          Dense(64, 10),
          softmax)

loss(x, y) = Flux.crossentropy(m(x), y)
accuracy(x, y) = mean(Flux.onecold(m(x)) .== Flux.onecold(y))
dataset = Iterators.repeated((X, Y), 200)
evalcb = () -> @show(loss(X, Y))
opt = ADAM()

Flux.train!(loss, params(m), dataset, opt, cb = Flux.throttle(evalcb, 10))
accuracy(X, Y)

tY = Flux.onehotbatch(TEST.class, 0:9);
tX = reduce(hcat, TEST.image);
accuracy(tX, tY)


@testset "laplacian" begin
    h = 1.0
    rho = 1000.0
    G = 9.8
    len_z = 16
    len_x = 32
    nz = Int(len_z/h + 1)
    nx = Int(len_x/h + 1)
    tf_h=constant(1.0)
    coef = rand(nz, nx)
    func = rand(nz, nx)

    tf_coef = constant(coef)
    tf_func = constant(func)

    # gradient check -- v
    function scalar_function(m)
        # return sum(tanh(laplacian(m, tf_func, tf_h, constant(rho*G))))
        return sum(tanh(laplacian(tf_coef, m, tf_h, constant(rho*G))))
    end

    # m_  = tf_coef
    m_ = tf_func
    v_ = 0.01*rand(nz, nx)
    y_ = scalar_function(m_)
    dy_ = gradients(y_, m_)
    ms_ = Array{Any}(undef, 5)
    ys_ = Array{Any}(undef, 5)
    s_ = Array{Any}(undef, 5)
    w_ = Array{Any}(undef, 5)
    gs_ =  @. 1 / 10^(1:5)

    for i = 1:5
        g_ = gs_[i]
        ms_[i] = m_ + g_*v_
        ys_[i] = scalar_function(ms_[i])
        s_[i] = ys_[i] - y_
        w_[i] = s_[i] - g_*sum(v_.*dy_)
    end

    sess = Session()
    init(sess)
    sval_ = run(sess, s_)
    wval_ = run(sess, w_)
    close("all")
    loglog(gs_, abs.(sval_), "*-", label="finite difference")
    loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
    loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
    loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

    plt[:gca]()[:invert_xaxis]()
    legend()
    xlabel("\$\\gamma\$")
    ylabel("Error")
    savefig("laplacian.png"); close("all")
end

@testset "poisson" begin
    # gradient check -- v
    h = 1.0
    rho = 1000.0
    G = 9.8
    len_z = 16
    len_x = 32
    nz = Int(len_z/h + 1)
    nx = Int(len_x/h + 1)
    tf_h=constant(1.0)
    coef = zeros(nz, nx)
    rhs  = zeros(nz, nx)
    for i = 1:nz
        for j = 1:nx
            rhs[i,j] = -sin(2*pi/len_z*(i-1)*h) * sin(2*pi/len_x*(j-1)*h)
            coef[i,j] = 1.0 - cos(2*pi/len_z*(i-1)*h) * sin(2*pi/len_x*(j-1)*h) * len_z / (2*pi*rho*G)

            # rhs[i,j] = 2.0*(i-1)*h*exp(-(((i-1)*h)^2) -(((j-1)*h)^2)) * rho * G
            # coef[i,j] = 1.0 + exp(-(((i-1)*h)^2) -(((j-1)*h)^2))
        end
    end

    tf_coef = constant(coef)
    tf_rhs = constant(rhs)
    function scalar_function(m)
        return sum(tanh(poisson_op(tf_coef,m,tf_h,constant(rho*G), constant(0))))
        # return sum(tanh(poisson_op(m,tf_rhs,tf_h, constant(rho*G), constant(0))))
    end

    m_ = tf_rhs
    # m_  = tf_coef
    v_ = 0.01*rand(nz, nx)
    y_ = scalar_function(m_)
    dy_ = gradients(y_, m_)
    ms_ = Array{Any}(undef, 5)
    ys_ = Array{Any}(undef, 5)
    s_ = Array{Any}(undef, 5)
    w_ = Array{Any}(undef, 5)
    gs_ =  @. 1 / 20^(1:5)

    for i = 1:5
        g_ = gs_[i]
        ms_[i] = m_ + g_*v_
        ys_[i] = scalar_function(ms_[i])
        s_[i] = ys_[i] - y_
        w_[i] = s_[i] - g_*sum(v_.*dy_)
    end

    sess = Session()
    init(sess)
    sval_ = run(sess, s_)
    wval_ = run(sess, w_)
    close("all")
    loglog(gs_, abs.(sval_), "*-", label="finite difference")
    loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
    loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
    loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

    plt.gca().invert_xaxis()
    legend()
    xlabel("\$\\gamma\$")
    ylabel("Error")
    savefig("poisson.png"); close("all")
end

@testset "sat_op" begin
    function ave_normal(quantity, m, n)
      aa = sum(quantity)
      return aa/(m*n)
    end

    # TODO:
    # const ALPHA = 0.006323996017182
    const ALPHA = 1.0
    const SRC_CONST = 86400.0
    const K_CONST =  9.869232667160130e-16 * 86400
    nz=20
    nx=30
    sw = constant(zeros(nz, nx))
    swref = constant(zeros(nz,nx))
    μw = constant(0.001)
    μo = constant(0.003)
    K = constant(100.0 .* ones(nz, nx))
    ϕ = constant(0.25 .* ones(nz, nx))
    dt = constant(30.0)
    h = constant(100.0 * 0.3048)
    q1 = zeros(nz,nx)
    q2 = zeros(nz,nx)
    q1[10,5] = 0.002 * (1/(100.0 * 0.3048)^2)/20.0/0.3048 * SRC_CONST
    q2[10,25] = -0.002 * (1/(100.0 * 0.3048)^2)/20.0/0.3048 * SRC_CONST
    qw = constant(q1)
    qo = constant(q2)

    λw = sw.*sw/μw
    λo = (1-sw).*(1-sw)/μo
    λ = λw + λo
    f = λw/λ
    q = qw + qo + λw/(λo+1e-16).*qo

    # Θ = laplacian_op(K.*λo, potential_c, h, constant(0.0))
    Θ = upwlap_op(K*K_CONST, λo, constant(zeros(nz,nx)), h, constant(0.0))

    load_normal = (Θ+q/ALPHA) - ave_normal(Θ+q/ALPHA, nz, nx)

    tf_comp_p0 = upwps_op(K*K_CONST, λ, load_normal, constant(zeros(nz,nx)), h, constant(0.0), constant(2))
    sess = Session()
    init(sess)
    p0 = run(sess, tf_comp_p0)
    tf_p0 = constant(p0)

    # s = sat_op(sw,p0,K,ϕ,qw,qo,sw,dt,h)

    # function step(sw)
    #     λw = sw.*sw
    #     λo = (1-sw).*(1-sw)
    #     λ = λw + λo
    #     f = λw/λ
    #     q = qw + qo + λw/(λo+1e-16).*qo

    #     # Θ = laplacian_op(K.*λo, constant(zeros(nz,nx)), h, constant(0.0))
    #     Θ = upwlap_op(K, λo, constant(zeros(nz,nx)), h, constant(0.0))
    #     # Θ = constant(zeros(nz,nx))

    #     load_normal = (Θ+q/ALPHA) - ave_normal(Θ+q/ALPHA, nz, nx)

    #     p = poisson_op(λ.*K, load_normal, h, constant(0.0), constant(0)) # potential p = pw - ρw*g*h
    #     # p = upwps_op(K, λ, load_normal, constant(zeros(nz,nx)), h, constant(0.0), constant(0))
        # sw = sat_op(sw,p,K,ϕ,qw,qo,μw,μo,sw,dt,h)
    #     return sw
    # end

    # NT=100
    # function evolve(sw, NT, qw, qo)
    #     # qw_arr = constant(qw) # qw: NT x m x n array
    #     # qo_arr = constant(qo)
    #     tf_sw = TensorArray(NT+1)
    #     function condition(i, ta)
    #         tf.less(i, NT+1)
    #     end
    #     function body(i, tf_sw)
    #         sw_local = step(read(tf_sw, i))
    #         i+1, write(tf_sw, i+1, sw_local)
    #     end
    #     tf_sw = write(tf_sw, 1, sw)
    #     i = constant(1, dtype=Int32)
    #     _, out = while_loop(condition, body, [i;tf_sw])
    #     read(out, NT+1)
    # end

    # s = evolve(sw, NT, qw, qo)

    # J = tf.nn.l2_loss(s)
    # tf_grad_K = gradients(J, K)
    # sess = Session()
    # init(sess)
    # # P = run(sess,p0)

    # # error("")
    # S=run(sess, s)
    # imshow(S);colorbar();

    # error("")

    # grad_K = run(sess, tf_grad_K)
    # imshow(grad_K);colorbar();
    # error("")
    # TODO:

    # gradient check -- v
    function scalar_function(m)
        # return sum(tanh(sat_op(m,tf_p0,K*K_CONST,ϕ,qw,qo,μw,μo,constant(zeros(nz,nx)),dt,h)))
        return sum(tanh(sat_op(sw,m,K*K_CONST,ϕ,qw,qo,μw,μo,constant(zeros(nz,nx)),dt,h)))
        # return sum(tanh(sat_op(sw,tf_p0,m,ϕ,qw,qo,μw,μo,constant(zeros(nz,nx)),dt,h)))
        # return sum(tanh(sat_op(sw,tf_p0,K*K_CONST,m,qw,qo,μw,μo,constant(zeros(nz,nx)),dt,h)))
    end

    # m_ = sw
    # v_ = 0.1 * rand(nz,nx)

    m_ = tf_p0
    v_ = 5e5 .* rand(nz,nx)

    # m_ = K*K_CONST
    # v_ = 10 .* rand(nz,nx) *K_CONST

    # m_ = ϕ
    # v_ = 0.1 * rand(nz,nx)

    y_ = scalar_function(m_)
    dy_ = gradients(y_, m_)
    ms_ = Array{Any}(undef, 5)
    ys_ = Array{Any}(undef, 5)
    s_ = Array{Any}(undef, 5)
    w_ = Array{Any}(undef, 5)
    gs_ =  @. 1 / 10^(1:5)

    for i = 1:5
        g_ = gs_[i]
        ms_[i] = m_ + g_*v_
        ys_[i] = scalar_function(ms_[i])
        s_[i] = ys_[i] - y_
        w_[i] = s_[i] - g_*sum(v_.*dy_)
    end

    sess = Session()
    init(sess)
    sval_ = run(sess, s_)
    wval_ = run(sess, w_)
    close("all")
    loglog(gs_, abs.(sval_), "*-", label="finite difference")
    loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
    loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
    loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

    plt.gca().invert_xaxis()
    legend()
    xlabel("\$\\gamma\$")
    ylabel("Error")
    savefig("sat_op.png"); close("all")
end

@testset "upwlap_op" begin
    h = 1.0
    rho = 1000.0
    G = 9.8
    len_z = 16
    len_x = 32
    nz = Int(len_z/h + 1)
    nx = Int(len_x/h + 1)
    tf_h=constant(1.0)

    perm = rand(nz, nx)
    mobi = rand(nz, nx)
    func = rand(nz, nx)

    tf_perm = constant(perm)
    tf_mobi = constant(mobi)
    tf_func = constant(func)

    # gradient check -- v
    function scalar_function(m)
        # return sum(tanh(upwlap_op(m, tf_mobi, tf_func, tf_h, constant(rho*G))))
        # return sum(tanh(upwlap_op(tf_perm, m, tf_func, tf_h, constant(rho*G))))
        return sum(tanh(upwlap_op(tf_perm, tf_mobi, m, tf_h, constant(rho*G))))
    end

    # m_ = constant(rand(10,20))
    # m_ = tf_perm
    # m_ = tf_mobi
    m_ = tf_func
    v_ = rand(nz, nx)
    y_ = scalar_function(m_)
    dy_ = gradients(y_, m_)
    ms_ = Array{Any}(undef, 5)
    ys_ = Array{Any}(undef, 5)
    s_ = Array{Any}(undef, 5)
    w_ = Array{Any}(undef, 5)
    gs_ =  @. 1 / 20^(1:5)

    for i = 1:5
        g_ = gs_[i]
        ms_[i] = m_ + g_*v_
        ys_[i] = scalar_function(ms_[i])
        s_[i] = ys_[i] - y_
        w_[i] = s_[i] - g_*sum(v_.*dy_)
    end

    sess = Session()
    init(sess)
    sval_ = run(sess, s_)
    wval_ = run(sess, w_)
    close("all")
    loglog(gs_, abs.(sval_), "*-", label="finite difference")
    loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
    loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
    loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

    plt.gca().invert_xaxis()
    legend()
    xlabel("\$\\gamma\$")
    ylabel("Error")
    savefig("upwlap_op.png"); close("all")
end

@testset "upwps_op" begin
    h = 20.0
    rho = 1000.0
    G = 9.8
    len_z = 16*20
    len_x = 32*20
    nz = Int(len_z/h + 1)
    nx = Int(len_x/h + 1)
    tf_h=constant(1.0)
    coef = zeros(nz, nx)
    ones_cons = ones(nz, nx)
    rhs  = zeros(nz, nx)
    for i = 1:nz
        for j = 1:nx
            rhs[i,j] = -sin(2*pi/len_z*(i-1)*h) * sin(2*pi/len_x*(j-1)*h)
            coef[i,j] = 1.0 - cos(2*pi/len_z*(i-1)*h) * sin(2*pi/len_x*(j-1)*h) * len_z / (2*pi*rho*G)

            # rhs[i,j] = 2.0*(i-1)*h*exp(-(((i-1)*h)^2) -(((j-1)*h)^2)) * rho * G
            # coef[i,j] = 1.0 + exp(-(((i-1)*h)^2) -(((j-1)*h)^2))
        end
    end
    # coef = 1.0 .+ rand(nz, nx)
    # rhs = rand(nz, nx)

    tf_coef = constant(coef)
    tf_rhs = constant(rhs)
    tf_funcref = constant(rand(nz, nx))
    tf_ones = constant(ones_cons)
    function scalar_function(m)
        # return sum(tanh(upwps_op(tf_coef, tf_ones, m, tf_funcref, tf_h, constant(rho*G), constant(0))))
        # return sum(tanh(upwps_op(m, tf_ones, tf_rhs, tf_ones, tf_h, constant(rho*G), constant(0))))
        return sum(tanh(upwps_op(tf_ones, m, tf_rhs, tf_ones, tf_h, constant(rho*G), constant(0))))
    end

    # m_ = tf_rhs
    m_  = tf_coef
    v_ = 0.01*rand(nz, nx)
    y_ = scalar_function(m_)
    dy_ = gradients(y_, m_)
    ms_ = Array{Any}(undef, 5)
    ys_ = Array{Any}(undef, 5)
    s_ = Array{Any}(undef, 5)
    w_ = Array{Any}(undef, 5)
    gs_ =  @. 1 / 10^(1:5)

    for i = 1:5
        g_ = gs_[i]
        ms_[i] = m_ + g_*v_
        ys_[i] = scalar_function(ms_[i])
        s_[i] = ys_[i] - y_
        w_[i] = s_[i] - g_*sum(v_.*dy_)
    end

    sess = Session()
    init(sess)
    sval_ = run(sess, s_)
    wval_ = run(sess, w_)
    close("all")
    loglog(gs_, abs.(sval_), "*-", label="finite difference")
    loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
    loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
    loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

    plt.gca().invert_xaxis()
    legend()
    xlabel("\$\\gamma\$")
    ylabel("Error")
    savefig("upwps_op.png"); close("all")
end


# This code presents an example to illustrate how NLboxsolve.jl can be used

using NLboxsolve

function test_nlboxsolve_outplace()

    function test6(x)

        f = similar(x)

        for i in eachindex(x)
            f[i] = cos(x[i]) - 1.0
        end

        return f

    end

    # Solve test6 model

    n = 20
    x0 = ones(n)*0.5
    l = ones(n)*-0.5
    u = ones(n)*0.7

    soln_newton     = constrained_newton(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm         = constrained_levenberg_marquardt(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_kyf     = constrained_levenberg_marquardt_kyf(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_ar      = constrained_levenberg_marquardt_ar(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg     = constrained_dogleg_solver(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_bmp = constrained_dogleg_bmp_solver(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_nk         = constrained_newton_krylov(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_jfnk       = constrained_jacobian_free_newton_krylov(test6,x0,l,u,ftol=1e-15,xtol=1e-15)

    soln_newton_sparse     = constrained_newton_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_sparse         = constrained_levenberg_marquardt_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_kyf_sparse     = constrained_levenberg_marquardt_kyf_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_ar_sparse      = constrained_levenberg_marquardt_ar_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_sparse     = constrained_dogleg_solver_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_bmp_sparse = constrained_dogleg_bmp_solver_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_nk_sparse         = constrained_newton_krylov_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)

    soln2 = nlboxsolve(test6,x0,l,u,xtol=1e-15,ftol=1e-15)

    test_one   = maximum(abs,soln_newton.zero       - soln2.zero) < 1e-8
    test_two   = maximum(abs,soln_lm.zero           - soln2.zero) < 1e-8
    test_three = maximum(abs,soln_lm_kyf.zero       - soln2.zero) < 1e-8
    test_four  = maximum(abs,soln_lm_ar.zero        - soln2.zero) < 1e-8
    test_five  = maximum(abs,soln_dogleg.zero       - soln2.zero) < 1e-8
    test_six   = maximum(abs,soln_dogleg_bmp.zero   - soln2.zero) < 1e-8
    test_seven = maximum(abs,soln_nk.zero           - soln2.zero) < 1e-8
    test_eight = maximum(abs,soln_jfnk.zero         - soln2.zero) < 1e-8

    test_nine     = maximum(abs,soln_newton_sparse.zero       - soln2.zero) < 1e-8
    test_ten      = maximum(abs,soln_lm_sparse.zero           - soln2.zero) < 1e-8
    test_eleven   = maximum(abs,soln_lm_kyf_sparse.zero       - soln2.zero) < 1e-8
    test_twelve   = maximum(abs,soln_lm_ar_sparse.zero        - soln2.zero) < 1e-8
    test_thirteen = maximum(abs,soln_dogleg_sparse.zero       - soln2.zero) < 1e-8
    test_fourteen = maximum(abs,soln_dogleg_bmp_sparse.zero   - soln2.zero) < 1e-8
    test_fifteen  = maximum(abs,soln_nk_sparse.zero           - soln2.zero) < 1e-8

    return test_one, test_two, test_three, test_four, test_five, test_six, test_seven, test_eight, test_nine, test_ten, test_eleven,
           test_twelve, test_thirteen, test_fourteen, test_fifteen

end

function test_nlboxsolve_inplace()

    function test6(f,x)

        for i in eachindex(x)
            f[i] = cos(x[i]) - 1.0
        end

    end

    # Solve test6 model

    n = 20
    x0 = ones(n)*0.5
    l = ones(n)*-0.5
    u = ones(n)*0.7

    soln_newton     = constrained_newton(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm         = constrained_levenberg_marquardt(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_kyf     = constrained_levenberg_marquardt_kyf(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_ar      = constrained_levenberg_marquardt_ar(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg     = constrained_dogleg_solver(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_bmp = constrained_dogleg_bmp_solver(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_nk         = constrained_newton_krylov(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_jfnk       = constrained_jacobian_free_newton_krylov(test6,x0,l,u,ftol=1e-15,xtol=1e-15)

    soln_newton_sparse     = constrained_newton_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_sparse         = constrained_levenberg_marquardt_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_kyf_sparse     = constrained_levenberg_marquardt_kyf_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_ar_sparse      = constrained_levenberg_marquardt_ar_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_sparse     = constrained_dogleg_solver_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_bmp_sparse = constrained_dogleg_bmp_solver_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_nk_sparse         = constrained_newton_krylov_sparse(test6,x0,l,u,ftol=1e-15,xtol=1e-15)

    soln2 = nlboxsolve(test6,x0,l,u,xtol=1e-15,ftol=1e-15)

    test_one   = maximum(abs,soln_newton.zero       - soln2.zero) < 1e-8
    test_two   = maximum(abs,soln_lm.zero           - soln2.zero) < 1e-8
    test_three = maximum(abs,soln_lm_kyf.zero       - soln2.zero) < 1e-8
    test_four  = maximum(abs,soln_lm_ar.zero        - soln2.zero) < 1e-8
    test_five  = maximum(abs,soln_dogleg.zero       - soln2.zero) < 1e-8
    test_six   = maximum(abs,soln_dogleg_bmp.zero   - soln2.zero) < 1e-8
    test_seven = maximum(abs,soln_nk.zero           - soln2.zero) < 1e-8
    test_eight = maximum(abs,soln_jfnk.zero         - soln2.zero) < 1e-8

    test_nine     = maximum(abs,soln_newton_sparse.zero       - soln2.zero) < 1e-8
    test_ten      = maximum(abs,soln_lm_sparse.zero           - soln2.zero) < 1e-8
    test_eleven   = maximum(abs,soln_lm_kyf_sparse.zero       - soln2.zero) < 1e-8
    test_twelve   = maximum(abs,soln_lm_ar_sparse.zero        - soln2.zero) < 1e-8
    test_thirteen = maximum(abs,soln_dogleg_sparse.zero       - soln2.zero) < 1e-8
    test_fourteen = maximum(abs,soln_dogleg_bmp_sparse.zero   - soln2.zero) < 1e-8
    test_fifteen  = maximum(abs,soln_nk_sparse.zero           - soln2.zero) < 1e-8

    return test_one, test_two, test_three, test_four, test_five, test_six, test_seven, test_eight, test_nine, test_ten, test_eleven,
           test_twelve, test_thirteen, test_fourteen, test_fifteen

end

test_nlboxsolve_outplace()
test_nlboxsolve_inplace()


module _DistEn2D
export DistEn2D
using StatsBase: fit, Histogram, skewness
    """
        Dist2D = DistEn2D(Mat)

    Returns the bidimensional distribution entropy estimate (`Dist2D`)
    estimated for the data matrix (`Mat`) using the default parameters:
    time delay = 1, histogram binning method = "sturges", logarithm = natural,
    template matrix size = [floor(H/10) floor(W/10)], (where H and W represent
    the height (rows) and width (columns) of the data matrix `Mat`) \n
    ** The minimum number of rows and columns of Mat must be > 10.**

        Dist2D = DistEn2D(Mat::AbstractArray{T,2} where T<:Real; m::Union{Int,Tuple{Int,Int}}=floor.(Int, size(Mat)./10), tau::Int=1,
                            Bins::Union{Int,String}="Sturges", Logx::Real=2, Norm::Int=2, Lock::Bool=true)

    Returns the bidimensional distribution entropy (`Dist2D`) estimate for
    the data matrix (`Mat`) using the specified 'keyword' arguments:

    # Arguments:
    `m`     - Template submatrix dimensions, an integer scalar (i.e. the same
              height and width) or a two-element tuple of integers
              [height, width] with a minimum value > 1.
              [default: [floor(H/10) floor(W/10)]] \n
    `tau`   - Time Delay, a positive integer       [default: 1]  \n
    `Bins`  - Histogram bin selection method for distance distribution,
              an integer > 1 indicating the number of bins, or one of the
              following strings {`"sturges", "sqrt", "rice", "doanes"``}
              [default: 'sturges']  \n
    `Logx`  - Logarithm base, a positive scalar        [default: natural]\n
              ** enter 0 for natural logarithm.**\n
    `Norm`  - Normalisation of `Dist2D` value, one of the following integers:
              [0]  no normalisation.
              [1]  normalises values of data matrix (`Mat`) to range [0 1].
              [2]  normalises values of data matrix (`Mat`) to range [0 1],
                   and normalises the distribution entropy value (`Dist2D`)
                   w.r.t the number of histogram bins.  [default]
              [3]  normalises the distribution entropy value
                   w.r.t the number of histogram bins, without normalising
                   data matrix values. \n
    `Lock`  - By default, DistEn2D only permits matrices with a maximum
              size of 128 x 128 to prevent memory errors when storing data on RAM.
              e.g. For Mat = [200 x 200], m = 3, and tau = 1, DistEn2D
              creates a vector of 753049836 elements. To enable matrices
              greater than [128 x 128] elements, set `Lock` to false.
              [default: 'true']
              `WARNING: unlocking the permitted matrix size may cause your Julia
               IDE to crash.`

    # See also `DistEn`, `XDistEn`, `SampEn2D`, `FuzzEn2D`, `MSEn`

    # References:
        [1] Hamed Azami, Javier Escudero and Anne Humeau-Heurtier,
            "Bidimensional distribution entropy to analyze the irregularity
            of small-sized textures."
            IEEE Signal Processing Letters
            24.9 (2017): 1338-1342.

    """
    function DistEn2D(Mat::AbstractArray{T,2} where T<:Real;
        m::Union{Int,Tuple{Int,Int}}=floor.(Int, size(Mat)./10), tau::Int=1,
        Bins::Union{Int,String}="Sturges", Logx::Real=2, Norm::Int=2, Lock::Bool=true)

    Logx == 0  ? Logx = exp(1) : nothing

    NL, NW = size(Mat)
    ((NL > 128 || NW > 128) && Lock) ?
        error("To prevent memory errors, matrix width & length must have <= 128 elements.
            To estimate DistEn2D  for the current matrix ($NL,$NW) change Lock to 'false'.
            Caution: unlocking the safe matrix size may cause the Julia IDE to crash.") :
            nothing

    length(m)==1 ?  (mL = m; mW = m) : (mL = m[1];  mW = m[2])
    (NL > 10 && NW > 10) ? nothing :
    error("Number of rows and columns in Mat must be > 10")
    (minimum(m) > 1) ? nothing : error("m:   must be an integer > 1")
    (tau >0) ? nothing : error("tau:   must be an integer > 0")
    (Logx>0) ? nothing : error("Logx:  must be a positive number > 0")
    if typeof(Bins)<:Int
        (Bins>1) ? nothing :
        error("Bins:  must be an integer > 1 (or name of binning method)")
    elseif typeof(Bins)<:String
        (lowercase(Bins) in ["sturges","sqrt","rice","doanes"]) ? nothing :
        error("Bins:    must be one of the following strings
                'sturges', 'sqrt', 'rice', 'doanes' (or an integer >1)")
    end
    (Norm in [0, 1, 2, 3]) ? nothing :  error("Norm:  must be an integer in the range [0 3]")

    Norm in [1, 2] ? Mat = (Mat.-minimum(Mat))./maximum(Mat.-minimum(Mat)) : nothing
    NL = NL - (mL-1)*tau
    NW = NW - (mW-1)*tau
    X = zeros(mL,mW,NL*NW)
    p = 0
    for k = 1:NL
        for n = 1:NW
            p += 1
            X[:,:,p] = Mat[k:tau:(mL-1)*tau+k,n:tau:(mW-1)*tau+n]
        end
    end

    p = size(X,3)
    p != NL*NW ? @warn("Potential error with submatrix division.") : nothing
    Ny = Int(p*(p-1)/2)
    Ny > 300000000 ? @warn("Number of pairwise distance calculations is $Ny") : nothing

    Y = zeros(Ny)
    for k = 1:p-1
        Ix = Int.([(k-1)*(p - k/2)+1, k*(p-((k+1)/2))])
        Y[Ix[1]:Ix[2]] = maximum(abs.(X[:,:,k+1:end] .- X[:,:,k]),dims=(1,2))
    end

    if eltype(Bins)<:Char
        if lowercase(Bins) == "sturges"
            Bx = ceil(log2(Ny) + 1)
        elseif lowercase(Bins) == "rice"
            Bx = ceil(2*(Ny^(1/3)))
        elseif lowercase(Bins) == "sqrt"
            Bx = ceil(sqrt(Ny))
        elseif lowercase(Bins) == "doanes"
            sigma = sqrt(6*(Ny-2)/((Ny+1)*(Ny+3)))
            Bx = ceil(1+log2(Ny)+log2(1+abs(skewness(convert(Array{Float64,1},Y))/sigma)))
        else
            error("Please enter a valid binning method")
        end
    else
        Bx = Bins
    end

    By = collect(range(minimum(Y),maximum(Y),length=Int(Bx+1)))
    By[end] += 1; By[1] -= 1

    Ppi = fit(Histogram, Y[:], By).weights/Ny
    if round(sum(Ppi),digits=6) != 1
        @warn("Potential error estimating probabilities (p = $(sum(Ppi)))")
        Ppi = Ppi[Ppi.!=0]
    elseif any(Ppi.==0)
        print("Note: $(sum(Ppi.==0))/$(length(Ppi)) bins were empty \n")
        Ppi = Ppi[Ppi.!=0]
    end
    Dist2D = -sum(Ppi.*log.(Logx, Ppi))
    Norm >= 2  ? Dist2D /= (log(Logx, Bx)) : nothing

    return Dist2D
    end

end

"""
Copyright 2021 Matthew W. Flood, EntropyHub

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub
"""

#= Y = []
for k = 1:p-1
    append!(Y,maximum(abs.(X[:,:,k+1:end] .- X[:,:,k]),dims=(1,2)))
end =#


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

module TestBunchKaufman

using Test, LinearAlgebra, Random
using LinearAlgebra: BlasComplex, BlasFloat, BlasReal, QRPivoted
using Base: getproperty

n = 10

# Split n into 2 parts for tests needing two matrices
n1 = div(n, 2)
n2 = 2*n1

Random.seed!(12343210)

areal = randn(n,n)/2
aimg  = randn(n,n)/2
a2real = randn(n,n)/2
a2img  = randn(n,n)/2
breal = randn(n,2)/2
bimg  = randn(n,2)/2

@testset "$eltya argument A" for eltya in (Float32, Float64, ComplexF32, ComplexF64, Int)
    a = eltya == Int ? rand(1:7, n, n) : convert(Matrix{eltya}, eltya <: Complex ? complex.(areal, aimg) : areal)
    a2 = eltya == Int ? rand(1:7, n, n) : convert(Matrix{eltya}, eltya <: Complex ? complex.(a2real, a2img) : a2real)
    asym = transpose(a) + a                  # symmetric indefinite
    aher = a' + a                  # Hermitian indefinite
    apd  = a' * a                  # Positive-definite
    for (a, a2, aher, apd) in ((a, a2, aher, apd),
                               (view(a, 1:n, 1:n),
                                view(a2, 1:n, 1:n),
                                view(aher, 1:n, 1:n),
                                view(apd , 1:n, 1:n)))
        ε = εa = eps(abs(float(one(eltya))))

        # check that factorize gives a Bunch-Kaufman
        @test isa(factorize(asym), LinearAlgebra.BunchKaufman)
        @test isa(factorize(aher), LinearAlgebra.BunchKaufman)
        @testset "$uplo Bunch-Kaufman factor of indefinite matrix" for uplo in (:L, :U)
            bc1 = bunchkaufman(Hermitian(aher, uplo))
            @test LinearAlgebra.issuccess(bc1)
            @test logabsdet(bc1)[1] ≈ log(abs(det(bc1)))
            if eltya <: Real
                @test logabsdet(bc1)[2] == sign(det(bc1))
            else
                @test logabsdet(bc1)[2] ≈ sign(det(bc1))
            end
            @test inv(bc1)*aher ≈ Matrix(I, n, n)
            @testset for rook in (false, true)
                @test inv(bunchkaufman(Symmetric(transpose(a) + a, uplo), rook))*(transpose(a) + a) ≈ Matrix(I, n, n)
                if eltya <: BlasFloat
                    # test also bunchkaufman! without explicit type tag
                    # no bunchkaufman! method for Int ... yet
                    @test inv(bunchkaufman!(transpose(a) + a, rook))*(transpose(a) + a) ≈ Matrix(I, n, n)
                end
                @test size(bc1) == size(bc1.LD)
                @test size(bc1, 1) == size(bc1.LD, 1)
                @test size(bc1, 2) == size(bc1.LD, 2)
                if eltya <: BlasReal
                    @test_throws ArgumentError bunchkaufman(a)
                end
                # Test extraction of factors
                if eltya <: Real
                    @test getproperty(bc1, uplo)*bc1.D*getproperty(bc1, uplo)' ≈ aher[bc1.p, bc1.p]
                    @test getproperty(bc1, uplo)*bc1.D*getproperty(bc1, uplo)' ≈ bc1.P*aher*bc1.P'
                end

                bc1 = bunchkaufman(Symmetric(asym, uplo))
                @test getproperty(bc1, uplo)*bc1.D*transpose(getproperty(bc1, uplo)) ≈ asym[bc1.p, bc1.p]
                @test getproperty(bc1, uplo)*bc1.D*transpose(getproperty(bc1, uplo)) ≈ bc1.P*asym*transpose(bc1.P)
                @test_throws ErrorException bc1.Z
                @test_throws ArgumentError uplo == :L ? bc1.U : bc1.L
            end
            # test Base.iterate
            ref_objs = (bc1.D, uplo == :L ? bc1.L : bc1.U, bc1.p)
            for (bki, bkobj) in enumerate(bc1)
                @test bkobj == ref_objs[bki]
            end
            if eltya <: BlasFloat
                @test convert(LinearAlgebra.BunchKaufman{eltya}, bc1) === bc1
                @test convert(LinearAlgebra.Factorization{eltya}, bc1) === bc1
                if eltya <: BlasReal
                    @test convert(LinearAlgebra.Factorization{Float16}, bc1) == convert(LinearAlgebra.BunchKaufman{Float16}, bc1)
                elseif eltya <: BlasComplex
                    @test convert(LinearAlgebra.Factorization{ComplexF16}, bc1) == convert(LinearAlgebra.BunchKaufman{ComplexF16}, bc1)
                end
            end
            @test Base.propertynames(bc1) == (:p, :P, :L, :U, :D)
        end

        @testset "$eltyb argument B" for eltyb in (Float32, Float64, ComplexF32, ComplexF64, Int)
            b = eltyb == Int ? rand(1:5, n, 2) : convert(Matrix{eltyb}, eltyb <: Complex ? complex.(breal, bimg) : breal)
            for b in (b, view(b, 1:n, 1:2))
                εb = eps(abs(float(one(eltyb))))
                ε = max(εa,εb)

                @testset "$uplo Bunch-Kaufman factor of indefinite matrix" for uplo in (:L, :U)
                    bc1 = bunchkaufman(Hermitian(aher, uplo))
                    @test aher*(bc1\b) ≈ b atol=1000ε
                end

                @testset "$uplo Bunch-Kaufman factors of a pos-def matrix" for uplo in (:U, :L)
                    @testset "rook pivoting: $rook" for rook in (false, true)
                        bc2 = bunchkaufman(Hermitian(apd, uplo), rook)
                        @test LinearAlgebra.issuccess(bc2)
                        bks = split(sprint(show, "text/plain", bc2), "\n")
                        @test bks[1] == summary(bc2)
                        @test bks[2] == "D factor:"
                        @test bks[4+n] == "$uplo factor:"
                        @test bks[6+2n] == "permutation:"
                        @test logdet(bc2) ≈ log(det(bc2))
                        @test logabsdet(bc2)[1] ≈ log(abs(det(bc2)))
                        @test logabsdet(bc2)[2] == sign(det(bc2))
                        @test inv(bc2)*apd ≈ Matrix(I, n, n)
                        @test apd*(bc2\b) ≈ b rtol=eps(cond(apd))
                        @test ishermitian(bc2) == !issymmetric(bc2)
                    end
                end
            end
        end
    end
end

@testset "Singular matrices" begin
    R = Float64[1 0; 0 0]
    C = ComplexF64[1 0; 0 0]
    for A in (R, Symmetric(R), C, Hermitian(C))
        @test_throws SingularException bunchkaufman(A)
        @test_throws SingularException bunchkaufman!(copy(A))
        @test_throws SingularException bunchkaufman(A; check = true)
        @test_throws SingularException bunchkaufman!(copy(A); check = true)
        @test !issuccess(bunchkaufman(A; check = false))
        @test !issuccess(bunchkaufman!(copy(A); check = false))
    end
    F = bunchkaufman(R; check = false)
    @test sprint(show, "text/plain", F) == "Failed factorization of type $(typeof(F))"
end

@testset "test example due to @timholy in PR 15354" begin
    A = rand(6,5); A = complex(A'*A) # to avoid calling the real-lhs-complex-rhs method
    F = cholesky(A);
    v6 = rand(ComplexF64, 6)
    v5 = view(v6, 1:5)
    @test F\v5 == F\v6[1:5]
end

@testset "issue #32080" begin
    A = Symmetric([-5 -9 9; -9 4 1; 9 1 2])
    B = bunchkaufman(A, true)
    @test B.U * B.D * B.U' ≈ A[B.p, B.p]
end

@test_throws DomainError logdet(bunchkaufman([-1 -1; -1 1]))
@test logabsdet(bunchkaufman([8 4; 4 2]; check = false))[1] == -Inf

@testset "0x0 matrix" begin
    for ul in (:U, :L)
        B = bunchkaufman(Symmetric(ones(0, 0), ul))
        @test isa(B, BunchKaufman)
        @test B.D == Tridiagonal([], [], [])
        @test B.P == ones(0, 0)
        @test B.p == []
        if ul == :U
            @test B.U == UnitUpperTriangular(ones(0, 0))
            @test_throws ArgumentError B.L
        else
            @test B.L == UnitLowerTriangular(ones(0, 0))
            @test_throws ArgumentError B.U
        end
    end
end

end # module TestBunchKaufman


ForwardDiff_gradient(x...) = ForwardDiff.gradient(x...)

function kr(psi, alpha, N)
	if psi < 0
		m = (N - 1) / N
		denom = 1 + abs(alpha * psi) ^ N
		numer = 1 - abs(alpha * psi) ^ (N - 1) * denom ^ (-m)
		return numer ^ 2 / denom ^ (m / 2)
	else
		return one(psi)
	end
end

kr(x::AbstractArray) = kr(x[1], x[2], x[3])

function Calculus.differentiate(x::Calculus.SymbolParameter{:kr}, args, wrt)
	chain_part = map(x->Calculus.simplify(Calculus.differentiate(x, wrt)), args)
	if chain_part == [0, 0, 0]
		return :(0)
	else
		return :(sum(ForwardDiff_gradient(kr, StaticArrays.SA[$(args...)]) .* StaticArrays.SA[$(chain_part...)]))
	end
end

function hm(x, y)
	return 2 / (1 / x + 1 / y)
end

hm(x::AbstractArray) = hm(x[1], x[2])

function Calculus.differentiate(x::Calculus.SymbolParameter{:hm}, args, wrt)
	chain_part = map(x->Calculus.simplify(Calculus.differentiate(x, wrt)), args)
	if chain_part == [0, 0]
		return :(0)
	else
		return :(sum(ForwardDiff_gradient(hm, StaticArrays.SA[$(args...)]) .* StaticArrays.SA[$(chain_part...)]))
	end
end

function Calculus.differentiate(x::Calculus.SymbolParameter{:abs}, args, wrt)
	if length(args) > 1
		error("Too many arguments passed to abs()")
	end
	arg = args[1]
	return :(ifelse($arg > 0, 1, -1) * $(Calculus.differentiate(arg, wrt)))
end

function effective_saturation(alpha::Number, psi::Number, N::Number)
	m = (N - 1) / N
	if psi < 0
		return (1 + abs(alpha * psi) ^ N) ^ (-m)
	else
		return 1
	end
end

@NonlinearEquations.equations exclude=(coords, dirichletnodes, neighbors, areasoverlengths) function richards(psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, specificstorage, volumes)
	NonlinearEquations.setnumequations(length(psi))
	dirichletnodes_set = Set(dirichletnodes)
	for i = 1:length(psi)
		if i in dirichletnodes_set
			NonlinearEquations.addterm(i, psi[i] - dirichletpsis[i])
		else
			NonlinearEquations.addterm(i, Qs[i] / (specificstorage[i] * volumes[i]))
		end
	end
	for (i, (node_a, node_b)) in enumerate(neighbors)
		for (node1, node2) in [(node_a, node_b), (node_b, node_a)]
			if !(node1 in dirichletnodes_set) && !(node2 in dirichletnodes_set)
				NonlinearEquations.addterm(node1, hm(kr(psi[node1], alphas[i], Ns[i]), kr(psi[node2], alphas[i], Ns[i])) * Ks[i] * (psi[node2] + coords[3, node2] - psi[node1] - coords[3, node1]) * areasoverlengths[i] / (specificstorage[node1] * volumes[node1]))
			elseif !(node1 in dirichletnodes_set) && node2 in dirichletnodes_set
				NonlinearEquations.addterm(node1, hm(kr(psi[node1], alphas[i], Ns[i]), kr(dirichletpsis[node2], alphas[i], Ns[i])) * Ks[i] * (dirichletpsis[node2] + coords[3, node2] - psi[node1] - coords[3, node1]) * areasoverlengths[i] / (specificstorage[node1] * volumes[node1]))
			end
		end
	end
end

function richards_steadystate(psi0, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs; kwargs...)
	args = (psi0, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, ones(length(Qs)), ones(length(Qs)))
	function residuals!(residuals, psi)
		myargs = (psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, ones(length(Qs)), ones(length(Qs)))
		copy!(residuals, DPFEHM.richards_residuals(myargs...))
	end
	function jacobian!(J, psi)
		myargs = (psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, ones(length(Qs)), ones(length(Qs)))
		DPFEHM.richards_psi!(J, myargs...)
	end
	J0 = DPFEHM.richards_psi(args...)
	residuals0 = DPFEHM.richards_residuals(args...)
	df = NLsolve.OnceDifferentiable(residuals!, jacobian!, psi0, residuals0, J0)
	soln = NLsolve.nlsolve(df, psi0; kwargs...)
	if !NLsolve.converged(soln)
		display(soln)
		error("solution did not converge")
	end
	return soln.zero
end

function ChainRulesCore.rrule(::typeof(richards_steadystate), psi0, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs; kwargs...)
	psi = richards_steadystate(psi0, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs; kwargs...)
	function pullback(delta)
		args = (psi, Ks, neighbors, areasoverlengths, dirichletnodes, dirichletpsis, coords, alphas, Ns, Qs, ones(length(Qs)), ones(length(Qs)))
		lambda = richards_psi(args...)' \ delta
		req_Ks = richards_Ks(args...)
		req_dirichletpsis = richards_dirichletpsis(args...)
		req_alphas = richards_alphas(args...)
		req_Ns = richards_Ns(args...)
		req_Qs = richards_Qs(args...)
		return (ChainRulesCore.NoTangent(),#step function
				@ChainRulesCore.thunk(zeros(length(psi0))),#psi0 -- should be insensitive to psi0
				@ChainRulesCore.thunk(-(req_Ks' * lambda)),#Ks
				@ChainRulesCore.thunk(ChainRulesCore.NoTangent()),#neighbors
				@ChainRulesCore.thunk(ChainRulesCore.NoTangent()),#areasoverlengths
				@ChainRulesCore.thunk(ChainRulesCore.NoTangent()),#dirichletnodes
				@ChainRulesCore.thunk(-(req_dirichletpsis' * lambda)),#dirichletpsis
				@ChainRulesCore.thunk(ChainRulesCore.NoTangent()),#coords
				@ChainRulesCore.thunk(-(req_alphas' * lambda)),#alphas
				@ChainRulesCore.thunk(-(req_Ns' * lambda)),#Ns
				@ChainRulesCore.thunk(-(req_Qs' * lambda)))#Qs
	end
	return psi, pullback
end


using SpectralMeasures, ValidatedNumerics, ValidatedNumerics.RootFinding,
            ComplexPhasePortrait

function tripleplot(a,b,Z=linspace(-3, 3, 300).+linspace(3,-3,300)'*im)
  # Build the measure and resolvents
  μ = spectralmeasure(a,b)
  R = principalresolvent(a,b)
  r = discresolvent(a,b)

  # Create the subplots
  p1 = plot(μ,xlims=(-2,2),ylims=(0,1.5))
  p2 = plot(portrait(R(Z),PTstepmod),xlims=(-3,3),ylims=(-3,3),aspect_ratio=1)
  plot!([-1.,1.],[0.,0.],linewidth=3,color=:black)
  p3 = plot(portrait(r(Z),PTstepmod),xlims=(-3,3),ylims=(-3,3),aspect_ratio=1)
  plot!(cos(linspace(0,2pi,100)),sin(linspace(0,2pi,100)),linewidth=3,color=:black)

  # Plot the subplots
  l = @layout [a{0.7w}; b c]
  plot(p1,p2,p3,layout=l)
end

function validatedspectrum(a,b)
  # Chop the a and b down
  a = chop!(a); b = .5+chop!(b-.5)
  n = max(2,length(a),length(b)+1)
  a = [a;zeros(n-length(a))]; b = [b;.5+zeros(n-length(b))]

  # Finds C such that J*C = C*Toeplitz([0,1/2])
  C = SpectralMeasures.connectioncoeffsoperator(a,b)
  c = Fun(Taylor,C.T.nonnegative)

  rts = find_roots(x->c(x),-1,1)
  if length(rts) > 0
     eigs=real(map(x->joukowsky(x.interval),rts))
     eigserrs = map(midpoint_radius,eigs)
     spectrum = ApproxFun.Interval(-1,1)
     maxerr = 0
     for (eig,err) in eigserrs
       maxerr = max(maxerr,err)
       spectrum = ApproxFun.Point(eig) ∪ spectrum
     end
  else
    spectrum = ApproxFun.Interval(-1,1)
    maxerr = 0
  end
  spectrum, maxerr
end


module Normalization

using Statistics
using JuliennedArrays
import LinearAlgebra:   normalize,
                        normalize!

export  fit,
        fit!,
        normalize!,
        normalize,
        denormalize!,
        denormalize,
        nansafe,
        ZScore,
        RobustZScore,
        Sigmoid,
        RobustSigmoid,
        MinMax

abstract type AbstractNormalization end
(𝒯::Type{<:AbstractNormalization})(dims) = 𝒯(;dims)
function (𝒯::Type{<:AbstractNormalization})(dims, p)
    isnothing(p) || (all(x->x==p[1], length.(p)) && error("Inconsistent parameter dimensions"))
    𝒯(;dims, p)
end
(T::AbstractNormalization)(dims) = dims == () || (T.dims = dims)

macro _Normalization(name, 𝑝, 𝑓, 𝑓⁻¹)
    :(mutable struct $(esc(name)) <: AbstractNormalization
        dims
        p::Union{Nothing, NTuple{length($𝑝), AbstractArray}}
        𝑝::NTuple{length($𝑝), Function}
        𝑓::Function
        𝑓⁻¹::Function
     end;
     ($(esc(name)))(; dims = nothing,
                         p = nothing,
                         𝑝 = $𝑝,
                         𝑓 = $𝑓,
                         𝑓⁻¹ = $𝑓⁻¹) = $(esc(name))(dims, p, 𝑝, 𝑓, 𝑓⁻¹)
     )
end

# * Common normalizations
@_Normalization ZScore (mean, std)         (x, 𝜇, 𝜎) -> x .= (x .- 𝜇)./𝜎  #=
                                        =# (y, 𝜇, 𝜎) -> y .= y.*𝜎 .+ 𝜇
@_Normalization Sigmoid (mean, std)        (x, 𝜇, 𝜎)->x.=1.0./(1 .+exp.(.-(x.-𝜇)./𝜎)) #=
                                        =# (y, 𝜇, 𝜎) -> y .= .-𝜎.*log.(1.0./y .- 1) .+ 𝜇
@_Normalization MinMax (minimum, maximum)  (x, l, u) -> x .= (x.-l)./(u-l) #=
                                        =# (y, l, u) -> y .= (u-l).*y .+ l
@_Normalization Center (mean,)             (x, 𝜇) -> x .= x .- 𝜇     (y, 𝜇) -> y .= y .+ 𝜇
@_Normalization RobustCenter (median,)     Centre().𝑓   Centre().𝑓⁻¹

# * Robust versions of typical 2-parameter normalizations
common_norms = [:ZScore, :Sigmoid,]
_iqr = x -> (quantile(x[:], 0.75) - quantile(x[:], 0.25))/1.35 # ? Divide by 1.35 so that std(x) ≈ _iqr(x) when x contains normally distributed values
_robustNorm(N::Symbol; name="Robust"*string(N)|>Symbol) = eval(:(@_Normalization $name (median, _iqr) ($N)().𝑓 ($N)().𝑓⁻¹))
_robustNorm.(common_norms)

# * Mixed versions of typical 2-parameter normalizations
mixedcenter(x) = (_iqr(x) == 0) ? mean(x) : median(x)
mixedscale(x) = (𝜎 = _iqr(x); 𝜎 == 0 ? std(x) : 𝜎)
_mixedNorm(N::Symbol; name="Mixed"*string(N)|>Symbol) = eval(:(@_Normalization $name (mixedcenter, mixedscale) ($N)().𝑓 ($N)().𝑓⁻¹))
_mixedNorm.(common_norms)

# * NaN-safe versions
_nansafe(p) = x -> p(filter(!isnan, x))
nansafe!(T::AbstractNormalization) = (T.𝑝=_nansafe.(T.𝑝); ())
nansafe(T::AbstractNormalization) = (N = deepcopy(T); nansafe!(N); N)
nansafe(𝒯::Type{<:AbstractNormalization}; dims=nothing) = dims |> 𝒯 |> nansafe

Base.reshape(x::Number, dims...) = reshape([x], dims...)
function fit!(T::AbstractNormalization, X::AbstractArray; dims=())
    T(dims)
    dims = isnothing(T.dims) ? (1:ndims(X)) : T.dims
    psz = size(X) |> collect
    psz[[dims...]] .= 1
    T.p = reshape.(map.(T.𝑝, (Slices(X, dims...),)), psz...)
end
fit(T::AbstractNormalization, X::AbstractArray; kw...)=(T=deepcopy(T); fit!(T, X; kw...); T)
fit(𝒯::Type{<:AbstractNormalization}, X::AbstractArray; dims=nothing) = (T = 𝒯(dims); fit!(T, X); T)

function normalize!(X::AbstractArray, T::AbstractNormalization)
    isnothing(T.p) && fit!(T, X)
    mapdims!(T.𝑓, X, T.p...; T.dims)
end
normalize!(X::AbstractArray, 𝒯::Type{<:AbstractNormalization}; dims=nothing) = normalize!(X, fit(𝒯, X; dims))
normalize(X::AbstractArray, T::Union{AbstractNormalization, Type{<:AbstractNormalization}}; kwargs...) = (Y=copy(X); normalize!(Y, T; kwargs...); Y)

function denormalize!(X::AbstractArray, T::AbstractNormalization)
    isnothing(T.p) && error("Cannot denormalize with an unfit normalization")
    mapdims!(T.𝑓⁻¹, X, T.p...; T.dims)
end
denormalize(X::AbstractArray, args...) = (Y=copy(X); denormalize!(Y, args...); Y)

"""
Map the function `f` over the `dims` of all of the arguments. `f` should accept the same number of arguments as there are variables in `x...`. The first element of `x` is the considered as the reference array, and all other arguments must have sizes consistent with the reference array, or equal to 1.
"""
function mapdims!(f, x...; dims)
    n = ndims(x[1])
    isnothing(dims) && (dims = 1:n)
    dims = sort([dims...])
    @assert max(dims...) <= n
    @assert unique(dims) == dims
    length(dims) == n && return f(x...) # Shortcut for global normalisation
    negdims = Base._negdims(n, dims)
    @assert all(all(size.(x[2:end], i) .== 1) for i ∈ dims)
    @assert all(all(size(x[1], i) .== size.(x, i)) for i ∈ negdims)
    idxs = Base.compute_itspace(x[1], (negdims...,)|>Val)
    Threads.@threads for i ∈ idxs # map(f!, Slices.(x, dims...)...)
        selectslice = x -> view(x, i...)
        f(selectslice.(x)...)
    end
end

end


n, p = 500, 5
((X, y, θ), (X1, y1, θ1)) = generate_continuous(n, p; seed=525)

# adding some outliers (both positive and negative)
y1a = outlify(y1, 0.1)

@testset "QuantileReg" begin
    δ = 0.5 # effectively LAD regression
    λ = 1.0
    rr = QuantileRegression(δ, lambda=λ, penalize_intercept=true,
                               scale_penalty_with_samples = false)
    J = objective(rr, X, y1a)
    o = RobustLoss(Quantile(δ)) + λ * L2Penalty()
    @test J(θ1) ≈ o(y1a, X1*θ1, θ1)
    ls = LinearRegression()
    θ_ls    = fit(ls, X, y1a)
    θ_lbfgs = fit(rr, X, y1a, solver=LBFGS())
    θ_iwls  = fit(rr, X, y1a, solver=IWLSCG())
    @test isapprox(J(θ1),      412.20773, rtol=1e-5)
    @test isapprox(J(θ_ls),    508.02443, rtol=1e-5)  # LS is crap bc outliers
    @test isapprox(J(θ_lbfgs), 411.98228, rtol=1e-5)
    @test isapprox(J(θ_iwls),  411.98,    rtol=1e-4)

    # NOTE: newton and newton-cg not available because ϕ = 0 identically
    # will throw an error if called.
    @test_throws ErrorException fit(rr, X, y1, solver=Newton())
    @test_throws ErrorException fit(rr, X, y1, solver=NewtonCG())

    # don't penalize intercept
    rr = QuantileRegression(δ, lambda=λ, scale_penalty_with_samples = false)
    J = objective(rr, X, y1a)
    ls = LinearRegression()
    θ_ls    = fit(ls, X, y1a)
    θ_lbfgs = fit(rr, X, y1a, solver=LBFGS())
    θ_iwls  = fit(rr, X, y1a, solver=IWLSCG())
    @test isapprox(J(θ1),      412.18594, rtol=1e-5)
    @test isapprox(J(θ_ls),    508.00993, rtol=1e-5)  # note that LS is crap due to outliers
    @test isapprox(J(θ_lbfgs), 411.95990, rtol=1e-5)
    @test isapprox(J(θ_iwls),  411.98,    rtol=1e-4)

    if DO_COMPARISONS
        # Compare with R's QuantReg package
        # NOTE: QuantReg doesn't allow for penalties so re-fitting with λ=0
        rr = QuantileRegression(δ, lambda=0)
        J  = objective(rr, X, y1a)
        θ_ls     = fit(LinearRegression(), X, y1a)
        θ_lbfgs  = fit(rr, X, y1a, solver=LBFGS())
        θ_iwls   = fit(rr, X, y1a, solver=IWLSCG())
        θ_qr_br  = rcopy(QUANTREG.rq_fit_br(X1, y1a))[:coefficients]
        θ_qr_fnb = rcopy(QUANTREG.rq_fit_fnb(X1, y1a))[:coefficients]
        # NOTE: we take θ_qr_br as reference point
        @test isapprox(J(θ_ls), 505.45286,  rtol=1e-5)
        @test J(θ_qr_br) ≈      409.570777 # <- ref value
        # Their IP algorithm essentially gives the same answer
        @test (J(θ_qr_fnb) - J(θ_qr_br)) ≤ 1e-10
        # Our algorithms are close enough
        @test isapprox(J(θ_lbfgs), 409.57154, rtol=1e-5)
        @test isapprox(J(θ_iwls),  409.59,    rtol=1e-4)
    end
end

###########################
## With Sparsity penalty ##
###########################

n, p = 500, 100
((X, y, θ), (X1, y1, θ1)) = generate_continuous(n, p;  seed=51112, sparse=0.1)
# pepper with outliers
y1a  = outlify(y1, 0.1)

@testset "LAD+L1" begin
    λ = 5.0
    γ = 10.0
    rr = LADRegression(λ, γ; penalize_intercept=true, scale_penalty_with_samples = false)
    J  = objective(rr, X, y1a)
    θ_ls    = X1 \ y1a
    θ_fista = fit(rr, X, y1a, solver=FISTA())
    θ_ista  = fit(rr, X, y1a, solver=ISTA())
    @test isapprox(J(θ_ls),    1058.6737, rtol=1e-5)
    @test isapprox(J(θ_fista), 454.23322, rtol=1e-5)
    @test isapprox(J(θ_ista),  454.27774, rtol=1e-5)
    @test nnz(θ_fista) == 43
    @test nnz(θ_ista)  == 43

    if DO_COMPARISONS
        # Compare with R's QuantReg package
        # NOTE: QuantReg doesn't apply the penalty on the intercept
        rr = LADRegression(5.0; penalty=:l1, scale_penalty_with_samples = false)
        J  = objective(rr, X, y1a)
        θ_ls       = X1 \ y1a
        θ_fista    = fit(rr, X, y1a, solver=FISTA())
        θ_ista     = fit(rr, X, y1a, solver=ISTA())
        θ_qr_lasso = rcopy(QUANTREG.rq_fit_lasso(X1, y1a))[:coefficients]
        @test isapprox(J(θ_ls),       888.3748, rtol=1e-5)
        @test isapprox(J(θ_qr_lasso), 425.5264, rtol=1e-5)
        # Our algorithms are close enough
        @test isapprox(J(θ_fista),    425.0526, rtol=1e-5)
        @test isapprox(J(θ_ista),     425.4113, rtol=1e-5)
        # in this case we do a fair bit better
        @test nnz(θ_qr_lasso) == 101
        @test nnz(θ_fista)    == 88
        @test nnz(θ_ista)     == 82
        # in this case fista is best
        @test J(θ_fista) < J(θ_ista) < J(θ_qr_lasso)
    end
end


using Parameters, StatsBase

"""
決定木関連
"""

function branch(x::Matrix,y::Vector,S,rf=0;loss_f=TSS)
    rf==0 && (m=size(x)[2])
    size(x)[2]==0 && return [0,0,0,0,0,0,0]
    best_score = Inf; i_1=1; j_1=1;
    left_1=[];right_1=[]
    left_score_1=Inf; right_score_1=Inf
    for j in 1:m
        for i in S
            left=[k for k in S if x[k,j]<x[i,j]]
            right=setdiff(S,left)
            left_score=loss_f(y[left]); right_score=loss_f(y[right])
            score=left_score + right_score
            if score<best_score
                best_score = score
                i_1=i; j_1=j
                left_1=left;right_1=right
                left_score_1=left_score;right_score_1=right_score
            end
        end
    end
    return i_1,j_1,left_1,right_1,best_score,left_score_1,right_score_1
end

mutable struct Stack
    parent::Union{Int,Vector,AbstractRange} #親ノード番号
    set::Union{Vector,AbstractRange} #教師データの行
    score::Number # 損失関数の値
end

mutable struct Node
    parent::Union{Int,Vector,AbstractRange} #親ノード番号
    j::Int #列番号
    th::Number #閾値
    set::Union{Vector,AbstractRange} #そのノードでの教師データの集合
    left::Int # 分岐後の列番号(左)
    right::Int # 分岐後の列番号(右)
    center::Number # yの平均値
end

function decisiontree(x::Matrix,y::Vector;α=0.0,n_min=1,rf=0,loss_f=TSS)
    rf==0 && (m=size(x)[2])
    stack = [Stack(0,1:size(x)[1],loss_f(y))]
    node=Node[]
    k=0 #Node番号
    while length(stack)>0
        popped=pop!(stack)
        k+=1　#Node番号を増やす。
        i,j,left,right,score,left_score,right_score = branch(x,y,popped.set,rf,;loss_f=loss_f)
        if popped.score-score<α || length(popped.set)<n_min || length(left)==0 || length(right)==0
            push!(node,Node(popped.parent,-1,0,popped.set,0,0,0))
        else
            push!(node,Node(popped.parent,j,x[i,j],popped.set,0,0,0))
            push!(stack,Stack(k,right,right_score))
            push!(stack,Stack(k,left,left_score))
        end
    end
    for h in k:-1:2
        pa=node[h].parent
        if node[pa].right==0
            node[pa].right=h
        else
            node[pa].left = h
        end
    end

    g = loss_f==TSS ? mean : mode_max

    for h in 1:k
        node[h].center =  node[h].j==-1 ? g(y[node[h].set]) : 0
    end
    return node
end

function dtvalue(u,node)
    r=1
    while node[r].j !=-1
        if u[node[r].j] < node[r].th
            r = node[r].left
        else
            r = node[r].right
        end
    end
    return node[r].center
end

function freq(y)
    S = Set(y) |> collect |> sort
    [count(==(i),y) for i in S]
end

mode(y) = length(y)==0 ? 0 : maximum(freq(y))

mis_match(y) = length(y) - mode(y)

function gini(y)
    n = length(y)
    n==0 && (return 0)
    fr = freq(y)
    sum(x->x*(n-x), fr)
end

function entropy(y)
    n = length(y)
    n==0 && (return 0)
    fr = freq(y)
    sum(x->-x*log(x/n), fr)
end

function mode_max(y)
    length(y)==0 && (return -Inf)
    counts(y) |> argmax
end


"""
### Nuclear attraction term
"""
function Aterm(i::Int64,r::Int64,u::Int64,l1::Int64,l2::Int64,ax::Float64,bx::Float64,cx::Float64,gamma::Float64)
    term1 = (-1)^(i+u)*binomial_prefactor(i,l1,l2,ax,bx)
    term2 = factorial(i)*cx^(i-2r-2u)
    term3 = (1/4/gamma)^(r+u)/factorial(r)/factorial(u)/factorial(i-2r-2u)
    return term1*term2*term3
end

function Aarray(l1::Int64,l2::Int64,a::Float64,b::Float64,c::Float64,g::Float64)
    Imax = l1+l2+1
    A = zeros(Float64,Imax)
    for i in 0:(Imax-1)
        for r in 0:div(i,2)
            for u in 0:div(i-2r,2)
                I = i-2r-u+1
                A[I] += Aterm(i,r,u,l1,l2,a,b,c,g)
            end
        end
    end
    return A
end

function nuclear_attraction(aexpn::Float64,ax::Float64,ay::Float64,az::Float64,
                            aI::Int64,aJ::Int64,aK::Int64,
                            bexpn::Float64,bx::Float64,by::Float64,bz::Float64,
                            bI::Int64,bJ::Int64,bK::Int64,
                            cx::Float64,cy::Float64,cz::Float64)
    px,py,pz = gaussian_product_center(aexpn,ax,ay,az,bexpn,bx,by,bz)
    gamma = aexpn+bexpn
    rab2 = dist2(ax-bx,ay-by,az-bz)
    rcp2 = dist2(cx-px,cy-py,cz-pz)
    Ax = Aarray(aI,bI,px-ax,px-bx,px-cx,gamma)
    Ay = Aarray(aJ,bJ,py-ay,py-by,py-cy,gamma)
    Az = Aarray(aK,bK,pz-az,pz-bz,pz-cz,gamma)
    total = 0.0
    for I in 0:(aI+bI)
        for J in 0:(aJ+bJ)
            for K in 0:(aK+bK)
                total += Ax[I+1]*Ay[J+1]*Az[K+1]*Fgamma(I+J+K,rcp2*gamma)
            end
        end
    end
    val=-2pi*exp(-aexpn*bexpn*rab2/gamma)*total/gamma
    #println(val)
    #println((Ax,Ay,Az,rcp2*gamma,Fgamma(0,rcp2*gamma)))
    return val
end

function nuclear_attraction(a::PGBF,b::PGBF,cx::Float64,cy::Float64,cz::Float64)
    return a.norm*b.norm*nuclear_attraction(a.expn,a.x,a.y,a.z,a.I,a.J,a.K,
                                            b.expn,b.x,b.y,b.z,b.I,b.J,b.K,cx,cy,cz)
end
nuclear_attraction(a::PGBF,b::PGBF,c::Atom) = c.atno*nuclear_attraction(a,b,c.x,c.y,c.z)
nuclear_attraction(a::PGBF,b::PGBF,m::Molecule) = sum([nuclear_attraction(a,b,c) for c in m.atomlist])

# Need a nested scope to squeeze this into the contract function
function nuclear_attraction(a::CGBF,b::CGBF,cx::Float64,cy::Float64,cz::Float64)
    na(a,b) = nuclear_attraction(a,b,cx,cy,cz)
    contract(na,a,b)
end
function nuclear_attraction(a::CGBF,b::CGBF,c::Atom)
    na(a,b) = nuclear_attraction(a,b,c)
    contract(na,a,b)
end
function nuclear_attraction(a::CGBF,b::CGBF,m::Molecule)
    na(a,b) = nuclear_attraction(a,b,m)
    contract(na,a,b)
end

function test_a_terms()
    @testset "a_terms" begin
    @test Aterm(0,0,0,0,0,0.,0.,0.,0.) == 1.0
    @test Aarray(0,0,0.,0.,0.,1.) == [1.0]
    @test Aarray(0,1,1.,1.,1.,1.) == [1.0, -1.0]
    @test Aarray(1,1,1.,1.,1.,1.) == [1.5, -2.5, 1.0]
    @test Aterm(0,0,0,0,0,0.,0.,0.,1.) == 1.0
    @test Aterm(0,0,0,0,1,1.,1.,1.,1.) == 1.0
    @test Aterm(1,0,0,0,1,1.,1.,1.,1.) == -1.0
    @test Aterm(0,0,0,1,1,1.,1.,1.,1.) == 1.0
    @test Aterm(1,0,0,1,1,1.,1.,1.,1.) == -2.0
    @test Aterm(2,0,0,1,1,1.,1.,1.,1.) == 1.0
    @test Aterm(2,0,1,1,1,1.,1.,1.,1.) == -0.5
    @test Aterm(2,1,0,1,1,1.,1.,1.,1.) == 0.5
    end #testset
end

function test_na()
    @testset "test_na" begin
    s = pgbf(1.0)
    c = cgbf(0.0,0.0,0.0)
    push!(c,1,1)
    @test amplitude(c,0,0,0) ≈ 0.71270547
    @test nuclear_attraction(s,s,0.,0.,0.) ≈ -1.59576912
    @test nuclear_attraction(c,c,0.,0.,0.) ≈ -1.59576912
    end #testset
end
#todo make into a test
function test_one()
    s1 = pgbf(1)
    s2 = pgbf(1,0,1,0)
    x=y=0.
    println("S: $(overlap(s1,s2))")
    println("T: $(kinetic(s1,s2))")
    for z in range(0,stop=1,length=5)
        println("V: $z $(nuclear_attraction(s1,s2,x,y,z))")
    end
end

function test_na2()
    @testset "test_na2" begin
    li,h = lih.atomlist
    bfs = build_basis(lih)
    s1,s2,x,y,z,h1s = bfs.bfs
    @test nuclear_attraction(s1,s1,lih) ≈ -8.307532656
    end #testset
end


using Test
using OpenQuantumSystems
using Random, SparseArrays, LinearAlgebra, StableRNGs

@testset "operators dense" begin

    Random.seed!(StableRNG(0), 1)

    D(op1::AbstractOperator, op2::AbstractOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))
    D(x1::StateVector, x2::StateVector) = norm(x2 - x1)
    D(op1::Array, op2::Array) = abs(norm(op1 - op2))

    b = GenericBasis([3])
    one_op = OneDenseOperator(b)
    one_op_ref = [
        1.0+0.0im 0.0+0.0im 0.0+0.0im
        0.0+0.0im 1.0+0.0im 0.0+0.0im
        0.0+0.0im 0.0+0.0im 1.0+0.0im
    ]
    @test 1e-12 > D(one_op.data, one_op_ref)

    annihilation_op = AnnihilationOperator(b)
    annihilation_op_ref = [0.0 1.0 0.0; 0.0 0.0 1.4142135623730951; 0.0 0.0 0.0]
    @test 1e-12 > D(annihilation_op.data, annihilation_op_ref)

    creation_op = CreationOperator(b)
    creation_op_ref = [0.0 0.0 0.0; 1.0 0.0 0.0; 0.0 1.4142135623730951 0.0]
    @test 1e-12 > D(creation_op.data, creation_op_ref)

    position_op = PositionOperator(b)
    position_op_ref = [0.0 0.7071067811865475 0.0; 0.7071067811865475 0.0 1.0; 0.0 1.0 0.0]
    @test 1e-12 > D(position_op.data, position_op_ref)

    momentum_op = MomentumOperator(b)
    momentum_op_ref = ComplexF64[0.0 + 0.0im -0.0 - 0.7071067811865475im 0.0 + 0.0im; 0.0 + 0.7071067811865475im 0.0 + 0.0im -0.0 - 1.0im; 0.0 + 0.0im 0.0 + 1.0im 0.0 + 0.0im]
    @test 1e-12 > D(momentum_op.data, momentum_op_ref)

    shift_op = ShiftOperator(b, 1.0)
    shift_op_ref = [0.779728662995649 + 0.0im -0.5431245605675081 + 0.0im 0.3115107121936076 + 0.0im; 0.5431245605675084 + 0.0im 0.33918598898694735 + 0.0im -0.7680941196124976 + 0.0im; 0.31151071219360765 + 0.0im 0.7680941196124975 + 0.0im 0.559457325991298 + 0.0im]
    @test 1e-12 > D(shift_op.data, shift_op_ref)

end # testset


export batched_mul, ⊠,  batched_vec
export batched_transpose, batched_adjoint

include("./batchedadjtrans.jl")

using LinearAlgebra: BlasFloat, Transpose, Adjoint, AdjOrTransAbsMat

_unbatch(A) = A
_unbatch(A::BatchedAdjOrTrans) = parent(A)

"""
    batched_mul(A, B) -> C
    A ⊠ B  # \\boxtimes

Batched matrix multiplication. Result has `C[:,:,k] == A[:,:,k] * B[:,:,k]` for all `k`.
If `size(B,3) == 1` then instead `C[:,:,k] == A[:,:,k] * B[:,:,1]`, and similarly for `A`.

To transpose each matrix, apply `batched_transpose` to the array,
or `batched_adjoint` for conjugate-transpose:

```jldoctest
julia> A, B = randn(2,5,17), randn(5,9,17);

julia> A ⊠ B |> size
(2, 9, 17)

julia> batched_adjoint(A) |> size
(5, 2, 17)

julia> batched_mul(A, batched_adjoint(randn(9,5,17))) |> size
(2, 9, 17)

julia> A ⊠ randn(5,9,1) |> size
(2, 9, 17)

julia> batched_transpose(A) == PermutedDimsArray(A, (2,1,3))
true
```

The equivalent `PermutedDimsArray` may be used in place of `batched_transpose`.
Other permutations are also handled by BLAS,
provided that the batch index `k` is not the first dimension of the underlying array.
Thus `PermutedDimsArray(::Array, (1,3,2))` and `PermutedDimsArray(::Array, (3,1,2))` are fine.

However, `A = PermutedDimsArray(::Array, (3,2,1))` is not acceptable to BLAS,
since the batch dimension is the contiguous one: `stride(A,3) == 1`.
This will be copied, as doing so is faster than `batched_mul_generic!`.

Both this `copy` and `batched_mul_generic!` produce `@debug` messages,
and setting for instance `ENV["JULIA_DEBUG"] = NNlib` will display them.
"""
function batched_mul(A::AbstractArray{T1, 3}, B::AbstractArray{T2, 3}) where {T1, T2}
    size(A, 3) == size(B, 3) || size(A, 3) == 1 || size(B, 3) == 1 ||
        throw(DimensionMismatch("batch size mismatch: A != B"))
    _batched_mul(storage_typejoin(A, B), A, B)
end

const ⊠ = batched_mul

function _batched_mul(::Type, A, B)
    T = promote_type(eltype(A), eltype(B))
    C = similar(A, T, (size(A, 1), size(B, 2), max(size(A, 3), size(B, 3))))
    batched_mul!(C, A, B)
    C
end
function _batched_mul(::Type{<:DenseArray{T}}, A, B) where {T<:BlasFloat}
    C = similar(A, T, (size(A, 1), size(B, 2), max(size(A, 3), size(B, 3))))
    batched_mul!(C, _copy_if_faster(A), _copy_if_faster(B))
    C
end

function _copy_if_faster(X::AbstractArray{<:Number, 3})
    is_strided(X) || return X
    if Base.stride(X, 3) == 1 && Base.stride(X, 1) != 1
        @debug "copying to avoid batched_mul_generic!" typeof(X) size(X) strides(X)
        return copy(X)
    end
    X
end
function _copy_if_faster(X::BatchedAdjoint{<:Complex})
    Xbase = _unbatch(X)
    is_strided(Xbase) || return X
    if Base.stride(Xbase, 1) != 1
        @debug "copying to avoid batched_mul_generic!" typeof(X) size(X) strides(_unbatch(X))
        return copy(X) # or batched_adjoint(copy(Xbase)), may be better on GPU?
    end
    X
end

# Gradient, allowing that size(A,3)==1 means it's "broadcasted" out to size(B,3)

function rrule(::typeof(batched_mul), A::AbstractArray{<:Any,3}, B::AbstractArray{<:Any,3})
    function batched_mul_pullback(Δ)
        Athunk = @thunk begin
            tmp = batched_mul(Δ, batched_adjoint(B))
            size(A,3) == 1 ? sum(tmp, dims=3) : tmp
        end
        Bthunk = @thunk begin
            tmp = batched_mul(batched_adjoint(A), Δ)
            size(B,3) == 1 ? sum(tmp, dims=3) : tmp
        end
        return (NoTangent(), Athunk, Bthunk)
    end
    batched_mul(A, B), batched_mul_pullback
end

"""
    batched_mul(A::Array{T,3}, B::Matrix)
    batched_mul(A::Matrix, B::Array{T,3})
    A ⊠ B

This is always matrix-matrix multiplication, but
either `A` or `B` may lack a batch index.

* When `B` is a matrix, result has `C[:,:,k] == A[:,:,k] * B[:,:]` for all `k`.

* When `A` is a matrix, then `C[:,:,k] == A[:,:] * B[:,:,k]`.
  This can also be done by reshaping and calling `*`,
  for instance `A ⊡ B` using TensorCore.jl, but is implemented here using
  `batched_gemm` instead of `gemm`.

```jldoctest
julia> randn(16,8,32) ⊠ randn(8,4) |> size
(16, 4, 32)

julia> randn(16,8,32) ⊠ randn(8,4,1) |> size  # equivalent
(16, 4, 32)

julia> randn(16,8) ⊠ randn(8,4,32) |> size
(16, 4, 32)
```

See also `batched_vec` to regard `B` as a batch of vectors, `A[:,:,k] * B[:,k]`.
"""
batched_mul(A::AbstractArray{T,3} where T, B::AbstractMatrix) = _semi_batched_mul(A,B)

# Simplify signature of batched_mul by hiding dispatch on Adjoint etc:

_semi_batched_mul(A::AbstractArray{<:Any,3}, B::AbstractMatrix) =
    batched_mul(A, reshape(B, size(B)..., 1))

_semi_batched_mul(A::AbstractArray{<:Any,3}, B::Adjoint{<:Number,<:AbstractMatrix}) =
    batched_mul(A, batched_adjoint(reshape(parent(B), size(parent(B))..., 1)))

_semi_batched_mul(A::AbstractArray{<:Any,3}, B::Transpose{<:Number,<:AbstractMatrix}) =
    batched_mul(A, batched_transpose(reshape(parent(B), size(parent(B))..., 1)))

batched_mul(A::AbstractMatrix, B::AbstractArray{T,3} where T) = _semi_batched_mul(A,B)

_semi_batched_mul(A::AbstractMatrix, B::AbstractArray{<:Any,3}) =
    batched_mul(reshape(A, size(A)..., 1), B)

_semi_batched_mul(A::Adjoint{<:Number,<:AbstractMatrix}, B::AbstractArray{<:Any,3}) =
    batched_mul(batched_adjoint(reshape(parent(A), size(parent(A))..., 1)), B)

_semi_batched_mul(A::Transpose{<:Number,<:AbstractMatrix}, B::AbstractArray{<:Any,3}) =
    batched_mul(batched_transpose(reshape(parent(A), size(parent(A))..., 1)), B)

"""
    batched_vec(A::Array{T,3}, B::Matrix)
    batched_vec(A::Array{T,3}, b::Vector)

Batched matrix-vector multiplication:
the result has `C[:,:,k] == A[:,:,k] * B[:,k]` for all `k`,
or else `C[:,:,k] == A[:,:,k] * b` for `b::Vector`.

With the same argument types, `batched_mul(A, B)` would regard `B` as
a fixed matrix, not a batch of vectors. Both reshape and then
call `batched_mul(::Array{T,3}, ::Array{T,3})`.

```jldoctest
julia> A, B, b = randn(16,8,32), randn(8,32), randn(8);

julia> batched_vec(A,B) |> size
(16, 32)

julia> batched_vec(A,b) |> size
(16, 32)
```
"""
function batched_vec(A::AbstractArray{T,3} where T, B::AbstractMatrix)
    # If B is transposed, then stride=1 is the batch dim, so we will end up copying anyway:
    if B isa AdjOrTransAbsMat{<:BlasFloat, <:StridedMatrix}
        return batched_vec(A, copy(B))
    end
    reshape(batched_mul(A, reshape(B, size(B,1), 1, size(B,2))), size(A,1), size(A,3))
end

batched_vec(A::AbstractArray{T,3} where T, b::AbstractVector) =
    reshape(batched_mul(A, reshape(b, length(b), 1, 1)), size(A,1), size(A,3))

"""
    batched_mul!(C, A, B) -> C
    batched_mul!(C, A, B, α=1, β=0)

In-place batched matrix multiplication, equivalent to
`mul!(C[:,:,k], A[:,:,k], B[:,:,k], α, β)` for all `k`.
If `size(B,3) == 1` then every batch uses `B[:,:,1]` instead.

This will call `batched_gemm!` whenever possible. For real arrays this means that,
for `X ∈ [A,B,C]`, either `strides(X,1)==1` or `strides(X,2)==1`, the latter may
be caused by `batched_transpose` or by for instance `PermutedDimsArray(::Array, (3,1,2))`.
Unlike `batched_mul` this will never make a copy.

For complex arrays, the wrapper made by `batched_adjoint` must be outermost to be seen.
In this case the strided accepted by BLAS are more restricted, if `stride(C,1)==1` then
only `stride(AorB::BatchedAdjoint,2) == 1` is accepted.
"""
function batched_mul!(C::AbstractArray{T,3}, A::AbstractArray{<:Any,3}, B::AbstractArray{<:Any,3},
        α::Number=one(T), β::Number=zero(T)) where {T}
    _batched_mul!(storage_typejoin(C,A,B), C, A, B, α, β)
    C
end

_batched_mul!(::Type, C, A, B, α::Number, β::Number) = batched_mul_generic!(C, A, B, α, β)

_batched_mul!(::Type{DT}, C, A, B, α::Number, β::Number) where {DT<:DenseArray{T}} where {T<:BlasFloat} =
    _batched_try_gemm!(DT, C, A, B, α, β)

function _batched_try_gemm!(::Type{DT}, C, A, B, α::Number, β::Number) where {DT<:DenseArray{T}} where {T<:BlasFloat}

    alpha, beta = promote(α, β, zero(T))
    alpha isa T && beta isa T || return batched_mul_generic!(C, A, B, α, β)

    are_strided(_unbatch(A), _unbatch(B)) || return batched_mul_generic!(C, A, B, α, β)
    C isa StridedArray || return batched_mul_generic!(C, A, B, α, β)

    blasA, transA = if A isa BatchedAdjoint && T <: Complex
        Base.stride(parent(A),1) == 1 || return batched_mul_generic!(C, A, B, α, β)
        parent(A), 'C'
    elseif Base.stride(A,2) == 1 && size(A,1) > 1
        batched_transpose(A), 'T'
    elseif Base.stride(A,1) == 1
        A, 'N'
    elseif Base.stride(A,2) == 1  # This is awful, but exhaustively tested. Issues 268, 282.
        batched_transpose(A), 'T'
    else
        return batched_mul_generic!(C, A, B, α, β)
    end

    blasB, transB = if B isa BatchedAdjoint && T <: Complex
        Base.stride(parent(B),1) == 1 || return batched_mul_generic!(C, A, B, α, β)
        parent(B), 'C'
    elseif Base.stride(B,2) == 1 && size(B,1) > 1
        batched_transpose(B), 'T'
    elseif Base.stride(B,1) == 1
        B, 'N'
    elseif Base.stride(B,2) == 1
        batched_transpose(B), 'T'
    else
        return batched_mul_generic!(C, A, B, α, β)
    end

    _batched_gemm!(DT, transA, transB, alpha, blasA, blasB, beta, C)
    C
end

_batched_gemm!(::Type{<:Array}, transA::Char, transB::Char, α::Number, A, B, β::Number, C) =
    batched_gemm!(transA, transB, α, A, B, β, C)

_BATCHED_LIST = [
    (:(AbstractArray{<:Any, 3}), :identity),
    (:BatchedTranspose,          :transpose),
    (:BatchedAdjoint,            :adjoint),
]
for (TA, fA) in _BATCHED_LIST, (TB, fB) in _BATCHED_LIST

    @eval function batched_mul_generic!(C::AbstractArray{T, 3}, A::$TA, B::$TB,
            α::Number=one(T), β::Number=zero(T)) where {T}

        size(A, 3) == size(C, 3) || size(A, 3) == 1 || throw(DimensionMismatch("batch size mismatch: A != C"))
        size(B, 3) == size(C, 3) || size(B, 3) == 1 || throw(DimensionMismatch("batch size mismatch: B != C"))
        @debug "calling fallback method for batched_mul!" typeof(A) size(A) typeof(B) size(B) typeof(C)

        Abase, Bbase = _unbatch(A), _unbatch(B)
        sA, oA = size(A,3) == 1 ? (0,1) : (1,0)
        sB, oB = size(B,3) == 1 ? (0,1) : (1,0)

        @inbounds for k in 1:size(C,3)
            @views mul!(C[:,:,k], $fA(Abase[:,:,k*sA+oA]), $fB(Bbase[:,:,k*sB+oB]), α, β)
        end
        C
    end

end

"""
    storage_type(A) -> Type

Removes all wrappers to return the `Array` or `CuArray` (or whatever) type within.
```
julia> view(reshape(ones(10)',2,5),:, 3:4) |> storage_type
Array{Float64,1}

julia> reshape(sparse(rand(10)), 5,2) |> storage_type
SparseVector{Float64,Int64}
```
"""
function storage_type(A::AbstractArray)
    P = parent(A)
    typeof(A) === typeof(P) ? typeof(A) : storage_type(P)
end
storage_type(A) = typeof(A)

"""
    storage_typejoin(A, B, C, ...) -> Type

Reduces with `Base.promote_typejoin`, in order that this conveys useful information
for dispatching to BLAS. It does not tell you what container to allocate:
```
julia> storage_typejoin(rand(2), rand(Float32, 2))
Array{T,1} where T

julia> eltype(ans) <: LinearAlgebra.BlasFloat
false

julia> storage_typejoin(rand(2), rand(2,3), rand(2,3,4))
Array{Float64,N} where N
```
"""
storage_typejoin(A, Bs...) = Base.promote_typejoin(storage_type(A), storage_typejoin(Bs...))
storage_typejoin(A) = storage_type(A)

"""
    is_strided(A::AbstractArray) -> Bool

This generalises `A isa StridedArray` to treat wrappers like `A::PermutedDimsArray`,
for which it returns `is_strided(parent(A))`.

It returns `true` for `CuArray`s, and `PermutedDimsArray`s of those.

Other wrappers (defined outside Base, LinearAlgebra) are assumed not to break
strided-ness, and hence also return `is_strided(parent(A))`.
This correctly handles things like `NamedDimsArray` wihch don't alter indexing.
However, it's a little pessimistic in that e.g. a `view` of such a container will return
`false`, even in cases where the same `view` of `parent(A)` would be a `StridedArray`.
"""
is_strided(A::StridedArray) = true
is_strided(A) = false
function is_strided(A::AbstractArray)
    M = parentmodule(typeof(A))
    if parent(A) === A # SparseMatrix, StaticArray, etc
        false
    elseif M === Base || M === Core || M ===LinearAlgebra
        # bad reshapes, etc, plus Diagonal, UpperTriangular, etc.
        false
    else
        is_strided(parent(A)) # PermutedDimsArray, NamedDimsArray
    end
end

is_strided(A::BatchedAdjoint) = eltype(A) <: Real && is_strided(parent(A))
is_strided(A::BatchedTranspose) = is_strided(parent(A))

is_strided(A::LinearAlgebra.Transpose) = is_strided(parent(A))
is_strided(A::LinearAlgebra.Adjoint) = eltype(A) <: Real && is_strided(parent(A))
# This needs Compat 3.14, for any Julia < 1.6

are_strided(As...) = mapfoldl(is_strided, &, As; init=true)


@testset "Schur Complement" begin
    T = Float64
    n = 11
    m = 16
    M = rand(n+m,n+m)
    S = ContactImplicitMPC.Schur(M, n=n, m=m)

    u = rand(T,n)
    v = rand(T,m)
    D = rand(m,m)
    ContactImplicitMPC.schur_factorize!(S, D)
    # @benchmark ContactImplicitMPC.schur_factorize!(S, D)
    ContactImplicitMPC.schur_solve!(S, u, v)
    # @benchmark ContactImplicitMPC.schur_solve!(S, u, v)
    M1 = [S.A S.B; S.C D]
    @test norm(M1*[S.x; S.y] - [u; v], Inf) < 1e-9
end

@testset "Schur Complement: ill-conditioning" begin
    T = Float64
    s = get_simulation("flamingo", "flat_2D_lc", "flat")
    ref_traj = deepcopy(get_trajectory(s.model, s.env,
        joinpath(module_dir(), "src/dynamics/flamingo/gaits/gait_forward_36_4.jld2"),
        load_type = :split_traj_alt))

    ix, iy1, iy2 = linearization_var_index(s.model, s.env)
    idyn, irst, ibil = linearization_term_index(s.model, s.env)
    nz = num_var(s.model, s.env)
    nθ = num_data(s.model)
    nx = length(ix)
    ny = length(iy1)

    r0 = zeros(T, nz)
    rz0 = zeros(T, nz, nz)
    rθ0 = zeros(T, nz, nθ)
    z0 = zeros(T, nz)
    θ0 = zeros(T, nθ)
    z0 .= max.(1e-6, ref_traj.z[10])
    θ0 .= ref_traj.θ[10]
    κ0 = 0.0
    s.res.r!(r0, z0, θ0, κ0)
    s.res.rz!(rz0, z0, θ0)
    s.res.rθ!(rθ0, z0, θ0)

    A = rz0[idyn, ix]
    B = rz0[idyn, iy1]
    C = rz0[irst, ix]
    D = rz0[irst, iy1] - rz0[irst, iy2] * Diagonal(z0[iy2] ./ z0[iy1])
    M = [A B ; C D]
    S = ContactImplicitMPC.Schur(M, n=nx, m=ny)

    u = r0[idyn]
    v = r0[irst] - r0[ibil] ./ z0[iy1]

    # @benchmark schur_factorize!(S, D)
    ContactImplicitMPC.schur_factorize!(S, D)
    # @benchmark ContactImplicitMPC.schur_solve!(S, u, v)
    ContactImplicitMPC.schur_solve!(S, u, v)
    M1 = [S.A S.B; S.C D]
    norm(M1*[S.x; S.y] - [u; v], Inf)
    @test norm(M1*[S.x; S.y] - [u; v], Inf) < 1e-7
end


# f(x)=A*X'
function FAXtranspose(A,X)
    if size(A)==(1,1)
        return (A[1].*X',nothing)
    elseif size(X)==(1,1)
        return (A.*X[1],nothing)
    else
        return (A*X',nothing)
    end
end

function FAXtranspose_inplace(value,auxvalue,A,X)
    if size(A)==(1,1)
        copy!(value,A[1]*X')
    elseif size(X)==(1,1)
        copy!(value,A*X[1])
    else
        BLAS.gemm!('N','T',1.0,A,X,0.0,value)
    end
end

function DAXtranspose(derivativeIDX,f_c,faux_c,grad_c,grad_n,A,X)
    if derivativeIDX==1
        if size(A)==(1,1)
            axpy!(1.0,[sum(X'.*grad_c)],grad_n)
        elseif size(X)==(1,1)
            axpy!(X[1],grad_c,grad_n)
        else
            BLAS.gemm!('N','N',1.0,grad_c,X,1.0,grad_n) # grad_c*X
        end
    elseif derivativeIDX==2
        if size(A)==(1,1)
            axpy!(A[1],grad_c',grad_n)
        elseif size(X)==(1,1)
            axpy!(1.0,[sum(A.*grad_c)],grad_n)
        else
            BLAS.gemm!('T','N',1.0,grad_c,A,1.0,grad_n) # grad_c'*A
        end
    end
end

if PROC=="GPU"
#    function FAXtranspose(A::CudaArray,X::CudaArray)
#        if size(A)==(1,1)
#            tmp=CudaArray(Float64,size(X)); gfill!(tmp,A)
#            return (CUBLAS.gemm('N','T',tmp,X),nothing)
#        elseif size(X)==(1,1)
#            tmp=CudaArray(Float64,size(A)); gfill!(tmp,X)
#            return (CUBLAS.gemm('N','N',A,tmp),nothing)
#        else
#            return (CUBLAS.gemm('N','T',A,X),[])
#        end
#    end

    function FAXtranspose_inplace(value::CudaArray,auxvalue,A::CudaArray,X::CudaArray)
        if size(A)==(1,1)
            CUBLAS.geam!('T','N',1.0,X,0.0,value,value) # nb argument converse of Base.scale!
            scale!(A,value) # nb argument converse of Base.scale!
        elseif size(X)==(1,1)
            copy!(value,A); scale!(X,value) # nb argument converse of Base.scale!
        else
            CUBLAS.gemm!('N','T',1.0,A,X,0.0,value)
        end
    end

    function DAXtranspose(derivativeIDX,f_c,faux_c,grad_c,grad_n,A::CudaArray,X::CudaArray)
        if derivativeIDX==1
            if size(A)==(1,1)
                tmp=CudaArray(Float64,size(X))
                CUBLAS.geam!('T','N',1.0,X,0.0,tmp,tmp) # tmp=X'
                vmult!(1.0,tmp,grad_c,tmp) # tmp=X'.*grad_c
                sum_update!(1.0,tmp,1.0,grad_n) # grad_n+=sum(X'.*grad_c)
                free(tmp)
            elseif size(X)==(1,1)
                alphaaxpy!(1.0,X,grad_c,grad_n)
            else
                CUBLAS.gemm!('N','N',1.0,grad_c,X,1.0,grad_n) # grad_c*X
            end
        elseif derivativeIDX==2
            if size(A)==(1,1)
                tmp=CudaArray(Float64,size(grad_c))
                CUBLAS.geam!('T','N',1.0,grad_c,0.0,tmp,tmp) # tmp=grad_c'
                alphaaxpy!(1.0,A,tmp,grad_n); free(tmp)
            elseif size(X)==(1,1)
                tmp=CudaArray(Float64,size(A))
                vmult!(1.0,A,grad_c,tmp)
                sum_update!(1.0,tmp,1.0,grad_n) # sum(A.*grad_c)
                free(tmp)
            else
                CUBLAS.gemm!('T','N',1.0,grad_c,A,1.0,grad_n) # grad_c'*A
            end
        end
    end
end

Derivative[FAXtranspose]=DAXtranspose
export FAXtranspose

Inplace[FAXtranspose]=FAXtranspose_inplace


function sumOfMult(a::Array{Int, 1}, left=1, right=length(a))
	if left == right
		return 0
	elseif right-left == 1
		return a[left] * a[right]
	end
	mid = (right+left) ÷ 2
	return sumOfMult(a, left, mid) + sumOfMult(a, mid, right)
end


abstract type Instruction end

struct Mask <: Instruction
    mask::AbstractString
end

struct Write <: Instruction
    address::Int64
    value::Int64
end

function read_data(lines)
    program = Instruction[]
    for line in lines
        first, v = split(line, " = ")
        if first == "mask"
            push!(program, Mask(v))
        else
            address = parse(Int64, first[5:end-1])
            value = parse(Int64, v)
            push!(program, Write(address, value))
        end
    end
    program
end

function initialization(program)
    memory = Dict{Int64, Int64}()
    mask = nothing
    for instruction in program
        if isa(instruction, Mask)
            mask = instruction.mask
        else
            s = bitstring(instruction.value)
            masked = join(map(((bit, m),) -> m == 'X' ? bit : m, zip(s[length(s)-35:end], mask)), "")
            masked_value = parse(Int64, masked; base=2)
            memory[instruction.address] = masked_value
        end
    end
    memory
end

function main()
    lines = readlines("../resources/input_14.txt")
    program = read_data(lines)
    memory = initialization(program)
    println(sum(values(memory)))
end

main()


"""
Mixed Matrix Exponential Function
"""

"""
mexpmix(f, Q, x; bounds = (0, Inf), transpose = :N, ufact = 1.01, eps = 1.0e-8, rmax = 500)
mexp(Q, x, dist, bounds = (minimum(dist), maximum(dist)), transpose = :N, ufact = 1.01, eps = 1.0e-8, rmax = 500)

Compute the probability vector for CTMC mixed with dist

int_bound[1]^bound[2] exp(tr(Q)*t) f(t) dt * x
int_bound[1]^bound[2] exp(tr(Q)*t) pdf(dist, t) dt * x

Parameters:
- Q: CTMC Kernel
- x: Array
- f: pdf of distribution
- dist: distribution (UnivariateDistribution)
- bounds: a tuple of domain of distribution
- transpose: forward or backward
- ufact: uniformization factor
- eps: tolerance error for Poisson p.m.f.
- rmax: The maximum number of uniformization steps

Return value:
- probability vector
"""

function mexpmix(f::Any, Q::AbstractMatrix{Tv}, x::ArrayT;
    bounds=(Tv(0.0), Tv(Inf)), transpose::Symbol=:N,
    ufact::Tv=Tv(1.01), eps::Tv=Tv(1.0e-8), rmax=500) where {Tv,ArrayT<:AbstractArray{Tv}}
    m, n = size(Q)
    @assert m == n
    de = deint(f, bounds[1], bounds[2])
    dt, maxt = itime(de.x)
    P, qv = unif(Q, ufact)
    right = rightbound(qv*maxt, eps)
    @assert right <= rmax "Time interval is too large. rmax should be changed: right = $right (rmax: $rmax)."
    prob = Vector{Tv}(undef, right+1)
    y0, y1 = copy(x), similar(x)
    result = zero(x)
    @inbounds for i in eachindex(dt)
        right = rightbound(qv*dt[i], eps)
        weight = poipmf!(qv*dt[i], prob; left=0, right=right)
        y1 .= Tv(0)
        unifstep!(transpose, P, prob, (0, right), weight, y0, y1)
        @axpy(de.w[i], y1, result)
        y0 .= y1
    end
    @scal(de.h, result)
    return result
end

function mexp(Q::AbstractMatrix{Tv}, x::ArrayT, dist::UnivariateDistribution;
    bounds=(minimum(dist), maximum(dist)), transpose::Symbol=:N,
    ufact::Tv=Tv(1.01), eps::Tv=Tv(1.0e-8), rmax=500) where {Tv,ArrayT<:AbstractArray{Tv}}
    mexpmix(Q, x, bounds=bounds, transpose=transpose, ufact=ufact, eps=eps, rmax=rmax) do x
        pdf(dist, x)
    end
end

"""
mexpcmix(f, Q, x; bounds = (0, Inf), transpose = :N, ufact = 1.01, eps = 1.0e-8, rmax = 500)
mexpc(Q, x, dist, bounds = (minimum(dist), maximum(dist)), transpose = :N, ufact = 1.01, eps = 1.0e-8, rmax = 500)

Compute the probability vector for CTMC and the cumulative value for time series.
Compute the probability vector for CTMC and the cumulative value which are mixed with dist

int exp(tr(Q)*t) * f(t) dt * x
int int_0^t exp(tr(Q)*u) * x du f(t) dt

Parameters:
- Q: CTMC Kernel
- x: Array
- f: pdf of distribution
- dist: distribution (UnivariateDistribution)
- bounds: a tuple of domain of distribution
- transpose: forward or backward
- ufact: uniformization factor
- eps: tolerance error for Poisson p.m.f.
- rmax: The maximum number of uniformization steps

Return value (tuple)
- probability vector
- cumulative value
"""

function mexpcmix(f::Any, Q::AbstractMatrix{Tv}, x::ArrayT;
    bounds=(Tv(0.0), Tv(Inf)), transpose::Symbol=:N,
    ufact::Tv=Tv(1.01), eps::Tv=Tv(1.0e-8), rmax=500) where {Tv,ArrayT<:AbstractArray{Tv}}
    m, n = size(Q)
    @assert m == n
    de = deint(f, bounds[1], bounds[2])
    dt, maxt = itime(de.x)
    P, qv = unif(Q, ufact)
    right = rightbound(qv*maxt, eps) + 1
    @assert right <= rmax "Time interval is too large. rmax should be changed: right = $right (rmax: $rmax)."
    prob = Vector{Tv}(undef, right+1)
    cprob = Vector{Tv}(undef, right+1)
    y0, y1 = copy(x), similar(x)
    cy = zero(x)
    tmp = similar(x)
    result, cresult = zero(x), zero(x)
    @inbounds for i in eachindex(dt)
        right = rightbound(qv*dt[i], eps) + 1
        weight = cpoipmf!(qv*dt[i], prob, cprob; left=0, right=right)
        tmp .= Tv(0)
        y1 .= Tv(0)
        cunifstep!(transpose, P, prob, cprob, (0, right), weight, qv*weight, y0, y1, tmp)
        cy .+= tmp
        @axpy(de.w[i], y1, result)
        @axpy(de.w[i], cy, cresult)
        y0 .= y1
    end
    @scal(de.h, result)
    @scal(de.h, cresult)
    return result, cresult
end

function mexpc(Q::AbstractMatrix{Tv}, x::ArrayT, dist::UnivariateDistribution;
    bounds = (minimum(dist), maximum(dist)), transpose::Symbol=:N,
    ufact::Tv=Tv(1.01), eps::Tv=Tv(1.0e-8), rmax=500) where {Tv,ArrayT<:AbstractArray{Tv}}
    mexpcmix(Q, x, bounds=bounds, transpose=transpose, ufact=ufact, eps=eps, rmax=rmax) do x
        pdf(dist, x)
    end
end