using ImputationQualityScores
using Base.Test

@test dosage_to_likelihoods(0.) == (1.,0.,0.)
@test dosage_to_likelihoods(1.) == (0.,1.,0.)
@test dosage_to_likelihoods(2.) == (0.,0.,1.)

@test likelihoods_to_dosage(1.,0.,0.) == (0.)
@test likelihoods_to_dosage(0.,1.,0.) == (1.)
@test likelihoods_to_dosage(1.,0.,1.) == (2.)

for i in 0:0.01:2
	@test i ≈ likelihoods_to_dosage(dosage_to_likelihoods(i))
end

@test proper_info([(0.3333333, 0.3333333, 0.3333333), (1.0, 0.0, 0.0)])≈0.428571367347


using Random, LinearAlgebra
using Stheno: WoodburyMat, Xt_invA_X, LazyPDMat

@testset "woodbury" begin

    let
        rng, N, D = MersenneTwister(123456), 10, 5
        X, Σ_, σ = randn(rng, D, N), randn(rng, D, D), exp(randn(rng))
        Σ = LazyPDMat(Σ_ * Σ_' + 1e-9I, 0)
        A = WoodburyMat(X, Σ, σ)
        B, b = randn(N, 2), randn(N)

        # Test that conversion and indexing are consistent with one another.
        Am, Al = Matrix(A), LazyPDMat(A)
        @test size(A) == (N, N)
        @test Matrix(Al) == Am
        @test Am[2, 3] ≈ A[2, 3]
        @test Am[3, 2] ≈ A[3, 2]

        # Check unary operations.
        @test logdet(A) ≈ logdet(Matrix(A))

        # Test basics of matrix multiplication.
        @test Matrix(A * B) ≈ Matrix(A) * B
        @test Vector(A * b) ≈ Matrix(A) * b
        Bt = B'
        @test Matrix(Bt * A) ≈ B' * Matrix(A)

        # Test basics of linear systems solving.
        @test Matrix(A \ B) ≈ Matrix(A) \ B
        @test Vector(A \ b) ≈ Matrix(A) \ b

        @test Matrix(B' / A) ≈ B' / Matrix(A)

        # Test quadratic form solving.
        @test Xt_invA_X(A, B) isa Matrix{<:Real}
        @test Xt_invA_X(A, B) ≈ B' * (Matrix(A) \ B)

        @test Xt_invA_X(A, b) isa Real
        @test Xt_invA_X(A, b) ≈ b' * (Matrix(A) \ b)
    end

end


include("L_DVBN.jl")

f = readcsv("data_auto_mpg.csv")

data = Array(Any,392,8)

discrete_index = [2,7,8]
continuous_index = [1,3,4,5,6]

for i = 1 : 392
    for j = 1 : 8
        if j in discrete_index
            data[i,j] = round(Int64,f[i,j])
        else
            data[i,j] = f[i,j]
        end
    end
end

##### Network Structure is known in advance #####
graph = [2,(2,3),(3,5),(5,1),(1,5,3,7),(3,4),(4,6),8]
order_continuous_order = graph_to_reverse_conti_order(graph,continuous_index)
cut_time = 8

e1 = Discretize_All(data,graph,continuous_index,cut_time)
##### Network structure is not known in advance #####

u = 2
cut_time = 5
times = 3

order = [1,2,3,4,5,6,7,8]
e2 = Learn_DVBN(data,continuous_index,order,u,cut_time)

e3 = Learn_Discrete_Bayesian_Net(data,continuous_index,u,cut_time,times)


using PyPlot
using WignerFamilies
using HalfIntegers
# compute all j₂=100, j₃=60, m₂=70, m₃=-55, m₁=-m₂-m₃
# bottom of page 14 of Raynal et al.
# On the zeros of 3j coefficients: polynomial degree versus recurrence order

X₀, Y₀ = 43, 25
# X₀, Y₀ = 109, 63
X(n) = iszero(n) ? X₀ : 7X(n-1) + 12Y(n-1)
Y(n) = iszero(n) ? Y₀ : 4(n-1) + 7Y(n-1)
n_seq = 2
(2X(n_seq)+1)/6, (X(n_seq)+2)/6
##
# j₂, j₃, m₂, m₃ = 100, 60, 70, -55
j₂, j₃, m₂, m₃ = (2X(n_seq)+1)/6, (X(n_seq)+2)/6, HalfInt(3/2), HalfInt(3/2)
# j₂, j₃, m₂, m₃ = 4, 4, 2, -2

# j₂, j₃, m₂, m₃ = HalfInt(23/2), HalfInt(21/2), HalfInt(3/2), HalfInt(3/2)

m₁ = -m₂ - m₃

w = WignerF(BigFloat, j₂, j₃, m₂, m₃)

nmid = Int(ceil((w.nₘᵢₙ + w.nₘₐₓ)/2))
nmid = iseven(nmid) ? nmid : nmid + 1

w3j = get_wigner_array(w)

n₊ = WignerFamilies.rψ!(w, nmid, w3j)
n₋ = WignerFamilies.sψ!(w, nmid, w3j)

for k in (n₊+1):w.nₘₐₓ  # product the ratios upwards
    w3j[k] *= w3j[k-1]
end
for k in (n₋-1):-1:w.nₘᵢₙ  # product the ratios downwards
    w3j[k] *= w3j[k+1]
end

ψn₋ = w3j[n₋]
ψn₊ = w3j[n₊]

WignerFamilies.ψauxplus!(w, n₋, n₊, w3j)

if n₋ > w.nₘᵢₙ
    scale_middle = ψn₋ / w3j[n₋]
    w3j[w.nₘᵢₙ:(n₋+1)] .*= scale_middle
end
if n₊ < w.nₘₐₓ
    scale_end = w3j[n₊] / ψn₊
    w3j[(n₊+1):w.nₘₐₓ] .*= scale_end
end

norm = WignerFamilies.normalization(w, w3j)
w3j ./= norm
if sign(w3j[w.nₘₐₓ]) != WignerFamilies.f_jmax_sgn(w)
    w3j .*= -1
end

##

plt.clf()
using WignerSymbols
exact = [wigner3j(BigFloat, n, w.j₂, w.j₃, m₁, m₂) for n in w.nₘᵢₙ:w.nₘₐₓ]
plt.plot(w.nₘᵢₙ:w.nₘₐₓ, Float64.(w3j))

plt.plot(w.nₘᵢₙ:w.nₘₐₓ, Float64.(exact), alpha=0.5)
plt.xlabel(raw"$j$")
plt.ylabel(raw"$f(j) \, \times \, 10^3$")
plt.tight_layout()
# plt.xlim(100,120)
gcf()

##
plt.clf()
plt.plot( Float64.(exact .- w3j.parent))
plt.gcf()

##
w = WignerF(BigFloat, 100, 100, 0, 0)
w3j = get_wigner_array(w)
@time WignerFamilies.classical_wigner3j_m0!(w, w3j)
##


### A Pluto.jl notebook ###
# v0.12.10

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 728d5546-7a96-11eb-1757-93681c4e7955
begin
	import Pkg
    Pkg.activate(".")
    Pkg.add("PlutoUI")
	Pkg.add("Plots")

	using PlutoUI
	using Plots
end

# ╔═╡ 58ba05d0-7a96-11eb-08da-41b782596598
md"# Euler Method"

# ╔═╡ 97f8a9f8-7b78-11eb-2f25-e5228f41db56
md"We implement a Euer Method to solve first oderd ODEs and use it to some relevant Biological Models."

# ╔═╡ ecad4e78-7a96-11eb-2b54-e7f910e83835
begin
	function euler(f,t0,tn,Δt,x₀)
		T = t0:Δt:tn
		F = Vector{typeof(x₀)}(undef,length(T))
		F[1] = x₀
		for (n,t) ∈ enumerate(T[2:end])
			F[n+1] = F[n] .+ f(F[n]) .* Δt
		end
		return F
	end

	close(f,par) = x -> f(x,par...)

	reshape_result(F) = [[f[i] for f ∈ F] for i ∈ 1:length(F[1])] ;
end

# ╔═╡ cd0f0650-7b78-11eb-12c0-cb427957850b
md"""
---
### Logistic Population Growth
"""

# ╔═╡ e9b52262-7b7d-11eb-21aa-2bf0d9f5ef50
md"""

$\frac{dP}{dt} = rP \left( 1- \frac{P}{K} \right)$

"""

# ╔═╡ dcffb046-7b78-11eb-333c-7df614c12617
begin
	LogisticGrowth(x,K,r) = r*x*(1-x/K)
	LG_0 = 1.0
end;

# ╔═╡ 04fdd000-7b79-11eb-029b-b3765364f627
md"""
growth rate r $(@bind r NumberField(0.0001:0.0001:0.1,default=0.01))

carrying capacity K $(@bind K NumberField(1:10^7,default=10_000))
"""

# ╔═╡ 5b3c4954-7b79-11eb-11ba-57b8ae51a6f5
let
	Δt = 0.01

	p = plot(framestlye=:zerolines,legend=:topleft)

	T=0:Δt:2500
	F = euler(close(LogisticGrowth,[K,r]),T[1],T[end],Δt,LG_0)
	plot!(p,T,F,label="Population Size")

	p
end

# ╔═╡ 7a010972-7b78-11eb-3478-f1fbc1f0645d
md"""
---
### Lotka-Volterra
"""

# ╔═╡ b963fdb2-7b7e-11eb-3184-1b216f77bd64
md"""

${\frac{dN_1}{dt} = N_1 (\varepsilon_1 - \gamma_1 N_2) \quad , \quad \frac{dN_2}{dt} = -N_2 (\varepsilon_2 - \gamma_2 N_1)}$

"""

# ╔═╡ 7fc751be-7a96-11eb-207d-d5d0142f43ba
begin
	LotkaVolterra(x,ε₁,δ₁,ε₂,δ₂) = [
			x[1]*(ε₁-δ₁*x[2]),
			-x[2]*(ε₂-δ₂*x[1])
	]
	LV_0 = [10.0,10.0]
end;

# ╔═╡ 07b392ee-7b77-11eb-00be-876772164be8
md"""
Reproduktionsraten
	ε₁ $(@bind ε₁ Slider(0.01:0.01:1,show_value=true, default=0.5))
	ε₂ $(@bind ε₂ Slider(0.01:0.01:1,show_value=true, default=0.25))

Sterberaten
	δ₁ $(@bind δ₁ Slider(0.01:0.01:1,show_value=true, default=0.07))
	δ₂ $(@bind δ₂ Slider(0.01:0.01:1,show_value=true, default=0.07))
"""

# ╔═╡ c080dd1a-7a96-11eb-345c-e364dc0f8ac9
let
	Δt = 0.01

	p = plot(framestlye=:zerolines)

	T=0:Δt:100
	F = euler(close(LotkaVolterra,[ε₁,δ₁,ε₂,δ₂]),T[1],T[end],Δt,LV_0)
	plot!(p,T,reshape_result(F),label=["N₁" "N₂"])

	p
end

# ╔═╡ 62645b8c-7b7a-11eb-1b35-a3957a0252fb
md"""
---
### SIR-Model
"""

# ╔═╡ 4341adf4-7b7f-11eb-25c1-e546febcad3e
md"""

${\frac{dS}{dt} = \nu N - \beta \frac{SI}{N} - \mu S}$

${\frac{dI}{dt} = \beta \frac{SI}{N} - \gamma I - \mu I}$

${\frac{dR}{dt} = \gamma I - \mu R \phantom{-\beta \frac{SI}{N}}}$

"""

# ╔═╡ 6b37eb20-7b7a-11eb-07de-a1042efa6721
begin
	function SIR(x,ν,β,μ,γ)
		N = sum(x)
		S, I, R = x[1], x[2], x[3]
		return [
			(ν*N - β*S*I/N - μ*S),	# dS/dt
			(β*S*I/N - γ*I - μ*I),	# dI/dt
			(γ*I - μ*R)				# dR/dt
		]
	end

	SIR_0 = [9999.9,0.1,0.0]

	Δt = 0.01
	T=0:Δt:150
end;

# ╔═╡ e64e9d62-7b7b-11eb-3620-a7a57b954fcc
md"""
recovery rate γ $(@bind γ Slider(0.0:0.01:1,show_value=true,default=0.15)) |
infection rate β $(@bind β Slider(0.0:0.01:1,show_value=true,default=0.7))

death rate μ $(@bind μ Slider(0.0:0.01:1,show_value=true)) |
birth rate ν $(@bind ν Slider(0.0:0.01:1,show_value=true))
"""

# ╔═╡ fa94e89e-7c1e-11eb-08a7-91210bc2fe60
begin
	F = euler(close(SIR,[ν,β,μ,γ]),T[1],T[end],Δt,SIR_0)
	rF = reshape_result(F)
	N = sum.(F)
end;

# ╔═╡ 68f66db2-7b7c-11eb-138d-2f40b3bfbb4e
let
	p = plot(framestlye=:zerolines,size=(680,450))

	plot!(p,T,rF,label=["S" "I" "R"])

	p
end

# ╔═╡ Cell order:
# ╟─58ba05d0-7a96-11eb-08da-41b782596598
# ╟─97f8a9f8-7b78-11eb-2f25-e5228f41db56
# ╠═ecad4e78-7a96-11eb-2b54-e7f910e83835
# ╟─cd0f0650-7b78-11eb-12c0-cb427957850b
# ╟─e9b52262-7b7d-11eb-21aa-2bf0d9f5ef50
# ╠═dcffb046-7b78-11eb-333c-7df614c12617
# ╟─04fdd000-7b79-11eb-029b-b3765364f627
# ╠═5b3c4954-7b79-11eb-11ba-57b8ae51a6f5
# ╟─7a010972-7b78-11eb-3478-f1fbc1f0645d
# ╟─b963fdb2-7b7e-11eb-3184-1b216f77bd64
# ╠═7fc751be-7a96-11eb-207d-d5d0142f43ba
# ╟─07b392ee-7b77-11eb-00be-876772164be8
# ╠═c080dd1a-7a96-11eb-345c-e364dc0f8ac9
# ╟─62645b8c-7b7a-11eb-1b35-a3957a0252fb
# ╟─4341adf4-7b7f-11eb-25c1-e546febcad3e
# ╠═6b37eb20-7b7a-11eb-07de-a1042efa6721
# ╠═fa94e89e-7c1e-11eb-08a7-91210bc2fe60
# ╟─e64e9d62-7b7b-11eb-3620-a7a57b954fcc
# ╟─68f66db2-7b7c-11eb-138d-2f40b3bfbb4e
# ╟─728d5546-7a96-11eb-1757-93681c4e7955


"""
    Restriction(base_freq::Vector{Float64},
                SubstitutionRates::Vector{Float64})::Tuple{Array{Float64,2}, Array{Float64,1}, Array{Float64,2}}

Calculate the eigenvalue decomposition of the Q matrix of the restriction site model.
The `SubstitutionRates` are ignored, and just for call stability.

The function returns the Eigenvectors, Eigenvalues, inverse of eigenvectors and
    the scale factor for expected number changes per site
"""
function Restriction(base_freq::Vector{Float64}, SubstitutionRates::Vector{Float64})::Tuple{Array{Float64,2}, Array{Float64,1}, Array{Float64,2}, Float64}
    Nbases = length(base_freq)
    Q::Array{Float64,2} = ones(Nbases,Nbases)
    Q[diagind(Nbases,Nbases)] .= -1
    Q .*= reverse(base_freq)
    D, U = eigen(Q)
    Uinv = inv(U)
    mu::Float64 =  1.0 / (2.0 * prod(base_freq))
    return U, D, Uinv, mu
end

"""
    JC(base_freq::Vector{Float64},
        SubstitutionRates::Vector{Float64})::Tuple{Array{Float64,2}, Array{Float64,1}, Array{Float64,2}}

Calculate the eigenvalue decomposition of the Q matrix of the Jukes-Cantor model.
The `SubstitutionRates` are ignored, and just for call stability.

The function returns the Eigenvectors, Eigenvalues and inverse of eigenvectors.
"""
function JC(base_freq::Vector{Float64}, SubstitutionRates::Vector{Float64})::Tuple{Array{Float64,2}, Array{Float64,1}, Array{Float64,2}, Float64}
    Nbases = length(base_freq)
    Q::Array{Float64,2} = ones(Nbases,Nbases)
    μ = SubstitutionRates[1]
    off_diag = μ/Nbases
    diag = off_diag * (Nbases-1)
    Q .= off_diag
    Q[diagind(Nbases,Nbases)] .= -diag
    D, U = eigen(Q)
    Uinv = inv(U)
    mu = 1.0
    return U, D, Uinv, mu
end

"""
    GTR(base_freq::Vector{Float64},
        SubstitutionRates::Vector{Float64})::Tuple{Array{Float64,2}, Array{Float64,1}, Array{Float64,2}}

Calculate the eigenvalue decomposition of the Q matrix of the General Time Reversible model.

The function returns the Eigenvectors, Eigenvalues and inverse of eigenvectors.
"""
function GTR(base_freq::Vector{Float64}, SubstitutionRates::Vector{Float64})::Tuple{Array{Float64,2}, Array{Float64,1}, Array{Float64,2}, Float64}
    Nbases = length(base_freq)
    Q::Array{Float64,2} = setmatrix(SubstitutionRates)
    Q = mapslices(x-> x.*base_freq, Q, dims=2)
    dia = sum(Q, dims=2)
    Q[diagind(Nbases,Nbases)] = -dia
    D, U = eigen(Q)
    Uinv = inv(U)
    return U, D, Uinv, 1/sum(dia)
end

"""
    function freeK(base_freq::Vector{Float64},SubstitutionRates::AbstractArray)::Tuple{Array{N,2}, Array{N,1}, Array{N,2}, M} where {N <: Number, M <: Number}

FreeK model of substitution.

"""
function freeK(
      base_freq::Vector{Float64},
      SubstitutionRates::AbstractArray,
      )::Tuple{Array, Array, Array, Float64}# where {N <: Number, M <: Number}
      Nrates = length(SubstitutionRates)
      Nbases = Int(ceil(sqrt(Nrates)))
      Q = zeros(Nbases, Nbases)
      counter = 1
      for i in 1:Nbases, j in 1:Nbases
            if i!=j
                  Q[j,i] = SubstitutionRates[counter]
                  counter += 1
            end
      end
      dia = sum(Q, dims = 2)
      Q[diagind(Nbases, Nbases)] = -dia
      D, U = eigen(Q)
      Uinv = inv(U)
      return U, D, Uinv, 1 / sum(dia)
end

## Helper Functions ##
function setmatrix(vec_vals::Array{T,1})::Array{T,2} where T
    n = length(vec_vals)
    s = round((sqrt(8n+1)+1)/2)
    s*(s-1)/2 == n || error("setmatrix: length of vector is not triangular")
    k = 0
    Q = [ i<j ? (k+=1; vec_vals[k]) : 0 for i=1:s, j=1:s ]
    Q.+transpose(Q)
end


function y_prob(X1::Array{T,3}, X2::Matrix{T}, D::Matrix{T}, α::Vector{T}, Π::Matrix{T},
    ξ_all::Vector{T}) where T
    N, J, K1, K2, L = get_dim(X1, X2, D)
    @assert length(ξ_all) == J
    y_prob = zeros(T, N, J)
    for (i, d) in enumerate(eachrow(D))
        index = exp.(X2*(Π*d) .+ X1[i,:,:]*α .+ ξ_all) # J-by-1
        y_prob[i,:] = index ./ sum(index)
    end
    return y_prob
end

# TODO: need to add a y_prob for hetero preference y_prob

"""
Sample likelihood when consumers face homogeneous choice set, no unobservable preference heterogeneity
"""
function llk_sample_fullset(Y::Vector, X1::Array, X2::Matrix, D::Matrix,
    α::Vector{T}, Π::Matrix{T}, ξ::Vector{T}) where T

    N, J, K1, K2, L = get_dim(X1, X2, D)
    ξ_all = vcat(zero(eltype(ξ)), ξ)
    @assert length(ξ_all) == J

    llk = 0.0
    for (i, d) in enumerate(eachrow(D))
        index = exp.(X2*(Π*d) .+ X1[i,:,:]*α .+ ξ_all) # J-by-1
        y_prob = index ./ sum(index)
        llk += log(y_prob[Y[i]])
    end
    return llk
end

"""
Sample likelihood when consumers face different choice sets, no unobservable preference heterogeneity

C::Matrix, N-by-J dummy matrix, determines the choice set of each individual
    This can be actual choice set restrictions, or purely random choice sets determined to facilitate estimation
"""
function llk_sample_subset(Y::Vector, X1::Array, X2::Matrix, D::Matrix, C::Vector{Vector{Int64}},
    α::Vector{T}, Π::Matrix{T}, ξ::Vector{T}) where T

    N, J, K1, K2, L = get_dim(X1, X2, D)
    ξ_all = vcat(zero(eltype(ξ)), ξ)
    @assert length(ξ_all) == J

    llk = 0.0
    for (i, d) in enumerate(eachrow(D))
        choice_id = findfirst(isequal(Y[i]), C[i])
        index = @views exp.(X2[C[i],:]*(Π*d) .+ X1[i,C[i],:]*α .+ ξ_all[C[i]]) # J-by-1
        y_prob = index ./ sum(index)
        llk += log(y_prob[choice_id])
    end
    return llk
end

function llk_sample(Y::Vector, X1::Array, X2::Matrix, D::Matrix,
    α::Vector{T}, Π::Matrix{T}, ξ::Vector{T}; subset = nothing) where T

    if isnothing(subset)
        llk = llk_sample_fullset(Y, X1, X2, D, α, Π, ξ)
    else
        llk = llk_sample_subset(Y, X1, X2, D, subset,  α, Π, ξ)
    end
    return llk
end

"""
Sample likelihood when there are unobservable preference heterogeneity
"""
function llk_sample(Y::Vector, X1::Array, X2::Matrix, D::Matrix,
    α::Vector{T}, Π::Matrix{T}, ξ::Vector{T}, σ::Vector{T}, n_sim::Integer) where T

    N, J, K1, K2, L = get_dim(X1, X2, D)
    ξ_all = vcat(zero(eltype(ξ)), ξ)
    @assert length(ξ_all) == J

    llk = 0.0
    for (i, d) in enumerate(eachrow(D))
        y_prob = zeros(T, n_sim, J)
        for i_sim = 1:n_sim
            index = exp.(X2*(Π*d .+ randn(K).*σ) .+ X1[i,:,:]*α .+ ξ_all) # J-by-1
            y_prob[i_sim,:] = index ./ sum(index)
        end
        y_prob = mean(y_prob, dims=1)
        llk += log(y_prob[Y[i]])
    end
    return llk
end

"""
CCP approach for homogeneous preference

    Note that the CCP method will only work for large number of commodities.
    Also, we will need IV for ξ to control endogeneity
"""
function ccp(Y::Vector, X::Matrix)
    # Y: N-by-1
    # X: J-by-K
    J, K = size(X)
    P = proportionmap(Y)
    δ = [P[j]-P[1] for j in 2:J]
    DX = zeros(J-1,K)
    for j = 2:J-1
        DX[j,:] = X[j,:] .- X[1,:]
    end
    β_hat = (DX'*DX)\(DX'*δ)
    return β_hat
end

function opt_homo(Y, X1, X2, D, init_guess; subset = nothing)
    N, J, K1, K2, L = get_dim(X1, X2, D)
    if isnothing(subset)
        f = para -> -llk_sample_fullset(Y, X1, X2, D,
            para[1:K1], # α
            reshape(para[K1+1:K1+K2*L], K2, L), # Π
            para[K1+K2*L+1:end]) # ξ
            # take negative to minimize
    else
        f = para -> -llk_sample_subset(Y, X1, X2, D, subset,
            para[1:K1], # α
            reshape(para[K1+1:K1+K2*L], K2, L), # Π
            para[K1+K2*L+1:end]) # ξ
            # take negative to minimize
    end
    opt = Optim.optimize(f, init_guess, LBFGS(); autodiff=:forward)
    return opt
end

function opt_hetero(Y, X, init_guess, n_sim)
    J, K = size(X)
    N = length(Y)
    f = para -> -llk_sample(Y, X, para[1:K], para[K+1:K+J-1], para[K+J:end]; n_sim=n_sim) # take negative to minimize, ξ only J-1 dim
    # g = TwiceDifferentiable(f, init_guess; autodiff=:forward)
    opt = Optim.optimize(f, init_guess, LBFGS(); autodiff=:forward)
    return opt
end

function estDCModel!(model::DCModel; init_guess::Vector, n_sim=5, subset=nothing)
    if !model.hetero_preference
        model.optResults = opt_homo(model.Y, model.X1, model.X2, model.D, init_guess; subset=subset)
    else
        model.optResults = opt_hetero(model.Y, model.X, init_guess, n_sim)
    end
    return model
end


include("operators.jl")

using GaussQuadrature
using OrdinaryDiffEq
using LinearAlgebra
using UnicodePlots

struct SpectralOP{T}
  ξ  :: Array{T, 1}
  ω  :: Array{T, 1}
  D  :: Array{T, 2}
  e0 :: Array{T, 1}
  eN :: Array{T, 1}
  function SpectralOP{T}(N) where T
    Nq = N + 1
    x, ω = legendre(T, Nq, both)
    D = derivative(x)
    e0 = [1; zeros(T, N)]
    eN = [zeros(T, N); 1]
    new{T}(x, ω, D, e0, eN)
  end
end

function dgodefun!(dq, q, p, t)
  o = p.o
  f = p.f
  Nq = length(o.x)
  Ne = div(q, 2Nq)
  v  = @views reshape(q[1:Nq*Ne], Nq, Ne)
  σ  = @views reshape(q[N+1:2Nq*Ne], Nq, Ne)
  dv = @views reshape(dq[1:Nq*Ne], Nq, Ne)
  dσ = @views reshape(dq[N+1:2Nq*Ne], Nq, Ne)

  D = o.D
  ω = o.ω

  # Volume terms (skew form)
  dv .= -D' * (ω .* σ)
  dσ .= ω .* (D * v)

  for e in 1:Ne
    # face 1 (left face)
    vm, σm = v[1, e], σ[1, e]
    vp, σp = begin
      if e == 1
      else
        v[Nq, e-1], σ[Nq, e-1]
      end
    end

    # face 2 (left face)
    vm, σm = v[Nq, e], σ[Nq, e]
    vp, σp = begin
      if e == Ne
      else
        v[1, e-1], σ[1, e-1]
      end
    end
  end

  e0 = o.e0
  dv .= (-D' * (ω .* D * u) - e0 * f(v[1])) ./ ω
  du .= v
end

let
  N = 100
  T = Float64

  tspan = (T(0), T(3//10))

  Nq = N + 1
  o = SpectralOP{T}(N)
  q = [zeros(T, Nq);sin.(2o.x * π)]
  f(v) = asinh(v)

  p = (o=o, f=f)
  prob = ODEProblem(spectralodefun!, q, tspan, p)
  sol = solve(prob, Tsit5();
              saveat=range(tspan[1], stop=tspan[2], length = 4),
              atol = 1e-6, rtol = 1e-3,
              internalnorm=(x, _)->norm(x, Inf))

  for n = 1:length(sol.t)
    display(lineplot(o.x, sol.u[n][Nq+1:2Nq], width=100, xlim=(0,1),
                     title="t = $(sol.t[n])"))
  end
end


# conditional neighbor probability for generally asymmetric SNE
# generalizes kernel and distance
function conditional_neighbor(X::AbstractMatrix, σ²::AbstractVector{<:Real},
                    kernel = gaussian, distance = euclidean)
    x = [c for c in eachcol(X)]
    xt, vt = permutedims(x), permutedims(σ²)
    K = @. kernel(distance(x, xt) / sqrt(vt)) # lazy?
    for i in diagind(K) # "Because we are only interested in modeling pairwise similarities, we set the value of pi|i to zero."
        K[i] = 0
    end
    sumK = sum(K, dims = 1) # this is an O(n²) operation!
    @. K /= sumK # columns are normalized probability distributions now
end

# adjusts variances of conditional neighbor distribution to achieve perplexity u
function optimize_variances(X::AbstractMatrix, u::Real, kernel = gaussian, distance = euclidean;
                            tol::Real = 1e-6, max_iter::Int = 32)
    n = size(X, 2)
    σ² = fill(10one(eltype(X)), n) # if this is too small initially, Newton has problems
    @threads for i in 1:n
        σ²[i] = optimize_variances(X, u, i, σ²[i], kernel, distance, tol = tol, max_iter = max_iter)
    end
    any(isnan, σ²) && throw(DomainError("Warning: variance optimization yielded NaN: $σ²"))
    return σ²
end

# optimizing variances using dense operations
function optimize_variances(X::AbstractMatrix, u::Real, i::Int, σ²::Real,
                            kernel = gaussian, distance = euclidean;
                            tol::Real = 1e-6, max_iter::Int = 32)
    x = [c for c in eachcol(X)]
    n = length(x)
    function objective(log_σ²::Real) # allocating version, is faster for high d, equal for low
        local σ² = exp(log_σ²)
        p = @. kernel(distance(x[i:i], x) / sqrt(σ²)) # this has sparse approximation (see sparse.jl)
        p[i] = 0
        p ./= sum(p)
        entropy(p) - log2(u)
    end
    log_σ² = newtons_method(objective, log(σ²), tol = tol, max_iter = max_iter)
    return exp(log_σ²)
end


function perform_wavelet_transf(f, Jmin, dir, filter = "9-7",separable = 0, ti = 0)

    """""
    perform_wavelet_transf - peform fast lifting transform

    y = perform_wavelet_transf(x, Jmin, dir, filter = "9-7",separable = 0, ti = 0);

    Implement 1D and 2D symmetric wavelets with symmetric boundary treatements, using
    a lifting implementation.

    filter gives the coefficients of the lifting filter.
    You can use h='linear' or h='7-9' to select automatically biorthogonal
    transform with 2 and 4 vanishing moments.

    You can set ti=1 to compute a translation invariant wavelet transform.

    You can set separable=1 to compute a separable 2D wavelet
    transform.

    Copyright (c) 2008 Gabriel Peyre
    """

    #copy f
    x = copy(f)

    #convert Jmin to int
    Jmin = Int(Jmin)

    # detect dimensionality
    d = ndims(x)
    # P/U/P/U/etc the last coefficient is scaling
    if filter in ["linear", "5-3"]
        h = [1/2, 1/4, sqrt(2)]

    elseif filter in ["9-7", "7-9"]
        h = [1.586134342, -.05298011854, -.8829110762, .4435068522, 1.149604398]

    else
        error("Unknown filter")

    end

    if (d == 2) & (separable == 1)
        ti = 0
        if ti == 1
            warn("Separable does not works for translation invariant transform")

        end

        # perform a separable wavelet transform
        n = size(x)[1]
        if dir == 1
            for i in 1 : n
                x[:, i] = perform_wavelet_transf(x[:, i], Jmin, dir, filter, separable, ti)
            end
            for i in 1 : n
                x[i, :] = transpose(perform_wavelet_transf(transpose(x[i,:]), Jmin, dir, filter, separable, ti))
            end
        else
            for i in 1 : n
                x[i, :] = transpose(perform_wavelet_transf(transpose(x[i, :]), Jmin, dir, filter, separable, ti))
            end
            for i in 1 : n
                x[:, i] = perform_wavelet_transf(x[:,i], Jmin, dir, filter, separable, ti)
            end
        end
    end

    # number of lifting steps
    if ndims(x) == 1
        n = length(x)
    else
        n = size(x)[2]
    end
    m = Base.div(length(h) - 1, 2)    #Base.div because I think div would be deprecated in the tutoriel
    Jmax = Int(log2(n) - 1)
    jlist = Jmax:-1:Jmin

    if dir == -1
        jlist = Jmin:1:Jmax
    end

    if ti == 0
        # subsampled
        for j in jlist
            if d == 1
                x[1:2^(j+1), :] = lifting_step(x[1:2^(j+1), :], h, dir)    #There may be a problem with x[1:2^(j+1)]
            else
                x[1:2^(j+1), 1:2^(j+1)] = lifting_step(x[1:2^(j+1), 1:2^(j+1)], h, dir)
                x[1:2^(j+1), 1:2^(j+1)] = transpose(lifting_step(transpose(x[1:2^(j+1), 1:2^(j+1)]), h, dir))
            end
        end

    else
        # TI
        nJ = Jmax - Jmin + 1  #It is maybe nJ = Jmax - Jmin
        if (dir == 1) & (d == 1)
            x = repeat(x, outer = [1, 1, nJ + 1])   #There is a slight difference with python's script, here I repeat the array with respect to the 3rd dimension.
        elseif (dir == 1) & (d == 2)
            x = repeat(x, outer = [1, 1, 3*nJ + 1])
        end
        #elif dir == 1:
        #    x = np.tile(x,(1,1,1))
        for j in jlist
            dist = 2^(Jmax - j)

            if d == 1
                if dir == 1
                    x[:,:,1:(j - Jmin + 2)] = lifting_step_ti(x[ :, :, 1], h, dir, dist)
                else
                    x[:, :, 1] = lifting_step_ti(x[:,:,1:(j - Jmin + 2)], h, dir, dist)
                end
            else
                dj = 3*(j - Jmin)

                if dir == 1
                    x[:, :, [1, dj + 2]] = lifting_step_ti(x[:, :, 1], h, dir, dist)

                    x[:, :,[1, dj + 3]] = lifting_step_ti(transpose(x[:, :,1]), h, dir, dist)
                    x[:, :, 1] = transpose(x[:, :, 1])
                    x[:, :, dj + 3] = transpose(x[:, :, dj + 3])

                    x[:, :, [2 + dj, 4 + dj]] = lifting_step_ti(transpose(x[:,:, dj + 2]), h, dir, dist)
                    x[:, :, dj + 2] = transpose(x[:, :, dj + 2])
                    x[:, :, dj + 4] = transpose(x[:, :, dj + 4])
                else

                    x[:, :, dj + 2] = transpose(x[:, :, dj + 2])
                    x[:, :, dj + 4] = transpose(x[:, :, dj + 4])

                    x[:, :, dj + 2] = transpose(lifting_step_ti(x[:, :, [2 + dj, 4 + dj]], h, dir, dist))

                    x[:, :, 1] = transpose(x[:, :, 1])
                    x[:, :, dj + 3] = transpose(x[:, :, dj+3])
                    x[:, :, 1] = transpose(lifting_step_ti(x[:, :, [1, dj + 3]], h, dir, dist))

                    x[:, :, 1] = lifting_step_ti(x[:, :, [1, dj + 2]], h, dir, dist)
                end
            end
        end

        if dir == -1
            x = x[:, :, 1]
        end
    end

    return x
end

###########################################################################
###########################################################################
###########################################################################

function lifting_step(x0, h, dir)

    #copy x
    x = copy(x0)

    # number of lifting steps
    m = Int(Base.div(length(h) - 1, 2))

    if dir == 1
        # split
        #d = x[1::2,]
        d = x[collect(i for i in 1:size(x)[1] if i%2 == 0), :]
        #x = x[0::2,]
        x = x[collect(i for i in 1:size(x)[1] if i%2 == 1), :]
        for i in 1:m
            d = d - h[2*i - 1] .* (x + vcat(x[2:end, :], x[end, :]'))  #I must check if h[2^i + 1] is multidimensional or not.
            x = x + h[2*i] .* (d + vcat(d[1, :]', d[1:(end-1), :]))   #I must check if h[2^i + 2] is multidimensional or not.
        end
        x = vcat(x.*h[end], d./h[end])

    else
        # retrieve detail coefs
        fin = size(x)[1]
        d = x[Int(fin/2) + 1 : end, :].*h[end]
        x = x[1:Int(fin/2), :]./h[end]
        for i in m:-1:1
            x = x - h[2*i] .* (d + vcat(d[1, :]', d[1 : end - 1, :])) #L'erreur est ici
            d = d + h[2*i - 1] .* (x + vcat(x[2 : end, :], x[end, :]'))
        end
        # merge
        x1 = vcat(x, x)
        x1[collect(i for i in 1:size(x1)[1] if i%2 == 1), :] = x
        x1[collect(i for i in 1:size(x1)[1] if i%2 == 0), :] = d
        x = x1
    end

    return x
end

###########################################################################
###########################################################################
###########################################################################
function lifting_step_ti(x0, h, dir, dist)

    #copy x
    x = copy(x0)

    # number of lifting steps
    m = Base.div(length(h) - 1, 2)
    n = size(x[:, :, 1])[1]

    s1 = collect(i for i in 1:n) + dist
    s2 = collect(i for i in 1:n) - dist

    # boundary conditions
    s1[s1 .> n] = 2*n .- s1[s1 .> n]
    s1[s1 .< 1] = 2 .- s1[s1 .< 1]

    s2[s2 .> n] = 2*n .- s2[s2 .> n]
    s2[s2 .< 1] = 2 .- s2[s2 .< 1]

    #indices in python start from 0
    # s1 = s1 - 1
    # s2 = s2 - 1

    if dir == 1
        # split
        d = x
        for i in 0 : m - 1
            if ndims(x) == 2
                x = repeat(x, outer = [1, 1, 1])
            end
            d = d - h[2*i + 1] .* (x[s1, :, :] + x[s2, :, :])
            x = x + h[2*i + 2] .* (d[s1, :, :] + d[s2, :, :])
        end

        #merge
        x = cat(3, x.*h[end], d./h[end])

    else
        # retrieve detail coefs

        d = x[:, :, 2].*h[end]
        x = x[:, :, 1]/h[end]

        for i in m:-1:1
            x = x - h[2*i] .* (d[s1, :] + d[s2, :])
            d = d + h[2*i - 1] .* (x[s1, :] + x[s2, :])
        end

        # merge
        x = (x + d)/2
    end

    return x
end


using Test
using AdvancedProbabilisticCircuits
using StatsFuns

const log2πhalf = log2π/2

@testset "nodes" begin

    @testset "constructors" begin
        @test Normal(1) isa AdvancedProbabilisticCircuits.AbstractLeaf
        @test Indicator(1) isa AdvancedProbabilisticCircuits.AbstractLeaf
        @test TruncatedNormal(1) isa AdvancedProbabilisticCircuits.AbstractLeaf
        @test Sum(Normal, 1, 2) isa AdvancedProbabilisticCircuits.Node
        @test Sum(Normal(1), Normal(1)) isa AdvancedProbabilisticCircuits.Node
        @test Prod(Normal, [1,2]) isa AdvancedProbabilisticCircuits.Node
        @test Prod(Normal(1), Normal(2)) isa AdvancedProbabilisticCircuits.Node
    end

    @testset "scope functions" begin
        @test scope(Normal(1)) == 1
        @test scope(Sum(Normal(1), Normal(1))) == 1
        @test scope(Prod(Normal(1), Normal(2))) == [1,2]
    end

    @testset "support functions" begin
        @test support(Normal(1)) == (1 => RealInterval(-Inf, Inf))
        @test support(Indicator(1)) == (1 => [1])
        @test support(Sum(Indicator(1, v = 1), Indicator(1, v = 2))) == (1 => [1,2])
    end

    @testset "logpdf functions" begin
        x = 0
        @test logpdf(Normal(1), x) ≈ -log2πhalf
        @test logpdf(Indicator(1), 1) ≈ 0
        @test logpdf(Indicator(1), 0) ≈ -Inf
        @test logpdf(TruncatedNormal(1), x) ≈ -log2πhalf
        @test logpdf(TruncatedNormal(1, min=1), x) ≈ -Inf
        @test logpdf(TruncatedNormal(1, max=-1), x) ≈ -Inf
        s = Sum(Normal(1), Normal(1))
        @test logpdf(s, x) ≈ logsumexp(-log2π/2 .+ params(s))
        @test logpdf(Prod(Normal(1), Normal(2)), [x,x]) ≈ -log2π
    end

end

@testset "properties" begin

end

@testset "util" begin

end

@testset "optimize" begin

end


# ============================================================
# Hydrodynamic Fluxes
# Reference: http://www.cfdbooks.com/
# ============================================================

"""
    flux_upwind(uL, uR, Ω::T, n::T, dt) where {T<:AbstractVector}

Upwind flux

"""
function flux_upwind(uL, uR, Ω::T, n::T, dt = 1.0) where {T<:AbstractVector}
    ip = dot(Ω, n) # inner product

    if ip > 0
        return dt * ip * uL
    else
        return dt * ip * uR
    end
end

"""
    flux_lax!(fw::AbstractArray{<:Real,1}, wL::AbstractArray{<:Real,1}, wR::AbstractArray{<:Real,1}, γ::Real, dt::Real, dx::Real)

Lax-Friedrichs flux

_P. D. Lax, Weak Solutions of Nonlinear Hyperbolic Equations and Their Numerical Computation,
Commun. Pure and Applied Mathematics, 7, 159-193, 1954._

"""
function flux_lax!(
    fw::X,
    wL::Y,
    wR::Y,
    γ,
    dt,
    dx,
) where {X<:AbstractArray{<:Real,1},Y<:AbstractArray{<:Real,1}}
    fw .= 0.5 * dt .* (euler_flux(wL, γ)[1] + euler_flux(wR, γ)[1] - dx / dt .* (wR - wL))
    return nothing
end

"""
    flux_hll!(fw::AbstractArray{<:Real,1}, wL::AbstractArray{<:Real,1}, wR::AbstractArray{<:Real,1}, γ::Real, dt::Real)

HLL flux for the Euler equations

- @args: variables at left & right sides of interface
- @args: specific heat ratio

"""
function flux_hll!(
    fw::X,
    wL::Y,
    wR::Y,
    γ,
    dt,
) where {X<:AbstractArray{<:Real,1},Y<:AbstractArray{<:Real,1}}

    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)

    aL = sound_speed(primL, γ)
    aR = sound_speed(primR, γ)

    λmin = primL[2] - aL
    λmax = primR[2] + aR

    if λmin >= 0.0
        fw .= euler_flux(wL, γ)[1]
    elseif λmax <= 0.0
        fw .= euler_flux(wR, γ)[1]
    else
        factor = 1.0 / (λmax - λmin)

        flux1 = euler_flux(wL, γ)[1]
        flux2 = euler_flux(wR, γ)[1]

        @. fw = factor * (λmax * flux1 - λmin * flux2 + λmax * λmin * (wR - wL))
    end

    fw .*= dt

    return nothing

end

"""
    flux_roe!(fw::AbstractArray{<:Real,1}, wL::AbstractArray{<:Real,1}, wR::AbstractArray{<:Real,1},
    γ::Real, dt::Real, n = [1.0, 0.0]::AbstractArray{<:Real,1})

Roe's flux with entropy fix

_P. L. Roe, Approximate Riemann Solvers, Parameter Vectors and Difference Schemes, Journal of Computational Physics, 43, pp. 357-372._
(_cf. http://cfdbooks.com/cfdcodes.html_)

* @args primL[1:4] = left state (rhoL, uL, vL, pL)
* @args primR[1:4] = right state (rhoR, uR, vR, pR)
* @args γ: specific heat ratio
* @args n[2]: unit face normal (L -> R)

"""
function flux_roe!(
    fw::X,
    wL::Y,
    wR::Y,
    γ,
    dt,
    n = [1.0, 0.0],
) where {X<:AbstractArray{<:Real,1},Y<:AbstractArray{<:Real,1}}

    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)

    if length(fw) == 3 # 1D

        # left state
        rhoL = primL[1]
        vL = primL[2]
        pL = 0.5 * primL[1] / primL[end]
        aL = sound_speed(primL, γ)
        HL = (wL[end] + pL) / rhoL

        # right state
        rhoR = primR[1]
        vR = primR[2]
        pR = 0.5 * primR[1] / primR[end]
        aR = sound_speed(primR, γ)
        HR = (wR[end] + pR) / rhoR

        # compute Roe averages
        RT = sqrt(rhoR / rhoL)
        rho = RT * rhoL
        v = (vL + RT * vR) / (1.0 + RT)
        H = (HL + RT * HR) / (1.0 + RT)
        a = sqrt((γ - 1.0) * (H - 0.5 * v^2))

        # differences in primitive variables
        drho = rhoR - rhoL
        du = vR - vL
        dP = pR - pL

        # wave strength (characteristic variables)
        dV = [
            0.5 * (dP - rho * a * du) / a^2,
            -(dP / a^2 - drho),
            0.5 * (dP + rho * a * du) / a^2,
        ]

        # absolute values of the wave speeds (eigenvalues)
        ws = [abs(v - a), abs(v), abs(v + a)]

        # modified wave speeds for nonlinear fields (to remove expansion shocks).
        Da = max(0.0, 4.0 * ((vR - aR) - (vL - aL)))
        if ws[1] < 0.5 * Da
            ws[1] = ws[1]^2 / Da + 0.25 * Da
        end
        Da = max(0.0, 4.0 * ((vR + aR) - (vL + aL)))
        if ws[3] < 0.5 * Da
            ws[3] = ws[3]^2 / Da + 0.25 * Da
        end

        # right eigenvectors
        R = zeros(3, 3)
        R[1, 1] = 1.0
        R[2, 1] = v - a
        R[3, 1] = H - v * a

        R[1, 2] = 1.0
        R[2, 2] = v
        R[3, 2] = 0.5 * v * v

        R[1, 3] = 1.0
        R[2, 3] = v + a
        R[3, 3] = H + v * a

        # compute average flux
        fw .= 0.5 .* (euler_flux(wL, γ)[1] + euler_flux(wR, γ)[1])

        # add matrix dissipation term to complete Roe flux
        for j = 1:3, k = 1:3
            fw[j] -= 0.5 * ws[k] * dV[k] * R[j, k]
        end

    elseif length(fw) == 4 # 2D

        # normal vector
        nx = n[1]
        ny = n[2]

        # tangent vector
        mx = -ny
        my = nx

        #--- primitive and other variables ---#
        # left state
        rhoL = primL[1]
        uL = primL[2]
        vL = primL[3]
        unL = uL * nx + vL * ny
        umL = uL * mx + vL * my
        pL = 0.5 * primL[1] / primL[4]
        aL = sound_speed(primL[4], γ)
        HL = aL^2 / (γ - 1.0) + 0.5 * (uL^2 + vL^2)

        # right state
        rhoR = primR[1]
        uR = primR[2]
        vR = primR[3]
        unR = uR * nx + vR * ny
        umR = uR * mx + vR * my
        pR = 0.5 * primR[1] / primR[4]
        aR = sound_speed(primR[4], γ)
        HR = aR^2 / (γ - 1.0) + 0.5 * (uR^2 + vR^2)

        # Roe averages
        RT = sqrt(rhoR / rhoL)
        rho = RT * rhoL
        u = (uL + RT * uR) / (1.0 + RT)
        v = (vL + RT * vR) / (1.0 + RT)
        H = (HL + RT * HR) / (1.0 + RT)
        a = sqrt((γ - 1.0) * (H - 0.5 * (u^2 + v^2)))
        un = u * nx + v * ny
        um = u * mx + v * my

        # wave strengths
        drho = rhoR - rhoL
        dp = pR - pL
        dun = unR - unL
        dum = umR - umL

        LdU = [
            (dp - rho * a * dun) / (2.0 * a^2),
            rho * dum,
            drho - dp / (a^2),
            (dp + rho * a * dun) / (2.0 * a^2),
        ]

        # wave speed
        ws = abs.([un - a, un, un, un + a])

        # Harten's entropy fix JCP(1983), 49, pp357-393
        # only for the nonlinear fields.
        if ws[1] < 0.2
            ws[1] = 0.5 * (ws[1]^2 / 0.2 + 0.2)
        end
        if ws[4] < 0.2
            ws[4] = 0.5 * (ws[4]^2 / 0.2 + 0.2)
        end

        # right eigenvectors
        Rv = zeros(4, 4)

        Rv[1, 1] = 1.0
        Rv[2, 1] = u - a * nx
        Rv[3, 1] = v - a * ny
        Rv[4, 1] = H - un * a

        Rv[1, 2] = 0.0
        Rv[2, 2] = mx
        Rv[3, 2] = my
        Rv[4, 2] = um

        Rv[1, 3] = 1.0
        Rv[2, 3] = u
        Rv[3, 3] = v
        Rv[4, 3] = 0.5 * (u * u + v * v)

        Rv[1, 4] = 1.0
        Rv[2, 4] = u + a * nx
        Rv[3, 4] = v + a * ny
        Rv[4, 4] = H + un * a

        # dissipation term
        diss = zeros(4)
        for i = 1:4, j = 1:4
            diss[i] += ws[j] * LdU[j] * Rv[i, j]
        end

        # compute fluxes
        fL = [
            rhoL * unL,
            rhoL * unL * uL + pL * nx,
            rhoL * unL * vL + pL * ny,
            rhoL * unL * HL,
        ]
        fR = [
            rhoR * unR,
            rhoR * unR * uR + pR * nx,
            rhoR * unR * vR + pR * ny,
            rhoR * unR * HR,
        ]

        @. fw = 0.5 * dt * (fL + fR - diss)

    end

    return nothing

end


# equation.jl

# test methods for equation solving

import Base: transpose
using Lazy

tolerance = 0.000001

@inline good_enough(x) = -tolerance ≤ x ≤ tolerance

function secant_solve(f::Function, a=-100, b=100)
    @rec function loop(a, b)
        let fa = f(a), fb = f(b)
            c = (fb*a-fa*b)/(fb-fa)
            fc = f(c)
            good_enough(fc) && return c
            fa*fc < 0 && return loop(a, c)
            fb*fc < 0 && return loop(b, c)
        end
        warn("error")
    end
    let fa = f(a), fb = f(b)
        if fa * fb > 0
            warn("Can't solve with initial interval [$a, $b]")
            return nothing
        end
        good_enough(fa) && return a
        good_enough(fb) && return b
        loop(a, b)
    end
end

function transpose(f::Function)
    const h = 1e-8
    return x -> (f(x+h) - f(x-h))/2h
end

@rec function newton_solve(f::Function, ξ)
    good_enough(ξ) && return ξ
    newton_solve(f, ξ-f(ξ)/f'(ξ))
    rand(0:100) == 0 && println("ξ = $ξ")
end


module constants

export ME, RE, G, MM, RM, MCM, DISTANCE_TO_MOON, MOON_PERIOD, MOON_INITIAL_ANGLE, ORIGIN
export TOTAL_DURATION, MARKER_TIME, TOLERANCE, INITIAL_POSITION, INITIAL_VELOCITY

const ME = 5.97e24 # mass of earth in kg
const RE = 6.378e6 # radius of earth in m (at equator)
const G = 6.67e-11 # gravitational constant in m3 / kg s2
const MM = 7.35e22 # mass of moon in kg
const RM = 1.74e6 # radius of moon in m
const MCM = 5000. # mass of command module in kg
const DISTANCE_TO_MOON = 400.5e6 # m (actually, not at all a constant)
const MOON_PERIOD = 27.3 * 24.0 * 3600. # s
const MOON_INITIAL_ANGLE = pi / 180. * -61. # radians
const ORIGIN = [0., 0.] # Vector that represents the origin

const TOTAL_DURATION = 12. * 24. * 3600.# s
const MARKER_TIME = 0.5 * 3600. # (used in error correction) s
const TOLERANCE = 100000. # (used in error correction) m

const INITIAL_POSITION = [-6.701e6, 0.] # Initial vector for the spacecraft in m
const INITIAL_VELOCITY = [0., -10.818e3] # Initial velocity for spacecraft in m/s

end


## Eigen Vectors/Values
## ====================
using Sigma

dist(A,B; epsilon = 1.0) = sum([(A[i] - B[i])^2 for i = 1:length(A)]) < epsilon

dist(A,B; epsilon = 1.0) = sum([abs(A[i] - B[i]) for i = 1:length(A)]) < epsilon
dist3(A,B; epsilon = 1.0) = (&)([abs(A[i] - B[i]) < epsilon for i = 1:length(A)]...)
dist2(A,B; epsilon = 1.0) = sum([abs(A[i] - B[i]) for i = 1:length(A)])

function issquare(x)
  size(x)[1] == size(x)[2]
end

function eigen(A, nsamples::Integer; precision=0.001, epsilon = 0.01)
  @assert issquare(A)
  len = size(A)[1]
  v = mvuniform(0.1,1.0, len)
  # V is nonzero
  λ = uniform(0, 1.0)

  sol = rand((A,v,λ), dist3(A*v, λ*v; epsilon=epsilon), nsamples; parallel=true, precision=precision, RandVarType=Sigma.SymbolicRandVar)
  return sol
end

A = mvuniform(0,1,5,5)

sol = eigen(A, 10; epsilon = 0.0001, precision=0.00001)
A = sol[1][1]
v = sol[2][1]
λ = sol[3][1]
dist2(A*v, λ*v)


include("../module/CustomStencil.jl")
include("../misc/misc.jl")
include("../misc/comp_count.jl")
include("../misc/cpu_stencils.jl")

## Star Stencil definition with radius = 1
coefs = [1;0.5;]
star_stencil = @def_stencil_expression c[0]D[x,y,z] + @sum(i, 1, 1, c[i]*(
            D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i] +
            D[x-i,y,z] + D[x,y-i,z] + D[x,y,z-i]))
st_def = CreateStencilDefinition(star_stencil, coefs)
st_single = NewStencilInstance(st_def, m_step=false)
sym = st_single.stencil_sym
countComp(sym,1,7)
## Dense Stencil definition with radius = 1
coefs = [0.5;0.25]
dense_stencil = @def_stencil_expression(
            @sum(i,-1,1,
                @sum(j,-1,1,
                    @sum(k,-1,1, c[max(abs(i),abs(j),abs(k))]*D[x+i,y+j,z+k]))))

st_def = CreateStencilDefinition(dense_stencil, coefs)
st_single = NewStencilInstance(st_def, m_step=false)
sym = st_single.stencil_sym
countComp(sym,1,4)
## Flux summation Stencil definition with radius = 1
coefs = [1;0.2]
flux_stencil = @def_stencil_expression c[0]D[x,y,z] + @sum(i, 1,1, c[i]*(
            D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i]))
st_def = CreateStencilDefinition(flux_stencil, coefs)
st_single = NewStencilInstance(st_def, m_step=false)
sym = st_single.stencil_sym
a = countComp(sym,1,8)


@testset  "Test wrapper for greedysolvers" begin

    ψ = randn(100, 20)
    u = randn(100)
    m, n = size(ψ)
    ctrue = ψ\u

    idxls, cls, ϵls = lsomp(ψ, u; invert = true, verbose = true, ϵrel = eps())
    idxqr, cqr, ϵqr = qromp(ψ, u; invert = true, verbose = true, ϵrel = eps())
    idxpivot, cpivot, ϵpivot = pivotedqr(ψ, u; invert = true, verbose = true, ϵrel = eps())

    # Test LS OMP
    idxgreedy, cgreedy, ϵgreedy = greedysolver("lsomp", ψ, u; invert = true, verbose = true, ϵrel = eps())

    @test norm(idxls - idxgreedy)<1e-14
    @test norm(cls - cgreedy)<1e-14
    @test norm(ϵls - ϵgreedy)<1e-14

    # Test QR OMP
    idxgreedy, cgreedy, ϵgreedy = greedysolver("qromp", ψ, u; invert = true, verbose = true, ϵrel = eps())

    @test norm(idxqr - idxgreedy)<1e-14
    @test norm(cqr - cgreedy)<1e-14
    @test norm(ϵqr - ϵgreedy)<1e-14

    # Test Pivoted QR
    idxgreedy, cgreedy, ϵgreedy = greedysolver("pivotedqr", ψ, u; invert = true, verbose = true, ϵrel = eps())

    @test norm(idxpivot - idxgreedy)<1e-14
    @test norm(cpivot - cgreedy)<1e-14
    @test norm(ϵpivot - ϵgreedy)<1e-14
end


@testset "isconvex" begin
    m = JuMP.Model()

    JuMP.@variables m begin
       x
       y
       z
   end

   # AffExpr
   @test MultilinearOpt.isconvex(x + y)
   @test MultilinearOpt.isconvex(x - z - 3)

   # QuadExpr
   @test MultilinearOpt.isconvex(x^2)
   @test MultilinearOpt.isconvex(x^2 + 0 * z^2) # test positive semidefinite gramian
   @test MultilinearOpt.isconvex(x^2 + 3 * y - z + 5)
   @test !MultilinearOpt.isconvex(-x^2)
   @test !MultilinearOpt.isconvex(x + y - z ^2 + x^2 + y^2)
   G = rand(3, 3)
   G = G * G'
   vars = [x, y, z]
   expr = dot(vars, G * vars)
   G_back, vars_back = MultilinearOpt.gramian(expr)
   @test vars_back == vars
   @test G_back ≈ G atol=1e-10
   @test MultilinearOpt.isconvex(expr)

   # LinearConstraint
   @test MultilinearOpt.isconvex(JuMP.constraint_object(JuMP.@constraint(m, x + 3 * y == 0)))
   @test MultilinearOpt.isconvex(JuMP.constraint_object(JuMP.@constraint(m, 2 * x - y >= z)))
   @test MultilinearOpt.isconvex(JuMP.constraint_object(JuMP.@constraint(m, 2 * x - y <= z)))

   # QuadConstr
   @test !MultilinearOpt.isconvex(JuMP.constraint_object(JuMP.@constraint(m, x == y * z)))

   @test MultilinearOpt.isconvex(JuMP.constraint_object(JuMP.@constraint(m, x^2 + y^2 <= z)))

   @test !MultilinearOpt.isconvex(JuMP.constraint_object(JuMP.@constraint(m, x^2 + y^2 <= z^2)))

   @test MultilinearOpt.isconvex(JuMP.constraint_object(JuMP.@constraint(m, -x^2 - y^2 >= -z)))
end


using OffsetArrays

@testset "Associated Legendre Functions" begin
    nmax = 1_000
    min_length = (nmax * (nmax + 3)) ÷ 2 + 1

    for T in [BigFloat, Float64, Float32]
        ϵ = eps(T)
        P̄ = zeros(T, min_length)
        recursion_coefficients = ALFRecursionCoefficients(nmax, T)

        θs = collect(zero(T):one(T)/10:T(π)/2)

        for (i, θ) in enumerate(θs)
            expiβ = exp(θ * 1im)
            ALFcompute!(P̄, expiβ, nmax, recursion_coefficients)
            offset = -1
            for n in 0:nmax
                offset -= n
                P̄ₙ = OffsetVector(P̄, offset)
                total = sum(P̄ₙ[m]^2 for m in 0:n)
                @test total≈T(2n+1) rtol=2ϵ*n
            end
        end
    end
end


using QuantumSpatialSearch
using LinearAlgebra, SparseArrays
using Test

include("rtn_test.jl")
include("grover_test.jl")


using FactorGraph, Random

# generate a random with 5 nodes (factors) and 10 edges (variables)
Random.seed!(2); rg = random_fg(5, 10)

# remove the first two edges from this incidence list, notice these labels (numbers) of edges are dynamicaly assigned.
rg = rg |> rem_edge(1) |> rem_edge(1)
# remove the first vertex from this incidence list
rg = rg |> rem_vertex(2)

# eliminate a variable (sum over an edge)
rg = rg |> contract(3) |> contract(1)

# the resulting tensors
rg.tensors .|> size

# the topology of the incidence list of the resulting V-3, E-4 graph.
rg.legmap


# Non-bipartite: https://theory.stanford.edu/~jvondrak/CS369P/lec5.pdf

mutable struct PerfectBipartiteMatchingLPSolver <: PerfectBipartiteMatchingSolver
  solver

  # Optimisation model (reused by solve_linear).
  model # ::Model
  x # ::Matrix{Variable}

  function PerfectBipartiteMatchingLPSolver(solver)
    return new(solver, nothing, nothing)
  end
end

function build!(solver::PerfectBipartiteMatchingLPSolver, n_arms::Int)
  # Build the optimisation model behind solve_linear.
  solver.model = Model(solver.solver)
  indices = collect((i, j) for i in 1:n_arms, j in 1:n_arms)
  solver.x = @variable(solver.model, [indices], lower_bound=0, upper_bound=1)

  for i in 1:n_arms # Left nodes.
    @constraint(solver.model, sum(solver.x[(i, j)] for j in 1:n_arms) <= 1)
  end
  for j in 1:n_arms # Right nodes.
    @constraint(solver.model, sum(solver.x[(i, j)] for i in 1:n_arms) <= 1)
  end
end

has_lp_formulation(::PerfectBipartiteMatchingLPSolver) = true

function get_lp_formulation(solver::PerfectBipartiteMatchingLPSolver, rewards::Dict{Tuple{Int, Int}, Float64})
  return solver.model,
    sum(rewards[(i, j)] * solver.x[(i, j)] for (i, j) in keys(rewards)),
    Dict{Tuple{Int, Int}, JuMP.VariableRef}((i, j) => solver.x[(i, j)] for (i, j) in keys(rewards))
end

function solve_linear(solver::PerfectBipartiteMatchingLPSolver, rewards::Dict{Tuple{Int, Int}, Float64})
  m, obj, vars = get_lp_formulation(solver, rewards)
  @objective(m, Max, obj)

  set_silent(m)
  optimize!(m)

  if termination_status(m) != MOI.OPTIMAL
    return Tuple{Int, Int}[]
  end
  return Tuple{Int, Int}[(i, j) for (i, j) in keys(rewards) if value(vars[(i, j)]) > 0.5]
end


#=
    Tiled LU factorization in pure Julia. The serial performance essentially
    comes from the `TriangularSolve` and `LoopVectorization`, and
    `RecursiveFactorization` packages. The parallelization is handled by
    `DataFlowTask`s.
=#

lu(A::Matrix,args...) = lu!(deepcopy(A),args...)

function lu!(A::Matrix,s=TILESIZE[],tturbo::Val{T}=Val(false)) where {T}
    _lu!(PseudoTiledMatrix(A,s),tturbo)
end

function _lu!(A::PseudoTiledMatrix,tturbo::Val{T}=Val(false)) where {T}
    m,n = size(A)
    for i in 1:m
        Aii = A[i,i]
        # TODO: for simplicity, no pivot is allowed. Pivoting the diagonal
        # blocks requires permuting the corresponding row/columns before continuining
        @dspawn RecursiveFactorization.lu!(Aii,NoPivot(),tturbo) (Aii,) (RW,)
        # @dspawn LinearAlgebra.lu!(Aii) (Aii,) (RW,)
        for j in i+1:n
            Aij = A[i,j]
            Aji = A[j,i]
            @dspawn begin
                TriangularSolve.ldiv!(UnitLowerTriangular(Aii),Aij,tturbo)
                TriangularSolve.rdiv!(Aji,UpperTriangular(Aii),tturbo)
            end (Aii,Aij,Aji) (R,RW,RW)
            # TriangularSolve.ldiv!(UnitLowerTriangular(Aii),Aij,tturbo)
            # TriangularSolve.rdiv!(Aji,UpperTriangular(Aii),tturbo)
        end
        for j in i+1:m
            for k in i+1:n
                Ajk = A[j,k]
                Aji = A[j,i]
                Aik = A[i,k]
                @dspawn schur_complement!(Ajk,Aji,Aik,tturbo) (Ajk,Aji,Aik) (RW,R,R)
                # schur_complement!(Ajk,Aji,Aik,tturbo)
            end
        end
    end
    # wait for all computations before returning
    DataFlowTasks.sync()
    return LU(A.data,BlasInt[],zero(BlasInt))
end

# a fork-join approach for comparison with the data-flow parallelism
function _lu_forkjoin!(A::PseudoTiledMatrix,tturbo::Val{T}=Val(false)) where {T}
    m,n = size(A)
    for i in 1:m
        Aii = A[i,i]
        # FIXME: for simplicity, no pivot is allowed. Pivoting the diagonal
        # blocks requires permuting the corresponding row/columns before continuining
        RecursiveFactorization.lu!(Aii,NoPivot(),tturbo)
        Threads.@threads for j in i+1:n
            Aij = A[i,j]
            Aji = A[j,i]
            TriangularSolve.ldiv!(UnitLowerTriangular(Aii),Aij,tturbo)
            TriangularSolve.rdiv!(Aji,UpperTriangular(Aii))
        end
        @sync for j in i+1:m
            for k in i+1:n
                Ajk = A[j,k]
                Aji = A[j,i]
                Aik = A[i,k]
                Threads.@spawn schur_complement!(Ajk,Aji,Aik,tturbo)
            end
        end
    end
    # wait for all computations before returning
    DataFlowTasks.sync()
    return LU(A.data,BlasInt[],zero(BlasInt))
end


# This program creates a simple 3D vel model for seismic modeling
#
# by GRAM | 8 Apr 2016

using Seismic

# def size
nx = 100
ny = 100
nz = 200
dz = 10 # depth increment
dx = 20
dy = 20

vel = zeros(Float32,nz,nx,ny)

# build model
for i = 1:nz, j = 1:nx, k = 1:ny # overburden 1
	vel[i,j,k] = 2000
end

for i = 1:nz, j = 1:nx, k = 1:ny # overburden 2
	if i > 30
		vel[i,j,k] = 2500
	end
end

for i = 1:nz, j = 1:nx, k = 1:ny # domal pyramid 1
	if ((i > floor(100 - 1.6 * j + 0.016 * j^2)) & (i > floor(100 - 1.6 * k + 0.016 * k^2)))
		vel[i,j,k] = 3000
	end
end

for i = 1:nz, j = 1:nx, k = 1:ny # domal pyramid 2
	if ((i > floor(140 - 1.6 * j + 0.016 * j^2)) & (i > floor(140 - 1.6 * k + 0.016 * k^2)))
		vel[i,j,k] = 4000
	end
end

for i = 1:nz, j = 1:nx, k = 1:ny # basement
	if i >= 170
		vel[i,j,k] = 5000
	end
end

# Here we "pick" erroneously high vels below the target
# zone, hypothetically because the geophysicist didn't
# want to pick a velocity inversion.

for i = 1:nz, j = 1:nx, k = 1:ny # target
	if ((i < 170) & (i > 130) & (i > floor(140 - 1.6 * j + 0.016 * j^2)) & (i > floor(140 - 1.6 * k + 0.016 * k^2)))
		vel[i,j,k] = 4050
	end
end

#= display
print(size(vel), "\n")
print(vel,"\n")
=#

#= write out text file for QC
out = open("../dat/vel_1.txt","w")
for k = 1:ny
	for i = 1:nz
		for j = 1:nx
			showcompact(out, vel[i,j,k])
			write(out," ")
		end
		write(out, "\r\n")
	end
	write(out, "\r\n")
end
close(out)
=#

# write out vel model for Seismic.jl
h = Array(Header,nx*ny)
ex = Seismic.Extent(convert(Int32,nz), convert(Int32,nx), convert(Int32,ny),
	1, 1, 0, 0, 0, 0, 0, convert(Float32,dz), convert(Float32,dx),
	convert(Float32,dy), 1, 1, "Depth", "mx", "my", "", "", "", "", "",
	"", "", "")

for i = 1 : ny
	for j = 1 : nx
		h[j + nx * (i - 1)] = Seismic.InitSeisHeader()
		h[j + nx * (i - 1)].tracenum = convert(Int32, j + nx * (i - 1))
		h[j + nx * (i - 1)].o1 = convert(Float32, 0)
		h[j + nx * (i - 1)].n1 = convert(Int32, nz)
		h[j + nx * (i - 1)].d1 = convert(Float32, dz)
		h[j + nx * (i - 1)].mx = convert(Float32, j)
		h[j + nx * (i - 1)].my = convert(Float32, i)
	end
end

SeisWrite("../dat/vel_incorrect",vel,h,ex)


### A Pluto.jl notebook ###
# v0.17.1

using Markdown
using InteractiveUtils

# ╔═╡ f019936e-8e7f-11ec-2c9e-3b68771ae51c
md"""
## The basics 2
We just learned how to use Bayes theorem to calculate the posterior probability distribution that allows us to make a quantitative assessment of our decision whether we should accept this gambling game or not.  Such decisions are constantly made in the financial industry and Bayesian methods are commonly used in this environment.  In the financial industry the goal is to maximize profit and to minimize losses, so it is imperative to use quantitative tools.  In science, we have the problem that scientists are not in the same situation.  For scientists it is more important to get more grants, to write more papers since this is what is seen as the deliverables of their efforts.  What I would like to teach you is that by following Bayes theorem, we can improve the reliability of our research without sacrificing academic productivity.  We need to know our data though.
"""

# ╔═╡ 88d6a70f-85aa-41a1-b159-ce7c16238e5c
md"""
# Some important properties of Bayes theorem

Before we continue, I would like to point out a few important properties of Bayes theorem and its application.  The first property has to do with the way the posterior is calculated.  Let us go back to the coin flipping experiment and assume that we did an experiment with the coin in question that initially involved N=10 coin flips.  We used our prior and by combining it with the Bionomial likelihood, we were able to calculate the posterior:

$$P(p \mid \alpha,\beta, N, k) =  \frac{\Gamma(\alpha + \beta +N)}{\Gamma(\alpha+k)\Gamma(\beta + N -k)}  p^{\alpha-1+k}(1-p)^{\beta-1+N-k}$$

Lets now say that we are not quite sure about what to do and we ask the Gambler to let us flip another 10 times.  What we would do is to take the posterior that we derived, use it as our prior and then again use the Bionomial distribution as likelihood.  When looking back at the equations, it is pretty obvious that it does not matter whether you analyze the data sequentially (flip 10 times, calculate the posterior, flip 10 more times, update the posterior) or analyze it in one shot (take result of 20 flips, use prior, calculate posterior).  This is an important property of Bayes theorem.  It should not matter how you do the analysis.  It would be perfectly fine to flip the coin and update the posterior each time.  The result should be the same.  This statement is pretty obvious in our case but it holds in general. Later in this class we will see examples that are much more complicated than simple coin flips and here it sometimes makes sense to perform a complicated analysis with parts of the data and then repeat the process, rather than take all the data at once.  Also, in science it often happens that you measure some data, and you want to see whether it makes sense to persue this direction further.  That is perfectly ok.  If your inital posterior looks promising then go ahead and do some more measurements.  Bayesian analysis allows you to improve your statistics by repeating the experiment.
"""

# ╔═╡ dbec33c5-8ae8-4c77-8adc-979259063dad
md"""
The next property has to do with the shape of the posterior.  We have seen last lecture that a typical outcome of your posterior can be described by the maximum of the distribution (or the 50th percentile) and its width.  In our coin flipping experiment, the maximum will tell you whether or not the gamble will be profitable for you and the width will tell you about the certainty or uncertainty that this will happen.  In the financial market you would be interested whether a stock is likely to go up or down, and what the volatility is.  So lets use our posterior to calculate these quantities.

The first thing that we want to do is to take the natural logarithm of the probability distribution.  This will not change the maximum of $P(x)$ since the transformation is monotonic.  This transformation is often using in computational approaches to probability because it turns multiplication into additions.  You will see another reason in the next paragraph.

$$logp(x) = \log(P(x | D)$$

Now we can do a Taylor expansion around $x=x_{0}$.  Since we are concerned about a peaked probability distribution we only have a constant term and a quadratic term.

$$logp(x) = logp(x_{0}) + \frac{1}{2}\frac{d^{2}logp(x)}{dx^{2}}\big | _{x_{0}}(x-x_{0})^{2}+\dots$$

Putting this back into the previous equation yields:

$$P(x|D) \propto A \exp\big[\frac{1}{2}\frac{d^{2}logp(x)}{dx^{2}}\big | _{x_{0}}(x-x_{0})^{2}\big]$$
"""

# ╔═╡ ec134312-fdcf-4ec4-990f-3156cd9d0095
md"""
If you look at this equation you will recognize that what we have just done is to approximate our pdf with a Gaussian distribution.  Gaussian distributions are very commonly found, and the reason for this is beyond this class, but suffice to say that in many cases posterior distributions tend to approach Gaussian distributions for large N (law of large numbers).  Lets look at a Gaussian distribution:

$$P(x \mid \mu, \sigma) = \frac{1}{{\sigma \sqrt {2\pi } }}e^{{{ - \left( {x - \mu }\right)^2 } /{2\sigma ^2 }}}$$
"""

# ╔═╡ 73af177b-3dd3-4785-af4d-3191ed96efde
md"""
The maximum of this distribution can be found by finding the root of the first derivative with respect to x, and the standard deviation can be found by taking the second derivative.  Here is the calculation.  You can check it if you want yourself.  In particuar,

$$ \sigma = (-\frac{d^{2}logp(x)}{dx^{2}}|_{x_{0}})^{-\frac{1}{2}}$$

Once we know $\mu$ and $\sigma$ we can report our posterior pdf as:

$$ x = x_{0} \pm \sigma$$

Since we approximated our pdf with a Gaussian we also know that true value of x lies in within $\pm \sigma$ $67\%$ of the time and within $\pm 2\sigma$ $95\%$ of the time.
"""

# ╔═╡ 4d018fd2-5971-48ba-896a-27170f138b82
md"""
Lets get back to our example posterior and see what happens

$$P(p \mid |N,k) \propto p^{k}(1-p)^{N-k}$$ with
$$logp = c + k\log(p)+(N-k)\log(1-p)$$

The derivatives are:

$$\frac{dlogp}{dp}=\frac{k}{p}-\frac{N-k}{1-p}$$
$$\frac{d^2 logp}{dp^2}=-\frac{k}{p^2}-\frac{N-k}{(1-p)^2}$$

The maximum can be found

$$\frac{dlogp}{dp}\big | _{p_{0}} = \frac{k}{p_{0}}-\frac{N-k}{1-p_{0}}=0$$
$$p_{0}=\frac{k}{N}$$

which is not surprising.  The second derivate is then:
$$\frac{d^2 logp}{dp^2}\big | _{p_{0}}=-\frac{N}{p_{0}(1-p_{0})}$$

$$\sigma = \sqrt{\frac{p_{0}(1-p_{0})}{N}}$$

For larger N, p_{0} will not change much and the numerator will be essentially constant.  In this regime, the standard deviation \mu is inversely proportional to the square root of the number of tries.
"""

# ╔═╡ Cell order:
# ╠═f019936e-8e7f-11ec-2c9e-3b68771ae51c
# ╠═88d6a70f-85aa-41a1-b159-ce7c16238e5c
# ╠═dbec33c5-8ae8-4c77-8adc-979259063dad
# ╠═ec134312-fdcf-4ec4-990f-3156cd9d0095
# ╠═73af177b-3dd3-4785-af4d-3191ed96efde
# ╠═4d018fd2-5971-48ba-896a-27170f138b82


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.8.1
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# +
using GLM
using CSV
using Random
using StatsBase
using DataFrames
using Dagitty
using Turing
using StatsPlots
using StatisticalRethinking, StructuralCausalModels
using StatisticalRethinkingPlots
using Cairo, Fontconfig
using Logging

default(labels = false)
Logging.disable_logging(Logging.Warn);
# -

# Code 6.1

# +
Random.seed!(1917)
N = 200   # grant proposals
p = 0.1   # proportion to select

# uncorrelated newsworthiness and trustworthiness
nw = rand(Normal(), N)
tw = rand(Normal(), N)

# select top 10% of combined score
s = nw .+ tw
q = quantile(s, 1 - p)
selected = s .>= q
cor(tw[selected], nw[selected])
# -

scatter(nw[.!selected], tw[.!selected]; xlab = "newsworthiness", ylab = "trustworthiness", label = "not selected")
scatter!(nw[selected], tw[selected]; label = "selected")

# # 6.1 Multicollinearity

# Code 6.2

Random.seed!(100)
N = 100
height = rand(Normal(10, 2), N)
leg_prop = rand(Uniform(0.4, 0.5), N)
leg_left = leg_prop .* height .+ rand(Normal(0, 0.02), N)
leg_right = leg_prop .* height .+ rand(Normal(0, 0.02), N)
d = DataFrame(:height => height, :leg_left => leg_left, :leg_right => leg_right);
first(d, 6)

# Code 6.3

# +
@model function model_m6_1(leg_left, leg_right, height)
    a ~ Normal(10, 100)
    bl ~ Normal(2, 10)
    br ~ Normal(2, 10)
    μ = @. a + bl * leg_left + br * leg_right
    σ ~ Exponential(1)
    height ~ MvNormal(μ, σ)
end

m6_1 = sample(model_m6_1(d.leg_left, d.leg_right, d.height), NUTS(), 1000)
m6_1_df = DataFrame(m6_1);
precis(m6_1_df)
# -

# Code 6.4

coeftab_plot(m6_1_df)

# Code 6.5

scatter(m6_1_df.br, m6_1_df.bl; alpha = 0.1)

# Code 6.6

@df m6_1_df density(:br + :bl; lw = 2, xlab = "sum of bl and br")

# Code 6.7

# +
@model function model_m6_2(leg_left, height)
    a ~ Normal(10, 100)
    bl ~ Normal(2, 10)
    μ = @. a + bl * leg_left
    σ ~ Exponential(1)
    height ~ MvNormal(μ, σ)
end

m6_2 = sample(model_m6_2(d.leg_left, d.height), NUTS(), 1000)
m6_2_df = DataFrame(m6_2);
precis(m6_2_df)
# -

std(m6_1_df.bl), std(m6_1_df.br), std(m6_1_df.bl + m6_1_df.br)

# Code 6.8

# +
d = DataFrame(CSV.File("data/milk.csv", missingstring = "NA"))

# get rid of dots in column names
rename!(n -> replace(n, "." => "_"), d)

d[!, :K] = standardize(ZScoreTransform, d.kcal_per_g)
d[!, :F] = standardize(ZScoreTransform, d.perc_fat)
d[!, :L] = standardize(ZScoreTransform, d.perc_lactose);
# -

# Code 6.9

# +
@model function model_m6_3(F, K)
    a ~ Normal(0, 0.2)
    bF ~ Normal(0, 0.5)
    μ = @. a + F * bF
    σ ~ Exponential(1)
    K ~ MvNormal(μ, σ)
end

m6_3 = sample(model_m6_3(d.F, d.K), NUTS(500, 0.65), 1000)
m6_3_df = DataFrame(m6_3);

@model function model_m6_4(L, K)
    a ~ Normal(0, 0.2)
    bL ~ Normal(0, 0.5)
    μ = @. a + L * bL
    σ ~ Exponential(1)
    K ~ MvNormal(μ, σ)
end

m6_4 = sample(model_m6_4(d.L, d.K), NUTS(), 1000)
m6_4_df = DataFrame(m6_4);

precis(m6_3_df)
precis(m6_4_df)

coeftab_plot(m6_3_df)
coeftab_plot(m6_4_df)
# -

# Code 6.10

# +
@model function model_m6_5(F, L, K)
    a ~ Normal(0, 0.2)
    bF ~ Normal(0, 0.5)
    bL ~ Normal(0, 0.5)
    μ = @. a + F * bF + L * bL
    σ ~ Exponential(1)
    K ~ MvNormal(μ, σ)
end

m6_5 = sample(model_m6_5(d.F, d.L, d.K), NUTS(), 1000)
m6_5_df = DataFrame(m6_5);
precis(m6_5_df)
# -

# Code 6.11

@df d corrplot([:kcal_per_g :perc_fat :perc_lactose]; seriestype = :scatter, bins = 10, grid = false)

# Code 6.12

# +
# get mean stderr for linear model's scale
function stderr_for_r(r)
    σ = sqrt(1 - r^2) * var(d.perc_fat)
    fat_scaled = r .* d.perc_fat
    stderr_x = [
        begin
            x = d.perc_fat .+ rand(MvNormal(fat_scaled, σ))
            # add the intercept to the model
            X = hcat(ones(length(x)), x)
            m = lm(X, d.kcal_per_g)
            stderror(m)[2]
        end
        for _ = 1:100
    ]
    s = mean(stderr_x)
end

r_seq = range(0, 0.99; step = 0.01)
s = stderr_for_r.(r_seq)
plot(r_seq, s; lw = 2, xlab = "correlation")
# -

# # 6.2 Post-treatment bias

# Code 6.13

# +
Random.seed!(70)
# number of plants
N = 100
h0 = rand(Normal(10, 2), N)
treatment = repeat(0:1, inner = div(N, 2))
fungus = [rand(Binomial(1, 0.5 - treat * 0.4)) for treat in treatment]
h1 = h0 .+ rand(MvNormal(5 .- 3 .* fungus, 1))

d = DataFrame(:h0 => h0, :h1 => h1, :treatment => treatment, :fungus => fungus)
precis(d)
# -

# Code 6.14

sim_p = rand(LogNormal(0, 0.25), 10_000)
precis(DataFrame(:sim_p => sim_p))

# Code 6.15

# +
@model function model_m6_6(h0, h1)
    p ~ LogNormal(0, 0.25)
    σ ~ Exponential(1)
    μ = h0 .* p
    h1 ~ MvNormal(μ, σ)
end

m6_6 = sample(model_m6_6(d.h0, d.h1), NUTS(), 1000)
m6_6_df = DataFrame(m6_6);
precis(m6_6_df)
# -

# Code 6.16

# +
@model function model_m6_7(h0, treatment, fungus, h1)
    a ~ LogNormal(0, 0.2)
    bt ~ Normal(0, 0.5)
    bf ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    p = @. a + bt * treatment + bf * fungus
    μ = h0 .* p
    h1 ~ MvNormal(μ, σ)
end

m6_7 = sample(model_m6_7(d.h0, d.treatment, d.fungus, d.h1), NUTS(), 1000)
m6_7_df = DataFrame(m6_7);
precis(m6_7_df)
coeftab_plot(m6_7_df)
# -

# Code 6.17

# +
@model function model_m6_8(h0, treatment, h1)
    a ~ LogNormal(0, 0.2)
    bt ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    p = @. a + bt * treatment
    μ = h0 .* p
    h1 ~ MvNormal(μ, σ)
end

m6_8 = sample(model_m6_8(d.h0, d.treatment, d.h1), NUTS(), 1000)
m6_8_df = DataFrame(m6_8);
precis(m6_8_df)
coeftab_plot(m6_8_df)
# -

# Code 6.18

plant_dag = Dagitty.DAG(:H₀ => :H₁, :F => :H₁, :T => :F)
drawdag(plant_dag, [2, 0, 1, 3], [0, 0, 0, 0])

# Code 6.19

implied_conditional_independencies_min(plant_dag)

# Code 6.20

# +
Random.seed!(70)
# number of plants
N = 1000
h0 = rand(Normal(10, 2), N)
treatment = repeat(0:1, inner = div(N, 2))
M = rand(Bernoulli(), N)
fungus = [
    rand(Binomial(1, 0.5 - treat * 0.4 + 0.4 * m))
    for (treat, m) ∈ zip(treatment, M)
]
h1 = h0 .+ rand(MvNormal(5 .+ 3 .* M, 1))

d2 = DataFrame(:h0 => h0, :h1 => h1, :treatment => treatment, :fungus => fungus);
precis(d2)
# -

m6_7 = sample(model_m6_7(d2.h0, d2.treatment, d2.fungus, d2.h1), NUTS(), 1000)
precis(DataFrame(m6_7))

m6_8 = sample(model_m6_8(d2.h0, d2.treatment, d2.h1), NUTS(), 1000)
precis(DataFrame(m6_8))

# # 6.3 Collider bias

# Code 6.21

d = sim_happiness(seed = 1977, n_years = 1000)
precis(d)

# +
d_m = d[d.married.==1, [:age, :happiness]]
d_u = d[d.married.==0, [:age, :happiness]]

scatter(d_m.age, d_m.happiness; label = "married", xlab = "age", ylab = "happiness")
scatter!(d_u.age, d_u.happiness; c = :white)
# -

# Code 6.22

d2 = d[d.age.>17, :]
d2[!, :A] = @. (d2.age - 18) / (65 - 18);

# Code 6.23

d2[!, :mid] = d2.married .+ 1;

# +
@model function model_m6_9(mid, A, happiness)
    a ~ MvNormal([0, 0], 1)
    bA ~ Normal(0, 2)
    μ = a[mid] .+ bA .* A
    σ ~ Exponential(1)
    happiness ~ MvNormal(μ, σ)
end

m6_9 = sample(model_m6_9(d2.mid, d2.A, d2.happiness), NUTS(), 1000)
m6_9_df = DataFrame(m6_9);
precis(m6_9_df)
# -

# Code 6.24

# +
@model function model_m6_10(A, happiness)
    a ~ Normal()
    bA ~ Normal(0, 2)
    μ = a .+ bA .* A
    σ ~ Exponential(1)
    happiness ~ MvNormal(μ, σ)
end

m6_10 = sample(model_m6_10(d2.A, d2.happiness), NUTS(), 1000)
m6_10_df = DataFrame(m6_10)
precis(m6_10_df)
# -

# Code 6.25

N = 200
b_GP = 1
b_GC = 0
b_PC = 1
b_U = 2

# Code 6.26

Random.seed!(42)
U = 2 .* rand(Bernoulli(), N) .- 1
G = rand(Normal(), N)
P = rand(MvNormal(@. b_GP * G + b_U * U))
C = rand(MvNormal(@. b_PC * P + b_GC * G + b_U * U))
d = DataFrame(:C => C, :P => P, :G => G, :U => U);

# Code 6.27

# +
@model function model_m6_11(P, G, C)
    a ~ Normal()
    b_PC ~ Normal()
    b_GC ~ Normal()
    μ = @. a + b_PC * P + b_GC * G
    σ ~ Exponential(1)
    C ~ MvNormal(μ, σ)
end

m6_11 = sample(model_m6_11(d.P, d.G, d.C), NUTS(), 1000)
m6_11_df = DataFrame(m6_11);
precis(m6_11_df)
# -

# Code 6.28

# +
@model function model_m6_12(P, G, U, C)
    a ~ Normal(0, 1)
    b_PC ~ Normal()
    b_GC ~ Normal()
    b_U ~ Normal()
    μ = @. a + b_PC * P + b_GC * G + b_U * U
    σ ~ Exponential(1)
    C ~ MvNormal(μ, σ)
end

m6_12 = sample(model_m6_12(d.P, d.G, d.U, d.C), NUTS(), 1000)
m6_12_df = DataFrame(m6_12);
precis(m6_12_df)

# NOTE: something not quite right here, as we should obtain the same values as the parameters used for generating the data. Does Turing scale by default?
# -

# # 6.4 Confronting confounding

# Code 6.29

# +
dag_61 = Dagitty.DAG(
    :X => :Y,
    :U => :X, :A => :U,
    :A => :C, :C => :Y,
    :U => :B, :C => :B,
)

all_backdoor_adjustment_sets(dag_61, :X, :Y)

dag_61 = StructuralCausalModels.DAG("dag_61",
    "DAG(
    U ~ A,
    C ~ A,
    B ~ U + C,
    X ~ U,
    Y ~ X + C
    )"
)
StructuralCausalModels.adjustment_sets(dag_61, :X, :Y)
# -

# Code 6.30

dag_62 = Dagitty.DAG(
    :A => :D,
    :A => :M, :M => :D,
    :S => :A, :S => :M,
    :S => :W, :W => :D,
)
all_backdoor_adjustment_sets(dag_62, :W, :D)

implied_conditional_independencies_min(dag_62)

#  StructuralCausalModels version:
dag_62b = StructuralCausalModels.DAG("dag_62b",
    "DAG(
        A ~ S,
        M ~ A+S,
        W ~ S,
        D ~ A + M + W
    )"
)
basis_set(dag_62b)
adjustment_sets(dag_62b, :W, :D)


using LogarithmicNumbers, Test

_approx(x,y) = (x ≈ y) || (isnan(x) && isnan(y))

# use these to check for type stability
_exp(args...) = @inferred exp(args...)
_log(args...) = @inferred log(args...)
_add(args...) = @inferred +(args...)
_sub(args...) = @inferred -(args...)
_mul(args...) = @inferred *(args...)
_div(args...) = @inferred /(args...)
_pow(args...) = @inferred ^(args...)
_float(args...) = @inferred float(args...)
_inv(args...) = @inferred inv(args...)
_prod(args...) = @inferred prod(args...)
_sum(args...) = @inferred sum(args...)

# sample values
vals = (-Inf, -100, -100.0, -2, -2.0, -1, -1.0, -0.5, 0, 0.0, 0.5, 1, 1.0, 2, 2.0, 100, 100.0, Inf)

# sample vectors
vecs = (
	Float64[x for x in vals if x isa Float64],
	Float64[x for x in vals if x isa Float64 && !isinf(x)],
	Float64[x for x in vals if x isa Float64 && !iszero(x)],
	Float64[x for x in vals if x isa Float64 && !isinf(x) && x ≥ 0],
	Int[x for x in vals if x isa Int],
	Int[x for x in vals if x isa Int && x ≥ 0],
)

atypes = (ULogarithmic, Logarithmic, CLogarithmic)

@testset "LogarithmicNumbers" begin

	@testset "types" begin
	    @test @isdefined ULogarithmic
	    @test @isdefined Logarithmic
	    @test @isdefined CLogarithmic
	end

	@testset "exp" begin
		for A in atypes, x in vals
			@test _exp(A, x) isa A{typeof(x)}
		end
	end

	@testset "convert" begin
	    for A in atypes, x in vals
	    	if x < 0 && A == ULogarithmic
	    		@test_throws DomainError A(x)
	    	else
	    		@test (@inferred A(x)) isa A
	    	end
	    end
	end

    @testset "log" begin
    	for A in atypes, x in vals
    		y = A==CLogarithmic ? Complex(x) : x
			@test _log(_exp(A, x)) === y
    	end
    end

	@testset "float" begin
	    for A in atypes, x in vals
	    	_float(_exp(A, x)) ≈ exp(x) || @info("float1",A,x)
	    	@test _float(_exp(A, x)) ≈ exp(x)
	    	if x < 0 && A == ULogarithmic
	    		nothing
	    	elseif x==-Inf && A == CLogarithmic
	    		@test AbstractFloat(A(x)) == x
	    	else
	    		@test _float(A(x)) ≈ x
	    	end
	    end
	end

    @testset "mul" begin
    	for A in atypes, x in vals, y in vals
    		@test _approx(_float(_mul(_exp(A, x), _exp(A, y))), exp(x) * exp(y))
    	end
    end

    @testset "div" begin
    	for A in atypes, x in vals, y in vals
    		@test _approx(_float(_div(_exp(A, x), _exp(A, y))), exp(x) / exp(y))
    	end
    end

    @testset "pow" begin
		for A in atypes, x in vals, n in (-2,-1,0,1,2,-1.1,0.0,2.3)
			if isinf(x) && iszero(n)
				# should be 1.0 but is NaN
				@test_broken _approx(_float(_pow(_exp(A, x), n)), exp(x)^n)
				@test isnan(_float(_pow(_exp(A, x), n)))
			else
			    @test _approx(_float(_pow(_exp(A, x), n)), exp(x)^n)
			end
        end
    end

    @testset "inv" begin
        for A in atypes, x in vals
            @test _approx(_float(_inv(_exp(A, x))), inv(exp(x)))
        end
    end

    @testset "add" begin
        for A in atypes, x in vals, y in vals
    		@test _approx(_float(_add(_exp(A, x), _exp(A, y))), exp(x)+exp(y))
        end
    end

    @testset "sub" begin
        for A in atypes, x in vals, y in vals
        	if A==ULogarithmic && x<y
	    		@test_throws DomainError _float(_sub(_exp(A, x), _exp(A, y)))
	    	elseif A==CLogarithmic && y==Inf && x<y
	    		# problem with complex logarithmics
	    		@test_broken _approx(_float(_sub(_exp(A, x), _exp(A, y))), exp(x)-exp(y))
	    	else
	    		@test _approx(_float(_sub(_exp(A, x), _exp(A, y))), exp(x)-exp(y))
	    	end
        end
    end

    @testset "prod" begin
        for A in atypes, xs in vecs
        	_approx(_float(_prod(map(x->_exp(A,x), xs))), prod(exp.(xs))) || @info("prod",A,xs)
        	@test _approx(_float(_prod(map(x->_exp(A,x), xs))), prod(exp.(xs)))
        end
    end

    @testset "sum" begin
        for A in atypes, xs in vecs
        	if eltype(xs) <: AbstractFloat
	            @test _approx(_float(_sum(map(x->exp(A,x),xs))), sum(exp.(xs)))
	        else
	        	# sum is not type-stable because typeof(xs[1]+xs[2]) != typeof(xs[1]).
	        	# hence this test is the same as above but without the stability check.
	        	# we don't promise type stability unless the base type is a float, so
	        	# this isn't a broken test.
	        	@test _approx(_float(sum(map(x->exp(A,x),xs))), sum(exp.(xs)))
	        end
        end
    end

end


function is_leap_year(year::Int)
  if year % 400 == 0
	return true
  end

  if year % 100 == 0
	return false
  end

  if year % 4 == 0
	return true
  end

  return false

end


function test_GaussianRewards()
    means = [x -> x[1], x -> 2 * x[1]]
    rdist = GaussianRewards(means, 2.0)
    seed!(123); r1 = sample([1.0], 2, rdist)
    seed!(123); r2 = rand(Normal(2.0, 2.0))
    @test r1 == r2
end

test_GaussianRewards()


module MP3
include("AudioDisplay.jl")

using Plots
using FFTW
using WAV

export AB1A1aPlay
export AB1Aufgabe1a
export AB1A1bPlay
export AB1Aufgabe1b
export AB1A1cPlay
export AB1Aufgabe1c
export AB1A1dPlay
export AB1Aufgabe1d
export AB1A2a1Play
export AB1A2a2Play
export AB1Aufgabe2a
export AB1A2bPlay
export AB1Aufgabe2b

function PlaySinusSound(frequency, amplitude, name)
    #Generate and play a sinus sound.
    eps = 1/44100
    x = range(0.0,stop=2,step=eps);
    y = zeros(Float64,length(x));
    for i = 1:length(y)
        y[i] = amplitude * sin(2*pi*frequency*x[i]);
    end
    #display(y)
    wavwrite(y, string(name, ".wav"), Fs=44100);
    #y, fs = wavread("SinustonAB1A1a.wav");
    #wavplay(y, fs)
	inline_audioplayer(string(name, ".wav"))
    plot(y)
end

function AB1A1aPlay()
    PlaySinusSound(200, 1, "SinustonAB1A1a")

    # Display 200 Hz sin sound.
    amplitude = 1
    frequenz = 200

    x = 0:0.0001:0.1;
    y = zeros(1001);
    for i = 1:1001
        y[i] = amplitude*sin(2*pi*frequenz*x[i]);
    end
    labels = ["Sinuston"];
    plot(x,y, label = labels)
end

function AB1Aufgabe1a(amplitude)
    # Check the amplitude was provided correctly
    if amplitude!=1
        @warn "Amplitude falsch!"
    else
        println("Richtig")
    end
end

function AB1A1bPlay(amplitude)
    PlaySinusSound(200, amplitude, "SinustonAB1A1b")
end

function AB1Aufgabe1b(amplitude)
    if amplitude == 1
        println("Probier doch auch mal eine andere Amplitude.")
    end
    # Display 200 Hz sin sound.
    frequenz = 200

    x = 0:0.0001:0.1;
    y = zeros(1001);
    for i = 1:1001
        y[i] = amplitude*sin(2*pi*frequenz*x[i]);
    end
    labels = ["Sinuston"];
    plot(x,y, label = labels)
end

function AB1A1cPlay(A)
    eps = 1/44100
    x = range(0.0,stop=2,step=eps);
    y = zeros(Float64,length(x));
    for i = 1:length(y)
        y[i] = A(x[i])*sin(2*pi*200*x[i]);
    end
    #display(ton)
    wavwrite(y, "SinustonAB1A1c1.wav", Fs=44100);
    #y, fs = wavread("SinustonAB1A1c1.wav");
    #wavplay(y, fs)
	inline_audioplayer("SinustonAB1A1c1.wav")
end

function AB1Aufgabe1c(A)
    if A(0) != 0
		@warn "Die Amplitude hat zur Zeit t = 0 nicht den Wert 0!"
    end
    m1 = A(1)-A(0);
    m2 = A(2)-A(1);
    if m1 != m2
		@warn "Die Funktion A besitzt kein lineares Wachstum!"
    end
    labels = ["dein Ton"];
    x = 0:0.0001:0.1;
    y = zeros(1001);
    for i = 1:1001
        y[i] = A(x[i])*sin(2*pi*200*x[i]);
    end
    plot(x,y, label = labels)
end

function AB1A1dPlay(A)
    eps = 1/44100
    x = range(0.0,stop=1,step=eps);
    y = zeros(Float64,length(x));
    for i = 1:length(y)
        y[i] = A(x[i])*sin(2*pi*200*x[i]);
    end
    #display(y)
    wavwrite(y, "SinustonAB1A1d.wav", Fs=44100);
    #y, fs = wavread("SinustonAB1A1d.wav");
    #wavplay(y, fs)
	inline_audioplayer("SinustonAB1A1d.wav")
end

function AB1Aufgabe1d(A)
    if A(6) != 0
		@warn "Bei t = 1 hat die Amplitude nicht den Wert 0!"
    end
    m1 = A(1)-A(0);
    m2 = A(2)-A(1);
    if abs(m1 - m2) > 0.0000000001
		@warn "Die Funktion A besitzt kein lineares Abklingverhalten!"
    end
    labels = ["dein Ton"];
    x = 0:0.0001:6;
    y = zeros(length(x));
    for i = 1:length(y)
        y[i] = A(x[i])*sin(2*pi*200*x[i]);
    end
    plot(x,y, label = labels)
end

function AB1A2a1Play()
    PlaySinusSound(500, 1, "SinustonAB1A2a")

      # Display 500 Hz sin sound.
      amplitude = 1
      frequenz = 500

      x = 0:0.0001:0.1;
      y = zeros(1001);
      for i = 1:1001
          y[i] = amplitude*sin(2*pi*frequenz*x[i]);
      end
      labels = ["Sinuston"];
      plot(x,y, label = labels)
end

function AB1A2a2Play(F_hoeher,F_tiefer)
    eps = 1/44100
    x = range(0.0,stop=2,step=eps);
    y = zeros(Float64,length(x));
    for i = 1:length(y)
        y[i] = sin(2*pi*200*x[i]);
    end
    z = zeros(Float64,length(x));
    for i = 1:length(z)
        z[i] = sin(F_hoeher*2*pi*200*x[i]);
    end
    w = zeros(Float64,length(x));
    for i = 1:length(w)
        w[i] = sin(F_tiefer*2*pi*200*x[i]);
    end
    ton = [y; z; w];
    #display(ton)
    wavwrite(ton, "SinustonAB1A2a.wav", Fs=44100);
    #ton, fs = wavread("SinustonAB1A2a.wav");
    #wavplay(ton, fs)
	inline_audioplayer("SinustonAB1A2a.wav")
end

function AB1Aufgabe2a(F_hoeher,F_tiefer)
    if F_hoeher >= 6 || F_hoeher <= 0
		@warn "F_hoeher liegt nicht zwischen 0 und 6!"
    end
    if F_tiefer >= 6 || F_tiefer <= 0
		@warn "F_tiefer liegt nicht zwischen 0 und 6!"
    end
    if F_hoeher <= 1
		@warn "Der erste Ton ist nicht höher als der Referenzton!"
    end
    if F_tiefer >= 1
		@warn "Der zweite Ton ist nicht tiefer als der Referenzton!"
    end
    x = 0:0.00001:0.01;
    y = zeros(length(x),3);
    for i = 1:length(x)
        y[i,1] = sin(2*pi*200*x[i]);
        y[i,2] = sin(F_hoeher*2*pi*200*x[i]);
        y[i,3] = sin(F_tiefer*2*pi*200*x[i]);
    end
    labels = ["Referenz" "hoeher" "tiefer"];
    plot(x,y, label = labels)
end

function AB1A2bPlay(F)
    eps = 1/44100
    x = range(0.0,stop=1,step=eps);
    y = zeros(Float64,length(x));
    for i = 1:length(y)
        y[i] = sin(F(x[i])*2*pi*200*x[i]);
    end
    #display(y)
    wavwrite(y, "SinustonAB1A2b.wav", Fs=44100);
    #y, fs = wavread("SinustonAB1A2b.wav");
    #wavplay(y, fs)
	inline_audioplayer("SinustonAB1A2b.wav")
end

function AB1Aufgabe2b(F)
    if F(0) != 1
		@warn "Zum Zeitpunkt t = 0 hat die Frequenz von g nicht den Wert 200 Hz!"
    end
    if F(0.1) != 2
		@warn "Zum Zeitpunkt t = 0.1 hat die Frequenz von g nicht den Wert 400Hz!"
    end
    m1 = F(1)-F(0);
    m2 = F(2)-F(1);
    if abs(m1-m2) > 0.00000001
		@warn "Die Funktion F besitzt kein lineares Wachstum!"
    end
    labels = ["dein Signal"];
    x = 0:0.0001:0.1;
    y = zeros(1001);
    for i = 1:1001
        y[i] = sin(F(x[i])*2*pi*200*x[i]);
    end
    plot(x,y, label = labels)
end

function AB1ZAPlay(A,a,b)
    eps = 1/44100
    a_ton = a - 1;
    b_ton = b + 1;
    x = range(a_ton,stop=b_ton,step=eps);
    y = zeros(Float64,length(x));
    for i = 1:length(y)
        y[i] = A(x[i])*sin(2*pi*200*x[i]);
    end
    #display(y)
    wavwrite(y, "SinustonAB1A2b.wav", Fs=44100);
    #y, fs = wavread("SinustonAB1A2b.wav");
    #wavplay(y, fs)
	inline_audioplayer("SinustonAB1A2b.wav")
end

function AB1ZusatzaufgabeA(A4,t_min,a,b)
    if t_min <= a || t_min >= b
		@warn " t_min liegt nicht zwischen a und b!"
    end
    min = 1;
    for i = -10:10
        if A4(t_min + i/100) < A4(t_min)
            min = 0;
        end
    end
    if min == 0
		@warn "Bei t = t_min ist die Amplitude nicht minimal!"
    end
    labels = ["dein Ton"];
    if t_min > a && t_min < b
        x = a:(b-a)/10000:b;
        y = zeros(length(x));
        for i = 1:length(y)
            y[i] = A4(x[i])*sin(2*pi*200*x[i]);
        end
        plot(x,y, label = labels)
    end
end

function AB1ZBPlay(A,a,b)
    eps = 1/44100
    x = range(a,stop=b,step=eps);
    y = zeros(Float64,length(x));
    for i = 1:length(y)
        y[i] = A(x[i])*sin(2*pi*50*x[i]);
    end
    #display(y)
    wavwrite(y, "SinustonAB1ZB.wav", Fs=44100);
    #y, fs = wavread("SinustonAB1ZB.wav");
    #wavplay(y, fs)
	inline_audioplayer("SinustonAB1ZB.wav")
end

function AB1ZusatzaufgabeB(A5,t_max,a,b)
    if t_max <= a || t_max >= b
		@warn "t_max liegt nicht zwischen a und b!"
    end
    if A5(a) != 0
		@warn "Bei t = a hat die Amplitude nicht den Wert 0!"
    end
    if A5(b) != 0
		@warn "Bei t = b hat die Amplitude nicht den Wert 0!"
    end
    max = 1;
    for i = -10:10
        if A5(t_max + i/100) > A5(t_max)
            max = 0;
        end
    end
    if max == 0
		@warn "Bei t = t_max ist die Amplitude nicht maximal!"
    end
    labels = ["dein Ton"];
    if t_max > a && t_max < b
        x = a:(b-a)/10000:b;
        y = zeros(length(x));
        for i = 1:length(y)
            y[i] = A5(x[i])*sin(2*pi*50x[i]);
        end
        plot(x,y, label = labels)
    end
end

function AB1ZCPlay(a,F)
    eps = 1/44100
    b = a + 1;
    x = range(a,stop=b,step=eps);
    y = zeros(Float64,length(x));
    for i = 1:length(y)
        y[i] = sin(F(x[i])*2*pi*200*x[i]);
    end
    #display(y)
    wavwrite(y, "SinustonAB1ZC.wav", Fs=44100);
    #y, fs = wavread("SinustonAB1ZC.wav");
    #wavplay(y, fs)
	inline_audioplayer("SinustonAB1ZC.wav")
end

function AB1ZusatzaufgabeC(a,b,F)
    if b <= a
		@warn "b muss größer sein als a!"
    end
    increase = 1;
    for i = 1:100
        if F(a+(b-a)/100*i) < F(a+(b-a)/100*(i+1))
            increase = 0;
        end
    end
    if increase == 0
		@warn "Die Frequenz fällt nicht auf dem Intervall [a,b]!"
    end
    labels = ["dein Signal"];
    if a < b
        x = a:(b-a)/10000:b;
        y = zeros(length(x));
        for i = 1:length(y)
            y[i] = sin(F(x[i])*2*pi*200*x[i]);
        end
        plot(x,y, label = labels)
    end
end

function AB1ZDPlay(a,b,F)
    eps = 1/44100
    x = range(a,stop=b,step=eps);
    y = zeros(Float64,length(x));
    for i = 1:length(y)
        y[i] = sin(F(x[i])*2*pi*200*x[i]);
    end
    #display(y)
    wavwrite(y, "SinustonAB1ZD.wav", Fs=44100);
    #y, fs = wavread("SinustonAB1ZD.wav");
    #wavplay(y, fs)
	inline_audioplayer("SinustonAB1ZD.wav")
end

function AB1ZusatzaufgabeD(a,b,F)
    if b <= a
		@warn "b muss größer sein als a!"
    end
    labels = ["dein Signal"];
    if a < b
        x = a:(b-a)/10000:b;
        y = zeros(length(x));
        for i = 1:length(y)
            y[i] = sin(F(x[i])*2*pi*200*x[i]);
        end
        plot(x,y, label = labels)
    end
end

end


"""
    SeisRadon_fx(In,Par,itype);
RADON_FX: Operators for f-x, tau-p forward and adjoint Radon transforms
          to compute linear or  parabolic Radon transform.
IN    Radon coefficients if itype =  1 (Forward transform)  with In(np,nt)
      CMP gather         if itype = -1 (Adjoint transform)  with In(nh,nt)

OUT   CMP gather         if itype =  1 (Forward transform)  with Out(nh,nt)
      Radon coeffcients  if itype = -1 (Adjoint transform)  with Out(np,nt)
# Parameters
- `Par.h`:vector containing the nh offsets
- `Par.p`:vector containing the np curvatures normalized by far offset (s) if itype='parab'
- `Par.p`:vector containing the np dips in s/m if itype='linear'
- `Par.dt`:sampling interval
- `Par.f`:frequency corners of BP operator - this acts like a zero phase wavelet.
- `Par.transf `:'parab', 'linear'
itype =  1  :  forward
itype = -1  :  adjoint
"""
function SeisRadon_fx(In,Par,itype);
    h = Par[1][:h];
    p = Par[2][:p];
    dt = Par[3][:dt];
    f  = Par[4][:f];
    f1 = f[1];
    f2 = f[2];
    f3 = f[3];
    f4 = f[4];
    np = length(p);
    nh = length(h);
    if itype == 1;
       m = In;
       (nt, np) = size(m);
       d = zeros(nt,nh);
       m = SeisBPFilter(m,dt,f1,f2,f3,f4);
    end;
    if itype == -1;
       d = In;
       (nt, nh) = size(d);
       m = zeros(nt,np);
    end;
    if cmp(Par[5][:transf],"parab")==0 ;
      hmax = maximum(abs.(h));
      h=(h ./hmax) .^2;
    end
    nfft = 2*nextpow(2,nt);
    ilow  = floor(Int,f[1]*dt*nfft)+1;
    ihigh = floor(Int,f[4]*dt*nfft)+1;
     A = h*p';

     if itype == 1;
         m = cat( m, zeros(nfft-size(m)[1], size(m)[2]),dims=1)
         M = fft(m,1);
         D = zeros(ComplexF64,nfft,nh);

        for ifreq = ilow:ihigh
             w = 2*pi*(ifreq-1)/(nfft*dt);
             L = exp.(im*w*A');
             x = M[ifreq,:];
             y = L'*x;
             D[ifreq,:] =y;
         end

         for k=Int(nfft/2+2):nfft
             D[k,:] = conj(D[nfft-k+2,:]);
         end

         d = ifft(D,1);
         d = real(d[1:nt,:]);

         Out = d;

     else
         d = cat(d,zeros(nfft-size(d)[1], size(d)[2]),dims=1)
         D = fft(d,1);
         M = zeros(ComplexF64,nfft,np);

         for ifreq = ilow:ihigh
             w = 2*pi*(ifreq-1)/(nfft*dt);
             L = exp.(im*w*A');
             y = D[ifreq,:];
             x = L*y;
             M[ifreq,:]=x;
         end

         for k = Int(nfft/2+2):nfft
             M[k,:] = conj(M[nfft-k+2,:]);
         end

         m = ifft(M,1);
         m = real( m[1:nt,:]);

         Out = m;

     end
     if itype == -1;
         m = SeisBPFilter(m,dt,f1,f2,f3,f4);
         Out = m;
     end;
     return Out
  end


# Algorithm 6 from Tight and rigourous error bounds. relative error < 3u²
@inline function add_dddd_dd(x::Tuple{T,T}, y::Tuple{T,T}) where T<:IEEEFloat
    xhi, xlo = x
    yhi, ylo = y
    hi, lo = two_sum(xhi, yhi)
    thi, tlo = two_sum(xlo, ylo)
    c = lo + thi
    hi, lo = two_hilo_sum(hi, c)
    c = tlo + lo
    hi, lo = two_hilo_sum(hi, c)
    return hi, lo
end

# Algorithm 6 from Tight and rigourous error bounds. relative error < 3u²
# reworked for subtraction
@inline function sub_dddd_dd(x::Tuple{T,T}, y::Tuple{T,T}) where T<:IEEEFloat
    xhi, xlo = x
    yhi, ylo = y
    hi, lo = two_diff(xhi, yhi)
    thi, tlo = two_diff(xlo, ylo)
    c = lo + thi# Algorithm 9 from Tight and rigourous error bounds. relative error <= 2u²
    hi, lo = two_hilo_sum(hi, c)
    c = tlo + lo
    hi, lo = two_hilo_sum(hi, c)
    return hi, lo
end

# Algorithm 12 from Tight and rigourous error bounds.  relative error <= 5u²
@inline function mul_dddd_dd(x::Tuple{T,T}, y::Tuple{T,T}) where T<:IEEEFloat
    xhi, xlo = x
    yhi, ylo = y
    hi, lo = two_prod(xhi, yhi)
    t = xlo * ylo
    t = fma(xhi, ylo, t)
    t = fma(xlo, yhi, t)
    t = lo + t
    hi, lo = two_hilo_sum(hi, t)
    return hi, lo
end

#=
# reltime 60
@inline function dvi_dddd_dd(x::Tuple{T,T}, y::Tuple{T,T}) where T<:IEEEFloat
    yinv = inv_dd_dd(y)
    hi, lo = mul_dddd_dd(x, yinv)
    return hi, lo
end
=#
# reltime 107

@inline function dvi_dddd_dd(x::Tuple{T,T}, y::Tuple{T,T}) where {T<:IEEEFloat}
    xhi, xlo = x
    yhi, ylo = y
    hi = xhi / yhi
    uh, ul = two_prod(hi, yhi)
    lo = ((((xhi - uh) - ul) + xlo) - hi*ylo)/yhi
    hi,lo = two_hilo_sum(hi, lo)
    return hi, lo
end

# reltime 40

@inline function dvi_dddd_dd_fast(x::Tuple{T,T}, y::Tuple{T,T}) where T<:IEEEFloat
    xhi, xlo = x
    yhi, ylo = y
    hi = xhi / yhi
    uh, ul = two_prod(hi, yhi)
    lo = ((((xhi - uh) - ul) + xlo) - hi*ylo)/yhi
    hi,lo = two_hilo_sum(hi, lo)
    return hi, lo
end


module timeevolution_simple

using ..states
using ..operators
using ODE

export master

"""
Evaluate master equation for diagonal jump operators.
"""
function dmaster(rho::DenseOperator, H::Operator, gamma::Vector, J::Vector, Jdagger::Vector)
    drho = -1im * (H*rho - rho*H)
    for n = 1:length(J)
        drho = drho + gamma[n]*(J[n]*rho*Jdagger[n] - Jdagger[n]*(J[n]*rho)/Complex(2) - rho*Jdagger[n]*J[n]/Complex(2))
    end
    return drho
end

"""
Evaluate master equation for non-diagonal jump operators.
"""
function dmaster(rho::DenseOperator, H::Operator, gamma::Matrix, J::Vector, Jdagger::Vector)
    drho = -1im * (H*rho - rho*H)
    for m=1:length(J), n=1:length(J)
       drho += gamma[m,n]*(J[m]*rho*Jdagger[n] - Jdagger[n]*(J[m]*rho)/Complex(2) - rho*Jdagger[n]*J[m]/Complex(2))
    end
    return drho
end

"""
Integrate master equation.
"""
function master(T::Vector, rho0::DenseOperator, H::Operator, J::Vector;
                    Jdagger=map(dagger,J),
                    gamma::Union{Real, Vector, Matrix}=ones(Int, length(J)),
                    kwargs...)
    operators.check_samebases(rho0, H)
    for j=J
        @assert typeof(j) <: Operator
        operators.check_samebases(rho0, j)
    end
    for j=Jdagger
        @assert typeof(j) <: Operator
        operators.check_samebases(rho0, j)
    end
    @assert length(J)==length(Jdagger)
    if typeof(gamma)<:Real
        gamma = ones(typeof(gamma), length(J))*gamma
    end
    nl = prod(rho0.basis_l.shape)
    nr = prod(rho0.basis_r.shape)
    N = nl*nr
    as_operator(x::Vector{Complex128}) = DenseOperator(rho0.basis_l, rho0.basis_r, reshape(x, nl, nr))
    as_vector(rho::DenseOperator) = reshape(rho.data, N)
    f(t::Float64, x::Vector{Complex128}) = as_vector(dmaster(as_operator(x), H, gamma, J, Jdagger))
    tout, x_t = ode45(f, as_vector(rho0), T; kwargs...)
    rho_t = DenseOperator[as_operator(x) for x=x_t]
    return tout, rho_t
end

master(T::Vector, psi0::Ket, H::Operator, J::Vector; kwargs...) = master(T, tensor(psi0, dagger(psi0)), H, J; kwargs...)

"""
Evaluate Schroedinger equation for ket states.
"""
function dschroedinger(psi::Ket, H::Operator)
    return -1im*H*psi
end

"""
Evaluate Schroedinger equation for bra states.
"""
function dschroedinger(psi::Bra, H::Operator)
    return 1im*psi*H
end

"""
Integrate Schroedinger equation.
"""
function schroedinger{T<:StateVector}(tspan::Vector, psi0::T, H::Operator; kwargs...)
    as_statevector(x::Vector{Complex128}) = T(psi0.basis, x)
    as_vector(psi::T) = psi.data
    f(t::Float64, x::Vector{Complex128}) = as_vector(dschroedinger(as_statevector(x), H))
    tout, x_t = ode45(f, as_vector(psi0), tspan; kwargs...)
    psi_t = T[as_statevector(x) for x=x_t]
    return tout, psi_t
end

end  # module


mutable struct SourceTerm
    λ::Float64
    μ::Float64
end

function (st::SourceTerm)(x::Float64,p::Float64)
    return st.λ*tanh(st.μ*(p-x))
end


# A walk-through example (using StanSample.jl)

using StanSample
using Distributions
using DataFrames
using MonteCarloMeasurements
using Tables

N = 100
df = DataFrame(
    height = rand(Normal(10, 2), N),
    leg_prop = rand(Uniform(0.4, 0.5), N),
)
df.leg_left = df.leg_prop .* df.height + rand(Normal(0, 0.02), N)
df.leg_right = df.leg_prop .* df.height + rand(Normal(0, 0.03), N)

stan6_1 = "
data {
  int <lower=1> N;
  vector[N] H;
  vector[N] LL;
  vector[N] LR;
}
parameters {
  real a;
  vector[2] b;
  real <lower=0> sigma;
}
model {
  vector[N] mu;
  mu = a + b[1] * LL + b[2] * LR;
  a ~ normal(10, 100);
  b ~ normal(2, 10);
  sigma ~ exponential(1);
  H ~ normal(mu, sigma);
}
";

m6_1s = SampleModel("m6.1s", stan6_1);
data = (H = df.height, LL = df.leg_left, LR = df.leg_right, N = size(df, 1))
rc6_1s = stan_sample(m6_1s; data);

if success(rc6_1s)
    chns = read_samples(m6_1s, :nesteddataframe)
    chns |> display

    # Obtain  b.1 , b,2 as a matrix

    chns = matrix(chns, :b)

    # Or use read_samples to only use chains 2 and 4 using the chains kwarg.

    chns2 = read_samples(m6_1s, :dataframe; chains=[2, 4])
    mean.(eachcol(chns2)) |> display

end

# Some more options:

init = (a = 2.0, b = [1.0, 2.0], sigma = 1.0)
#rc6_2s = stan_sample(m6_1s; data, init);
rc6_2s = stan_sample(m6_1s; data);

if success(rc6_2s)

    read_summary(m6_1s, true)
    println()

    post6_1s_df = read_samples(m6_1s, :dataframe)

    part6_1s = read_samples(m6_1s, :particles)
    part6_1s |> display
    println()

    nt6_1s = read_samples(m6_1s, :namedtuple)
    mean(nt6_1s.b, dims=2)
end


#=

  Euler #7 in Julia.

  Problem 7
  """
  By listing the first six prime numbers: 2, 3, 5, 7, 11, and 13, we can see
  that the 6^(th) prime is 13.

  What is the 10001^(st) prime number?
  """

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

=#

include("Euler.jl")

# 0.00416286s
function euler7a()
    i = 11
    c = 4
    while c < 10001
        if isPrime(i)
            c += 1
        end
        i += 2
    end
    return i-2
end

# 0.00981122s
function euler7b()
    p = primes(200000) # slightly cheating...
    return p[10001];
end

# run_euler(euler7a)
run_euler(euler7b)


module TestSNE

using StochasticNeighborhoodEmbeddings
using StochasticNeighborhoodEmbeddings: conditional_neighbor, joint_neighbor,
        two_bump_data, gaussian, euclidean, k_nearest_neighbors

using Test

d = 16
n = 32
X = two_bump_data(d, n)
x = [c for c in eachcol(X)]
# TODO: standardize
@testset "neighbor matrices" begin
    σ² = 1:n
    P = conditional_neighbor(X, σ²)
    @test vec(sum(P, dims = 1)) ≈ ones(n)

    # test if sigma is constant for each column
    p1 = @. gaussian(euclidean(x[1:1], x[2:end]) / σ²[1])
    p1 ./= sum(p1)
    @test P[2:end, 1] ≈ p1
end

using StochasticNeighborhoodEmbeddings: optimize_variances, perplexity, entropy
@testset "variance optimization" begin
    u = 30
    σ² = optimize_variances(X, u)
    @test all(!isnan, σ²)
    P = conditional_neighbor(X, σ²)
    # println([perplexity(p) for p in eachcol(P)].-u)
    @test all(eachcol(P)) do p isapprox(perplexity(p), u, atol = 1e-3) end
end

using SparseArrays
@testset "sparse" begin
    # now sparse case
    σ² = 1:n
    u = 5 # perplexity
    neighbors = k_nearest_neighbors(X, floor(3u))
    P = conditional_neighbor(X, σ², neighbors)
    @test P isa SparseMatrixCSC
    @test vec(sum(P, dims = 1)) ≈ ones(n)

    σ² = optimize_variances(X, u, neighbors)
    P = conditional_neighbor(X, σ², neighbors)
    # println([perplexity(p) for p in eachcol(P)].-u)
    @test all(eachcol(P)) do p isapprox(perplexity(p), u, atol = 1e-3) end
end

@testset "gradient" begin

end

# u = 30 # perplexity: higher value means distribution P[:, i] has larger entropy -> more uneven
# Y = tsne(X, u)

end # module


# Function for running Hamiltonian Monte Carlo (HMC)
function hmc(model::Model, opts::HmcOpts)
  mcmc = Array(Float64, opts.mcmc.nPostBurnin, model.nPars)
  z = Array(Float64, opts.mcmc.nPostBurnin, model.nPars)
  proposed, accepted = 0., 0.

  println("Running burn-in iterations...")

  currentPars = model.randPrior()
  currentLogPosterior = model.logPosterior(currentPars)
  currentGradLogPosterior = model.gradLogPosterior(currentPars)
  invMass = inv(opts.mass)
  leapStep = opts.setLeapStep(1, 0., opts.mcmc.n, opts.mcmc.nBurnin, 0.)

  for i = 1:opts.mcmc.n
    proposed += 1

    proposedPars = copy(currentPars)
    momentum = opts.mass'*randn(model.nPars)
    currentHamiltonian = -currentLogPosterior+(momentum'*(invMass*momentum))/2
    randomSteps = ceil(rand()*opts.nLeaps)
    proposedGradLogPosterior = copy(currentGradLogPosterior)
    for j = 1:randomSteps
      momentum = momentum+(leapStep/2)*proposedGradLogPosterior
      proposedPars = proposedPars+leapStep*momentum
      proposedGradLogPosterior = model.gradLogPosterior(proposedPars)
      momentum = momentum+(leapStep/2)*proposedGradLogPosterior
    end
    proposedLogPosterior = model.logPosterior(proposedPars)
    proposedHamiltonian = -proposedLogPosterior+(momentum'*(invMass*momentum))/2

    ratio = (currentHamiltonian-proposedHamiltonian)[1]

    if ratio > 0 || (ratio > log(rand()))
      accepted += 1

      currentPars = copy(proposedPars)
      currentLogPosterior = copy(proposedLogPosterior)
      currentGradLogPosterior = copy(proposedGradLogPosterior)
    end

    if i > opts.mcmc.nBurnin
      mcmc[i-opts.mcmc.nBurnin, :] = currentPars
      z[i-opts.mcmc.nBurnin, :] = -currentGradLogPosterior/2
    end

    if mod(i, opts.mcmc.monitorRate) == 0
      acceptanceRatio = accepted/proposed

    leapStep = opts.setLeapStep(i, acceptanceRatio, opts.mcmc.n,
      opts.mcmc.nBurnin, leapStep)

      println("Iteration $i of $(opts.mcmc.n): ", round(100*acceptanceRatio, 2),
        " % acceptance ratio")

      proposed, accepted = 0., 0.
    end

    if i == opts.mcmc.nBurnin
      println("Burn-in completed...\n\nRunning post burn-in MCMC...");
    end
  end

  return mcmc, z
end