export pick_inN2, pick_inN, pick_inN3

"""
pick\\_inN2 is the picking scheme which picks the Newton point whenever possible.

"""
function pick_inN2(a, b, ϕ, ϕa, ϕb, dϕa, dϕb, forward)
    γ = 0.9
    L = b - a
    if forward
        dN = - (dϕa / hess(ϕ, a))
        if (dN > 0) & (dN < γ*L)
            return a + dN
        elseif b<Inf
            return (a + b)/2.0
        else
            return pick_inInf(a)
        end
    else
        dN =  - (dϕb / hess(ϕ, b))
        if (dN < 0) & (dN > -γ*L)
            return b + dN
        else
            return (a + b)/2.0
        end
    end
end

"""
pick\\_inN is the picking scheme which picks the Newton point whenever possible but
         uses pick\\_inInf whenever b=∞

Pick\\_inN2 is more efficient.

"""
function pick_inN(a, b, ϕ, ϕa, ϕb, dϕa, dϕb, forward)
    γ = 0.9
    if b<Inf
        L = b - a
        if forward
            dN = - (dϕa / hess(ϕ, a))
            if (dN > 0) & (dN < γ*L)
                return a + dN
            else
                return (a + b)/2.0
            end
        else
            dN =  - (dϕb / hess(ϕ, b))
            if (dN < 0) & (dN > -γ*L)
                return b + dN
            else
                return (a + b)/2.0
            end
        end
    else
        return pick_inInf(a)
    end
end

# Tricky implementation to keep the predecessor iterate for inclusion in the secant formula
#
#

let (tpred, ϕpred, dϕpred, d2ϕpred) = (Inf, Inf, Inf, Inf)
    γ = 0.9
    global function resetN3!(;t::Float64 = Inf, vϕ::Float64 = Inf, dϕ::Float64 = Inf, d2ϕ::Float64 = Inf)
        (tpred, ϕpred, dϕpred, d2ϕpred) = (t, vϕ, dϕ, d2ϕ)
    end

    # based on pick_inS2Opt: filters first the first Newton step and then tries the second.

"""
pick\\_inN3 is the picking scheme which picks the Newton point whenever possible and then eventually
     performs a further pseudo-Newton iteration using a quintic Hermite interpolant

"""
    global function pick_inN3(a, b, ϕ, ϕa, ϕb, dϕa, dϕb, forward)
        γ = 0.9
        L = b - a
        if forward
            d2ϕa = hess(ϕ, a)
            dN = - (dϕa / d2ϕa)
            t2 = a + dN
            # first the Newton's improvement
            if (dN > 0) & (dN < γ*L)
                #prepare to improve the Newton step
                α = coefHIN5([a;tpred],[ϕa dϕa d2ϕa; ϕpred dϕpred d2ϕpred])
                p = Polynomial(α)
                dp = Polynomials.derivative(p)
                d2p = Polynomials.derivative(dp)
                d2 = -(dp(dN) / d2p(dN))
                dN2 = dN + d2
                d3 = -(dp(dN2) / d2p(dN2))
                dN3 = dN2 + d3
                d4 = -(dp(dN3) / d2p(dN3))
                dN4 = dN3 + d4
               if (dN4 > 0) & (dN4 < γ*L)
                    t2 = a + dN4
                end
            elseif b<Inf
                t2 = (a + b)/2.0
            else
                t2 = pick_inInf(a)
            end
            tpred, ϕpred, dϕpred, d2ϕpred = a, ϕa, dϕa, d2ϕa
            return t2
        else
            d2ϕb = hess(ϕ, b)
            dN =  - (dϕb / d2ϕb)
            t2 = b + dN
            # first the Newton's improvement
            if (dN < 0) & (dN > -γ*L)
                 #prepare to improve the Newton step
                α = coefHIN5([b;tpred],[ϕb dϕb d2ϕb; ϕpred dϕpred d2ϕpred])
                p = Polynomial(α)
                dp = Polynomials.derivative(p)
                d2p = Polynomials.derivative(dp)
                d2 = -(dp(dN) / d2p(dN))
                dN2 = dN + d2
                d3 = -(dp(dN2) / d2p(dN2))
                dN3 = dN2 + d3
                d4 = -(dp(dN3) / d2p(dN3))
                dN4 = dN3 + d4
                if (dN4 > 0) & (dN4 < γ*L)
                    t2 = b + dN4
                end
            else
                t2 = (a + b)/2.0
            end
            tpred, ϕpred, dϕpred, d2ϕpred = b, ϕb, dϕb, d2ϕb
            return t2
        end
    end

end

"""
pick\\_inN3 is the picking scheme which picks the Newton point whenever possible and then eventually
     performs a further pseudo-Newton iteration using a quintic Hermite interpolant

"""
pick_inN3


module ParallelAnalysis

using MultivariateStats, StatsFuns, Statistics, ProgressMeter
import Statistics: cov
import MultivariateStats: loadings
import Base: show
using StatsBase: sample
# Polychoric
using Trapz: trapz#, @trapz
using Optim: optimize, Brent
using LinearAlgebra: LowerTriangular, diagind, Symmetric
using QuadGK: quadgk
using Distributions, DataFrames
# Plot recipes
using RecipesBase

include("Polychoric.jl")
export polycor, cov, loadings, communalities

include("RandomMatrix.jl")
export random_matrix, random_sample

include("FA.jl")
export fa

include("IRT.jl")
export generate_response, v2vv, heuristicIRT

include("PA.jl")
export parallel

export show

end


import Base.conj
export conj
export sign, curvature, monotonicity, evaluate, conic_form!
struct ConjugateAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int, Int}

    function ConjugateAtom(x::AbstractExpr)
        children = (x,)
        return new(:conj, hash(children), children, (x.size[1], x.size[2]))
    end
end

function sign(x::ConjugateAtom)
    return sign(x.children[1])
end

function monotonicity(x::ConjugateAtom)
    return (Nondecreasing(),)
end

function curvature(x::ConjugateAtom)
    return ConstVexity()
end

function evaluate(x::ConjugateAtom)
    return conj(evaluate(x.children[1]))
end

function conic_form!(x::ConjugateAtom, unique_conic_forms::UniqueConicForms)
    if !has_conic_form(unique_conic_forms, x)
        objective = conic_form!(x.children[1], unique_conic_forms)
        for var in keys(objective)
            x1 = conj(objective[var][1])
            x2 = conj(objective[var][2])
            objective[var] = (x1,x2)
        end
        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

conj(x::AbstractExpr) = ConjugateAtom(x)
conj(x::Constant) = Constant(conj(x.value))


using ModelingToolkitStandardLibrary.Thermal, ModelingToolkit, OrdinaryDiffEq, Test
@parameters t
D = Differential(t)

# Modelica example
begin
    @named mass1 = HeatCapacitor(C=15, T_start=373.15)
    @named mass2 = HeatCapacitor(C=15, T_start=273.15)
    @named conduction = ThermalConductor(G=10)
    @named Tsensor1 = TemperatureSensor()
    @named Tsensor2 = TemperatureSensor()

    connections = [
        connect(mass1.port, conduction.port_a),
        connect(conduction.port_b, mass2.port),
        connect(mass1.port, Tsensor1.port),
        connect(mass2.port, Tsensor2.port),
    ]

    @named model = ODESystem(connections, t, systems=[mass1, mass2, conduction, Tsensor1, Tsensor2])
    sys = structural_simplify(model)
    prob = DAEProblem(sys, D.(states(sys)) .=> 0.0, [mass1.der_T => 1.0, mass2.der_T => 1.0], (0, 3.0))
    sol = solve(prob, DFBDF())
    @test sol.retcode == :Success
end


using RiemannHilbert, ApproxFun, ApproxFunBase, LinearAlgebra, SingularIntegralEquations, Test
import ApproxFunBase: pieces

@testset "NLS" begin
    x,t = 0.0,0.0
    r = k -> 0.5exp(-k^2)
    r̄ = k -> conj(r(conj(k)))
    θ = k -> 4t*k^2+2x*k
    Gf = k -> [1-abs2(r(k))                      -conj(r(k))*exp(-im*θ(k));
                    r(k)*exp(im*θ(k))      1  ]

    g = Fun(Gf, -5.5..5.5)
    @test first(g) ≈ last(g) ≈ [1 0; 0 1]
    φ = transpose(rhsolve(transpose(g), 500))
    istieltjes(φ)(0.1)
    @test φ(1+im) ≈ [0.9847854664616504-0.014938097892888065im -0.0778564768587141-0.05302708037085426im;
                0.07824233623082648+0.05395699475793745im 1.0120649351483406+0.006873030868410332im] # mathematica

    x,t = 1.0,1.0

    g = Fun(Gf, -5.5..5.5)
    @test first(g) ≈ last(g) ≈ [1 0; 0 1]
    φ = transpose(rhsolve(transpose(g), 1000))
    @test φ(1+im) ≈ [0.9952708020858667-0.00828828357987084im -0.036773740429971044-0.01579754685573759im;
                0.009328821936648003+0.04514458247615043im 1.0050689684367242+ 0.006553783105674419im]

    @test norm(φ.coefficients[end-50:end]) ≤ 1E-14
    @test φ(0.1⁺) ≈ φ(0.1⁻)g(0.1)

    z₀ = -x/(4t)
    Γ = Segment(z₀,z₀+3exp(3im*π/4)) ∪ Segment(z₀,-5.5) ∪ Segment(z₀,z₀+3exp(-3im*π/4)) ∪ Segment(z₀,z₀+3exp(-im*π/4)) ∪ Segment(z₀,z₀+3exp(im*π/4))

    U = z -> [1 -r̄(z)/(1-r(z)r̄(z))*exp(-im*θ(z)); 0 1]
    D = z -> [(1-r(z)r̄(z)) 0; 0 1/(1-r(z)r̄(z))]
    L = z -> [1 0; r(z)/(1-r(z)r̄(z))*exp(im*θ(z)) 1]
    M = z -> [1 -r̄(z)*exp(-im*θ(z)); 0 1]
    P = z -> [1 0; r(z)*exp(im*θ(z)) 1]
    G = Fun(function(z)
            z ∈ component(Γ,1) && return inv(U(z))
            z ∈ component(Γ,2) && return inv(D(z))
            z ∈ component(Γ,3) && return inv(L(z))
            z ∈ component(Γ,4) && return M(z)
            z ∈ component(Γ,5) && return P(z)
            error("Not in contour")
        end, Γ)

    @test prod(first.(pieces(G))) ≈ [1 0; 0 1]
    @test norm(last.(pieces(G)) .- Ref(I)) ≤ 1E-14

    Φ = transpose(rhsolve(transpose(G), 3000))
    @test norm(Φ.coefficients[end-50:end]) ≤ 10E-14

    @test Φ(2im) ≈ φ(2im)

    z = Fun(ℂ)
    2im*(z*Φ[1,2])(Inf)≈ 2im*(z*φ[1,2])(Inf)

    # make evaluation faster
    φ2 = Matrix(φ)
    φ̃ = z -> [φ2[1,1](z) φ2[1,2](z); φ2[2,1](z) φ2[2,2](z)]

    @time V = Fun(function(z)
            z ∈ component(Γ,1) && return φ̃(z)*(inv(U(z))-I)
            z ∈ component(Γ,2) && return φ̃(z-10eps()im)L(z)-φ̃(z+10eps()im)inv(U(z))
            z ∈ component(Γ,3) && return φ̃(z)*(I-L(z))
            z ∈ component(Γ,4) && return φ̃(z)*(M(z)-I)
            z ∈ component(Γ,5) && return φ̃(z)*(I-inv(P(z)))
            error("Not in contour")
        end, Γ, 1000)

    @test norm(V.coefficients[end-20:end]) ≤ 1E-15
    @test P(z₀+3exp(im*π/4)) ≈ [1 0; 0 1]
    Φ = I+cauchy(V)
    Φ̃ = function(z)
            3π/4  < angle(z-z₀) < π     && return φ(z)inv(U(z))
            -π    < angle(z-z₀) < -3π/4 && return φ(z)L(z)
            -3π/4 < angle(z-z₀) < -π/4  && return φ(z)
            -π/4  < angle(z-z₀) <  0    && return φ(z)M(z)
            0    < angle(z-z₀) <  π/4  && return φ(z)inv(P(z))
            π/4  < angle(z-z₀) < 3π/4  && return φ(z)
            error("Not defined")
        end

    @test Φ̃(0.1+10eps()im) ≈ Φ̃(0.1-10eps()im)

    z = z₀+0.1exp(im*π/4); @test Φ(z-10eps())-Φ(z+10eps()) ≈ Φ̃(z-10eps())-Φ̃(z+10eps()) ≈ φ̃(z)*(I-inv(P(z)))
    z = z₀+0.1exp(3im*π/4); @test Φ(z-10eps())-Φ(z+10eps()) ≈ Φ̃(z-10eps())-Φ̃(z+10eps()) ≈ φ̃(z)*(inv(U(z))-I)
    z = z₀-0.1+0im; @test Φ(z-10im*eps())-Φ(z+10im*eps()) ≈ Φ̃(z-10im*eps())-Φ̃(z+10im*eps()) ≈ φ̃(z-10im*eps())*L(z)-φ̃(z+10im*eps())*inv(U(z))
    z = z₀+0.1exp(-3im*π/4); @test Φ(z+10eps())-Φ(z-10eps()) ≈ Φ̃(z+10eps())-Φ̃(z-10eps()) ≈ φ̃(z)*(I-L(z))
    z = z₀+0.1exp(-im*π/4); @test Φ(z+10eps())-Φ(z-10eps()) ≈ Φ̃(z+10eps())-Φ̃(z-10eps()) ≈ φ̃(z)*(M(z)-I)
    z = z₀+0.1+0im; @test Φ(z-10im*eps()) ≈ Φ(z+10im*eps()) ≈ Φ̃(z-10im*eps()) ≈ Φ̃(z+10im*eps())
end


# ------------------------------------------------
# import
# ------------------------------------------------
using Random
using Printf
#using Distributions

# ------------------------------------------------
# main
# ------------------------------------------------
function main()
    # -------------------------------------
    # file
    # -------------------------------------
    f_train = "train.txt"
    f_test  = "test.txt"

    f_re_train = "corr_history_train"
    f_re_test  = "corr_history_test"

    # -------------------------------------
    # input & output
    # -------------------------------------
    in_n       = 10           # +1 for b
    fout_n     = 2            # final output
    out_n_at_m = [in_n, 6, fout_n]          # number of cell at middle layers

    # -------------------------------------
    # nn para
    # -------------------------------------
    num_mlayers = 1   # middle layer
    mu = 0.9
    ep = 0.5

    # -------------------------------------
    # init weight
    # -------------------------------------
    mean  = 0.0    # average
    stdev = 1.0    # variance value

    # -------------------------------------
    # nn loop
    # -------------------------------------
    main_ln  = 1000
    online_ln = 10

    # -------------------------------------
    # cal
    # -------------------------------------
    for i in 1:length(out_n_at_m)
        out_n_at_m[i] += 1
    end
    out_n = maximum(out_n_at_m)
    num_layers = num_mlayers + 2

    # ---------------------------------------------------------------------------------------
    # -------------------------------------
    # setup w
    # -------------------------------------
    w_old  = ones(num_layers-1, out_n, out_n)
    w_new  = ones(num_layers-1, out_n, out_n)
    w_temp = ones(num_layers-1, out_n, out_n)

    # dis of gauss
    for i in 1:num_layers-1
        for j in 2:out_n
            for k in 1:out_n
                #rng = MersenneTwister(1234)
                w_old[i,j,k] = rand_normal(mean, stdev)
                w_new[i,j,k] = rand_normal(mean, stdev)
            end
        end
    end

    # -------------------------------------
    # setup for nn
    # -------------------------------------
    u     = zeros(num_layers  , out_n)   # sum result
    z     = zeros(num_layers  , out_n)   # output
    temp  = zeros(out_n)                 # output
    delta = zeros(num_layers  , out_n)
    dEdw  = zeros(num_layers-1, out_n, out_n)

    roca_test  = zeros(main_ln)
    roca_train = zeros(main_ln)

    # -------------------------------------
    # read file
    # -------------------------------------
    train_re, train_data, train_ans = read_file(f_train)
    test_re,  test_data,  test_ans  = read_file(f_test)
    s = size(train_re)[1]

    # -------------------------------------
    # machine learning by NN
    # -------------------------------------
    for t in 1:main_ln
        for l in 1:online_ln
            n = rand(1:s)

            # store init input para
            z[1,1] = 0.0
            for j in 1:in_n-1
                z[1,j+1] = train_data[n,j]
            end

            # Forward propagation
            for i in 2:num_layers

                for k in 1:out_n_at_m[i]
                    u[i,k] = 1.0
                    for j in 1:out_n_at_m[i-1]
                        u[i,k] += w_new[i-1,j,k] * z[i-1,j]
                    end
                end
                """
                u[0] = 1.0 + train_data[i][0] * w_new[0,1] + train_data[i][1] * w_new[0,2]
                u[1] = 1.0 + train_data[i][0] * w_new[1,1] + train_data[i][1] * w_new[1,2]
                """

                # activation function
                if i != num_layers
                    # middle layers
                    #rectifier_func(z, u, i)
                    z = sigm_func(z, u, i)
                else
                    # output layer
                    z = softmax_func(z, u, i)
                end

                # set z[i,0]
                z[i,1] = 0.0
            end

            # Back propagation
            # delta at output layer
            i = num_layers - 1
            for k in 2:out_n_at_m[num_layers]
                delta[i,k] = z[i,k] - train_re[n,k-1]
            end

            # delta
            for ii in 1:num_layers-1
                i = num_layers - ii+1
                for k in 1:out_n_at_m[i]
                    #df = rectifier_func_derivation(u[i,k])
                    df = sigm_func_derivation(u[i,k])
                    for j in 1:out_n_at_m[i-1]
                        delta[i-1,k] += delta[i,j] *(w_new[i-1,j,k] * df )
                    end
                end
            end
            """
            delta[0] = 0.0
            delta[1] = z[1] - train_re[i,0]
            delta[2] = z[2] - train_re[i,1]
            """

            for ii in 1:num_layers-1
                i = num_layers - ii+1

                # slope
                for k in 1:out_n_at_m[i]
                    for j in 1:out_n_at_m[i-1]
                        dEdw[i-1,j,k] += delta[i,j] * z[i,k]
                    end
                end
                """
                dEdw[0,0] += delta[0] * 0.0
                dEdw[0,1] += delta[0] * train_data[i,0]
                dEdw[0,2] += delta[0] * train_data[i,1]
                dEdw[1,0] += delta[1] * 0.0
                dEdw[1,1] += delta[1] * train_data[i,0]
                dEdw[1,2] += delta[1] * train_data[i,1]
                """
            end

            # update w
            for i in 1:num_layers-1
                for j in 2:out_n
                    for k in 1:out_n
                        w_temp[i,j,k] = w_old[i,j,k] + mu*(w_new[i,j,k] - w_old[i,j,k]) - ep*dEdw[i,j,k] / online_ln
                    end
                end
            end
            #w_temp = w_old + mu*(w_new - w_old) - ep*dEdw / online_ln
            w_old  = copy(w_new)
            w_new  = copy(w_temp)

            # reset
            for i in 1:num_layers-1
                for j in 1:out_n
                    for k in 1:out_n
                        dEdw[i,j,k] = 0.0
                    end
                end
            end
            for i in 1:num_layers
                for k in 1:out_n
                    delta[i,k] = 0.0
                end
            end
        end

        # check rate of correct answer
        roca_test[t]  = test_nn( test_data,  test_ans, w_new, in_n, out_n, out_n_at_m, num_layers,  f_re_test)
        roca_train[t] = test_nn(train_data, train_ans, w_new, in_n, out_n, out_n_at_m, num_layers, f_re_train)
    end

    # output
    output_result(f_re_train, main_ln, roca_train)
    output_result(f_re_test, main_ln, roca_test)
end

# ------------------------------------
# read file
# ------------------------------------
function read_file(infile)
    skipnum = 1

    fff = []
    open(infile, "r") do f
        fff = read(f,String)
    end
    fff = split(fff, "\n", keepempty=false)
    num_data = length(fff) - skipnum

    temp   = split(fff[2],",")
    data   = zeros(num_data, length(temp)-2)
    result = zeros(num_data, 2)
    ans    = zeros(num_data)

    for i in 1+skipnum:length(fff)
        ii = i - 1
        fff[i] = replace(fff[i]," \r" => "")

        temp = split(fff[i],",")

        result[ii, (parse(Int,temp[2]) + 1)] = 1
        for j in 1:length(temp)-2
            data[ii,j] = parse(Float64, temp[j+2])
        end

        ans[ii] = parse(Int, temp[2])
    end

    #println(size(data))
    #println(size(result))
    #println(size(ans))

    return result, data, ans
end

# ------------------------------------
# output file
# ------------------------------------
function output_result(outf, main_ln, roca)
    open(outf,"w") do f
        for t in 1:main_ln
            a = @sprintf("%8.8e", roca[t])
            write(f, a)
            write(f, "\n")
        end
    end
end

# ------------------------------------------------
# whitening
# ------------------------------------------------
function whitening(data)
    nl = length(data)
    nd = length(data[1])

    # ave = 0
    for i in 1:nl
        ave = 0.0
        for j in 1:nd
            ave += data[i,j]
        end
        ave /= nd
        for j in 1:nd
            data[i,j] -= ave
        end
    end

    # variance value = 0
    for i in 1:nl
        sigma = 0.0
        for j in 1:nd
            sigma += data[i,j]^2
        end
        sigma /= nd
        sigma = (sigma)^0.5

        if sigma == 0.0
            sigma = 1.0
        end
        for j in 1:nd
            data[i,j] /= sigma
        end
    end

    return data
end

# ------------------------------------------------
# random on gauss distribution
# ------------------------------------------------
function rand_normal(mean, stdev)
    if stdev <= 0.0
        error("standard deviation must be positive")
    end
    u1 = rand()
    u2 = rand()
    r  = sqrt( -2.0*log(u1) )
    theta = 2.0*pi*u2
    rg = mean + stdev*r*sin(theta)
    return rg
end
# ------------------------------------------------
# logistic sigmoid function
# ------------------------------------------------
function sigm_func(z, u, i)
    ld = length(u[i])
    for j in 1:ld
        z[i,j] = 1.0 / ( 1.0 + exp(-u[i,j]) )
    end
    return z
end

# ------------------------------------------------
# Derivation of logistic sigmoid function
# ------------------------------------------------
function sigm_func_derivation(ui)
    y = 1 / ( 1 + exp(-ui) )
    y = (1-y)*y
    return y
end

# ------------------------------------------------
# hyperbolic tangent function
# ------------------------------------------------
function hyperbolic_tangent_func(z, u, i)
    ld = length(u[i])
    for j in 1:ld
        z[i,j] = ( exp(u[i,j]) - exp(-u[i,j]) ) / ( exp(u[i,j]) + exp(-u[i,j]) )
    end
    return z
end

# ------------------------------------------------
# Rectified linear function
# ------------------------------------------------
function rectifier_func(z, u, i)
    ld = length(u[i])
    for j in 1:ld
        z[i,j] = max( u[i,j], 0.0 )
    end
    return z
end

# ------------------------------------------------
# Derivation of Rectified linear function
# ------------------------------------------------
function rectifier_func_derivation(ui)
    fd = 1.0

    if ui <= 0.0
        fd = 0.0
    end
    return fd
end

# ------------------------------------------------
# linear mapping ( function )
# ------------------------------------------------
function linear_func(z, u, i)
    ld = length(u[i])
    for j in 1:ld
        z[i,j] = u[i,j]
    end
    return z
end

# ------------------------------------------------
# softmax function
# ------------------------------------------------
function softmax_func(z, u, i)
    ld = length(u[i])

    s = 0.0
    for j in 1:ld
        s += exp(u[i,j])
    end
    for j in 1:ld
        z[i,j] = exp(u[i,j]) / s
    end
    return z
end

# ------------------------------------------------
# test
# ------------------------------------------------
function test_nn(test_data, corr_data, w, in_n, out_n, out_n_at_m, num_layers, of)
    size = Int(length(test_data) / in_n)
    ans  = zeros(size)
    u    = zeros(num_layers, out_n) # sum result
    z    = zeros(num_layers, out_n) # output

    point = 0
    for n in 1:size
        # store init input para
        z[1,1] = 0.0
        for j in 1:in_n
            z[1,j+1] = test_data[n,j]
        end

        # Forward propagation
        for i in 2:num_layers

            for k in 1:out_n_at_m[i]
                u[i,k] = 1.0
                for j in 1:out_n_at_m[i-1]
                    u[i,k] += w[i-1,j,k] * z[i-1,j]
                end
            end
            """
            u[0] = 1.0 + train_data[i][0] * w_new[0,1] + train_data[i][1] * w_new[0,2]
            u[1] = 1.0 + train_data[i][0] * w_new[1,1] + train_data[i][1] * w_new[1,2]
            """

            # activation function
            if i != num_layers
                # middle layers
                z = rectifier_func(z, u, i)
            else
                # output layer
                z = softmax_func(z, u, i)
            end

            # set z[i,0]
            z[i,1] = 0.0
        end

        # if corr or not
        max_index = argmax(z[num_layers-1])
        ans[n] = max_index - 1

        if ans[n] == corr_data[n]
            point += 1
        end
        """
        print(z[num_layers-1])
        print(ans)
        raise ValueError("error!")
        """
    end

    # rate of correct answer
    roca = point / size

    """
    print("------------------------------")
    print(ans[:50])
    print(corr_data[:50])
    print(point)
    print(size)
    print(roca)
    """

    return roca
end

# ---------------------------------------------
# ---------------------------------------------
main()


""" Combinatorial maps and related structures, as C-sets.

In topological graph theory and graph drawing, an **embedded graph** is a
combinatorial structure representing a graph embedded in an (oriented) surface,
up to equivalence under (orientation-preserving) homeomorphism. This module
defines data structures for rotation systems, combinatorial maps, and other
combinatorial objects describing embedded graphs.
"""
module CombinatorialMaps
export σ, α, ϕ, trace_vertices, trace_edges, trace_faces,
  AbstractRotationGraph, RotationGraph, add_corolla!, pair_half_edges!,
  AbstractRotationSystem, RotationSystem,
  AbstractCombinatorialMap, CombinatorialMap

using Catlab, Catlab.CategoricalAlgebra.CSets, Catlab.Graphs
using Catlab.Graphs.BasicGraphs: TheoryHalfEdgeGraph
using Catlab.Permutations: cycles

# General properties
####################

""" Vertex permutation of combinatorial map or similar structure.
"""
σ(x::ACSet, args...) = subpart(x, args..., :σ)

""" Edge permutation of combinatorial map or similar structure.
"""
α(x::ACSet, args...) = subpart(x, args..., :α)

""" Face permutation of combinatorial map or similar structure.
"""
ϕ(x::ACSet, args...) = subpart(x, args..., :ϕ)

""" Trace vertices of combinatorial map or similar, returning a list of cycles.
"""
trace_vertices(x::ACSet) = cycles(σ(x))

""" Trace edges of combinatorial map or similar, return a listing of cycles.

Usually the cycles will be pairs of half edges but in a hypermap the cycles can
be arbitrary.
"""
trace_edges(x::ACSet) = cycles(α(x))

""" Trace faces of combinatorial map or similar, returning list of cycles.
"""
trace_faces(x::ACSet) = cycles(ϕ(x))

# Rotation graphs
#################

@present TheoryRotationGraph <: TheoryHalfEdgeGraph begin
  σ::Hom(H,H)

  compose(σ, vertex) == vertex
end

@abstract_acset_type AbstractRotationGraph <: AbstractHalfEdgeGraph
@acset_type RotationGraph(TheoryRotationGraph,
                          index=[:vertex]) <: AbstractRotationGraph

α(g::AbstractRotationGraph) = inv(g)
ϕ(g::AbstractRotationGraph) = sortperm(inv(g)[σ(g)]) # == (σ ⋅ inv)⁻¹

""" Add corolla to rotation graph, rotation system, or similar structure.

A *corolla* is a vertex together with its incident half-edges, the number of
which is its *valence*. The rotation on the half-edges is the consecutive one
induced by the half-edge part numbers.
"""
function add_corolla!(g::AbstractRotationGraph, valence::Int; kw...)
  v = add_vertex!(g; kw...)
  n = nparts(g, :H)
  add_parts!(g, :H, valence; vertex=v, σ=circshift((n+1):(n+valence), -1))
end

""" Pair together half-edges into edges.
"""
pair_half_edges!(g::AbstractRotationGraph, h, h′) =
  set_subpart!(g, [h; h′], :inv, [h′; h])

# Rotation systems
##################

@present TheoryRotationSystem(FreeSchema) begin
  H::Ob
  σ::Hom(H,H)
  α::Hom(H,H)

  compose(α, α) == id(H)
end

@abstract_acset_type AbstractRotationSystem
@acset_type RotationSystem(TheoryRotationSystem) <: AbstractRotationSystem

# ϕ == (σ⋅α)⁻¹ == α⁻¹ ⋅ σ⁻¹
ϕ(sys::AbstractRotationSystem) = sortperm(α(sys)[σ(sys)])

function add_corolla!(sys::AbstractRotationSystem, valence::Int)
  n = nparts(sys, :H)
  add_parts!(sys, :H, valence; σ=circshift((n+1):(n+valence), -1))
end

pair_half_edges!(sys::AbstractRotationSystem, h, h′) =
  set_subpart!(sys, [h; h′], :α, [h′; h])

# Combinatorial maps
####################

@present TheoryHypermap(FreeSchema) begin
  H::Ob
  σ::Hom(H,H)
  α::Hom(H,H)
  ϕ::Hom(H,H)

  compose(σ, α, ϕ) == id(H)
end

@present TheoryCombinatorialMap <: TheoryHypermap begin
  compose(α, α) == id(H)
end

@abstract_acset_type AbstractCombinatorialMap
@acset_type CombinatorialMap(TheoryCombinatorialMap) <: AbstractCombinatorialMap

# TODO: What kind of interface should we have for maps and hypermaps?

end


"""
    PushBot
"""
mutable struct PushBot{T} <: Model{T}
    nq::Int
	nu::Int
	nw::Int
	nc::Int

    mb::T # mass
	ma::T
	l::T # length

    μ_world::T  # coefficient of friction
    μ_joint::T  # gravity
	g::T

	base::BaseMethods
	dyn::DynamicsMethods

	joint_friction::SVector
end

# Kinematics
function _kinematics(model::PushBot, q; mode = :com)
	θ, d  = q
	if mode == :d
		return [-model.l * sin(θ) + d * cos(θ);
				 model.l * cos(θ) + d * sin(θ)]
	elseif mode == :com
		return [-1.0 * model.l * sin(θ);
				 1.0 * model.l * cos(θ)]
	elseif mode == :ee
		return [-model.l * sin(θ);
				 model.l * cos(θ)]
	else
		@error "incorrect mode"
	 	return zeros(2)
	end
end

function _jacobian(model::PushBot, q; mode = :com)
	θ, d = q

	if mode == :d
 		return [(-model.l * cos(θ) - d * sin(θ)) cos(θ);
		        (-model.l * sin(θ) + d * cos(θ)) sin(θ)]
 	elseif mode == :com
 		return [-1.0 * model.l * cos(θ) 0.0;
		        -1.0 * model.l * sin(θ) 0.0]
	elseif mode == :ee
		return [-model.l * cos(θ) 0.0;
		        -model.l * sin(θ) 0.0]
 	else
 		@error "incorrect mode"
 	 	return
 	end
end

function kinematics(model::PushBot, q)
	[_kinematics(model, q, mode = :d);
	 _kinematics(model, q, mode = :d)]
end

function lagrangian(model::PushBot, q, q̇)
	L = 0.0

	vθ = _jacobian(model, q, mode = :com) * q̇
	L += 0.5 * model.mb * transpose(vθ) * vθ
	L -= model.mb * model.g * _kinematics(model, q, mode = :com)[2]

	vd1 = _jacobian(model, q, mode = :d) * q̇
	L += 0.5 * model.ma * transpose(vd1) * vd1
	L -= model.ma * model.g * _kinematics(model, q, mode = :d)[2]

	return L
end

function M_func(model::PushBot, q)
	Jθ = _jacobian(model, q, mode = :com)
	Jd = _jacobian(model, q, mode = :d)

	return model.mb * transpose(Jθ) * Jθ + model.ma * transpose(Jd) * Jd
end

function ϕ_func(model::PushBot, env::Environment, q)
	# walls at x = -0.5, x = 0.5
    SVector{2}([_kinematics(model, q, mode = :d)[1] + 0.5;
	            0.5 - _kinematics(model, q, mode = :d)[1]])
end

function rot2D(x)
	[cos(x) sin(x);
    -sin(x) cos(x)]
end

function J_func(model::PushBot, env::Environment, q)
	r1 = [0.0 -1.0; 1.0 0.0]
	r2 = [0.0 1.0; -1.0 0.0]

    SMatrix{4, 2}([r1 * _jacobian(model, q, mode = :d);
		           r2 * _jacobian(model, q, mode = :d)])
end

function B_func(model::PushBot, q)
	@SMatrix [1.0 * model.l 1.0;
	          1.0 1.0 / model.l]
end

function A_func(::PushBot, q)
	@SMatrix [1.0 0.0;
			  0.0 1.0]
end

# Parameters
g = 9.81 # gravity
μ_world = 0.5  # coefficient of friction
μ_joint = 10.0

mb = 1.0 # body mass
ma = 0.01 # mass of ee # CHANGED
l = 1.0  # leg mass

# Dimensions
nq = 2
nu = 2
nw = nq
nc = 2
nquat = 0

pushbot = PushBot(nq,nu,nw,nc,
			   mb, ma, l,
			   μ_world, μ_joint, g,
			   BaseMethods(), DynamicsMethods(),
			   SVector{2}(μ_joint * [1.0; 1.0]))

function friction_coefficients(model::PushBot)
	return [model.μ_world]
end

function initialize_z!(z, model::PushBot, idx::RoboDojo.IndicesZ, q)
    z .= 1.0
    z[idx.q] .= q
end


function run_test(M,net_param,time_param,dt)

N = length(M[1,:]);
extinput = zeros(N);
u = zeros(N);

x, r, b, stim = net_param[1:N], net_param[N+1:2*N], net_param[2*N+1], net_param[2*N+2:3*N+1];
stim_on, stim_off, test_time = time_param[1], time_param[2], time_param[3];

#===== Run simulation =====#
dtheta(x_var,I_var) = 1. - cos(x_var) + I_var.*(1. + cos(x_var));
dr(r_var) = -b*r_var;

t = 0.;

utest = zeros(Int(test_time/dt),N);

for i=1:Int(test_time/dt)

      if i > Int(stim_on/dt) && i <= Int(stim_off/dt)
            extinput[:] = stim[:];
      else
            extinput[:] = zeros(N)[:];
      end

      # Define synaptic activity
      u = M*r;

      # Update neuron phase
      k1 = dt*dtheta(x,      u+extinput);
      k2 = dt*dtheta(x+k1/2, u+extinput);
      k3 = dt*dtheta(x+k2/2, u+extinput);
      k4 = dt*dtheta(x+k3,   u+extinput);
      xnext  = x + (k1+2*k2+2*k3+k4)/6;

      # Update filtered spikes
      l1 = dt*dr(r);
      l2 = dt*dr(r+l1/2);
      l3 = dt*dr(r+l2/2);
      l4 = dt*dr(r+l3);

      # Spike detection.
      idx1 = xnext - x .> 0.;
      idx2 = xnext - x .> mod(pi - mod(x,2*pi),2*pi); # distance to next pi.
      idx  = idx1.*idx2;
      ind  = collect(1:N)[idx[:]]';

      # Update each neuron's spikes
      if !isempty(ind)
          r[idx[:]] = r[idx[:]] + b;
      end

      # Update x, r
      x = x + (k1+2*k2+2*k3+k4)/6;
      r = r + (l1+2*l2+2*l3+l4)/6;
      u = M*r;
      t = t + dt;

      # save test data
      utest[i,:] = u[:];
end

return utest

end


"""
    findthomsen(p::Vector{VMeasure},
                     sv::Vector{VMeasure},
                     sh::Vector{VMeasure},
                     vp₀::Tuple{Real,Real},
                     vs₀::Tuple{Real,Real},
                     ϵ₀::Tuple{Real,Real},
                     δ₀::Tuple{Real,Real},
                     γ₀::Tuple{Real,Real})

    Estimate Thomsen parameters, wavespeeds and phase angles from measurements of group P, SV ans SH wave velocities, using quasinewton method. A priori values ()₀ are tuples with value and standard deviation (a priori gaussian).

Return (vp0, vs0, ϵ, δ, γ), θp, θsv, θsh, CMpost

"""
function findthomsen(p::Vector{VMeasure},
                     sv::Vector{VMeasure},
                     sh::Vector{VMeasure},
                     vp₀::Tuple{Real,Real},
                     vs₀::Tuple{Real,Real},
                     ϵ₀::Tuple{Real,Real},
                     δ₀::Tuple{Real,Real},
                     γ₀::Tuple{Real,Real})

    dobs, CDi, np, nsv, nsh = inputdobs(p,sv,sh)

    #make vector of model parameters (first guess)
    m₀ = [vp₀[1],vs₀[1],ϵ₀[1],δ₀[1],γ₀[1]]
    σm = [vp₀[2],vs₀[2],ϵ₀[2],δ₀[2],γ₀[2]]
    for x in p
        push!(m₀, phaseangleP(x.angle, vp₀[1]/vs₀[1], ϵ₀[1], δ₀[1]))
        push!(σm, 100.)
    end
    for x in sv
        push!(m₀, phaseangleSV(x.angle, vp₀[1]/vs₀[1], ϵ₀[1], δ₀[1]))
        push!(σm, 100.)
    end
    for x in sh
        push!(m₀, phaseangleSH(x.angle, γ₀[1]))
        push!(σm, 100.)
    end

    CM = Diagonal(σm.^2)
    CMi = inv(CM)

    m, CMpost = quasinewton(dobs, _gt, m₀, CMi, CDi, (np,nsv,nsh))

    return m[1:5], m[6:5+np], m[6+np:5+np+nsv], m[6+np+nsv:5+np+nsv+nsh], CMpost

end

"""
    findthomsen(p::Vector{VMeasure},
                     sv::Vector{VMeasure},
                     sh::Vector{VMeasure},
                     vp₀::Real,
                     vs₀::Real,
                     ϵ₀::Real,
                     δ₀::Real,
                     γ₀::Real)

    Method using just numbers as input for a priori values, assuming std=10.0 for all parameters.

"""
function findthomsen(p::Vector{VMeasure},
                     sv::Vector{VMeasure},
                     sh::Vector{VMeasure},
                     vp₀::Real,
                     vs₀::Real,
                     ϵ₀::Real,
                     δ₀::Real,
                     γ₀::Real)

    return findthomsen(p,sv,sh,(vp₀,10.),(vs₀,10.),(ϵ₀,10.),(δ₀,10.),(γ₀,10.))

end

"""
    findmoduli(p::Vector{VMeasure},
                    sv::Vector{VMeasure},
                    sh::Vector{VMeasure},
                    ρ::AbastractFloat,
                    c11₀::Tuple{Real,Real},
                    c13₀::Tuple{Real,Real},
                    c33₀::Tuple{Real,Real},
                    c55₀::Tuple{Real,Real},
                    c66₀::Tuple{Real,Real}

    Estimate stiffness matrix and phase angles from measurements of group P, SV ans SH wave velocities, using quasinewton method. A priori values ()₀ are tuples with value and standard deviation (a priori gaussian).

Return Cijkl, θp, θsv, θsh, CMpost

"""
function findmoduli(p::Vector{VMeasure},
                    sv::Vector{VMeasure},
                    sh::Vector{VMeasure},
                    ρ::Real,
                    c11₀::Tuple{Real,Real},
                    c13₀::Tuple{Real,Real},
                    c33₀::Tuple{Real,Real},
                    c55₀::Tuple{Real,Real},
                    c66₀::Tuple{Real,Real})

    dobs, CDi, np, nsv, nsh = inputdobs(p,sv,sh)

    #make vector of model parameters (first guess)
    m₀ = [c11₀[1],c13₀[1],c33₀[1],c55₀[1],c66₀[1]]
    σm = [c11₀[2],c13₀[2],c33₀[2],c55₀[2],c66₀[2]]

    C₀ = stiffmatrixTI(m₀...)
    ϵ₀, δ₀, γ₀ = moduli2thomsen(C₀)
    vp,vs = velocities0(C₀,ρ) # here density does not matter

    for x in p
        push!(m₀, phaseangleP(x.angle, vp/vs, ϵ₀, δ₀))
        push!(σm, 100.)
    end
    for x in sv
        push!(m₀, phaseangleSV(x.angle, vp/vs, ϵ₀, δ₀))
        push!(σm, 100.)
    end
    for x in sh
        push!(m₀, phaseangleSH(x.angle, γ₀))
        push!(σm, 100.)
    end

    CM = Diagonal(σm.^2)
    CMi = inv(CM)

    m, CMpost = quasinewton(dobs, _gc, m₀, CMi, CDi, (ρ,np,nsv,nsh))

    return stiffmatrixTI(m[1:5]...), m[6:5+np], m[6+np:5+np+nsv], m[6+np+nsv:5+np+nsv+nsh], CMpost
end

"""
	findmoduli(p::Vector{VMeasure},
                   sv::Vector{VMeasure},
                   sh::Vector{VMeasure},
                   ρ::Real,
                   vp₀::Real,
                   vs₀::Real,
                   errlogC::Real)

Method with isotropic vp,vs as input.
"""
function findmoduli(p::Vector{VMeasure},
                    sv::Vector{VMeasure},
                    sh::Vector{VMeasure},
                    ρ::Real,
                    vp₀::Real,
                    vs₀::Real,
                    errlogC::Real)

    C0 = stiffmatrixTI(vpvsYoungsPoisson(vp₀,vs₀,ρ))

    return findmoduli(p,sv,sh,ρ,
                      (C0[1,1],errlogC),
                      (C0[1,3],errlogC),
                      (C0[3,3],errlogC),
                      (C0[5,5],errlogC),
                      (C0[6,6],errlogC))

end
"""
    findmoduli(p::Vector{VMeasure},
                    sv::Vector{VMeasure},
                    sh::Vector{VMeasure},
                    ρ::Real,
                    c11₀::Real,
                    c13₀::Real,
                    c33₀::Real,
                    c55₀::Real,
                    c66₀::Real,
                    errlogC::Real)

Method using just a unique error on moduli.
"""
function findmoduli(p::Vector{VMeasure},
                    sv::Vector{VMeasure},
                    sh::Vector{VMeasure},
                    ρ::Real,
                    c11₀::Real,
                    c13₀::Real,
                    c33₀::Real,
                    c55₀::Real,
                    c66₀::Real,
                    errlogC::Real)

    return findmoduli(p,sv,sh,ρ,
                      (c11₀,errlogC),
                      (c13₀,errlogC),
                      (c33₀,errlogC),
                      (c55₀,errlogC),
                      (c66₀,errlogC))

end

"""
	inputdobs(p,sv,sh)
"""
function inputdobs(p,sv,sh)
    np = length(p)
    nsv = length(sv)
    nsh = length(sh)

    # make vector of data and data covariance
    dobs = Float64[]
    σd = Float64[]
    for x in vcat(p,sv,sh)
        append!(dobs, [x.angle, log(x.value)])
        append!(σd, [x.err_angle, x.err_value])
    end

    CD = Diagonal(σd.^2)
    CDi = inv(CD)

    return dobs, CDi, np, nsv, nsh
end

function _gc(m,ρ,np,nsv,nsh)
    C₀ = stiffmatrixTI(m[1:5]...)
    ϵ, δ, γ = moduli2thomsen(C₀)
    vp0,vs0 = velocities0(C₀,ρ)
    mt = vcat([vp0,vs0,ϵ,δ,γ], m[6:end])
    return _gt(mt, np, nsv, nsh)
end

function _gt(m,np,nsv,nsh)
    vp0 = m[1]
    vs0 = m[2]
    ϵ = m[3]
    δ = m[4]
    γ = m[5]
    θp = m[6:5+np]
    θsv= m[6+np:5+np+nsv]
    θsh= m[6+np+nsv:5+np+nsv+nsh]

    d = zeros(eltype(m), 2np+2nsv+2nsh)  #NEED eltype HERE to work with ForwardDiff!
    c = 1
    for x in θp
        vp = vphaseP(x,vp0,vs0,ϵ,δ)
        dvvp = dvPdthetav(x,vp0/vs0,ϵ,δ)
        d[c] = groupangle(x,dvvp)
        d[c+1] = log(vgroup(vp, dvvp))
        c+=2
    end
    for x in θsv
        vsv = vphaseSV(x,vp0,vs0,ϵ,δ)
        dvvsv = dvSVdthetav(x,vp0/vs0,ϵ,δ)
        d[c] = groupangle(x,dvvsv)
        d[c+1] = log(vgroup(vsv, dvvsv))
        c+=2
    end
    for x in θsh
        d[c] = groupangleSH(x,γ)
        d[c+1] = log(vgroupSH(d[c], vs0, γ))
        c+=2
    end

    return d
end


# matrices of Gensys
G0 = zeros(24,24)
G1 = zeros(24,24)
Psi = zeros(24,8)
Pi = zeros(24,6)

# Equation 1 (Euler)

G0[1,1] = -1
G0[1,2] = -(1-h)/sig
G0[1,3] = (1-h)/sig
G0[1,7] = 1+h
G0[1,8] = (1-h)/sig

G1[1,3] = (1+h)/sig
G1[1,7] = h

Pi[1,1] = -1
Pi[1,2] = -(1-h)/sig
Pi[1,3] = (1-h)/sig

# Equation 2 (goods market equilibrium)

G0[2,9]  = 1
G0[2,7]  = -(1-alfa)
G0[2,10] = -alfa*eta*(2-alfa)
G0[2,11] = -alfa*eta
G0[2,17] = -alfa

# Equation 3 (tot)

G0[3,10] = 1
G0[3,13] = -1
G0[3,12] = 1

G1[3,10] = 1

# Equation 4 (relationship q and tot)

G0[4,10] = -(1-alfa)
G0[4,14] = 1
G0[4,11] = -1

# Equation 5 (relationship q, s, pi)

G0[5,14] = 1
G0[5,15] = -1
G0[5,18] = -1

G1[5,2]  = -1
G1[5,14] = 1
G1[5,15] = -1

# Equation 6 (phillips curve - domestic)

G0[6,5]  = -bet
G0[6,12] = 1+bet*delta_D
G0[6,24] = -bet*(1-theta_D)*(1-theta_D*bet)/theta_D

G1[6,12] = delta_D

Pi[6,4] = -bet

# Equation 7 (marginal cost)

G0[7,24] = 1
G0[7,10] = -alfa
G0[7,9]  = -phi
G0[7,7]  = -sig/(1-h)
G0[7,20] = 1+phi

G1[7,7] = -(sig*h)/(1-h)

# Equation 8  (phillips curve - importing)

G0[8,6]  = -bet
G0[8,13] = 1+bet*delta_I
G0[8,11] = -(1-theta_I)*(1-theta_I*bet)/theta_I
G0[8,22] = -1

G1[8,13] = delta_I

Pi[8,5]  = -bet

# Equation 9 (relationship pi and tot)

G0[9,12] = -1
G0[9,10] = -alfa

G1[9,10] = -alfa
G1[9,2]  = -1

# Equation 10 (bugdet constraint)

G0[10,21] = 1
G0[10,10] = alfa
G0[10,11] = alfa
G0[10,9]  = -1
G0[10,7]  = 1

G1[10,21] = 1/bet

# Equation 11 (UIP)

G0[11,4]  = -1
G0[11,15] = 1
G0[11,8]  = 1
G0[11,19] = -1
G0[11,21] = chi
G0[11,23] = -1

Pi[11,6] = -1

# Equation 12 (Taylor RMI)

G0[12,8]  = 1
G0[12,16] = -rho_1
G0[12,9]  = -(1-rho_1-rho_2)*lambda_y
G0[12,15] = -(1-rho_1-rho_2)*lambda_s

G1[12,16] = rho_2
G1[12,2]  = (1-rho_1-rho_2)*lambda_pi
G1[12,15] = -(1-rho_1-rho_2)*lambda_s

Psi[12,5] = sigma_i

###################
## Foreign Block ##
###################

# System (Equations 13, 14 and 15)

G0[13,17] = 1
G0[14,17] = a_0piy
G0[14,18] = 1
G0[15,17] = a_0iy
G0[15,18] = a_0ipi
G0[15,19] = 1

G1[13,17] = a_1yy
G1[13,18] = a_1ypi
G1[13,19] = a_1yi
G1[14,17] = a_1piy
G1[14,18] = a_1pipi
G1[14,19] = a_1pii
G1[15,17] = a_1iy
G1[15,18] = a_1ipi
G1[15,19] = a_1ii

Psi[13,6] = sigma_y_star
Psi[14,7] = sigma_pi_star
Psi[15,8] = sigma_i_star

#################
## Shocks block #
#################

# Equation 16 (a)

G0[16,20] = 1

G1[16,20] = rho_a

Psi[16,1] = sigma_a

# Equation 17 (gamma)

G0[17,3] = 1

G1[17,3] = rho_gamma

Psi[17,2] = sigma_gamma

# Equation 18 (epsilon_cp)

G0[18,22] = 1

G1[18,22] = rho_cp

Psi[18,3] = sigma_cp

# Equation 19 (phi)

G0[19,23] = 1

G1[19,23] = rho_phi

Psi[19,4] = sigma_phi

####################
## Identity block ##
####################

# Equation 20 (c)

G0[20,7] = 1
G1[20,1] = 1

# Equation 21 (s)

G0[21,15] = 1
G1[21,4]  = 1

# Equation 22 (pi_D)

G0[22,12] = 1
G1[22,5]  = 1

# Equation 23 (pi_I)

G0[23,13] = 1
G1[23,6]  = 1

# Equation 24 (i)

G0[24,16] = 1
G1[24,8]  = 1

# G0[findall(G0 .== nothing)] .= 0
# G1[findall(G1 .== nothing)] .= 0
# Psi[findall(Psi .== nothing)] .= 0
# Pi[findall(Pi .== nothing)] .= 0


import Pkg; Pkg.activate("../../")

include("GPs.jl")

using Distributions
using PyPlot
using BenchmarkTools
import Random

Random.seed!(4)
n = 10
x = randn(n)
f(x) =  sin(3 * x) * sin(x) * (-1)^(x > 0)
sigma = 0
y = f.(x) + randn(n) * sigma;

xgrid = range(minimum(x) - 0.3, maximum(x) + 0.3, length=100)
rhos = rand(Uniform(0.4, 0.6), 50)

p = [let
    gp = GPs.GP(y, x[:,:], GPs.SqExpKernel(1, rho), sigma=0);
    post = GPs.posterior(gp, xgrid[:,:])
    rand(post)
end for rho in rhos]
p = hcat(p...)

plt.scatter(x, y)
p_lower, p_upper = GPs.ci(p)
plt.plot(xgrid, mean(p, dims=2))
plt.fill_between(xgrid, p_lower, p_upper, alpha=0.2, color="blue")

for i in sample(1:size(p, 2), 10)
    plt.plot(xgrid, p[:, i], alpha=0.2, color="black")
end

plt.savefig("out/post.pdf")
plt.close()


import LinearAlgebra: I
import ForneyLab: unsafeMean, unsafePrecision, unsafeCov, eye
import StatsFuns: @irrational
@irrational log2π  1.8378770664093454836 log(big(2.)*π)

function eye(dim)
    return Matrix{Float64}(I,dim,dim)
end

function mn(d::ProbabilityDistribution{Univariate, T}) where T<:Union{Gaussian,Gamma}
    return ForneyLab.unsafeMean(d)
end

function mn(d::ProbabilityDistribution{Multivariate, T}) where T<:Union{Gaussian,Gamma}
    return ForneyLab.unsafeMean(d)
end

function pc(d::ProbabilityDistribution{Univariate, T}) where T<:Union{Gaussian,Gamma}
    return ForneyLab.unsafePrecision(d)
end

function pc(d::ProbabilityDistribution{Multivariate, T}) where T<:Union{Gaussian,Gamma}
    return ForneyLab.unsafePrecision(d)
end

function cv(d::ProbabilityDistribution{Univariate, T}) where T<:Union{Gaussian,Gamma}
    return ForneyLab.unsafeCov(d)
end

function cv(d::ProbabilityDistribution{Multivariate, T}) where T<:Union{Gaussian,Gamma}
    return ForneyLab.unsafeCov(d)
end


function diversificationCore(wgts::Array{Float64, 2})
    ## calculates several diversification measures
    ## Note: low-level function not to be called directly!
    ##
    ## Inputs:
    ## 	wgts 	nObs x nAss matrix of portfolio weights
    ##
    ## Outputs:
    ## 	DataFrame containing information on number of assets
    ## 	divIndicators table containing information on number of
    ## 	assets, number of significant assets, maximum weight,
    ## 	standard deviation of weights, number of negative
    ## 	weights and gini coefficient (herfindahl) for each day

    ## tolerance level for significant assets
    sigLev = 0.0001;

    ## transpose if column vector
    if size(wgts, 2) == 1
        wgts = wgts'
    end

    ## basic statistics
    nAss = sum(wgts .!= 0, 2);
    nSignAss = sum(abs(wgts .- 0) .> sigLev, 2)
    maxWgt = maximum(wgts, 2)
    stdWgt = std(wgts, 2)
    nShort = sum(wgts .< 0, 2)

    ## gini coefficient with negative weights
    absExposure = sum(abs(wgts), 2);
    relExposures = wgts ./ repmat(absExposure, 1, size(wgts, 2))
    gini = sum(relExposures.^2, 2);

    divIndicators = DataFrame()
    divIndicators[:nAss] = nAss[:]
    divIndicators[:nSignAss] = nSignAss[:]
    divIndicators[:maxWgt] = maxWgt[:]
    divIndicators[:stdWgt] = stdWgt[:]
    divIndicators[:nShort] = nShort[:]
    divIndicators[:Gini] = gini[:]

    return divIndicators
end

function diversification(invs::Investments)
    ## calculates several diversification measures
    ##
    ## Inputs:
    ## 	wgts 	nObs x nAss matrix of portfolio weights
    ##
    ## Outputs:
    ## 	DataFrame containing information on number of assets
    ## 	divIndicators table containing information on number of
    ## 	assets, number of significant assets, maximum weight,
    ## 	standard deviation of weights, number of negative
    ## 	weights and gini coefficient (herfindahl) for each day

    wgts = AssetMgmt.core(invs)

    divIndicators = diversificationCore(wgts)
    return Timematr(divIndicators, AssetMgmt.idx(invs))
end


@testset "1-Particle Basis"  begin

##

using ACE
using Printf, Test, LinearAlgebra, StaticArrays
using ACE: evaluate, evaluate_d, Rn1pBasis, Ylm1pBasis,
      EuclideanVectorState, Product1pBasis
using Random: shuffle
using ACEbase.Testing: dirfdtest, fdtest, print_tf

##

@info "Build a 1p basis from scratch"

maxdeg = 5
r0 = 1.0
rcut = 3.0

trans = PolyTransform(1, r0)   # r -> x = 1/r^2
J = transformed_jacobi(maxdeg, trans, rcut; pcut = 2)   #  J_n(x) * (x - xcut)^pcut
Rn = Rn1pBasis(J)
Ylm = Ylm1pBasis(maxdeg)
B1p = Product1pBasis( (Rn, Ylm) )
ACE.init1pspec!(B1p)

nX = 10
Xs = rand(EuclideanVectorState, Rn, nX)
cfg = ACEConfig(Xs)

A = evaluate(B1p, cfg)

@info("test against manual summation")
A1 = sum( evaluate(B1p, X) for X in Xs )
println(@test A1 ≈ A)

@info("test permutation invariance")
println(@test A ≈ evaluate(B1p, ACEConfig(shuffle(Xs))))

##

@info("Ylm1pBasis gradients")
Y = ACE.alloc_B(Ylm)
dY = ACE.alloc_dB(Ylm)
tmpd = ACE.alloc_temp_d(Ylm)
ACE.evaluate_ed!(Y, dY, tmpd, Ylm, Xs[1])

for ntest = 1:30
   x0 = randn(3)
   c = rand(length(Y))
   F = x -> sum(ACE.evaluate(Ylm, EuclideanVectorState(SVector{3}(x))) .* c)
   dF = x -> Vector(sum(ACE.evaluate_d(Ylm, EuclideanVectorState(SVector{3}(x))) .* c))
   print_tf(@test fdtest(F, dF, x0; verbose=false))
end
println()

##

@info("Rn1pBasis gradients")

for ntest = 1:30
   x0 = randn(3)
   c = rand(length(Rn))
   F = x -> sum(ACE.evaluate(Rn, EuclideanVectorState(SVector{3}(x))) .* c)
   dF = x -> Vector(sum(ACE.evaluate_d(Rn, EuclideanVectorState(SVector{3}(x))) .* c))
   print_tf(@test fdtest(F, dF, x0; verbose=false))
end
println()

##

@info("Product basis evaluate_ed! tests")

tmp_d = ACE.alloc_temp_d(B1p)
A1 = ACE.alloc_B(B1p)
A2 = ACE.alloc_B(B1p)
ACE.evaluate!(A1, tmp_d, B1p, cfg)
dA = ACE.alloc_dB(B1p, length(cfg))
ACE.evaluate_ed!(A2, dA, tmp_d, B1p, cfg)
println(@test A1 ≈ A2)

##
@info("Product basis gradient test")

for ntest = 1:30
   x0 = randn(3)
   c = rand(length(B1p))
   F = x -> sum(ACE.evaluate(B1p, EuclideanVectorState(SVector{3}(x))) .* c)
   dF = x -> Vector(sum(ACE.evaluate_ed(B1p, EuclideanVectorState(SVector{3}(x)))[2] .* c))
   F(x0)
   dF(x0)
   print_tf(@test fdtest(F, dF, x0; verbose=false))
end
println()
##

end


# sqrt example

using Distributions
using KernelDensityEstimate, KernelDensityEstimatePlotting
using IncrementalInference
using Gadfly, DataFrames

# requried for overloading with new factors
import IncrementalInference: getSample

include(joinpath(@__DIR__, "SquareRootTypes.jl"))

## Direct example

fg = initfg()

addVariable!(fg, :x, ContinuousScalar)
# addVariable!(fg, :y, x0,N=N)

# TODO perhaps make Mixture instead for multiple computation
pts = rand(Distributions.Normal(4.0,0.05),100)   #;rand(Distributions.Normal(144.0,0.05),N)]
md = kde!(pts)

npx = NumbersPrior(md)
addVariable!(fg, :xy, ContinuousScalar)

addFactor!(fg, [:xy;], npx)

xty = Square(Distributions.Normal(0.0,0.01))
addFactor!(fg, [:x,:xy], xty)

# drawGraph(fg)

# initialize from prior
doautoinit!(fg, :xy)
# initialize any random numbers for the square root initial value
manualinit!(fg, :x, randn(1,100))

# find solution
tree, smt, hist = solveTree!(fg)

## plot the result
plotKDE(map(x->getKDE(fg,x), [:x; :xy]))


import Base.Matrix
import Base.+
import Base.*

export SU2

"""
    Represents the group SU(2), {[α -̅β; β ̅α] | α, β ∈ ℂ, |α|² + |β|² = 1}.

fields: α and β.
"""
struct SU2 <: Group
    α::U1
    β::U1
    function SU2(α::U1, β::U1)
        new(α, β)
    end
    function SU2(a::Array{U1,2})
        @assert(size(a) == (2, 2), "The size of Matrix must be 2×2.")
        new(a[1], a[2])
    end
    function SU2(a::Array{Complex,2})
        @assert(size(a) == (2, 2), "The size of Matrix must be 2×2.")
        new(U1(a[1]), U1(a[2]))
    end
    function SU2(a::Array{Complex{Float64},2})
        @assert(size(a) == (2, 2), "The size of Matrix must be 2×2.")
        new(U1(a[1]), U1(a[2]))
    end
end

"""
    Matrix(a::SU2)

Construct a Matrix with the given number.
"""
function Matrix(a::SU2)
    convert(Array{U1,2}, [a.α -conj(a.β); a.β conj(a.α)])
end

"""
    Base.show(io::IO, a::SU2)

Print a string representation of the given number.
"""
function Base.show(io::IO, a::SU2)
    print(io, Matrix(a))
end

"""
    +(a::SU2, b::SU2)

Add two numbers.
"""
function +(a::SU2, b::SU2)
    convert(Array{Complex,2}, Matrix(a) + Matrix(b))
end

"""
    +(a::SU2, b::Array{U1,2})

Add two numbers.
"""
function +(a::SU2, b::Array{U1,2})
    convert(Array{Complex,2}, Matrix(a) + b)
end

"""
    +(a::Array{U1,2}, b::SU2)

Add two numbers.
"""
function +(a::Array{U1,2}, b::SU2)
    convert(Array{Complex,2}, a + Matrix(b))
end

"""
    +(a::SU2, b::Array{Complex,2})

Add two numbers.
"""
function +(a::SU2, b::Array{Complex,2})
    Matrix(a) + b
end

"""
    +(a::Array{Complex,2}, b::SU2)

Add two numbers.
"""
function +(a::Array{Complex,2}, b::SU2)
    a + Matrix(b)
end

"""
    *(a::SU2, b::SU2)

Multiply two numbers.
"""
function *(a::SU2, b::SU2)
    SU2(Matrix(a) * Matrix(b))
end

"""
    *(a::SU2, b::Array{U1,2})

Multiply two numbers.
"""
function *(a::SU2, b::Array{U1,2})
    a * SU2(b)
end

"""
    *(a::Array{U1,2}, b::SU2)

Multiply two numbers.
"""
function *(a::Array{U1,2}, b::SU2)
    SU2(a) * b
end

"""
    *(a::SU2, b::Array{Complex,2})

Multiply two numbers.
"""
function *(a::SU2, b::Array{Complex,2})
    Matrix(a) * b
end

"""
    *(a::Array{Complex,2}, b::SU2)

Multiply two numbers.
"""
function *(a::Array{Complex,2}, b::SU2)
    s * Matrix(b)
end

"""
    *(a::SU2, b::U1)

Multiply two numbers.
"""
function *(a::SU2, b::U1)
    SU2(Matrix(a) .* b)
end

"""
    *(a::U1, b::SU2)

Multiply two numbers.
"""
function *(a::U1, b::SU2)
    SU2(a .* Matrix(b))
end


# Butcher Tableau
const B = [[0.161],
        [-0.008480655492356992, 0.3354806554923570],
        [2.8971530571054944, -6.359448489975075, 4.362295432869581],
        [5.32586482843926, -11.74888356406283, 7.495539342889836, -0.09249506636175525],
        [5.86145544294642, -12.92096931784711, 8.159367898576159,  -0.07158497328140100,
        -0.02826905039406838],
        [0.09646076681806523, 0.01, 0.4798896504144996, 1.379008574103742,
        -3.290069515436081, 2.324710524099774]
        ]

# Step size fractions
const nodes = [0.161, 0.327, 0.9, 0.9800255409045097, 1, 1]

# Weights for 5th order method
const b5 = [0.001780011052226, 0.000816434459657, -0.007880878010262, 0.144711007173263,
            -0.582357165452555, 0.458082105929187, 1/66]
# Weights for 4th order method
const b4 = [0.09646076681806523, 0.01, 0.4798896504144996, 1.379008574103742,
                -3.290069515436081, 2.324710524099774, 0]
# Error estimate
const errest = b5 .- b4

#Interpolation coefficients
const interpC = hcat(
    [1.0, -2.763706197274826, 2.9132554618219126, -1.0530884977290216],
    [0, 0.13169999999999998, -0.2234, 0.1017],
    [0, 3.9302962368947516, -5.941033872131505, 2.490627285651253],
    [0, -12.411077166933676, 30.33818863028232, -16.548102889244902],
    [0, 37.50931341651104, -88.1789048947664, 47.37952196281928],
    [0, -27.896526289197286, 65.09189467479366, -34.87065786149661],
    [0, 1.5, -4, 2.5]
)


include("../RLAlgorithms.jl")
using .RLAlgorithms
using Plots
using Flux
using StatsBase: sample, Weights

function episode!(actor, episode, env, maxt)
    reset!(env)
    for _ in 1:maxt
        s = copy(env.current)
        a = sample(Weights(actor(s)))
        _, r, done = step!(env, a)
        append!(episode, s, a, r)
        done && break
    end
end

function run!(actor, opt, env; epochs=500, maxt=100, γ=99f-2)
    rewards = zeros(Float32, div(epochs, 10))
    episode = PGEpisode{Float32, Int64, Float32}(length(env.observationspace), 1, maxt, γ=γ)

    for i in 1:epochs
        episode!(actor, episode, env, maxt)
        discount!(episode)

        Flux.train!(params(actor), [batch(episode)], opt) do s, a, r, _
            -sum(log.(gather(actor(s), a)) .* r)
        end
        i % 10 == 0 && (rewards[div(i, 10)] = test(s -> Flux.onecold(actor(s)), env))
    end
    return rewards
end

env = CartPole{Float32}()

actor = Chain(
    Dense(length(env.observationspace), 32, relu),
    Dense(32, length(env.actionspace)),
    softmax
)

r = run!(actor, ADAM(0.01, (0.99, 0.999)), env)
plt = plot(r, labels="Reward")
display(plt)
#savefig(plt, "reinforce")


using Pkg
Pkg.activate(".")
using Psychometrics
using Distributions
using LinearAlgebra
using Dates
using Random
#using StatsPlots

function est()
    I_total = 40
    N = 500

    # ITEM PARAMETERS AND LATENTS
    items = [Item(i, string("item_",i), ["math"], Parameters2PL(Product([LogNormal(0.3, 0.2), Normal(0,1)]), [1e-5,5.0], [-6.0, 6.0])) for i = 1 : I_total];
    examinees = [Examinee(e, string("examinee_",e), Latent1D(Normal(0,1), [-6.0, 6.0])) for e = 1 : N];

    # RESPONSES

    responses = answer(examinees, items);

    # Get response matrix (zeros are for wrong answers and for missing answers)
    response_matrix = get_response_matrix(responses, I_total, N);

    #Set Seeds for Random Generation
    #Random.seed!()

    # function tr_norm_gen(number, mu, v)
    #     temp_u = Distributions.rand(Distributions.Uniform(zero(Float64),one(Float64)), number)
    #     norm_dist = Distributions.Normal(mu, sqrt(v))
    #     norm_cdf = Distributions.cdf(norm_dist, zero(Float64))
    #     tr_norm_value = temp_u .- (temp_u .* norm_cdf) .+ norm_cdf
    #     return Distributions.quantile(norm_dist, tr_norm_value)
    # end

    # using StatsPlots
    # A short example of sampling truncated standard normal distribution
    # histogram(tr_norm_gen(800,0,1))
    # #histogram(truncate_rand(Normal(0.0,1.0), [0.0, Inf], n=800))
    # histogram(rand(Distributions.TruncatedNormal(0.0,1.0, 0.0, Inf), 800))

    # Simulation code by Jiang and Templin
    # Gibbs Samplers for Logistic Item Response Models via the Pólya–GammaDistribution:
    # A Computationally Efficient Data-Augmentation StrategyArticleinPsychometrika
    # October 2018 DOI: 10.1007/s11336-018-9641-xCITATIONS2READS692

    ##################################################################################################
    ###########################      Data and Settings    ############################################
    ##################################################################################################

    #Specified by users:Response matrix X
    X = permutedims(response_matrix);
    ##################################################################################################
    ###########################      PolyGamma MCMC sampler   ########################################
    ##################################################################################################

    #K matrix for the coming PolyaGamma samplings
    K = X .- 0.5;

    #Initial Values, need to make sure that all Variance needs to be positive
    ini_a = ones(I_total);
    ini_b = zeros(I_total);
    ini_theta = zeros(N);
    ini_kernel = zeros(N, I_total);

    ini_w = [Distributions.rand(Psychometrics.PolyaGamma(1.0, ini_a[i] * (ini_theta[p] - ini_b[i]))) for p = 1:N, i = 1:I_total] ;

    function theta_sampler(
        vec_a, # I_total x 1
        vec_b, # I_total x 1
        mat_w_id, # N x 1
        mat_k_id, # N x 1
        theta_prior_mu, # 1 X 1
        theta_prior_var,# 1 X 1
    )
        omega_theta = LinearAlgebra.Diagonal(mat_w_id) # I_total x I_total
        theta_variance = 1 / ((vec_a' * omega_theta * vec_a) + (1 / theta_prior_var)) # 1 x 1
        z_theta = vec_a .* vec_b .* mat_w_id + mat_k_id # I_total x 1 # times w_e
        theta_mu = theta_variance * ((vec_a' * z_theta) + (theta_prior_mu / theta_prior_var))
        rand(Distributions.Normal(theta_mu, sqrt(theta_variance)))
    end
    function a_sampler(
        vec_theta, # N x 1
        vec_b_id, # I_total x 1
        mat_w_id, # N x I_total
        mat_k_id, # N x I_total
        a_prior_mu, # 1 x 1
        a_prior_var, # 1 x 1
    )
        omega_a = LinearAlgebra.Diagonal(mat_w_id) # N x N
        b_omega_a = vec_theta .- vec_b_id
        a_variance =
            1 / (
                (b_omega_a' * omega_a * b_omega_a) +
                (1 / a_prior_var)
            ) # 1 x 1
        #z_a = mat_k_id ./ mat_w_id # don't need it
        a_mu = a_variance * (b_omega_a' * mat_k_id + (a_prior_mu / a_prior_var)) # ok
        #return tr_norm_gen(1, a_mu, a_variance)[1]
        return  rand(Distributions.TruncatedNormal(a_mu, sqrt(a_variance), 0.0, Inf))
    end
    #  ( T   -   b   ) ^2 *  w
    # N x 1   I_total x 1        N x I_total

    function b_sampler(vec_theta, vec_a_id, mat_w_id, mat_k_id, b_prior_mu, b_prior_var, N)
        omega_b = LinearAlgebra.Diagonal(mat_w_id) # N x N
        vec_a_id_vec = fill(vec_a_id, N) #N
        b_variance = 1 / (vec_a_id_vec' * omega_b * vec_a_id_vec + (1 / b_prior_var))  # 1 x 1 ok, vec_a_id_vec can be positive because it is squared anyway
        z_b = mat_k_id .- (vec_a_id .* vec_theta .* mat_w_id) # N x 1 # times w_i
        b_mu = b_variance * (((.-vec_a_id_vec)' * z_b) + (b_prior_mu / b_prior_var)) # ok
        return rand(Distributions.Normal(b_mu, sqrt(b_variance)))
    end
    function i_sampler(vec_theta, vec_a_id, vec_b_id, mat_w_id, mat_k_id, a_prior_mu, a_prior_var, b_prior_mu, b_prior_var)
        sigma2 = mapreduce(
            (e, w) -> [(e - vec_b_id)^2, vec_a_id^2 ] .* w,
            +,
            vec_theta,
            mat_w_id,
        )
        sigma2 = 1 ./ (sigma2 + (1 ./ [a_prior_var, b_prior_var]))
        mu = mapreduce(
            (e, w, r) -> [
                (e - vec_b_id) *
                (r),
                -vec_a_id * (
                    #(get_responses_by_examinee_id(e.id, responses)[1].val - 0.5) -
                    (r) -
                    (vec_a_id * e * w)
                ),
            ],
            +,
            vec_theta,
            mat_w_id,
            mat_k_id
        )
        mu =
            sigma2 .* (
                mu + (
                    [a_prior_mu, b_prior_mu] ./
                    [a_prior_var, b_prior_var]
                )
            )
        return rand(Distributions.Product([
            Distributions.TruncatedNormal(mu[1], sqrt(sigma2[1]), 0.0, Inf),
            Distributions.Normal(mu[2], sqrt(sigma2[2])),
        ]))
    end
    # using RCall
    # R"""
    # rpg_sp <- function(x){
    # BayesLogit::rpg(1,1,x)
    # }
    # """
    function w_sampler(vec_theta, vec_a, vec_b, N, I_total)
        #w_temp = [vec_a[i] * (vec_theta[p] - vec_b[i]) for p = 1:N, i = 1:I_total]
        #pg = rcopy(R"apply($w_temp, c(1,2), rpg_sp)")
        pg =[rand(PolyaGamma(1.0, vec_a[i] * (vec_theta[p] - vec_b[i]))) for p = 1:N, i = 1:I_total]
        return pg
    end

    ##################################################################################################
    ###########################   Estimation Starts from here  #######################################
    ##################################################################################################

    #Specified by users:Construct 2PL IRT Estimator

    Iter = 4000
    theta_prior_mu = 0
    theta_prior_var = 1
    a_prior_mu = 1
    a_prior_var = 5
    b_prior_mu = 0
    b_prior_var = 5

    sav_a = zeros(Float64, Iter, I_total);
    sav_b = zeros(Float64, Iter, I_total);
    sav_theta = zeros(Float64, Iter, N);

    sav_a[1, :] = ini_a
    sav_b[1, :] = ini_b
    sav_theta[1, :] = ini_theta
    #sav.iter.hist<-matrix(0,N,I_total)

    for iter = 2:Iter
        sav_w = w_sampler(sav_theta[iter-1, :], sav_a[iter-1, :], sav_b[iter-1, :], N, I_total)
        if mod(iter,100)==0
            println(iter)
        end
        sav_theta_iter = map(
            p -> begin
                no_missing_i = findall(.!ismissing.(K[p, :]))
                theta_sampler(
                sav_a[iter-1, no_missing_i],
                sav_b[iter-1, no_missing_i],
                sav_w[p, no_missing_i],
                K[p, no_missing_i],
                theta_prior_mu,
                theta_prior_var,
            )
        end
        ,
        1:N
        )
        res_iter = map(
            i -> begin
                no_missing_e = findall(.!ismissing.(K[:, i]))
                i_sampler(
                    sav_theta_iter,
                    sav_a[iter-1, i],
                    sav_b[iter-1, i],
                    sav_w[no_missing_e, i],
                    K[no_missing_e, i],
                    a_prior_mu,
                    a_prior_var,
                    b_prior_mu,
                    b_prior_var
                )
            end,
            1:I_total,
        )
        sav_a_iter = map( r -> r[1], res_iter)
        sav_b_iter = map( r -> r[2], res_iter)

        # sav_a_iter = map(
        #     i -> a_sampler(
        #         sav_theta_iter,
        #         sav_b[iter-1, i],
        #         sav_w[:, i],
        #         K[:, i],
        #         a_prior_mu,
        #         a_prior_var,
        #     ),
        #     1:I_total,
        # )
        # sav_b_iter = map(
        #     i -> b_sampler(
        #         sav_theta_iter,
        #         sav_a_iter[i],
        #         sav_w[:, i],
        #         K[:, i],
        #         b_prior_mu,
        #         b_prior_var,
        #         N,
        #     ),
        #     1:I_total,
        # )
        sav_theta[iter, :] .= copy(sav_theta_iter)
        sav_a[iter, :] .= copy(sav_a_iter)
        sav_b[iter, :] .= copy(sav_b_iter)
    end

    sampled = (mcmc_iter-1000):mcmc_iter
    mean_a = [mean(sav_a[sampled, i]) for i = 1:I_total];
    mean_b = [mean(sav_b[sampled, i]) for i = 1:I_total];
    mean_theta = [mean(sav_theta[sampled, i]) for i = 1:N];

    # hcat(map(i -> i.parameters.a, items), mean_a)
    # hcat(map(i -> i.parameters.b, items), mean_b)

    # RMSEs
    println(sqrt(sum((map(i -> i.parameters.a, items) .- mean_a).^2)/I_total))
    println(sqrt(sum((map(i -> i.parameters.b, items) .- mean_b).^2)/I_total))
    println(sqrt(sum((map(e -> e.latent.val, examinees) .- mean_theta).^2)/N))

    return examinees, examinees_est, items, items_est, responses
end
examinees, examinees_est, items, items_est, responses = est();

# using StatsPlots
# plot(sav_a[:,1])

# RMSE_a_500 = zeros(Iter-500);
# RMSE_b_500 = zeros(Iter-500);
# RMSE_t_500 = zeros(Iter-500);
# for j in 1:(Iter-500)
# mean_a = [mean(sav_a[j:(j+499), i]) for i = 1:I_total];
# mean_b = [mean(sav_b[j:(j+499), i]) for i = 1:I_total];
# mean_theta = [mean(sav_theta[j:(j+499), i]) for i = 1:N];
# # RMSEs
# RMSE_a_500[j]=sqrt(sum((map(i -> i.parameters.a, items) .- mean_a).^2)/I_total)
# RMSE_b_500[j]=sqrt(sum((map(i -> i.parameters.b, items) .- mean_b).^2)/I_total)
# RMSE_t_500[j]=sqrt(sum((map(e -> e.latent.val, examinees) .- mean_theta).^2)/N)
# end

# plot(hcat(RMSE_a_500, RMSE_b_500, RMSE_t_500))

# RMSE_a_100 = zeros(Iter-100);
# RMSE_b_100 = zeros(Iter-100);
# RMSE_t_100 = zeros(Iter-100);
# for j in 1:(Iter-100)
# mean_a = [mean(sav_a[j:(j+99), i]) for i = 1:I_total];
# mean_b = [mean(sav_b[j:(j+99), i]) for i = 1:I_total];
# mean_theta = [mean(sav_theta[j:(j+99), i]) for i = 1:N];
# # RMSEs
# RMSE_a_100[j]=sqrt(sum((map(i -> i.parameters.a, items) .- mean_a).^2)/I_total)
# RMSE_b_100[j]=sqrt(sum((map(i -> i.parameters.b, items) .- mean_b).^2)/I_total)
# RMSE_t_100[j]=sqrt(sum((map(e -> e.latent.val, examinees) .- mean_theta).^2)/N)
# end

# plot(hcat(RMSE_a_100, RMSE_b_100, RMSE_t_100))


"""
Computes the Fourier representation of a Gaussian filter with width ``R_S``.
You would multiply this with the FFT'd field and then IFFT to apply the filter.

# Arguments:
- `R_S::T`: filter radius
- `kv`: wavevector, tuple of (kx, ky, kz)

# Returns:
- `Array{T,3}`: fourier representation of gaussian filter
"""
function kspace_gaussian_filter(R_S::T, kv) where T
    kx, ky, kz = kv
    nx, ny, nz = length(kx), length(ky), length(kz)
    filter_k = zeros(T, (nx, ny, nz))
    Threads.@threads for ix = 1:nx
        for iy = 1:ny, iz = 1:nz
            filter_k[ix, iy, iz] = exp(
                -(kx[ix]^2 + ky[iy]^2 + kz[iz]^2) * R_S^2 / 2)
        end
    end
    return filter_k
end

"""
Apply a gaussian filter to a density field with smoothing radius `R_S`
making use of wavevectors `kv`

# Arguments:
- `f::AbstractArray{T,3}`: density field
- `R_S::T`: filter radius
- `kv`: wavevector, tuple of (kx, ky, kz)

# Returns:
- Array{T,3}: smoothed density field
"""
function smooth_gauss(f::AbstractArray{T,3}, R_S::T, kv) where T

    GF = kspace_gaussian_filter(R_S, kv)  # get filter in Fourier space
    f_Rn = (real(ifft(GF .* fft(f))))
    f_Rn = f_Rn .* (sum(f)/sum(f_Rn))  # normalize
    return f_Rn
end

"""
Apply a gaussian filter to a density field with smoothing radius `R_S`
making use of wavevectors `kv` using log smoothing

# Arguments:
- `f::AbstractArray{T,3}`: density field
- `R_S::T`: filter radius
- `kv`: wavevector, tuple of (kx, ky, kz)

# Returns:
- Array{T,3}: log-smoothed density field
"""
function smooth_loggauss(f::AbstractArray{T,3}, R_S::T, kv) where {AA, T}
    # get filter in Fourier space
    GF = kspace_gaussian_filter(R_S, kv)

    # # simple version of what follows
    # f_Rn = 10 .^(real(ifft(GF .* fft(log10.(f)))))

    buffer = similar(f, Complex{T})   # make a complex version of the array
    @strided buffer .= log10.(f)      # threaded log10 the entire array
    fft!(buffer)                      # FFT
    @strided buffer .= GF .* buffer   # convolve
    ifft!(buffer)                     # IFFT
    f_result = real(buffer)           # get the real part
    @strided f_result .= (10 .^ f_result)  # threaded exp
    norm = sum(f)/sum(f_result)       # compute normalization
    @strided f_result .= f_result .* norm  # normalize
    return f_result
end

# THESE FUNCTIONS ONLY WORK FOR CUBES. From arxiv:1209.2043

# Computes Hessian of log-gaussian-smoothed density field
# f_Rn is the smoothed field
function slow_hessian_from_smoothed(f_Rn::AbstractArray{T,3}, R_S, kv) where T
    dims = size(f_Rn)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))
    f_Rn_hat = fft(f_Rn)

    kx = kv[1]
    ky = kv[2]
    kz = kv[3]

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (1,1)
                hessian[x,y,z,1] =
                    - kx[x] * kx[x] * R_S^2 * f_Rn_hat[x,y,z]
                # (1,2)
                hessian[x,y,z,2] =
                    - kx[x] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
                # (1,3)
                hessian[x,y,z,3] =
                    - kx[x] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
                # (2,2)
                hessian[x,y,z,4] =
                    - ky[y] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
                # (2,3)
                hessian[x,y,z,5] =
                    - ky[y] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
                # (3,3)
                hessian[x,y,z,6] =
                    - kz[z] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end
    hessian[:,:,:,1] .= ifft(hessian[:,:,:,1])
    hessian[:,:,:,2] .= ifft(hessian[:,:,:,2])
    hessian[:,:,:,3] .= ifft(hessian[:,:,:,3])
    hessian[:,:,:,4] .= ifft(hessian[:,:,:,4])
    hessian[:,:,:,5] .= ifft(hessian[:,:,:,5])
    hessian[:,:,:,6] .= ifft(hessian[:,:,:,6])
    real(R_S^2 .* hessian)
end

function fast_hessian_from_smoothed(f_Rn::AbstractArray{T,3}, R_S, kv) where T
    dims = size(f_Rn)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))
    f_Rn_hat = fft(f_Rn)

    kx = kv[1]
    ky = kv[2]
    kz = kv[3]

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (1,1)
                hessian[x,y,z,1] =
                    - kx[x] * kx[x] * R_S^2 * f_Rn_hat[x,y,z]
                # (1,2)
                hessian[x,y,z,2] =
                    - kx[x] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
                # (1,3)
                hessian[x,y,z,3] =
                    - kx[x] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
                # (2,2)
                hessian[x,y,z,4] =
                    - ky[y] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
                # (2,3)
                hessian[x,y,z,5] =
                    - ky[y] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
                # (3,3)
                hessian[x,y,z,6] =
                    - kz[z] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end

    for j in 1:6
        ifft!(@view (hessian[:,:,:,j]))
    end
    real(R_S^2 .* hessian)
end

function reduce_RAM_hessian_from_smoothed(f_Rn::AbstractArray{T,3}, R_S, kv, output_directory) where T
    dims = size(f_Rn)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))
    f_Rn_hat = fft(f_Rn)
    id_strings = ["_1", "_2", "_3", "_4", "_5", "_6"]

    kx = kv[1]
    ky = kv[2]
    kz = kv[3]

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (1,1)
                hessian[x,y,z] =
                    - kx[x] * kx[x] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[1], hessian)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                #(1,2)
                hessian[x,y,z] =
                    - kx[x] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[2], hessian)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (1,3)
                hessian[x,y,z] =
                    - kx[x] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[3], hessian)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (2,2)
                hessian[x,y,z] =
                    - ky[y] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[4], hessian)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (2,3)
                hessian[x,y,z] =
                    - ky[y] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]

            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[5], hessian)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (3,3)
                hessian[x,y,z] =
                    - kz[z] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]

            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[6], hessian)
    hessian=nothing

end

function signatures_from_hessian(hes::AbstractArray{T,4}) where T
    """
    Function to calculate the signatures from a given hessian.
    """

    hsize = size(hes)[1:3]
    sigs = zeros(T, (hsize[1], hsize[2], hsize[3], 3) )
    Threads.@threads for i in 1:hsize[1]
        for j in 1:hsize[2]
            for k in 1:hsize[3]

                # StaticArray (SA) here for speed, we know the dimensions are 3x3
                #the hessian matrix is symmetric
                hes_slice = SA[
                    hes[i,j,k,1]  hes[i,j,k,2]  hes[i,j,k,3];
                    hes[i,j,k,2]  hes[i,j,k,4]  hes[i,j,k,5];
                    hes[i,j,k,3]  hes[i,j,k,5]  hes[i,j,k,6]
                ]

                e1, e2, e3 = sort(real(eigvals(hes_slice)), rev=false)

                sigs[i,j,k,1] = (abs(e3 / e1) *
                    abs(e3)) * (θ(-e1) * θ(-e2) * θ(-e3))
                sigs[i,j,k,2] = (abs(e2 / e1) *
                    xθ(1-abs(e3 / e1))) * (abs(e2) * θ(-e1) * θ(-e2) )
                sigs[i,j,k,3] = (xθ(1-abs(e2 / e1)) *
                    xθ(1-abs(e3 / e1))) * (abs(e1) * θ(-e1))
            end
        end
    end

    # category shape strength 𝒮 (arxiv:1209.2043 eq 7)
    return sigs
end

"""
In this function we start at minimum length scale to calculate signatures
then iterate to the maximum scale of the problem, calculating the signatures
for each smoothing scale. We then chose the maximum signature over all scales
at each point to categorize whether that point is a wall, filament, or cluster.
"""
function maximum_signature(Rs, field::AbstractArray{T,3}; alg=:NEXUSPLUS) where T

    if alg ∉ (:NEXUS, :NEXUSPLUS)
        throw(ArgumentError("alg must be either :NEXUS or :NEXUSPLUS"))
    end

    nx,ny,nz = size(field)

    #we need to make sure the field has no 0 values
    field = field .+ 0.0001

    #calculate wave vectors for our field
    wave_vecs = wavevectors3D(T, (nx,ny,nz))

    #calculate signatures at each scale Rn, determine max
    sigmax = zeros(T,nx,ny,nz,3)

    for (R_index, R) in enumerate(Rs)

        if alg == :NEXUS
            f_Rn = smooth_gauss(field, R, wave_vecs)
        elseif alg == :NEXUSPLUS
            f_Rn = smooth_loggauss(field, R, wave_vecs)
        end

        H_Rn = fast_hessian_from_smoothed(f_Rn, R, wave_vecs)

        sigs_Rn = signatures_from_hessian(H_Rn)

        for ix = 1:nx, iy = 1:ny, iz = 1:nz, sigtype = 1:3
            sigmax[ix, iy, iz, sigtype] = max(
                sigmax[ix, iy, iz, sigtype], sigs_Rn[ix, iy, iz, sigtype])
        end
    end

    return sigmax
end

"""
In this function we start at minimum length scale to calculate signatures
then iterate to the maximum scale of the problem, calculating the signatures
for each smoothing scale. We then chose the maximum signature over all scales
at each point to categorize whether that point is a wall, filament, or cluster.
"""
function reduce_RAM_maximum_signature(Rs, output_directory, field::AbstractArray{T,3}; alg=:NEXUSPLUS) where T

    if alg ∉ (:NEXUS, :NEXUSPLUS)
        throw(ArgumentError("alg must be either :NEXUS or :NEXUSPLUS"))
    end

    nx,ny,nz = size(field)

    #we need to make sure the field has no 0 values
    field = field .+ 0.0001

    #calculate wave vectors for our field
    wave_vecs = wavevectors3D(T, (nx,ny,nz))

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

    print("\n")
    print("Computing All Hessian Values: ")
    print("\n")

    #create all hessians
    for (R_index, R) in enumerate(Rs)

        if alg == :NEXUS
            f_Rn = smooth_gauss(field, R, wave_vecs)
        elseif alg == :NEXUSPLUS
            f_Rn = smooth_loggauss(field, R, wave_vecs)
        end

        #calculate hessian subvolumes
        reduce_RAM_hessian_from_smoothed(f_Rn, R, wave_vecs, output_directory)

    end

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

    print("\n")
    print("Compiling Hessians & Calculating Max Signatures: ")
    print("\n")

    #subvolumes
    subvolumes = ["000", "001", "010", "100", "011", "110", "101", "111"]
    nx_sub,ny_sub,nz_sub = nx ÷ 2,ny ÷ 2,nz ÷ 2

    for (sub, subvolume) in enumerate(subvolumes)

        #calculate signatures at each scale Rn, determine max
        sigmax_sub = zeros(T,nx_sub,ny_sub,nz_sub,3)

        for (R_index, R) in enumerate(Rs)

            if alg == :NEXUS
                f_Rn = smooth_gauss(field, R, wave_vecs)
            elseif alg == :NEXUSPLUS
                f_Rn = smooth_loggauss(field, R, wave_vecs)
            end

            #combine hessian components
            compile_sub_hessians(f_Rn, R, subvolume, output_directory)

            #load back in necessary hessian subvolume
            @load output_directory * "hessian_component_"*subvolume*"_full.jld2" hessian_sub

            #calculate signatures on subvolume
            sigs_Rn = signatures_from_hessian(hessian_sub)

            #determine maximum signatures
            for ix = 1:nx_sub, iy = 1:ny_sub, iz = 1:nz_sub, sigtype = 1:3
                sigmax_sub[ix, iy, iz, sigtype] = max(
                    sigmax_sub[ix, iy, iz, sigtype], sigs_Rn[ix, iy, iz, sigtype])
            end

        end

        @save output_directory * "max_sigs_"*subvolume*".jld2" sigmax_sub

        #delete used hessian subvolume
        file = output_directory * "hessian_component_"*subvolume*"_full.jld2"
        run(`rm $file`)

    end
end


using AbstractDifferentiation
using Distributions
using ForwardDiff
using GalacticOptim
using LinearAlgebra
using Optim
using Pathfinder
using Random
using ReverseDiff
using SciMLBase
using Test
using Transducers

@testset "single path pathfinder" begin
    @testset "IsoNormal" begin
        # here pathfinder finds the exact solution after 1 iteration
        logp(x) = -sum(abs2, x) / 2
        ∇logp(x) = -x
        ndraws = 100
        rngs = if VERSION ≥ v"1.7"
            [MersenneTwister(), Random.default_rng()]
        else
            [MersenneTwister()]
        end
        seed = 42
        @testset for dim in [1, 5, 10, 100], rng in rngs
            executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx()

            init = randn(dim)
            Random.seed!(rng, seed)
            result = @inferred pathfinder(logp, ∇logp; init, ndraws, rng, executor)
            @test result isa PathfinderResult
            @test result.input === (logp, ∇logp)
            @test result.optim_prob isa SciMLBase.OptimizationProblem
            @test result.logp(init) ≈ logp(init)
            @test result.rng === rng
            @test result.optimizer ===
                Pathfinder.default_optimizer(Pathfinder.DEFAULT_HISTORY_LENGTH)
            fit_distribution = result.fit_distribution
            @test fit_distribution isa MvNormal
            @test fit_distribution.μ ≈ zeros(dim)
            @test fit_distribution.Σ isa Pathfinder.WoodburyPDMat
            @test fit_distribution.Σ ≈ I
            @test size(fit_distribution.Σ.B) == (dim, 2) # history contains only 1 iteration
            @test result.draws isa AbstractMatrix
            @test size(result.draws) == (dim, ndraws)
            @test result.fit_distribution_transformed === result.fit_distribution
            @test result.draws_transformed === result.draws
            @test result.num_tries ≥ 1
            @test result.optim_solution isa SciMLBase.OptimizationSolution
            @test result.optim_trace isa Pathfinder.OptimizationTrace
            @test result.fit_distributions isa Vector{typeof(fit_distribution)}
            @test length(result.fit_distributions) == length(result.optim_trace)
            @test result.fit_distributions[result.fit_iteration + 1] == fit_distribution
            @test result.fit_iteration ==
                argmax(getproperty.(result.elbo_estimates, :value))

            Random.seed!(rng, seed)
            result2 = pathfinder(logp, ∇logp; init, ndraws, rng, executor)
            @test result2.fit_iteration == result.fit_iteration
            @test result2.draws == result.draws
            @test getproperty.(result2.elbo_estimates, :value) ==
                getproperty.(result.elbo_estimates, :value)

            ndraws = 2
            result3 = pathfinder(logp, ∇logp; init, ndraws, executor)
            @test size(result3.draws) == (dim, ndraws)
        end
    end
    @testset "MvNormal" begin
        #! format: off
        Σ = [
            2.71   0.5    0.19   0.07   1.04
            0.5    1.11  -0.08  -0.17  -0.08
            0.19  -0.08   0.26   0.07  -0.7
            0.07  -0.17   0.07   0.11  -0.21
            1.04  -0.08  -0.7   -0.21   8.65
        ]
        #! format: on
        P = inv(Symmetric(Σ))
        logp(x) = -dot(x, P, x) / 2
        ∇logp(x) = -(P * x)
        dim = 5
        ad_backend = AD.ReverseDiffBackend()
        ndraws_elbo = 100
        rngs = if VERSION ≥ v"1.7"
            [MersenneTwister(), Random.default_rng()]
        else
            [MersenneTwister()]
        end
        seed = 38
        optimizer = Optim.LBFGS(; m=6)
        @testset for rng in rngs
            executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx()

            Random.seed!(rng, seed)
            result = @inferred pathfinder(
                logp; rng, dim, optimizer, ndraws_elbo, ad_backend, executor
            )
            @test result.input === logp
            @test result.fit_distribution.Σ ≈ Σ rtol = 1e-1
            @test result.optimizer == optimizer
            Random.seed!(rng, seed)
            result2 = pathfinder(
                logp; rng, dim, optimizer, ndraws_elbo, ad_backend, executor
            )
            @test result2.fit_distribution == result.fit_distribution
            @test result2.draws == result.draws
            @test getproperty.(result2.elbo_estimates, :value) ==
                getproperty.(result.elbo_estimates, :value)
        end
        @testset "kwargs forwarded to solve" begin
            Random.seed!(42)
            i = 0
            callback = (args...,) -> (i += 1; false)
            pathfinder(logp; dim, callback)
            @test i ≠ 6

            Random.seed!(42)
            i = 0
            pathfinder(logp; dim, callback, maxiters=5)
            @test i == 6
        end
    end
    @testset "retries" begin
        @testset "logp returning NaN" begin
            dim = 5
            nfail = 20
            logp(x) = i ≤ nfail ? NaN : -sum(abs2, x) / 2
            callback = (args...,) -> (i += 1; true)
            i = 1
            result = pathfinder(logp; dim, callback)
            @test result.fit_distribution.μ ≈ zeros(dim) atol = 1e-6
            @test result.fit_distribution.Σ ≈ diagm(ones(dim)) atol = 1e-6
            @test result.num_tries == nfail + 1
            @test result.optim_prob.u0 == result.optim_trace.points[1]
            i = 1
            init = randn(dim)
            result2 = pathfinder(logp; init, callback, ntries=nfail)
            @test !isapprox(result2.fit_distribution.μ, zeros(dim); atol=1e-6)
            @test result2.fit_iteration == 0
            @test isempty(result2.elbo_estimates)
            @test result2.num_tries == nfail
            @test result2.optim_prob.u0 == result2.optim_trace.points[1]
        end
    end
    @testset "UniformSampler" begin
        @test_throws DomainError Pathfinder.UniformSampler(-1.0)
        @test_throws DomainError Pathfinder.UniformSampler(0.0)
        @testset for scale in [1, 2], seed in [42, 38]
            sampler_fun = Pathfinder.UniformSampler(scale)
            rng = Random.seed!(Random.default_rng(), seed)
            x = zeros(100)
            sampler_fun(rng, x)
            @test all(-scale .≤ x .≤ scale)
            Random.seed!(rng, seed)
            x2 = zeros(100)
            x2 .= rand.(rng) .* 2scale .- scale
            @test x2 == x
        end
    end
    @testset "errors if no gradient provided" begin
        logp(x) = -sum(abs2, x) / 2
        init = randn(5)
        prob = SciMLBase.OptimizationProblem(logp, init, nothing)
        @test_throws ArgumentError pathfinder(prob)
        fun = SciMLBase.OptimizationFunction(logp, GalacticOptim.AutoForwardDiff())
        prob = SciMLBase.OptimizationProblem(fun, init, nothing)
        @test_throws ArgumentError pathfinder(prob)
    end
    @testset "errors if neither dim nor init valid" begin
        logp(x) = -sum(abs2, x) / 2
        @test_throws ArgumentError pathfinder(logp)
        @test_throws ArgumentError pathfinder(logp; dim=0)
        pathfinder(logp; dim=3)
        pathfinder(logp; init=randn(5))
    end
end


function test_integral_form()

    using EffectiveWaves

    # physical parameters
    θin = 0.0
    k=1.;
    ho = 2;

    medium = Medium(1.0,1.0+0.0im)
    ω = real(k*medium.c)
    # specie = Specie(ρ=0.1,r=1.0, c=0.5, volfrac=0.2)
    specie = Specie(ρ=0.1,r=1.0, c=0.5, volfrac=0.1)
    # specie = Specie(ρ=0.1,r=0.1, c=0.5, volfrac=0.1)
    specie2 = Specie(ρ=2.0, r=1.0, c=0.1, volfrac=0.15)

    # From effective wave theory
    k_eff0 = wavenumber_low_volfrac(ω, medium, [specie]; hankel_order = ho)
    k_effs = wavenumbers(ω, medium, [specie]; mesh_points = 10, tol = 1e-8, hankel_order = ho)
    k_effs = sort(k_effs, by=imag)

    eff_medium = effective_medium(medium, [specie])
    ω/eff_medium.c
    k_effs2 = wavenumbers(ω, medium, [specie2]; tol = 1e-8, hankel_order = ho)
    k_effs2 = sort(k_effs2, by=imag)

    using OffsetArrays, ApproxFun, IterTools
    include("src/integral_form/integral_form.jl")

    # (x, (MM_quad,b_mat)) = average_wave_system(ω, medium, specie; θin = θin, mesh_points = 501, hankel_order=ho);
    X = 0.0:0.005:30.0
    (MM_quad,b_mat) = average_wave_system(ω, X, medium, specie;θin = θin, hankel_order=ho);

    # discretization parameters
    J = length(collect(x)) - 1

    len = (J + 1) * (2ho + 1)
    MM_mat = reshape(MM_quad, (len, len));
    b = reshape(b_mat, (len));

    As = MM_mat\b;
    As_mat = reshape(As, (J+1, 2ho+1));

    amps_eff1 = scattering_amplitudes_average(ω, x, medium, [specie];
            k_eff = k_effs[1], hankel_order = ho, θin=θin, tol=1e-8)
    amps_eff2 = scattering_amplitudes_average(ω, x, medium, [specie];
            k_eff = k_effs[2], hankel_order = ho, θin=θin, tol=1e-8)
    amps_eff3 = scattering_amplitudes_average(ω, x, medium, [specie];
            k_eff = k_effs[end], hankel_order = ho, θin=θin, tol=1e-8)

    amps0_eff = scattering_amplitudes_average(ω, x, medium, [specie];
            k_eff = k_eff0, hankel_order = ho, θin=θin)
    amps2_eff = scattering_amplitudes_average(ω, x, medium, [specie2];
            k_eff = k_effs2[2], hankel_order = ho, θin=θin, tol=1e-8)
    # sanity check: the abs of reflection coefficients should always be smaller than one.
    reflection_coefficient(ω, medium, specie2; amps = amps2_eff, θin = θin)
    reflection_coefficient(ω, medium, specie; amps = amps_eff2, θin = θin)
    reflection_coefficient(ω, medium, specie; amps = amps_eff3, θin = θin)
    reflection_coefficient(ω, medium, specie; amps = amps_eff1, θin = θin)

    error0_eff = reshape( abs.(MM_mat*amps0_eff.amplitudes[:] .- b), (J+1, 2ho+1))
    error_eff1 = reshape( abs.(MM_mat*amps_eff1.amplitudes[:] .- b), (J+1, 2ho+1))
    error_eff2 = reshape( abs.(MM_mat*amps_eff2.amplitudes[:] .- b), (J+1, 2ho+1))
    error_eff3 = reshape( abs.(MM_mat*amps_eff3.amplitudes[:] .- b), (J+1, 2ho+1))

    error2_eff = reshape( abs.(MM_mat*amps2_eff.amplitudes[:] .- b), (J+1, 2ho+1))
    # With a completely wrong wavenumber I can still get a small error by just scalling the amplitudes
    amps2_eff = scattering_amplitudes_average(ω, x, medium, [specie];
            k_eff = k_effs2[2], hankel_order = ho, θin=θin)
    error2_eff = reshape( abs.(MM_mat*amps2_eff.amplitudes[:] .- b), (J+1, 2ho+1))
    # scale_amplitudes_effective(ω, k_effs[1], amps2_eff.amplitudes, medium, [specie]; θin = θin)

    error = reshape( abs.((MM_mat*As)./b .- 1.0+0.0im), (J+1, 2ho+1))

    using Plots; pyplot(linewidth=2) # unicodeplots()
    plot(xlabel = "depth (1 wavelength = 2π )", ylabel = "error %", ylims=(-0.1,0.5), title="Transmitted wave errors")
    plot!(x,error_eff1[:,ho+1], label = "Eff. error")
    plot!(x,error0_eff[:,ho+1], linestyle=:dash, label = "Eff. low φ error")
    plot!(x,error2_eff[:,ho+1], linestyle=:dot, label = "Eff. wrong k_eff error")
    plot!(x,error[:,ho+1], linestyle=:dashdot, label = "Integral method error")

    plot(x, [real.(As_mat[:,ho+1]),imag.(As_mat[:,ho+1])], labels = ["real sol." "imag sol."])
    plot!(x, [real.(amps_eff1.amplitudes[:,ho+1]),imag.(amps_eff1.amplitudes[:,ho+1])],
        labels = ["real eff. 1" "imag eff. 1"], linestyle=:dash)
    plot!(x, [real.(amps_eff2.amplitudes[:,ho+1]),imag.(amps_eff2.amplitudes[:,ho+1])],
        labels = ["real eff. 2" "imag eff. 2"], linestyle=:dash)

    plot(x, [abs.(As_mat[:,ho+1] .- amps_eff1.amplitudes[:,ho+1])],
            labels = ["diff abs eff. " "diff imag eff."])

    plot(x, [real.(As_mat[:,ho+2]),imag.(As_mat[:,ho+2])], labels = ["real sol." "imag sol."])
    plot!(x, [real.(amps_eff1.amplitudes[:,ho+2]),imag.(amps_eff1.amplitudes[:,ho+2])],
        labels = ["real eff. 1" "imag eff. 1"], linestyle=:dash)

    # plot!(x, [real.(amps_eff3.amplitudes[:,ho+1]),imag.(amps_eff3.amplitudes[:,ho+1])],
    #     labels = ["real eff. 3" "imag eff. 3"], linestyle=:dash)
    # plot!(x, [real.(amps0_eff.amplitudes[:,ho+1]),imag.(amps0_eff.amplitudes[:,ho+1])],
    #     labels = ["real φ eff." "imag φ eff."], linestyle=:dot)

    is = 250:(length(collect(x))-1)
    plot(x[is], log.(abs.(As_mat[is,ho+1])), labels = "abs sol.")
    plot!(x[is], log.(abs.(As_eff_mat[is,ho+1])), labels = "abs eff.")
    # b_mat ≈ [ sum(MM_quad[l,m,j,n]*A_mat[j,n] for j=1:(J+1), n=1:(2M+1)) for l=1:(J+1), m=1:(2M+1)]
end

# ints = [ check_integration(1.0+1.0im; h = 1./n) for n=10:30:310]

function test_check_integration()
#from Mathematics, had several errors when running
    math = [
            [65.0544 - 146.891im, 24.1498 - 8.37225im, -5.10012 + 3.74221im
                , 0.940428 - 3.4201im, 7.25312 - 0.0534466im, 22.9058 + 38.434im, -49.5574 + 92.4377im]
            ,[-32.7224 + 1.52201im, -12.033 - 16.2499im, 5.93884 - 10.8663im
                , 8.51838 + 3.68509im, -4.94285 + 9.32462im, -12.0994 - 5.48456im, 3.6755 - 13.9906im]
            ,[-9.82076 - 18.6418im, 10.2188 - 12.2782im, 11.8699 + 5.34705im
                , -2.698 + 12.0431im, -13.1158 - 0.218682im, -2.92199 - 13.8296im, 12.9993 - 6.4169im]
         ];
    # result of ints = check_integration(h = 1/400., max_x = 8.0)
    julia = [
            [64.9882-146.903im, 24.1561-8.37069im, -5.10105+3.74215im
                , 0.940588-3.4198im, 7.25286-0.0526824im, 22.9026+38.4402im, -49.5882+92.5011im]
            ,[-32.8359+1.66127im, -12.0424-16.2309im, 5.93782-10.8642im
                , 8.51816+3.68541im, -4.94155+9.32627im, -12.1151-5.49728im, 3.77515-13.8555im]
            ,[-9.88661-18.6492im, 10.2111-12.2772im, 11.8693+5.34716im
                , -2.69806+12.0429im, -13.1162-0.21808im, -2.92135-13.8369im, 12.9773-6.359im]
        ];

    using Plots; pyplot()

    scatter(-3:3, abs.(julia[1]./math[1] .- 1), xlab = "hankel_order", ylab = "rel. error", lab ="", title="Mathematica vs Julia for x = 0.0")
    scatter(-3:3, abs.(julia[2]./math[2] .- 1), xlab = "hankel_order", ylab = "rel. error", lab ="", title="Mathematica vs Julia for x = 1.0")
    scatter(-3:3, abs.(julia[3]./math[3] .- 1), xlab = "hankel_order", ylab = "rel. error", lab ="", title="Mathematica vs Julia for x = 2.0")

#from julia
# using JLD
# ints =  first(values(load("integrated_As.jld")))
# data = [ [ints[i][j][k] for i in eachindex(ints)] for j=1:3, k=1:7];
# # i=2;j=3;k=5; ints[i][j][k] == data[j,k][i]
# # x = 0. => j = 1
# M = 3;
#
# hs =  [ 1./n for n=10:30:310]
# j = 3
#  plot()
#  for k=1:7
#      plot!(hs[4:end], abs.(data[j,k][4:end]), label = "m = $(k - M - 1)", xlab = "mesh element size")
#      scatter!([hs[end]], [abs(math[j][k])])
#  end
#   gui()
#
#  plot()
#  for k=1:7
#      plot!(hs[4:end], 1 - abs.(data[j,k][4:end])./abs(data[j,k][end]), label = "m = $(k - M - 1)")
#  end
#   gui()
end

# tests the integration scheme
function check_integration(k_eff::Complex{Float64} = 1.0+1.0im; k=1.,a=1., h = a*k/55., max_x = 6.0)

    A(n,x) = im^Float64(n)*exp(-im*n*θin)*exp(im*x*k_eff)
    k=1.; a=1.;
    # physical parameters
    θin = 0.3

    # discretization parameters
    M = 3;
    J = Int(round(max_x/h)) # choose an even number for Integration schemes
    x = OffsetArray{Float64}(undef, 0:J)
    x[0:J] = (0:J)*h

    # trapezoidal
    σ =  OffsetArray(trap_scheme(collect(x); xn=max_x), 0:J)

    PQ_quad = intergrand_kernel(x, a*k; θin = θin, M = M)

    # Apply integration scheme
    for j=0:J
        PQ_quad[:,:,j+1,:] *= σ[j]
    end
    # PQ_quad = [σ[j]*PQ_quad[l+1,m+M+1,j+1,n+M+1] for l=0:J, m=-M:M, j=0:J, n=-M:M]

    len = (J + 1) * (2M + 1)
    PQ_mat = reshape(PQ_quad, (len, len))

    A_mat = [ A(m,x[l]) for l = 0:J, m = -M:M]
    As = reshape(A_mat, (len))
    As_integrated = reshape(PQ_mat*As, (J+1,2M+1))

    i1 = Int(round(1/h)) # for x1 = 1
    i2 = Int(round(2/h)) # for x1 = 2
    [As_integrated[1,:], As_integrated[1+i1,:], As_integrated[1+i2,:]]
end


module TwoSourceSound
	using QuadGK
	using Cubature
	using PyCall
	const SPEED_OF_SOUND = 343

	struct TwoSourceOpts
		f0::Float64
		f1::Float64
		d::Float64
		center::Tuple{Float64, Float64}
	end

	export TwoSourceOpts

	function calcAvgAmplitudeIntegral(a::Float64, b::Float64, f::Float64, df::Float64)
		k = 2 * pi / SPEED_OF_SOUND

		return -(b*cos(a*f*k) + a*cos(b*f*k))/(a*b*k*df)
	end

	function calcAvgAmplitude(opts::TwoSourceOpts, p::Tuple{Float64, Float64})
		df = opts.f1 - opts.f0
		spk0x = opts.center[1] - opts.d / 2
		spk1x = opts.center[1] + opts.d / 2
		a = sqrt((p[1] - spk0x)^2 + (p[2] - opts.center[2])^2)
		b = sqrt((p[1] - spk1x)^2 + (p[2] - opts.center[2])^2)

		return calcAvgAmplitudeIntegral(a, b, opts.f1, df) - calcAvgAmplitudeIntegral(a, b, opts.f0, df)
	end

	function calcAvgAmplitudeDevIntegral(a::Float64, b::Float64, avgAmp::Float64, f::Float64, df::Float64)
		k = 2 * pi / SPEED_OF_SOUND

		m = a*avgAmp*b*f*k + a*cos(b*f*k) + b*cos(a*f*k)
		n = sign(sin(a*f*k) + sin(b*f*k) - avgAmp)
		l = a * b * k * df

		return (-m)/l
	end

	function calcAvgAmplitudeDev(opts::TwoSourceOpts, p::Tuple{Float64, Float64})
		df = opts.f1 - opts.f0
		spk0x = opts.center[1] - opts.d / 2
		spk1x = opts.center[1] + opts.d / 2
		a = sqrt((p[1] - spk0x)^2 + (p[2] - opts.center[2])^2)
		b = sqrt((p[1] - spk1x)^2 + (p[2] - opts.center[2])^2)
		avgAmp = calcAvgAmplitude(opts, p)

		f(x) = abs(sin(2*pi*a*x/SPEED_OF_SOUND) + sin(2*pi*b*x/SPEED_OF_SOUND) - avgAmp)

		return (1/df) * quadgk(f, opts.f0, opts.f1, rtol=0.001)[1]
	end

	function calcAvgPlaneAmplitudeDev(d::Float64, fRange::Tuple{Float64, Float64}, xRange::Tuple{Float64, Float64}, yRange::Tuple{Float64, Float64})
		df = fRange[2] - fRange[1]
		dx = xRange[2] - xRange[1]
		dy = yRange[2] - yRange[1]

		A_0(f, x, y) = sin(2*pi*sqrt(x^2 + y^2)*f/SPEED_OF_SOUND) + sin(2*pi*sqrt((x - d)^2 + y^2)*f/SPEED_OF_SOUND)
		A_bar_indef(f, x, y) = -(SPEED_OF_SOUND/(2*pi*df)) * ((cos(2*pi*f*sqrt((x-d)^2 + y^2)))/(sqrt((x-d)^2 + y^2)) + (cos(2*pi*f*sqrt(x^2 + y^2)))/(sqrt(x^2 + y^2)))
		A_bar(x, y) = A_bar_indef(fRange[2], x, y) - A_bar_indef(fRange[1], x, y)

		func(y, x, f) = abs(A_0(f, x, y) - A_bar(x, y))
		func_cubature(x) = func(x[1], x[2], x[3])

		return 1/(dy * dx * df) * hcubature(func_cubature, (yRange[1], xRange[1], fRange[1]), (yRange[2], xRange[2], fRange[2]), reltol=0.01)[1]
	end

	export calcAvgAmplitudeDev
end


function snr(data,boxcar_width)
    first_quant = quantile(data |> vec, 0.25)
    noise_std = std(filter(x->x<first_quant,data))
    data ./ noise_std
end


export hessian_dense
"""Compute the full hessian of an nlp (AbstractNLPModel). Should only be used
for lower dimensions problems"""
function hessian_dense(nlp, x)
    n = length(x)
    H = hess(nlp, x)
    tempH = (H + tril(H, -1)')
    H = Matrix(Symmetric(tempH, :L))
    return H
end