function mh(
    likelihood::Function,
    prior::Function,
    propdistpdf::Function,
    propdistsample::Function,
    startval::Vector,
    mcit::Int64,
    burn::Int64
    )

    sample = zeros(mcit,length(startval))
    sample[1,:] = startval

    for i in 1:mcit-1
        curr = sample[i,:]

        #create new sample
        proposalval = curr + propdistsample()

        #calculate acceptance probability
        A = (likelihood(proposalval)/likelihood(curr)*prior(proposalval)/
            prior(curr)*propdistpdf(proposalval-curr)/
            propdistpdf(curr-proposalval))[1]

        #accept or reject sample
        if A>=rand(Uniform(0,1))
            sample[i+1,:] = proposalval
        else
            sample[i+1,:] = curr
        end
    end
    return sample[burn+1:mcit,:]
end

function gibbssample(
    condist::Vector,
    startval::Vector,
    mcit::Int64,
    burn::Int64
    )

    sample = zeros(mcit, length(startval))
    sample[1,:] = startval

    for i in 1:mcit-1
        sample[i+1,:] = sample[i,:]
        for j in 1:size(sample, 2)
            #update each variable by itself
            sample[i+1,j] = condist[j](sample[i+1,:])
        end
    end
    return sample[burn+1:mcit,:]
end

function tresample(
    sset::Array,
    hatw::Array
    )
    snum = size(sset,1)
    rsset = zeros(snum, size(sset,2))
    hatwn = zeros(snum, 1)
    number = [1:snum...]

    for i in 1:snum
        #resample according to weight
        resample = StatsBase.sample(number, ProbabilityWeights(vec(hatw)))
        rsset[i,:] = sset[resample, :]
        hatwn[i] = hatw[resample]
    end
    return rsset, hatwn
end

function calcbeta(
    likelihood::Function,
    betaj::Float64,
    sset::Array,
    snum::Int64
    )

    betaj1 = betaj + (1.1 - betaj)*0.5
    likelihoodcov = zeros(snum, 1)
    step = (1.1 - betaj)*0.25
    cov = 0

    for i in 1:snum
        likelihoodcov[i] = likelihood(sset[i,:])
    end
    #beta approximates so that coefficient of variance --> 1
    while cov> 1.0001 || cov < 0.9999
        cov = std(likelihoodcov.^(betaj1-betaj))/mean(likelihoodcov.^(betaj1-betaj))
        if cov < 1
            betaj1 += step
            step = 0.5*step
        else
            betaj1 -= step
            step = 0.5*step
        end
        if betaj1>1
            break
        end
    end
    betaj1 = min(1, betaj1)
    return betaj1, likelihoodcov.^(betaj1-betaj)
end

function calcsigma(
    sset::Array,
    hatw::Array,
    gamma::Real
    )
    sdim = size(sset, 2)
    snum = size(sset, 1)

    thet = zeros(sdim, 1)
    sig = zeros(sdim, sdim)

    #formula for calculating Sigma in Ching and Chen 2007
    for i in 1:snum
        thet = thet + sset[i,:].*hatw[i]
    end

    for i in 1:snum
        sig = sig + ((sset[i,:]-thet)*(sset[i,:]-thet)').*hatw[i]
    end
    sig = sig.*gamma^2

    return sig
end

function tmcmc(
    likelihood::Function,
    prior::Function,
    priorsample::Function,
    snum::Int64,
    gamma::Real
    )

    j = 0
    betaj = 0.0
    sset = priorsample(snum)

    while betaj < 1
        betaj1, w = calcbeta(likelihood, betaj, sset, snum)
        hatw = w./sum(w)#normalize

        sset, w = tresample(sset, hatw)
        hatw = w./sum(w)

        sig = calcsigma(sset, hatw, gamma)

        for i in 1:snum
            #MH step for each sample
            mvpropdist = MvNormal(zeros(size(sset, 2)), sig)
            propdistsample() = rand(mvpropdist)
            propdistpdf(x) = pdf(mvpropdist, x)
            likelihoodbeta(x) = likelihood(x)^(betaj1)
            sset[i,:] = mh(likelihoodbeta, prior, propdistpdf, propdistsample, sset[i,:], 2,1)
        end
        betaj = betaj1
    end

    return sset
end

function resample(
    sset::Array,
    hatw::Array
    )
    snum = size(sset,1)
    rsset = zeros(snum, size(sset,2))
    number = [1:snum...]

    for i in 1:snum
        #weighted resampling
        resample = StatsBase.sample(number, ProbabilityWeights(vec(hatw)))
        rsset[i,:] = sset[resample, :]
    end

    return rsset
end

function smc(
    likelihood::Function,
    prior::Function,
    priorsample::Function,
    propdistpdf::Function,
    propdistsample::Function,
    snum::Int64,
    term::Function
    )

    j = 0
    sset = priorsample(snum)
    w = zeros(snum,1)

    #calculate weights
    for i in 1:snum
        w[i] = likelihood(sset[i,:])*prior(sset[i,:])
    end

    hatw = w./sum(w)#normalize

    while term(sset, j)
        j += 1
        Neff = 1/sum(hatw.^2)#efficient sample size
        if Neff < snum/2 #test for degeneracy
            sset = resample(sset, hatw)
            hatw = ones(snum,1)./snum
        end

        rsset = zeros(snum, size(sset,2))
        w = zeros(snum, 1)

        for i in 1:snum
            #one step in a random direction
            rsset[i,:] = sset[i,:] + propdistsample()
            w[i] = hatw[i].*likelihood(rsset[i,:])/likelihood(sset[i,:])*
                prior(rsset[i,:])/prior(sset[i,:])*
                propdistpdf(rsset[i,:]-sset[i,:])/propdistpdf(sset[i,:]-rsset[i,:])
        end

        sset = rsset
        hatw = w./sum(w)
    end
    return sset
end

function calcequaltails(
    dsample::Array,
    quantile::Float64,
    )
    snum = size(dsample, 1)
    dsample = sort(dsample, dims = 1)

    #calc cutoff points
    low = floor(Int32,snum*(1-quantile)/2)
    high = floor(Int32,snum - snum*(1-quantile)/2)

    rv = zeros(2, size(dsample,2))

    for i in 1:size(dsample,2)
        rv[:,i] = [dsample[low,i]; dsample[high,i]]
    end
    return rv
end

function grconvergence(
    likelihood::Function,
    prior::Function,
    priorsample::Function,
    propdistpdf::Function,
    propdistsample::Function,
    mcit::Int64,
    nchains::Int64
    )

    startval = priorsample(nchains)
    dims = size(startval, 2)
    rhat = 0
    mcitb = floor(Int, mcit/2)
    samples = zeros(floor(Int, mcit/2), size(startval,2), nchains)

    #create n Markov chains
    for i in 1:nchains
        samples[:,:,i] = mh(likelihood, prior, propdistpdf, propdistsample,startval[i,:], mcit, floor(Int, mcit/2))
    end

    if dims == 1
        chainvariance = zeros(nchains)
        chainmean = zeros(nchains)
        overallmean = 0
        for i in 1:nchains
            chainmean[i] = mean(samples[:,1,i], dims=1)[1]
            overallmean += chainmean[i]/nchains
            for j in 1:floor(Int,mcit/2) chainvariance[i] += (samples[j,1,i] - chainmean[i])^2 end
        end
        #vector of variances for each chain
        chainvariance = chainvariance./(mcitb-1)

        #variance of all chains
        mvariance = 0
        for i in 1:nchains mvariance += chainvariance[i]/nchains end

        #variance bias estimator
        Bn = 0
        for i in 1:nchains Bn += (chainmean[i]-overallmean)^2 end
        Bn = Bn/(nchains-1)

        #over corrected variance
        sig2 = ((mcitb-1)/mcitb)*mvariance + Bn

        rhat = sqrt(sig2/mvariance)
    else
        chainvariance = zeros(dims, dims, nchains)
        chainmean = zeros(dims, nchains)
        overallmean = zeros(dims)

        for i in 1:nchains
            chainmean[:, i] = mean(samples[:,:,i], dims=1)
            overallmean += chainmean[:, i]./nchains
            for j in 1:floor(Int, mcit/2)
                v = samples[j,:,i]-chainmean[:,i]
                chainvariance[:,:,i] += v*v'
            end
        end
        #array of covariance matrices for each chain
        chainvariance = chainvariance./(mcitb-1)

        #covariance matrix of all chains
        mvariance = zeros(dims, dims)
        for i in 1:nchains mvariance += chainvariance[:,:,i]./nchains end

        #bias estimator
        Bn = zeros(dims, dims)
        for i in 1:nchains
            v = chainmean[:,i] - overallmean
            Bn += v*v'
        end

        Bn = Bn*mcitb/(nchains-1)
        mvariance = inv(mvariance)
        maxeig = maximum(eigvals(mvariance*Bn))

        rhat = sqrt((mcitb-1)/mcitb + maxeig/mcitb)
    end

    return rhat
end


@testset "GeometricBM" begin

using DiffEqNoiseProcess, DiffEqBase, DiffEqMonteCarlo, Test, Statistics

μ = 1.0
σ = 2.0

W = GeometricBrownianMotionProcess(μ,σ,0.0,1.0,1.0)
dt = 0.1
calculate_step!(W,dt)
for i in 1:10
  accept_step!(W,dt)
end

W = GeometricBrownianMotionProcess(μ,σ,0.0,1.0,1.0)
prob = NoiseProblem(W,(0.0,1.0))
sol = solve(prob;dt=0.1)

W = GeometricBrownianMotionProcess(μ,σ,0//1,1.0,1.0)
prob_rational = NoiseProblem(W,(0,1))
dt = 1//10
sol = solve(prob_rational;dt=dt)

dt = dt/100
for t in dt:dt:1-dt
  sol(t)
end

t = 1.0
u0 = 1.0
expected_mean = u0*exp(μ*t)
expected_variance = u0^2*exp(2μ*t)*(exp(σ^2*t)-1)
monte_prob = MonteCarloProblem(prob;output_func = (sol,i)-> (sol[end],false))
sol = solve(monte_prob;dt=0.1,num_monte=40000)
@test abs(mean(sol) - expected_mean) < 0.4
#abs(var(sol) - expected_variance) < 0.4 # Converges slowly

W = GeometricBrownianMotionProcess!(μ,σ,0.0,ones(2),ones(2))
prob = NoiseProblem(W,(0.0,1.0))
sol = solve(prob;dt=0.1)

end


#=
Square-function with gradient
=#
function f_square(x)
    return x^2
end

function g_square(x)
    return 2*x
end

#=
Logistic-regression loss-function, gradient and sub-gradient
=#

function f_logreg()
    return function(x_i,y_i,θ,λ)
        return log(1+exp(-y_i*dot(θ,x_i)))
    end
end

function f_logreg(X,y,θ,λ)
    return function(θ)
        data = X
        labels = y
        λ = λ
        funcVal = 0
        # size(data,1) = 1000 datapoints
        for i in 1:size(data,1)
            # x = 784*1 column vector
            x_i = data[i:i,:]'
            funcVal += log( 1+ exp(-labels[i] * dot(θ,x_i) ) )
        end
        funcVal += λ/2*(norm(θ,2)^2)
        return funcVal
    end
end

function g_logreg(X,y,θ,λ)
    return function(θ)
        data = X
        labels = y
        λ = λ
        # n-dim-column-vector
        array = zeros(size(data,2))
        # add l2 norm
        array += λ * θ
        # iterate datapoints
        for i in 1:size(data,1)
            # pick i-th datapoint
        	x_i = data[i:i,:]'
            exponential = exp(-labels[i]*dot(θ,x_i))
            array += (-labels[i]*x_i * exponential)/(1+exponential)
        end
        return array
    end
end

function sub_g_logreg()
        return function(x_i,y_i,θ,λ)
              exponential = exp(-y_i*dot(θ,x_i))
              return (-y_i * x_i * exponential) / ( 1 + exponential ) + λ * θ
        end
end

#=
hinge-loss - Loss-function and gradient
=#
function f_svm(X,y,θ,λ)
    return function(θ)
        X = X
        y = y
        λ = λ
        fval = 0
        # datapoints
        for i in 1:size(X,1)
            x_i = X[i:i,:]'
            fval += max(0, (1 - y[i]*dot(θ,x_i)))
        end
        fval = fval / size(X,1)
        fval += λ/2 * norm(θ,2)
        return fval
    end
end

#=
Compute the gradient for one datapoint for hinge-loss.
=#
function sub_g_svm()
    return function(x_i,y_i,θ,λ)
      val = 1 - (y_i * dot(θ,x_i))
      if val < 0
          return λ * θ;
      else
          return -y_i * x_i + λ * θ;
      end
  end
end

#=
Rosenbrock-function and its gradient.
=#
function f_rosenbrock(x_k)
    return 100*(x_k[2]-x_k[1]^2)^2+(1-x_k[1])^2
end

function g_rosenbrock(x_k)
    return [ -400*x_k[1]*(x_k[2]-x_k[1]^2)-2(1-x_k[1]) ; 200*x_k[1]*(x_k[2]-x_k[1]^2) ]
end


function exp!(X::IntervalBox)
    x, y = X

    y = y ∩ exp(x)
    x = x ∩ log(y)

    return IntervalBox(x, y)
end

"""
Reverse function for `exp`.
"""
function exp_rev(y::Interval, x::Interval)
    y_new = y ∩ (0..∞)
    x_new = x ∩ log(y)
    return y_new, x_new
end

"""
Reverse function for `exp2`.
"""
function exp2_rev(y::Interval, x::Interval)
    y_new = y ∩ (0..∞)
    x_new = x ∩ log2(y)

    return y_new, x_new
end

"""
Reverse function for `exp10`.
"""
function exp10_rev(y::Interval, x::Interval)
    y_new = y ∩ (0..∞)
    x_new = x ∩ log10(y)

    return y_new, x_new
end

"""
Reverse function for `expm1`.
"""
function expm1_rev(y::Interval, x::Interval)
    y_new = y ∩ (-1.0..∞)
    x_new = x ∩ log1p(y)

    return y_new, x_new
end

function log!(X::IntervalBox)  # y = log(x)
    x, y = X

    x = x ∩ exp(y)
    y = y ∩ log(x)

    return IntervalBox(x, y)
end

"""
Reverse function for `log`: ``y = \\log(x)``
"""
function log_rev(y::Interval, x::Interval)
    x_new = x ∩ exp(y)

    return y, x_new
end

"""
Reverse function for `log2`: ``y = \\log2(x)``
"""
function log2_rev(y::Interval, x::Interval)
    x_new = x ∩ exp2(y)

    return y, x_new
end

"""
Reverse function for `log10`: ``y = \\log10(x)``
"""
function log10_rev(y::Interval, x::Interval)
    x_new = x ∩ exp10(y)

    return y, x_new
end

"""
Reverse function for `log1p`: ``y = \\log1p(x)``
"""
function log1p_rev(y::Interval, x::Interval)
    x_new = x ∩ expm1(y)

    return y, x_new
end


using ReachabilityAnalysis, Symbolics, Plots

@variables x y z
const positive_orthant = HPolyhedron([x >= 0, y >= 0, z >= 0], [x, y, z])

function prod_dest_verif(sol; T=100.0, target=10.0)

    # convert to a zonotopic flowpipe
    solz = overapproximate(sol, Zonotope)

    # project the final reach-set onto the space variables x, y, z
    X = project(solz(T), vars=(1, 2, 3))

    # check that all variables are nonnegative
    nonnegative = X ⊆ positive_orthant

    # compute the volume of the last reach-set
    H = overapproximate(X, Hyperrectangle)
    vol = volume(H)

    # check that that target belongs to the minkowski sum of the reach-sets projected in each coordinate
    B = convert(IntervalBox, H) # get the product-of-intervals representation
    contains_target = target ∈ sum(B)

    return nonnegative && contains_target, vol
end

@taylorize function prod_dest_I!(du, u, params, t)
    local a = 0.3
    x, y, z = u[1], u[2], u[3]

    du[1] = - (x * y) / (1 + x)
    du[2] = (x * y) / (1 + x) - a * y
    du[3] = a * y
    return du
end

X0 = (9.5 .. 10.0) × (0.01 .. 0.01) × (0.01 .. 0.01)
prob = @ivp(x'= prod_dest_I!(x), dim:3, x(0) ∈ X0)

solI = solve(prob, T=100.0, alg=TMJets(abstol=1e-11, orderT=7, orderQ=1));

property, vol = prod_dest_verif(solI)

plot(solI, vars=(0, 3), linecolor=:orange, color=:orange, alpha=0.3, lab="I")

@taylorize function prod_dest_IP!(du, u, params, t)
    x, y, z, a = u[1], u[2], u[3], u[4]

    du[1] = - (x * y) / (1 + x)
    du[2] = (x * y) / (1 + x) - a * y
    du[3] = a * y
    du[4] = zero(x)
    return du
end

X0 = (9.98 .. 9.98) × (0.01 .. 0.01) × (0.01 .. 0.01) × (0.296 .. 0.304)
prob = @ivp(x'= prod_dest_IP!(x), dim:4, x(0) ∈ X0)

solP = solve(prob, T=100.0, alg=TMJets(abstol=1e-12, orderT=7, orderQ=1));

property, vol = prod_dest_verif(solP)

plot(solP, vars=(0, 3), linecolor=:blue, color=:blue, alpha=0.3, lab="P")

X0 = (9.5 .. 10.0) × (0.01 .. 0.01) × (0.01 .. 0.01) × (0.296 .. 0.304)
prob = @ivp(x'= prod_dest_IP!(x), dim:4, x(0) ∈ X0)

solIP = solve(prob, T=100.0, alg=TMJets(abstol=1e-11, orderT=7, orderQ=1));

property, vol = prod_dest_verif(solIP)

plot(solIP, vars=(0, 3), linecolor=:red, color=:red, alpha=0.3, lab="I & P")

@taylorize function prod_dest_I_optimized!(du, u, params, t)
    local a = 0.3
    x, y, z = u[1], u[2], u[3]

    num = x * y
    den = 1 + x
    aux = num/den
    aux2 = a * y
    du[1] = -aux
    du[2] = aux - aux2
    du[3] = aux2
    return du
end

@taylorize function prod_dest_IP_optimized!(du, u, params, t)
    x, y, z, a = u[1], u[2], u[3], u[4]

    num = x * y
    den = 1 + x
    aux = num/den
    aux2 = a * y
    du[1] = -aux
    du[2] = aux - aux2
    du[3] = aux2
    du[4] = zero(x)
    return du
end


######### CmdStan program example  ###########

using CmdStan

ProjDir = @__DIR__
cd(ProjDir) do

  bernoullimodel = "
  data {
    int<lower=1> N;
    int<lower=0,upper=1> y[N];
  }
  parameters {
    real<lower=0,upper=1> theta;
  }
  model {
    theta ~ beta(1,1);
    y ~ bernoulli(theta);
  }
  ";

  observeddata = Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1])

  # Preserve these variables outside cd/do block.
  global stanmodel, rc, samples, cnames, df

  stanmodel = Stanmodel(name="bernoulli", model=bernoullimodel,
    printsummary=false);

  rc, samples, cnames = stan(stanmodel, observeddata, ProjDir, CmdStanDir=CMDSTAN_HOME);

  if rc == 0
    # Fetch cmdstan summary as a DataFrame
    df = read_summary(stanmodel)

    df |> display
    println()

    df[df.parameters .== :theta, [:mean, :ess, :r_hat]] |> display
  end

end # cd block


using MultivariateStats
using Base.Test
import StatsBase

srand(34568)

## PCA with zero mean

X = randn(5, 10)
Y = randn(3, 10)

W = qr(randn(5, 5))[1][:, 1:3]
σ² = 0.1
M = PPCA(Float64[], W, σ²)

@test indim(M) == 5
@test outdim(M) == 3
@test mean(M) == zeros(5)
@test loadings(M) == W
@test var(M) == σ²

T = inv(W'*W .+ σ²*eye(3))*W'
@test transform(M, X[:,1]) ≈ T * X[:,1]
@test transform(M, X) ≈ T * X

R = W*inv(W'W)*(W'W .+ σ²*eye(3))
@test reconstruct(M, Y[:,1]) ≈ R * Y[:,1]
@test reconstruct(M, Y) ≈ R * Y

## PCA with non-zero mean

mv = rand(5)
M = PPCA(mv, W, σ²)

@test indim(M) == 5
@test outdim(M) == 3
@test mean(M) == mv
@test loadings(M) == W
@test var(M) == σ²

@test transform(M, X[:,1]) ≈ T * (X[:,1] .- mv)
@test transform(M, X) ≈ T * (X .- mv)

@test reconstruct(M, Y[:,1]) ≈ R * Y[:,1] .+ mv
@test reconstruct(M, Y) ≈ R * Y .+ mv

## prepare training data

d = 5
n = 1000

R = qr(randn(d, d))[1]
@test R'R ≈ eye(5)
scale!(R, sqrt.([0.5, 0.3, 0.1, 0.05, 0.05]))

X = R'randn(5, n) .+ randn(5)
mv = vec(mean(X, 2))
Z = X .- mv

M0 = fit(PCA, X; mean=mv, maxoutdim = 4)

## ppcaml (default)

M = fit(PPCA, X)
P = projection(M)
W = loadings(M)

@test indim(M) == 5
@test outdim(M) == 4
@test mean(M) == mv
@test P'P ≈ eye(4)
@test reconstruct(M, transform(M, X)) ≈ reconstruct(M0, transform(M0, X))

M = fit(PPCA, X; mean=mv)
@test loadings(M) ≈ W

M = fit(PPCA, Z; mean=0)
@test loadings(M) ≈ W

M = fit(PPCA, X; maxoutdim=3)
P = projection(M)
W = loadings(M)

@test indim(M) == 5
@test outdim(M) == 3
@test P'P ≈ eye(3)

# ppcaem

M = fit(PPCA, X; method=:em)
P = projection(M)
W = loadings(M)

@test indim(M) == 5
@test outdim(M) == 4
@test mean(M) == mv
@test P'P ≈ eye(4)
@test all(isapprox.(reconstruct(M, transform(M, X)), reconstruct(M0, transform(M0, X)), atol=1e-3))

M = fit(PPCA, X; method=:em, mean=mv)
@test loadings(M) ≈ W

M = fit(PPCA, Z; method=:em, mean=0)
@test loadings(M) ≈ W

M = fit(PPCA, X; method=:em, maxoutdim=3)
P = projection(M)

@test indim(M) == 5
@test outdim(M) == 3
@test P'P ≈ eye(3)

@test_throws StatsBase.ConvergenceException fit(PPCA, X; method=:em, tot=1)

# bayespca
M0 = fit(PCA, X; mean=mv, maxoutdim = 3)

M = fit(PPCA, X; method=:bayes)
P = projection(M)
W = loadings(M)

@test indim(M) == 5
@test outdim(M) == 3
@test mean(M) == mv
@test P'P ≈ eye(3)
@test reconstruct(M, transform(M, X)) ≈ reconstruct(M0, transform(M0, X))

M = fit(PPCA, X; method=:bayes, mean=mv)
@test loadings(M) ≈ W

M = fit(PPCA, Z; method=:bayes, mean=0)
@test loadings(M) ≈ W

M = fit(PPCA, X; method=:em, maxoutdim=2)
P = projection(M)

@test indim(M) == 5
@test outdim(M) == 2
@test P'P ≈ eye(2)

@test_throws StatsBase.ConvergenceException fit(PPCA, X; method=:bayes, tot=1)

# test that fit works with Float32 values
X2 = convert(Array{Float32,2}, X)
# Float32 input, default pratio
M = fit(PPCA, X2; maxoutdim=3)


export Poisson

"""
Description:

    Poisson factor node

    Real scalars
    l > 0 (rate)

    f(out, l) = Poisson(out|l) = 1/(x!) * l^x * exp(-l)

Interfaces:

    1. out
    2. l

Construction:

    Poisson(id=:some_id)
"""
mutable struct Poisson <: SoftFactor
    id::Symbol
    interfaces::Vector{Interface}
    i::Dict{Symbol,Interface}

    function Poisson(out, l; id=generateId(Poisson))
        @ensureVariables(out, l)
        self = new(id, Array{Interface}(undef, 2), Dict{Symbol, Interface}())
        addNode!(currentGraph(), self)
        self.i[:out] = self.interfaces[1] = associate!(Interface(self), out)
        self.i[:l] = self.interfaces[2] = associate!(Interface(self), l)
        return self
    end
end

slug(::Type{Poisson}) = "Poisson"

format(dist::ProbabilityDistribution{Univariate, Poisson}) = "$(slug(Poisson))(l=$(format(dist.params[:l])))"

ProbabilityDistribution(::Type{Univariate}, ::Type{Poisson}; l=1.0) = ProbabilityDistribution{Univariate, Poisson}(Dict(:l=>l))
ProbabilityDistribution(::Type{Poisson}; l=1.0) = ProbabilityDistribution{Univariate, Poisson}(Dict(:l=>l))

dims(dist::ProbabilityDistribution{Univariate, Poisson}) = 1

vague(::Type{Poisson}) = ProbabilityDistribution(Univariate, Poisson, l=huge)

isProper(dist::ProbabilityDistribution{Univariate, Poisson}) = (0 < dist.params[:l] < huge)

unsafeMean(dist::ProbabilityDistribution{Univariate, Poisson}) = Float64(dist.params[:l])

unsafeVar(dist::ProbabilityDistribution{Univariate, Poisson}) = Float64(dist.params[:l])

logPdf(dist::ProbabilityDistribution{Univariate, Poisson}, x) = x*log(dist.params[:l]) - dist.params[:l] - logfactorial(x)

sample(dist::ProbabilityDistribution{Univariate, Poisson}) = poisinvcdf(dist.params[:l], rand())

# ∑ [λ^k*log(k!)]/k! from k=0 to inf
# Approximates the above sum for calculation of averageEnergy and differentialEntropy
# @ref https://arxiv.org/pdf/1708.06394.pdf
function apprSum(l, j=100)
    sum([(l)^(k)*logfactorial(k)/exp(logfactorial(k)) for k in collect(0:j)])
end

# Entropy functional
# @ref https://en.wikipedia.org/wiki/Poisson_distribution
function differentialEntropy(dist::ProbabilityDistribution{Univariate, Poisson})
    l = clamp(dist.params[:l], tiny, huge)
    l*(1-log(l)) + exp(-l)*apprSum(l)
end

# Average energy functional
function averageEnergy(::Type{Poisson}, marg_out::ProbabilityDistribution{Univariate}, marg_l::ProbabilityDistribution{Univariate})
    unsafeMean(marg_l) -
    unsafeMean(marg_out)*unsafeLogMean(marg_l) +
    exp(-unsafeMean(marg_out))*apprSum(unsafeMean(marg_out))
end


function assemble_system!(problem::AllenCahnProblem1D{TLBC, TRBC, TA}) where {TLBC <: NeumannBC, TRBC <: NeumannBC, TA <: SparseMatrixCSC}
    # Construct first row corresponding to uₓ(a) = vₗ
    problem.A[1,1] = -2
    problem.A[1,2] = 2

    # Construct rows corresponding to interior mesh points
    for i ∈ 2:problem.nₓ + 1
        problem.A[i, i - 1] = 1
        problem.A[i,i] = -2
        problem.A[i, i + 1] = 1
    end

    # Construct last row corresponding to uₓ(b) = vᵣ
    problem.A[end, end] = -2
    problem.A[end, end - 1] = 2

    @. problem.A *= 1/(problem.Δx)^2

    return problem
end

function assemble_system!(problem::AllenCahnProblem1D{TLBC, TRBC, TA}) where {TLBC <: PeriodicBC, TRBC <: PeriodicBC, TA <: SparseMatrixCSC}
    # Construct first row corresponding to u₀
    problem.A[1,1] = -2
    problem.A[1,2] = 1
    problem.A[1, end] = 1

    # Construct rows corresponding to interior mesh points
    for i ∈ 2:problem.nₓ
        problem.A[i, i - 1] = 1
        problem.A[i,i] = -2
        problem.A[i, i + 1] = 1
    end

    # Construct last row corresponding to uₙ
    problem.A[end, 1] = 1
    problem.A[end, end - 1] = 1
    problem.A[end, end] = -2

    @. problem.A *= 1/(problem.Δx)^2

    return problem
end


using Giac
using Test

using SpecialFunctions

#For elementary functions we check whether the Giac/Function of a
#numerical value computes (approximately) the same result as the
#corresponding Julia function

for (ex, res) in [
    (sqrt(giac(0.9)), sqrt(0.9)),
    (exp(giac(0.9)), exp(0.9)),
    (log(giac(0.9)), log(0.9)),
    (log(3,giac(0.9)), log(3,0.9)),
    (log10(giac(0.9)), log10(0.9)), # log10 not available in Giac
    (sin(giac(0.9)), sin(0.9)),
    (cos(giac(0.9)), cos(0.9)),
    (tan(giac(0.9)), tan(0.9)),
    (cot(giac(0.9)), cot(0.9)),
    (sec(giac(0.9)), sec(0.9)),
    (csc(giac(0.9)), csc(0.9)),
    (sinh(giac(0.9)), sinh(0.9)),
    (cosh(giac(0.9)), cosh(0.9)),
    (tanh(giac(0.9)), tanh(0.9)),
    (asin(giac(0.9)), asin(0.9)),
    (acos(giac(0.9)), acos(0.9)),
    (atan(giac(0.9)), atan(0.9)),
    (acot(giac(0.9)), acot(0.9)),
    (asec(giac(1.9)), asec(1.9)), # |arg|>=1
    (acsc(giac(1.9)), acsc(1.9)), # |arg|>=1
    (asinh(giac(0.9)), asinh(0.9)),
    (acosh(giac(1.9)), acosh(1.9)), #arg>=1
    (atanh(giac(0.9)), atanh(0.9)),
    (expm1(giac(0.9)), expm1(0.9)),
    (log1p(giac(0.9)), log1p(0.9)),

    (erf(giac(0.9)), erf(0.9)),
    (erfc(giac(0.9)), erfc(0.9)),
    (gamma(giac(0.9)), gamma(0.9)),
    (beta(giac(0.9),giac(1.9)), beta(0.9,1.9)),
    (zeta(giac(0.9)), zeta(0.9)),
    (factorial(giac(13)), factorial(13)),
    (binomial(giac(13),giac(7)), binomial(13,7)),
    (airyai(giac(0.9)), airyai(0.9)), # Giac/Airy_Ai does not work as expected
    (airybi(giac(0.9)), airybi(0.9)), # Giac/Airy_Bi does not work as expected
    (besselj(2,giac(0.9)), besselj(2,0.9)),
    (bessely(2,giac(0.9)), bessely(2,0.9)),
    (digamma(giac(0.9)), digamma(0.9)),
    (polygamma(2,giac(0.9)), polygamma(2,0.9)),
    (sinint(giac(0.9)), sinint(0.9)),
    (cosint(giac(0.9)), cosint(0.9)),
    (Ei(giac(0.9)), 1.6228117136968674413),
]
    @test convert(Float64,ex) ≈ res
end

@test convert(Float64, sign(giac(-12.3))) == -1.0
@test convert(Float64, sign(giac(0.0))) == 0.0
@test convert(Float64, sign(giac(1.3))) == 1.0
@test convert(Float64, round(giac(-1.7))) == -2.0
@test convert(Float64, round(giac(1.4999))) == 1.0
@test convert(Float64, ceil(giac(1.1))) == 2.0
@test convert(Float64, ceil(giac(-1.1))) == -1.0
@test convert(Float64, floor(giac(1.1))) == 1.0
@test convert(Float64, floor(giac(-1.1))) == -2.0
@test convert(Float64, trunc(giac(1.1))) == 1.0
@test convert(Float64, trunc(giac(-1.1))) == -1.0

#Note: We don't check whether the underlying Giac C++ library computes
#correct results. We only test if the Julia functions in Giac.jl/src/library.jl
#are correctly implemented. We assume that this is the case, if the Julia
#function computes the same (or at least an equivalent) result as the
#corresponding Giac function called via the Giac parser.

@giac x y z n

ex1 = simplify(4*atan(giac(1//5))-atan(giac(1//239)))
ex2 = giac("simplify(4*atan(1/5)-atan(1/239))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = collect(x+2*x+1-4)
ex2 = giac("collect(x+2*x+1-4)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = expand((x+y)*(z+1))
ex2 = giac("expand((x+y)*(z+1))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = factor(x^4-4)
ex2 = giac("factor(x^4-4)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = partfrac((x^2-2*x+3)/(x^2-3*x+2))
ex2 = giac("partfrac((x^2-2*x+3)/(x^2-3*x+2))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = numerator((x^3-1)/(x^2-1))
ex2 = giac("numer((x^3-1)/(x^2-1))")
#Note: Giac/numer -> Julia/numerator
@test evaluate(ex1) == evaluate(ex2)

ex1 = denominator((x^3-1)/(x^2-1))
ex2 = giac("denom((x^3-1)/(x^2-1))")
#Note: Giac/denom -> Julia/denominator
@test evaluate(ex1) == evaluate(ex2)

ex1 = left(x^2-1 ⩦ 2*x+3)
ex2 = giac("left(x^2-1 = 2*x+3)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = right(x^2-1 ⩦ 2*x+3)
ex2 = giac("right(x^2-1 = 2*x+3)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = diff(x^3-x,x)
ex2 = giac("diff(x^3-x,x)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = diff(exp(x*y),x,x,x,y,y)
ex2 = giac("diff(exp(x*y),x,x,x,y,y)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = integrate(1/(4+z^2),z)
ex2 = giac("integrate(1/(4+z^2),z)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = integrate(1/(1-x^4),x,2,3)
ex2 = giac("integrate(1/(1-x^4),x,2,3)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = limit(sin(x)/(x^2-3*x),x=>0)
ex2 = giac("limit(sin(x)/(x^2-3*x),x=0)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = limit((2*x-1)/exp(1/(x-1)),x=>plus_inf)
ex2 = giac("limit((2*x-1)/exp(1/(x-1)),x=+infinity)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = limit(sign(x),x=>0,1)
ex2 = giac("limit(sign(x),x=0,1)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = series((x^4+x+2)/(x^2+1),x=>0, 5)
ex2 = giac("series((x^4+x+2)/(x^2+1),x=0, 5)")
@test evaluate(ex1) == evaluate(ex2)

#ex1 = series(1/(1+x+y),[x,y],[0,0],5)
#ex2 = giac("series(1/(1+x+y),[x,y],[0,0],5)")
#@test evaluate(ex1) == evaluate(ex2)

ex1 = sum(1/n^2,n,1,17)
ex2 = giac("sum(1/n^2,n,1,17)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = sum(1/n^2,n,1,17,2)
ex2 = giac("sum(1/n^2,n,1,17,2)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = sum(1/n^2,n,1,infinity)
ex2 = giac("sum(1/n^2,n,1,infinity)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = sum(cos(n*x),n)
ex2 = giac("sum(cos(n*x),n)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = curl([2*x*y,x*z,y*z],[x,y,z])
ex2 = giac("curl([2*x*y,x*z,y*z],[x,y,z])")
@test evaluate(ex1) == evaluate(ex2)

ex1 = divergence([x^2+y,x+z+y,z^3+x^2],[x,y,z])
ex2 = giac("divergence([x^2+y,x+z+y,z^3+x^2],[x,y,z])")
@test evaluate(ex1) == evaluate(ex2)

ex1 = grad(2*x^2*y-x*z^3,[x,y,z])
ex2 = giac("grad(2*x^2*y-x*z^3,[x,y,z])")
@test evaluate(ex1) == evaluate(ex2)

ex1 = hessian(2*x^2*y-x*z,[x,y,z])
ex2 = giac("hessian(2*x^2*y-x*z,[x,y,z])")
@test evaluate(ex1) == evaluate(ex2)

ex1 = taylor(sin(x)/x,x)
ex2 = giac("taylor(sin(x)/x,x)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = taylor(sin(x)/x,x=>0, 7)
ex2 = giac("taylor(sin(x)/x,x=0, 7)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = solve(x^2-3⩦1)
ex2 = giac("solve(x^2-3=1)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = solve([y-z⩦0,z-x⩦0,x-y⩦0,x-1+y+z⩦0],[x,y,z])
ex2 = giac("solve([y-z=0,z-x=0,x-y=0,x-1+y+z=0],[x,y,z])")
@test evaluate(ex1) == evaluate(ex2)

ex1 = zeros(x^2-4,x)
ex2 = giac("zeros(x^2-4,x)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = zeros([x^2-1,x^2-y^2],[x,y])
ex2 = giac("zeros([x^2-1,x^2-y^2],[x,y])")
@test evaluate(ex1) == evaluate(ex2)

ex1 = cSolve(x^4-1⩦0, x)
ex2 = giac("cSolve(x^4-1=0, x)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = cSolve([x^4-y^4,x+y⩦0,x^2⩦2*x],[x,y])
ex2 = giac("cSolve([x^4-y^4,x+y=0,x^2=2*x],[x,y])")
@test evaluate(ex1) == evaluate(ex2)

ex1 = cZeros(x^4-1, x)
ex2 = giac("cZeros(x^4-1, x)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = cZeros([x^2-1,x^2-y^2],[x,y])
ex2 = giac("cZeros([x^2-1,x^2-y^2],[x,y])")
@test evaluate(ex1) == evaluate(ex2)

ex1 = fSolve(cos(x)⩦x,x=>0)
ex2 = giac("fsolve(cos(x)=x,x=0)")
#Note: renamed fsolve->fSolve, because that is the name the HP Prime uses
@test evaluate(ex1) == evaluate(ex2)

ex1 = deSolve(diff(y(x),x,x)+y(x)⩦0, y)
ex2 = giac("deSolve(diff(y(x),x,x)+y(x)=0, y)")
@test evaluate(ex1) == evaluate(ex2)

#ex1 = deSolve(y''(x)+y(x)⩦0, y)
#ex2 = giac("deSolve(y''+y=0, y)")
##Note: standalone y,y',y'' in this context not yet defined
#@test evaluate(ex1) == evaluate(ex2)

ex1 = linsolve([x+y+z⩦1,x-y⩦2,2*x-z⩦3],[x,y,z])
ex2 = giac("linsolve([x+y+z=1,x-y=2,2*x-z=3],[x,y,z])")
@test evaluate(ex1) == evaluate(ex2)

ex1 = logcollect(log(x)+2*log(y))
ex2 = giac("lncollect(log(x)+2*log(y))")
#Note: renamed lncollect->logcollect because Giac/ln is Julia/log
@test evaluate(ex1) == evaluate(ex2)

ex1 = powexpand(2^(x+y))
ex2 = giac("powexpand(2^(x+y))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = texpand(sin(2*x)+exp(x+y))
ex2 = giac("texpand(sin(2*x)+exp(x+y))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = exp2pow(exp(3*log(x)))
ex2 = giac("exp2pow(exp(3*log(x)))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = pow2exp(x^y)
ex2 = giac("pow2exp(x^y)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = exp2trig(exp(im*x))
ex2 = giac("exp2trig(exp(i*x))")
#Note: Julia/im corresponds to Giac/i
@test evaluate(ex1) == evaluate(ex2)

ex1 = expexpand(exp(3*x+2*y))
ex2 = giac("expexpand(exp(3*x+2*y))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = asin2acos(acos(x)+asin(x))
ex2 = giac("asin2acos(acos(x)+asin(x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = asin2atan(2*asin(x))
ex2 = giac("asin2atan(2*asin(x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = sin2costan(sin(x))
ex2 = giac("sin2costan(sin(x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = acos2asin(acos(x)+asin(x))
ex2 = giac("acos2asin(acos(x)+asin(x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = acos2atan(2*acos(x))
ex2 = giac("acos2atan(2*acos(x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = cos2sintan(cos(x))
ex2 = giac("cos2sintan(cos(x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = atan2asin(atan(2*x))
ex2 = giac("atan2asin(atan(2*x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = atan2acos(atan(2*x))
ex2 = giac("atan2acos(atan(2*x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = tan2sincos(tan(x))
ex2 = giac("tan2sincos(tan(x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = halftan(sin(x))
ex2 = giac("halftan(sin(x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = trigsin(cos(x)^4+sin(x)^2)
ex2 = giac("trigsin(cos(x)^4+sin(x)^2)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = trigcos(cos(x)^4+sin(x)^2)
ex2 = giac("trigcos(cos(x)^4+sin(x)^2)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = trigtan(cos(x)^4+sin(x)^2)
ex2 = giac("trigtan(cos(x)^4+sin(x)^2)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = atrig2log(atan(x))
ex2 = giac("atrig2ln(atan(x))")
#Note: renamed atrig2ln->atrig2log because Giac/ln is Julia/log
@test evaluate(ex1) == evaluate(ex2)

ex1 = tlin(sin(x)^3)
ex2 = giac("tlin(sin(x)^3)")
@test evaluate(ex1) == evaluate(ex2)

ex1 = tcollect(sin(x)+cos(x))
ex2 = giac("tcollect(sin(x)+cos(x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = trigexpand(sin(3*x))
ex2 = giac("trigexpand(sin(3*x))")
@test evaluate(ex1) == evaluate(ex2)

ex1 = trig2exp(sin(x))
ex2 = giac("trig2exp(sin(x))")
@test evaluate(ex1) == evaluate(ex2)

# to be continued ...


"""
A SYK model solver based on a forward fixed-point iteration method.

 The self-energy of the SYK model is given by,

    Σ(τ) = J² * G(τ) * G(τ) * G(β-τ),

 where Green's function of the SYK model is given by the Dyson equation,

    G(iωₙ) = -1/(iωₙ -μ + Σ(iωₙ))

 We solve the Dyson equation self-consistently by a weighted fixed point iteration,
 with weight `mix` assigned to the new iterate and weight `1-mix` assigned to the previous iterate.

 The self-energy is evaluated in the imaginary time domain,
 and the Dyson equation is solved in the Matsubara frequency domain.

 The SYK Green's function has particle-hole symmetry when μ=0.
 You may enforce such symmetry by setting `symmetry = :ph` when initialize the DLR grids.
 A symmetrized solver tends to be more robust than a unsymmetrized one.
"""

using Lehmann
using Printf

diff(a, b) = maximum(abs.(a - b)) # return the maximum deviation between a and b

conformal_tau(τ, β) = π^(1 / 4) / sqrt(2β) * 1 / sqrt(sin(π * τ / β))

function syk_sigma_dlr(d, G_x, J = 1.0; sumrule = nothing, verbose = false)

    tau_k = d.τ # DLR imaginary time nodes
    tau_k_rev = d.β .- tau_k # Reversed imaginary time nodes

    G_x_rev = tau2tau(d, G_x, tau_k_rev, sumrule = sumrule, verbose = verbose) # G at beta - tau_k

    Sigma_x = J .^ 2 .* G_x .^ 2 .* G_x_rev # SYK self-energy in imaginary time

    return Sigma_x
end

function dyson(d, sigma_q, mu)
    if d.symmetry == :ph #symmetrized G
        @assert mu ≈ 0.0 "Only the case μ=0 enjoys the particle-hole symmetry."
        return 1im * imag.(-1 ./ (d.ωn * 1im .- mu .+ sigma_q))
    elseif d.symmetry == :none
        return -1 ./ (d.ωn * 1im .- mu .+ sigma_q)
    else
        error("Not implemented!")
    end
end

function solve_syk_with_fixpoint_iter(d, mu, tol = d.rtol * 10; mix = 0.1, maxiter = 5000, G_x = zeros(ComplexF64, length(d)), sumrule = nothing, verbose = true)

    for iter in 1:maxiter

        Sigma_x = syk_sigma_dlr(d, G_x, sumrule = sumrule, verbose = verbose)

        G_q_new = dyson(d, tau2matfreq(d, Sigma_x), mu)

        G_x_new = matfreq2tau(d, G_q_new, sumrule = sumrule, verbose = verbose)

        if verbose
            if iter % (maxiter / 10) == 0
                println("round $iter: change $(diff(G_x_new, G_x))")
            end
        end
        if maximum(abs.(G_x_new .- G_x)) < tol && iter > 10
            break
        end

        G_x = mix * G_x_new + (1 - mix) * G_x # Linear mixing
    end
    return G_x
end

function printG(d, G_x)
    @printf("%15s%40s%40s%40s\n", "τ", "DLR imag", "DLR real", "asymtotically exact")
    for i in 1:d.size
        if d.τ[i] <= d.β / 2
            @printf("%15.8f%40.15f%40.15f%40.15f\n", d.τ[i], imag(G_x[i]), real(G_x[i]), conformal_tau(d.τ[i], d.β))
        end
    end
    println()
end

verbose = false

printstyled("=====    Prepare the expected Green's function of the SYK model     =======\n", color = :yellow)
dsym_correct = DLRGrid(Euv = 5.0, β = 10000.0, isFermi = true, rtol = 1e-14, symmetry = :ph) # Initialize DLR object
G_x_correct = solve_syk_with_fixpoint_iter(dsym_correct, 0.00, mix = 0.1, verbose = false)
printG(dsym_correct, G_x_correct)

printstyled("=====    Test Symmetrized and Unsymmetrized DLR solver for SYK model     =======\n", color = :yellow)

@printf("%30s%30s%30s%30s%20s\n", "Euv", "sym_solver", "unsym_solver", "unsym_solver+sum_rule", "good or bad")
for Euv in LinRange(5.0, 10.0, 50)

    rtol = 1e-10
    β = 10000.0
    # printstyled("=====     Symmetrized DLR solver for SYK model     =======\n", color = :yellow)
    mix = 0.01
    dsym = DLRGrid(Euv = Euv, β = β, isFermi = true, rtol = rtol, symmetry = :ph, rebuild = true, verbose = false) # Initialize DLR object
    G_x_ph = solve_syk_with_fixpoint_iter(dsym, 0.00, mix = mix, sumrule = nothing, verbose = verbose)

    # printstyled("=====     Unsymmetrized DLR solver for SYK model     =======\n", color = :yellow)
    mix = 0.01
    dnone = DLRGrid(Euv = Euv, β = β, isFermi = true, rtol = rtol, symmetry = :none, rebuild = true, verbose = false) # Initialize DLR object
    G_x_none = solve_syk_with_fixpoint_iter(dnone, 0.00, mix = mix, sumrule = nothing, verbose = verbose)

    # printstyled("=====     Unsymmetrized DLR solver for SYK model     =======\n", color = :yellow)
    mix = 0.01
    G_x_none_sumrule = solve_syk_with_fixpoint_iter(dnone, 0.00, mix = mix, sumrule = 1.0, verbose = verbose)
    # printG(dnone, G_x_none)

    # printstyled("=====     Unsymmetrized versus Symmetrized DLR solver    =======\n", color = :yellow)
    # @printf("%15s%40s%40s%40s\n", "τ", "sym DLR (interpolated)", "unsym DLR", "difference")
    # G_x_interp = tau2tau(dsym_correct, G_x_correct, dnone.τ)
    # for i in 1:dnone.size
    #     if dnone.τ[i] <= dnone.β / 2
    #         @printf("%15.8f%40.15f%40.15f%40.15f\n", dnone.τ[i], real(G_x_interp[i]), real(G_x_none[i]), abs(real(G_x_interp[i] - G_x_none[i])))
    #     end
    # end

    G_x_interp_ph = tau2tau(dsym_correct, G_x_correct, dsym.τ)
    G_x_interp_none = tau2tau(dsym_correct, G_x_correct, dnone.τ)
    G_x_interp_none_sumrule = tau2tau(dsym_correct, G_x_correct, dnone.τ)
    d_ph = diff(G_x_interp_ph, G_x_ph)
    d_none = diff(G_x_interp_none, G_x_none)
    d_none_sumrule = diff(G_x_interp_none_sumrule, G_x_none_sumrule)
    flag = (d_ph < 100rtol) && (d_none < 100rtol) && (d_none_sumrule < 100rtol) ? "good" : "bad"

    @printf("%30.15f%30.15e%30.15e%30.15e%20s\n", Euv, d_ph, d_none, d_none_sumrule, flag)
    # println("symmetric Euv = $Euv maximumal difference: ", diff(G_x_interp, G_x_ph))
    # println("non symmetric Euv = $Euv maximumal difference: ", diff(G_x_interp, G_x_none))

end


#=
The decimal number, 585 = 1001001001 (binary), is palindromic in both bases.

Find the sum of all numbers, less than one million, which are palindromic in base 10 and base 2.

(Please note that the palindromic number, in either base, may not include leading zeros.)
=#

function isdecpalindrome(a)
  digits(a) == reverse(digits(a))
end

function isbinarypalindrome(a)
  digits(a, base=2) == reverse(digits(a, base=2))
end

function calc()
  total = 0
  for i in 1:999_999
    if isdecpalindrome(i) && isbinarypalindrome(i)
      total += i
    end
  end
  total
end
@time println(calc())


include("preamble.jl")

# ----------------------------
# Plotting a 3D polyhedron
# ----------------------------

# Note: in this figure we use non-latex fonts for the ticks labels because
# the Makie plot ignores L"x_3"

using LazySets, Symbolics
import Polyhedra, CDDLib, Plots
using Colors

# projected polyhedron from symbolic constraints
vars = @variables x1, x2, x3  # create symbols
P = HPolyhedron([2*x1>=0, 3*x2+1.7*x3>=0], vars)
Q = project(P, [2, 3])  # 2D projection (x2 and x3)

fig = p(xlim=(-4, 4), ylim=(-4, 4), ratio=1, xlab="x₂", ylab="x₃",
           legendfontsize=20,
           tickfont=font(20), # , "Times"),
           guidefontsize=20,
           xguidefont=font(20), # , "Times"),
           yguidefont=font(20), # , "Times"),
           xtick = ([-4, 0, 4], ["-4", "0", "4"]),
           ytick = ([-4, 0, 4], ["-4", "0", "4"]),
           bottom_margin=0mm, left_margin=0mm, right_margin=0mm, top_margin=2mm, size=(500, 380))

plot!(fig, Q, color=:deepskyblue)  # plot the 2D projection

savefig("polyhedron2D.pdf")

# plot the 3D set directly
using GLMakie

# intersection with a bounding box
B = BallInf(zeros(3), 5.0)
R = intersection(P, B)

# obtain the mesh
backend = LazySets.default_polyhedra_backend(R);
P_poly_mesh = LazySets._plot3d_helper(R, backend);

fig3d = Figure()
ax = Axis3(fig3d[1, 1], xlabel="x₁", xlabelsize=40, xticklabelsize=40,
                    ylabel="x₂", ylabelsize=40, yticklabelsize=40,
                    zlabel="x₃", zlabelsize=40, zticklabelsize=40,
                    aspect=:data,
                    height=400)

mesh!(P_poly_mesh, alpha=1.0, color=:deepskyblue, colormap=:viridis, colorrange=Makie.Automatic(), interpolate=false,
                   linewidth=1, overdraw=false, shading=true, transparency=true, visible=true)

wireframe!(P_poly_mesh, color = :red, linewidth=1.5, linestyle=:dash, axis=(type=Axis3,))

Makie.save("polyhedron3D.png", fig3d)

# f = Figure()
# fig = plot3d(R)  # plot in three dimensions
# current_figure()
# display(fig) # display in a pop-up window


# indicator of a simplex

"""
  IndSimplex(a::Union{Real, Integer}=1.0)

Returns the function `g = ind{x : x ⩾ 0, sum(x) = a}`.
"""

immutable IndSimplex{T <: Union{Real, Integer}} <: IndicatorConvex
  a::T
  function IndSimplex(a::T)
    if a <= 0
      error("parameter a must be positive")
    else
      new(a)
    end
  end
end

IndSimplex{T <: Union{Real, Integer}}(a::T=1.0) = IndSimplex{T}(a)

function (f::IndSimplex){T <: Real}(x::AbstractArray{T,1})
  if all(x .>= 0) && abs(sum(x)-f.a) <= 1e-14
    return 0.0
  end
  return +Inf
end

function prox!{T <: Real}(f::IndSimplex, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
# Implements Algorithm 1 in Condat, "Fast projection onto the simplex and the l1 ball", Mathematical Programming, 158:575–585, 2016.
# We should consider implementing the other algorithms reviewed there, and the one proposed in the paper.
  n = length(x)
  p = []
  if ndims(x) == 1
    p = sort(x, rev=true)
  else
    p = sort(x[:], rev=true)
  end
  s = 0
  for i = 1:n-1
    s = s + p[i]
    tmax = (s - f.a)/i
    if tmax >= p[i+1]
      @inbounds for j in eachindex(y)
        y[j] = x[j] < tmax ? 0.0 : x[j] - tmax
      end
      return 0.0
    end
  end
  tmax = (s + p[n] - f.a)/n
  @inbounds for j in eachindex(y)
    y[j] = x[j] < tmax ? 0.0 : x[j] - tmax
  end
  return 0.0
end

fun_name(f::IndSimplex) = "indicator of the probability simplex"
fun_dom(f::IndSimplex) = "AbstractArray{Real}"
fun_expr(f::IndSimplex) = "x ↦ 0 if x ⩾ 0 and sum(x) = a, +∞ otherwise"
fun_params(f::IndSimplex) = "a = $(f.a)"

function prox_naive{T <: Real}(f::IndSimplex, x::AbstractArray{T}, gamma::Real=1.0)
  low = minimum(x)
  upp = maximum(x)
  v = x
  s = Inf
  for i = 1:100
    if abs(s)/f.a <= 1e-15
      break
    end
    alpha = (low+upp)/2
    v = max.(x - alpha, 0.0)
    s = sum(v) - f.a
    if s <= 0
      upp = alpha
    else
      low = alpha
    end
  end
  return v, 0.0
end


#=
  Test of Bernoulli

  See
     ~/blog/bernoulli_test.blog
     ~/webppl/bernoulli_test.wppl
=#

# using Memoization
using Turing, StatsPlots, DataFrames
using ReverseDiff, Zygote, Tracker
# Turing.setadbackend(:reversediff)
# Turing.setadbackend(:zygote)
# Turing.setadbackend(:tracker)
include("jl_utils.jl")

@model function bernoulli_test(y)
    prob ~ Beta(4,4)
    n = length(y)
    for i in 1:n
        y[i] ~ Bernoulli(prob)
    end

end

# Generated by R:
# > rbinom(120,1,0.3)
# I.e. prob should be around 0.3
#=
y = [0,0,1,0,0,0,0,0,0,0,1,0,0,0,0,1,0,0,0,0,1,0,0,0,0,0,0,0,1,1,0,1,1,0,1,0,
    1,0,0,0,0,1,0,0,0,1,0,1,0,0,0,0,0,1,1,0,1,0,1,0,0,1,1,0,0,0,1,0,1,0,0,0,
    0,0,1,0,0,1,0,1,0,0,0,0,0,0,0,0,0,1,0,1,1,0,0,1,0,1,1,0,0,1,0,0,1,0,1,1,
    0,1,0,1,0,0,0,0,0,1,0,0]
=#
# Generate data
y = rand(Bernoulli(0.3),1_000)
println("length(y):$(length(y))")
# println("y:$y")
model = bernoulli_test(y)
num_chns = 4

# chns = sample(model, Prior(), MCMCThreads(), 1000, num_chns)

# chns = sample(model, MH(), 10_000)
chns = sample(model, MH(), MCMCThreads(), 10_000, num_chns)

# chns = sample(model, PG(20), MCMCThreads(), 10_000, num_chns)
# chns = sample(model, PG(20), 1_000)
# chns = sample(model, IS(), MCMCThreads(), 10_000, num_chns) # Nope: does not handle the observation
# chns = sample(model, IS(), 10_000)

# chns = sample(model, SMC(), 10_000)

# chns = sample(model, NUTS(1000,0.65), MCMCThreads(), 40_000, num_chns)
# chns = sample(model, Gibbs(MH(:zlabels),NUTS(1000,0.65,:m,:b,:sigma)), MCMCThreads(), 40_000, num_chns)

display(chns)


# These examples illustrate common operations on polyhedra using [Polyhedra.jl](https://github.com/JuliaPolyhedra/Polyhedra.jl):
#
# - [Convex hull](#Convex-hull).
# - [Intersection](#Intersection).
#
# We start by choosing a polyhedral library that will be used for computing the H-representation from the V-representation and vice-versa as well as removing redundant points.
# In these example, we use the default library available in Polyhedra but it can be replaced by any other library listed [here](https://juliapolyhedra.github.io/), e.g. by changing the last two lines below by `import CDDLib` and `lib = CDDLib.Library()` to use [CDDLib](https://github.com/JuliaPolyhedra/CDDLib.jl).

using Test #jl
using Polyhedra
import GLPK
lib = DefaultLibrary{Float64}(GLPK.Optimizer)

# ### Convex hull
#
# The binary convex hull operation between two polyhedra is obtained with the [`convexhull`](https://juliapolyhedra.github.io/Polyhedra.jl/latest/utilities.html#Polyhedra.convexhull) function.
#
# Below we compute the convex hull of the union of two polygons from their V-representation.

P1 = polyhedron(vrep([
    -1.9 -1.7
    -1.8  0.5
     1.7  0.7
     1.9 -0.3
     0.9 -1.1
]), lib)

P2 = polyhedron(vrep([
    -2.5 -1.1
    -0.8  0.8
     0.1  0.9
     1.8 -1.2
     1.3  0.1
]), lib)

Pch = convexhull(P1, P2)

# Note that the convex hull operation is done in the V-representation so no representation conversion is needed for this operation since `P1` and `P2` where constructed from their V-representation:

hrepiscomputed(P1), hrepiscomputed(P2), hrepiscomputed(Pch) #!jl
@test (hrepiscomputed(P1), hrepiscomputed(P2), hrepiscomputed(Pch)) == (false, false, false) #jl

# Let us note that the `convexhull` of a V-representation contains points and rays and represents the convex hull of the points together with the conic hull of the rays. So, `convexhull(P1, P2)` does the union of the vertices:

npoints(Pch) #!jl
@test npoints(Pch) == 10 #jl

# However, if we want to remove the redundant points we can use `removevredundancy!`:

removevredundancy!(Pch)
npoints(Pch) #!jl
@test npoints(Pch) == 8 #jl

# We can plot the polygons and the convex hull of their union using the `plot` function. For further plotting options see the [Plots.jl](http://docs.juliaplots.org/latest/) documentation.
# We can see below the 8 redundant points highlighted with green dots, the two points that are not highlighted are the redundant ones.

using Plots                          #!jl
plot(P1, color="blue", alpha=0.2)    #!jl
plot!(P2, color="red", alpha=0.2)    #!jl
plot!(Pch, color="green", alpha=0.1) #!jl
scatter!(Pch, color="green")         #!jl

# ### Intersection
#
# Intersection of polyhedra is obtained with the [`intersect`](https://juliapolyhedra.github.io/Polyhedra.jl/latest/utilities.html#Base.intersect) function.
#
# Below we compute the intersection of the two polygons from the previous example.

Pint = intersect(P1, P2)
@test nhalfspaces(Pint) == 10 #jl

# While `P1` and `P2` have been constructed from their V-representation, their H-representation has been computed to build the intersection `Pint`.

hrepiscomputed(P1), vrepiscomputed(P1), hrepiscomputed(P2), vrepiscomputed(P2) #!jl
@test (hrepiscomputed(P1), vrepiscomputed(P1), hrepiscomputed(P2), vrepiscomputed(P2)) == (true, true, true, true) #jl

# On the other hand, `Pint` is constructed from its H-representation hence its V-representation has not been computed yet.

hrepiscomputed(Pint), vrepiscomputed(Pint) #!jl
@test (hrepiscomputed(Pint), vrepiscomputed(Pint)) == (true, false) #jl

# We can obtain the number of points in the intersection with `npoints` as follows:

npoints(Pint) #!jl
@test npoints(Pint) == 8 #jl

# Note that this triggers the computation of the V-representation:

hrepiscomputed(Pint), vrepiscomputed(Pint) #!jl
@test (hrepiscomputed(Pint), vrepiscomputed(Pint)) == (true, true) #jl

# We can plot the polygons and their intersection using the `plot` function.

using Plots                            #!jl
plot(P1, color="blue", alpha=0.2)      #!jl
plot!(P2, color="red", alpha=0.2)      #!jl
plot!(Pint, color="yellow", alpha=0.6) #!jl

examples/Convex hull and intersection.jl


"""
filter_matrix_meshwidth(f, domain, N)

Assemble discrete filtering matrix from a continuous filter `f` width constant width
\$h(x) = \\Delta x / 2\$.
"""
function filter_matrix_meshwidth(f::TopHatFilter, domain::PeriodicIntervalDomain, N)
    x = discretize(domain, N)
    Δx = 1 // N * (domain.right - domain.left)
    h = f.width
    h₀ = h(x[1])

    all(≈(h(x[1])), h.(x)) || error("Filter width must be constant")
    if Δx .≈ 2h₀
        # Three point stencil
        inds = [-1, 0, 1]
        stencil = [1 // 24, 11 // 12, 1 // 24]
    elseif Δx .≈ h₀
        # Three point stencil
        inds = [-1, 0, 1]
        stencil = [1 // 6, 4 // 6, 1 // 6]
    else
        error("Filter width must be equal to mesh width")
    end

    # Five point stencil
    # inds = [-2, -1, 0, 1, 2]
    # stencil = [-6 // 2033, 77 // 1440, 863 // 960, 77 // 1440, -6 // 2033]

    # Construct banded matrix
    diags = [i => fill(s, N - abs(i)) for (i, s) ∈ zip(inds, stencil)]
    W = spdiagm(diags...)

    # Periodic extension of three point stencil
    W[1, end] = stencil[1]
    W[end, 1] = stencil[end]

    # Periodic extension of five point stencil
    # W[1, [end - 1, end]] = stencil[[1, 2]]
    # W[2, end] = stencil[1]
    # W[end - 1, 1] = stencil[end]
    # W[end, [1, 2]] = stencil[[end-1, end]]

    W
end

function filter_matrix_meshwidth(f::TopHatFilter, domain::ClosedIntervalDomain, N)
    x = discretize(domain, N)
    Δx = 1 // N * (domain.right - domain.left)
    h = f.width
    h₀ = h(x[1])

    all(≈(h(x[1])), h.(x)) || error("Filter width must be constant")
    Δx ≈ 2h₀ || error("Filter width must be equal to mesh width")

    # Three point stencil
    inds = [-1, 0, 1]
    stencil = [1 / 24, 11 / 12, 1 / 24]

    # Five point stencil
    # inds = [-2, -1, 0, 1, 2]
    # stencil = [-6 / 2033, 77 / 1440, 863 / 960, 77 / 1440, -6 / 2033]

    # Construct banded matrix
    diags = [i => fill(s, N + 1 - abs(i)) for (i, s) ∈ zip(inds, stencil)]
    W = spdiagm(diags...)

    # Boundary weights for three point stencil
    W[1, [1, 2]] = [3 / 4, 1 / 4]
    W[end, [end, end - 1]] = [3 / 4, 1 / 4]

    # Boundary weights for five point stencil
    # W[1, [1, 2, 3]] = [2 / 3, 5 / 12, -1 / 12]
    # W[2, [1, 2, 3, 4]] = [1/24, 11/12, 1/24, 0]
    # W[end - 1, [end, end - 1, end - 2, end - 3]] = [1/24, 11/12, 1/24, 0]
    # W[end, [end, end - 1, end - 2]] = [2 / 3, 5 / 12, -1 / 12]

    W
end


function euclidean_distance(vec1, vec2)
    sqrt(sum([(vec1[i]-vec2[i])^2 for i in 1:length(vec1)]))
end

function k_nearest_neighbour(eigenface_table, new_eigenface, k)
    n=length(eigenface_table)
    distances = [euclidean_distance(eigenface_table[i][2], new_eigenface) for i in 1:n]
    nearest_neighbours = [eigenface_table[i] for i in sortperm(distances)[1:k]]
    classes = unique([neighbour[1] for neighbour in nearest_neighbours])
    nearest = nearest_neighbours[1][1]
    inverse_distance(neighbour) = 1/(euclidean_distance(neighbour, new_eigenface))
    min = 0
    for class in classes
        freq = sum([inverse_distance(neighbour[2]) for neighbour in filter(x->x[1]==class, nearest_neighbours)])
        if freq > min
            min = freq
            nearest = class
        end
    end
    nearest
end

function test_model(proj, images, d, n, k)
    labels = unique(map(x->x[2], images))
    training_set = reduce(vcat,[map(x->x[1], filter(x->x[2] == labels[i], images))[1:n+1] for i in 1:size(labels, 1)])
    image_matrix = images_to_image_matrix(training_set)
    transformed = transpose(transpose(proj[:,1:d])*image_matrix)
    eigenface_table = reduce(vcat, [[[labels[i+1] [transformed[j,:]]] for j in (i*(n+1)+1):(i*(n+1)+n+1)] for i in 0:size(labels, 1)-1])
    correct = 0
    for i in 1:length(training_set)
        test_eigenface = eigenface_table[i]
        data = reduce(vcat, [eigenface_table[1:i-1], eigenface_table[i+1:end]])
        if k_nearest_neighbour(data, test_eigenface[2], k) == test_eigenface[1]
            correct = correct + 1
        end
    end
    return correct/length(training_set)
end


# tests based on AMLCR 2nd ed,
@testset "AMLCR Math" begin

    # §5.7 payment frequency and certain (Tables 5.1 and 5.2)
    function f(age)
        life = SingleLife(mort = AMLCR_std, issue_age = age)
        int = Yields.Constant(0.05)
        lc = [
            AnnuityImmediate(life,int)
            AnnuityImmediate(life,int,frequency=4)
            AnnuityDue(life,int)
            AnnuityDue(life,int,frequency=4)
            AnnuityImmediate(life,int,n=10)
            AnnuityImmediate(life,int,n=10,frequency=4)
            AnnuityDue(life,int,n=10,frequency=4)
        ]

        return present_value.(lc)
    end

    # atol only to two decimals b/c AMLCR uses rounded rates?
    @test all(isapprox.(f(20), [18.966, 19.338, 19.966, 19.588, 7.711, 7.855, 7.952],atol=.01))
    @test all(isapprox.(f(80), [ 7.548,  7.917,  8.548,  8.167, 6.128, 6.373, 6.539],atol=.01))

    life = SingleLife(mort = AMLCR_std, issue_age = 20)
    int = Yields.Constant(0.05)

    # property tests

    @test present_value(AnnuityDue(life,int)) > present_value(AnnuityDue(life,int,start_time=5))

    @test present_value(AnnuityDue(life,int,n=5,certain=5,start_time=5)) == sum([1 for t in 5:9] .* [1/1.05 ^t for t in 5:9])

    # relation on pg 124
    @test issorted(present_value.([
        AnnuityImmediate(life,int)
        AnnuityImmediate(life,int,frequency=4)
        AnnuityDue(life,int,frequency=4)
        AnnuityDue(life,int)
    ]))

end#


# generate a random pos def matrix with eigenvalues between 0.1 and 2
function generate_pos_def_matrix(n::Int64,rng)
	X = rand(rng, n, n)
	# any real square matrix can be QP decomposed into a orthogonal matrix and an uppertriangular matrix R
	Q, R = qr(X)
	eigs = rand(rng ,n) .* (2 .- 0.1) .+ 0.1
	X = Q * Matrix(Diagonal(eigs)) * Q'
	X = 0.5 * (X + X')
	return X
end

function is_numerically_pos_sem_def(X, atol)
	X = X ./ 2
	X = X + X'

	F = eigfact(X)
	if size(find( x-> x < -atol, F[:values]), 1) == 0
		return true
	else
		return false
	end
end

function is_numerically_symmetric(X,atol)
	n = size(X, 2)
	for i = 1:n-1, j = i+1:n
		if abs(X[i, j] - X[j, i]) >= atol
			return false
		end
	end
	return true
end

function find_nonsymmetric_component(X)
	for i = 2:size(X, 1), j = 1:(i - 1)
		if abs(X[i, j] - X[j, i]) > 0.0
			return i, j, abs(X[i, j] - X[j, i])
		end
	end
end

function find_different_elements(A, B)
	if size(A) != size(B)
		error("Matrices are not the same size")
	end
	m, n = size(A)
	diff_el = Array[]
	for iii = 1:m, jjj = 1:n
		if A[iii, jjj] != B[iii, jjj]
			push!(diff_el, [iii, jjj])
		end
	end
	return diff_el
end

function duplicate_sparsity_pattern(A)
	m, n = size(A)
	B = zeros(m, n)
	for iii = 1:m, jjj = 1:n
		if A[iii, jjj] != 0
			B[iii, jjj] = 1
		end
	end
	return B

end

function recreate_sparse_matrix(A)
	rowInd = A.rowval
	colPtr = A.colptr
	val = A.nzval

	#compute column indices
	colInd = zeros(Int64, length(rowInd))
	cval = 1
	for iii = 2:length(colPtr)
		currentPtr = colPtr[iii]
		prevPtr = colPtr[iii - 1]
		colInd[prevPtr:currentPtr - 1] = cval
		cval += 1
	end

	# sort rowInd and vals
	p = sortperm(rowInd)
	return sparse(rowInd, colInd, val, size(A, 1), size(A, 2))
end

# Geometric mean from https://github.com/JuliaStats/StatsBase.jl
function gmean(a::AbstractArray{T}) where T<:Real
	s = 0.0
	n = length(a)
	for i in 1:n
		tmp = a[i]
		if tmp < 0.0
			throw(DomainError())
		elseif tmp == 0.0
			return 0.0
		else
			s += log(tmp)
		end
	end
	return exp(s / n)
end


#
# Restricted Two-body Problem models
#

"""
A `ModelingToolkit.ODESystem` for the Restricted Two-body Problem.

The order of the states follows: `[x, y, z, ẋ, ẏ, ż]`.

The order of the parameters follows: `[μ]`.

# Extended Help
The Restricted Two-body Problem is a simplified dynamical model
describing one small body (spacecraft, etc.) and one celestial
body. The gravity of the celestial body exhibits a force on the
small body. This model is commonly used as a simplification to
descibe our solar systems' planets orbiting our sun, or a
spacecraft orbiting Earth.

### Usage

```julia
model = R2BP()
```
"""
@memoize function R2BP(; stm=false, structural_simplify=true, name=:R2BP)

    @parameters t μ
    @variables x(t) y(t) z(t) ẋ(t) ẏ(t) ż(t)
    δ = Differential(t)
    r = @SVector [x,y,z]
    v = @SVector [ẋ,ẏ,ż]

    eqs = vcat(
        δ.(r) .~ v,
        δ.(v) .~ -μ .* (r ./ norm(r)^3)
    )

    if stm
        @variables Φ[1:6,1:6](t)
        Φ = Symbolics.scalarize(Φ)
        A = Symbolics.jacobian(map(el -> el.rhs, eqs), vcat(r,v))

        LHS = map(δ, Φ)
        RHS = map(simplify, A * Φ)

        eqs = vcat(eqs, [LHS[i] ~ RHS[i] for i in 1:length(LHS)])
    end

    if string(name) == "R2BP" && stm
        modelname = Symbol("R2BPWithSTM")
    else
        modelname = name
    end

    sys = ODESystem(
        eqs, t, stm  ? vcat(r,v,Φ...) : vcat(r,v), [μ];
        name = modelname
    )
    return structural_simplify ? ModelingToolkit.structural_simplify(sys) : sys
end

"""
Returns an `ODEFunction` for R2BP dynamics.
Results are cached with `Memoize.jl`.

The order of the states follows: `[x, y, z, ẋ, ẏ, ż]`.

The order of the parameters follows: `[μ]`.

# Extended Help

### Usage

The `stm`, `structural_simplify`, and `name` keyword arguments
are passed to `R2BP`. All other keyword arguments are passed
directly to `SciMLBase.ODEFunction`.

```julia
f = R2BPFunction(; stm=false, structural_simplify=true, name=:R2BP, jac=true)
let u = randn(6), p = randn(1), t = 0
    f(u, p, t)
end
```
"""
@memoize function R2BPFunction(; stm=false, structural_simplify=true, name=:R2BP, kwargs...)
    defaults = (; jac=true)
    options  = merge(defaults, kwargs)
    return ODEFunction(
        R2BP(; stm=stm, structural_simplify=structural_simplify, name=name);
        options...
    )
end


using SpecialPolynomials; const SP=SpecialPolynomials
using Polynomials
import LinearAlgebra: eigvals, det, dot, norm
import SpecialFunctions: gamma
using Test

@testset "Orthogonal" begin include("Orthogonal.jl") end
@testset "Orthogonal compare with tables" begin include("Orthogonal-compare.jl") end
#@testset "ChebyshevT" begin include("ChebyshevT.jl") end
#@testset "ChebyshevU" begin include("ChebyshevU.jl") end
#@testset "Laguerre{α}" begin include("GeneralizedLaguerre.jl") end
#@testset "Gegenbauer{α}" begin include("Gegenbauer.jl") end

@testset "Lagrange" begin include("Lagrange.jl") end
@testset "Newton" begin include("Newton.jl") end

@testset "Bernstein" begin include("Bernstein.jl") end


"""
    count_total_bits(n::Int64)

Calculates the total number of bits from 1 to `n` for a given integer `n` by combing a
while-loop for building the binary subsets with a range for-loop for counting the bits.

# Arguments
- `n::Int64`: Integer value which has to be evaluated.

# Examples
```julia-repl
julia> import ClassicAlgorithmsCollections
julia> ClassicAlgorithmsCollections.count_total_bits(17)
35
```
"""
function count_total_bits(n::Int64)

    i = 0
    result = 0

    while n > (1 << i)

        k =0
        change = 1 << i

        for _ in 1:n+1
            result += k

            if change == 1
                # Flipping the bit
                if k == 0
                    k = 1
                else
                    k = 0
                end
                change = 1 << i

            else
                change -= 1
            end
        end
        i += 1

    end
    return result
end


function deflate_leading_zeros(ps::Vector{T}) where {T}
    ## trim any 0s from the end of ps
    N = findlast(!iszero, ps)
    K = findfirst(!iszero, ps)

    N == 0 && return(zeros(T,0), length(ps))
    ps = ps[K:N]
    ps, K-1
end

## take poly [p0, p1, ..., pn] and return
## [q_m-1, q_m-2, ..., q0], k
## where we trim of k roots of 0, and then make p monic, then reverese
## monomial x^5
function reverse_poly(ps::Vector{T}) where {T}
    # assume we have called deflate_leading_zeros
    qs = reverse(ps./ps[end])[2:end]
    qs
end

#
function quadratic_equation(a::T, b::T, c::T) where {T <: Real}
    qdrtc(a, -(0.5)*b, c)
end

## make more robust
function quadratic_equation(a::Complex{T}, b::Complex{T}, c::Complex{T}) where {T}
    d = sqrt(b^2 - 4*a*c)
    e1 = (-b + d)/(2a); e2 = (-b-d)/(2a)
    return (real(e1), imag(e1), real(e2), imag(e2))

end

## Kahan quadratic equation with fma
##  https://people.eecs.berkeley.edu/~wkahan/Qdrtcs.pdf

## solve ax^2 - 2bx + c
function qdrtc(a::T, b::T, c::T) where {T <: Real}
    # z1, z2 roots of ax^2 - 2bx + c
    d = discr(a,b,c)  # (b^2 - a*c), as 2 removes 4

    if d <= 0
        r = b/a  # real
        s = sqrt(-d)/a #imag
        return (r,s,r,-s)
    else
        r = sqrt(d) * (sign(b) + iszero(b)) + b
        return (r/a, zero(T), c/r, zero(T))
    end
end

## more work could be done here.
function discr(a::T,b::T,c::T) where {T}
    pie = 3.0 # depends on 53 or 64 bit...
    d = b*b - a*c
    e = b*b + a*c

    pie*abs(d) > e && return d

    p = b*b
    dp = muladd(b,b,-p)
    q = a*c
    dq = muladd(a,c,-q)

    (p-q) + (dp - dq)
end

##
## solve degree 2 or less case
function solve_simple_cases(ps::Vector{T}) where {T}
    S = T <: Complex ? T : Complex{T}

    N = length(ps)

    if N <= 1
        return S[]
    elseif N == 2
        return S[-ps[1]/ps[2]]
    elseif N == 3
        c,b,a = ps
        r1,i1, r2,i2 = quadratic_equation(a,b,c)
        S[complex(r1, i1), complex(r2, i2)]
    end
end


using Distributions
dists = [
    DiscreteUniform(10,20),
    Binomial(10,0.5),
    Geometric(0.5),
    NegativeBinomial(10,0.5),
    Hypergeometric(30, 40, 10),
    Poisson(5.5)]

println("Distribution \t\t\t\t\t\t Parameters \t Support")
reshape([dists ;  params.(dists) ;
		((d)->(minimum(d),maximum(d))).(dists) ],
		length(dists),3)


using SurrealNumbers
out_dir = "Data/"

# addition table
x = convert.(SurrealFinite, [0, 1/2, 1, 2])
A = zeros(SurrealFinite, length(x), length(x))
B = zeros(Int, length(x), length(x))
for i=1:length(x)
    for j=1:length(x)
        A[i,j] = x[i] + x[j]
        B[i,j] = generation(x[i]) + generation(x[j])
    end
end
float.(A)
A
generation.(A)
# A .== canonicalise.(A)
iscanonical.(A)
B

a10 = convert(SurrealFinite, 3/2)
a11 = convert(SurrealFinite, 1/2) + convert(SurrealFinite, 1)
a12 = convert(SurrealFinite, 3/4) + convert(SurrealFinite, 3/4)

file = "$(out_dir)addition_ex_10.dot"
FID = open(file, "w")
surreal2dag(FID, a10)
close(FID)
run(`dot -Tpdf -O $file`)

file = "$(out_dir)addition_ex_11.dot"
FID = open(file, "w")
surreal2dag(FID, a11)
close(FID)
run(`dot -Tpdf -O $file`)

file = "$(out_dir)addition_ex_12.dot"
FID = open(file, "w")
surreal2dag(FID, a12)
close(FID)
run(`dot -Tpdf -O $file`)


using JuMP, Gurobi, Distributions, StatsFuns, StatsBase

const gurobi_env = Gurobi.Env()

const numVariables = 100
const numJCC = 100
const conRHS = (numVariables*1.0)^2

xi_mean = zeros(Float64,numVariables)
xi_covariance = 0.5*ones(Float64,numJCC,numJCC)

for j = 1:numVariables
	xi_mean[j] = j*1.0/numVariables
end

for i = 1:numJCC
	xi_covariance[i,i] = 1.0
end

xi_cov_chol_tmp = cholfact(xi_covariance)
xi_cov_chol = xi_cov_chol_tmp[:L]

xi_mean_mat = zeros(Float64,numJCC,numVariables)
for j = 1:numVariables
	xi_mean_mat[:,j] = xi_mean[j]*ones(numJCC)
end

# tailored implementation of multivariate normal distribution
function myMvNormal(numSamples::Int64)

	xi = zeros(Float64,numJCC,numVariables,numSamples)

	xi_tmp = rand(Normal(0.0,1.0),numJCC,numVariables,numSamples)
	for samp = 1:numSamples
		xi[:,:,samp] = xi_cov_chol*xi_tmp[:,:,samp] + xi_mean_mat
	end

	return xi
end

# generate random samples from multivariate normal distribution
function generateRandomSamples(numSamples::Int64)

	xi = myMvNormal(numSamples)

	return xi
end

srand(1234)

# generate samples to estimate actual risk level of solution
const numAnalyticalSamples = 100000
const xi_analytical = generateRandomSamples(numAnalyticalSamples)

srand()

# compute conservative bound on the probability that
# constraints aren't satisfied
function boundRiskLevel(x::Array{Float64},reliabilityLevel::Float64=1E-06)

	const numSamples = size(xi_analytical,3)
	simpleEst::Bool = false

	numViolated::Int64 = checkScenarioConstraints(x,xi_analytical)

	riskLevel::Float64 = Inf

	if(simpleEst)
		riskLevel = numViolated*1.0/numSamples
		return riskLevel
	else
		if(numViolated == 0)
			riskLevel = -log(reliabilityLevel)/numSamples
			return riskLevel
		end

		 # termination criterion for bisection
		gamma_tolerance::Float64 = 0.1/numSamples

		# use tail bounds on binomial distribution to estimate lower and upper bounds
		gamma_low::Float64 = (numViolated - sqrt(-numSamples*log(reliabilityLevel)/2.0))/numSamples
		if(gamma_low < 0.0)
			gamma_low = 0.0
		end
		gamma_up::Float64 = (numViolated + sqrt(-numSamples*log(reliabilityLevel)/2.0))/numSamples
		if(gamma_up > 1.0)
			gamma_up = 1.0
		end

		gamma::Float64 = 0.0
		func_value::Float64 = 0.0

		# use bisection to solve the nonlinear equation
		while (gamma_up - gamma_low > gamma_tolerance)

			gamma = (gamma_low + gamma_up)/2.0
			func_value = -reliabilityLevel
			for i = 1:numViolated
				tmp_func = i*log(gamma) + (numSamples-i)*log(1-gamma) + lfact(numSamples) - lfact(numSamples-i) - lfact(i)
				func_value += exp(tmp_func)
				if(func_value > 0)
					break
				end
			end
			if(func_value > 0)
				gamma_low = (gamma_low + gamma_up)/2.0
			else
				gamma_up = (gamma_low + gamma_up)/2.0
			end

		end

		riskLevel = gamma_up
		return riskLevel
	end

end

# determines if all constraints are satisfied for a given single scenario
# returns one if scenario constraints are all satisfied, zero otherwise
function checkScenarioConstraints(x::Array{Float64},xi::Array{Float64})

	numSamples::Int64 = size(xi,3)

	numViolated::Int64 = 0
	for samp = 1:numSamples
		for i = 1:numJCC
			con::Float64 = evaluateConstraint(i,x,xi[i,:,samp])
			if(con > 0)
				numViolated += 1
				break
			end
		end
	end

	return numViolated
end

# evaluate the constraint values for given decision vector x,
# realization of the random variables xi, and constraint scalings
function evaluateConstraint(index::Int64,x::Array{Float64},xi::Array{Float64})

	con::Float64 = norm(xi.*x)^2 - conRHS

	return con
end

NumSamples = 50

# solution file name
const baseDirName = "C:/Users/rkannan/Desktop/SA for CCP/Scenario approximation/experiments/normopt_noniid_scenapprox/"

for samp = 1:NumSamples

	const dirName = baseDirName * string(samp) * "/"

	const fileName = dirName * "solution.txt"

	const outputFile = dirName * "trueRiskLevels.txt"
	open(outputFile, "w") do f
	end

	const timeFile = dirName * "trueRiskTime.txt"
	open(timeFile, "w") do f
	end

	soln = readdlm(fileName)

	const numPoints = size(soln,1)

	for iter = 1:numPoints
		tic()
		trueRiskLevel::Float64 = boundRiskLevel(soln[iter,:])
		riskTime = toq()
		println("Risk level #",iter," : ",trueRiskLevel,"  time: ",riskTime)
		open(outputFile, "a") do f
			write(f,"$trueRiskLevel \n")
		end
		open(timeFile, "a") do f
			write(f,"$riskTime \n")
		end
	end

end


function create_testfuncs(g::OneDGrid{Tg,<:Array}) where Tg
    g.nx > 8 || error("nx must be > 8")
    m = 5
    k₀ = g.k[2] # Fundamental wavenumber

    # Analytic function
    φ = π/3
    f₁ = @. cos(m*k₀*g.x + φ)

    # Transform of f₁ three ways
     f₁h = fft(f₁)
    f₁hr = rfft(f₁)
    f₁hr_mul = zeros(Complex{eltype(g.x)}, g.nkr)
    mul!(f₁hr_mul, g.rfftplan, f₁)

    # Analytical values of the fft and rfft
     f₁h_analytical = zeros(Complex{Float64}, size(f₁h))
    f₁hr_analytical = zeros(Complex{Float64}, size(f₁hr))

    for i in 1:g.nk
      if abs(real(g.k[i])) == m*k₀
        f₁h_analytical[i] = -exp(sign(real(g.k[i]))*im*φ)*g.nx/2
      end
    end

    for i in 1:g.nkr
      if abs(real(g.k[i])) == m*k₀
        f₁hr_analytical[i] = -exp(sign(real(g.kr[i]))*im*φ)*g.nx/2
      end
    end

    f₁, f₁h, f₁hr, f₁hr_mul, f₁h_analytical, f₁hr_analytical
end

function create_testfuncs(g::TwoDGrid{Tg,<:Array}) where Tg
    g.nx > 8 || error("nx must be > 8")
    m, n = 5, 2
    k₀ = g.k[2]
    l₀ = g.l[2]

    # Analytic functions
    f₁ = @. cos(m*k₀*g.x) * cos(n*l₀*g.y)
    f₂ = @. sin(m*k₀*g.x + n*l₀*g.y)

     f₁h = fft(f₁)
     f₂h = fft(f₂)
    f₁hr = rfft(f₁)
    f₂hr = rfft(f₂)

    f₁hr_mul = zeros(Complex{eltype(g.x)}, (g.nkr, g.nl))
    f₂hr_mul = zeros(Complex{eltype(g.x)}, (g.nkr, g.nl))
    mul!(f₁hr_mul, g.rfftplan, f₁)
    mul!(f₂hr_mul, g.rfftplan, f₂)

    # Theoretical results
     f₁h_analytical = zeros(Complex{eltype(g.x)}, size(f₁h))
     f₂h_analytical = zeros(Complex{eltype(g.x)}, size(f₂h))
    f₁hr_analytical = zeros(Complex{eltype(g.x)}, size(f₁hr))
    f₂hr_analytical = zeros(Complex{eltype(g.x)}, size(f₂hr))

    for j in 1:g.nl, i in 1:g.nk
      if ( abs(real(g.k[i])) == m*k₀ && abs(real(g.l[j])) == n*l₀ )
        f₁h_analytical[i, j] = - g.nx*g.ny/4
      end
      if ( real(g.k[i]) == m*k₀ && real(g.l[j]) == n*l₀ )
        f₂h_analytical[i, j] = -g.nx*g.ny/2
      elseif ( real(g.k[i]) == -m*k₀ && real(g.l[j]) == -n*l₀ )
        f₂h_analytical[i, j] = g.nx*g.ny/2
      end
    end
    f₂h_analytical = -im*f₂h_analytical;

    for j in 1:g.nl, i in 1:g.nkr
      if ( abs(g.kr[i])==m*k₀ && abs(g.l[j])==n*l₀ )
        f₁hr_analytical[i, j] = - g.nx*g.ny/4
      end
      if ( real(g.kr[i]) == m*k₀ && real(g.l[j]) == n*l₀ )
        f₂hr_analytical[i, j] = -g.nx*g.ny/2
      elseif ( real(g.kr[i]) == -m*k₀ && real(g.l[j]) == -n*l₀ )
        f₂hr_analytical[i, j] = g.nx*g.ny/2
      end
    end
    f₂hr_analytical = -im*f₂hr_analytical;

    f₁, f₂, f₁h, f₂h, f₁hr, f₂hr, f₁hr_mul, f₂hr_mul, f₁h_analytical, f₁hr_analytical, f₂h_analytical, f₂hr_analytical
end

@has_cuda begin
  function create_testfuncs(g::OneDGrid{Tg, <:CuArray}) where Tg
    cpugrid = OneDGrid(g.nx, g.Lx)
    out = create_testfuncs(cpugrid)
    return map(x->CuArray(x), out)
  end

  function create_testfuncs(g::TwoDGrid{Tg, <:CuArray}) where Tg
    cpugrid = TwoDGrid(g.nx, g.Lx, g.ny, g.Ly)
    out = create_testfuncs(cpugrid)
    return map(x->CuArray(x), out)
  end
end


losses = [[10, 0],
          [1, 50],
          [0, 200]]

num_actions = length(losses)

function expected_loss_of_action(prob_spam, action)
    #TODO: Return expected loss over a Bernoulli random variable
    #      with mean prob_spam.
    #      Losses are given by the table above.
    losses = [[10, 0],
              [1, 50],
              [0, 200]]
    loss = losses[action]
    length_prob_spam = length(prob_spam)

    non_prob = 1 .- prob_spam
    expect_loss = [prob_spam[i]*loss[1]+non_prob[i]*loss[2] for i in 1:length_prob_spam]
    return expect_loss
end

prob_range = range(0., stop=1., length=500)

using Plots
for action in 1:num_actions
  display(plot!(prob_range, expected_loss_of_action(prob_range, action)))
end

function optimal_action(prob_spam)
    #TODO: return best action given the probability of spam.
    # Hint: Julia's findmin function might be helpful.

    losses = [[10, 0],[1, 50],[0, 200]]
    non_prob_spam = 1-prob_spam
    expected_loss = []
    for i in 1:3
        loss = losses[i]
        push!(expected_loss, loss[1]*prob_spam+loss[2]*non_prob_spam)
    end
    return findmin(expected_loss)
end

prob_range = range(0, stop=1., length=500)
optimal_losses = []
optimal_actions = []
for p in prob_range
    # TODO:  Compute the optimal action and its expected loss for
    # probability of spam given by p.
    res = optimal_action(p)
    push!(optimal_losses,res[1])
    push!(optimal_actions,res[2])

end
optimal_actions
plot(prob_range, optimal_losses, linecolor=optimal_actions)


"""
```julia
randToeplitz(d, n;  norm, hermitian)
```
- `d` : entry distribution
- `n` : dimension
- `norm` : default `false`;  if `norm` set to `true`, then the matrix will be normlaized with ``n^{-1/2}``.
- `hermitian`: default `true`; if `true` the matrix will be Hermitian

# Examples

Generate a ``4 \\times 4`` random Hermitian Toeplitz matrix with entries Standard Normal.
```julia
randToeplitz(4)

4×4 Matrix{Float64}:
  1.10207   -0.47292   -0.745498   1.06809
 -0.47292    1.10207   -0.47292   -0.745498
 -0.745498  -0.47292    1.10207   -0.47292
  1.06809   -0.745498  -0.47292    1.10207
```
Generate a ``4 \\times 4`` normalized random Toeplitz matrix with entries `Exponential(1)`.
```julia
using Distributions
randToeplitz(Exponential(1),4, norm = true, hermitian = false)

4×4 Matrix{Float64}:
 0.667888  0.260045  1.48812   0.477305
 1.50374   0.667888  0.260045  1.48812
 1.1475    1.50374   0.667888  0.260045
 0.363966  1.1475    1.50374   0.667888
```
"""
function  randToeplitz(d::D, n::Int;  norm = false::Bool, hermitian=true::Bool)  where D<:S
    if hermitian
      t = rand(d,n)
      t = [t[n:-1:2]'..., t...]
    else
      t = rand(d,2n-1)
    end
    M = zeros(eltype(d),n,n)
    for i in 1:n,j in 1:n
        M[i,j]=t[i-j+n]
    end

    if norm  return M/sqrt(n) end
    M

end

function randToeplitz(n::Int; norm=false::Bool, hermitian=true::Bool)
  return randToeplitz(Normal(),n,norm=norm,hermitian=hermitian)
end

"""
```julia
randHankel(d, n;  norm )

randHankel(n;  norm)
```
- `d` : entry distribution
- `n` : dimension
- `norm` : default `false`,  if `norm` set to `true`, then the matrix will be normlaized with ``n^{-1/2}``.

# Examples

Generate a ``5\\times 5`` random Hankel matrix with entries uniformly distributed on ``\\{1, i, \\pi \\}``
```julia
randHankel((1,im,pi),5)

5×5 Matrix{Number}:
  1   1  im  1   1
  1  im   1  1   π
 im   1   1  π   π
  1   1   π  π   π
  1   π   π  π  im
```

"""
function randHankel(d::D, n::Int;  norm = false::Bool)  where D<:S
  h = rand(d,2n-1)
  M = zeros(eltype(d),n,n)
  for i in 1:n
    M[i,:] = h[i:i+n-1]
  end

  if norm  return M/sqrt(n) end
  M
end

function randHankel(n::Int;  norm = false::Bool)
  return randHankel(Normal(),n,norm=norm)
end