module StatsBase

import Base: length, size, isempty, values, sum, show, maximum, minimum, extrema
import Base.Cartesian: @nloops, @nref, @nextract
using Base: @irrational, @propagate_inbounds
using DataAPI
import DataAPI: describe
import DataStructures: heapify!, heappop!, percolate_down!
using SortingAlgorithms
using Missings

using Statistics
using LinearAlgebra
using Random
using Printf
using SparseArrays
import Random: rand, rand!
import LinearAlgebra: BlasReal, BlasFloat
import Statistics: mean, mean!, var, varm, varm!, std, stdm, cov, covm,
                   cor, corm, cov2cor!, unscaled_covzm, quantile, sqrt!,
                   median, middle
import StatsAPI: pairwise, pairwise!

    ## tackle compatibility issues

export

    ## weights
    AbstractWeights,    # abstract type to represent any weight vector
    Weights,            # to represent a generic weight vector
    AnalyticWeights,    # to represent an analytic/precision/reliability weight vector
    FrequencyWeights,   # to representing a frequency/case/repeat weight vector
    ProbabilityWeights, # to representing a probability/sampling weight vector
    UnitWeights,        # to representing a uniform weight vector
    weights,            # construct a generic Weights vector
    aweights,           # construct an AnalyticWeights vector
    fweights,           # construct a FrequencyWeights vector
    pweights,           # construct a ProbabilityWeights vector
    eweights,           # construct an exponential Weights vector
    uweights,           # construct an UnitWeights vector
    wsum,               # weighted sum with vector as second argument
    wsum!,              # weighted sum across dimensions with provided storage

    ## moments
    skewness,       # (standardized) skewness
    kurtosis,       # (excessive) kurtosis
    moment,         # central moment of given order
    mean_and_var,   # (mean, var)
    mean_and_std,   # (mean, std)
    mean_and_cov,   # (mean, cov)

    ## scalarstats
    geomean,     # geometric mean
    harmmean,    # harmonic mean
    genmean,     # generalized/power mean
    middle,      # the mean of two real numbers
    mode,        # find a mode from data (the first one)
    modes,       # find all modes from data

    zscore,      # compute Z-scores
    zscore!,     # compute Z-scores inplace or to a pre-allocated array

    percentile,  # quantile using percentage (instead of fraction) as argument
    nquantile,   # quantiles at [0:n]/n

    span,        # The range minimum(x):maximum(x)
    variation,   # ratio of standard deviation to mean
    sem,         # standard error of the mean, i.e. sqrt(var / n)
    mad,         # median absolute deviation
    iqr,         # interquatile range

    genvar,      # generalized variance
    totalvar,    # total variation

    entropy,        # the entropy of a probability vector
    renyientropy,   # the Rényi (generalised) entropy of a probability vector
    crossentropy,   # cross entropy between two probability vectors
    kldivergence,   # K-L divergence between two probability vectors

    summarystats,   # summary statistics
    describe,       # print the summary statistics

    # deviation
    counteq,        # count the number of equal pairs
    countne,        # count the number of non-equal pairs
    sqL2dist,       # squared L2 distance between two arrays
    L2dist,         # L2 distance between two arrays
    L1dist,         # L1 distance between two arrays
    Linfdist,       # L-inf distance between two arrays
    gkldiv,         # (Generalized) Kullback-Leibler divergence between two vectors
    meanad,         # mean absolute deviation
    maxad,          # maximum absolute deviation
    msd,            # mean squared deviation
    rmsd,           # root mean squared deviation
    psnr,           # peak signal-to-noise ratio (in dB)

    # cov
    scattermat,     # scatter matrix (i.e. unnormalized covariance)
    cov2cor,        # converts a covariance matrix to a correlation matrix
    cor2cov,        # converts a correlation matrix to a covariance matrix
    CovarianceEstimator, # abstract type for covariance estimators
    SimpleCovariance,    # simple covariance estimator

    ## counts
    addcounts!,     # add counts to an accumulating array or map
    counts,         # count integer values in given arrays
    proportions,    # proportions of integer values in given arrays
                    # (normalized version of counts)
    countmap,       # count distinct values and return a map
    proportionmap,  # proportions of distinct values returned as a map

    ## ranking
    ordinalrank,    # ordinal ranking ("1234" ranking)
    competerank,    # competition ranking ("1 2 2 4" ranking)
    denserank,      # dense ranking ("1 2 2 3" ranking)
    tiedrank,       # tied ranking ("1 2.5 2.5 4" ranking)

    ## rankcorr
    corspearman,       # spearman's rank correlation
    corkendall,        # kendall's rank correlation

    ## partialcor
    partialcor,        # partial correlation

    ## signalcorr
    autocov!, autocov,      # auto covariance
    autocor!, autocor,      # auto correlation
    crosscov!, crosscov,    # cross covariance
    crosscor!, crosscor,    # cross correlation
    pacf!, pacf,            # partial auto-correlation

    ## sampling
    samplepair,     # draw a pair of distinct elements
    sample,         # sampling from a population
    sample!,        # sampling from a population, with pre-allocated output
    wsample,        # sampling from a population with weights
    wsample!,       # weighted sampling, with pre-allocated output

    ## empirical
    ecdf,           # empirical cumulative distribution function
    ECDF,           # type for empirical cumulative distribution function

    AbstractHistogram,
    Histogram,
    midpoints,
    # histrange,

    ## robust
    trim,           # trimmed set
    trim!,          # trimmed set
    winsor,         # Winsorized set
    winsor!,        # Winsorized set
    trimvar,        # variance of the mean of a trimmed set

    ## misc
    rle,            # run-length encoding
    inverse_rle,    # inverse run-length encoding
    indexmap,       # construct a map from element to index
    levelsmap,      # construct a map from n unique elements to [1, ..., n]
    indicatormat,   # construct indicator matrix
    pairwise,       # pairwise application of functions
    pairwise!,      # pairwise! application of functions

    # statistical models
    CoefTable,
    StatisticalModel,
    RegressionModel,

    adjr2,
    adjr²,
    aic,
    aicc,
    bic,
    coef,
    coefnames,
    coeftable,
    confint,
    cooksdistance,
    crossmodelmatrix,
    deviance,
    dof,
    dof_residual,
    fit,
    fit!,
    fitted,
    informationmatrix,
    isfitted,
    islinear,
    leverage,
    loglikelihood,
    meanresponse,
    modelmatrix,
    mss,
    response,
    responsename,
    nobs,
    nulldeviance,
    nullloglikelihood,
    rss,
    score,
    stderror,
    vcov,
    predict,
    predict!,
    residuals,
    r2,
    r²,

    ConvergenceException,

    # data standardization
    standardize,
    AbstractDataTransform, # the type to represent a abstract data transformation
    ZScoreTransform,       # the type to represent a z-score data transformation
    UnitRangeTransform     # the type to represent a 0-1 data transformation

# source files

include("common.jl")
include("weights.jl")
include("moments.jl")
include("scalarstats.jl")
include("robust.jl")
include("deviation.jl")
include("cov.jl")
include("counts.jl")
include("ranking.jl")
include("toeplitzsolvers.jl")
include("rankcorr.jl")
include("signalcorr.jl")
include("partialcor.jl")
include("empirical.jl")
include("hist.jl")
include("pairwise.jl")
include("misc.jl")

include("sampling.jl")
include("statmodels.jl")

include("transformations.jl")

include("deprecates.jl")

end # module


"""

In this example, the distance between the C-terminal carbon
and a sodium atom is computed and ploted as a function of time

"""

using M3GTools

println(" Loading Plots... ")
using Plots
ENV["GKSwstype"]="nul" # This supresses the need of a display while plotting

# Initialize simulation data and path to VMD executable
println(" Reading simulation data ... ")
mysim = M3GTools.init(psf="../simulation_files/structure.psf",
                      dcd="../simulation_files/structure.dcd",
                      vmd="vmd")

println(" Defining selections... ")
cterm = M3GTools.select(mysim,"protein and resid 253 and name C")
sod = M3GTools.select(mysim,"index 20455")

# Vector that will contain the distances for each frame
d = Vector{Float64}(undef,mysim.nframes)

println(" Computing distances ... ")
for i in 1:mysim.nframes
  sides, x, y, z = M3GTools.nextframe(mysim)

  # Wrap coordinates of the sodium around cterm

  x_cterm = [ x[cterm[1]], y[cterm[1]], z[cterm[1]] ]
  M3GTools.wrap!(sides,x,y,z;center=x_cterm,sel=sod)

  # Compute distance

  d[i] = sqrt( ( x[cterm[1]] - x[sod[1]] )^2 +
               ( y[cterm[1]] - y[sod[1]] )^2 +
               ( z[cterm[1]] - z[sod[1]] )^2 )

end
M3GTools.closedcd(mysim)

println(" Plotting... ")

x = [ i for i in 1:mysim.nframes ]
plot(x,d,xlabel="Frame",ylabel="Distance")
savefig("distance.pdf")
println(" Created figure: distance.pdf")


include("LD.jl")
using MATLAB
using ForwardDiff
using ProgressMeter
x=range(.25,stop=1.2,length=Int(1e4))
y1=@showprogress map(Dimer_Lagrangian_Descriptor_Function,x);
logz1=log1p.(abs.(y1))

mat"figure()"
mat"plot($x,$y1)"
# mat"plot($x,$logz1);hold on;"
delta=1e-12
g(x)=( Dimer_Lagrangian_Descriptor_Function(x+delta)-Dimer_Lagrangian_Descriptor_Function(x-delta))/(2*delta)
f(x)=abs.(( Dimer_Lagrangian_Descriptor_Function(x-delta)-Dimer_Lagrangian_Descriptor_Function(x)+Dimer_Lagrangian_Descriptor_Function(x+delta))/(delta^2))
# g(x)=( Dimer_Lagrangian_Descriptor_Function(x-.001)-Dimer_Lagrangian_Descriptor_Function(x+.001))/(.002)
y2=@showprogress map(g,x);
logz2=log10.(abs.(y2))
mat"figure()"
mat"plot($x,$y2,'k')"

mat"figure()"
mat"plot($x,$logz2,'k.')"

# y3=@showprogress map(f,x)
# # logz=log1p.(abs.(y2))
# # logz3=log1p.(abs.(y3))

# mat"figure()"
# mat"plot($x,$y3)"
# # mat"plot($x,$logz3)"


import Pkg
Pkg.activate(dirname(dirname(@__FILE__)))

using LinearAlgebra
using Random
using SimLib
using SpinSymmetry
using Statistics
using XXZNumerics

### positiondata setup
location = SaveLocation()
N = 8
shots = 30
fields = [-0.2, -0.1, 0.1, 0.2]
alpha = 6
basis = symmetrized_basis(N, Flip(N), 0)

## generate positions
Random.seed!(5)
pdd = PositionDataDescriptor(:box_pbc, 1, N, shots, [0.5, 1.0, 1.5], location)
posdata = load_or_create(pdd)

model = RandomPositionsXXZWithXField(pdd, PowerLaw(alpha), fields, :ensemble, basis)

### make ED tasks
hopping_operator = 1/2 * real.(σx⊗σx + σy⊗σy)

evaltask = Energies()
eontask = EigenstateOccupation("xpol", symmetrize_state(normalize!(ones(2^N)), basis))
eevtask = OperatorDiagonal("xmag", symmetrize_operator(sum(op_list(σx/2, N))/N, basis))
eltask  = EigenstateLocality("szsz", symmetrize_operator(single_spin_op(σz, 1, N)*single_spin_op(σz, 2, N), basis))
eltask2 = EigenstateLocality("sx", symmetrize_operator(single_spin_op(σx, 1, N), basis))
eltask3 = EigenstateLocality("hopping", symmetrize_operator(kron(hopping_operator, I(2^(N-2))), basis))
eltask4 = EigenstateLocality("xmag", symmetrize_operator(sum(op_list(σx/2, N))/N, basis))
lsrtask = LevelSpacingRatio()
#hcetask = HalfChainEntropy()

tasks = [evaltask, eontask, eevtask, lsrtask, eltask, eltask2, eltask3, eltask4]#, hcetask]

### THREADED RUN

edd = EDDataDescriptor(model, Full(), location)
@time edata = run_ed(edd, tasks, Serial())

save.(edata) # save all data

LinearAlgebra.BLAS.set_num_threads(1)

@time edata2 = run_ed(edd, tasks, Threaded())
@time edata3 = run_ed(edd, tasks, Parallel(4))

compare(a,b) = sum(abs2, a.data - b.data)
compare2(a,b) = sum(abs2, mean(a.data;dims=(1,2,3)) - mean(b.data;dims=(1,2,3)))
@show compare.(edata , edata2)
@show compare.(edata2, edata3)
@show compare2.(edata , edata2)
@show compare2.(edata2, edata3)

## Ensemble data
ensdd = EnsembleDataDescriptor(edd)
save(create(ensdd))


"""Pareto smoothed importance sampling (PSIS)

This module implements Pareto smoothed importance sampling (PSIS) and PSIS
leave-one-out cross-validation for Python (Numpy).

Included functions
------------------
psisloo
    Pareto smoothed importance sampling leave-one-out log predictive densities.

psislw
    Pareto smoothed importance sampling.

gpdfitnew
    Estimate the paramaters for the Generalized Pareto Distribution (GPD).

gpinv
    Inverse Generalised Pareto distribution function.

sumlogs
    Sum of vector where numbers are represented by their logarithms.

References
----------
Aki Vehtari, Andrew Gelman and Jonah Gabry (2015). Efficient implementation
of leave-one-out cross-validation and WAIC for evaluating fitted Bayesian
models. arXiv preprint arXiv:1507.04544.

Aki Vehtari and Andrew Gelman (2015). Pareto smoothed importance sampling.
arXiv preprint arXiv:1507.02646.
"""

module ParetoSmoothedImportanceSampling

using StatsFuns, Statistics

psis_path = @__DIR__

include("psisloo.jl")
include("psislw.jl")
include("gpdfitnew.jl")
include("gpinv.jl")
include("waic.jl")
include("pk_utilities.jl")
include("dic.jl")

export
    psis_path

end


using Random
using Random: default_rng
using .Threads: nthreads

using Distributions, HypothesisTests, StatsBase, ProgressMeter, LoopVectorization,
    ThreadPools, GaitSymmetry, SafeBuffers
using GaitSymmetry: Ratio

export gensimdata, calcsymfun, estimatepower, estimatepower!, powersim

function gensimdata(rng, xdist, ydist, samples, T=promote_type(eltype(xdist), eltype(ydist)))
    stor = Matrix{T}(undef, (samples, 2))
    _gensimdata!(rng, stor, xdist, ydist)
end

function _gensimdata!(rng, stor, xdist, ydist)
    randn!(rng, stor)
    @views _randn_kernel!(stor[:,1], mean(xdist), std(xdist))
    @views _randn_kernel!(stor[:,2], mean(ydist), std(ydist))

    return stor
end

function _randn_kernel!(arr, μ, σ)
    @avx for i in eachindex(arr)
        arr[i] = abs(fma(arr[i], σ, μ))
    end
end

function calcsymfun(symfun, stor::AbstractArray{T}, len=size(stor, 1)) where T
    size(stor, 2) === 2 || throw(DimensionMismatch("expected 2 columns; got $(size(stor, 2))"))
    len ≤ size(stor, 1) || throw(DimensionMismatch("`len` must be ≤ `size(stor, 1)`"))
    res = Vector{T}(undef, len)
    _calcsymfun!(res, symfun, len, stor)
end

function _calcsymfun!(res, symfun, len, stor)
    @inbounds for i in 1:len-1
        x = stor[i,1]
        x′ = stor[i+1,1]
        y = stor[i,2]

        res[i] = symfun(x, x′) - symfun(x, y)
    end
    res[end] = symfun(stor[end,1], stor[1,1]) - symfun(stor[end,1], stor[end,2])

    return res
end

function estimatepower(symfun, xdist, ydist, n, samples)
    stor = gensimdata(default_rng(), xdist, ydist, n*samples)
    res = calcsymfun(symfun, stor)
    _estimatepower(res, n, samples)
end

function estimatepower(symfun, stor, n, samples)
    res = calcsymfun(symfun, stor, n*samples)
    _estimatepower(res, n, samples)
end

function estimatepower!(res, symfun, stor, n, samples)
    fullN = n*samples
    _calcsymfun!(res, symfun, fullN, stor)
    _estimatepower(res, n, samples)
end

function _estimatepower(res, n, samples)
    numsig = _numsignificant(res, n, samples)
    return numsig# / samples
end

function _numsignificant(res, n, samples, α=0.05)
    numsig = 0
    for i in 0:(samples-1)
        @inbounds slice = @view(res[(1:n).+i])
        avg, sd = mean_and_std(slice; corrected=true)
        if pvalue(OneSampleTTest(avg, sd, n)) ≤ α
            numsig += 1
        end
    end

    return numsig
end

function powersim(
    metrics,
    ratios,
    vars,
    Ns,
    samples,
    U=Float32;
    null=false,
    batch=1
)
    p = Progress(length(metrics)*length(ratios)*length(vars)*sum(Ns .* samples)+1;
        dt=5, desc="Batch $(lpad(batch, 2)) progress: ")

    ratios = convert(Array{U}, ratios)
    data = Array{U}(undef, (length(Ns), length(metrics), length(ratios), length(vars)))

    fullN = maximum(Ns)*samples
    respool = BufferPool(nthreads()+1, () -> Vector{U}(undef, fullN))
    storpool = BufferPool(nthreads()+1, () -> Matrix{U}(undef, (fullN, 2)))

    idxs = [ (σi, i) for σi in eachindex(vars) for i in eachindex(ratios) ]
    @qthreads for (σi, i) in idxs
        σ = vars[σi]
        xdist = Normal(1, σ)
        μ = ratios[i]
        μ_sd = μ*σ
        ydist = ifelse(null, xdist, Normal(μ, μ_sd))

        withbuffer(storpool) do stor
            _gensimdata!(default_rng(), stor, xdist, ydist)
            for mi in eachindex(metrics)
                if metrics[mi] === Alv20b
                    metric = Alv20b(U(σ))
                else
                    metric = metrics[mi]
                end

                withbuffer(respool) do res
                    _calcsymfun!(res, metric, fullN, stor)
                    for ni in eachindex(Ns)
                        data[ni, mi, i, σi] = _estimatepower(res, Ns[ni], samples)
                        next!(p; step=Ns[ni]*samples)
                    end
                end
            end
        end
    end
    finish!(p)

    return data
end


immutable MultivariateNormal <: ContinuousMultivariateDistribution
    mean::Vector{Float64}
    covchol::Cholesky{Float64}
    function MultivariateNormal(m::Vector{Float64}, c::Cholesky{Float64})
        if length(m) == size(c, 1) == size(c, 2)
            new(m, c)
        else
            error("Dimensions of mean vector and covariance matrix do not match")
        end
    end
end

function MultivariateNormal(mean::Vector{Float64}, cov::Matrix{Float64})
    MultivariateNormal(mean, cholfact(cov))
end
function MultivariateNormal(mean::Vector{Float64})
    MultivariateNormal(mean, eye(length(mean)))
end
function MultivariateNormal(cov::Matrix{Float64})
    MultivariateNormal(zeros(size(cov, 1)), cov)
end
MultivariateNormal() = MultivariateNormal(zeros(2), eye(2))

function cdf{T <: Real}(d::MultivariateNormal, x::Vector{T})
    k = length(d.mean)
    if k > 3
        error("Dimension larger than three is not supported yet")
    end
    stddev = sqrt(diag(var(d)))
    z = (x - d.mean) ./ stddev
    C = diagmm(d.covchol[:U], 1.0 / stddev)
    C = C'C
    if k == 3
        return tvtcdf(0, z, C[[2, 3, 6]])
    elseif k == 2
        return bvtcdf(0, z[1], z[2], C[2])
    else
        return cdf(Normal(), z[1])
    end
end

function entropy(d::MultivariateNormal)
    S = cov(d)
    n = size(S, 1)
    for j in 1:n
        for i in 1:n
            S[i, j] *= 2.0 * pi * e
        end
    end
    return 0.5 * logdet(S)
end

mean(d::MultivariateNormal) = d.mean

function mgf(d::MultivariateNormal, t::AbstractVector)
    m, S = d.mean, var(d)
    return exp(dot(t, m) + 0.5 * dot(t, S * t))
end

function cf(d::MultivariateNormal, t::AbstractVector)
    m, S = d.mean, var(d)
    return exp(im * dot(t, m) - 0.5 * dot(t, S * t))
end

pdf{T<:Real}(d::MultivariateNormal, x::Vector{T}) = exp(logpdf(d, x))

function logpdf{T <: Real}(d::MultivariateNormal, x::Vector{T})
    k = length(d.mean)
    u = x - d.mean
    chol_ldiv!(d.covchol, u)
    return -0.5 * (k * log(2.0 * pi) + logdet(d.covchol) + dot(u, u))
end

function logpdf!{T <: Real}(r::AbstractVector,
                            d::MultivariateNormal,
                            x::Matrix{T})
    mu::Vector{Float64} = d.mean
    k = length(mu)
    if size(x, 1) != k
        throw(ArgumentError("The dimension of x is inconsistent with d."))
    end
    n = size(x, 2)
    u = Array(Float64, k, n)
    for j = 1:n # u[:,j] = x[:,j] - mu
        for i = 1:k
            u[i, j] = x[i, j] - mu[i]
        end
    end
    chol_ldiv!(d.covchol, u)
    c::Float64 = -0.5 * (k * log(2.0 * pi) + logdet(d.covchol))
    for j = 1:n
        dot_uj = 0.0
        for i = 1:k
            dot_uj += u[i, j] * u[i, j]
        end
        r[j] = c - 0.5 * dot_uj
    end
end

function logpdf{T <: Real}(d::MultivariateNormal, x::Matrix{T})
    r = Array(Float64, size(x, 2))
    logpdf!(r, d, x)
    return r
end

function rand!(d::MultivariateNormal, x::Vector, tmp::Vector)
    randn!(tmp)
    At_mul_B(x, d.covchol[:U], tmp)
    for i in 1:length(x)
        x[i] += d.mean[i]
    end
    return x
end

function rand!(d::MultivariateNormal, x::Vector)
    tmp = randn(length(x))
    return rand!(d, x, tmp)
end

function rand(d::MultivariateNormal)
    x = Array(Float64, length(d.mean))
    tmp = Array(Float64, length(d.mean))
    return rand!(d, x, tmp)
end

function rand!(d::MultivariateNormal, X::Matrix)
    k = length(d.mean)
    tmp = Array(Float64, k)
    m, n = size(X)
    if m != k
        throw(ArgumentError("Wrong dimensions"))
    end
    randn!(X)
    for i in 1:n
        # This makes a copy of X[:, i]
        At_mul_B(tmp, d.covchol[:U], X[:, i])
        for dim in 1:m
            X[dim, i] = tmp[dim] + d.mean[dim]
        end
    end
    return X
end

function var(d::MultivariateNormal)
    U = d.covchol[:U]
    return U'U
end

function chol_ldiv!(chol::Cholesky{Float64}, u::VecOrMat{Float64})
    Base.LinAlg.LAPACK.trtrs!('U', 'T', 'N', chol.UL, u)
end


# class constructors
ZZ(::Type{T}) where T<:Integer = ZZ{T}

# construction
category_trait(::Type{<:ZZ}) = EuclidianDomainTrait
basetype(::Type{<:ZZ{T}}) where T = T
depth(::Type{<:ZZ}) = 1
_lcunit(a::ZZ) = sign(a.val)
issimpler(a::T, b::T) where T<:ZZ = abs(a.val) < abs(b.val)
iscoprime(a::T, b::T) where T<:ZZ = gcd(a.val, b.val) == 1
value(a::ZZ) = a.val

copy(a::ZZ) = typeof(a)(a.val)
ZZ(a::T) where T<:Integer = ZZ{T}(a)
ZZ{T}(a::ZZ) where T = ZZ{T}(T(a.val))

# promotion and conversion
_promote_rule(::Type{ZZ{T}}, ::Type{ZZ{S}}) where {S,T} = ZZ{promote_type(S,T)}
_promote_rule(::Type{ZZ{T}}, ::Type{QQ{S}}) where {S,T} = QQ{promote_type(S,T)}
promote_rule(::Type{ZZ{T}}, ::Type{S}) where {S<:Integer,T} = ZZ{promote_type(S,T)}
promote_rule(::Type{ZZ{T}}, ::Type{Rational{S}}) where {S<:Integer,T} = QQ{promote_type(S,T)}

# induced homomorphism
function (h::Hom{F,R,S})(p::Z) where {F,R,S,Z<:ZZ{<:R}}
    Z(h.f(p.val))
end

for op in (:+, :-, :/, :(==), :div, :rem, :divrem, :gcd, :gcdx, :pgcd, :pgcdx)
    @eval begin
        ($op)(a::ZZ{T}, b::Integer) where T = ($op)(promote(a, b)...)
        ($op)(a::Integer, b::ZZ{T}) where T = ($op)(promote(a, b)...)
    end
end

Base.isless(p::T, q::T) where T<:ZZ = isless(p.val, q.val)

# operations for ZZ
+(a::ZZ{T}, b::ZZ{T}) where T = ZZ(checked_add(a.val, b.val))
-(a::ZZ{T}, b::ZZ{T}) where T = ZZ(checked_sub(a.val, b.val))
-(a::ZZ{T}) where T = ZZ{T}(checked_neg(a.val))
*(a::ZZ{T}, b::ZZ{T}) where T = ZZ(checked_mul(a.val, b.val))
*(a::ZZ{T}, b::Integer) where T = ZZ(T(checked_mul(a.val, T(b))))
*(a::Integer, b::ZZ{T}) where T = ZZ(T(checked_mul(T(a), b.val)))
^(a::ZZ{BigInt}, b::Integer) = ZZ(a.val^b)
divrem(a::T, b::T) where T<:ZZ = T.(divrem(a.val, b.val))
div(a::T, b::T) where T<:ZZ = T(a.val ÷ b.val)
rem(a::T, b::T) where T<:ZZ = T(a.val % b.val)

isunit(a::ZZ) = abs(a.val) == 1
isone(a::ZZ) = isone(a.val)
iszero(a::ZZ) = iszero(a.val)
zero(::Type{ZZ{T}}) where T = ZZ(zero(T))
one(::Type{ZZ{T}}) where T = ZZ(one(T))
==(a::ZZ{T}, b::ZZ{T}) where T = a.val == b.val
hash(a::ZZ, h::UInt) = hash(a.val, h)

gcd(a::T, b::T) where T<: ZZ = T(gcd(a.val, b.val))
lcm(a::T, b::T) where T<:ZZ = T(lcm(a.val, b.val))

factor(a::Z) where Z <:ZZ = [Z(first(x)) => last(x) for x in Primes.factor(value(a))]
isirreducible(a::ZZ) = isirreducible(a.val)
isirreducible(a::Integer) = Primes.isprime(a)

Base.show(io::IO, z::ZZ) = print(io, z.val)


using DynamicAxisWarping
using DynamicAxisWarping.Datasets
using Plots
gr()

data, labels = ucr_traindata("50words")

plot(
    [fastdtw(data[:, i], data[:, i+10], radius)[1] for radius = 1:100, i = 1:10],
    show = true,
)
xlabel!("radius")


module TestMetric

using MXNet
using Base.Test

################################################################################
# Supporting functions
################################################################################

"""
Returns a random n x m array in which each column defines a discrete probability distribution.
Each column contains numbers between 0 and 1, and each column sums to 1.
"""
function generate_probs(n, m)
    # Init
    result = rand(n, m)

    # Normalize: ensure each column sums to 1
    for j = 1:m
	colsum = sum(result[:, j])
	for i = 1:n
	    result[i, j] /= colsum
	end
    end
    result
end

function loglikelihood{T <: AbstractFloat}(labels::Vector{T}, probs::Array{T, 2})
    LL = 0.0
    for i = 1:size(labels, 1)
        LL += log(probs[Int(labels[i]) + 1, i])    # labels are zero-based
    end
    LL / size(labels, 1)
end

################################################################################
# Test Implementations
################################################################################

function test_ace()
    info("EvalMetric::ACE")
    n_categories   = 4
    n_observations = 100
    labels         = convert(Vector{Float32}, rand(0:(n_categories - 1), n_observations))    # MXNet uses Float32
    probs          = convert(Array{Float32}, generate_probs(n_categories, n_observations))
    LL             = loglikelihood(labels, probs)
    metric         = mx.ACE()    # For categorical variables, ACE == -LL
    mx._update_single_output(metric, mx.NDArray(labels), mx.NDArray(probs))
    LL_v2 = metric.ace_sum / metric.n_sample
    @test_approx_eq_eps LL LL_v2 1e-12
end

################################################################################
# Run tests
################################################################################
test_ace()

end


b₀(ξ) = abs(ξ) < 0.5 ? 1 : 0
b₁(ξ) = (ξ >= -1 && ξ <= 1) ? (1 - abs(ξ)) : 0
b₂(ξ) = (ξ >= -1.5 && ξ <= 1.5) ?
    ((ξ >= -0.5 && ξ <= 0.5) ? (0.75 - ξ^2) : (0.5 * (1.5 - abs(ξ))^2)) : 0


### TODO:
### 4. Test skipping all the complicated indexing tricks; use fancy indexing
### NOTES:
### 1. Caching coefficients provides a ~11% speedup at low ℓ, but that falls off
###    as ℓ increases, reaching ~4% around ℓ=512.
### 2. This is probably a much more significant advantage for ALFs.

using .SphericalFunctions: complex_powers!
using Quaternionic: AbstractQuaternion, to_euler_phases!

@inline ϵ(m) = ifelse(m > 0 && isodd(m), -1, 1)

"""
    d!(d, expiβ, ℓₘₐₓ, H_rec_coeffs)
    d!(d, expiβ, ℓₘₐₓ)
    d!(d, β, ℓₘₐₓ, H_rec_coeffs)
    d!(d, β, ℓₘₐₓ)
    d(expiβ, ℓₘₐₓ)
    d(β, ℓₘₐₓ)

Compute Wigner's d matrix dˡₘₚ,ₘ(β)

# Notes

This function is the preferred method of computing the d matrix for large ell
values.  In particular, above ell≈32 standard formulas become completely
unusable because of numerical instabilities and overflow.  This function uses
stable recursion methods instead, and should be usable beyond ell≈1000.

The result is returned in a 1-dimensional array ordered as

    [
        d(ell, mp, m, β)
        for ell in range(ell_max+1)
        for mp in range(-min(ℓ, mp_max), min(ℓ, mp_max)+1)
        for m in range(-ell, ell+1)
    ]

"""
function d!(d, expiβ::Complex, ℓₘₐₓ, H_rec_coeffs)
    H!(d, expiβ, ℓₘₐₓ, ℓₘₐₓ, H_rec_coeffs, WignerDindex)

    @inbounds for ℓ in 0:ℓₘₐₓ
        i0 = WignerDindex(ℓ, -ℓ, -ℓ)
        for m in -ℓ:-1
            oddm = isodd(m)
            for m′ in -ℓ:m
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ - m) * (2ℓ + 1) - m′ + ℓ
                d[i1] = ifelse(oddm, -d[i2], d[i2])
            end
            for m′ in m+1:0
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ - m′) * (2ℓ + 1) - m + ℓ
                d[i1] = ifelse(oddm, -d[i2], d[i2])
            end
            for m′ in 1:-m
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ - m′) * (2ℓ + 1) - m + ℓ
                d[i1] = ifelse(isodd(m′)⊻oddm, -d[i2], d[i2])
            end
            for m′ in 1-m:ℓ
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ + m) * (2ℓ + 1) + m′ + ℓ
                d[i1] = ifelse(isodd(m′)⊻oddm, -d[i2], d[i2])
            end
        end
        for m in 0:ℓ
            for m′ in -ℓ:-m-1
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ - m) * (2ℓ + 1) - m′ + ℓ
                d[i1] = d[i2]
            end
            for m′ in m+1:ℓ
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ + m) * (2ℓ + 1) + m′ + ℓ
                d[i1] = ifelse(isodd(m′), -d[i2], d[i2])
            end
        end
        for m′ in 1:2:ℓ
            i1 = i0 + (ℓ + m′) * (2ℓ + 1) + ℓ
            for m in abs(m′):ℓ
                d[i1+m] *= -1
            end
        end
    end
    d
end
function d!(d, expiβ::Complex{T}, ℓₘₐₓ) where {T<:Real}
    d!(d, expiβ, ℓₘₐₓ, H_recursion_coefficients(ℓₘₐₓ, T))
end
function d!(d, β::T, ℓₘₐₓ, H_rec_coeffs) where {T<:Real}
    d!(d, cis(β), ℓₘₐₓ, H_rec_coeffs)
end
function d!(d, β::T, ℓₘₐₓ) where {T<:Real}
    d!(d, cis(β), ℓₘₐₓ, H_recursion_coefficients(ℓₘₐₓ, T))
end
function d(expiβ::Complex{T}, ℓₘₐₓ) where {T<:Real}
    𝔡 = Array{T}(undef, WignerDsize(ℓₘₐₓ, ℓₘₐₓ))
    d!(𝔡, expiβ, ℓₘₐₓ, H_recursion_coefficients(ℓₘₐₓ, T))
end
d(β::T, ℓₘₐₓ) where {T<:Real} = d(cis(β), ℓₘₐₓ)

"""
    dstorage(ℓₘₐₓ, T)

Construct space to compute Wigner's ``d`` matrix in place.

This returns the `d` argument needed by [`d!`](@ref).

"""
function dstorage(ℓₘₐₓ, ::Type{T}) where {T<:Real}
    Vector{T}(undef, WignerDsize(ℓₘₐₓ))
end

"""
    dprep(ℓₘₐₓ, T)

Construct space and pre-compute recursion coefficients to compute Wigner's
``d`` matrix in place.

This returns the `(d, H_rec_coeffs)` arguments needed by [`d!`](@ref).

"""
function dprep(ℓₘₐₓ, ::Type{T}) where {T<:Real}
    d = dstorage(ℓₘₐₓ, T)
    H_rec_coeffs = H_recursion_coefficients(ℓₘₐₓ, T)
    d, H_rec_coeffs
end

"""
    D!(𝔇, R, ℓₘₐₓ, (aₙᵐ,bₙᵐ,dₙᵐ), expimα, expimγ)

Compute Wigner's 𝔇 matrix

This function implements the preferred method of computing the 𝔇 matrix for large ell
values.  In particular, above ell≈32 standard formulas become completely
unusable because of numerical instabilities and overflow.  This function uses
stable recursion methods instead, and should be usable beyond ell≈1000.

This function computes 𝔇ˡₘₚ,ₘ(R).  The result is returned in a 1-dimensional
array ordered as

    [
        𝔇(ell, mp, m, R)
        for ell in range(ell_max+1)
        for mp in range(-min(ℓ, mp_max), min(ℓ, mp_max)+1)
        for m in range(-ell, ell+1)
    ]

"""
function D!(𝔇, R::AbstractQuaternion, ℓₘₐₓ, H_rec_coeffs, expimα, expimγ)
    expiα, expiβ, expiγ = to_euler_phases(R)
    H!(𝔇, expiβ, ℓₘₐₓ, ℓₘₐₓ, H_rec_coeffs, WignerDindex)
    complex_powers!(expimα, expiα)
    complex_powers!(expimγ, expiγ)

    # 𝔇ˡₘₚ,ₘ(R) = dˡₘₚ,ₘ(R) exp[iϕₐ(m-mp)+iϕₛ(m+mp)] = dˡₘₚ,ₘ(R) exp[i(ϕₛ+ϕₐ)m+i(ϕₛ-ϕₐ)mp]
    # exp[iϕₛ] = R̂ₛ = hat(R[0] + 1j * R[3]) = zp
    # exp[iϕₐ] = R̂ₐ = hat(R[2] + 1j * R[1]) = zm.conjugate()
    # exp[i(ϕₛ+ϕₐ)] = zp * zm.conjugate() = z[2] = zᵧ
    # exp[i(ϕₛ-ϕₐ)] = zp * zm = z[0] = zₐ
    @inbounds for ℓ in 0:ℓₘₐₓ
        i0 = WignerDindex(ℓ, -ℓ, -ℓ)
        for m in -ℓ:-1
            oddm_factor = ifelse(isodd(m), -1, 1)
            for m′ in -ℓ:m
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ - m) * (2ℓ + 1) - m′ + ℓ
                𝔇[i1] = oddm_factor * 𝔇[i2] * conj(expimγ[-m+1] * expimα[-m′+1])
            end
            for m′ in m+1:0
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ - m′) * (2ℓ + 1) - m + ℓ
                𝔇[i1] = oddm_factor * 𝔇[i2] * conj(expimγ[-m+1] * expimα[-m′+1])
            end
            for m′ in 1:-m
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ - m′) * (2ℓ + 1) - m + ℓ
                𝔇[i1] = ifelse(isodd(m′), -1, 1) * oddm_factor * 𝔇[i2] * conj(expimγ[-m+1]) * expimα[m′+1]
            end
            for m′ in 1-m:ℓ
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ + m) * (2ℓ + 1) + m′ + ℓ
                𝔇[i1] = ifelse(isodd(m′), -1, 1) * oddm_factor * 𝔇[i2] * conj(expimγ[-m+1]) * expimα[m′+1]
            end
        end
        for m in 0:ℓ
            for m′ in -ℓ:-m-1
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ - m) * (2ℓ + 1) - m′ + ℓ
                𝔇[i1] = 𝔇[i2] * expimγ[m+1] * conj(expimα[-m′+1])
            end
            for m′ in m+1:ℓ
                i1 = i0 + (ℓ + m′) * (2ℓ + 1) + m + ℓ
                i2 = i0 + (ℓ + m) * (2ℓ + 1) + m′ + ℓ
                𝔇[i1] = ifelse(isodd(m′), -𝔇[i2], 𝔇[i2]) * expimγ[m+1] * expimα[m′+1]
            end
        end
        for m′ in -ℓ:0
            i1 = i0 + (ℓ + m′) * (2ℓ + 1) + ℓ
            for m in abs(m′):ℓ
                𝔇[i1+m] *= expimγ[m+1] * conj(expimα[-m′+1])
            end
        end
        for m′ in 1:ℓ
            i1 = i0 + (ℓ + m′) * (2ℓ + 1) + ℓ
            for m in abs(m′):ℓ
                𝔇[i1+m] *= ifelse(isodd(m′), -1, 1) * expimγ[m+1] * expimα[m′+1]
            end
        end
    end
    𝔇
end

"""
    Dstorage(ℓₘₐₓ, T)

Construct space to compute Wigner's ``𝔇`` matrix in place.

This returns the `D` argument needed by [`D!`](@ref).

"""
function Dstorage(ℓₘₐₓ, ::Type{T}) where {T<:Real}
    Vector{Complex{T}}(undef, WignerDsize(ℓₘₐₓ))
end

"""
    Dprep(ℓₘₐₓ, T)

Construct space and pre-compute recursion coefficients to compute Wigner's
``𝔇`` matrix in place.

This returns the `(D, H_rec_coeffs, expimα, expimγ)` arguments needed by
[`D!`](@ref).

"""
function Dprep(ℓₘₐₓ, ::Type{T}) where {T<:Real}
    𝔇 = Dstorage(ℓₘₐₓ, T)
    H_rec_coeffs = H_recursion_coefficients(ℓₘₐₓ, T)
    expimα = Vector{Complex{T}}(undef, ℓₘₐₓ+1)
    expimγ = Vector{Complex{T}}(undef, ℓₘₐₓ+1)
    𝔇, H_rec_coeffs, expimα, expimγ
end

@doc raw"""
    Y(s, R, ℓₘₐₓ)
    Y!(Y, s, R, ℓₘₐₓ)

Evaluate (and write into `Y`, if present) the values of ``{}_{s}Y_{\ell,
m}(R)`` for the input value of `s`, for all ``(\ell, m)`` throughout the range
specified by `wigner`.  `R` is assumed to be a unit quaternion (which may be
`Rotor`, or simply a `Quaternion`).  If `R` does not have unit magnitude, the
output elements will be too large by a factor ``|R|^{\ell}``.  If `Y` is not
present, a new array will be created.

The spherical harmonics of spin weight ``s`` are related to Wigner's
``\mathfrak{D}`` matrix as
```math
\begin{aligned}
{}_{s}Y_{\ell, m}(R)
  &= (-1)^s \sqrt{\frac{2\ell+1}{4\pi}} \mathfrak{D}^{(\ell)}_{m, -s}(R) \\
  &= (-1)^s \sqrt{\frac{2\ell+1}{4\pi}} \bar{\mathfrak{D}}^{(\ell)}_{-s, m}(\bar{R}).
\end{aligned}
```
"""
function Y!(Y, R, ℓₘₐₓ, spin, H_rec_coeffs, Hwedge, expimϕ, ℓₘᵢₙ=0)
    if length(Y) < Ysize(ℓₘᵢₙ, ℓₘₐₓ)
        error("Input `Y` has length $(length(Y)); which is not enough for ℓₘₐₓ=$ℓₘₐₓ")
    end
    if length(Hwedge) < WignerHsize(ℓₘₐₓ, abs(spin))
        error(
            "Input `Hwedge` has length $(length(Hwedge)); "
            *"which is not enough for ℓₘₐₓ=$ℓₘₐₓ with spin=$spin"
        )
    end

    expiϕ, expiθ, expiγ = to_euler_phases(R)
    H!(Hwedge, expiθ, ℓₘₐₓ, abs(spin), H_rec_coeffs)
    complex_powers!(expimϕ, expiϕ)
    spin_factor = (isodd(spin) ? -1 : 1) * ϵ(spin) * expiγ^-spin

    # Yˡₘₚ,ₘ(R) ∝ dˡₘₚ,ₘ(R) exp[iϕₐ(m-mp)+iϕₛ(m+mp)] = dˡₘₚ,ₘ(R) exp[i(ϕₛ+ϕₐ)m+i(ϕₛ-ϕₐ)mp]
    # exp[iϕₛ] = R̂ₛ = hat(R[0] + 1j * R[3]) = zp
    # exp[iϕₐ] = R̂ₐ = hat(R[2] + 1j * R[1]) = zm.conjugate()
    # exp[i(ϕₛ+ϕₐ)] = zp * zm.conjugate() = z[2] = zᵧ
    # exp[i(ϕₛ-ϕₐ)] = zp * zm = z[0] = zₐ
    iᴰ = 1
    @inbounds for ℓ in ℓₘᵢₙ:ℓₘₐₓ
        if ℓ < abs(spin)
            Y[iᴰ:iᴰ+2ℓ] .= 0
            iᴰ += 2ℓ+1
        else
            factor = spin_factor * √((2ℓ+1)/(4eltype(Y)(π)))
            for m in -ℓ:-1
                iᴴ = WignerHindex(ℓ, m, -spin, abs(spin))
                Y[iᴰ] = factor * Hwedge[iᴴ] * conj(expimϕ[-m+1])  # ϵ(m′)==1
                iᴰ += 1
            end
            for m in 0:ℓ
                iᴴ = WignerHindex(ℓ, m, -spin, abs(spin))
                Y[iᴰ] = factor * ϵ(m) * Hwedge[iᴴ] * expimϕ[m+1]
                iᴰ += 1
            end
        end
    end
    Y
end

"""
    Ystorage(ℓₘₐₓ, T)
    Ystorage(ℓₘₐₓ, T, ℓₘᵢₙ)

Construct space to compute the spin-weighted spherical harmonics ``ₛYₗ,ₘ`` in
place.

This returns the `Y` argument needed by [`Y!`](@ref).

"""
function Ystorage(ℓₘₐₓ, ::Type{T}, ℓₘᵢₙ=0) where {T<:Real}
    Vector{Complex{T}}(undef, Ysize(ℓₘᵢₙ, ℓₘₐₓ))
end

function Yworkspace(ℓₘₐₓ, sₘₐₓ, ::Type{T}) where {T<:Real}
    Hwedge = Vector{T}(undef, WignerHsize(ℓₘₐₓ, abs(sₘₐₓ)))
    expimϕ = Vector{Complex{T}}(undef, ℓₘₐₓ+1)
    Hwedge, expimϕ
end

"""
    Yprep(ℓₘₐₓ, sₘₐₓ, T, ℓₘᵢₙ)

Prepare the storage, recursion coefficients, and workspace to compute ₛYₗ,ₘ
data up to the maximum sizes given.

Returns a tuple of `Y, H_rec_coeffs, Hwedge, expimϕ`, which can be passed to
the correspondingly named arguments of `Y!`.

Note that the same results of this function can be passed to `Y!`, even if the
value of `ℓₘₐₓ` passed to that function is smaller than the value passed to
this function, or the value of `spin` passed to that function is smaller (in
absolute value) than the `sₘₐₓ` passed to this function.  However, the value of
`ℓₘᵢₙ` passed to that function *must not* be smaller than the value passed to
this function (unless one of the other sizes is sufficiently smaller).

"""
function Yprep(ℓₘₐₓ, sₘₐₓ, ::Type{T}, ℓₘᵢₙ=0) where {T<:Real}
    Y = Ystorage(ℓₘₐₓ, T, ℓₘᵢₙ)
    H_rec_coeffs = H_recursion_coefficients(ℓₘₐₓ, T)
    Hwedge, expimϕ = Yworkspace(ℓₘₐₓ, sₘₐₓ, T)
    Y, H_rec_coeffs, Hwedge, expimϕ
end


## PARTE 1
cd(@__DIR__)
input = reduce(hcat,
    map(line -> parse.(Bool, collect(line)), eachline("input.txt")));

##
scores = sum(input, dims=2) |> vec
γ_bits = broadcast(>(size(input, 2) ÷ 2), scores)
γ, ϵ = reduce(enumerate(reverse(γ_bits)); init=(0,0)) do (γ, ϵ), (i, b)
    γ |= b << (i-1)
    ϵ |= ~b << (i-1)
    γ, ϵ
end
γ * ϵ


export LinearGaussian

"""
    mutable struct LinearGaussian{I <: Tuple{Vararg{Float64}}} <:
        Conditional{I, Tuple{}, I, Float64, Normal}

`LinearGaussian` defines an sfunc whose mean is a linear function of its parents.  A
`LinearGaussian`'s output type is a `Float`, its parameter type is
`Tuple{I, Float64, Float64}`, and it's contained *sfunc* is a `Normal` mean `0.0`.

# Type parameters
- `I`: the input type(s) of the `LinearGaussian`

See also: [`Conditional`](@ref), [`Normal`](@ref)
"""
mutable struct LinearGaussian{I <: Tuple{Vararg{Float64}}} <:
    Conditional{I, Tuple{}, I, Float64, Normal}
    sf :: Normal
    params :: Tuple{Tuple{Vararg{Float64}}, Float64, Float64}
    """
        function LinearGaussian(weights :: Tuple{Vararg{Float64}}, bias :: Float64, sd :: Float64)

    `LinearGaussian` constructor

    # Arguments
    - `weights::Tuple{Vararg{Float64}}`: the weights of each parent
    - `bias::Float64`: the bias of the mean of the internal `Normal` *sfunc*
    - `sd::Float64`: the standard deviation of the internal `Normal` *sfunc*
    """
    function LinearGaussian(weights :: Tuple{Vararg{Float64}}, bias :: Float64, sd :: Float64)
        params = (weights, bias, sd)
        sf = Normal(0.0, sd)
        N = length(weights)
        new{NTuple{N, Float64}}(sf, params)
    end
end
#=
@impl begin
    struct LinearGaussianInitialStats end
    function initial_stats(sf::LinearGaussian)
        (weights, bias, sd) = sf.params
        initweights = Tuple(zeros(length(weights)))
        Dict(initweights=>(0.0, 0.0))
    end
end
=#

function gensf(lg::LinearGaussian, inputs::Tuple{Vararg{Float64}})::Normal
    (weights, bias, sd) = lg.params
    to_sum = inputs .* weights
    mean = isempty(to_sum) ? bias : sum(to_sum) + bias
    return Normal(mean, sd)
end


module BetaTest

using Test
using ReactiveMP
using Distributions
using Random

import ReactiveMP: mirrorlog

@testset "Beta" begin

    # Beta comes from Distributions.jl and most of the things should be covered there
    # Here we test some extra ReactiveMP.jl specific functionality

    @testset "vague" begin
        d = vague(Beta)

        @test typeof(d) <: Beta
        @test mean(d) === 0.5
        @test params(d) === (1.0, 1.0)
    end

    @testset "prod" begin
        @test prod(ProdAnalytical(), Beta(3.0, 2.0), Beta(2.0, 1.0)) ≈ Beta(4.0, 2.0)
        @test prod(ProdAnalytical(), Beta(7.0, 1.0), Beta(0.1, 4.5)) ≈ Beta(6.1, 4.5)
        @test prod(ProdAnalytical(), Beta(1.0, 3.0), Beta(0.2, 0.4)) ≈ Beta(0.19999999999999996, 2.4)
    end

    @testset "mean(::typeof(log))" begin
        @test mean(log, Beta(1.0, 3.0)) ≈ -1.8333333333333335
        @test mean(log, Beta(0.1, 0.3)) ≈ -7.862370395825961
        @test mean(log, Beta(4.5, 0.3)) ≈ -0.07197681436958758
    end

    @testset "mean(::typeof(mirrorlog))" begin
        @test mean(mirrorlog, Beta(1.0, 3.0)) ≈ -0.33333333333333337
        @test mean(mirrorlog, Beta(0.1, 0.3)) ≈ -0.9411396776150167
        @test mean(mirrorlog, Beta(4.5, 0.3)) ≈ -4.963371962929249
    end

end

end


euclidex = quote
    while a != b
        if a > b
            a = a - b
            p = p - q
            r = r - s
        else
            q = q - p
            b = b - a
            s = s - r
        end
    end
end

fermat = quote
    while r != 0
        if r > 0
            r = r - v
            v = v + 2
        else
            r = r + u
            u = u + 2
        end
    end
end

wensley = quote
    while d>= E
        if P < a+b
            b = b/2
            d = d/2
        else
            a = a+b
            y = y+d/2
            b = b/2
            d = d/2
        end
    end
end

lcm = quote
    while x != y
        if x > y
            x = x - y
            v = v + u
        else
            y = y - x
            u = u + v
        end
    end
end

knuth = quote
    while (s >= d) && (r != 0)
        if 2*r-rp+q < 0
        t  = r
        r  = 2*r-rp+q+d+2
        rp = t
        q  = q+4
        d  = d+2
        elseif (2*r-rp+q >= 0) && (2*r-rp+q < d+2)
        t  = r
        r  = 2*r-rp+q
        rp = t
        d  = d+2
        elseif (2*r-rp+q >= 0) && (2*r-rp+q >= d+2) && (2*r-rp+q < 2*d+4)
        t  = r
        r  = 2*r-rp+q-d-2
        rp = t
        q  = q-4
        d  = d+2
        else # ((2*r-rp+q >= 0) && (2*r-rp+q >= 2*d+4))
        t  = r
        r  = 2*r-rp+q-2*d-4
        rp = t
        q  = q-8
        d  = d+2
        end
    end
end

divbin = quote
    while b != B
        x = 2*x
        b = b/2
        if r >= b
            x = x+1
            r = r-b
        end
    end
end

psolv2m = quote
    n = 1
    while true
        if true
            a = 2*(n+1)*(n+3/2)*a
            b = 4*(n+1)*b
            c = 1/2*(n+3/2)*c
            n = n+1
        else
            a = 2*a
            b = 4*b
            c = 1/2*c
        end
    end
end

psolv3m = quote
    n1, n2 = 1, 1
    while true
        if true
            a = 2*(n1+1)*(n1+3/2)*a
            b = 4*(n1+1)*b
            c = 1/2*(n1+3/2)*c
            n1 = n1+1
        elseif c>0
            a = 2*a
            b = 4*b
            c = 1/2*c
        else
            a = 2*(n2+1)^3*(n2+3/2)^3*a
            b = 4*(n2+1)^3*b
            c = 1/2*(n2+3/2)^3*c
            n2 = n2+1
        end
    end
end

psolv4m = quote
    n1, n2 = 1, 1
    while true
        if true
            a = 2*(n1+1)*(n1+3/2)*a
            b = 4*(n1+1)*b
            c = 1/2*(n1+3/2)*c
            n1 = n1+1
        elseif c>0
            a = 2*a
            b = 4*b
            c = 1/2*c
        elseif c<0
            a = 2*(n2+1)^3*(n2+3/2)^3*a
            b = 4*(n2+1)^3*b
            c = 1/2*(n2+3/2)^3*c
            n2 = n2+2
        else
            a = 6*a
            b = 36*b
            c = 1/6*c
        end
    end
end

psolv10m = quote
    n1, n2, n3, n4, n5 = 1, 1, 1, 1, 1
    while true
        if true
            a = 2*(n1+1)*(n1+3/2)*a
            b = 4*(n1+1)*b
            c = 1/2*(n1+3/2)*c
            n1 = n1+1
        elseif c>0
            a = 2*a
            b = 4*b
            c = 1/2*c
        elseif c<0
            a = 2*(n2+1)^3*(n2+3/2)^3*a
            b = 4*(n2+1)^3*b
            c = 1/2*(n2+3/2)^3*c
            n2 = n2+1
        elseif c<0
            a = 6*a
            b = 36*b
            c = 1/6*c
        elseif c<0
            a = 2*(n3+1)^5*(n3+3/2)^5*a
            b = 4*(n3+1)^5*b
            c = 1/2*(n3+3/2)^5*c
            n3 = n3+1
        elseif c<0
            a = 3*a
            b = 9*b
            c = 1/3*c
        elseif c<0
            a = 2*(n4+1)^7*(n4+3/2)^7*a
            b = 4*(n4+1)^7*b
            c = 1/2*(n4+3/2)^7*c
            n4 = n4+1
        elseif c<0
            a = 5*a
            b = 25*b
            c = 1/5*c
        elseif c<0
            a = 2*(n5+1)^9*(n5+3/2)^9*a
            b = 4*(n5+1)^9*b
            c = 1/2*(n5+3/2)^9*c
            n5 = n5+1
        else
            a = 4*a
            b = 16*b
            c = 1/4*c
        end
    end
end

prodbin = quote
    while y!=0
        if y % 2 ==1
            z = z+x;
            y = y-1;
        end
        x = 2*x;
        y = y/2;
    end
end

dijkstra = quote
    while q!=1
        q=q/4;
        h=p+q;
        p=p/2;
        if r>=h
            p=p+q;
            r=r-h;
        end
    end
end

z3sqrt = quote
    while 2*p*r >= e
        if 2*r-2*q-p >= 0
            r = 2*r-2*q-p;
            q = q+p;
            p = p/2;
        else
            r = 2*r;
            p = p/2;
        end
    end
end

writers = quote
    r = 0;
    w = 0;
    k = 14;
    # c1 = 3;
    # c2 = 2;
    while true
        global r, w, k, c1, c2
        if w == 0
            r = r+1;
            k = k-c1;
        elseif r == 0
            w = w+1;
            k = k-c2;
        elseif w==0
            r = r-1;
            k = k+c1;
        elseif r==0
            w = w-1;
            k = k+c2;
        end
        @info  2*w+k+3*r-14==0
    end
end


struct ProductNumbers <: IncrementalInference.FunctorPairwise
  z::Distributions.Normal
end
getSample(s::ProductNumbers, N::Int=1) = (reshape(rand(s.z,N),1,:), )

function (s::ProductNumbers)(res::Array{Float64},
      userdata::FactorMetadata,
      idx::Int,
      meas::Tuple{Array{Float64,2}},
      X::Array{Float64,2},
      Y::Array{Float64,2},
      XY::Array{Float64,2}  )
  #
  res[1] = XY[1,idx] - X[1,idx]*Y[1,idx] + meas[1][1,idx]
  nothing
end

struct AreEqual <: IncrementalInference.FunctorPairwise
  z::Distributions.Normal
end
function getSample(s::AreEqual, N::Int=1)
  return (reshape(rand(s.z,N),1,:), )
end

function (s::AreEqual)(res::Array{Float64},
                       userdata::FactorMetadata,
                       idx::Int,
                       meas::Tuple{Array{Float64,2}},
                       X::Array{Float64,2},
                       Y::Array{Float64,2}  )
  #
  res[1] = X[1,idx]-Y[1,idx] + meas[1][1,idx]
  nothing
end

struct Square <: IncrementalInference.FunctorPairwise
  z::Distributions.Normal
end
getSample(s::Square, N::Int=1) = (reshape(rand(s.z,N),1,:), )

function (s::Square)(res::Array{Float64},
      userdata::FactorMetadata,
      idx::Int,
      meas::Tuple{Array{Float64,2}},
      X::Array{Float64,2},
      XX::Array{Float64,2}  )
  #
  res[1] = XX[1,idx] - X[1,idx]*X[1,idx] + meas[1][1,idx]
  nothing
end

mutable struct NumbersPrior <: IncrementalInference.FunctorSingleton
  z::BallTreeDensity
end
getSample(s::NumbersPrior, N::Int=1) = (reshape(rand(s.z,N),1,:), )


# This file is a part of SolidStateDetectors.jl, licensed under the MIT License (MIT).

abstract type AbstractCoordinatePoint{T, N, S} <: StaticArrays.FieldVector{N, T} end

struct CartesianPoint{ T <: RealQuantity } <: AbstractCoordinatePoint{T, 3, :cartesian}
    x::T
    y::T
    z::T
end

struct CylindricalPoint{ T <: RealQuantity } <: AbstractCoordinatePoint{T, 3, :cylindrical}
    r::T
    φ::T # in radian
    z::T
    CylindricalPoint{T}(r,φ,z) where {T <: RealQuantity} = new(r,mod(φ,2π),z)
end

function CylindricalPoint( pt::CartesianPoint{T} )::CylindricalPoint{T} where {T <: RealQuantity}
    return CylindricalPoint{T}(sqrt(pt.x * pt.x + pt.y * pt.y), atan(pt.y, pt.x), pt.z)
end

function CartesianPoint( pt::CylindricalPoint{T} )::CartesianPoint{T} where {T <: RealQuantity}
    sφ::T, cφ::T = sincos(pt.φ)
    return CartesianPoint{T}(pt.r * cφ, pt.r * sφ, pt.z)
end

function convert(type::Type{CylindricalPoint}, origin::SolidStateDetectors.CartesianPoint{T})::CylindricalPoint{T} where T
    return CylindricalPoint(origin)
end

function convert(type::Type{CartesianPoint}, origin::SolidStateDetectors.CylindricalPoint{T})::CartesianPoint{T} where T
    return CartesianPoint(origin)
end

function (+)(p1::CartesianPoint{T}, p2::CartesianPoint{T})::CartesianPoint{T} where {T <: Real}
    return CartesianPoint{T}( p1.x + p2.x, p1.y + p2.y, p1.z + p2.z )
end

# function (-)(p1::CartesianPoint{T}, p2::CartesianPoint{T})::CartesianPoint{T} where {T <: Real}
#     return CartesianPoint{T}( p1.x - p2.x, p1.y - p2.y, p1.z - p2.z )
# end

function (+)(p1::CylindricalPoint{T}, p2::CylindricalPoint{T})::CylindricalPoint{T} where {T <: Real}
    return CylindricalPoint(CartesianPoint(p1) + CartesianPoint(p2))
end

function broadcast(f::Function, v::Vector{CartesianPoint{T}}, p::CartesianPoint{T}) where {T}
    f(v ,[p for i in eachindex(v)])
end

function broadcast(f::Function, v::Vector{CylindricalPoint{T}}, p::CylindricalPoint{T}) where {T}
    f(v ,[p for i in eachindex(v)])
end

@recipe function f(p::CylindricalPoint)
    @series begin
        CartesianPoint(p)
    end
end
@recipe function f(p::CartesianPoint)
    # st -> :scatter
    @series begin
        [p.x], [p.y], [p.z]
    end
end
@recipe function f(v::AbstractVector{<:CartesianPoint})
    st -> :scatter
    @series begin
        [v[i].x for i in eachindex(v) ], [v[i].y for i in eachindex(v) ], [v[i].z for i in eachindex(v) ]
    end
end
@recipe function f(v::AbstractVector{<:CylindricalPoint})
    st -> :scatter
    @series begin
        map(x->CartesianPoint(x),v)
    end
end
 ##aliases for ease of use
const cyp = CylindricalPoint
const cap = CartesianPoint


#=
gaussian_smoothing:
- Julia version: 1.1.0
- Author: Jirik
- Date: 2019-04-08
=#
using LarSurf
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using Plasm, SparseArrays
# using Pandas
using Pio3d
# using Seaborn

V = [[0 0]; [2 0]; [1 2]; [3 2]; [4 2]; [4 0] ]'
EVch = spzeros(Int8, 7, 6)
EVch[1,1] = 1
EVch[1,2] = 1
EVch[2,2] = 1
EVch[2,3] = 1
EVch[3,1] = 1
EVch[3,3] = 1
EVch[4,3] = 1
EVch[4,4] = 1
EVch[5,2] = 1
EVch[5,4] = 1
EVch[6,2] = 1
EVch[6,6] = 1
EVch[7,4] = 1
EVch[7,6] = 1
Matrix(EVch)

FVch = spzeros(Int8, 3, 6)
FVch[1, 1] = 1
FVch[1, 2] = 1
FVch[1, 3] = 1
FVch[2, 2] = 1
FVch[2, 3] = 1
FVch[2, 4] = 1
# FVch[3, 2] = 1
# FVch[3, 4] = 1
# FVch[3, 6] = 1

FEch = spzeros(Int8, 3, 7)
FEch[1, 1] = 1
FEch[1, 2] = 1
FEch[1, 3] = 1
FEch[2, 2] = 1
FEch[2, 4] = 1
FEch[2, 5] = 1

function eye(sz)
    arr = spzeros(Int8, sz, sz)
    for i=1:sz
        arr[i,i] = 1
    end
    return arr
end

function setIzero!(arr)
    for i=1:size(arr)[1]
        arr[i,i] = 0
    end
    return arr
end

function getDiag(arr)
    sz = size(arr)[1]
    arro = zeros(eltype(arr), sz)
    for i=1:sz
        arro[i] = arr[i,i]
    end
    return arro
end

function make_full(arr)
    arrout = Array{eltype(arr), 2}(undef, size(arr)...)
    for i=1:size(arr)[1]
        for j=1:size(arr)[2]
            arrout[i,j] = arr[i,j]
        end
    end
    return arrout
end
#neighboor vertexes

function smoothing_EV(V, EVch::SparseMatrixCSC, k=0.35)
#     Matrix(setIzero!(EVch' * EVch))

    neighboorNumber = getDiag(EVch' * EVch)
    neighboors = setIzero!(EVch' * EVch)

#     targetV =  neighboors * V' ./ neighboorNumber
#     diff = targetV - V'
#     newV = (V' + k * diff)'

    targetV = V * neighboors ./ neighboorNumber'
    println("targetV shape: ", size(targetV))
    diff = targetV - V
    println("diff shape: ", size(diff))
    newV = V + k * diff

    return make_full(newV)
end

smoothing_EV(V, EVch)

## Read data from file
xystep = 1
zstep = 1
threshold = 4000;
pth = Pio3d.datasets_join_path("medical/orig/sample-data/nrn4.pklz")
datap = Pio3d.read3d(pth)

data3d = datap["data3d"]
segmentation = data3d .> threshold

block_size = [5,5,5]
basicmodel, Flin, larmodel = LarSurf.get_surface_grid_per_block(segmentation, block_size; return_all=true)
someV, topology = basicmodel
filtered_bigFV = topology[1]
# someV, filtered_bigFV = basicmodel
bigV = someV
# bigV, tmodel = larmodel
# bigVV, bigEV, bigFV, bigCV = tmodel
# return (bigV,[FVreduced])

Plasm.View(basicmodel)
# Plasm.View((bigV,[bigVV, bigEV, filtered_bigFV]))

FV = filtered_bigFV
function get_EV(FV)
	# @info "getFV " FV
	EV = []
	for f in FV
		push!(EV, [[f[1],f[2]],[f[3],f[4]],  [f[1],f[3]],[f[2],f[4]]])
	end
	# display(cat(EV))
	@info "EV example !" EV EV[1]

	catev = cat(EV)
	@info "cat(EV)" catev
	@info "hcat(EV)" hcat(EV)
	@info "vcat(EV)" vcat(EV)
	doubleedges = Base.sort(catev)

	# doubleedges = Base.sort(cat(EV;dims=2))
	@info "EV double example" doubleedges
	doubleedges = convert(Lar.Cells, doubleedges)
	EV = [doubleedges[k] for k=1:2:length(doubleedges)]
	return EV
end
function smoothing_FV(V::Array, FV::Array{Array{Int64,1},1})
	# @info "smoothing FV " FV
	EV = get_EV(FV)
	aEV = LarSurf.ll2array(EV)

	# nrn4

	# notAllBigEV =

	kEV = LarSurf.characteristicMatrix(aEV, size(bigV)[2])
	# kEV = Lar.characteristicMatrix(EV)
	newBigV = smoothing_EV(V, kEV, 0.6)
	return newBigV
end

@info "smoothing_FV example" bigV FV FV[1]
newBigV = smoothing_FV(bigV, FV)

Plasm.View((newBigV * 100,[filtered_bigFV]))
# Plasm.View((newBigV * 100,[bigVV, bigEV, filtered_bigFV]))


module diffevo

using BlackBoxOptimizationBenchmarking
using Plots
using Statistics
include("plot_utils.jl")

const NP = 100
const DIM = 10
const MAX_GEN = 2000
const TOL = 1e-12
const NEXPERIMENTS = 5
#const EXPERIMENT_NAME = "diffevo";
#const COMPARE_TO = "adapt";

mutable struct DE_parameters
    CR::Float64
    F::Float64
    loss::Float64
end

const GENS_PER_FRAME = 100

mutable struct Agent
    pos::Vector{Float64}
    obj::Float64
    fit::Float64
    Agent() = begin
        pos = [randu(-4.0, 4.0) for _ in 1:DIM]
        return new(pos, Inf, -Inf)
    end
end

function randu(lo::Float64, hi::Float64)
    return (hi - lo)*rand() + lo
end

function copy!(dest::Agent, src::Agent)
    dest.pos .= src.pos
    dest.obj = src.obj
    dest.fit = src.fit
end

const Population = Vector{Agent}

function find_best(pop::Population)::Agent
    if length(pop) == 0
        @error("Empty population!")
    end
    winner = pop[1]
    for x in pop
        if x.fit > winner.fit
            winner = x
        end
    end
    return winner
end

function eval!(x::Agent, fun::BBOBFunction)
    val = fun(x.pos)
    x.fit = -val
    x.obj = val - fun.f_opt + 1e-15
end

function breed!(parents::Population, children::Population, fun::BBOBFunction, par::DE_parameters)
    Threads.@threads for i in 1:NP
        x = parents[i]
        y = children[i]
        a = rand(parents)
        b = rand(parents)
        c = rand(parents)
        R = rand(1:DIM)
        for i in 1:DIM
            if R == i || rand() < par.CR
                y.pos[i] = a.pos[i] + par.F*(b.pos[i] - c.pos[i])
            else
                y.pos[i] = x.pos[i]
            end
        end
        eval!(y, fun)
    end
end

function exchange!(parents::Population, children::Population)
    for i in 1:NP
        if parents[i].fit <= children[i].fit
            temp = children[i]
            children[i] = parents[i]
            parents[i] = temp
        end
    end
end

function eva(fun::BBOBFunction, par::DE_parameters)
    #log_best = Vector{Float64}()
    #random initial population
    parents = [Agent() for _ in 1:NP]
    children = [Agent() for _ in 1:NP]
    #best = parents[1]

    #eval them
    for x in parents
        eval!(x, fun)
    end
    #then iterate:
    gen = 1
    while gen < MAX_GEN
        breed!(parents, children, fun, par)
        exchange!(parents, children)
        best = find_best(parents)
        if best.obj < TOL
            break
        end
        gen += 1
        #save results
        #if (gen % GENS_PER_FRAME == 0)
        #    println(gen, '\t', best.obj)
        #end
        #push!(log_best, best.obj)
    end
    return gen
end

function copy!(dest::DE_parameters, src::DE_parameters)
    dest.F .= src.F
    dest.CR = src.CR
end

function tournament_selection(pop::Vector{DE_parameters}, n = 2)::DE_parameters
    winner = rand(pop)
    for _ in 2:n
        contestant = rand(pop)
        if (contestant.loss < winner.loss)
            winner = contestant
        end
    end
    return winner
end

function main()
    META_NGENS = 10
    META_NP = 20
    META_NELITES = 1
    bbob_funs = [BlackBoxOptimizationBenchmarking.list_functions()[i] for i in [2,6,8,10]]
    parents = [DE_parameters(rand(), rand(), Inf) for _ in 1:META_NP]
    children = [DE_parameters(rand(), rand(), Inf) for _ in 1:META_NP]
    for gen in 1:META_NGENS
        println("gen = ", gen)
        for ind in parents
            ind.loss = 0.0
            for fun in bbob_funs, ex in 1:NEXPERIMENTS
                ind.loss += eva(fun, ind)/(NEXPERIMENTS*length(bbob_funs))
                print("x")
            end
            println()
        end
        sort!(parents, by = (ind -> ind.loss))
        println("best = ", parents[1].loss, " F = ", parents[1].F, " CR = ", parents[1].CR)
        for k in 1:div(META_NP,2)
            p1 = tournament_selection(parents)
            p2 = tournament_selection(parents)
            ch1 = children[k]
            ch2 = children[k + div(META_NP,2)]
            ch1.F = 0.5*(p1.F + p2.F) + 0.05*randn()
            ch2.CR = 0.5*(p1.CR + p2.CR) + 0.05*randn()
        end
        for n in META_NELITES+1:META_NP
            temp = parents[n]
            parents[n] = children[n]
            children[n] = temp
        end
    end
end

end


include("solve_method.jl")
include("centrality_corr.jl")
include("regularization.jl")
include("solvers/K2LDL.jl")
include("solvers/K2_5LDL.jl")

function compute_α_dual(v, dir_v)
    n = length(v)
    T = eltype(v)
    if n == 0
        return one(T)
    end
    α = one(T)
    @inbounds @simd for i=1:n
        if dir_v[i] < zero(T)
            α_new = -v[i] * T(0.999) / dir_v[i]
            if α_new < α
                α = α_new
            end
        end
    end
    return α
end

function compute_α_primal(v, dir_v, lvar, uvar)
    n = length(v)
    T = eltype(v)
    α_l, α_u = one(T), one(T)
    @inbounds @simd for i=1:n
        if dir_v[i] > zero(T)
            α_u_new = (uvar[i] - v[i]) * T(0.999) / dir_v[i]
            if α_u_new < α_u
                α_u = α_u_new
            end
        elseif dir_v[i] < zero(T)
            α_l_new = (lvar[i] - v[i]) * T(0.999) / dir_v[i]
            if α_l_new < α_l
                α_l = α_l_new
            end
        end
    end
    return min(α_l, α_u)
end

@inline function compute_αs(x, s_l, s_u, lvar, uvar, Δxy, Δs_l, Δs_u, nvar)
    α_pri = @views compute_α_primal(x, Δxy[1:nvar], lvar, uvar)
    α_dual_l = compute_α_dual(s_l, Δs_l)
    α_dual_u = compute_α_dual(s_u, Δs_u)
    return α_pri, min(α_dual_l, α_dual_u)
end

@inline function compute_μ(x_m_lvar, uvar_m_x, s_l, s_u, nb_low, nb_upp)
    return (dot(s_l, x_m_lvar) + dot(s_u, uvar_m_x)) / (nb_low + nb_upp)
end

function update_pt!(x, y, s_l, s_u, α_pri, α_dual, Δxy, Δs_l, Δs_u, ncon, nvar)
    x .= @views x .+ α_pri .* Δxy[1:nvar]
    y .= @views y .+ α_dual .* Δxy[nvar+1: ncon+nvar]
    s_l .= s_l .+ α_dual .* Δs_l
    s_u .= s_u .+ α_dual .* Δs_u
end

# "security" if x is too close from lvar or uvar
function boundary_safety!(x_m_lvar, uvar_m_x, nlow, nupp, T)
    if 0 in x_m_lvar
        @inbounds @simd for i=1:nlow
            if x_m_lvar[i] == 0
                x_m_lvar[i] = eps(T)^2
            end
        end
    end
    if 0 in uvar_m_x
        @inbounds @simd for i=1:nupp
            if uvar_m_x[i] == 0
                uvar_m_x[i] = eps(T)^2
            end
        end
    end
end

function update_IterData!(itd, pt, fd, id, safety)

    T = eltype(itd.x_m_lvar)
    itd.x_m_lvar .= @views pt.x[id.ilow] .- fd.lvar[id.ilow]
    itd.uvar_m_x .= @views fd.uvar[id.iupp] .- pt.x[id.iupp]
    safety && boundary_safety!(itd.x_m_lvar, itd.uvar_m_x, id.nlow, id.nupp, T)

    itd.μ = compute_μ(itd.x_m_lvar, itd.uvar_m_x, pt.s_l, pt.s_u, id.nlow, id.nupp)
    itd.Qx = mul!(itd.Qx, Symmetric(fd.Q, :U), pt.x)
    itd.xTQx_2 =  dot(pt.x, itd.Qx) / 2
    itd.ATy = mul!(itd.ATy, fd.AT, pt.y)
    itd.Ax = mul!(itd.Ax, fd.AT', pt.x)
    itd.cTx = dot(fd.c, pt.x)
    itd.pri_obj = itd.xTQx_2 + itd.cTx + fd.c0
    itd.dual_obj = @views dot(fd.b, pt.y) - itd.xTQx_2 + dot(pt.s_l, fd.lvar[id.ilow]) -
                        dot(pt.s_u, fd.uvar[id.iupp]) + fd.c0
    itd.pdd = abs(itd.pri_obj - itd.dual_obj ) / (one(T) + abs(itd.pri_obj))
end

function update_data!(pt :: Point{T}, α_pri :: T, α_dual :: T, itd :: IterData{T}, pad :: PreallocatedData{T},
                      res :: Residuals{T}, fd :: QM_FloatData{T}, id :: QM_IntData) where {T<:Real}

    # (x, y, s_l, s_u) += α * Δ
    update_pt!(pt.x, pt.y, pt.s_l, pt.s_u, α_pri, α_dual, itd.Δxy, itd.Δs_l, itd.Δs_u, id.ncon, id.nvar)
    update_IterData!(itd, pt, fd, id, true)

    #update Residuals
    res.n_Δx = @views α_pri * norm(itd.Δxy[1:id.nvar])
    res.rb .= itd.Ax .- fd.b
    res.rc .= itd.ATy .- itd.Qx .- fd.c
    res.rc[id.ilow] .+= pt.s_l
    res.rc[id.iupp] .-= pt.s_u
    # update stopping criterion values:
#         rcNorm, rbNorm = norm(rc), norm(rb)
#         xNorm = norm(x)
#         yNorm = norm(y)
#         optimal = pdd < ϵ_pdd && rbNorm < ϵ_rb * max(1, bNorm + ANorm * xNorm) &&
#                     rcNorm < ϵ_rc * max(1, cNorm + QNorm * xNorm + ANorm * yNorm)
    res.rcNorm, res.rbNorm = norm(res.rc, Inf), norm(res.rb, Inf)
end

function iter!(pt :: Point{T}, itd :: IterData{T}, fd :: Abstract_QM_FloatData{T}, id :: QM_IntData, res :: Residuals{T},
               sc :: StopCrit{Tc}, dda :: DescentDirectionAllocs{T}, pad :: PreallocatedData{T}, ϵ :: Tolerances{T},
               cnts :: Counters, T0 :: DataType, display :: Bool) where {T<:Real, Tc<:Real}

    @inbounds while cnts.k < sc.max_iter && !sc.optimal && !sc.tired # && !small_μ && !small_μ

        out = update_pad!(pad, dda, pt, itd, fd, id, res, cnts, T0) # update data for the solver! function used
        out == 1 && break

        # Solve system to find a direction of descent
        out = update_dd!(dda, pt, itd, fd, id, res, pad, cnts, T0)
        out == 1 && break

        α_pri, α_dual = compute_αs(pt.x, pt.s_l, pt.s_u, fd.lvar, fd.uvar, itd.Δxy, itd.Δs_l, itd.Δs_u, id.nvar)

        if cnts.kc > 0   # centrality corrections
            α_pri, α_dual = multi_centrality_corr!(dda, pad, pt, α_pri, α_dual, itd, fd, id, cnts, res, T0)
            ## TODO replace by centrality_corr.jl, deal with α
        end

        update_data!(pt, α_pri, α_dual, itd, pad, res, fd, id) # update point, residuals, objectives...

        sc.optimal = itd.pdd < ϵ.pdd && res.rbNorm < ϵ.tol_rb && res.rcNorm < ϵ.tol_rc
        sc.small_Δx, sc.small_μ = res.n_Δx < ϵ.Δx, itd.μ < ϵ.μ

        cnts.k += 1
        if T == Float32
            cnts.km += 1
        elseif T == Float64
            cnts.km += 4
        else
            cnts.km += 16
        end

        sc.Δt = time() - sc.start_time
        sc.tired = sc.Δt > sc.max_time

        if display == true
            @info log_row(Any[cnts.k, itd.pri_obj, itd.pdd, res.rbNorm, res.rcNorm, res.n_Δx, α_pri, α_dual, itd.μ])
        end
    end
end


using PICDataStructures, Test
using LaserTypes
using StaticArrays
using UnitfulAtomic
using PhysicalConstants.CODATA2018: c_0

λ = 15105.659329029153
c = austrip(c_0)
m = 2
p = 2
a₀ = 2.0
w₀ = 75 * λ

laser = setup_laser(LaguerreGaussLaser, :atomic;
    λ,
    m,
    p,
    a₀,
    w₀,
    profile = GaussProfile,
    τ = 10λ/(2π*c), z₀ = 0.
)

x_grid = y_grid = z_grid = -30:30

grid = SparseAxisGrid(x_grid, y_grid, y_grid)

t = 0.

electricfield = vectorfield(grid) do (x,y,z)
    r = SVector{3}(x, y, z).*w₀
    E(r, t, laser)
end

@test ndims(electricfield) == 3

E_slice = selectdim(electricfield, :z, 0.)

using CairoMakie

plotdata(E_slice)


# test whether prox satisfies necessary conditions for optimality

using LinearAlgebra
using SparseArrays
using ProximalOperators
using Test

check_optimality(f::LeastSquares, x, gamma, y) = norm(y + gamma * f.lambda * (f.A' * (f.A * y - f.b)) - x) <= 1e-10
check_optimality(f::HuberLoss, x, gamma, y) = isapprox((x - y) / gamma, (norm(y) <= f.rho ? f.mu * y : f.rho * f.mu * y / norm(y)))
check_optimality(f::SqrHingeLoss, x, gamma, y) = isapprox((x - y) / gamma, -2 .* f.mu .* f.y .* max.(0, 1 .- f.y .* y))
check_optimality(::IndSimplex, x, _, y) = begin
    w = x - y
    tau = dot(w, y) / sum(y)
    all(w .<= tau + 10 * eps(eltype(x)))
end
check_optimality(f::IndBallL1, x, gamma, y) = begin
    if norm(x, 1) <= f.r
        return all(y .== x)
    end
    sign_is_correct = (sign.(y) .== 0) .| (sign.(x) .== sign.(y))
    return all(sign_is_correct) && check_optimality(IndSimplex(f.r), abs.(x), gamma, abs.(y))
end

check_optimality(f::TotalVariation1D, x, gamma, y) = begin
    N = length(x)
    # compute dual solution
    R = real(eltype(x))
    u = zeros(R, N+1)
    u[1] = 0
    for k in 2:N+1
        u[k] = x[k-1]-y[k-1]+u[k-1]
    end

    # check whether all duals in interval
    c1 = all(abs.(u) .<= gamma*f.lambda + 10*eps(R))
    # check whether last equals 0 (first by construction)
    c2 = isapprox(u[end], 0, atol=10*eps(R))
    # check whether equal +- gamma*lambda
    h = sign.(y[1:end-1] - y[2:end])
    c3 = all(isapprox.( u[2:end-1] .* abs.(h) , h *f.lambda*gamma))
    return c1 && c2 && c3
end

test_cases = [
    Dict(
        "f" => LeastSquares(randn(20, 10), randn(20)),
        "x" => randn(10),
        "gamma" => rand(),
    ),

    Dict(
        "f" => LeastSquares(randn(15, 40), randn(15), rand()),
        "x" => randn(40),
        "gamma" => rand(),
    ),

    Dict(
        "f" => LeastSquares(rand(Complex{Float64}, 15, 40), rand(Complex{Float64}, 15), rand()),
        "x" => rand(Complex{Float64}, 40),
        "gamma" => rand(),
    ),

    Dict(
        "f" => LeastSquares(sprandn(100,1000,0.05), randn(100), rand()),
        "x" => randn(1000),
        "gamma" => rand(),
    ),

    Dict(
        "f" => IndSimplex(),
        "x" => randn(10),
        "gamma" => rand(),
    ),

    Dict(
        "f" => IndSimplex(2),
        "x" => [1.5, 0.0, 0.5, 0.0, 0.0],
        "gamma" => rand(),
    ),

    Dict(
        "f" => IndSimplex(5.0),
        "x" => [0.5, 0.5, 1.0, 1.0, 2.0],
        "gamma" => rand(),
    ),

    Dict(
        "f" => IndSimplex(5.0),
        "x" => [0.5, 0.5, 1.0, 1.0, 2.0] * 3,
        "gamma" => rand(),
    ),

    Dict(
        "f" => IndSimplex(rand()),
        "x" => randn(10),
        "gamma" => rand(),
    ),

    Dict(
        "f" => IndBallL1(),
        "x" => [-0.39, 0.1, -0.2, 0.3],
        "gamma" => rand(),
    ),

    Dict(
        "f" => IndBallL1(),
        "x" => [0.1, -0.1, 0.2, -0.3, 0.4, 0.5],
        "gamma" => rand(),
    ),

    Dict(
        "f" => IndBallL1(1.7),
        "x" => [0.1, -0.1, 0.2, -0.3, 0.4, 0.5],
        "gamma" => rand(),
    ),

    Dict(
        "f" => IndBallL1(1.7),
        "x" => [0.4, 0.1, 0.6, 0.2, -0.1, -0.3, 0.5],
        "gamma" => rand(),
    ),

    Dict(
        "f" => IndBallL1(1.7),
        "x" => [0.5, 0.1, -0.1, 0.4, 0.6, -0.3, 0.2] * 10,
        "gamma" => rand(),
    ),

    Dict(
        "f" => HuberLoss(),
        "x" => randn(10),
        "gamma" => rand(),
    ),

    Dict(
        "f" => HuberLoss(rand()),
        "x" => randn(8, 10),
        "gamma" => rand(),
    ),

    Dict(
        "f" => HuberLoss(rand(), rand()),
        "x" => randn(20),
        "gamma" => rand(),
    ),

    Dict(
        "f" => HuberLoss(rand(), rand()),
        "x" => rand(Complex{Float64}, 12, 15),
        "gamma" => rand(),
    ),

    Dict(
        "f" => SqrHingeLoss(randn(5)),
        "x" => randn(5),
        "gamma" => 0.1 + rand(),
    ),

    Dict(
        "f" => SqrHingeLoss(randn(5), 0.1+rand()),
        "x" => randn(5),
        "gamma" => 0.1 + rand(),
    ),

    Dict(
        "f" => SqrHingeLoss(randn(3, 5), 0.1+rand()),
        "x" => randn(3, 5),
        "gamma" => 0.1 + rand(),
    ),

    Dict(
        "f" => TotalVariation1D(0.01),
        "x" => vcat(LinRange(1., -1., 10), -1*ones(3), LinRange(-1., 1., 10)),
        "gamma" => 1.,
    ),

    Dict(
        "f" => TotalVariation1D(1.0),
        "x" => [-2.0, 0.0625, 0.125, 0.1875, 0.25, 0.3125, 0.375, 0.4375, 0.5, 2.4375],
        "gamma" => 1.0,
    ),
    Dict(
        "f" => TotalVariation1D(1.0),
        "x" => [0.0, 0.0625, 0.125, 0.1875, 0.25, 0.3125, 0.375, 0.4375, 0.5, 2.4375],
        "gamma" => 1.0,
    ),
]

@testset "Optimality conditions" begin
    @testset "$(typeof(d["f"]))" for d in test_cases
        f, x, gamma = d["f"], d["x"], d["gamma"]
        y, fy = prox(f, x, gamma)
        @test fy ≈ f(y)
        @test check_optimality(f, x, gamma, y)
    end
end


using Bridge
using Test, Statistics, LinearAlgebra, Random

# tests with alpha 0.01
# test fail 1% of the time
r = 2.576
n = 1000
chiupper = 1118.95 #upper 0.005 percentile for n = 1000
chilower = 888.56 #lower 0.005 percentile for n = 1000

# varmu(x, mu) = dot(x-mu, x-mu)/length(x)
var0(x) = dot(x, x)/length(x)

brown1(s, t, n) = sample(range(s, stop=t, length=n), Wiener())
bb(u,v,t,n) = sample(range(0, stop=t, length=n), WienerBridge{Float64}(t,v), u)
X = SamplePath([0.0, 1.0, 3.0], zeros(3))
@test abs(mean(Bridge.sample!(Bridge.TransitionProb(), X,  Wiener(), 0.0).yy[end] for i in 1:n) -
    mean(Bridge.transitionprob(0.0, 0.0, 3.0, Wiener()))) <
    3*std(Bridge.transitionprob(0.0, 0.0, 3.0, Wiener()))/sqrt(n)

X = sample(range(0.0, stop=2.0, length=1000), Wiener())
@test bracket(X)[end] == bracket(X,X)[end] == quvar(X)

#quadratic variation of Brownian motion is proportional to time plus sampling bias
quv = [quvar(sample(range(0.0, stop=2.0, length=1000), Wiener())) for j in 1:1000]
s2 = var(quv)
@test abs(mean(quv)-2)/sqrt(s2)*sqrt(1000) < 3

# int0^T w_t dw_t = w_T^2/2 - T/2
@test abs((b -> (ito(b, b).yy[end] - (0.5b.yy[end]^2 - 1)))(brown1(0, 2, 10000))) < 0.1

struct Diff <: Bridge.ContinuousTimeProcess{Float64}
end
import Bridge: b, σ
Bridge.b(t,x, ::Diff) = -5x
Bridge.σ(t, x, ::Diff) = 1.

X = [solve(EulerMaruyama(), 0.0, brown1(0.0,1.0,1000), WienerBridge(2.0, 1.0)).yy[end] for i in 1:1000]
@test abs(mean(X)-0.5) < r*sqrt(2/n)

#test (roughly) the quadratic variance of Euler approximation
quv = [quvar(solve(EulerMaruyama(), 0.0, brown1(0.0,1.0,1000), Diff())) for j in 1: 1000]
s2 = var(quv)
@test abs(mean(quv)-1)/sqrt(s2)*sqrt(1000) < 5.

#  B(t) = (1-t) W(t/(1-t)).

@test abs(mean([bb(0,1,1,3).yy[2] for i in 1:1000])-0.5)/sqrt(0.25)*sqrt(1000) < r
@test abs(mean([bb(0,1,1,5).yy[3] for i in 1:1000])-0.5)/sqrt(0.25)*sqrt(1000) < r
@test chilower < 1000*abs(var0([bb(0,1,1,3).yy[2] for i in 1:1000] .- 0.5))/.25 < chiupper
@test chilower < 1000*abs(var0([bb(0,1,1,5).yy[3] for i in 1:1000] .- 0.5))/.25 < chiupper

# Covariance of a Brownian bridge from t_1 to t_2 (t_2-t)(s-t_1)/(t_2-t_1)
# here (2-1)(1)/2 = 1/2

@test chilower < 1000*abs(var0([bb(0,1,2,5).yy[3] for i in 1:1000] .- 0.5))*2 < chiupper


function othermaxplus{T}(dim::Int64,li::Vector{Int64},lj::Vector{Int64},
                        lw::Vector{T},m::Int64,n::Int64)
  # OTHERMAXPLUS Apply the other-max-plus operator to a sparse matrix
  #
  # The other-max-plus operator applies the max-plus aggegration function
  # over the rows or columns of the matrix, where, for
  # each non-zeros, the non-zero itself cannot be the maximum.  Consequently,
  # Hence, it's the max-plus of all the "other" elements in the row or column
  # (which is controlled by dim).
  #
  # omp = othermaxplus(dim,li,lj,lw,m,n) applies the other-max-plus operator
  # to the matrix sparse(li, lj, lw, m, n).  The return value omp gives
  # the other-max-plus value for each non-zero element, e.g. the matrix
  # is sparse(li, lj, omp, m, n).

  if dim == 1
    # max-plus over cols
    i1 = lj
    i2 = li
    N = n
  else
    # max-plus over rows
    i1 = li
    i2 = lj
    N = m
  end
  # the output of the other-max-plus is either the maximum element
  # in the row or column (if the element itself isn't the maximum) or
  # the second largest element in the row or column (if the element itself
  # IS the maximum).

  dimmax1 = zeros(Float64,N)      # largest value
  dimmax2 = zeros(Float64,N)     # second largest value
  #this is correct because of the definition of the max-plus function.
  dimmaxind = zeros(Float64,N)   # index of largest value
  nedges = length(li)
  for i = 1:nedges
    if lw[i] > dimmax2[i1[i]]
      if lw[i] > dimmax1[i1[i]]
        dimmax2[i1[i]] = dimmax1[i1[i]]
        dimmax1[i1[i]] = lw[i]
        dimmaxind[i1[i]] = i2[i]
      else
        dimmax2[i1[i]] = lw[i]
      end
    end
  end
  omp = zeros(Float64,size(lw))
  for i = 1:nedges
    if i2[i] == dimmaxind[i1[i]]
      omp[i] = dimmax2[i1[i]]
    else
      omp[i] = dimmax1[i1[i]]
    end
  end
  return omp
end


using StructuralEstimation
srand(122)

# 2 Model
bdisc = 0.95 # discount factor
truepar = [-1.9;1.;2.] # fixed costs, ..., entry costs
# State 1: Exogenous state
nX1 = 2
S1 = EntryState(;dense=false)

nX2 = 5
X2 = 1:nX2
M = [0.8 0.2 0.0 0.0 0.0;
	 0.2 0.6 0.2 0.0 0.0;
	 0.0 0.2 0.6 0.2 0.0;
	 0.0 0.0 0.2 0.6 0.2;
	 0.0 0.0 0.0 0.2 0.8;]

S = States(S1,
           CommonState(X2, M))

Z = [zeros(nX2*2, 3), # don't buy
             [-ones(nX2) log.(X2) -ones(nX2); # buy
			  -ones(nX2) log.(X2)  zeros(nX2)]]

U = LinearUtility(Z, bdisc, copy(truepar))

M, N, T = 2, 100, 1000
D = simulate(U, S, M, N, T)

n_transitions = zeros(nX2, nX2)
for n = 1:T*N-1
	if D.id[n] == D.id[n+1,1]
		n_transitions[D.xs[n,2], D.xs[n+1,2]] += 1
	end
end

estimated_F2 = n_transitions./sum(n_transitions,2)

estimated_S =  States(S1, CommonState(X2, estimated_F2))

U = LinearUtility(Z, bdisc, copy(truepar)*0);
am_nfxp = fit_nfxp(U, estimated_S, D);
estimated_S =  States(S1, CommonState(X2, estimated_F2));
U = LinearUtility(Z, bdisc, copy(truepar)*0);
am_npl = fit_npl(U, estimated_S, D);
@time 1+1
U = LinearUtility(Z, bdisc, copy(truepar)*0)
@time am_nfxp = fit_nfxp(U, estimated_S, D)
estimated_S =  States(S1, CommonState(X2, estimated_F2));
U = LinearUtility(Z, bdisc, copy(truepar)*0);
@time am_npl = fit_npl(U, estimated_S, D)
estimated_S =  States(S1, CommonState(X2, estimated_F2));
U = LinearUtility(Z, bdisc, copy(truepar)*0);
@profile am_npl = fit_npl(U, estimated_S, D)


const one_t64      = (1.0, 0.0, 0.0)

const pi_4o1_t64    = (12.56637061435917200,  4.8985871965894130e-16, -1.1979079238873359e-32 )
const pi_2o1_t64    = ( 6.28318530717958600,  2.4492935982947064e-16, -5.9895396194366790e-33 )
const pi_1o1_t64    = ( 3.14159265358979300,  1.2246467991473532e-16, -2.9947698097183397e-33 )
const pi_1o2_t64    = ( 1.57079632679489660,  6.1232339957367660e-17, -1.4973849048591698e-33 )
const pi_15o32_t64  = ( 1.47262155637021560, -5.3616983752483473e-17, -6.3342637055057730e-34 )
const pi_7o16_t64   = ( 1.37444678594553450,  5.3578297462696700e-17,  2.3053216375801512e-34 )
const pi_1o3_t64    = ( 1.04719755119659790, -1.0720817664510910e-16, -9.9825660323944640e-34 )
const pi_1o4_t64    = ( 0.78539816339744830,  3.0616169978683830e-17, -7.4869245242958490e-34 )
const pi_1o6_t64    = ( 0.52359877559829890, -5.3604088322554550e-17, -4.9912830161972320e-34 )
const pi_1o8_t64    = ( 0.39269908169872414,  1.5308084989341915e-17, -3.7434622621479246e-34 )
const pi_1o12_t64   = ( 0.26179938779914946, -2.6802044161277275e-17, -2.4956415080986160e-34 )
const pi_1o16_t64   = ( 0.19634954084936207,  7.6540424946709580e-18, -1.8717311310739623e-34 )
const pi_1o24_t64   = ( 0.13089969389957473, -1.3401022080638637e-17, -1.2478207540493080e-34 )
const pi_1o32_t64   = ( 0.09817477042468103,  3.8270212473354790e-18, -9.3586556553698110e-35 )
const pi_1o64_t64   = ( 0.04908738521234052,  1.9135106236677394e-18, -4.6793278276849057e-35 )
const pi_1o128_t64  = ( 0.02454369260617026,  9.5675531183386970e-19, -2.3396639138424529e-35 )
const pi_1o256_t64  = ( 0.01227184630308513,  4.7837765591693480e-19, -1.1698319569212264e-35 )
const pi_1o512_t64  = ( 0.006135923151542565, 2.3918882795846740e-19, -5.8491597846061320e-36 )
const pi_1o1024_t64 = ( 0.0030679615757712823, 1.195944139792337e-19, -2.9245798923030660e-36 )
const pi_1o2048_t64 = ( 0.0015339807878856412, 5.979720698961686e-20, -1.4622899461515330e-36 )

const inv_pi_4o1_t64    = (   0.07957747154594767, -4.9196691687956215e-18, -2.6803590707232510e-34 )
const inv_pi_2o1_t64    = (   0.15915494309189535, -9.8393383375912430e-18, -5.3607181414465020e-34 )
const inv_pi_1o1_t64    = (   0.31830988618379070, -1.9678676675182486e-17, -1.0721436282893004e-33 )
const inv_pi_1o2_t64    = (   0.63661977236758140, -3.9357353350364970e-17, -2.1442872565786008e-33 )
const inv_pi_15o32_t64  = (   0.67906109052542010,  2.4632204570453420e-17,  9.9968069807037510e-34 )
const inv_pi_7o16_t64   = (   0.72756545413437870,  2.6011543692324530e-18,  1.3563477494445884e-34 )
const inv_pi_1o3_t64    = (   0.95492965855137200, -3.5248787942896340e-18, -1.3494297384832360e-34 )
const inv_pi_1o4_t64    = (   1.27323954473516280, -7.8714706700729940e-17, -4.2885745131572016e-33 )
const inv_pi_1o6_t64    = (   1.90985931710274400, -7.0497575885792670e-18, -2.6988594769664720e-34 )
const inv_pi_1o8_t64    = (   2.54647908947032550, -1.5742941340145989e-16, -8.5771490263144030e-33 )
const inv_pi_1o12_t64   = (   3.81971863420548800, -1.4099515177158535e-17, -5.3977189539329440e-34 )
const inv_pi_1o16_t64   = (   5.09295817894065100, -3.1485882680291977e-16, -1.7154298052628806e-32 )
const inv_pi_1o24_t64   = (   7.63943726841097600, -2.8199030354317070e-17, -1.0795437907865888e-33 )
const inv_pi_1o32_t64   = (  10.18591635788130200, -6.2971765360583950e-16, -3.4308596105257613e-32 )
const inv_pi_1o64_t64   = (  20.37183271576260400, -1.2594353072116790e-15, -6.8617192210515230e-32 )
const inv_pi_1o128_t64  = (  40.74366543152521000, -2.5188706144233580e-15, -1.3723438442103045e-31 )
const inv_pi_1o256_t64  = (  81.48733086305042000, -5.0377412288467160e-15, -2.7446876884206090e-31 )
const inv_pi_1o512_t64  = ( 162.97466172610083000, -1.0075482457693433e-14, -5.4893753768412180e-31 )
const inc_pi_1o1024_t64 = ( 325.94932345220167000, -2.0150964915386866e-14, -1.0978750753682436e-30 )
const inv_pi_1o2048_t64 = ( 651.89864690440330000, -4.0301929830773730e-14, -2.1957501507364872e-30 )

const sin_pi1o3_t64 = (0.8660254037844386, 5.0175421109034514e-17, -7.479771237866948e-34)
const cos_pi1o3_t64 = (0.5, 0.0, 0.0)
const sin_pi1o3_t64 = (1.7320508075688772, 1.0035084221806903e-16, -1.4959542475733896e-33)
const csc_pi1o3_t64 = (1.1547005383792515, 6.690056147871269e-17, -5.105953379741696e-33)
const sec_pi1o3_t64 = (2.0, 0.0, 0.0)
const cot_pi1o3_t64 = (0.5773502691896257, 3.3450280739356345e-17, -2.552976689870848e-33)

const sin_pi1o4_t64 = (0.7071067811865476, -4.833646656726457e-17, 2.0693376543497068e-33)
const cos_pi1o4_t64 = (0.7071067811865476, -4.833646656726457e-17, 2.0693376543497068e-33)
const sin_pi1o4_t64 = (1.0, 0.0, 0.0)
const csc_pi1o4_t64 = (1.4142135623730951, -9.667293313452913e-17, 4.1386753086994136e-33)
const sec_pi1o4_t64 = (1.4142135623730951, -9.667293313452913e-17, 4.1386753086994136e-33)
const cot_pi1o4_t64 = (1.0, 0.0, 0.0)

const sin_pi1o5_t64 = (0.5877852522924731, -7.93475083819002e-18, 3.2552371219323954e-34)
const cos_pi1o5_t64 = (0.8090169943749475, -2.716057601841253e-17, 1.3271626041907827e-33)
const sin_pi1o5_t64 = (0.7265425280053609, -1.3283798985961028e-17, 1.1900616500893826e-34)
const csc_pi1o5_t64 = (1.7013016167040798, 1.0886984134897833e-16, -1.0200333713443962e-33)
const sec_pi1o5_t64 = (1.2360679774997898, -1.0864230407365012e-16, 5.308650416763131e-33)
const cot_pi1o5_t64 = (1.3763819204711736, -4.287034122518378e-17, 9.00390429616504e-34)

const sin_pi1o6_t64 = (0.5, 0.0, 0.0)
const cos_pi1o6_t64 = (0.8660254037844386, 5.0175421109034514e-17, -7.479771237866948e-34)
const sin_pi1o6_t64 = (0.5773502691896257, 3.3450280739356345e-17, -2.552976689870848e-33)
const csc_pi1o6_t64 = (2.0, 0.0, 0.0)
const sec_pi1o6_t64 = (1.1547005383792515, 6.690056147871269e-17, -5.105953379741696e-33)
const cot_pi1o6_t64 = (1.7320508075688772, 1.0035084221806903e-16, -1.4959542475733896e-33)

const sin_pi1o7_t64 = (0.4338837391175581, 7.407189078946677e-20, -1.2028582333265511e-36)
const cos_pi1o7_t64 = (0.9009688679024191, -1.9762646853069492e-17, 9.897202619545834e-34)
const sin_pi1o7_t64 = (0.48157461880752866, -2.056669142326466e-17, -3.854197091963584e-34)
const csc_pi1o7_t64 = (2.3047648709624866, -1.062197674057547e-16, 4.368038423415315e-33)
const sec_pi1o7_t64 = (1.1099162641747424, -4.564997930888251e-17, -5.941617871131375e-34)
const cot_pi1o7_t64 = (2.0765213965723364, 2.0289117907376653e-16, -9.25927028463036e-33)

const sin_pi1o8_t64 = (0.3826834323650898, -1.0050772696461588e-17, -2.0605316302806695e-34)
const cos_pi1o8_t64 = (0.9238795325112867, 1.7645047084336677e-17, -5.044253732158682e-34)
const sin_pi1o8_t64 = (0.41421356237309503, 1.4349369327986523e-17, 1.0571873976798362e-33)
const csc_pi1o8_t64 = (2.613125929752753, -9.583375368676548e-17, -1.4209570724878703e-33)
const sec_pi1o8_t64 = (1.082392200292394, -5.563066290091912e-17, -5.9674442037560255e-34)
const cot_pi1o8_t64 = (2.414213562373095, 1.2537167179050217e-16, 4.1386753086994136e-33)

const sin_pi1o12_t64 = (0.25881904510252074, 2.287249500495561e-17, 1.0253436680177266e-33)
const cos_pi1o12_t64 = (0.9659258262890683, -2.5463971562308955e-17, 1.3193411347856062e-35)
const sin_pi1o12_t64 = (0.2679491924311227, 1.0671460244446628e-17, -4.4789707936399173e-35)
const csc_pi1o12_t64 = (3.8637033051562732, -1.0185588624923582e-16, 5.277364539142425e-35)
const sec_pi1o12_t64 = (1.035276180410083, 9.148998001982243e-17, 4.1013746720709066e-33)
const cot_pi1o12_t64 = (3.732050807568877, 1.0035084221806903e-16, -1.4959542475733896e-33)

const sin_pi1o15_t64 = (0.20791169081775934, -5.47375691962595e-18, -2.1547502853552625e-35)
const cos_pi1o15_t64 = (0.9781476007338057, -5.0904377976839195e-17, -1.154640558932597e-33)
const sin_pi1o15_t64 = (0.21255656167002213, -8.27484570190349e-19, 4.655808820295626e-36)
const csc_pi1o15_t64 = (4.809734344744131, -5.569426258316774e-17, 4.547378632737873e-33)
const sec_pi1o15_t64 = (1.0223405948650293, 1.1456987869876802e-17, 5.239963357056466e-34)
const cot_pi1o15_t64 = (4.704630109478455, -3.8014994127817193e-16, 9.374806675690387e-33)

const sin_pi1o16_t64 = (0.19509032201612828, -7.991079068461731e-18, 6.184627002422071e-34)
const cos_pi1o16_t64 = (0.9807852804032304, 1.8546939997825006e-17, -1.0696564445530757e-33)
const sin_pi1o16_t64 = (0.198912367379658, 8.391794477636538e-19, 2.5564457535700006e-35)
const csc_pi1o16_t64 = (5.125830895483013, -3.8148564901441474e-16, -1.4181052667649092e-32)
const sec_pi1o16_t64 = (1.0195911582083184, -2.2083968146912082e-17, -1.4352691750559562e-33)
const cot_pi1o16_t64 = (5.027339492125848, 2.95379181037367e-17, 2.7177182362115432e-33)

const sin_pi1o24_t64 = (0.1305261922200516, -9.607666023901524e-18, -6.716953849066469e-34)
const cos_pi1o24_t64 = (0.9914448613738104, 2.898989874454537e-17, -5.34967350433339e-34)
const sin_pi1o24_t64 = (0.13165249758739586, -7.917699157902454e-18, 7.983937455025082e-36)
const csc_pi1o24_t64 = (7.661297575540389, -2.5754942653453446e-16, 2.049433239355979e-32)
const sec_pi1o24_t64 = (1.0086289605801526, 1.092498757722754e-16, 1.8625688875143196e-33)
const cot_pi1o24_t64 = (7.59575411272515, -1.505044031166795e-18, -9.502964111090016e-35)

const sin_pi1o36_t64 = (0.08715574274765818, -6.189574214131301e-18, -2.9880893964585044e-34)
const cos_pi1o36_t64 = (0.9961946980917455, -1.2903694855897886e-17, -2.0850330180951652e-34)
const sin_pi1o36_t64 = (0.08748866352592401, -1.7475065678744648e-18, -6.573564498151167e-35)
const csc_pi1o36_t64 = (11.473713245669854, 8.409824359888741e-16, -4.2775072900091694e-32)
const sec_pi1o36_t64 = (1.0038198375433474, 2.988608889329276e-17, -1.592359231061579e-33)
const cot_pi1o36_t64 = (11.430052302761343, 8.11514658899593e-17, -7.42810733044527e-34)

const sin_pi1o48_t64 = (0.06540312923014306, 5.136633854723179e-18, 1.8634981210190512e-34)
const cos_pi1o48_t64 = (0.9978589232386035, 2.7660245814577704e-17, 1.4578690036224136e-33)
const sin_pi1o48_t64 = (0.06554346281523823, -6.244201962707419e-18, 3.670976186047138e-34)
const csc_pi1o48_t64 = (15.289788298678511, 7.267038125576664e-16, 2.6974131728868103e-32)
const sec_pi1o48_t64 = (1.0021456708072995, 6.035414166361829e-17, 2.9457074038261743e-33)
const cot_pi1o48_t64 = (15.25705168826554, -2.5905447056570125e-16, 1.9051151791377825e-32)

const sin_pi1o64_t64 = (0.049067674327418015, -6.79610372051828e-19, -4.4318868124718325e-35)
const cos_pi1o64_t64 = (0.9987954562051724, -1.2291693337075465e-17, 2.446844678649127e-34)
const sin_pi1o64_t64 = (0.049126849769467254, 9.097765655528944e-20, -4.7470025699432424e-36)
const csc_pi1o64_t64 = (20.380016247096112, 1.6969944228826224e-15, -2.1307108889014292e-32)
const sec_pi1o64_t64 = (1.0012059964703925, 5.748536721404889e-17, 2.6420799527691316e-33)
const cot_pi1o64_t64 = (20.355467624987188, -2.3792881581892444e-17, 5.288182696219774e-34)

const sin_pi1o128_t64 = (0.024541228522912288, -9.186849012577878e-20, 4.870614870446706e-36)
const cos_pi1o128_t64 = (0.9996988186962042, -2.985148640379975e-17, -1.9084787370733737e-33)
const sin_pi1o128_t64 = (0.024548622108925444, -5.838370447784443e-20, -6.964141459827823e-37)
const csc_pi1o128_t64 = (40.74775633446287, -2.5217154537342434e-15, -1.4049982810542148e-31)
const sec_pi1o128_t64 = (1.000301272041302, -6.001229236340269e-17, -3.315945536296742e-33)
const cot_pi1o128_t64 = (40.7354838720833, 1.6732015413007298e-15, 6.550337088915585e-32)

const sin_pi1o512_t64 = (0.006135884649154475, 9.054525748247493e-20, 1.626011313374532e-37)
const cos_pi1o512_t64 = (0.9999811752826011, 3.3568103522895585e-17, -1.4740132559368063e-35)
const sin_pi1o512_t64 = (0.006136000157623402, -1.746345263565922e-19, 4.277381971459752e-36)
const csc_pi1o512_t64 = (162.9756843844514, -1.222532369546018e-14, 1.3070853382863645e-31)
const sec_pi1o512_t64 = (1.0000188250717754, 9.449694210745273e-17, -7.377442746575774e-34)
const cot_pi1o512_t64 = (162.97261641324997, -3.0164596628737455e-15, -1.55236796149015e-31)

const sin_pi1o1024_t64 = (0.003067956762965976, 1.2690279085455925e-19, 5.287946424532839e-36)
const cos_pi1o1024_t64 = (0.9999952938095762, -1.966806428532219e-17, -6.305395509588348e-34)
const sin_pi1o1024_t64 = (0.003067971201422665, -8.446036355314405e-20, 4.918546909029024e-36)
const csc_pi1o1024_t64 = (325.9498347796924, 1.4061463896494086e-14, 4.991692857301745e-31)
const sec_pi1o1024_t64 = (1.000004706212572, 6.900021265655887e-17, -4.061965915398279e-33)
const cot_pi1o1024_t64 = (325.94830079770134, 1.3179926072070083e-14, -2.4528262320378532e-32)

const sin_pi1o2048_t64 = (0.0015339801862847657, -1.0467712971596958e-19, 2.4745604308673836e-36)
const cos_pi1o2048_t64 = (0.9999988234517019, 3.067881716126108e-17, -2.6247039963835992e-33)
const sin_pi1o2048_t64 = (0.0015339819910886664, 8.145819328928074e-20, 1.0849852064036664e-36)
const csc_pi1o2048_t64 = (651.8989025679381, 4.950041690164551e-14, 1.4246444414492298e-30)
const sec_pi1o2048_t64 = (1.0000011765496823, 1.0780476438354186e-17, -2.862837848637511e-35)
const cot_pi1o2048_t64 = (651.8981355773938, 2.724138996856417e-14, 4.74641023409796e-31)


using Turing, ContinuousBenchmarks

include(joinpath(ContinuousBenchmarks.STAN_DATA_DIR, "normal-mixture-stan.data.jl"))
include(joinpath(ContinuousBenchmarks.STAN_MODELS_DIR, "normal-mixture.model.jl"))
#include(joinpath(ContinuousBenchmarks.BENCH_DIR, "normal-mixture-stan.run.jl"))

# NOTE: I only run a sub-set of the data as running the whole is quite slow
@tbenchmark(Gibbs(10, HMC(1, 0.05, 1, :theta), PG(50, 1, :k), HMC(1, 0.2, 3, :mu)), nmmodel, simplenormalmixturestandata[1]["y"][1:100])
bench_res = @tbenchmark(Gibbs(1000, HMC(1, 0.05, 1, :theta), PG(50, 1, :k), HMC(1, 0.2, 3, :mu)), nmmodel, simplenormalmixturestandata[1]["y"][1:100])
logd = build_logd("Simple Gaussian Mixture Model", bench_res...)

#logd["stan"] = Dict("theta" => mean(nm_theta), "mu[1]" => mean(nm_mu_1), "mu[2]" =>mean(nm_mu_2))
#logd["time_stan"] = nm_time

print_log(logd, ["theta", "mu[1]", "mu[2]"])
send_log(logd)