function h_layout_spring_adj(locs_r::Array{Float64,1}, angs::Array{Float64,1}, adj_matrix::Array{T,2}; C=2.0, MAXITER=100, INITTEMP=2.0) where T

    size(adj_matrix, 1) != size(adj_matrix, 2) && error("Adj. matrix must be square.")
    N = size(adj_matrix, 1)

    # The optimal distance bewteen vertices
    K = C * sqrt(4.0 / N)

    # Store forces and apply at end of iteration all at once
    force = zeros(N)

    function scaler(z, a, b)
        2.0*((z - a)/(b - a)) - 1.0
    end

    # Iterate MAXITER times
    @inbounds for iter = 1:MAXITER
        # Calculate forces
        for i = 1:N
            force_vec = 0.0
            for j = 1:N
                i == j && continue
                P = HPoint(locs_r[i],angs[i])
                Q = HPoint(locs_r[j],angs[j])
                d = dist(P,Q)
                if adj_matrix[i,j] != zero(eltype(adj_matrix)) || adj_matrix[j,i] != zero(eltype(adj_matrix))
                    # F = d^2 / K - K^2 / d
                    F_d = d / K - K^2 / d^2
                else
                    # Just repulsive
                    # F = -K^2 / d^
                    F_d = -K^2 / d^2
                end
                # d  /          sin θ = d_y/d = fy/F
                # F /| dy fy    -> fy = F*d_y/d
                #  / |          cos θ = d_x/d = fx/F
                # /---          -> fx = F*d_x/d
                # dx fx
                d_theta = abs(angs[i] - angs[j])
                force_vec += F_d*d * cos(d_theta)

            end
            force[i] = force_vec
        end
        # Cool down
        TEMP = INITTEMP / iter
        # Now apply them, but limit to temperature
        for i = 1:N
            force_mag  = force[i]
            scale      = min(force_mag, TEMP)/force_mag
            locs_r[i] += force[i] * scale
        end
        # Scale to unit circle each time
        min_r, max_r = minimum(locs_r), maximum(locs_r)
        locs_r = map(z -> scaler(z, min_r, max_r), locs_r)
    end

    return locs_r, angs
end


"""
Computes random samples of functions from aligned data using Gaussian model

    gauss_model(warp_data; n=1, sort_samples=false)
    :param warp_data: fdawarp type from srsf_align of aligned data
    :param n: number of samples
    :param sort_samples: sort samples

    Returns warp_data containing
    :return fs: random aligned functions
    :return gams: random warping functions
    :return ft: random functions
"""
function gauss_model(warp_data::fdawarp; n=1, sort_samples=false)
    fn = warp_data.fn
    timet = warp_data.time
    qn = warp_data.qn
    gam = warp_data.gam
    # Parameters
    eps1 = eps(Float64);
    binsize = mean(diff(timet));
    M = length(timet);

    # compute mean and covariance in q-domain
    mq_new = mean(qn,dims=2);
    mididx = round(Integer, M/2);
    m_new = sign.(fn[mididx,:]) .* sqrt.(abs.(fn[mididx,:]));
    mqn = vec([mq_new; mean(m_new)]);
    qn2 = vcat(qn, m_new');
    C = Statistics.covm(qn2, mean(qn2,dims=2), 2);

    tmp = cholesky(C,check=false)
    if (!issuccess(tmp))
        for i = 1:M+1
            C[i,i] = C[i,i] + 1e-7
        end
    end
    di = MvNormal(mqn, C);
    q_s = rand(di, n);

    # compute the correspondence to the orignal function domain
    fs = zeros(M, n);
    for k in 1:n
        fs[:, k] = cumtrapzmid(timet, q_s[1:M,k].*abs.(q_s[1:M,k]),
                               sign(q_s[M+1,k])*(q_s[M+1,k]^2), mididx);
    end
    fbar = mean(fn,dims=2)
    fsbar = mean(fs,dims=2)
    err = repeat(fbar-fsbar,1,n)
    fs = fs + err

    # random warping generation
    rgam = random_gamma(gam, n);
    gams = zeros(M, n);
    for k in 1:n
        gams[:,k] = invert_gamma(rgam[:, k]);
    end

    if sort_samples
        mx = maximum(fs,1);
        seq1 = sortperm(vec(mx));

        # compute the psi-function
        psi = zeros(size(rgam));
        for ii in 1:size(rgam, 2)
            fy = gradient(rgam[:, ii], binsize);
            psi[:,ii] = fy ./ sqrt(abs(fy) + eps1);
        end
        ip = zeros(n);
        len = zeros(n);

        for i in 1:n
            tmp = ones(1,M);
            psi_tmp = (psi[:,i]/M);
            ip[i] = (tmp * psi_tmp)[1];
            len[i] = acos(ip[i]);
        end

        seq2 = sortperm(len);

        # combine x-variability and y-variability
        ft = zeros(M, n);
        time0 = [0:(M-1)]./(M-1);
        for k in 1:n
            ft[:,k] = approx(time0, fs[:,seq1[k]], gams[:,seq2[k]]);
            tmp = isnan(ft[:, k]);
            while any(tmp)
                rgam2 = random_gamma(gam, 1);
                gam_tmp = invert_gamma(rgam2);
                ft[:, k] = approx(time0, fs[:,seq1[k]], gam_tmp);
            end
        end
    else
        # combine x-variability and y-variability
        ft = zeros(M, n);
        time0 = collect(0:(M-1))./(M-1);
        for k in 1:n
            ft[:,k] = approx(time0, fs[:,k], gams[:,k]);
            tmp = isnan.(ft[:, k]);
            while any(tmp)
                rgam2 = random_gamma(gam, 1);
                gam_tmp = invert_gamma(rgam2);
                ft[:, k] = approx(time0, fs[:,k], gam_tmp);
            end
        end
    end

    out = fdawarp(warp_data.f, timet, fn, qn, warp_data.q0, warp_data.fmean,
                  warp_data.mqn, gam, warp_data.orig_var, warp_data.amp_var,
                  warp_data.phase_var, warp_data.cost, warp_data.lambda,
                  warp_data.method, warp_data.omethod, warp_data.gamI, true,
                  fs, rgam, ft, q_s)

    return out

end


module GXBeam

using ArnoldiMethod
using FLOWMath
using LinearAlgebra
using LinearMaps
using NLsolve
using SparseArrays
using StaticArrays
using WriteVTK

export curve_length, discretize_beam
export System, Assembly, PrescribedConditions, DistributedLoads
export reset_state!
export static_analysis, steady_state_analysis, eigenvalue_analysis, time_domain_analysis
export static_analysis!, steady_state_analysis!, eigenvalue_analysis!, time_domain_analysis!
export AssemblyState, left_eigenvectors, wiener_milenkovic, write_vtk, correlate_eigenmodes

# Constant used for scaling forces/moments
# this is needed because sparse arrays don't seem to be able to handle
# mismatched units very well
const FORCE_SCALING = nextpow(2, 1e3)

const GAUSS_NODES = SVector(-0.8611363115940526, -0.3399810435848563, 0.3399810435848563, 0.8611363115940526)
const GAUSS_WEIGHTS = SVector(0.34785484513745385, 0.6521451548625462, 0.6521451548625462, 0.34785484513745385)

include("math.jl")
include("element.jl")
include("point.jl")
include("loads.jl")
include("assembly.jl")
include("system.jl")
include("analyses.jl")
include("postprocess.jl")

end


using MiraculousNN
using Test

using MLDatasets

@testset "Basic neural net workflow" begin
    ninstances = 1000
    nclasses = 10

    x, y_orig = MNIST.traintensor(Float64, 1:ninstances), MNIST.trainlabels(1:ninstances)

    x = reshape(x, (:, ninstances))
    y = MiraculousNN.utils.to_categorical(y_orig, nclasses)

    input_size = size(x, 1)
    output_size = size(y, 1)

    nn_params = MiraculousNN._init_layers(typeof(x), [input_size, 20, 10, output_size])

    @testset "init_layers" begin
        @test keys(nn_params) |> length == 6

        @test size(nn_params["W1"]) == (20, 784)
        @test size(nn_params["B1"]) == (20, 1)

        @test size(nn_params["W2"]) == (10, 20)
        @test size(nn_params["B2"]) == (10, 1)

        @test size(nn_params["W3"]) == (output_size, 10)
        @test size(nn_params["B3"]) == (output_size, 1)

        @test nn_params["W1"] != nn_params["W2"]
        @test nn_params["W1"] != nn_params["W3"]

        @test nn_params["B1"] != nn_params["B2"]
        @test nn_params["B1"] != nn_params["B3"]
    end

    nn_values = MiraculousNN._forward(x, nn_params)

    @testset "forward" begin
        @test keys(nn_values) |> length == 6

        @test size(nn_values["Z1"]) == (20, ninstances)
        @test size(nn_values["A1"]) == (20, ninstances)

        @test size(nn_values["Z2"]) == (10, ninstances)
        @test size(nn_values["A2"]) == (10, ninstances)

        @test size(nn_values["Z3"]) == (output_size, ninstances)
        @test size(nn_values["A3"]) == (output_size, ninstances)
    end

    grads = MiraculousNN._backward(nn_params, nn_values, x, y)

    @testset "backward" begin
        @test keys(grads) |> length == 6

        @test size(grads["W1"]) == (20, 784)
        @test size(grads["B1"]) == (20, 1)

        @test size(grads["W2"]) == (10, 20)
        @test size(grads["B2"]) == (10, 1)

        @test size(grads["W3"]) == (output_size, 10)
        @test size(grads["B3"]) == (output_size, 1)
    end

    MiraculousNN._update_weights!(nn_params, grads, 0.1)

    @testset "update_weights" begin
        @test keys(nn_params) |> length == 6

        @test size(nn_params["W1"]) == (20, 784)
        @test size(nn_params["B1"]) == (20, 1)

        @test size(nn_params["W2"]) == (10, 20)
        @test size(nn_params["B2"]) == (10, 1)

        @test size(nn_params["W3"]) == (output_size, 10)
        @test size(nn_params["B3"]) == (output_size, 1)

        @test nn_params["W1"] != nn_params["W2"]
        @test nn_params["W1"] != nn_params["W3"]

        @test nn_params["B1"] != nn_params["B2"]
        @test nn_params["B1"] != nn_params["B3"]
    end
end


function ETDRK4(SLPDE::SemiLinearPDE{1, S}, T::Number, n::Int) where S
    U0 = SLPDE.U0[1]
    U = deepcopy(U0) # solution spherical harmonic coefficients
    NU = zero(U) # nonlinearity spherical harmonic coefficients

    # ETDRK4 stages
    A = zero(U)
    B = zero(U)
    C = zero(U)
    NA = zero(U)
    NB = zero(U)
    NC = zero(U)

    P = plan_sph2fourier(U)
    PS = plan_sph_synthesis(U)
    PA = plan_sph_analysis(U)

    L = create_linear_operator(SLPDE.L, size(U, 1), size(U, 2)÷2 + 1)
    h = S(T)/n

    # Create all the matrices of diagonal scalings used in ETDRK4. These operate
    # component-wise on the spherical harmonic coefficients.

    Z = L*h
    ez = sph_zero_spurious_modes!(exp.(Z))
    ez2 = sph_zero_spurious_modes!(exp.(Z./2))
    h2ez2m1 = sph_zero_spurious_modes!((h/2.0).*φ1.(Z./2))
    hezα = sph_zero_spurious_modes!(h.*expα.(Z))
    h2ezβ = sph_zero_spurious_modes!(2h.*expβ.(Z))
    hezγ = sph_zero_spurious_modes!(h.*expγ.(Z))
    ez2u = zero(Z)

    for k=1:n
        ez2u .= ez2.*U
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NU .= U)
        A .= ez2u .+ h2ez2m1.*NU
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NA .= A)
        B .= ez2u .+ h2ez2m1.*NA
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NB .= B)
        C .= ez2.*A .+ h2ez2m1.*(2.0.*NB .- NU)
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NC .= C)
        U .= ez.*U .+ hezα.*NU .+ h2ezβ.*(NA .+ NB) .+ hezγ.*NC
    end

    U
end

function ETDRK4(SLPDE::SemiLinearPDE{2, S}, T::Number, n::Int) where S
    U0 = SLPDE.U0
    U = deepcopy(U0) # solution spherical harmonic coefficients
    NU = zero(U[1]), zero(U[2]) # nonlinearity spherical harmonic coefficients
    TEMP = deepcopy(NU[1])

    # ETDRK4 stages
    A = zero(U[1]), zero(U[2])
    B = zero(U[1]), zero(U[2])
    C = zero(U[1]), zero(U[2])
    NA = zero(U[1]), zero(U[2])
    NB = zero(U[1]), zero(U[2])
    NC = zero(U[1]), zero(U[2])

    P = plan_sph2fourier(U[1])
    PS = plan_sph_synthesis(U[1])
    PA = plan_sph_analysis(U[1])

    L = create_linear_operator(SLPDE.L[1], size(U[1], 1), size(U[1], 2)÷2 + 1), create_linear_operator(SLPDE.L[2], size(U[2], 1), size(U[2], 2)÷2 + 1)
    h = S(T)/n

    # Create all the matrices of diagonal scalings used in ETDRK4. These operate
    # component-wise on the spherical harmonic coefficients.

    Z = L[1]*h, L[2]*h
    ez = sph_zero_spurious_modes!(exp.(Z[1])), sph_zero_spurious_modes!(exp.(Z[2]))
    ez2 = sph_zero_spurious_modes!(exp.(Z[1]./2)), sph_zero_spurious_modes!(exp.(Z[2]./2))
    h2ez2m1 = sph_zero_spurious_modes!((h/2.0).*φ1.(Z[1]./2)), sph_zero_spurious_modes!((h/2.0).*φ1.(Z[2]./2))
    hezα = sph_zero_spurious_modes!(h.*expα.(Z[1])), sph_zero_spurious_modes!(h.*expα.(Z[2]))
    h2ezβ = sph_zero_spurious_modes!(2h.*expβ.(Z[1])), sph_zero_spurious_modes!(2h.*expβ.(Z[2]))
    hezγ = sph_zero_spurious_modes!(h.*expγ.(Z[1])), sph_zero_spurious_modes!(h.*expγ.(Z[2]))
    ez2u = zero(Z[1]), zero(Z[2])

    for k=1:n
        ez2u[1] .= ez2[1].*U[1]
        ez2u[2] .= ez2[2].*U[2]
        NU[1] .= U[1]
        NU[2] .= U[2]
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NU, TEMP)
        A[1] .= ez2u[1] .+ h2ez2m1[1].*NU[1]
        A[2] .= ez2u[2] .+ h2ez2m1[2].*NU[2]
        NA[1] .= A[1]
        NA[2] .= A[2]
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NA, TEMP)
        B[1] .= ez2u[1] .+ h2ez2m1[1].*NA[1]
        B[2] .= ez2u[2] .+ h2ez2m1[2].*NA[2]
        NB[1] .= B[1]
        NB[2] .= B[2]
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NB, TEMP)
        C[1] .= ez2[1].*A[1] .+ h2ez2m1[1].*(2.0.*NB[1] .- NU[1])
        C[2] .= ez2[2].*A[2] .+ h2ez2m1[2].*(2.0.*NB[2] .- NU[2])
        NC[1] .= C[1]
        NC[2] .= C[2]
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NC, TEMP)
        U[1] .= ez[1].*U[1] .+ hezα[1].*NU[1] .+ h2ezβ[1].*(NA[1] .+ NB[1]) .+ hezγ[1].*NC[1]
        U[2] .= ez[2].*U[2] .+ hezα[2].*NU[2] .+ h2ezβ[2].*(NA[2] .+ NB[2]) .+ hezγ[2].*NC[2]
    end

    U
end

function fillF!(Ft, F)
    m, n = size(F)
    FN = 0.0
    FS = 0.0
    @inbounds for j = 1:n
        FN += F[1, j]
        FS += F[m, j]
    end
    FN /= n
    FS /= n
    @inbounds for j = 1:n
        for i = 1:m
            Ft[i+1, j] = F[i, j]
        end
        Ft[1, j] = FN
        Ft[m+2, j] = FS
    end
    @inbounds for i = 1:m+2
        Ft[i, n+1] = Ft[i, 1]
    end
end

function ETDRK4(SLPDE::SemiLinearPDE{1, S}, T::Number, n::Int, surf, io; colorrange::Symbol=:fixed) where S
    U0 = SLPDE.U0[1]
    U = deepcopy(U0) # solution spherical harmonic coefficients
    NU = zero(U) # nonlinearity spherical harmonic coefficients
    F = zero(U) # solution function values
    Ft = zeros(Float32, size(F, 1)+2, size(F, 2)+1)

    # ETDRK4 stages
    A = zero(U)
    B = zero(U)
    C = zero(U)
    NA = zero(U)
    NB = zero(U)
    NC = zero(U)

    P = plan_sph2fourier(U)
    PS = plan_sph_synthesis(U)
    PA = plan_sph_analysis(U)

    L = create_linear_operator(SLPDE.L, size(U, 1), size(U, 2)÷2 + 1)
    h = S(T)/n

    # Create all the matrices of diagonal scalings used in ETDRK4. These operate
    # component-wise on the spherical harmonic coefficients.

    Z = L*h
    ez = sph_zero_spurious_modes!(exp.(Z))
    ez2 = sph_zero_spurious_modes!(exp.(Z./2))
    h2ez2m1 = sph_zero_spurious_modes!((h/2.0).*φ1.(Z./2))
    hezα = sph_zero_spurious_modes!(h.*expα.(Z))
    h2ezβ = sph_zero_spurious_modes!(2h.*expβ.(Z))
    hezγ = sph_zero_spurious_modes!(h.*expγ.(Z))
    ez2u = zero(Z)

    for k=1:n
        ez2u .= ez2.*U
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NU .= U)
        A .= ez2u .+ h2ez2m1.*NU
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NA .= A)
        B .= ez2u .+ h2ez2m1.*NA
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NB .= B)
        C .= ez2.*A .+ h2ez2m1.*(2.0.*NB .- NU)
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NC .= C)
        U .= ez.*U .+ hezα.*NU .+ h2ezβ.*(NA .+ NB) .+ hezγ.*NC

        # plotting
        F .= U
        lmul!(P, F)
        lmul!(PS, F)
        fillF!(Ft, F)
        surf[:color] = Ft # animate scene
        if colorrange != :fixed
            surf[:colorrange] = extrema(Ft)
        end
        recordframe!(io) # record a new frame
    end

    U
end

function ETDRK4(SLPDE::SemiLinearPDE{2, S}, T::Number, n::Int, surf, io; colorrange::Symbol=:fixed) where S
    U0 = SLPDE.U0
    U = deepcopy(U0) # solution spherical harmonic coefficients
    NU = zero(U[1]), zero(U[2]) # nonlinearity spherical harmonic coefficients
    TEMP = deepcopy(NU[1])
    F = zero(U[1]) # solution function values
    Ft = zeros(Float32, size(F, 1)+2, size(F, 2)+1)

    # ETDRK4 stages
    A = zero(U[1]), zero(U[2])
    B = zero(U[1]), zero(U[2])
    C = zero(U[1]), zero(U[2])
    NA = zero(U[1]), zero(U[2])
    NB = zero(U[1]), zero(U[2])
    NC = zero(U[1]), zero(U[2])

    P = plan_sph2fourier(U[1])
    PS = plan_sph_synthesis(U[1])
    PA = plan_sph_analysis(U[1])

    L = create_linear_operator(SLPDE.L[1], size(U[1], 1), size(U[1], 2)÷2 + 1), create_linear_operator(SLPDE.L[2], size(U[2], 1), size(U[2], 2)÷2 + 1)
    h = S(T)/n

    # Create all the matrices of diagonal scalings used in ETDRK4. These operate
    # component-wise on the spherical harmonic coefficients.

    Z = L[1]*h, L[2]*h
    ez = sph_zero_spurious_modes!(exp.(Z[1])), sph_zero_spurious_modes!(exp.(Z[2]))
    ez2 = sph_zero_spurious_modes!(exp.(Z[1]./2)), sph_zero_spurious_modes!(exp.(Z[2]./2))
    h2ez2m1 = sph_zero_spurious_modes!((h/2.0).*φ1.(Z[1]./2)), sph_zero_spurious_modes!((h/2.0).*φ1.(Z[2]./2))
    hezα = sph_zero_spurious_modes!(h.*expα.(Z[1])), sph_zero_spurious_modes!(h.*expα.(Z[2]))
    h2ezβ = sph_zero_spurious_modes!(2h.*expβ.(Z[1])), sph_zero_spurious_modes!(2h.*expβ.(Z[2]))
    hezγ = sph_zero_spurious_modes!(h.*expγ.(Z[1])), sph_zero_spurious_modes!(h.*expγ.(Z[2]))
    ez2u = zero(Z[1]), zero(Z[2])

    for k=1:n
        ez2u[1] .= ez2[1].*U[1]
        ez2u[2] .= ez2[2].*U[2]
        NU[1] .= U[1]
        NU[2] .= U[2]
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NU, TEMP)
        A[1] .= ez2u[1] .+ h2ez2m1[1].*NU[1]
        A[2] .= ez2u[2] .+ h2ez2m1[2].*NU[2]
        NA[1] .= A[1]
        NA[2] .= A[2]
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NA, TEMP)
        B[1] .= ez2u[1] .+ h2ez2m1[1].*NA[1]
        B[2] .= ez2u[2] .+ h2ez2m1[2].*NA[2]
        NB[1] .= B[1]
        NB[2] .= B[2]
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NB, TEMP)
        C[1] .= ez2[1].*A[1] .+ h2ez2m1[1].*(2.0.*NB[1] .- NU[1])
        C[2] .= ez2[2].*A[2] .+ h2ez2m1[2].*(2.0.*NB[2] .- NU[2])
        NC[1] .= C[1]
        NC[2] .= C[2]
        evaluate_nonlinear_operator!(SLPDE.N, P, PS, PA, NC, TEMP)
        U[1] .= ez[1].*U[1] .+ hezα[1].*NU[1] .+ h2ezβ[1].*(NA[1] .+ NB[1]) .+ hezγ[1].*NC[1]
        U[2] .= ez[2].*U[2] .+ hezα[2].*NU[2] .+ h2ezβ[2].*(NA[2] .+ NB[2]) .+ hezγ[2].*NC[2]

        # plotting
        F .= U[1]
        lmul!(P, F)
        lmul!(PS, F)
        fillF!(Ft, F)
        surf[:color] = Ft # animate scene
        if colorrange != :fixed
            surf[:colorrange] = extrema(Ft)
        end
        recordframe!(io) # record a new frame
    end

    U
end

function sph_zero_spurious_modes!(A::AbstractMatrix)
    M, N = size(A)
    n = N÷2
    @inbounds for j = 1:n-1
        @simd for i = M-j+1:M
            A[i,2j] = 0
            A[i,2j+1] = 0
        end
    end
    @inbounds @simd for i = M-n+1:M
        A[i,2n] = 0
        2n < N && (A[i,2n+1] = 0)
    end
    A
end

#
# These formulæ, appearing in Eq. (2.5) of:
#
# A.-K. Kassam and L. N. Trefethen, Fourth-order time-stepping for stiff PDEs, SIAM J. Sci. Comput., 26:1214--1233, 2005,
#
# are derived to implement ETDRK4 in double precision without numerical instability from cancellation.
#

expα_asy(x) = (exp(x)*(4-3x+x^2)-4-x)/x^3
expβ_asy(x) = (exp(x)*(x-2)+x+2)/x^3
expγ_asy(x) = (exp(x)*(4-x)-4-3x-x^2)/x^3

# TODO: General types

expα_taylor(x::Union{Float64,ComplexF64}) = @evalpoly(x,1/6,1/6,3/40,1/45,5/1008,1/1120,7/51840,1/56700,1/492800,1/4790016,11/566092800,1/605404800,13/100590336000,1/106748928000,1/1580833013760,1/25009272288000,17/7155594141696000,1/7508956815360000)
expβ_taylor(x::Union{Float64,ComplexF64}) = @evalpoly(x,1/6,1/12,1/40,1/180,1/1008,1/6720,1/51840,1/453600,1/4435200,1/47900160,1/566092800,1/7264857600,1/100590336000,1/1494484992000,1/23712495206400,1/400148356608000,1/7155594141696000,1/135161222676480000)
expγ_taylor(x::Union{Float64,ComplexF64}) = @evalpoly(x,1/6,0/1,-1/120,-1/360,-1/1680,-1/10080,-1/72576,-1/604800,-1/5702400,-1/59875200,-1/691891200,-1/8717829120,-1/118879488000,-1/1743565824000,-1/27360571392000,-1/457312407552000,-1/8109673360588800)

expα(x::Float64) = abs(x) < 17/16 ? expα_taylor(x) : expα_asy(x)
expβ(x::Float64) = abs(x) < 19/16 ? expβ_taylor(x) : expβ_asy(x)
expγ(x::Float64) = abs(x) < 15/16 ? expγ_taylor(x) : expγ_asy(x)

expα(x::ComplexF64) = abs2(x) < (17/16)^2 ? expα_taylor(x) : expα_asy(x)
expβ(x::ComplexF64) = abs2(x) < (19/16)^2 ? expβ_taylor(x) : expβ_asy(x)
expγ(x::ComplexF64) = abs2(x) < (15/16)^2 ? expγ_taylor(x) : expγ_asy(x)

expα(x) = expα_asy(x)
expβ(x) = expβ_asy(x)
expγ(x) = expγ_asy(x)

φ1(x) = iszero(x) ? one(x) : expm1(x)/x


using Random
using ForwardDiff
using LinearAlgebra

const ∇ = ForwardDiff.gradient

function DSGD!(net::Chain, c::Function, traindata::AbstractArray, testdata::AbstractArray, epochs::Integer, minibatchsize::Integer; η_max::Real=50)
    t = test(net, testdata)
    n = length(traindata)
    η = η_max
    for i in 1:epochs
        minibatchs = [shuffle!(traindata)[k-minibatchsize+1:k] for k in minibatchsize:minibatchsize:n]
        for minibatch in minibatchs
            GD!(net, c, minibatch, η=η)
        end
        nt = test(net, testdata)
        println(nt)
        if nt <= t
            η /= 2
        end
        t = nt
    end
    η
end

function test(net::Chain, testdata::AbstractArray)
    count = 0
    for xy in testdata
        py = net(xy[1])
        findmax(xy[2])[2] == findmax(py)[2] && (count+=1)
    end
    return count/length(testdata)
end

function SGD!(net::Chain, c::Function, traindata::AbstractArray, epochs::Integer, minibatchsize::Integer; η::Real=1)
    n = length(traindata)
    for i in 1:epochs
        minibatchs = [shuffle!(traindata)[k-minibatchsize+1:k] for k in minibatchsize:minibatchsize:n]
        for minibatch in minibatchs
            GD!(net, c, minibatch, η=η)
        end
    end
end

function GD!(net::Chain, c::Function, traindata::AbstractArray, epochs::Integer; η::Real=1)
    for i in 1:epochs
        GD!(net, c, traindata, η=η)
    end
end

function GD!(net::Chain, c::Function, traindata::AbstractArray; η::Real=1)
    dW = [zeros(size(l.W)) for l in net.layers]
    db = [zeros(size(l.b)) for l in net.layers]
    for xy in traindata
        ∂b, ∂W = backprop(net, c, xy...)
        for i in eachindex(dW)
            dW[i] += ∂W[i]
            db[i] += ∂b[i]
        end
    end
    for i in eachindex(dW)
        net.layers[i].W .-= dW[i]*η/length(traindata)
        net.layers[i].b .-= db[i]*η/length(traindata)
    end
end

function backprop(net::Chain, c::Function, x::Vector{R}, y::Vector{R})where R<:Real
    a = x
    as = Vector{Vector{R}}()
    zs = Vector{Vector{R}}()
    push!(as, deepcopy(a))
    for l in net.layers
        z = l(a)
        a = l.f.(z)
        push!(zs, z)
        push!(as, deepcopy(a))
    end
    δs = similar(zs)
    δs[end] = ∇(a->cost(c, a, y), as[end]) .* dσ.(zs[end])
    for i in length(net.layers)-1
        δs[end-i] = transpose(net.layers[end-i+1].W) * δs[end-i+1] .* dσ.(zs[end-i])
    end
    ∂W = Vector{Matrix{R}}()
    for k in eachindex(δs)
        push!(∂W, zeros(length(δs[k]), length(as[k])))
        for i in eachindex(δs[k]), j in eachindex(as[k])
            ∂W[k][i, j] = δs[k][i] * as[k][j]
        end
    end
    return δs, ∂W # δs equals to δs
end

function cost(c::Function, a::Vector{R}, y::Vector{R})where R<:Real
    return sum(c, y.-a)/length(a)
end


function createstiffnessmatrices!(params::Parameters)
  noofnodes = length(params.nodes)
  sb0 = zeros(6, 6)
  for i in 1:noofnodes
    if 1 < i < noofnodes
      ler = params.elements[i-1].er
      lea = params.elements[i-1].ea
      lgj = params.elements[i-1].gj
      l1 = params.nodes[i].z - params.nodes[i-1].z
      l2 = params.nodes[i+1].z - params.nodes[i].z
      params.sa[:, :, i] = stiffnessa!(i, l1, ler, lea, lgj)
      params.sb1[:, :, i] = stiffnessb1!(i, l1, ler, lea, lgj)
      params.sb2[:, :, i] = stiffnessb2!(i, noofnodes, l2, ler, lea, lgj)
      params.sc[:, :, i] = stiffnessc!(i, noofnodes, l2, ler, lea, lgj)
      params.si[:, :, i] = stiffnessi!(i, params.sb1[:, :, i], params.sb2[:, :, i])
    else
      if i == 1
        ler = params.elements[i].er
        lea = params.elements[i].ea
        lgj = params.elements[i].gj
        l2 = params.nodes[i+1].z - params.nodes[i].z
        params.sb2[:, :, i] = stiffnessb2!(i, noofnodes, l2, ler, lea, lgj)
        params.sc[:, :, i] = stiffnessc!(i, noofnodes, l2, ler, lea, lgj)
        params.si[:, :, i] = stiffnessi!(i, sb0, params.sb2[:, :, i])
      else
        ler = params.elements[i-1].er
        lea = params.elements[i-1].ea
        lgj = params.elements[i-1].gj
        l1 = params.nodes[i].z - params.nodes[i-1].z
        params.sa[:, :, i] = stiffnessa!(i, l1, ler, lea, lgj)
        params.sb1[:, :, i] = stiffnessb1!(i, l1, ler, lea, lgj)
        params.si[:, :, i] = stiffnessi!(i, params.sb1[:, :, i], sb0)
      end
    end
  end

  params.sainit[:, :, :] = copy(params.sa)
  params.sb1init[:, :, :] = copy(params.sb1)
  params.sb2init[:, :, :] = copy(params.sb2)
  params.scinit[:, :, :] = copy(params.sc)

  nothing
end

function stiffnessa!(node, l, ler, lea, lgj)
  s = zeros(6, 6)
  if node == 1
    return
  else
    s[1, 1] = -12.0 * ler / l^3
    s[1, 5] =  -6.0 * ler / l^2
    s[2, 2] = -12.0 * ler / l^3
    s[2, 4] =   6.0 * ler / l^2
    s[3, 3] =        -lea / l
    s[4, 2] =  -6.0 * ler / l^2
    s[4, 4] =   2.0 * ler / l
    s[5, 1] =   6.0 * ler / l^2
    s[5, 5] =   2.0 * ler / l
    s[6, 6] =        -lgj / l
  end
  s
end

function stiffnessb1!(node, l, ler, lea, lgj)
  s = zeros(6, 6)
  if node == 1
    return
  else
    s[1, 1] =  12.0 * ler / l^3
    s[1, 5] =  -6.0 * ler / l^2
    s[2, 2] =  12.0 * ler / l^3
    s[2, 4] =   6.0 * ler / l^2
    s[3, 3] =         lea / l
    s[4, 2] =   6.0 * ler / l^2
    s[4, 4] =   4.0 * ler / l
    s[5, 1] =  -6.0 * ler / l^2
    s[5, 5] =   4.0 * ler / l
    s[6, 6] =         lgj / l
  end
  s
end

function stiffnessb2!(node, noofnodes, l, ler, lea, lgj)
  s = zeros(6, 6)
  if node == noofnodes
    return
  else
    s[1, 1] =  12.0 * ler / l^3
    s[1, 5] =   6.0 * ler / l^2
    s[2, 2] =  12.0 * ler / l^3
    s[2, 4] =  -6.0 * ler / l^2
    s[3, 3] =         lea / l
    s[4, 2] =  -6.0 * ler / l^2
    s[4, 4] =   4.0 * ler / l
    s[5, 1] =   6.0 * ler / l^2
    s[5, 5] =   4.0 * ler / l
    s[6, 6] =         lgj / l
  end
  s
end

function stiffnessc!(node, noofnodes, l, ler, lea, lgj)
  s = zeros(6, 6)
  if node == noofnodes
    return
  else
    s[1, 1] = -12.0 * ler / l^3
    s[1, 5] =   6.0 * ler / l^2
    s[2, 2] = -12.0 * ler / l^3
    s[2, 4] =  -6.0 * ler / l^2
    s[3, 3] =        -lea / l
    s[4, 2] =   6.0 * ler / l^2
    s[4, 4] =   2.0 * ler / l
    s[5, 1] =  -6.0 * ler / l^2
    s[5, 5] =   2.0 * ler / l
    s[6, 6] =        -lgj / l
  end
  s
end

function stiffnessi!(node, sb1, sb2)
  s = zeros(6, 6)
  b = sb1 + sb2
  h = b[1, 1] * b[5, 5] - b[1, 5]^2
  ff = b[5, 5] / h
  g = -b[1, 5] / h
  h = b[1, 1] / h
  s[1, 1] = ff
  s[1, 5] = g
  s[2, 2] = ff
  s[2, 4] = -g
  s[3, 3] = node == 1 ? 0.0 : 1.0 / b[3, 3]
  s[4, 2] = -g
  s[4, 4] = h
  s[5, 1] = g
  s[5, 5] = h
  s[6, 6] = node == 1 ? 0.0 : 1.0 / b[6, 6]
  s
end


using Huggett890, Test

function euler_test()
    @testset "Euler" begin
        m = init_model();
        t = 2;  # test other ages +++++
        k = 5.0;
        eIdx = 3;
        kPrimeTomorrow = init_test_worker_solution(eIdx + 1);

        kPrime = Huggett890.solve_one_point(m, t, k, eIdx, kPrimeTomorrow);
        dev = Huggett890.euler_dev_one_point(m, t, k, eIdx, kPrime, kPrimeTomorrow);
        @test kPrime >= Huggett890.kprime_min(m.kGrid, t);
        @test kPrime <= Huggett890.kprime_max(m.kGrid, t);
        @test abs(dev) < 1e-4
        # dev2 = Huggett890.euler_dev_one_point2(m, t, k, kPrime, kPrimeTomorrow);
        # @test abs(dev2) < 1e-4

        # Euler needs more direct test +++++
    end
end

function solve_worker_test()
    @testset "Solve worker" begin
        m = init_model();

        # Just a syntax test. Need to solve t and t+1 for substantive testing.
        t = 3; # test other ages +++++
        k = 5.0;
        eIdx = 2;
        nEff = 4;
        kPrimeTomorrow = init_test_worker_solution(nEff);
        sol = Huggett890.solve_one_period(m, kPrimeTomorrow, t);

        Huggett890.solve_worker(m);
        # +++++  @test Huggett890.check_retired_solution(m, sol)
    end
end

function solve_period_test()
    @testset "One period" begin
        m = init_model();
        t = 2;
        k = 5.0;
        nEff = 4;
        eIdx = 2;
        kPrimeTomorrow = init_test_worker_solution(nEff);
        kPrimeFct, cFct = hg.solve_one_efficiency(m, t, eIdx, kPrimeTomorrow);
        kPrime = kPrimeFct(k);
        eMu = hg.e_uprime_c(m, t, eIdx, kPrime, kPrimeTomorrow);
        c = cFct(k);
        muC = hg.marg_utility(m.util, c);
        dev2 = muC / eMu - hg.betar(m);
        # Surprisingly imprecise +++
        @test abs(dev2) < 1e-2

        eff = hg.efficiency(m.endow, eIdx);
        kPrimeMin, kPrimeMax = hg.kprime_range(m, t, k, eff);
        @test kPrimeMax > kPrimeMin
        kPrime2 = hg.solve_one_point(m, t, k, eIdx, kPrimeTomorrow);
        @test isapprox(kPrime2, kPrime)
    end
end

function solve_test()
    @testset "Solve" begin
        m = init_model();
        sol = solve_worker(m);
        @test sol isa Solution
    end
end

@testset "Worker" begin
    euler_test();
    solve_period_test();
    solve_worker_test();
    solve_test();
end

# -----------


# # One-dimensional cellular automata
# Agents.jl provides a module (CA1D) to create and plot 1D cellular automata.

using Agents
using Agents.CA1D
using AgentsPlots

# ## 1. Define the rule
# Here is Wolfram's rule 22

rules = Dict(
    "111" => "0",
    "110" => "0",
    "101" => "0",
    "100" => "1",
    "011" => "0",
    "010" => "1",
    "001" => "1",
    "000" => "0",
);

# ## 2. Build the model
# All the cells are by default initially "off"
model = CA1D.build_model(rules = rules, ncols = 21)

# This turns on the middle cell
model.agents[11].status = "1";

# ## 3. Run the model and collect data.
runs = 20
data, _ = CA1D.ca_run(model, runs);

# ## 4. Visualize the data
p = plot_CA1D(data, nodesize = 2)


"""
compute_scc(m::Model=get_model(); year::Int = nothing, last_year::Int = 2595), prtp::Float64 = 0.03)

Computes the social cost of CO2 for an emissions pulse in `year` for the provided MimiRICE2010 model.
If no model is provided, the default model from MimiRICE2010.get_model() is used.
Constant discounting is used from the specified pure rate of time preference `prtp`.
"""
function compute_scc(m::Model=get_model(); year::Union{Int, Nothing} = nothing, last_year::Int = model_years[end], prtp::Float64 = 0.015, eta::Float64=1.5)
    year === nothing ? error("Must specify an emission year. Try `compute_scc(m, year=2015)`.") : nothing
    !(last_year in model_years) ? error("Invlaid value of $last_year for last_year. last_year must be within the model's time index $model_years.") : nothing
    !(year in model_years[1]:10:last_year) ? error("Cannot compute the scc for year $year, year must be within the model's time index $(model_years[1]):10:$last_year.") : nothing

    mm = get_marginal_model(m; year = year)

    return _compute_scc(mm, year=year, last_year=last_year, prtp=prtp, eta=eta)
end

"""
compute_scc_mm(m::Model=get_model(); year::Int = nothing, last_year::Int = 2595, prtp::Float64 = 0.03)

Returns a NamedTuple (scc=scc, mm=mm) of the social cost of carbon and the MarginalModel used to compute it.
Computes the social cost of CO2 for an emissions pulse in `year` for the provided MimiRICE2010 model.
If no model is provided, the default model from MimiRICE2010.get_model() is used.
Constant discounting is used from the specified pure rate of time preference `prtp`.
"""
function compute_scc_mm(m::Model=get_model(); year::Union{Int, Nothing} = nothing, last_year::Int = model_years[end], prtp::Float64 = 0.015, eta::Float64=1.5)
    year === nothing ? error("Must specify an emission year. Try `compute_scc_mm(m, year=2015)`.") : nothing
    !(last_year in model_years) ? error("Invlaid value of $last_year for last_year. last_year must be within the model's time index $model_years.") : nothing
    !(year in model_years[1]:10:last_year) ? error("Cannot compute the scc for year $year, year must be within the model's time index $(model_years[1]):10:$last_year.") : nothing

    mm = get_marginal_model(m; year = year)
    scc = _compute_scc(mm; year=year, last_year=last_year, prtp=prtp, eta=eta)

    return (scc = scc, mm = mm)
end

# helper function for computing SCC from a MarginalModel, not to be exported or advertised to users
function _compute_scc(mm::MarginalModel; year::Int, last_year::Int, prtp::Float64, eta::Float64)
    ntimesteps = findfirst(isequal(last_year), model_years)     # Will run through the timestep of the specified last_year
    run(mm, ntimesteps=ntimesteps)

    marginal_damages = -1 * mm[:neteconomy, :C][1:ntimesteps, :] * 10.0^12 * 12/44 # convert from trillion $/ton C to $/ton CO2; multiply by -1 to get positive value for damages
    global_marginal_damages = dropdims(sum(marginal_damages, dims = 2), dims=2)

    global_c = dropdims(sum(mm.base[:neteconomy, :C], dims = 2), dims=2)
    global_pop = dropdims(sum(mm.base[:neteconomy, :l], dims = 2), dims=2)
    global_cpc = 1000 .* global_c ./ global_pop

    year_index = findfirst(isequal(year), model_years)

    df = [zeros(year_index-1)..., ((global_cpc[year_index]/global_cpc[i])^eta * 1/(1+prtp)^(t-year) for (i,t) in enumerate(model_years) if year<=t<=last_year)...]
    scc = sum(df .* marginal_damages * 10)  # currently implemented as a 10year step function; so each timestep of discounted marginal damages is multiplied by 10
    return scc
end

"""
get_marginal_model(m::Model = get_model(); year::Int = nothing)

Creates a Mimi MarginalModel where the provided m is the base model, and the marginal model has additional emissions of CO2 in year `year`.
If no Model m is provided, the default model from MimiDICE2010.get_model() is used as the base model.
"""
function get_marginal_model(m::Model=get_model(); year::Union{Int, Nothing} = nothing)
    year === nothing ? error("Must specify an emission year. Try `get_marginal_model(m, year=2015)`.") : nothing
    !(year in model_years) ? error("Cannot add marginal emissions in $year, year must be within the model's time index $(model_years[1]):10:$last_year.") : nothing

    mm = create_marginal_model(m, 1e10) # Pulse has a value of 1GtC per year for ten years
    add_marginal_emissions!(mm.modified, year)

    return mm
end

"""
Adds a marginal emission component to year m which adds 1Gt of additional C emissions per year for ten years starting in the specified `year`.
"""
function add_marginal_emissions!(m::Model, year::Int)
    add_comp!(m, Mimi.adder, :marginalemission, before=:co2cycle)

    time = Mimi.dimension(m, :time)
    addem = zeros(length(time))
    addem[time[year]] = 1.0     # 1 GtC per year for ten years

    set_param!(m, :marginalemission, :add, addem)
    connect_param!(m, :marginalemission, :input, :emissions, :E)
    connect_param!(m, :co2cycle, :E, :marginalemission, :output)
end

# OLD previously used marginal damages functions

function getmarginal_rice_models(;emissionyear=2005,datafile=joinpath(@__DIR__, "..", "data", "RICE_2010_base_000.xlsm"))

    RICE = get_model()
    run(RICE)

    mm = MarginalModel(RICE)
    m1 = mm.base
    m2 = mm.modified

    add_comp!(m2, Mimi.adder, :marginalemission, before=:co2cycle)

    time = Mimi.dimension(m1, :time)
    addem = zeros(length(time))
    addem[time[emissionyear]] = 1.0

    set_param!(m2,:marginalemission,:add,addem)
    connect_param!(m2,:marginalemission,:input,:emissions,:E)
    connect_param!(m2, :co2cycle,:E,:marginalemission,:output)

    run(m1)
    run(m2)

    return m1, m2
end

# This function returns a matrix of marginal damages per one t of carbon emission in the
# emissionyear parameter year.
function getmarginaldamages_rice(;emissionyear=2005,datafile=joinpath(@__DIR__, "..", "data", "RICE_2010_base_000.xlsm"))
    m1, m2 = getmarginal_rice_models(emissionyear=emissionyear, datafile=datafile)

    run(m1)
    run(m2)

    damage1 = m1[:damages,:DAMAGES]
    damage2 = m2[:damages,:DAMAGES]

    # Calculate the marginal damage between run 1 and 2 for each
    # year/region
    marginaldamage = (damage2.-damage1) .* 1e12 / 1e9 / 10.

    return marginaldamage
end


using Test, MathOptInterface
using LinQuadOptInterface

const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities
const MOIT = MathOptInterface.Test
const MOIU = MathOptInterface.Utilities
const LQOI = LinQuadOptInterface

# Test that using LinQuadOptInterface passes without error. We have to do this
# here because types can't be created in testsets. If it can be defined without
# error, we're okay. This is most likely to error if fields are added to
# @LinQuadOptimizerBase without proper module prefixes.
mutable struct OptimizerTest <: LinQuadOptInterface.LinQuadOptimizer
    LinQuadOptInterface.@LinQuadOptimizerBase
    OptimizerTest() = new()
end

@testset "LinQuadOptInterface" begin
    solver = LQOI.MockLinQuadOptimizer()

    @testset "Printing" begin
        @test sprint(show, solver) == "A LinQuadOptInterface model with " *
            "backend:\nMockLinQuadModel\n    Sense       minimize\n    " *
            "Variables   0\n    Constraints 0\n"
    end

    config = MOIT.TestConfig(solve=false)

    @testset "Unit Tests" begin
        config = MOIT.TestConfig(solve=false)
        MOIT.basic_constraint_tests(solver, config)
        MOIT.unittest(solver, config, [
            "solve_affine_interval",
            "solve_qp_edge_cases",
            "solve_qcp_edge_cases",
            "solve_affine_deletion_edge_cases"
        ])
    end

    @testset "Linear tests" begin
        config = MOIT.TestConfig(solve=false)
        # Exclude partial_start since VariablePrimalStart not implemented.
        MOIT.contlineartest(solver, config, ["partial_start"])
        config = MOIT.TestConfig()
        include("contlinear.jl")
        set_linear1test_solutions!(solver)
        MOIT.linear1test(solver, config)
        set_linear2test_solutions!(solver)
        MOIT.linear2test(solver, config)
        set_linear3test_solutions!(solver)
        MOIT.linear3test(solver, config)
        set_linear4test_solutions!(solver)
        MOIT.linear4test(solver, config)
        set_linear5test_solutions!(solver)
        MOIT.linear5test(solver, config)
        set_linear6test_solutions!(solver)
        MOIT.linear6test(solver, config)
        set_linear7test_solutions!(solver)
        MOIT.linear7test(solver, config)
        set_linear8atest_solutions!(solver)
        MOIT.linear8atest(solver, config)
        set_linear8btest_solutions!(solver)
        MOIT.linear8btest(solver, config)
        set_linear8ctest_solutions!(solver)
        MOIT.linear8ctest(solver, config)
        set_linear9test_solutions!(solver)
        MOIT.linear9test(solver, config)
        set_linear10test_solutions!(solver)
        MOIT.linear10test(solver, config)
        set_linear11test_solutions!(solver)
        MOIT.linear11test(solver, config)
        set_linear12test_solutions!(solver)
        MOIT.linear12test(solver, MOIT.TestConfig(atol=1e-3,rtol=1e-3))
        set_linear13test_solutions!(solver)
        MOIT.linear13test(solver, config)
        set_linear14test_solutions!(solver)
        MOIT.linear14test(solver, config)
        # set_linear15test_solutions!(solver)
        # MOIT.linear15test(solver, config)
    end

    @testset "Quadratic tests" begin
        config = MOIT.TestConfig(solve=false)
        MOIT.contquadratictest(solver, config)
        include("contquadratic.jl")
        config = MOIT.TestConfig(atol=1e-3, rtol=1e-3, duals=false)
        set_qp1test_solutions!(solver)
        MOIT.qp1test(solver, config)
        set_qp2test_solutions!(solver)
        MOIT.qp2test(solver, config)
        set_qp3test_solutions!(solver)
        MOIT.qp3test(solver, config)
        set_qcp1test_solutions!(solver)
        MOIT.qcp1test(solver, config)
        set_qcp2test_solutions!(solver)
        MOIT.qcp2test(solver, config)
        set_qcp3test_solutions!(solver)
        MOIT.qcp3test(solver, config)
        set_socp1test_solutions!(solver)
        MOIT.socp1test(solver, config)
    end

    @testset "Linear Conic tests" begin
        config = MOIT.TestConfig(solve=false)
        MOIT.lintest(solver, config)
        include("contconic.jl")
        config = MOIT.TestConfig()
        set_lin1test_solutions!(solver)
        MOIT.lin1vtest(solver, config)
        set_lin1test_solutions!(solver)
        MOIT.lin1ftest(solver, config)
        set_lin2test_solutions!(solver)
        MOIT.lin2vtest(solver, config)
        set_lin2test_solutions!(solver)
        MOIT.lin2ftest(solver, config)
        set_lin3test_solutions!(solver)
        MOIT.lin3test(solver, config)
        set_lin3test_solutions!(solver)
        MOIT.lin4test(solver, config)
    end

    @testset "Integer Linear tests" begin
        config = MOIT.TestConfig(solve=false)
        MOIT.intlineartest(solver, config, ["int2"])
        include("intlinear.jl")
        config = MOIT.TestConfig()
        set_knapsacktest_solutions!(solver)
        MOIT.knapsacktest(solver, config)
        set_int1test_solutions!(solver)
        MOIT.int1test(solver, config)
        set_int3test_solutions!(solver)
        MOIT.int3test(solver, config)
    end

    @testset "ModelLike tests" begin
        config = MOIT.TestConfig(solve=false)
        @testset "default_objective_test" begin
            MOIT.default_objective_test(LQOI.MockLinQuadOptimizer())
        end
        @testset "default_status_test" begin
            MOIT.default_status_test(LQOI.MockLinQuadOptimizer())
        end
        @testset "nametest" begin
            MOIT.nametest(LQOI.MockLinQuadOptimizer())
        end
        @testset "validtest" begin
            MOIT.validtest(LQOI.MockLinQuadOptimizer())
        end
        @testset "emptytest" begin
            MOIT.emptytest(LQOI.MockLinQuadOptimizer())
        end
        @testset "orderedindicestest" begin
            MOIT.orderedindicestest(LQOI.MockLinQuadOptimizer())
        end
        @testset "copytest" begin
            MOIT.copytest(LQOI.MockLinQuadOptimizer(),
                          LQOI.MockLinQuadOptimizer())
            MOIT.failcopytestc(LQOI.MockLinQuadOptimizer())
            MOIT.failcopytestia(LQOI.MockLinQuadOptimizer())
            MOIT.failcopytestva(LQOI.MockLinQuadOptimizer())
            MOIT.failcopytestca(LQOI.MockLinQuadOptimizer())
        end
    end
end

MOIU.@model(ModelComplete,
    (MOI.ZeroOne, MOI.Integer),
    (MOI.EqualTo, MOI.GreaterThan, MOI.LessThan, MOI.Interval,
     MOI.Semicontinuous, MOI.Semiinteger),
    (MOI.Reals, MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives,
     MOI.SecondOrderCone, MOI.RotatedSecondOrderCone, MOI.GeometricMeanCone,
     MOI.ExponentialCone, MOI.DualExponentialCone,
     MOI.PositiveSemidefiniteConeTriangle, MOI.PositiveSemidefiniteConeSquare,
     MOI.RootDetConeTriangle, MOI.RootDetConeSquare, MOI.LogDetConeTriangle,
     MOI.LogDetConeSquare),
    (MOI.PowerCone, MOI.DualPowerCone, MOI.SOS1, MOI.SOS2),
    (MOI.SingleVariable,),
    (MOI.ScalarAffineFunction, MOI.ScalarQuadraticFunction),
    (MOI.VectorOfVariables,),
    (MOI.VectorAffineFunction, MOI.VectorQuadraticFunction))
@testset "Copy from/to @Model" begin
    MOIT.copytest(LQOI.MockLinQuadOptimizer(), ModelComplete{Float64}())
    MOIT.copytest(ModelComplete{Float64}(), LQOI.MockLinQuadOptimizer())
end

@testset "Issue #52" begin
    @testset "scalaraffine" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        f1 =  MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0], [x]), 1.5)
        s1 = MOI.LessThan(2.0)
        c = MOI.add_constraint(model, f1, s1)
        f2 = MOI.get(model, MOI.ConstraintFunction(), c)
        @test f1 ≈ f2
        s2 = MOI.get(model, MOI.ConstraintSet(), c)
        @test typeof(s1) == typeof(s2)
        @test s1.upper == s2.upper
    end
    @testset "vectoraffine" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        f1 =  MOI.VectorAffineFunction(
            MOI.VectorAffineTerm.([1], MOI.ScalarAffineTerm.([1.0], [x])), [1.5])
        s1 = MOI.Zeros(1)
        c = MOI.add_constraint(model, f1, s1)
        f2 = MOI.get(model, MOI.ConstraintFunction(), c)
        @test f1 ≈ f2
        s2 = MOI.get(model, MOI.ConstraintSet(), c)
        @test typeof(s1) == typeof(s2)
        @test s1.dimension == s2.dimension
    end
end

@testset "Issue #54" begin
    @testset "scalar, one-sided" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        f =  MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0], [x]), 1.5)
        s = MOI.LessThan(2.0)
        c = MOI.add_constraint(model, f, s)
        # Change the constraint set, and verify that we get the same set
        # when we retrieve it:
        s2 = MOI.LessThan(1.0)
        MOI.set(model, MOI.ConstraintSet(), c, s2)
        s3 = MOI.get(model, MOI.ConstraintSet(), c)
        @test typeof(s2) == typeof(s3)
        @test s2.upper == s3.upper
    end
    @testset "scalar, interval" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        f =  MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0], [x]), 1.5)
        s = MOI.Interval(2.0, 3.0)
        c = MOI.add_constraint(model, f, s)
        # Change the constraint set, and verify that we get the same set
        # when we retrieve it:
        s2 = MOI.Interval(1.0, 2.0)
        MOI.set(model, MOI.ConstraintSet(), c, s2)
        s3 = MOI.get(model, MOI.ConstraintSet(), c)
        @test typeof(s2) == typeof(s3)
        @test s2.lower == s3.lower
        @test s2.upper == s3.upper
    end
end

@testset "Issue #32" begin
    check_set(actual::MOI.GreaterThan, expected::MOI.GreaterThan) =
        @test actual.lower == expected.lower
    check_set(actual::MOI.LessThan, expected::MOI.LessThan) =
        @test actual.upper == expected.upper
    function check_set(actual::MOI.Interval, expected::MOI.Interval)
        @test actual.lower == expected.lower
        @test actual.lower == expected.lower
    end
    function check_set(actual::S, expected::S) where {S <: Union{MOI.Nonnegatives, MOI.Nonpositives, MOI.Zeros}}
        @test actual.dimension == expected.dimension
    end

    function check_row(model, constraint_index, expected_function, expected_set)
        actual_function = MOI.get(model, MOI.ConstraintFunction(), constraint_index)
        @test actual_function ≈ expected_function
        actual_set = MOI.get(model, MOI.ConstraintSet(), constraint_index)
        check_set(actual_set, expected_set)
    end

    @testset "Scalar affine function" begin
        @testset "Single variable" begin
            model = LQOI.MockLinQuadOptimizer()
            x = MOI.add_variable(model)

            # Add the constraint and verify that we can retrieve it
            f =  MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0], [x]), 0.1)
            s = MOI.GreaterThan(0.5)
            c = MOI.add_constraint(model, f, s)
            check_row(model, c, f, s)

            # Replace the constraint function with itself and verify that
            # the problem is unchanged
            MOI.set(model, MOI.ConstraintFunction(), c, f)
            check_row(model, c, f, s)

            # Replace the constraint function with a new function, verify
            # that the replacement occurred and that the set is unchanged
            f2 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0], [x]), 0.5)
            MOI.set(model, MOI.ConstraintFunction(), c, f2)
            check_row(model, c, f2, s)

            # Replace the constraint function with a new function which is not
            # in canonical form:
            f3 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 0.5], [x]), 0.5)
            MOI.set(model, MOI.ConstraintFunction(), c, f3)
            check_row(model, c, MOIU.canonical(f3), s)

            # Replace the constraint function with a new function whose sparsity
            # pattern is different:
            f4 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{Float64}[], 2.0)
            MOI.set(model, MOI.ConstraintFunction(), c, f4)
            check_row(model, c, f4, s)
        end

        @testset "Multiple variables" begin
            model = LQOI.MockLinQuadOptimizer()
            x = MOI.add_variable(model)
            y = MOI.add_variable(model)

            # Add the constraint and verify that we can retrieve it
            f =  MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 2.0], [x, y]), 0.1)
            s = MOI.GreaterThan(0.5)
            c = MOI.add_constraint(model, f, s)
            check_row(model, c, f, s)

            # Replace the constraint function with itself and verify that
            # the problem is unchanged
            MOI.set(model, MOI.ConstraintFunction(), c, f)
            check_row(model, c, f, s)

            # Replace the constraint function with a new function, verify
            # that the replacement occurred and that the set is unchanged
            f2 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([0.5, 1.0], [x, y]), 0.5)
            MOI.set(model, MOI.ConstraintFunction(), c, f2)
            check_row(model, c, f2, s)

            # Replace the constraint function with a new function which is not
            # in canonical form:
            f3 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-0.1, 0.2], [y, x]), 0.5)
            MOI.set(model, MOI.ConstraintFunction(), c, f3)
            check_row(model, c, MOIU.canonical(f3), s)

            # Replace the constraint function with a new function whose sparsity
            # pattern is different:
            f4 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-2.0], [y]), 2.0)
            MOI.set(model, MOI.ConstraintFunction(), c, f4)
            check_row(model, c, f4, s)
        end
    end

    @testset "Vector affine function" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        y = MOI.add_variable(model)

        # Add the constraint and verify that we can retrieve it
        f = MOI.VectorAffineFunction(
            [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(1.0, x)),
             MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(2.0, y)),
             MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(-1.0, x))],
            [0.1, 0.2])
        s = MOI.Nonpositives(2)
        c = MOI.add_constraint(model, f, s)
        check_row(model, c, f, s)

        # Replace the constraint function with itself and verify that
        # the problem is unchanged
        MOI.set(model, MOI.ConstraintFunction(), c, f)
        check_row(model, c, f, s)

        # Replace the constraint function with a new function, verify
        # that the replacement occurred and that the set is unchanged
        f2 = MOI.VectorAffineFunction(
            [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(3.0, x)),
             MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(-2.0, y)),
             MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(-1.5, x))],
            [0.5, 2.0])
        MOI.set(model, MOI.ConstraintFunction(), c, f2)
        check_row(model, c, f2, s)

        # Replace the constraint function with a new function which is not
        # in canonical form:
        f3 = MOI.VectorAffineFunction(
            [MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(2.0, x)),
             MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(-1.5, x)),
             MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(-2.5, y))],
            [0.1, 0.2])
        MOI.set(model, MOI.ConstraintFunction(), c, f3)
        check_row(model, c, MOIU.canonical(f3), s)

        # Replace the constraint function with a new function whose sparsity
        # pattern is different:
        f4 = MOI.VectorAffineFunction(
            [MOI.VectorAffineTerm(2, MOI.ScalarAffineTerm(-1.5, y)),
             MOI.VectorAffineTerm(1, MOI.ScalarAffineTerm(2.0, x))],
            [0.1, 0.2])
        MOI.set(model, MOI.ConstraintFunction(), c, f4)
        check_row(model, c, f4, s)
    end
end

@testset "Issue #67" begin
    model = LQOI.MockLinQuadOptimizer()
    function objective_type_test(f::MOI.AbstractScalarFunction)
        MOI.set(model, MOI.ObjectiveFunction{typeof(f)}(), f)
        @test MOI.get(model, MOI.ObjectiveFunctionType()) == typeof(f)
    end
    x = MOI.add_variable(model)
    f = MOI.SingleVariable(x)
    objective_type_test(f)
    objective_type_test(convert(MOI.ScalarAffineFunction{Float64}, f))
    objective_type_test(convert(MOI.ScalarQuadraticFunction{Float64}, f))
end

@testset "Conflicting SingleVariable constraints" begin
    @testset "ZeroOne" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        MOI.add_constraint(model, MOI.SingleVariable(x), MOI.ZeroOne())
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Integer())
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Semiinteger(2.0, 3.0))
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Semicontinuous(2.0, 3.0))
        c1 = MOI.add_constraint(model, MOI.SingleVariable(x), MOI.GreaterThan(0.5))
        c2 = MOI.add_constraint(model, MOI.SingleVariable(x), MOI.LessThan(0.5))
        @test MOI.is_valid(model, c1)
        @test MOI.is_valid(model, c2)
    end
    @testset "Integer" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Integer())
        MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Interval(2.0, 3.0))
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.ZeroOne())
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Semiinteger(2.0, 3.0))
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Semicontinuous(2.0, 3.0))
    end
    @testset "Semiinteger" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Semiinteger(2.0, 3.0))
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Interval(2.0, 3.0))
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.ZeroOne())
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Integer())
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Semicontinuous(2.0, 3.0))
    end
    @testset "Semicontinuous" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Semicontinuous(2.0, 3.0))
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Interval(2.0, 3.0))
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.ZeroOne())
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Integer())
        @test_throws Exception MOI.add_constraint(model, MOI.SingleVariable(x), MOI.Semiinteger(2.0, 3.0))
    end
end

@testset "SemiXXX variables" begin
    @testset "Semiinteger" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        c = MOI.add_constraint(
            model, MOI.SingleVariable(x), MOI.Semiinteger(1.0, 4.5))
        @test model.inner.vartype[1] == Cchar('N')
        @test MOI.get(model, MOI.ConstraintSet(), c) == MOI.Semiinteger(1.0, 4.5)
        @test MOI.get(model, MOI.ConstraintFunction(), c) == MOI.SingleVariable(x)
        MOI.delete(model, c)
        @test model.inner.vartype[1] == Cchar('C')
    end
    @testset "Semicontinuous" begin
        model = LQOI.MockLinQuadOptimizer()
        x = MOI.add_variable(model)
        c = MOI.add_constraint(
            model, MOI.SingleVariable(x), MOI.Semicontinuous(1.0, 4.0))
        @test model.inner.vartype[1] == Cchar('S')
        @test MOI.get(model, MOI.ConstraintSet(), c) == MOI.Semicontinuous(1.0, 4.0)
        @test MOI.get(model, MOI.ConstraintFunction(), c) == MOI.SingleVariable(x)
        MOI.delete(model, c)
        @test model.inner.vartype[1] == Cchar('C')
    end
end

@testset "VariablePrimalStart" begin
    model = LQOI.MockLinQuadOptimizer()
    x = MOI.add_variable(model)
    @test_throws MOI.UnsupportedAttribute MOI.set(model, MOI.VariablePrimalStart(), x, 1.0)
end

@testset "Copy names" begin
    model = LQOI.MockLinQuadOptimizer()
    x = MOI.add_variable(model)
    MOI.set(model, MOI.VariableName(), x, "x")
    @test MOI.get(model, MOI.VariableName(), x) == "x"
    c = MOI.add_constraint(model, MOI.SingleVariable(x), MOI.ZeroOne())
    MOI.set(model, MOI.ConstraintName(), c, "c")
    @test MOI.get(model, MOI.ConstraintName(), c) == "c"

    model2 = LQOI.MockLinQuadOptimizer()
    copy_map = MOI.copy_to(model2, model)
    @test MOI.get(model2, MOI.VariableName(), copy_map.varmap[x]) == "x"
    @test MOI.get(model2, MOI.ConstraintName(), copy_map.conmap[c]) == "c"
end


### A Pluto.jl notebook ###
# v0.14.6

using Markdown
using InteractiveUtils

# ╔═╡ 943cd3a0-b9a7-11eb-0852-fb73c86a7ebe
using CUDA, Zygote, Optim, Napari, TestImages, LinearAlgebra, FFTW, Noise, LineSearches, Tullio, Statistics, DeconvOptim, FourierTools, Colors, ImageShow, Images

# ╔═╡ 11c1d382-a7b0-4c84-b2fa-ee21ea2da201
md"## Import Packages"

# ╔═╡ 26439a7f-d22c-4b1e-a29c-5124c0a68094
FFTW.set_num_threads(6)

# ╔═╡ 0554827f-6a54-42d6-8a03-fd3a2d1112f5
md"## How does CUDA work?"

# ╔═╡ a14bacc8-a080-430b-a8fc-8084ce7a0e4d
begin
	x = rand(10, 10, 10)
	x_c = CuArray(x)
end;

# ╔═╡ e3d6933a-3ef4-4067-bc6a-570b5f5b1bb0
typeof(x)

# ╔═╡ b5e85f9d-e389-456a-b415-6303cffe7bdb
typeof(x_c)

# ╔═╡ 178b4a21-dd8b-4858-9337-5cf062fe4b09
fft(x);

# ╔═╡ d7347964-eda2-4ae2-80f9-ea2eeca640d8
fft(x_c);

# ╔═╡ fd0311fb-ea86-4e07-b3f9-0038fb539cc8
begin
	# prevent slow scalar indexing on GPU
	CUDA.allowscalar(false);

	# we need to fix some operations so that they are fast o GPUs
	# Reference: https://discourse.julialang.org/t/cuarray-and-optim/14053
	LinearAlgebra.norm1(x::CUDA.CuArray{T,N}) where {T,N} = sum(abs, x); # specializes the one-norm
	LinearAlgebra.normInf(x::CUDA.CuArray{T,N}) where {T,N} = maximum(abs, x); # specializes the one-norm
	Optim.maxdiff(x::CUDA.CuArray{T,N},y::CUDA.CuArray{T,N}) where {T,N} = maximum(abs.(x-y));
end

# ╔═╡ 7fcfc7b7-74b9-436b-9de9-e4b13da0087b
testimage("simple_3d_ball")

# ╔═╡ 8e481367-6fd5-4f63-b4e8-e5f84a70acd5
md"## Load example 3D image"

# ╔═╡ 5a242207-d570-43c7-bfef-20dc02b5a9ea
begin
	img = Float32.(testimage("simple_3d_ball"))
	img = abs.(resample(img, (128, 128, 128)))
	psf = permutedims(ifftshift(Float32.(testimage("simple_3d_psf").parent)), (3, 1, 2))
	psf = abs.(resample(psf, (128, 128, 128)))
	psf ./= sum(psf)
end;

# ╔═╡ b20ca70f-a979-4a53-ba9c-bf3b91dc5261
@view_image Float32.(testimage("simple_3d_psf").parent)

# ╔═╡ c979f867-0830-4e0b-aaca-3ccdc3ac2bc9
md"## Define a Loss Function"

# ╔═╡ eeb36b76-e9bc-4734-95d7-b1296f4050b1
function loss(rec, meas)
	return sum(abs2, rec .- meas)
end

# ╔═╡ 0f29fcc1-90f5-4446-b563-7765de99c177
md"## Create a forward deconvolution model

$$Y = (S * \text{PSF}) + N$$

measurement $Y$, sample $S$,  point spread function $\text{PSF}$, Noise $N$
"

# ╔═╡ d2dbe922-9dc1-4996-b753-62c8f7ff5a6b
begin
	otf = rfft(psf)

	function forward(img)
		return irfft(rfft(img) .* otf, size(img, 1))
	end

	otf_cuda = CuArray(otf)
	p = plan_rfft(CuArray(psf))
	p_inv = inv(p)

	function forward(img::CuArray)
		return p_inv.scale .* (p_inv.p * ((p * img) .* otf_cuda))
	end

	#return conv

end

# ╔═╡ cfc95c98-3f72-4b3a-86d6-6c73806330c6
forward(img);

# ╔═╡ 67e0992e-c504-4631-a510-3adb564a434b
img_c = CuArray(img);

# ╔═╡ 877a5075-0dac-456b-9539-3712abd884e0
forward(img_c);

# ╔═╡ ecf713de-7c69-4fc2-8bf2-00bdff7ddf04
md"## Create a noisy measurement"

# ╔═╡ a4a8e9e9-d68e-43f2-b575-f6b30861753b
begin
	measurement = poisson(forward(img), 50)
	measurement_cuda = CuArray(measurement)
end;

# ╔═╡ 09b5ae2e-0031-44fe-a29d-2c0b0105cb65
@view_image measurement

# ╔═╡ 88498069-35df-489d-9973-57695163fcdb
md"## Regularizer"

# ╔═╡ 7f9b85e3-a3e6-4120-8fd5-ce30debfbf82
reg(rec) = @tullio r = sqrt(1f-8 + abs2(rec[i, j, k] - rec[i + 1, j, k]) + abs2(rec[i, j, k] - rec[i, j + 1, k]) + abs2(rec[i, j, k] - rec[i, j, k + 1]))

# ╔═╡ fdb9ea82-2965-4993-9872-919bc2ed53e8
reg(rec::CuArray) = DeconvOptim.TV_cuda(num_dims=3)(rec)

# ╔═╡ dcc35fc0-b894-42aa-b67e-61bd66f1d7a2
begin
	f(rec) = loss(forward(rec.^2), measurement) + 0.001 * reg(rec.^2)
	f(rec::CuArray) = loss(forward(rec.^2), measurement_cuda) + 0.001 * reg(rec.^2)
	g!(G, rec) = G .= gradient(f, rec)[1]
end

# ╔═╡ 043efd53-e032-4f01-ab13-37a17b9d944c
begin
	rec0 = mean(measurement) .* ones(Float32, size(measurement));
	rec0_cuda = CuArray(rec0);
end;

# ╔═╡ f198cd1d-1fb1-4e00-b135-a69b9c54872b
begin
	f(rec0)
	f(rec0_cuda)
	g!(copy(rec0), rec0)
	g!(copy(rec0_cuda), rec0_cuda)
end;

# ╔═╡ eb9aaea3-b521-4b40-a3c8-04c893c1bf9c
md"## Optimizer"

# ╔═╡ beef10c8-2c24-4be2-8396-3024ff899d09
@time res = Optim.optimize(f, g!, rec0, LBFGS(linesearch=LineSearches.BackTracking()), Optim.Options(iterations=40))

# ╔═╡ f55b0aa2-4de9-43ab-9a30-3d1a7bc85984
CUDA.@time res_cuda = Optim.optimize(f, g!, rec0_cuda, LBFGS(linesearch=LineSearches.BackTracking()), Optim.Options(iterations=40))

# ╔═╡ 214fa627-9dd0-45c1-b766-77cbc2494a7d
begin
	rec = Optim.minimizer(res);
	rec_cuda = Optim.minimizer(res_cuda);
end

# ╔═╡ 4fc1d818-f2e8-428d-bd96-0fbfa9af3954
md"## Viewer"

# ╔═╡ de6df8ce-a6de-4910-a423-ec4e058bd79f
begin
	vr= napari.view_image(rec.^2)
	@add_image vr Array(rec_cuda.^2)
	@add_image vr measurement
	@add_image vr img
	vr.grid.enabled = true
end

# ╔═╡ Cell order:
# ╟─11c1d382-a7b0-4c84-b2fa-ee21ea2da201
# ╠═943cd3a0-b9a7-11eb-0852-fb73c86a7ebe
# ╠═26439a7f-d22c-4b1e-a29c-5124c0a68094
# ╠═0554827f-6a54-42d6-8a03-fd3a2d1112f5
# ╠═a14bacc8-a080-430b-a8fc-8084ce7a0e4d
# ╠═e3d6933a-3ef4-4067-bc6a-570b5f5b1bb0
# ╠═b5e85f9d-e389-456a-b415-6303cffe7bdb
# ╠═178b4a21-dd8b-4858-9337-5cf062fe4b09
# ╠═d7347964-eda2-4ae2-80f9-ea2eeca640d8
# ╟─fd0311fb-ea86-4e07-b3f9-0038fb539cc8
# ╠═7fcfc7b7-74b9-436b-9de9-e4b13da0087b
# ╟─8e481367-6fd5-4f63-b4e8-e5f84a70acd5
# ╠═5a242207-d570-43c7-bfef-20dc02b5a9ea
# ╠═b20ca70f-a979-4a53-ba9c-bf3b91dc5261
# ╟─c979f867-0830-4e0b-aaca-3ccdc3ac2bc9
# ╠═eeb36b76-e9bc-4734-95d7-b1296f4050b1
# ╟─0f29fcc1-90f5-4446-b563-7765de99c177
# ╠═d2dbe922-9dc1-4996-b753-62c8f7ff5a6b
# ╠═cfc95c98-3f72-4b3a-86d6-6c73806330c6
# ╠═877a5075-0dac-456b-9539-3712abd884e0
# ╠═67e0992e-c504-4631-a510-3adb564a434b
# ╠═ecf713de-7c69-4fc2-8bf2-00bdff7ddf04
# ╠═a4a8e9e9-d68e-43f2-b575-f6b30861753b
# ╠═09b5ae2e-0031-44fe-a29d-2c0b0105cb65
# ╠═88498069-35df-489d-9973-57695163fcdb
# ╠═7f9b85e3-a3e6-4120-8fd5-ce30debfbf82
# ╠═fdb9ea82-2965-4993-9872-919bc2ed53e8
# ╠═dcc35fc0-b894-42aa-b67e-61bd66f1d7a2
# ╠═f198cd1d-1fb1-4e00-b135-a69b9c54872b
# ╠═043efd53-e032-4f01-ab13-37a17b9d944c
# ╠═eb9aaea3-b521-4b40-a3c8-04c893c1bf9c
# ╠═beef10c8-2c24-4be2-8396-3024ff899d09
# ╠═f55b0aa2-4de9-43ab-9a30-3d1a7bc85984
# ╠═214fa627-9dd0-45c1-b766-77cbc2494a7d
# ╠═4fc1d818-f2e8-428d-bd96-0fbfa9af3954
# ╠═de6df8ce-a6de-4910-a423-ec4e058bd79f


module DynMultiply

export dynmultiply, dynmultiply_plan, dynmultiplystr

import Base.getindex, Base.setindex!

# The table contains at index N + ... + (N - l + 2) + s the entry
# corresponding to a subsequence with start index s and length l
type DynMultTable
    data::Array{Int, 1}
    N::Int
    DynMultTable(N) = new(zeros(Int, div(N*(N+1), 2)), N)
end

function getdataindex(table::DynMultTable, s::Int, l::Int)
    return (l-1)*table.N - div((l-1)*(l-2), 2) + s
end

# Accessor functions
function getindex(table::DynMultTable, s::Int, l::Int)
    return table.data[getdataindex(table, s, l)]
end

function setindex!(table::DynMultTable, val::Int, s::Int, l::Int)
    table.data[getdataindex(table, s, l)] = val
end

# Multiplication cost for multiplying group s, ..., j-1 and
# j, ..., f-1
function multiplycost(arrays, s::Int, j::Int, f::Int)
    size(arrays[s], 1) * size(arrays[j], 1) * size(arrays[f-1], 2)
end

# Compute multiplication order of matrices
function dynmultiply_plan(arrays...)
    N = length(arrays)
    costs = DynMultTable(N)
    plan = DynMultTable(N)
    for s=1:N
        costs[s, 1] = 0
    end

    # calculate costs and subsequence lengths
    for l= 2:N          # sequence length
        for s=1:N-l+1   # start index
            mincost = typemax(Int)
            l1opt = 0
            for l1=1:l-1    # first sequence length
                multcost = multiplycost(arrays, s, s+l1, s+l)
                cost = costs[s, l1] + costs[s+l1, l-l1] + multcost

                if cost < mincost
                    mincost = cost
                    l1opt = l1
                end
            end

            costs[s, l] = mincost
            plan[s, l] = l1opt
        end
    end

    return plan
end

# Multiply arrays in optimal order
function dynmultiply(arrays...)
    plan = dynmultiply_plan(arrays...)
    N = length(arrays)
    return dynmultiply_part(arrays, 1, N, plan)
end

# Multiply subsequence s, ..., s+l corresponding to plan
function dynmultiply_part(arrays, s::Int, l::Int, plan)
    if l == 1
        return arrays[s]
    else
        l1 = plan[s, l]
        return dynmultiply_part(arrays, s, l1, plan) * dynmultiply_part(arrays, s+l1, l-l1, plan)
    end
end

# Create string representation of optimal multiplication order
function dynmultiplystr(arrays...)
    plan = dynmultiply_plan(arrays...)
    N = length(arrays)
    return dynmultiplystr_part(1, N, plan)
end

# Create string representation of multiplication order of subsequence
# s, ..., s+l corresponding to plan
function dynmultiplystr_part(s::Int, l::Int, plan)
    if l == 1
        return "A$s"
    else
        l1 = plan[s, l]
        return "(" * dynmultiplystr_part(s, l1, plan) * "*" * dynmultiplystr_part(s+l1, l-l1, plan) * ")"
    end
end

end


export  computeSP,
        computeSP2

function computeSP(a_::AbstractVector{<:Int},
                   α::AbstractVector{<:AbstractVector{<:Real}},β::AbstractVector{<:AbstractVector{<:Real}},
                   nodes::AbstractVector{<:AbstractVector{<:Real}},weights::AbstractVector{<:AbstractVector{<:Real}},
                   ind::Matrix{<:Int};
                   issymmetric::BitArray=falses(length(α)),
                   zerotol::Float64=1e-10)
    minimum(a_) < 0 && throw(DomainError(minimum(a_),"no negative degrees allowed"))
    l, p = size(ind) # p-variate basis
    p == 1 && computeSP(a_,α[1],β[1],nodes[1],weights[1];issymmetric=issymmetric[1])
    l -= 1
    maximum(a_) > l && throw(DomainError(maximum(a_), "not enough elements in multi-index (requested: $(maximum(a_)), max: $l)"))
    !(length(α)==length(β)==length(nodes)==length(weights)==length(issymmetric)==p) && throw(InconsistencyError("inconsistent number of recurrence coefficients and/or nodes/weights"))

    a = Vector{Int64}()

    for aa in a_
        !iszero(aa) && push!(a,aa)
    end

    if iszero(length(a))
        return prod(β[i][1] for i in 1:p)
    elseif isone(length(a))
        return 0.
    else
        inds_uni = multi2uni(a,ind)
        val = 1.
        for i in 1:p
            # compute univariate scalar product
            @inbounds v = computeSP(inds_uni[i,:],α[i],β[i],nodes[i],weights[i];issymmetric=issymmetric[i])
            isapprox(v,0,atol=zerotol) ? (return 0.) : (val*=v)
        end
    end
    return val
end

function computeSP(a::AbstractVector{<:Int}, op::AbstractVector, ind::Matrix{<:Int})
    α, β = coeffs(op)
    nodes, weights = nw(op)
    computeSP(a, α, β, nodes, weights, ind; issymmetric=issymmetric.(op))
end

computeSP(a::AbstractVector{<:Int},mop::MultiOrthoPoly) = computeSP(a, mop.uni, mop.ind)

"""
__Univariate__
```
computeSP(a::AbstractVector{<:Int},α::AbstractVector{<:Real},β::AbstractVector{<:Real},nodes::AbstractVector{<:Real},weights::AbstractVector{<:Real};issymmetric::Bool=false)
computeSP(a::AbstractVector{<:Int},op::AbstractOrthoPoly;issymmetric=issymmetric(op))
```
__Multivariate__
```
computeSP( a::AbstractVector{<:Int},
           α::AbstractVector{<:AbstractVector{<:Real}},β::AbstractVector{<:AbstractVector{<:Real}},
           nodes::AbstractVector{<:AbstractVector{<:Real}},weights::AbstractVector{<:AbstractVector{<:Real}},
           ind::Matrix{<:Int};
           issymmetric::BitArray=falses(length(α)))
computeSP(a::AbstractVector{<:Int},op::AbstractVector,ind::Matrix{<:Int})
computeSP(a::AbstractVector{<:Int},mOP::MultiOrthoPoly)
```

Computes the scalar product
```math
\\langle \\phi_{a_1},\\phi_{a_2},\\cdots,\\phi_{a_n} \\rangle,
```
where `n = length(a)`.
This requires to provide the recurrence coefficients `(α,β)` and the quadrature rule
`(nodes,weights)`, as well as the multi-index `ind`.
If provided via the keyword `issymmetric`, symmetry of the weight function is exploited.
All computations of the multivariate scalar products resort back to efficient computations
of the univariate scalar products.
Mathematically, this follows from Fubini's theorem.

The function is dispatched to facilitate its use with `AbstractOrthoPoly` and its quadrature rule `Quad`.

!!! note
    - Zero entries of ``a`` are removed automatically to simplify computations, which follows from
    ```math
    \\langle \\phi_i, \\phi_j, \\phi_0,\\cdots,\\phi_0 \\rangle = \\langle \\phi_i, \\phi_j \\rangle,
    ```
    because `\\phi_0 = 1`.
    - It is checked whether enough quadrature points are supplied to solve the integral exactly.
"""
function computeSP(a_::AbstractVector{<:Int},α::AbstractVector{<:Real},β::AbstractVector{<:Real},nodes::AbstractVector{<:Real},weights::AbstractVector{<:Real};issymmetric::Bool=false)
    minimum(a_) < 0 && throw(DomainError(minimum(a_), "no negative degrees allowed"))
    # this works because ``<Φ_i,Φ_j,Φ_0,...,Φ_0> = <Φ_i,Φ_j>``
    # hence, avoiding unnecessary operations that introduce numerical gibberish
    # a = Vector{Int64}()
    # for aa in a_
    #     !iszero(aa) && push!(a,aa)
    # end
    a = filter(x -> x > 0, a_)
    # simplification in case the measure is symmetric w.r.t t=0
    # exploit symmetry of the density, see Theorem 1.17 in W. Gautschi's book
    # and exploit the fact that
    (issymmetric && isodd(sum(a))) && return 0.
    # Gauss quadrature rules have exactness 2N-1
    !(Int(ceil(0.5*(sum(a)+1)))<=length(nodes)) && throw(InconsistencyError("not enough nodes to integrate exactly ($(length(nodes)) provided, where $(Int(ceil(0.5*(sum(a)+1)))) are needed)"))

    res =
    if iszero(length(a))
        first(β)
    elseif isone(length(a))
        0.
    elseif length(a) == 2
        @inbounds a[1] == a[2] ? computeSP2(first(a),β)[end] : 0.
    else
        f = ones(Float64,length(nodes))
        @simd for i = 1:length(a)
            @inbounds f = f.*evaluate(a[i],nodes,α,β)
        end
        dot(weights,f)
    end
end

function computeSP(a::AbstractVector{<:Int},op::AbstractOrthoPoly)
    computeSP(a,op.α,op.β,op.quad.nodes,op.quad.weights;issymmetric=issymmetric(op))
end

"""
    computeSP2(n::Int,β::AbstractVector{<:Real})
    computeSP2(n::Int,op::AbstractOrthoPoly) = computeSP2(n,op.β)
    computeSP2(op::AbstractOrthoPoly) = computeSP2(op.deg,op.β)

Computes the `n` *regular* scalar products aka 2-norms of the orthogonal polynomials, namely
```math
\\|ϕ_i\\|^2 = \\langle \\phi_i,\\phi_i\\rangle \\quad \\forall i \\in \\{ 0,\\dots,n \\}.
```
Notice that only the values of `β` of the recurrence coefficients `(α,β)` are required.
The computation is based on equation (1.3.7) from Gautschi, W. "Orthogonal Polynomials: Computation and Approximation".
Whenever there exists an analytic expressions for `β`, this function should be used.

The function is multiply dispatched to facilitate its use with `AbstractOrthoPoly`.
"""
function computeSP2(n::Int,β::AbstractVector{<:Real})
    n < 0 && throw(DomainError(n, "can only compute scalar products for non-negative degrees"))
    length(β) < n + 1 && throw(InconsistencyError("inconsistent length of β"))
    n == 0 && return β[1]
    s = ones(Float64, n+1)
    @inbounds s[1] = β[1]
    for i in 2:n+1
        @inbounds s[i] = s[i-1]*β[i]
    end
    return s
end
computeSP2(n::Int,op::AbstractOrthoPoly) = computeSP2(n,op.β)
computeSP2(op::AbstractOrthoPoly) = computeSP2(op.deg,op.β)


using ProblemParser
using Utils

using Lazy
using Underscores

const file = find_input(@__FILE__)
const input = @>> file slurp parse(Lines())
@assert @_ map(length(_), input) |> Set |> length == 1

function part1(input)
    matrix = @_ input |>
        map(map(==('1'), collect(_)), __) |>
        reduce(hcat, __) |>
        transpose
    len = size(matrix, 1)

    gamma =
        @_ [sum(matrix[:,r]) > len ÷ 2 for r in axes(matrix, 2)] |>
        map(curry(+, '0'), __) |>
        String |>
        parse(Int, __; base=2)
    epsilon = gamma ⊻ (2^size(matrix, 2) - 1)
    gamma * epsilon
end
assertequal(part1(input), 693486)

function sortdown(digit1, input)
    input = sort(input)
    pos = 1
    while length(input) > 1
        first1 = @_ findfirst(_[pos] == '1', input)
        num0 = first1 - 1
        if digit1 ⊻ (num0 <= length(input) ÷ 2)
            input = input[begin:num0]
        else
            input = input[first1:end]
        end
        pos += 1
    end
    parse(Int, input[begin]; base=2)
end

function part2(input)
    oxygen = sortdown(true, input)
    co2 = sortdown(false, input)
    oxygen * co2
end
assertequal(part2(input), 3379326)


# example modified from  mthelm85
# by lazarusA # HIDE
# using GLMakie # HIDE
using GLMakie, ForwardDiff
GLMakie.activate!()
let
    f(x,y) = -5*x*y*exp(-x^2-y^2)
    x = -1:0.05:1.0
    y = -1:0.05:1.0
    z = [f(i,j) for i in x, j in y]
    # This is the same function as above, just modified so that it will
    # work with ForwardDiff
    g(x,y) = [-5*x*y*exp(-x^2-y^2)]
    J(xx,yy) = ForwardDiff.jacobian(x -> g(x[1], x[2]), [xx,yy])
    field(i,j) = Point(J(i,j)[1], J(i,j)[2])

    zmin, zmax = minimum(z), maximum(z)
    cmap = :viridis
    function plot()
        fig = Figure(resolution = (1200,600))
        ax1 = Axis3(fig, aspect = (1,1,1), perspectiveness = 0.5,
            elevation = π/3.5, azimuth = 0.1π, )
        ax2 = Axis(fig, aspect = DataAspect(), xlabel = "x", ylabel = "y")
        surface!(ax1, x, y, z, colormap = cmap, colorrange = (zmin, zmax))
        contour3d!(ax1, x, y, z .+ 0.005, levels = 15, linewidth = 2, color=(:white, 0.5))
        wireframe!(ax1, x, y, z, overdraw = false, transparency = true, color = (:black, 0.1))
        streamplot!(ax1, field, -1..1, -1..1, colormap = cmap, gridsize= (40,40),
            arrow_size = 0.05,linewidth=1, transformation = (:xy, -zmax))
        streamplot!(ax2, field, -1..1, -1..1, colormap = cmap, gridsize= (40,40),
                arrow_size = 0.03,linewidth=1)
        fig[1,1] = ax1
        fig[1,2] = ax2
        fig
    end
    fig = with_theme(plot, theme_dark())
    save(joinpath(@__DIR__, "output", "theme_dark_Contour3d.png"), fig, px_per_unit = 2) # HIDE
end
using Pkg # HIDE
Pkg.status(["GLMakie","ForwardDiff"]) # HIDE


using CmdStan, DiffEqBayes, OrdinaryDiffEq, ParameterizedFunctions,
      ModelingToolkit, RecursiveArrayTools, Distributions, Random, Test

# Uncomment for local testing only, make sure MCMCChains and StatsPlots are available
using MCMCChains, StatsPlots

#Random.seed!(123)
cd(@__DIR__)
isdir("tmp") && rm("tmp", recursive=true)

include("../../stan_inference.jl")

println("\nOne parameter case 1\n")
f1 = @ode_def begin
  dx = a*x - x*y
  dy = -3y + x*y
end a
u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5]
prob1 = ODEProblem(f1,u0,tspan,p)
sol = solve(prob1,Tsit5())
t = collect(range(1,stop=10,length=10))
randomized = VectorOfArray([(sol(t[i]) + .5randn(2)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [truncated(Normal(0.7,1),0.1,2)]

bayesian_result = stan_inference(prob1,t,data,priors;num_samples=2000, nchains=4,
  num_warmup=1000,likelihood=Normal)

sdf  = CmdStan.read_summary(bayesian_result.model)
@test sdf[sdf.parameters .== :theta1, :mean][1] ≈ 1.5 atol=3e-1

# Uncomment for local chain inspection
chn = CmdStan.convert_a3d(bayesian_result.chains, bayesian_result.cnames, Val(:mcmcchains))
plot(chn)
!isdir("tmp") && mkdir("tmp")
savefig("$(@__DIR__)/one_par_case_01.png")

v1 = [sdf[sdf.parameters .== Symbol("u_hat[$i,1]"), :mean][1] for i in 1:10]
v2 = [sdf[sdf.parameters .== Symbol("u_hat[$i,2]"), :mean][1] for i in 1:10]
qa_prey = zeros(10, 3);
qa_pred = zeros(10, 3);
achn = Array(chn);
for i in 1:10
  qa_prey[i, :] = quantile(achn[:, 4+i], [0.055, 0.5, 0.945])
  qa_pred[i, :] = quantile(achn[:, 14+i], [0.055, 0.5, 0.945])
end
p1 = plot(1:10, v1; ribbon=(qa_prey[:, 1], qa_prey[:, 3]), color=:lightgrey, leg=false)
title!("Prey u_hat 89% quantiles")
plot!(v1, lab="prey", xlab="time", ylab="prey", color=:darkred)

p2 = plot(1:10, v2; ribbon=(qa_pred[:, 1], qa_pred[:, 3]), color=:lightgrey, leg=false)
title!("Preditors u_hat 89% quantiles")
plot!(v2, lab="pred", xlab="time", ylab="preditors", color=:darkblue)

plot(p1, p2, layout=(2,1))
savefig("$(@__DIR__)/one_par_pred_prey_01.png")

println("\nOne parameter case 2\n")
priors = [
  truncated(Normal(1.,0.5), 0.1, 3),
  truncated(Normal(1.,0.5), 0.1, 3),
  truncated(Normal(1.5,0.5), 0.1, 4)
]
bayesian_result = stan_inference(prob1,t,data,priors;num_samples=2000, nchains=4,
  num_warmup=1000,likelihood=Normal,sample_u0=true)

sdf  = CmdStan.read_summary(bayesian_result.model)
@test sdf[sdf.parameters .== :theta1, :mean][1] ≈ 1. atol=3e-1
@test sdf[sdf.parameters .== :theta2, :mean][1] ≈ 1. atol=3e-1
@test sdf[sdf.parameters .== :theta3, :mean][1] ≈ 1.5 atol=3e-1

# Uncomment for local chain inspection
chn = CmdStan.convert_a3d(bayesian_result.chains, bayesian_result.cnames, Val(:mcmcchains))
plot(chn)
!isdir("tmp") && mkdir("tmp")
savefig("$(@__DIR__)/one_par_case_02.png")

v1 = [sdf[sdf.parameters .== Symbol("u_hat[$i,1]"), :mean][1] for i in 1:10]
v2 = [sdf[sdf.parameters .== Symbol("u_hat[$i,2]"), :mean][1] for i in 1:10]
qa_prey = zeros(10, 3);
qa_pred = zeros(10, 3);
achn = Array(chn);
for i in 1:10
  qa_prey[i, :] = quantile(achn[:, 8+i], [0.055, 0.5, 0.945])
  qa_pred[i, :] = quantile(achn[:, 18+i], [0.055, 0.5, 0.945])
end
p1 = plot(1:10, v1; ribbon=(qa_prey[:, 1], qa_prey[:, 3]), color=:lightgrey, leg=false)
title!("Prey u_hat 89% quantiles")
plot!(v1, lab="prey", xlab="time", ylab="prey", color=:darkred)

p2 = plot(1:10, v2; ribbon=(qa_pred[:, 1], qa_pred[:, 3]), color=:lightgrey, leg=false)
title!("Preditors u_hat 89% quantiles")
plot!(v2, lab="pred", xlab="time", ylab="preditors", color=:darkblue)

plot(p1, p2, layout=(2,1))
savefig("$(@__DIR__)/one_par_pred_prey_02.png")

println("\nOne parameter case 3\n")

sol = solve(prob1,Tsit5(),save_idxs=[1])
randomized = VectorOfArray([(sol(t[i]) + .01 * randn(1)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [truncated(Normal(1.5,0.5), 0.1,2)]
bayesian_result = stan_inference(prob1,t,data,priors;num_samples=2000, nchains=4,
  num_warmup=1000,likelihood=Normal,save_idxs=[1])

sdf  = CmdStan.read_summary(bayesian_result.model)
@test sdf[sdf.parameters .== :theta1, :mean][1] ≈ 1.5 atol=3e-1

# Uncomment for local chain inspection
chn = CmdStan.convert_a3d(bayesian_result.chains, bayesian_result.cnames, Val(:mcmcchains))
plot(chn)
savefig("$(@__DIR__)/one_par_case_03.png")

v1 = [sdf[sdf.parameters .== Symbol("u_hat[$i,1]"), :mean][1] for i in 1:10]
v2 = [sdf[sdf.parameters .== Symbol("u_hat[$i,2]"), :mean][1] for i in 1:10]
qa_prey = zeros(10, 3);
qa_pred = zeros(10, 3);
achn = Array(chn);
for i in 1:10
  qa_prey[i, :] = quantile(achn[:, 3+i], [0.055, 0.5, 0.945])
  qa_pred[i, :] = quantile(achn[:, 13+i], [0.055, 0.5, 0.945])
end
p1 = plot(1:10, v1; ribbon=(qa_prey[:, 1], qa_prey[:, 3]), color=:lightgrey, leg=false)
title!("Prey u_hat 89% quantiles")
plot!(v1, lab="prey", xlab="time", ylab="prey", color=:darkred)

p2 = plot(1:10, v2; ribbon=(qa_pred[:, 1], qa_pred[:, 3]), color=:lightgrey, leg=false)
title!("Preditors u_hat 89% quantiles")
plot!(v2, lab="pred", xlab="time", ylab="preditors", color=:darkblue)

plot(p1, p2, layout=(2,1))
savefig("$(@__DIR__)/one_par_pred_prey_03.png")

println("\nOne parameter case 4\n")

priors = [truncated(Normal(1.,0.5), 0.1, 3),
  truncated(Normal(1.5,0.5), 0.1, 4)]
bayesian_result = stan_inference(prob1,t,data,priors;num_samples=2000,
  nchains=4, num_warmup=1000,likelihood=Normal,save_idxs=[1],sample_u0=true)

sdf  = CmdStan.read_summary(bayesian_result.model)
@test sdf[sdf.parameters .== :theta1, :mean][1] ≈ 1. atol=3e-1
@test sdf[sdf.parameters .== :theta2, :mean][1] ≈ 1.5 atol=3e-1

# Uncomment for local chain inspection
chn = CmdStan.convert_a3d(bayesian_result.chains, bayesian_result.cnames, Val(:mcmcchains))
plot(chn)
savefig("$(@__DIR__)/one_par_case_04.png")

v1 = [sdf[sdf.parameters .== Symbol("u_hat[$i,1]"), :mean][1] for i in 1:10]
v2 = [sdf[sdf.parameters .== Symbol("u_hat[$i,2]"), :mean][1] for i in 1:10]
qa_prey = zeros(10, 3);
qa_pred = zeros(10, 3);
achn = Array(chn);
for i in 1:10
  qa_prey[i, :] = quantile(achn[:, 4+i], [0.055, 0.5, 0.945])
  qa_pred[i, :] = quantile(achn[:, 14+i], [0.055, 0.5, 0.945])
end
p1 = plot(1:10, v1; ribbon=(qa_prey[:, 1], qa_prey[:, 3]), color=:lightgrey, leg=false)
title!("Prey u_hat 89% quantiles")
plot!(v1, lab="prey", xlab="time", ylab="prey", color=:darkred)

p2 = plot(1:10, v2; ribbon=(qa_pred[:, 1], qa_pred[:, 3]), color=:lightgrey, leg=false)
title!("Preditors u_hat 89% quantiles")
plot!(v2, lab="pred", xlab="time", ylab="preditors", color=:darkblue)

plot(p1, p2, layout=(2,1))
savefig("$(@__DIR__)/one_par_pred_prey_04.png")


import GeometricFlux: message, update, propagate

in_channel = 10
out_channel = 5
N = 6
adj = [0. 1. 0. 0. 0. 0.;
       1. 0. 0. 1. 1. 1.;
       0. 0. 0. 0. 0. 1.;
       0. 1. 0. 0. 1. 0.;
       0. 1. 0. 1. 0. 1.;
       0. 1. 1. 0. 1. 0.]
ne = [[2], [1,4,5,6], [6], [2,5], [2,4,6], [2,3,5]]

struct NewLayer <: MessagePassing
    adjlist::AbstractVector{<:AbstractVector}
    weight
end

NewLayer(adjm::AbstractMatrix, m, n) = NewLayer(neighbors(adjm), randn(m,n))

(l::NewLayer)(X) = propagate(l, X=X, aggr=:add)
message(n::NewLayer; x_i=zeros(0), x_j=zeros(0)) = x_j
update(::NewLayer; X=zeros(0), M=zeros(0)) = M

X = Array(reshape(1:N*in_channel, in_channel, N))
l = NewLayer(adj, in_channel, out_channel)

@testset "Test neighboring" begin
    Y, clst = neighboring(l; X=X)
    @test adjlist(l) == ne
    @test size(Y) == (in_channel, 14)
    @test Y[1, :] == [11, 1, 31, 41, 51, 51, 11, 41, 11, 31, 51, 11, 21, 41]
    @test clst == [1, 2, 2, 2, 2, 3, 4, 4, 5, 5, 5, 6 ,6, 6]
end

message(n::NewLayer; x_i=zeros(0), x_j=zeros(0)) = n.weight' * x_j

@testset "Test MessagePassing layer" begin
    Y = l(X)
    @test size(Y) == (out_channel, N)
end

in_channel = 100
X = Array(reshape(1:N*in_channel, in_channel, N))
l = NewLayer(adj, in_channel, out_channel)

@testset "Test multi-thread MessagePassing layer" begin
    Y = l(X)
    @test size(Y) == (out_channel, N)
end


"""
ConstantsFS object for FAM with feature selection.
Augments Constants object with new priors:
W_star_prior, p_prior, and omega_prior.
Note that only one of p_prior or omega_prior will be used.
"""
mutable struct ConstantsFS{T <: ContinuousDistribution}
  W_star_prior::Gamma # W*_ik ~ Gamma(a_W, 1)
  p_prior::Beta # p_c ~ Beta(a_p, b_p)
  omega_prior::T # omage_each ~ Normal(m_omega , s_omega)
  constants::Constants
end

function ConstantsFS(c::Constants)
  # ws_prior = Gamma(1. / c.K, 1.)  # Old version.
  ws_prior = Gamma(10., 1.)
  p_prior = Beta(.1, .1)
  omega_prior = Normal(-5, 1.)

  return ConstantsFS(ws_prior, p_prior, omega_prior, c)
end

function printConstants(c::ConstantsFS, preprintln::Bool=true)
  if preprintln
    println("ConstantsFS:")
  end

  _fieldnames = collect(fieldnames(typeof(c)))

  for fname in filter(fn -> fn != :constants, _fieldnames)
    x = getfield(c, fname)
    T = typeof(x)
    if T <: Number
      println("$fname: $x")
    elseif T <: Vector
      N = length(x)
      for i in 1:N
        println("$(fname)[$i]: $(x[i])")
      end
    elseif T <: Dict
      for (k, v) in x
        println("$(fname)[$k]: $v")
      end
    else
      println("$fname: $x")
    end
  end

  printConstants(c.constants, true)
end