# test operators with pd matrix types
using LinearAlgebra, PDMats
using Test

@testset "scalar multiplication" begin
    printstyled("Testing scalar multiplication\n"; color=:blue)
    pm1 = PDMat(Matrix(1.0I, 3, 3))
    pm2 = PDiagMat(ones(3))
    pm3 = ScalMat(3, 1)

    pm1a = PDMat(Matrix(3.0I, 3, 3))
    pm2a = PDiagMat(3.0 .* ones(3))
    pm3a = ScalMat(3, 3)

    pmats = Any[pm1, pm2, pm3]
    pmatsa = Any[pm1a, pm2a, pm3a]

    for i in 1:length(pmats)
        @test Matrix(3.0 * pmats[i]) == Matrix(pmatsa[i])
        @test Matrix(pmats[i] * 3.0) == Matrix(pmatsa[i])
        @test Matrix(3 * pmats[i]) == Matrix(pmatsa[i])
        @test Matrix(pmats[i] * 3) == Matrix(pmatsa[i])
    end
end

# issue #121
@test isposdef(PDMat([1.0 0.0; 0.0 1.0]))
@test isposdef(PDiagMat([1.0, 1.0]))
@test isposdef(ScalMat(2, 3.0))


include("../server/simulation/stimulus/stimulus.jl")

function test_generating()
    stream = PoissonStream(UInt64(13163))
    println(next(stream))
    println(next(stream))
    println(next(stream))
    println(next(stream))
end

const TOTAL_SPIKES = 25

function test_stream_reset()
    println("Testing PoissonStream with reset")

    # Create a stream that generates spikes
    stream = PoissonStream(UInt64(13163), 0.2)
    # Capture isi value for testing after a reset
    isi = stream.isi
    # println("isi: ", stream.isi)

    # Create array to capture spike stream for testing after reset
    spikes = []

    # Fill array with output of stream
    for i in 1:TOTAL_SPIKES
        step!(stream)
        push!(spikes, output(stream))
        # println(output(stream))
    end

    # Now reset stream. This should reset the stream back to its begining.
    reset!(stream)

    # They should be the same
    @assert isi == stream.isi "ISI should be equal"

    # println("isi: ", stream.isi)
    # And the spikes should be too.
    for i in 1:TOTAL_SPIKES
        step!(stream)
        @assert output(stream) == spikes[i] "Invalid spike at index: $i"
        # println(output(stream))
    end

    println("Passed")
end

function test_stream()
    println("Testing PoissonStream with no reset")
    stream = PoissonStream(UInt64(13163), 0.2)
    isi = stream.isi
    # println("isi: ", stream.isi)
    spikes = []
    for i in 1:TOTAL_SPIKES
        step!(stream)
        push!(spikes, output(stream))
        # println(output(stream))
    end

    nisi = stream.isi
    @assert isi != nisi "ISI should NOT be equal"
    # println("isi: ", stream.isi)
    atleast_one_ne = false
    for i in 1:TOTAL_SPIKES
        step!(stream)

        if output(stream) != spikes[i]
            atleast_one_ne = true
            break;
        end
    end
    @assert atleast_one_ne "Atleast one spike should be different"
    println("Passed")
end

function test_stream_two()
    println("Testing two different poisson streams")

    stream = PoissonStream(UInt64(13163), 0.2)
    isi = stream.isi

    spikes = []
    for i in 1:TOTAL_SPIKES
        step!(stream)
        # println(output(stream))
        push!(spikes, output(stream))
    end

    stream2 = PoissonStream(UInt64(11862), 0.2)
    atleast_one_ne = false
    # println("----------")

    for i in 1:TOTAL_SPIKES
        step!(stream2)
        # println(output(stream))
        if output(stream2) != spikes[i]
            atleast_one_ne = true
            break;
        end
    end
    @assert atleast_one_ne "Atleast one spike should be different between streams"

    println("Passed")
end

test_stream_two()
test_stream()
test_stream_reset()


"""
    compute_bin_averages(O::Array{Float64, 1}, n::Int)

Computes the bin averages, used in computing the autocorrelation time.

Arguments
* `O`       : Array representing an observable, e.g. E or M
* `n`       : Number of observations per bin
Returns
* `bin_avgs`: Bin averages
"""
function compute_bin_averages(O::Array{Float64,1}, n::Int)
    N = length(O)
    N_B = Int(floor(N / n))
    bin_avgs = Array{Float64,1}(undef, N_B)

    for k = 1:N_B
        a = (k - 1) * n + 1
        b = k * n
        bin_avgs[k] = mean(O[a:b])
    end

    return bin_avgs
end

"""
    compute_autocorrelation_time(model::Union{Ising, Potts}, ns::Array{Int, 1})

Computes the integrated autocorrelation time using the binning method. The autocorrelation
time is the value at which the output sequence (either energy `:e` or magnetization `:m` in
the return dictionary) asymptotes at.

Arguments
* `model`: A Potts or Ising type model
* `ns`   : Array of bin sizes
Returns
* Dictionary with keys `[:e, :m]` with bin method values to be plotted against `ns`
"""
function compute_autocorrelation_time(model::Union{Ising,Potts}, ns::Array{Int,1})
    V = length(model.σ)
    e = model.observables.E ./ V
    m = abs.(model.observables.M) ./ V
    N = length(e)
    σ²_e = mean(e .^ 2) - mean(e)^2
    σ²_m = mean(m .^ 2) - mean(m)^2
    τ_int_e = []
    τ_int_m = []

    for n in ns
        N_B = floor(Int, N / n)
        e_bin_avgs = compute_bin_averages(e, n)
        m_bin_avgs = compute_bin_averages(m, n)
        ε²_e = sum((e_bin_avgs .- mean(e_bin_avgs)) .^ 2) / (N_B * (N_B - 1))
        ε²_m = sum((m_bin_avgs .- mean(m_bin_avgs)) .^ 2) / (N_B * (N_B - 1))
        push!(τ_int_e, ε²_e * N / 2σ²_e)
        push!(τ_int_m, ε²_m * N / 2σ²_m)
    end

    return Dict(:e => τ_int_e, :m => τ_int_m)
end

"""
    compute_autocorrelation_time(model::XY, ns::Array{Int, 1})

Computes the integrated autocorrelation time using the binning method. The autocorrelation
time is the value at which the output sequence (either energy `:e` or magnetization `:m` in
the return dictionary) asymptotes at.

Arguments
* `model`: An XY type model
* `ns`   : Array of bin sizes
Returns
* Dictionary with keys `[:e, :mx, my]` with bin method values to be plotted against `ns`
"""
function compute_autocorrelation_time(model::XY, ns::Array{Int,1})
    V = length(model.σ)
    e = model.observables.E ./ V
    mx = abs.(model.observables.Mx) ./ V
    my = abs.(model.observables.My) ./ V
    N = length(e)
    σ²_e = mean(e .^ 2) - mean(e)^2
    σ²_mx = mean(mx .^ 2) - mean(mx)^2
    σ²_my = mean(my .^ 2) - mean(my)^2
    τ_int_e = []
    τ_int_mx = []
    τ_int_my = []

    for n in ns
        N_B = floor(Int, N / n)
        e_bin_avgs = compute_bin_averages(e, n)
        mx_bin_avgs = compute_bin_averages(mx, n)
        my_bin_avgs = compute_bin_averages(my, n)
        ε²_e = sum((e_bin_avgs .- mean(e_bin_avgs)) .^ 2) / (N_B * (N_B - 1))
        ε²_mx = sum((mx_bin_avgs .- mean(mx_bin_avgs)) .^ 2) / (N_B * (N_B - 1))
        ε²_my = sum((my_bin_avgs .- mean(my_bin_avgs)) .^ 2) / (N_B * (N_B - 1))
        push!(τ_int_e, ε²_e * N / 2σ²_e)
        push!(τ_int_mx, ε²_mx * N / 2σ²_mx)
        push!(τ_int_my, ε²_my * N / 2σ²_my)
    end

    return Dict(:e => τ_int_e, :mx => τ_int_mx, :my => τ_int_my)
end

"""
    sim_dict_to_df(sims::Dict)

Converts a dictionary where the values are models into a DataFrame object.

Arguments
* `sims`: Dictionary where the values are simulations
Returns
* `df`  : DataFrame containing the simulation data
"""
function sim_dict_to_df(sims::Dict; include_magnetization::Bool = true)
    if include_magnetization
        df = DataFrame(
            L = Int[],
            β = Float64[],
            e = Float64[],
            c = Float64[],
            m = Float64[],
            χ = Float64[],
            U = Float64[],
        )

        for sim in values(sims)
            push!(
                df,
                [
                    sim.params.L,
                    sim.params.β,
                    sim.observables.statistics[:e],
                    sim.observables.statistics[:c],
                    sim.observables.statistics[:m],
                    sim.observables.statistics[:χ],
                    sim.observables.statistics[:U],
                ],
            )
        end
    else
        df = DataFrame(L = Int[], β = Float64[], e = Float64[], c = Float64[])

        for sim in values(sims)
            push!(
                df,
                [
                    sim.params.L,
                    sim.params.β,
                    sim.observables.statistics[:e],
                    sim.observables.statistics[:c],
                ],
            )
        end
    end

    sort!(df, [:L, :β])

    return df
end


@doc """
Enumerates all connected neutral networks for NKqLandscapes and all
approximately neutral networks for NKLandscapes, and summarizes their
properties.  The objective is to see if there are "giant" connected neutral
networks that contain almost all genotypes at a given fitness level as
predicted by the "percolation" model of Sergey Gavrilets for uncorrelated
landscapes.  (Reference needed.) An integer representation of genotypes
is used extensively.  The functions `gtoi` and
`itog` convert between the integer representation and the
array representation used elsewhere in NKLandscapes.jl.  If the arity
of the landscape is 2, then the integer representation can be
considered as a bit-string representation, and some functions require
that the arity of the landscape is 2.
"""
using NKLandscapes
using DataStructures

@doc """ function neutral_net_summary(counts, num_intervals)
Returns a list of information about the connected neutral nets at each non-empty fitness level
Each list element is a 5-tuple with the following components
*  The fit level (empty fit levels are not included)
*  The lower bound of the fit level
*  The number of genotypes at the fit level
*  The maximum size of a neutral net at the fit level
*  The ratio of maximum size neutral net over the number of genotypes
Note that if this ratio is close to 1.0, then a high proportion of the genotypes at the fitness level
  are in a "giant component" neutral net.
"""
function neutral_net_summary(counts, num_intervals)
  fit_increment = 1.0/num_intervals
  nn3 = map(x->x[3],counts)
  lb = minimum(nn3)
  ub = maximum(nn3)
  summary = []
  for i = lb:ub
    filtered_nn = [y[2] for y in filter(x->x[3]==i,counts)]
    if length(filtered_nn) > 0
      max_nn = maximum(filtered_nn)
      sum_nn = sum(filtered_nn)
      push!(summary,(i, i*fit_increment, sum_nn, max_nn, max_nn/sum_nn))
    end
  end
  return summary
end

@doc """ function neutral_net_summary(ls::NKLandscape)

Produces the same results as the previous function, but starting with an NK or NKq landscape.
"""

function neutral_net_summary(ls::Landscape, num_intervals::Int64=ls.n)
  lsf = lsfits(ls);
  flevs = fitlevs(ls, num_intervals, lsf);
  nn = neutralnets(ls,flevs)
  ncounts = netcounts(nn,flevs)
  neutral_net_summary(ncounts,num_intervals)
end

@doc """
Prints the the list of fitness levels returned by neutral_net_summary in CSV format.
"""
function print_neutral_nets(summary)
  @printf("fit lev\tfit_lb\t total count\t   max count\tpct max\n")
  for s in summary
    @printf("%2d\t%.3f\t%12d\t%12d\t%.5f\n",s[1],s[2],convert(Int,s[3]),convert(Int,s[4]),s[5])
  end
end


"""
`annualtoquarter(v)`

Convert from annual to quarter frequency... by dividing by 4.
"""
function annualtoquarter(v)
    v / 4
end

"""
`nominal_to_real(col, df; deflator_mnemonic=:GDPCTPI)`

Converts nominal to real values using the specified deflator.

## Arguments
- `col`: symbol indicating which column of `df` to transform
- `df`: DataFrame containining series for proper population measure and `col`

## Keyword arguments

- `deflator_mnemonic`: indicates which deflator to use to calculate real values. Default
  value is the FRED GDP Deflator mnemonic.
"""
function nominal_to_real(col, df; deflator_mnemonic=:GDPCTPI)
    return df[col] ./ df[deflator_mnemonic]
end

"""
```
percapita(m, col, df)
percapita(col, df, population_mnemonic)
```

Converts data column `col` of DataFrame `df` to a per-capita value.

## Arguments
- `col`: symbol indicating which column of data to transform
- `df`: DataFrame containining series for proper population measure and `col`
- `population_mnemonic`: a mnemonic found in df for some population measure.
"""
function percapita(m::AbstractModel, col::Symbol, df::DataFrame)
    population_mnemonic = get_setting(m, :population_mnemonic)
    percapita(col, df, population_mnemonic)
end
function percapita(col::Symbol, df::DataFrame, population_mnemonic::Symbol)
    df[col] ./ df[population_mnemonic]
end

"""
```
yt, yf = hpfilter(y, λ::Real)
```

Applies the Hodrick-Prescott filter ("H-P filter"). The smoothing parameter `λ` is applied
to the columns of `y`, returning the trend component `yt` and the cyclical component `yf`.
For quarterly data, one can use λ=1600.

Consecutive missing values at the beginning or end of the time series are excluded from the
filtering. If there are missing values within the series, the filtered values are all NaN.

See also:
```
Hodrick, Robert; Prescott, Edward C. (1997). "Postwar U.S. Business Cycles: An Empirical
Investigation". Journal of Money, Credit, and Banking 29 (1): 1–16.
```
"""
function hpfilter(y, λ::Real)
    # Convert y to vector
    if !isa(y, Vector)
        try
            y = vec(y)
        catch
            error("Series must be convertible to Vector")
        end
    end

    # Indices of consecutive NaN elements at beginning
    i = 1
    j = length(y)
    while isnan(y[i])
        i = i+1
    end
    while isnan(y[j])
        j = j-1
    end

    # Filter and adjust for NaNs
    yt_, yf_ = hpfilter_(y[i:j], λ)
    yt = [fill(NaN, i-1); yt_; fill(NaN, length(y)-j)]
    yf = [fill(NaN, i-1); yf_; fill(NaN, length(y)-j)]

    return yt, yf
end

function hpfilter_{T<:Real}(y::Vector{T}, λ::Real)
    n = length(y)
    a = spzeros(n,n)
    for i = 3:n-2
        a[i,i]   = 6λ+1
        a[i,i+1] = -4λ
        a[i,i+2] = λ
        a[i,i-2] = λ
        a[i,i-1] = -4λ
    end

    a[2,2] = 1+5λ
    a[2,3] = -4λ
    a[2,4] = λ
    a[2,1] = -2λ
    a[1,1] = 1+λ
    a[1,2] = -2λ
    a[1,3] = λ

    a[n-1,n-1] = 1+5λ
    a[n-1,n-2] = -4λ
    a[n-1,n-3] = λ
    a[n-1,n]   = -2λ
    a[n,n]     = 1+λ
    a[n,n-1]   = -2λ
    a[n,n-2]   = λ

    yt = a\y
    yf = y-yt

    return yt, yf
end

"""
```
difflog(x::Vector{AbstractFloat})
```
"""
function difflog{T<:AbstractFloat}(x::Vector{T})
    [NaN; log(x[2:end]) - log(x[1:end-1])]
end

"""
```
difflog(x::DataArray{AbstractFloat})
```
"""
function difflog(x::DataArray)
    DSGE.na2nan!(x)
    y = convert(Vector{Float64}, x)
    return difflog(y)
end

"""
```
oneqtrpctchange(y)
```

Calculates the quarter-to-quarter percentage change of a series.
"""
function oneqtrpctchange(y)
    100 * difflog(y)
end

"""
```
hpadjust(y, df)
```

Adjust series to compensate for differences between filtered and unfiltered population.
## Arguments
- `y`: A vector of data
- `df`: DataFrame containing both a filtered and unfiltered population growth series
"""
function hpadjust(y, df; filtered_mnemonic=:filtered_population_growth,
                         unfiltered_mnemonic=:unfiltered_population_growth)
    y + 100 * (df[unfiltered_mnemonic] - df[filtered_mnemonic])
end


using SparseArrays
using LinearAlgebra
using Krylov
const newaxis = [CartesianIndex()]

include("./utils.jl")
include("./utils_dec.jl")

@dec _time function aq_update_codebook(data::AbstractArray{<:AbstractFloat,2}, aq_codes_B::AbstractArray{<:Integer,2},K::Integer)
    """
    data size = (n, D)
    aq_codes_B size= (n, M)

    return
        aq_codebook size = (MK, D)
    """
    (n,M) = size(aq_codes_B)
    (n,D) = size(data)

    B = zeros(Float32, n, M*K)
    for i in 1:n
        for j in 1:M
            B[i,aq_codes_B[i,j]] = 1
        end
    end

    B = sparse(B)
    # M1 = Matrix(transpose(B)*B) + I(M*K)
    # aq_codebook = inv(M1)*transpose(B)*data

    atol=1e-6
    btol=1e-6
    atol = convert(Float32, atol)
    btol = convert(Float32, btol)

    λ = 1.0e-6
    λ = convert(Float32, λ)
    λ > 0.0 && (N = I / λ)
    data = Float32.(data)

    aq_codebook = zeros(M*K,D)
    for dim in 1:D
        (aq_codebook[:,dim],~) = lsmr(B, data[:,dim], λ=λ, sqd=λ > 0, atol=atol, btol=btol, N=N)
    end
    # aq_codebook = data \ B
    return aq_codebook
end

function aq_single_encode_CoordinatesDown(
                point::AbstractArray{<:AbstractFloat,1},
                aq_codebook::AbstractArray{<:AbstractFloat, 3},
                M::Integer,
                iter=3
                )
    # point size = (D)
    # aq_codebook size= (M, K, D)
    (M,K,D) = size(aq_codebook)
    code = zeros(Int, M)
    code .+= 1

    compress_vec = zeros(D)
    for i in 1:M
        compress_vec += aq_codebook[i,code[i],:]
    end

    for t in 1:iter
        for i in 1:M
            compress_vec_deli = compress_vec - aq_codebook[i,code[i],:]
            r_i = point - compress_vec_deli

            r_vec = broadcast(-, r_i[newaxis,:], aq_codebook[i,:,:]) # (K,D)
            error = sum(r_vec.^2,dims=2)
            code[i] = argmin(error)[1]
            compress_vec = compress_vec_deli + aq_codebook[i,code[i],:]
        end
    end
    r = point - compress_vec
    error = sum(r.^2)
    return code, error
end

@dec _time function aq_encode(
    data::AbstractArray{<:AbstractFloat,2},
    aq_codebook::AbstractArray{<:AbstractFloat, 3},
    M::Integer,
)
    (n, D) = size(data)
    aq_codes = zeros(Int, n, M)
    loss = 0
    Threads.@threads for i in 1:n
        point = data[i,:]
        (code, error) = aq_single_encode_CoordinatesDown(point,aq_codebook,M)
        loss += error
        aq_codes[i,:] = code
    end
    # loss += sum(aq_codebook.^2)
    return aq_codes,loss
end

function decode(aq_code::AbstractArray{<:Integer,1},aq_codebook::AbstractArray{<:AbstractFloat, 3})
    (M,K,D) = size(aq_codebook)
    compress_vec = zeros(D)
    for i in 1:M
        compress_vec += aq_codebook[i, aq_code[i], :]
    end
    return compress_vec
end

function compute_loss(aq_codes, aq_codebook, data)
    (M,K,D) = size(aq_codebook)
    (n,D) = size(data)

    loss = 0
    for i in 1:n
        compress_vec = decode(aq_codes[i,:], aq_codebook)
        loss += sum((data[i,:] - compress_vec).^2)
    end
    return loss
end

function aq_codebook_init(M::Integer, K::Integer, data::AbstractArray{<:AbstractFloat,2})
    (n,D) = size(data)
    aq_codebook = zeros(M, K, D)
    randix = rand(1:n, M, K)
    for i in 1:M
        for j in 1:K
            aq_codebook[i,j,:] = data[randix[i,j],:]
        end
    end
    return aq_codebook
end

function aq_train(
    M::Integer,
    K::Integer,
    data::AbstractArray{<:AbstractFloat,2},
    maxiter=20)

    aq_codebook = aq_codebook_init(M,K,data) # ndim=3
    aq_codes,loss = aq_encode(data,aq_codebook,M)
    println("loss = ",loss)
    loss_change = 1

    i = 0
    while (i<=maxiter) & (abs(loss_change) > 0.01)
        i += 1
        println("iter_num:",i)
        aq_codes_B = aq_codesTo_codesB(aq_codes, K)
        aq_codebook = aq_update_codebook(data, aq_codes_B, K) # ndim = 2
        aq_codebook = codebook_trans2To3(aq_codebook, M)

        old_loss = loss

        aq_codes, loss = aq_encode(data, aq_codebook, M)

        new_loss = loss
        loss_change = round((old_loss - new_loss )/old_loss, digits = 2)
        println("newloss = ",new_loss, ",loss_change = ",loss_change)
    end

    aq_codebook = codebook_trans3To2(aq_codebook)
    return aq_codebook,aq_codes

end


# OptimalTransport.jl -- optimal transportation algorithms for Julia
# See prettyprinted documentation at https://juliaoptimaltransport.github.io/OptimalTransport.jl/dev/
#

module OptimalTransport

using Distances
using LinearAlgebra
using IterativeSolvers, SparseArrays
using LogExpFunctions: LogExpFunctions
using MathOptInterface
using Distributions
using PDMats
using QuadGK
using NNlib: NNlib
using StatsBase: StatsBase

export SinkhornGibbs, SinkhornStabilized, SinkhornEpsilonScaling
export SinkhornBarycenterGibbs
export QuadraticOTNewton

export sinkhorn, sinkhorn2
export emd, emd2
export sinkhorn_stabilized, sinkhorn_stabilized_epsscaling, sinkhorn_barycenter
export sinkhorn_unbalanced, sinkhorn_unbalanced2
export quadreg
export ot_cost, ot_plan, wasserstein, squared2wasserstein

const MOI = MathOptInterface

include("distances/bures.jl")
include("utils.jl")
include("exact.jl")
include("wasserstein.jl")

include("entropic/sinkhorn.jl")
include("entropic/sinkhorn_gibbs.jl")
include("entropic/sinkhorn_stabilized.jl")
include("entropic/sinkhorn_epsscaling.jl")
include("entropic/sinkhorn_unbalanced.jl")
include("entropic/sinkhorn_barycenter.jl")
include("entropic/sinkhorn_barycenter_gibbs.jl")
include("entropic/sinkhorn_solve.jl")

include("quadratic.jl")
include("quadratic_newton.jl")

end


## main file for Markov chain related structs and functions

"Check if all the probabilities are greater equal than 0."
function check_prob(mat_prob::Union{Array{Float64,2}, Array{Int64,1}})
    if all(mat_prob .<= 0)
        throw("All probabilities must be greater equal than 0.")
    end
end

"Check if all the summations of rows in stpm are 1."
function check_prob_sum(stpm::Array{Float64,2})
    n_row = size(stpm)[1]
    if all(stpm * ones(n_row, 1) .== 1)
        throw("The sum of probabilities in $i-th row doesn't equal to 1.")
    end
end

"Calculate limiting probability distribution of a Markov chain family."
function cal_vec_lpd(stpm::Array{Float64,2})

    hcat([1; 1; 1])
    return vec_lpd
end

"Define the non-mutable struct for Markov chain family."
struct MarkovChainFamily
    stpm::Array{Float64,2}  # [Stationary Transition Probability Matrix]
    vec_lpd::Array{Float64,2}  # [Limiting Probability Distribution]

    function MarkovChainFamily(stpm::Array{Float64,2})
        check_prob(stpm)
        check_prob_sum(stpm)

        vec_lpd = cal_vec_lpd(stpm)

        new(stpm, vec_lpd)
    end
end


immutable SharedData{T<:FloatingPoint}
	Lambda::Diagonal{T}
	sigma2::T
	diagZ::Array{T, 1}
	Z::Union(Base.LinAlg.Cholesky{T}, Array{T,2})
end
function get_shared(m, k, lik)
	Lambda = Diagonal(Float64[S(k, sqrt(lam)) for lam in m.lambdas])
	sigma2 = exp(2lik.log_sigma)
	diagZ = sigma2 ./ diag(Lambda) + diag(m.PhiTPhi)
	Z = sigma2 * inv(Lambda) + m.PhiTPhi
	try
		Z = cholfact(Z)
	catch
	end
	return SharedData(Lambda, sigma2, diagZ, Z)
end

function nll(m::LaplacianModel{OneD})
	s = get_shared(m, m.k, m.lik)
	nll(m, m.k, m.lik, s)
end

function nll(m, k, lik, s)
	n = size(m.X, 1)
	yTQy(m, k, s) / 2 + logdetQ(m, k, s) / 2 + n*log(2pi)
end

function dnlldl(m, k, lik)
	s = get_shared(m, k, lik)
	dnlldl(m, k, lik, s)
end
function dnlldnu(m, k, lik)
	s = get_shared(m, k, lik)
	dnlldnu(m, k, lik, s)
end

function dnlldsigma2(m, k, lik)
	s = get_shared(m, k, lik)
	dnlldsigma2(m, k, lik, s)
end

function dnlldl(m, k, lik, s)
	dyTQydl(m, k, s) / 2 + dlogdetQdl(m, k, s) / 2
end
function dnlldnu(m, k, lik, s)
	dyTQydnu(m, k, s) / 2 + dlogdetQdnu(m, k, s) / 2
end

function dnlldsigma2(m, k, lik, s)
	dyTQydsigma2(m, k, s) / 2 + dlogdetQdsigma2(m, k, s) / 2
end

function yTQy(m, k, s)
	(dot(m.y,m.y) - dot(m.PhiTy, s.Z \ m.PhiTy)) / s.sigma2
end

function dyTQydl(m, k, s)
	dLambdadl = Diagonal([dSdl(k, sqrt(lam)) for lam in m.lambdas])
	- dot(m.PhiTy, s.Z \ (inv(s.Lambda)^2 * dLambdadl *  (s.Z \ m.PhiTy)))
end

function dyTQydnu(m, k, s)
	dLambdadnu = Diagonal([dSdnu(k, sqrt(lam)) for lam in m.lambdas])
	- dot(m.PhiTy, s.Z \ (inv(s.Lambda)^2 * dLambdadnu *  (s.Z \ m.PhiTy)))
end

function dyTQydsigma2(m, k, s)
	2s.sigma2 * (dot(m.PhiTy, s.Z \ (inv(s.Lambda) * (s.Z \ m.PhiTy))) - yTQy(m, k, s)) / s.sigma2
end

function logdetQ(m, k, s)
	logdetZ = logdet(s.Z)
	n = size(m.X, 1)
	(n - m.jmax)*log(s.sigma2) + logdetZ + sum(log(diag(s.Lambda)))
end

function dlogdetQdl(m, k, s)
	dLambdadl = Diagonal([dSdl(k, sqrt(l)) for l in m.lambdas])
	trace(dLambdadl / s.Lambda) - s.sigma2 * sum(1 ./ (s.diagZ  .* inv(inv(s.Lambda)^2 * dLambdadl).diag))
end

function dlogdetQdnu(m, k, s)
	dLambdadnu = Diagonal([dSdnu(k, sqrt(l)) for l in m.lambdas])
	trace(dLambdadnu / s.Lambda) - s.sigma2 * sum(1 ./ (s.diagZ  .* inv(inv(s.Lambda)^2 * dLambdadnu).diag))
end

function dlogdetQdsigma2(m, k, s)
	n = size(m.X, 1)
	((n - m.jmax) / (s.sigma2) + sum(1./ (s.diagZ .* diag(s.Lambda)))) * 2s.sigma2
end


export Register

"""
Creates register initializing with the |0...0⟩ state

- `N`     : Total number of qubits in circuit
- `state` : State vector for the 2^N Hilbert space

"""
mutable struct Register
    N::Int
    state::Array{ComplexF64}

    # Construct with state |0...0⟩
    function Register(N::Int)
        x = zeros(ComplexF64,2^N)
        x[1] = complex(1.0)
        new(N,x)
    end
end


#############################################################################
# ECOS.jl
# Wrapper around the ECOS solver https://github.com/ifa-ethz/ecos
# See http://github.com/JuliaOpt/ECOS.jl
#############################################################################
# test/mpb.jl
# Test the MathProgBase.jl interface for the ECOS.jl solver wrapper
#############################################################################

using Base.Test
importall MathProgBase
using ECOS

s = ECOSSolver()

# Problem 1 - all vars in nonneg cone
# min -3x - 2y - 4z
# st    x +  y +  z == 3
#            y +  z == 2
#       x>=0 y>=0 z>=0
# Opt solution = -11
# x = 1, y = 0, z = 2
m = MathProgBase.model(s)
MathProgBase.loadconicproblem!(m,
                    [-3.0, -2.0, -4.0],
                    [ 1.0   1.0   1.0;
                      0.0   1.0   1.0],
                    [ 3.0,  2.0],
                    [(:NonNeg, 1:3)])
MathProgBase.optimize!(m)
@test MathProgBase.status(m) == :Optimal
@test_approx_eq_eps MathProgBase.getobjval(m) -11 1e-6
@test_approx_eq_eps MathProgBase.getsolution(m)[1] 1.0 1e-6
@test_approx_eq_eps MathProgBase.getsolution(m)[2] 0.0 1e-6
@test_approx_eq_eps MathProgBase.getsolution(m)[3] 2.0 1e-6

# Problem 2 - mixed free, nonneg, nonpos, zero, shuffled cones
# min  3x + 2y - 4z + 0s
# st    x           -  s  == -4    (i.e. x >= -4)
#            y            == -3
#       x      +  z       == 12
#       x free
#       y <= 0
#       z >= 0
#       s zero
# Opt solution = -82
# x = -4, y = -3, z = 16, s == 0
m = MathProgBase.model(s)
MathProgBase.loadconicproblem!(m,
                    [ 3.0,  2.0, -4.0,  0.0],
                    [ 1.0   0.0   0.0  -1.0;
                      0.0   1.0   0.0   0.0;
                      1.0   0.0   1.0   0.0],
                    [-4.0, -3.0, 12.0],
                    [(:Free,1), (:NonNeg,3), (:Zero,4), (:NonPos,2)])
MathProgBase.optimize!(m)
@test MathProgBase.status(m) == :Optimal
@test_approx_eq_eps MathProgBase.getobjval(m) -82 1e-6
@test_approx_eq_eps MathProgBase.getsolution(m)[1] -4.0 1e-6
@test_approx_eq_eps MathProgBase.getsolution(m)[2] -3.0 1e-6
@test_approx_eq_eps MathProgBase.getsolution(m)[3] 16.0 1e-6
@test_approx_eq_eps MathProgBase.getsolution(m)[4]  0.0 1e-6

# Problem 3 - SOC
# min 0x - 1y - 1z
#  st  x            == 1
#      x >= ||(y,z)||
m = MathProgBase.model(s)
MathProgBase.loadconicproblem!(m,
                    [ 0.0, -1.0, -1.0],
                    [ 1.0   0.0   0.0],
                    [ 1.0],
                    [(:SOC,1:3)])
MathProgBase.optimize!(m)
@test MathProgBase.status(m) == :Optimal
@test_approx_eq_eps MathProgBase.getobjval(m) -sqrt(2.0) 1e-6
@test_approx_eq_eps MathProgBase.getsolution(m)[1] 1.0 1e-6
@test_approx_eq_eps MathProgBase.getsolution(m)[2] 1.0/sqrt(2.0) 1e-6
@test_approx_eq_eps MathProgBase.getsolution(m)[3] 1.0/sqrt(2.0) 1e-6


@doc """

    ar2(y, x, p, βbar, Bbarinv, ϕbar, Vbarinv, υbar, δbar, βold, ϕold, σ2old, varind, lev2ind, factors, numobs, varassign)

Description:
Estimate observed variable hyperparameters conditional on the factors.
Refer to pg. 1002 of Otrok and Whiteman (1998), or Chib and Greenberg (1994).

Inputs:
- y             = observed series vector.
- x             = data matrix of regressors in the observation equation (intercept + factors).
- p             = number of idiosyncratic error lags.
- βbar          = prior mean of factor loading (factor coefficient in measurement equation).
- Bbarinv       = prior variance of factor loading (factor coefficient in measurement equation).
- ϕbar          = prior mean of phi (idiosyncratic AR polynomial).
- Vbarinv       = prior precision of phi (idiosyncratic AR polynomial).
- υbar          = inverted gamma parameter of innovation variances.
- δbar          = inverted gamma parameter of innovation variances.
- βold          = old draw of factor loading.
- ϕold          = old draw of phi (idiosyncratic AR polynomial).
- σ2old         = old draw of innovation variances.
- varind        = measurement eq. regressor index.
- lev2ind       = level-2 factor index.
- factors       = factor draw.
- numobs        = total number of time periods / observations in the sample.
- varassign     = list containing vector of observed series indeces assigned to each level-2 factor.

Outputs:
- βnew      = β coefficient hyperparameter estimates.
- ϕnew      = ϕ coefficient hyperparameter estimates.
- σ2new     = σ^2 coefficient hyperparameter estimate.
- factors   = factor estimates.
"""

function ar2(y, x, p, βbar, Bbarinv, ϕbar, Vbarinv, υbar, δbar, βold, ϕold, σ2old, varind, lev2ind, factors, numobs, varassign)

    # Guarantee that y is a vector
    y = vec(y)

    # Define ỹ = first p observations of series vector
    ỹ = y[1:p]

    # Define x̃ = first p observations of regressor matrix
    x̃ = x[1:p, :]

    # Create Φ matrix (companion matrix of idiosyncratic error autoregression)
    Φ = [ϕold'; I(p - 1) zeros(p - 1)]

    # Covariance matrix component of the first p errors (Σ in σ2Σ)
    vecΣ = inv(I - Φ ⊗ Φ) * vec([1; zeros(p - 1)] * [1; zeros(p - 1)]')
    Σ = reshape(vecΣ, p, p)
    Σ = Hermitian(Σ)

    # Compute the Cholesky factor of Σ
    # notation is such that Q Q' = Σ
    # so that Q is the lower factor of Σ
    Q = cholesky(Σ).L

    # Define ỹ∗1 = Q^{-1} ỹ and x̃∗1 = Q^{-1} x̃
    ỹstar1 = inv(Q) * ỹ
    x̃star1 = inv(Q) * x̃

    # Define ỹ∗2 as a (T-p × 1) with t-th row ϕ(L)y
    # and x̃∗2 as a (T-p × 3) with t-th row ϕ(1)
    ỹstar2 = gendiff(y, ϕold)
    x̃star2 = [gendiff(ones(numobs), ϕold) gendiff(factors[:, 1], ϕold) gendiff(factors[:, lev2ind], ϕold)]

    # Define e = last T-p "observed" idiosyncratic errors
    e = y - x * βold
    e = e[(p+1):numobs]

    # Define E = [ lag(e,1) lag(e,2) ... lag(e,p) ]
    E = zeros(numobs, p)
    for j = 1:p
        E[:, j] = lag(y - x * βold, j, default = 0.0)
    end
    E = E[(p+1):numobs, :]

    # Define x̃∗ = [x̃∗1 ; x̃∗2] and ỹ∗ = [ỹ∗1 ; ỹ∗2]
    x̃star = [x̃star1; x̃star2]
    ỹstar = [ỹstar1; ỹstar2]

    #############################################
    #############################################
    ## Create objects for posterior densities

    # Define V̅ and B̅
    Vbar = inv(Vbarinv)
    Bbar = inv(Bbarinv)

    # Define V = V̅ + σ²E'E
    V = Vbar + σ2old * E' * E

    # Define B = B̅ + ̅σ²x̃∗'x̃∗
    B = Bbar + σ2old * x̃star' * x̃star

    # Define ̂ϕ = inv(V)(V̅ ̅ϕ + σ²E'e)
    ϕhat = inv(V) * (Vbar * ϕbar + σ2old * E' * e)

    #############################################
    #############################################
    ## Draw β

    # Initialize sign of level-1 factor loading
    # of first series in dataset (must be positive)
    β2sign = false

    # Initialize sign of level-2 factor loading
    # of first series in corresponding level-2 group
    # (must be positive)
    β3sign = false

    # Initialize β
    βnew = βold

    # Keep track of wrong sign persistence
    β2signmax = 0
    β3signmax = 0

    whileind = 0
    while β2sign == false || β3sign == false

        whileind += 1
        #=println("Series index: $varind")
        println("While loop index: $whileind")
        println("β2sigmax: $β2signmax")
        println("β3signmax: $β3signmax")
        =#

        # Draw new β
        #βnew = sim_MvNormal(inv(B) * (Bbar * βbar + inv(σ2old) * x̃star' * ỹstar), inv(B))
        β0 = zeros(size(x)[2])
        Σ0 = Matrix(I, size(β0)[1], size(β0)[1]) .* 1000.0
        ## Posterior parameters in N(β1,Σ1)
        β1 = inv(inv(Σ0) + inv(σ2old) * transpose(x) * x) * (inv(Σ0) * β0 + inv(σ2old) * transpose(x) * y)
        β1 = vec(β1)
        Σ1 = inv(inv(Σ0) + inv(σ2old) * transpose(x) * x)
        Σ1 = Hermitian(Σ1)

        ## Generate new β
        βnew = sim_MvNormal(β1, Σ1)

        # Check if new β draw satisfies
        # sign identification restrictions
        if varind == 1
            β2sign = (βnew[2] > 0)
            if β2sign == false
                β2signmax += 1
            end
        else
            β2sign = true
        end
        if varind == varassign[lev2ind][1]
            β3sign = (βnew[3] > 0)
            if β3sign == false
                β3signmax += 1
            end
        else
            β3sign = true
        end

        if β2signmax > 100

            # Reflect level-1 factor over x-axis
            factors[:, 1] = -factors[:, 1]

            # Redefine x
            x = [ones(numobs) factors[:, 1] factors[:, 1+lev2ind]]

            # Define x̃ = first p observations of regressor matrix
            x̃ = x[1:p, :]

            # Create Φ matrix (companion matrix of idiosyncratic error autoregression)
            Φ = [ϕold'; I(p - 1) zeros(p - 1)]

            # Covariance matrix component of the first p errors (Σ in σ2Σ)
            vecΣ = inv(I - Φ ⊗ Φ) * vec([1; zeros(p - 1)] * [1; zeros(p - 1)]')
            Σ = reshape(vecΣ, p, p)
            Σ = Hermitian(Σ)

            # Compute the Cholesky factor of Σ
            # notation is such that Q Q' = Σ
            # so that Q is the lower factor of Σ
            Q = cholesky(Σ).L

            # Define ỹ∗1 = Q^{-1} ỹ and x̃∗1 = Q^{-1} x̃
            ỹstar1 = inv(Q) * ỹ
            x̃star1 = inv(Q) * x̃

            # Define ỹ∗2 as a (T-p × 1) with t-th row ϕ(L)y
            # and x̃∗2 as a (T-p × 3) with t-th row ϕ(1)
            ỹstar2 = gendiff(y, ϕold)
            x̃star2 = [gendiff(ones(numobs), ϕold) gendiff(factors[:, 1], ϕold) gendiff(factors[:, lev2ind], ϕold)]

            # Define e = last T-p "observed" idiosyncratic errors
            e = y - x * βold
            e = e[(p+1):numobs]

            # Define E = [ lag(e,1) lag(e,2) ... lag(e,p) ]
            E = zeros(numobs, p)
            for j = 1:p
                E[:, j] = lag(y - x * βold, j, default = 0.0)
            end
            E = E[(p+1):numobs, :]

            # Define x̃∗ = [x̃∗1 ; x̃∗2] and ỹ∗ = [ỹ∗1 ; ỹ∗2]
            x̃star = [x̃star1; x̃star2]
            ỹstar = [ỹstar1; ỹstar2]

            # Reset `β2signmax`
            β2signmax = 0
        end

        if β3signmax > 1000

            # Reflect level-1 factor over x-axis
            factors[:, 1+lev2ind] = -factors[:, 1+lev2ind]

            # Redefine x
            x = [ones(numobs) factors[:, 1] factors[:, 1+lev2ind]]

            # Define x̃ = first p observations of regressor matrix
            x̃ = x[1:p, :]

            # Create Φ matrix (companion matrix of idiosyncratic error autoregression)
            Φ = [ϕold'; I(p - 1) zeros(p - 1)]

            # Covariance matrix component of the first p errors (Σ in σ2Σ)
            vecΣ = inv(I - Φ ⊗ Φ) * vec([1; zeros(p - 1)] * [1; zeros(p - 1)]')
            Σ = reshape(vecΣ, p, p)
            Σ = Hermitian(Σ)

            # Compute the Cholesky factor of Σ
            # notation is such that Q Q' = Σ
            # so that Q is the lower factor of Σ
            Q = cholesky(Σ).L

            # Define ỹ∗1 = Q^{-1} ỹ and x̃∗1 = Q^{-1} x̃
            ỹstar1 = inv(Q) * ỹ
            x̃star1 = inv(Q) * x̃

            # Define ỹ∗2 as a (T-p × 1) with t-th row ϕ(L)y
            # and x̃∗2 as a (T-p × 3) with t-th row ϕ(1)
            ỹstar2 = gendiff(y, ϕold)
            x̃star2 = [gendiff(ones(numobs), ϕold) gendiff(factors[:, 1], ϕold) gendiff(factors[:, lev2ind], ϕold)]

            # Define e = last T-p "observed" idiosyncratic errors
            e = y - x * βold
            e = e[(p+1):numobs]

            # Define E = [ lag(e,1) lag(e,2) ... lag(e,p) ]
            E = zeros(numobs, p)
            for j = 1:p
                E[:, j] = lag(y - x * βold, j, default = 0.0)
            end
            E = E[(p+1):numobs, :]

            # Define x̃∗ = [x̃∗1 ; x̃∗2] and ỹ∗ = [ỹ∗1 ; ỹ∗2]
            x̃star = [x̃star1; x̃star2]
            ỹstar = [ỹstar1; ỹstar2]

            # Reset `β2signmax`
            β3signmax = 0
        end

    end

    #############################################
    #############################################
    ## Draw ϕ

    #=
    # Draw new ϕ candidate
    ϕnew = sim_MvNormal(ϕhat, inv(V))

    # Check for stationary of new ϕ draw
    coef = [-reverse(ϕnew); 1]
    coef = vec(coef)
    root = roots(Polynomial(reverse(coef)))
    rootmod = abs.(root)
    accept = min(rootmod...) >= 1.001

    # Complete
    if accept == 0
        ϕnew = ϕold
    else
        # Define Ψ(ϕold) (pg. 1002)
        Ψold = det(sigmat(ϕold, p))^(-1 / 2) * exp(-(1 / (2 * σ2old)) * (ỹ - x̃ * βnew)' * inv(sigmat(ϕold, p)) * (ỹ - x̃ * βnew))
        Ψold = first(Ψold)

        # Define Ψ(ϕnew) (pg. 1002)
        Ψnew = det(sigmat(ϕnew, p))^(-1 / 2) * exp(-(1 / (2 * σ2old)) * (ỹ - x̃ * βnew)' * inv(sigmat(ϕnew, p)) * (ỹ - x̃ * βnew))
        Ψnew = first(Ψnew)

        # Determine acceptance probability
        if Ψold == 0
            accept = 1
        else
            u = rand()
            accept = (u <= Ψnew / Ψold)
        end

        # Pick new ϕ
        ϕnew = ϕnew * accept + ϕold * (1 - accept)
    end
    =#

    # Generate ϕ^j
    ## Prior parameters in N(c0,B0)
    c0 = zeros(size(ϕold)[1])
    B0 = Matrix(I, size(c0)[1], size(c0)[1]) .* 1000.0
    ## Generate e^⋆
    e_star = y - x * βnew
    ## Generate E^⋆
    E_star = zeros(numobs, length(ϕold))
    E_star = E_star[(1+length(ϕold)):end, :]
    for i = 1:(size(E_star)[2]) # iterate over variables in X
        e_temp = lag(e_star, i)
        E_star[:, i] = e_temp[(1+length(ϕold)):end]
    end
    e_star = e_star[(1+length(ϕold)):end]
    ## Posterior parameters in N(c1,B1)
    c1 = inv(inv(B0) + inv(σ2old) * transpose(E_star) * E_star) * (inv(B0) * c0 + inv(σ2old) * transpose(E_star) * e_star)
    B1 = inv(inv(B0) + inv(σ2old) * transpose(E_star) * E_star)
    B1 = Hermitian(B1)
    ## Generate new ϕ
    ϕnew = sim_MvNormal(c1, B1)

    #############################################
    #############################################
    ## Draw σ2

    # Draw new σ2 from inverse gamma distribution
    #σ2new = rand(InverseGamma((υbar + numobs) / 2, (δbar + norm(ỹstar - x̃star * βnew)^2) / 2))
    ν0 = 0.002
    δ0 = 0.002
    ## Posterior parameters in IG(ν1/2, δ1/2)
    ν1 = ν0 + numobs
    δ1 = δ0 + norm(y - x * βnew)^2
    ## Generate new σ2
    σ2new = rand(InverseGamma(ν1 / 2, δ1 / 2))

    #############################################
    #############################################
    ## Return new hyperparameter draws
    return βnew, σ2new, ϕnew, factors
end


if !@isdefined TEST
    using Revise
end
using Mitosis
using Random, Test, LinearAlgebra, Statistics
using Soss
using StatsBase
DIAG = true
if DIAG
    using MCMCChains
end

N = 50_000 # samples

# Define some vectors and matrices
ξ0 = [1., 0.]
x = [1.2, 0.1]
P0 = Matrix(1.0I, 2, 2)

Φ = [0.8 0.5; -0.1 0.8]
β = [0.1, 0.2]
Q = [0.2 0.0; 0.0 1.0]

H = [1.0 0.0]
R = Matrix(1.0I, 1, 1)

# and a nonlinear vector function
f(x) = [0.8 0.5; -0.1 0.8]*[atan(x[1]), atan(x[2])] + [0.1, 0.2]

# We define the equivalent of the Soss model
(m = @model ξ0, H, f, P0, R, Q begin
    a ~ MvNormal(ξ0, P0) # priortransition
    b ~ MvNormal(H*a, R) # partial observation
    c ~ MvNormal(f(a), Q) # nonlineartransition
    d ~ MvNormal(H*c, R) # partial observation
    e ~ MvNormal(c, P0) # full observation
    # return b, d, e
end)

# Define some transition kernels.

# We use AffineMap, LinearMap and ConstantMap
# For example
@test AffineMap(Φ, β)(x) == Φ*x + β

# Prior
# a ~ N(ξ0, P0)
k_0a = Mitosis.kernel(Gaussian; μ=LinearMap(I(2)), Σ=ConstantMap(P0))

# Partial observation
# b ~ N(H*a, R)
k_ab = Mitosis.kernel(Gaussian; μ=LinearMap(H), Σ=ConstantMap(R))

# Nonlinear transition and linear approximation
# c ~ N(f(a), Q)
k_ac = Mitosis.kernel(Gaussian; μ=f, Σ=ConstantMap(Q))
k_ac_approx = Mitosis.kernel(Gaussian; μ=AffineMap(Φ, β), Σ=ConstantMap(Q))

# d ~ N(H*c, R)
k_cd = Mitosis.kernel(Gaussian; μ=LinearMap(H), Σ=ConstantMap(R))

# e ~ N(c, P0)
k_ce = Mitosis.kernel(Gaussian; μ=LinearMap(I(2)), Σ=ConstantMap(P0))

# And a deterministic copy kernel
cp2 = Mitosis.Copy{2}() # copy kernel
@test cp2(1.1) == (1.1, 1.1)

# Forward sample a Bayes net # rand(m(ξ0=ξ0))
a = rand(k_0a(ξ0))
b = rand(k_ab(a))
c = rand(k_ac(a)) # forward model with nonlinear transition
d = rand(k_cd(c))
e = rand(k_ce(c))

# We actually want to write down the
# forward sample with explicit copies
# so every state is only used as input of a single kernel
a = rand(k_0a(ξ0))
a1, a2 = cp2(a)
b = rand(k_ab(a1))
c = rand(k_ac(a2)) # forward model with nonlinear transition
c1, c2 = cp2(c)
d = rand(k_cd(c1))
e = rand(k_ce(c2))

# sample posterior with dynamicHMC via Soss for comparison
# monkey patch for xform
Soss.xform(d::MvNormal, _data) = Soss.as(Array, Soss.asℝ, size(d))
@time posterior = dynamicHMC(m(ξ0=ξ0, H=H, f=f, P0=P0, R=R, Q=Q), (b=b, d=d, e=e,), N)

# observations are b d e (all leaves)
mc2, pc2 = backward(BFFG(), k_ce, e; unfused=true) # child of copy
mc1, pc1 = backward(BFFG(), k_cd, d; unfused=true) # child of copy
mc, pc = backward(BFFG(), cp2, pc1, pc2) # copy
ma2, pa2 = backward(BFFG(), k_ac_approx, pc; unfused=true) # # child of copy
ma1, pa1 = backward(BFFG(), k_ab, b; unfused=true) # child of copy
ma, pa = backward(BFFG(), cp2, pa1, pa2) # copy
m0, p0 = backward(BFFG(), k_0a, pa) # not a child of a copy

samples = []
lls = []
@time for i in 1:N
    a = rand(forward(BFFG(), k_0a, m0, weighted(ξ0)))
    a1, a2 = rand(forward(BFFG(), cp2, ma, a))
    b = rand(forward(BFFG(), k_ab, ma1, a1))
    c = rand(forward(BFFG(), (k_ac, k_ac_approx), ma2, a2)) # forward model with nonlinear transition
    c1, c2 = rand(forward(BFFG(), cp2, mc, c))
    d = rand(forward(BFFG(), k_cd, mc1, c1))
    e = rand(forward(BFFG(), k_ce, mc2, c2))

    # Sample and weight (sum over the weights of the leafs.)
    push!(samples, (a=a.x,b=b.x,c=c.x,d=d.x,e=e.x))
    push!(lls, b.ll + d.ll + e.ll)

end

# Importance weights
w = exp.(lls)/mean(exp.(lls))

# Estimate Mitosis
â1 = mean(getfield.(samples, :a).*w)
v1 = mean(Mitosis.outer.(getfield.(samples, :a).-Ref(â1)).*w)

# Estimate Soss
â2 = mean(getfield.(posterior, :a))
v2 = cov(getfield.(posterior, :a))

if DIAG
    # Kish's Effective Sample Size
    ess1 = sum(w).^2/sum(w.^2)
    ess2 = MCMCChains.ess(Chains(reshape(first.(getfield.(posterior, :a)), :,1,1)))
end

[â1 â2]


function sparklineit(arr::Vector{<:Integer})
    sparkchars = '\u2581':'\u2588'
    dyn = length(sparkchars)
    lo, hi = extrema(arr)
    b = @. max(ceil(Int, dyn * (arr - lo) / (hi - lo)), 1)
    return join(sparkchars[b])
end

v = rand(0:10, 10)
println("$v → ", sparklineit(v))
v = 10rand(10)
println("$(round.(v, 2)) → ", sparklineit(v))


"""
    struct VariableSparsity <: Sparsity end

Variable or correlative sparsity as developed in [WSMM06].

[WSMM06] Waki, Hayato, Sunyoung Kim, Masakazu Kojima, and Masakazu Muramatsu. "Sums of squares and semidefinite program relaxations for polynomial optimization problems with structured sparsity." SIAM Journal on Optimization 17, no. 1 (2006): 218-242.
"""
struct VariableSparsity <: Sparsity end

const CEG = ChordalExtensionGraph

function csp_graph(poly::MP.APL, ::FullSpace)
    G = CEG.LabelledGraph{MP.variable_union_type(poly)}()
    for mono in MP.monomials(poly)
        CEG.add_clique!(G, MP.effective_variables(mono))
    end
    return G
end

function csp_graph(poly::MP.APL, domain::AbstractAlgebraicSet)
    G = csp_graph(poly, FullSpace())
    for p in equalities(domain)
        CEG.add_clique!(G, MP.effective_variables(p))
    end
    return G
end

function csp_graph(poly::MP.APL, domain::BasicSemialgebraicSet)
    G = csp_graph(poly, domain.V)
    for p in inequalities(domain)
        CEG.add_clique!(G, MP.effective_variables(p))
    end
    return G
end

function chordal_csp_graph(poly::MP.APL, domain::AbstractBasicSemialgebraicSet)
    H, cliques = CEG.chordal_extension(csp_graph(poly, domain), CEG.GreedyFillIn())
    for clique in cliques
        sort!(clique, rev=true)
        unique!(clique)
    end
    return H, cliques

end

function sparsity(poly::MP.AbstractPolynomial, domain::BasicSemialgebraicSet,
                  sp::VariableSparsity, certificate::Putinar)
    H, cliques = chordal_csp_graph(poly, domain)
    function bases(q)
        return [
            maxdegree_gram_basis(certificate.basis, clique,
                                 multiplier_maxdegree(certificate.maxdegree, q))
            for clique in cliques if variables(q) ⊆ clique
        ]
    end
    return bases(poly), map(bases, domain.p)
end


module bushing_examples
using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsDeforLinear.AlgoDeforLinearModule
using FinEtools.MeshExportModule
using Statistics: mean

E = 3.0;
nu = 0.4999999; # This is the correct Poisson ratio
Ri = 0.25;
Re = 1.0;
L = Ri/2;
ang = 180/180*pi;
iuz = Ri/15
p = 0.27;
tolerance = min(L, Re-Ri, ang)/1000;
nR, nc, nt = 7, 14, 1

function bushing_h8_full()

	fens, fes  = H8block(ang, L, Re-Ri, nc, nt, nR)
	internal_fenids = selectnode(fens, box = [0 ang 0 L 0 0], inflate = tolerance);

	MR = DeforModelRed3D
	material = MatDeforElastIso(MR, E, nu)
	femm = FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(3, 2)), material)

	boundaryfes = meshboundary(fes);
	ibcl = selectelem(fens, boundaryfes, box = [0 ang 0 L 0 0], inflate = tolerance);
	ebcl = selectelem(fens, boundaryfes, box = [0 ang 0 L Re-Ri Re-Ri], inflate = tolerance);

	for i in 1:count(fens)
		a = fens.xyz[i,1]; y = fens.xyz[i,2]; r = fens.xyz[i,3]+Ri;
		fens.xyz[i,:] .= (r*sin(a), y, r*cos(a));
	end

	geom = NodalField(fens.xyz)
	u = NodalField(zeros(size(fens.xyz,1), 3)) # displacement field

	x0nl = selectnode(fens, box = [0 0 -Inf Inf -Inf Inf], inflate = tolerance);
	setebc!(u, x0nl, true, 1, 0.0)

	y0nl = selectnode(fens, box = [-Inf Inf 0 0 -Inf Inf], inflate = tolerance);
	setebc!(u, y0nl, true, 2, 0.0)
	yLnl = selectnode(fens, box = [-Inf Inf L L -Inf Inf], inflate = tolerance);
	setebc!(u, yLnl, true, 2, 0.0)

	enl = connectednodes(subset(boundaryfes, ebcl))
	setebc!(u, enl, true, [1; 2; 3]; val = 0.0)
	inl = connectednodes(subset(boundaryfes, ibcl))
	setebc!(u, inl, true, [1; 2]; val = 0.0)
	setebc!(u, inl, true, [3]; val = iuz)

	applyebc!(u)
	numberdofs!(u)

	associategeometry!(femm, geom)

	K = stiffness(femm, geom, u)
	F = nzebcloadsstiffness(femm, geom, u)

	scattersysvec!(u, K\F)

	File =  "bushing_h8_full_u.vtk"
	vtkexportmesh(File, fens, fes; vectors = [("u", u.values)])
	@async run(`"paraview.exe" $File`)

    true
end # bushing_h8_full

function bushing_h8_algo_full()

	fens, fes  = H8block(ang, L, Re-Ri, nc, nt, nR)
	internal_fenids = selectnode(fens, box = [0 ang 0 L 0 0], inflate = tolerance);

	MR = DeforModelRed3D
	material = MatDeforElastIso(MR, E, nu)
	femm = FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(3, 2)), material)
	regions = FDataDict[FDataDict("femm"=>femm)]

	boundaryfes = meshboundary(fes);
	ibcl = selectelem(fens, boundaryfes, box = [0 ang 0 L 0 0], inflate = tolerance);
	ebcl = selectelem(fens, boundaryfes, box = [0 ang 0 L Re-Ri Re-Ri], inflate = tolerance);

	for i in 1:count(fens)
		a = fens.xyz[i,1]; y = fens.xyz[i,2]; r = fens.xyz[i,3]+Ri;
		fens.xyz[i,:] .= (r*sin(a), y, r*cos(a));
	end

	essential_bcs = FDataDict[]
	x0nl = selectnode(fens, box = [0 0 -Inf Inf -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>x0nl))

	y0nl = selectnode(fens, box = [-Inf Inf 0 0 -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>y0nl))
	yLnl = selectnode(fens, box = [-Inf Inf L L -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>yLnl))

	enl = connectednodes(subset(boundaryfes, ebcl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>enl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>enl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 3, "node_list"=>enl))
	inl = connectednodes(subset(boundaryfes, ibcl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>inl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>inl))
	push!(essential_bcs, FDataDict("displacement"=> iuz, "component"=> 3, "node_list"=>inl))

	modeldata = FDataDict("fens"=>fens,
	"regions"=>regions,
	"essential_bcs"=>essential_bcs
	)
	modeldata = AlgoDeforLinearModule.linearstatics(modeldata)

	filebase = "bushing_h8_algo_full"
	modeldata["postprocessing"] = FDataDict("file"=>filebase * "-s", "quantity"=>:pressure, "component"=>collect(1:1), "nodevalmethod"=>:averaging, "reportat"=>:extraptrend)
	modeldata = AlgoDeforLinearModule.exportstress(modeldata)

    true
end # bushing_h8_algo_full

function bushing_h8_algo_ms()

	fens, fes  = H8block(ang, L, Re-Ri, nc, nt, nR)
	internal_fenids = selectnode(fens, box = [0 ang 0 L 0 0], inflate = tolerance);

	MR = DeforModelRed3D
	material = MatDeforElastIso(MR, E, nu)
	femm = FEMMDeforLinearMSH8(MR, IntegDomain(fes, GaussRule(3, 2)), material)
	regions = FDataDict[FDataDict("femm"=>femm)]

	boundaryfes = meshboundary(fes);
	ibcl = selectelem(fens, boundaryfes, box = [0 ang 0 L 0 0], inflate = tolerance);
	ebcl = selectelem(fens, boundaryfes, box = [0 ang 0 L Re-Ri Re-Ri], inflate = tolerance);

	for i in 1:count(fens)
		a = fens.xyz[i,1]; y = fens.xyz[i,2]; r = fens.xyz[i,3]+Ri;
		fens.xyz[i,:] .= (r*sin(a), y, r*cos(a));
	end

	essential_bcs = FDataDict[]
	x0nl = selectnode(fens, box = [0 0 -Inf Inf -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>x0nl))

	y0nl = selectnode(fens, box = [-Inf Inf 0 0 -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>y0nl))
	yLnl = selectnode(fens, box = [-Inf Inf L L -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>yLnl))

	enl = connectednodes(subset(boundaryfes, ebcl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>enl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>enl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 3, "node_list"=>enl))
	inl = connectednodes(subset(boundaryfes, ibcl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>inl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>inl))
	push!(essential_bcs, FDataDict("displacement"=> iuz, "component"=> 3, "node_list"=>inl))

	modeldata = FDataDict("fens"=>fens,
	"regions"=>regions,
	"essential_bcs"=>essential_bcs
	)
	modeldata = AlgoDeforLinearModule.linearstatics(modeldata)

	filebase = "bushing_h8_algo_ms"
	modeldata["postprocessing"] = FDataDict("file"=>filebase * "-nodal-", "quantity"=>:pressure, "component"=>collect(1:1), "nodevalmethod"=>:averaging, "reportat"=>:extraptrend)
	modeldata = AlgoDeforLinearModule.exportstress(modeldata)
	modeldata["postprocessing"] = FDataDict("file"=>filebase * "-elwise-", "quantity"=>:pressure, "component"=>collect(1:1), "nodevalmethod"=>:averaging, "reportat"=>:extraptrend)
	modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)

    true
end # bushing_h8_algo_ms

function bushing_h8u_algo_ms()

	fens, fes  = T4block(ang, L, Re-Ri, Int.(round.((nc, nt, nR)./2).+1)...)
	fens, fes  = T4toH8(fens, fes)
	internal_fenids = selectnode(fens, box = [0 ang 0 L 0 0], inflate = tolerance);

	MR = DeforModelRed3D
	material = MatDeforElastIso(MR, E, nu)
	femm = FEMMDeforLinearMSH8(MR, IntegDomain(fes, GaussRule(3, 2)), material)
	regions = FDataDict[FDataDict("femm"=>femm)]

	boundaryfes = meshboundary(fes);
	ibcl = selectelem(fens, boundaryfes, box = [0 ang 0 L 0 0], inflate = tolerance);
	ebcl = selectelem(fens, boundaryfes, box = [0 ang 0 L Re-Ri Re-Ri], inflate = tolerance);

	for i in 1:count(fens)
		a = fens.xyz[i,1]; y = fens.xyz[i,2]; r = fens.xyz[i,3]+Ri;
		fens.xyz[i,:] .= (r*sin(a), y, r*cos(a));
	end

	essential_bcs = FDataDict[]
	x0nl = selectnode(fens, box = [0 0 -Inf Inf -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>x0nl))

	y0nl = selectnode(fens, box = [-Inf Inf 0 0 -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>y0nl))
	yLnl = selectnode(fens, box = [-Inf Inf L L -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>yLnl))

	enl = connectednodes(subset(boundaryfes, ebcl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>enl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>enl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 3, "node_list"=>enl))
	inl = connectednodes(subset(boundaryfes, ibcl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>inl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>inl))
	push!(essential_bcs, FDataDict("displacement"=> iuz, "component"=> 3, "node_list"=>inl))

	modeldata = FDataDict("fens"=>fens,
	"regions"=>regions,
	"essential_bcs"=>essential_bcs
	)
	modeldata = AlgoDeforLinearModule.linearstatics(modeldata)

	filebase = "bushing_h8u_algo_ms"
	modeldata["postprocessing"] = FDataDict("file"=>filebase * "-nodal-", "quantity"=>:pressure, "component"=>collect(1:1), "nodevalmethod"=>:averaging, "reportat"=>:extraptrend)
	modeldata = AlgoDeforLinearModule.exportstress(modeldata)
	modeldata["postprocessing"] = FDataDict("file"=>filebase * "-elwise-", "quantity"=>:pressure, "component"=>collect(1:1), "nodevalmethod"=>:averaging, "reportat"=>:extraptrend)
	modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)

    true
end # bushing_h8u_algo_ms

function bushing_h8_export()
    nt = 1
    for nperradius in [2, 4, 8]
    	nt=nt+1;
    	fens, fes  = Q4circlen(R, nperradius)
    	fens, fes = H8extrudeQ4(fens, fes, nt, (x, k) -> [x[1],x[2],k*thickness/nt])

    	MR = DeforModelRed3D
    	material = MatDeforElastIso(MR, E, nu)
    	femm = FEMMDeforLinearMSH8(MR, IntegDomain(fes, GaussRule(3, 2)), material)

    	boundaryfes = meshboundary(fes);
    	topl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf thickness thickness], inflate =   tolerance);
    	botl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf 0.0 0.0], inflate =   tolerance);
    	x0l = selectelem(fens,boundaryfes, box =  [0.0 0.0 -Inf Inf 0.0 thickness], inflate =   tolerance);
    	y0l = selectelem(fens,boundaryfes, box =  [-Inf Inf 0.0 0.0 0.0 thickness], inflate =   tolerance);
    	cyll = setdiff(1:count(boundaryfes),topl,botl,x0l,y0l);

    	geom = NodalField(fens.xyz)
    	u = NodalField(zeros(size(fens.xyz,1), 3)) # displacement field

    	x0nl = connectednodes(subset(boundaryfes, x0l))
    	setebc!(u,x0nl,true,1,0.0)

    	y0nl = connectednodes(subset(boundaryfes, y0l))
    	setebc!(u,y0nl,true,2,0.0)

    	cylnl = connectednodes(subset(boundaryfes, cyll))
    	setebc!(u,cylnl,true,1,0.0)
    	setebc!(u,cylnl,true,2,0.0)
    	setebc!(u,cylnl,true,3,0.0)

    	applyebc!(u)
    	numberdofs!(u)

    	cnl = selectnode(fens; box = [0 0 0 0 0 thickness], inflate  =  tolerance)

    	AE = AbaqusExporter("clcircularplatecl_h8_export_$(nperradius)");
    	HEADING(AE, "Clamped square plate with concentrated force");
    	PART(AE, "part1");
    	END_PART(AE);
    	ASSEMBLY(AE, "ASSEM1");
    	INSTANCE(AE, "INSTNC1", "PART1");
    	NODE(AE, fens.xyz);
    	ELEMENT(AE, "c3d8rh", "AllElements", 1, connasarray(femm.integdomain.fes))
    	NSET_NSET(AE, "cnl", cnl)
    	NSET_NSET(AE, "x0nl", x0nl)
    	NSET_NSET(AE, "y0nl", y0nl)
    	NSET_NSET(AE, "cylnl", cylnl)
    	ORIENTATION(AE, "GlobalOrientation", vec([1. 0 0]), vec([0 1. 0]));
    	SOLID_SECTION(AE, "elasticity", "GlobalOrientation", "AllElements", "Hourglassctl");
    	END_INSTANCE(AE);
    	END_ASSEMBLY(AE);
    	MATERIAL(AE, "elasticity")
    	ELASTIC(AE, E, nu)
    	SECTION_CONTROLS(AE, "Hourglassctl", "HOURGLASS=ENHANCED")
    	STEP_PERTURBATION_STATIC(AE)
    	BOUNDARY(AE, "ASSEM1.INSTNC1.x0nl", 1)
    	BOUNDARY(AE, "ASSEM1.INSTNC1.y0nl", 2)
    	BOUNDARY(AE, "ASSEM1.INSTNC1.cylnl", 1)
    	BOUNDARY(AE, "ASSEM1.INSTNC1.cylnl", 2)
    	BOUNDARY(AE, "ASSEM1.INSTNC1.cylnl", 3)
    	CLOAD(AE, "ASSEM1.INSTNC1.cnl", 3, Magnitude/4/length(cnl))
    	END_STEP(AE)
    	close(AE)

    	# File =  "clcircularplatecl_h8_export_$(nperradius).vtk"
    	# vtkexportmesh(File, fens, fes; vectors = [("u", u.values)])
    	# @async run(`"paraview.exe" $File`)

    end

    true
end # bushing_h8_export

function bushing_t10_algo_ms()

	fens, fes  = T10block(ang, L, Re-Ri, nc, nt, nR)
	internal_fenids = selectnode(fens, box = [0 ang 0 L 0 0], inflate = tolerance);

	MR = DeforModelRed3D
	material = MatDeforElastIso(MR, E, nu)
	femm = FEMMDeforLinearMST10(MR, IntegDomain(fes, TetRule(4)), material)
	regions = FDataDict[FDataDict("femm"=>femm)]

	boundaryfes = meshboundary(fes);
	ibcl = selectelem(fens, boundaryfes, box = [0 ang 0 L 0 0], inflate = tolerance);
	ebcl = selectelem(fens, boundaryfes, box = [0 ang 0 L Re-Ri Re-Ri], inflate = tolerance);

	for i in 1:count(fens)
		a = fens.xyz[i,1]; y = fens.xyz[i,2]; r = fens.xyz[i,3]+Ri;
		fens.xyz[i,:] .= (r*sin(a), y, r*cos(a));
	end

	essential_bcs = FDataDict[]
	x0nl = selectnode(fens, box = [0 0 -Inf Inf -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>x0nl))

	y0nl = selectnode(fens, box = [-Inf Inf 0 0 -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>y0nl))
	yLnl = selectnode(fens, box = [-Inf Inf L L -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>yLnl))

	enl = connectednodes(subset(boundaryfes, ebcl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>enl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>enl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 3, "node_list"=>enl))
	inl = connectednodes(subset(boundaryfes, ibcl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>inl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>inl))
	push!(essential_bcs, FDataDict("displacement"=> iuz, "component"=> 3, "node_list"=>inl))

	modeldata = FDataDict("fens"=>fens,
	"regions"=>regions,
	"essential_bcs"=>essential_bcs
	)
	modeldata = AlgoDeforLinearModule.linearstatics(modeldata)

	filebase = "bushing_T10_algo_ms"
	modeldata["postprocessing"] = FDataDict("file"=>filebase * "-nodal-", "quantity"=>:pressure, "component"=>collect(1:1), "nodevalmethod"=>:averaging, "reportat"=>:extraptrend)
	modeldata = AlgoDeforLinearModule.exportstress(modeldata)
	modeldata["postprocessing"] = FDataDict("file"=>filebase * "-elwise-", "quantity"=>:pressure, "component"=>collect(1:1), "nodevalmethod"=>:averaging, "reportat"=>:extraptrend)
	modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)

    true
end # bushing_t10_algo_ms

function bushing_t10_algo()

	fens, fes  = T10block(ang, L, Re-Ri, nc, nt, nR)
	internal_fenids = selectnode(fens, box = [0 ang 0 L 0 0], inflate = tolerance);

	MR = DeforModelRed3D
	material = MatDeforElastIso(MR, E, nu)
	femm = FEMMDeforLinear(MR, IntegDomain(fes, TetRule(4)), material)
	regions = FDataDict[FDataDict("femm"=>femm)]

	boundaryfes = meshboundary(fes);
	ibcl = selectelem(fens, boundaryfes, box = [0 ang 0 L 0 0], inflate = tolerance);
	ebcl = selectelem(fens, boundaryfes, box = [0 ang 0 L Re-Ri Re-Ri], inflate = tolerance);

	for i in 1:count(fens)
		a = fens.xyz[i,1]; y = fens.xyz[i,2]; r = fens.xyz[i,3]+Ri;
		fens.xyz[i,:] .= (r*sin(a), y, r*cos(a));
	end

	essential_bcs = FDataDict[]
	x0nl = selectnode(fens, box = [0 0 -Inf Inf -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>x0nl))

	y0nl = selectnode(fens, box = [-Inf Inf 0 0 -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>y0nl))
	yLnl = selectnode(fens, box = [-Inf Inf L L -Inf Inf], inflate = tolerance);
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>yLnl))

	enl = connectednodes(subset(boundaryfes, ebcl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>enl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>enl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 3, "node_list"=>enl))
	inl = connectednodes(subset(boundaryfes, ibcl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 1, "node_list"=>inl))
	push!(essential_bcs, FDataDict("displacement"=> 0.0, "component"=> 2, "node_list"=>inl))
	push!(essential_bcs, FDataDict("displacement"=> iuz, "component"=> 3, "node_list"=>inl))

	modeldata = FDataDict("fens"=>fens,
	"regions"=>regions,
	"essential_bcs"=>essential_bcs
	)
	modeldata = AlgoDeforLinearModule.linearstatics(modeldata)

	filebase = "bushing_T10_algo"
	modeldata["postprocessing"] = FDataDict("file"=>filebase * "-nodal-", "quantity"=>:pressure, "component"=>collect(1:1), "nodevalmethod"=>:averaging, "reportat"=>:extraptrend)
	modeldata = AlgoDeforLinearModule.exportstress(modeldata)
	modeldata["postprocessing"] = FDataDict("file"=>filebase * "-elwise-", "quantity"=>:pressure, "component"=>collect(1:1), "nodevalmethod"=>:averaging, "reportat"=>:extraptrend)
	modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)

    true
end # bushing_t10_algo_ms

#     true
# end # twisted_beam_export

function allrun()
    println("#####################################################")
    println("# bushing_h8_full ")
    bushing_h8_full()
    println("#####################################################")
    println("# bushing_h8_algo_full ")
    bushing_h8_algo_full()
    println("#####################################################")
    println("# bushing_h8_algo_ms ")
    bushing_h8_algo_ms()
    println("#####################################################")
    println("# bushing_h8u_algo_ms ")
    bushing_h8u_algo_ms()
    println("#####################################################")
    println("# bushing_t10_algo_ms")
    bushing_t10_algo_ms()
    return true
end # function allrun

end # module bushing_examples


#=

  Euler #40 in Julia.

  Problem 40:
  """
  An irrational decimal fraction is created by concatenating the positive integers:

  0.123456789101112131415161718192021...

  It can be seen that the 12th digit of the fractional part is 1.

  If dn represents the nth digit of the fractional part, find the
  value of the following expression.

  d1 × d10 × d100 × d1000 × d10000 × d100000 × d1000000
  """

  This Julia model was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/Julia_progs/

=#

include("Euler.jl")

# 0.00941826s
function euler40a()
    i = 1
    dlen = 1
    p = 1
    index = 10; # Index = 10, 100, 1000, ..., 1000000
    while dlen <= 1000000
        i += 1
        istr = string(i)
        istrlen = length(istr)
        if dlen+istrlen>=index
            p *= parse(Int,istr[index-dlen])
            index *= 10
        end
        dlen += istrlen
    end

    return p
end

run_euler(euler40a)


include("Graph.jl")

using SparseArrays
using LinearAlgebra
import Laplacians

function RoleSim(G; delta = 1e-6, beta = 0.15, maxIt = 100)
    n = G.n
    g = adjacentList(G)
    d = degrees(G)
    Spre = zeros(n, n)
    S = zeros(n, n)
    Sp = ones(n, n)

    calc(x, y) = begin
        hx = zeros(Bool, d[x])
        hy = zeros(Bool, d[y])
        A = Vector{Tuple{Int32, Int32, Float64}}()
        for i = 1 : d[x], j = 1 : d[y]
            push!(A, (i, j, Sp[g[x][i], g[y][j]]))
        end
        sort!(A, by=x->-x[3])
        ret = 0.0
        for (u, v, w) in A
            if !hx[u] && !hy[v]
                hx[u] = true
                hy[v] = true
                ret += w
            end
        end
        return ret
    end

    ts = 0
    while true
        ts += 1
        for i = 1 : n, j = 1 : n
            S[i, j] = (1-beta)*calc(i, j)/max(d[i], d[j]) + beta
        end
        if (norm(S - Sp) < delta) || (ts > maxIt)
            break
        end
        for i = 1 : n, j = 1 : n
            Sp[i, j] = S[i, j]
        end
    end
    return S
end

function getW(G)
    IpL = zeros(G.n, G.n)
    for (u, v) in G.E
        IpL[u, u] += 1
        IpL[v, v] += 1
        IpL[u, v] -= 1
        IpL[v, u] -= 1
    end
    foreach(i -> IpL[i, i] += 1, 1 : G.n)
    return inv(IpL)
end

function ForestSim(G)
    n = G.n
    W = getW(G)
    S = zeros(n, n)
    for i = 1 : n, j = 1 : n
        S[i, j] = min(W[i, i], W[j, j]) / max(W[i, i], W[j, j])
    end
    return S
end

function getSparseIpL(G)
    d = ones(G.n)
    for (u, v) in G.E
        d[u] += 1
        d[v] += 1
    end
    Is = zeros(Int32, G.m*2+G.n)
    Js = zeros(Int32, G.m*2+G.n)
    Vs = zeros(G.m*2+G.n)
    ID = 0
    for (u, v) in G.E
        ID += 1
        Is[ID] = u
        Js[ID] = v
        Vs[ID] = -1
        Is[ID + G.m] = v
        Js[ID + G.m] = u
        Vs[ID + G.m] = -1
    end
    for i = 1 : G.n
        Is[G.m + G.m + i] = i
        Js[G.m + G.m + i] = i
        Vs[G.m + G.m + i] = d[i]
    end
    return sparse(Is, Js, Vs, G.n, G.n)
end

function getSparseB(G)
    n = G.n
    m = G.m
    Is = zeros(Int32, m*2)
    Js = zeros(Int32, m*2)
    Vs = zeros(m*2)
    ID = 0
    for (u, v) in G.E
        ID += 1
        Is[ID] = ID
        Js[ID] = u
        Vs[ID] = 1
        Is[ID + m] = ID
        Js[ID + m] = v
        Vs[ID + m] = -1
    end
    return sparse(Is, Js, Vs, m, n)
end

function approxForest(G; eps = 0.1)
    n = G.n
    m = G.m
    IpL = getSparseIpL(G)
    B = getSparseB(G)
    k = round(Int, log2(G.n) / eps^2)
    f = Laplacians.approxchol_sddm(IpL, tol=1e-12)
    ret = zeros(G.n)
    for i = 1 : k
        y1 = B' * randn(m)
        y2 = rand(n)
        z1 = f(y1)
        z2 = f(y2)
        foreach(j -> ret[j] += (z1[j]^2 + z2[j]^2), 1 : G.n)
    end
    ret ./= k
    return ret
end

function calcBins(G, maxk)
    maxbin = ceil(Int32, log2(G.n))
    n = G.n
    g = adjacentList(G)
    d = degrees(G)
    B = zeros(Int32, maxk, n, maxbin)
    dk = zeros(Int32, maxk, n)
    for s = 1 : n
        dist = zeros(Int32, n)
        fill!(dist, 1048576)
        dist[s] = 0
        Q = zeros(Int32, n)
        front = 1
        rear = 1
        Q[1] = s
        while front <= rear
            u = Q[front]
            front += 1
            if dist[u] == maxk
                continue
            end
            for v in g[u]
                if dist[u]+1 < dist[v]
                    dist[v] = dist[u] + 1
                    B[dist[v], s, floor(Int32, log2(d[v]))+1] += 1
                    dk[dist[v], s] += 1
                    rear += 1
                    Q[rear] = v
                end
            end
        end
    end
    return B, dk
end

function StructSim(G; maxk = 3, wi = 0.5)
    n = G.n
    d = degrees(G)
    B, dk = calcBins(G, maxk)
    maxb = size(B, 3)
    S = zeros(n, n)
    for i = 1 : n, j = 1 : n
        S[i, j] = min(d[i], d[j]) / max(d[i], d[j])
    end
    for k = 1 : maxk
        for i = 1 : n, j = 1 : n
            delta = 0.0
            foreach(z -> delta += min(B[k, i, z], B[k, j, z]), 1 : maxb)
            delta /= max(dk[k, i], dk[k, j])
            S[i, j] = (1.0 - wi) * S[i, j] + wi * delta
        end
    end
    return S
end


using Statistics

abstract type CorMethod end
struct CorAll <: CorMethod end
struct CorComplete <: CorMethod end
struct CorPairwise <: CorMethod end

function Statistics.cor(df::DataFrame; method::CorMethod=CorAll())
    cor1(i, j) = nrow(df2) == 0 ? missing : cor(df2[i], df2[j])

    function cor2(i, j)
       x = dropmissing(DataFrame([df2[i], df2[j]]))
       nrow(x) == 0 ? missing : cor(x[1], x[2])
    end

    use_cor = method == CorPairwise() ? cor2 : cor1
    df2 = method == CorComplete() ? dropmissing(df) : df

    m = Matrix{Union{Float64, Missing}}(undef, ncol(df), ncol(df))
    for i in 1:ncol(df), j in i:ncol(df)
        m[i, j] = use_cor(i, j)
        m[j, i] = m[i, j]
    end
    m
end

Chapter07/05. Handling missing data/cor.jl


export Material
export matparam, kottke_avg_param
import Base:string

# Note that the output of tensorize is always a matrix: for K = 1, it is a 1×1 matrix.
tensorize(x::Number) = tensorize(@SVector(fill(complex(x), 3)))  # SComplexF{3}(x,x,x): vector
tensorize(v::AbsVecNumber) = diagm(Val(0)=>SComplexF{3}(v))  # complex(v) is allocating, so use SComplexF{K}(v)
tensorize(m::AbsMatNumber) = S²ComplexF3(m)  # complex(m) is allocating, so use S²ComplexF{K,K*K}(m)

struct Material
    name::String
    param::Tuple{S²ComplexF3,S²ComplexF3}  # (material parameter interacting with E, and H)
end

Material(name::String, param::Tuple2{MatParam}) = Material(name, tensorize.(param))
Material(name::String; ε::MatParam=1, μ::MatParam=1) = Material(name, (ε,μ))  # ε, μ: keyword arguments
# Consider implementing Material2 and Material3 for 2D and 3D solvers.

string(m::Material) = m.name
matparam(m::Material, ft::FieldType) = m.param[Int(ft)]

kottke_avg_param(param1::AbsMatNumber, param2::AbsMatNumber, n12::AbsVecReal, rvol1::Real) =
    (K = length(n12); kottke_avg_param(S²ComplexF{K,K^2}(param1), S²ComplexF{K,K^2}(param2), SFloat{K}(n12), rvol1))

# Implement the averaging scheme between two local material parameter tensors developed in
# the paper by Kottke, Farjadpour, Johnson entitled "Perturbation theory for anisotropic
# dielectric interfaces, and application to subpixel smoothing of discretized numerical
# methods", Physical Review E 77 (2008): 036611.
function kottke_avg_param(param1::S²ComplexF{K}, param2::S²ComplexF{K}, n12::SFloat{K}, rvol1::Real) where {K}
    Scomp = @SMatrix rand(K,K-1)  # directions complementary to n12; works even for K = 1
    Stemp = [n12 Scomp]  # SMatrix{K,K}
    S = qr(Stemp).Q  # nonallocating; 1st column is normalized n12

    τ1 = τ_trans(transpose(S) * param1 * S)  # express param1 in S coordinates, and apply τ transform
    τ2 = τ_trans(transpose(S) * param2 * S)  # express param2 in S coordinates, and apply τ transform

    τavg = τ1 .* rvol1 + τ2 .* (1-rvol1)  # volume-weighted average

    return S * τ⁻¹_trans(τavg) * transpose(S)  # apply τ⁻¹ and transform back to global coordinates
end

# Kottke's averaging scheme reduced for the field dimension orthogonal to the shape
# dimension (such as Hz in the TE mode).  In that case, the field is always parallel to the
# shape boundaries, so the averaging scheme reduces to simple arithmetic averaging.  For
# example for the case with Shape{2} and Kf = 1 (like Hz in the TE mode), the material
# parameter is an 1×1 tensor (= scalar) and therefore the situation is the same as averaging
# isotropic ε between E-field voxels, which is arithmetic averaging.  The only other case is
# the case with Shape{1} and Kf = 1  (like Ex in the slapes in the z-direction).  We can
# easily prove that in these cases again the averaging reduces to simple arithmetic
# averaging between material parameter tensors.  See Agenda > MaxwellBase > Feb/22/2020.
kottke_avg_param(param1::AbsMatNumber, param2::AbsMatNumber, rvol1::Real) =
    param1 .* rvol1 + param2 .* (1-rvol1)

# Equation (4) of the paper by Kottke et al.  Need to verify for K = 1 and 2.
function τ_trans(ε::S²ComplexF3)
    ε₁₁, ε₂₁, ε₃₁, ε₁₂, ε₂₂, ε₃₂, ε₁₃, ε₂₃, ε₃₃ = ε

    return S²ComplexF3(
        -1/ε₁₁, ε₂₁/ε₁₁, ε₃₁/ε₁₁,
        ε₁₂/ε₁₁, ε₂₂ - ε₂₁*ε₁₂/ε₁₁, ε₃₂ - ε₃₁*ε₁₂/ε₁₁,
        ε₁₃/ε₁₁, ε₂₃ - ε₂₁*ε₁₃/ε₁₁, ε₃₃ - ε₃₁*ε₁₃/ε₁₁
    )
    # t = zeros(3,3)
    # t[2:3, 2:3] = s[2:3, 2:3]
    # s11 = s[1,1]
    # s[1,1] = -1
    #
    # t = t - (s(:,1) * s(1,:)) ./ s11
    # return t
end

function τ_trans(ε::S²ComplexF2)
    ε₁₁, ε₂₁, ε₁₂, ε₂₂ = ε

    return S²ComplexF2(
        -1/ε₁₁, ε₂₁/ε₁₁,
        ε₁₂/ε₁₁, ε₂₂ - ε₂₁*ε₁₂/ε₁₁
    )
end

τ_trans(ε::S²ComplexF1) = -1 ./ ε  # S²ComplexF1

# function τ_trans(ε::S²ComplexF{K}) where {K}
#     τ_temp1 = @SMatrix [ε[i,j] for i = 2:K, j = 2:K]  # cannot use type parameter K here: error is generated
#     τ_temp2 = [@SMatrix(zeros(ComplexF,1,K-1)); τ_temp1]
#     τ_temp = [@SVector(zeros(ComplexF,K)) τ_temp2]
#
#     σcol_temp = @SVector [ε[i,1] for i = 2:K]
#     σcol = [SVec{1,ComplexF}(-1); σcol_temp]
#
#     σrow_temp = @SVector [ε[1,j] for j = 2:K]
#     σrow = [SVec{1,ComplexF}(-1); σrow_temp]
#
#     τ = τ_temp - (σcol * transpose(σrow)) ./ ε[1,1]
#
#     return τ
# end

# Equation (23) of the paper by Kottke et al.
function τ⁻¹_trans(τ::S²ComplexF3)
    τ₁₁, τ₂₁, τ₃₁, τ₁₂, τ₂₂, τ₃₂, τ₁₃, τ₂₃, τ₃₃ = τ

    return S²ComplexF3(
        -1/τ₁₁, -τ₂₁/τ₁₁, -τ₃₁/τ₁₁,
        -τ₁₂/τ₁₁, τ₂₂ - τ₂₁*τ₁₂/τ₁₁, τ₃₂ - τ₃₁*τ₁₂/τ₁₁,
        -τ₁₃/τ₁₁, τ₂₃ - τ₂₁*τ₁₃/τ₁₁, τ₃₃ - τ₃₁*τ₁₃/τ₁₁
    )

    # s = zeros(3,3)
    # s[2:3, 2:3] = t[2:3, 2:3]
    # t11 = t[1,1]
    # t[1,1] = 1
    #
    # s = s - (t(:,1) * t(1,:)) ./ t11
    # return s
end

function τ⁻¹_trans(τ::S²ComplexF2)
    τ₁₁, τ₂₁, τ₁₂, τ₂₂ = τ

    return S²ComplexF2(
        -1/τ₁₁, -τ₂₁/τ₁₁,
        -τ₁₂/τ₁₁, τ₂₂ - τ₂₁*τ₁₂/τ₁₁
    )
end

τ⁻¹_trans(τ::S²ComplexF1) = -1 ./ τ  # S²ComplexF1