# burgers.jl
##########################
using FFTW, DSP, LinearAlgebra, Plots

# initial data
N=256;
n=128;
M=1000;
u=zeros(ComplexF64, N,M);
v=zeros(ComplexF64, N,M);
dv=zeros(ComplexF64, N,M);
z=zeros(ComplexF64, 2*N,M);
w=zeros(ComplexF64, N,M);
U=zeros(ComplexF64, N,N);
V=zeros(ComplexF64, N,M);

for p=1:192
    u[p,1]=8*sin(pi*p/192);
end

# setting
E=zeros(ComplexF64, N, N);
for k=1:N
    E[k,k]=1;
end

t=1/100;

# sqrt(M) times Fourier matrix
omega=exp(im*2*pi/N);
F=zeros(ComplexF64, N, N);
for k=1:N
    for l=1:N
    F[k,l]=omega^(-(k-1)*(l-1));
    end
end

D1=zeros(ComplexF64, N, N)
D2=zeros(ComplexF64, N, N)
for k=1:n
    D1[k,k]=k-1;
    D1[k+n,k+n]=k-n;
    D2[k,k]=(k-1)^2;
    D2[k+n,k+n]=(k-n)^2;
end

G=-(4*pi^2*t*N^(-2))*D2;
A=E+G+G^2/2+G^3/6;
B=-t*N^(-1)*(E+G/2+G^2/6+G^3/24);

v[:,1]=F*u[:,1];
dv[:,1]=2*pi*im*D1*v[:,1];
z[:,1]=[v[:,1];v[:,1]];

for k=1:N
    for q=1:N
        w[k,1]=w[k,1]+z[N+k-q+1,1]*dv[q,1]/N;
    end
end

for l=2:M
    v[:,l]=A*v[:,l-1]+B*w[:,l-1];
    dv[:,l]=2*pi*im*D1*v[:,l];
    z[:,l]=[v[:,l];v[:,l]];
    for k=1:N
        for q=1:N
            w[k,l]=w[k,l]+z[N+k-q+1,l]*dv[q,l]/N;
        end
    end
end

U=F'*v/N;

V=zeros(ComplexF64, N,M)
for l=1:M
    V[:,l]=U[:,M-l+1];
end

p=plot(real.(U'),
   title="IVP for Burgers' equation",
    st=:heatmap,
    color=:jet,
    yaxis="time t",
    xaxis="position x")
plot!(yticks = ([0 1000;], [0 10]))
plot!(xticks = ([0 256;], [0 256]))

q=surface(real.(V),
        title="IVP for Burgers' equation",
        camera=(80,60),
        color=:jet,
        xaxis="time t",
        yaxis="position x")
plot!(xticks = ([0 1000;], [10 0]))
plot!(yticks = ([0 256;], [0 256]))

plot(p, q,
     layout=(1,2),
     size=(1600,600),
     margin=Plots.Measures.Length(:mm, 10.0))
savefig("burgers.png")


using DataDrivenDiffEq
using OrdinaryDiffEq
using Plots
gr()

# Define measurements from unstable system with known control input
X = [4 2 1 0.5 0.25; 7 0.7 0.07 0.007 0.0007]
U = [-4 -2 -1 -0.5]
B = Float32[1; 0]

# See fail with unknown input
sys = DMDc(X, U)
# But with a little more knowledge
sys = DMDc(X, U, B = B)

# Extract the DMD from inside DMDc
get_dynamics(sys)
# Acess all the other stuff
eigen(sys)
eigvals(sys)
eigvecs(sys)
isstable(sys)
# Get unforced dynamics
dudt_ = dynamics(sys)
prob = DiscreteProblem(dudt_, X[:, 1], (0., 10.))
sol_unforced = solve(prob,  FunctionMap())
plot(sol_unforced)
sol_unforced[:,:]
# Create a system with cos control input to stabilize
dudt_ = dynamics(sys, control = (u, p, t) -> -0.5u[1])
prob = DiscreteProblem(dudt_, X[:, 1], (0., 10.))
sol = solve(prob, FunctionMap())

plot!(sol)


#=
# 3D Poisson, Dirichlet bc
# Uses Dendro imported as a custom gen target
=#
if !@isdefined(Femshop)
    include("../Femshop.jl");
    using .Femshop
end
init_femshop("poissondendro");

useLog("poissondendrolog", level=3)

# default values (max_depth=6, wavelet_tol = 0.1, partition_tol = 0.3, solve_tol = 1e-6, max_iters = 100)
generateFor("target_dendro_cg.jl", params=(7, 0.01, 0.3, 0.000001, 1000))

domain(3)
functionSpace(order=4)

u = variable("u")
testSymbol("v")

boundary(u, 1, DIRICHLET, "0")

f = coefficient("f", "-30*pi*pi*sin(5*pi*x)*sin(2*pi*y)*sin(pi*z)")
weakForm(u, "-dot(grad(u),grad(v)) - f*v")

build_octree_with(f);
solve(u);

finalize_femshop()


using Qaintessent

"""
    qft_circuit(N)

Construct the quantum Fourier transform circuit for `N` qubits.
"""
function qft_circuit(N)
    main = vcat([
        CircuitGate[j == i ?
             circuit_gate(i, HadamardGate()) :
             circuit_gate(j, PhaseShiftGate(2π/2^(j-i+1)), i) for j in N:-1:i]
        for i in N:-1:1]...)
    swap = [circuit_gate(i, N-i+1, SwapGate()) for i in 1:(N÷2)]
    return N > 1 ? CircuitGate[main; swap] : main
end

"""
    toffoli_circuit(trg, cntrl, N)

construct the circuit for decomposing the Toffoli gate in Figure 4.9 of Nielsen and Chuang (2000). the constructed toffoli gate acts in circuit of `N` qubits, has controls
    on wires in Tuple `cntrl` and has target on wire `trg`.
returns a `CircuitGateChain{N}` object.
"""
function toffoli_circuit(trg::Integer, cntrl::Tuple{<:Integer,<:Integer})
    [
        circuit_gate(trg, HadamardGate()),
        circuit_gate(trg, X, cntrl[2]),
        circuit_gate(trg, TdagGate()),
        circuit_gate(trg, X, cntrl[1]),
        circuit_gate(trg, TGate()),
        circuit_gate(trg, X, cntrl[2]),
        circuit_gate(trg, TdagGate()),
        circuit_gate(trg, X, cntrl[1]),
        circuit_gate(trg, TGate()),
        circuit_gate(cntrl[2], TGate()),
        circuit_gate(cntrl[2], X, cntrl[1]),
        circuit_gate(trg, HadamardGate()),
        circuit_gate(cntrl[1], TGate()),
        circuit_gate(cntrl[2], TdagGate()),
        circuit_gate(cntrl[2], X, cntrl[1]),
    ]
end

"""
    vbe_adder_circuit(N)

Construct an in-place adder for 2 integers represented by `N` qubits.
Based on ripple-carry adder circuit in Vedral et. al (Phys. Rev. A 54, 147 (1996), arXiv:quant-ph/9511018)
Returns a CircuitGateChain{3N+1} as there are N+1 ancillary wires

If the two added integers are represented as:

``A = a_{0}\\times 2^{0} + a_{1} \\times 2^{1} + a_{2} \\times 2^{2} + .. + a_{N} \\times 2^{N} \\\\ B = b_{0}\\times 2^{0} + b_{1} \\times 2^{1} + b_{2} \\times 2^{2} + .. + b_{N} \\times 2^{N}``

The input index should be ``a_{0}a_{1}a_{2}..a_{N}b_{0}b_{1}b_{2}..b_{N} + 1`` with ``a_{0}`` as the fastest running index

The output index will be in the form ``a_{0}a_{1}a_{2}..a_{N}c_{0}c_{1}c_{2}..c_{N} + 1`` where:

``C = (A+B) \\% (2^{N} + 1) = c_{0}\\times 2^{0} + c_{1} \\times 2^{1} + c_{2} \\times 2^{2} + ... + c_{N} \\times 2^{N}``
"""
function vbe_adder_circuit(N::Integer)
    M = 3N + 1
    carry(a,b,c,d) =
    [
        circuit_gate(d, X, (b, c)),
        circuit_gate(c, X, b),
        circuit_gate(d, X, (a, c)),
    ]
    rcarry(a,b,c,d) =
    [
        circuit_gate(d, X, (a, c)),
        circuit_gate(c, X, b),
        circuit_gate(d, X, (b, c)),
    ]
    sum(a,b,c) =
    [
        circuit_gate(c, X, b),
        circuit_gate(c, X, a),
    ]
    rsum(a,b,c) =
    [
        circuit_gate(c, X, a),
        circuit_gate(c, X, b),
    ]
    cgc = carry(2N+1, 1, N+1, 2N+2)
    for i in 2:N
        append!(cgc, carry(2N+i, i, N+i, 2N+i+1))
    end

    append!(cgc, [circuit_gate(2N, X, N)])
    append!(cgc, sum(3N, N, 2N))
    for i in N-1:-1:1
        append!(cgc, rcarry(2N+i, i, N+i, 2N+i+1))
        append!(cgc, sum(2N+i, i, N+i))
    end
    return cgc
end

"""
    p_round(P, M, N)
        execute a P round for qcla algorithms
"""
function p_round(P::Function, ::Integer, N::Integer)
    # Add gates for P rounds
    pchain = CircuitGate[]
    tmax = floor(Int, log(2, N)) - 1
    for t in 1:tmax
        mmax = floor(Int, N/2^t) - 1
        for m in 1:mmax
            push!(pchain, circuit_gate(P(t, m), X, (P(t-1, 2m), P(t-1, 2m+1))))
        end
    end
    return pchain
end

"""
    p_round(P, M, N)
        execute a P round for qcla algorithms
"""
function p_round(P::Function, ::Integer, ::Integer, Ñ::Integer)
    # Add gates for P rounds
    pchain = CircuitGate[]
    tmax = floor(Int, log(2, Ñ)) - 1
    for t in 1:tmax
        mmax = floor(Int, Ñ/2^t) - 1
        for m in 1:mmax
            push!(pchain, circuit_gate(P(t, m), X, (P(t-1, 2m), P(t-1, 2m+1))))
        end
    end
    return pchain
end

"""
    g_round(P, M, N)
        execute a P round for qcla algorithms
"""
function g_round(P::Function, G::Function, ::Integer, N::Integer)
    # Add gates for G rounds
    gchain = CircuitGate[]
    tmax = floor(Int, log(2, N))
    for t in 1:tmax
        mmax = floor(Int, N/2^t) - 1
        for m in 0:mmax
            push!(gchain, circuit_gate(G(2^t*m + 2^t), X, (G(2^t*m+2^(t-1)) , P(t-1, 2m+1))))
        end
    end
    return gchain
end

function g_round(P::Function, G::Function, M::Integer, ::Integer, Ñ::Integer)
    # Add gates for G rounds
    gchain = CircuitGate[]
    tmax = floor(Int, log(2, Ñ))
    for t in 1:tmax
        mmax = floor(Int, Ñ/2^t) - 1
        for m in 0:mmax
            push!(gchain, circuit_gate(G(2^t*m + 2^t), X, (G(2^t*m+2^(t-1)), P(t-1, 2m+1))))
        end
    end
    return gchain
end

"""
    c_round(P, G, M, N)
"""
function c_round(P::Function, G::Function, M::Integer, N::Integer)
    cchain = CircuitGate[]
    tmax = floor(Int, log(2, 2N/3))
    for t in tmax:-1:1
        mmax = floor(Int, (N-2^(t-1))/2^t)
        for m in 1:mmax
            push!(cchain, circuit_gate(G(2^t*m + 2^(t-1)), X, (G(2^t*m), P(t-1, 2m))))
        end
    end
    return cchain
end

function c_round(P::Function, G::Function, M::Integer, N::Integer, Ñ::Integer)
    cchain = CircuitGate[]
    tmax = floor(Int, log(2, 2Ñ/3))
    for t in tmax:-1:1
        mmax = floor(Int, (Ñ-2^(t-1))/2^t)
        for m in 1:mmax
            push!(cchain, circuit_gate(G(2^t*m + 2^(t-1)), X, (G(2^t*m), P(t-1, 2m))))
        end
    end
    return cchain
end

"""
    qcla_out_adder_circuit(N)
Construct an out-of-place adder for 2 integers represented by `N` qubits. returns a `CircuitGateChain{3N+1}` object.
Based on quantum carry-lookahead adder circuit by Draper et. al (Quant. Inf. Comp. 6, 351-369 (2006), arXiv:quant-ph/0406142)
Returns a CircuitGateChain{3N+1} as there are N+1 ancillary wires
If the two added integers are represented as:
``A = a_{0}\\times 2^{0} + a_{1} \\times 2^{1} + a_{2} \\times 2^{2} + .. + a_{N} \\times 2^{N} \\\\ B = b_{0}\\times 2^{0} + b_{1} \\times 2^{1} + b_{2} \\times 2^{2} + .. + b_{N} \\times 2^{N}``
The input index should be ``a_{0}a_{1}a_{2}..a_{N}b_{0}b_{1}b_{2}..b_{N} + 1`` as Julia starts indexing at `1` and ``a_{0}`` as the fastest running index
The output index will be in the form ``a_{0}a_{1}a_{2}..a_{N}b_{0}b_{1}b_{2}..b_{N}c_{0}c_{1}c_{2}..c_{N+1} + 1`` where:
``C = A+B = c_{0}\\times 2^{0} + c_{1} \\times 2^{1} + c_{2} \\times 2^{2} + ... + c_{N+1} \\times 2^{N+1}``
"""
function qcla_out_adder_circuit(N)

    anc = N - count_ones(N) - floor(Int, log(N))
    M = 3N + anc + 1

    function a(m::Integer)
        m + 1
    end

    function b(m::Integer)
        N + m + 1
    end

    function s(m::Integer)
        2N + m + 1
    end

    function G(m::Integer)
        s(m)
    end

    function P(l::Integer, m::Integer)
        if l == 0
            return b(m)
        end

        while l > 1
            m += floor(Int, N/2^(l-1)-1)
            l -= 1
        end
        3N + 1 + m
    end

    setup = CircuitGate[]

    push!(setup, circuit_gate(s(1), X, (a(0), b(0))))
    for i in 1:N-1
        push!(setup, circuit_gate(s(i+1), X, (a(i), b(i))))
        push!(setup, circuit_gate(b(i), X, a(i)))
    end
    cgc = setup

    # Add gates for P rounds
    pchain = p_round(P, M, N)
    append!(cgc, pchain)

    # Add gates for G rounds
    gchain = g_round(P, G, M, N)
    append!(cgc, gchain)

    #  Add gates for C rounds
    cchain = c_round(P, G, M, N)
    append!(cgc, cchain)

    # Remove effect of P rounds
    append!(cgc, reverse(pchain))

    # Teardown
    teardown = CircuitGate[]

    push!(teardown, circuit_gate(s(0), X, b(0)))
    push!(teardown, circuit_gate(s(0), X, a(0)))
    for i in 1:N-1
        push!(teardown, circuit_gate(s(i), X, b(i)))
        push!(teardown, circuit_gate(b(i), X, a(i)))
    end
    append!(cgc, teardown)
    return cgc
end

# TODO: Add comparator circuit for QCLA adder
"""
    qcla_inplace_adder_circuit(N)

Construct an in-place adder for 2 integers represented by `N` qubits.
Based on quantum carry-lookahead adder circuit by Draper et. al (Quant. Inf. Comp. 6, 351-369 (2006), arXiv:quant-ph/0406142)
Returns a CircuitGateChain{3N+1} as there are N+1 ancillary wires

If the two added integers are represented as:

``A = a_{0}\\times 2^{0} + a_{1} \\times 2^{1} + a_{2} \\times 2^{2} + .. + a_{N} \\times 2^{N} \\\\ B = b_{0}\\times 2^{0} + b_{1} \\times 2^{1} + b_{2} \\times 2^{2} + .. + b_{N} \\times 2^{N}``

The input index should be ``a_{0}a_{1}a_{2}..a_{N}b_{0}b_{1}b_{2}..b_{N} + 1`` as Julia starts indexing at `1` and ``a_{0}`` as the fastest running index

The output index will be in the form ``a_{0}a_{1}a_{2}..a_{N}c_{0}c_{1}c_{2}..c_{N+1} + 1`` where:

``C = A+B = c_{0}\\times 2^{0} + c_{1} \\times 2^{1} + c_{2} \\times 2^{2} + ... + c_{N+1} \\times 2^{N+1}``
"""
function qcla_inplace_adder_circuit(N)

    n = 1
    anc = 0
    while 2^n < N
        anc += N ÷ 2^n -1
        n += 1
    end
    M = 3N + anc

    function a(m::Integer)
        m < N || error("a only takes m < N")
        m + 1
    end

    function b(m::Integer)
        m < N || error("b only takes m < N")
        N + m + 1
    end

    function G(m::Integer)
        m > 0 || error("G only takes positive integers")
        M - anc - m + 1
        # 2N + m
    end

    function s(m::Integer)
        m <= N || error("a only takes m <= N")
        if m < N
            return b(m)
        end
        G(m)
    end

    function P(l::Integer, m::Integer)
        if l == 0
            return b(m)
        end

        while l > 1
            m += floor(Int, N/2^(l-1)-1)
            l -= 1
        end
        3N + m
    end

    # setup
    setup = CircuitGate[]
    for i in 0:N-1
        push!(setup, circuit_gate(G(i+1), X, (a(i), b(i))))
        push!(setup, circuit_gate(b(i), X, a(i)))
    end

    # Add gates for P rounds
    pchain = p_round(P, M, N)
    cgc = deepcopy(pchain)

    # Add gates for G rounds
    gchain = g_round(P, G, M, N)
    append!(cgc, gchain)

    #  Add gates for C rounds
    cchain = c_round(P, G, M, N)
    append!(cgc, cchain)

    # Remove effect of P rounds
    append!(cgc, reverse(pchain))

    # Run intermediate circuit
    intermediate = CircuitGate[]

    push!(intermediate, circuit_gate(b(0), X))
    for i in 1:N-2
        push!(intermediate, circuit_gate(b(i), X, G(i)))
        push!(intermediate, circuit_gate(b(i), X))
        push!(intermediate, circuit_gate(b(i), X, a(i)))
    end
    if N > 1
        push!(intermediate, circuit_gate(b(N-1), X, G(N-1)))
    end

    if N > 1
        pchain2 = p_round(P, M, N, N-1)
        gchain2 = g_round(P, G, M, N, N-1)
        cchain2 = c_round(P, G, M, N, N-1)
        rpchain2 = reverse(pchain2)
        r = vcat(pchain2, gchain2, cchain2, rpchain2)
    else
        r = CircuitGate[]
    end

    append!(cgc, intermediate)
    append!(cgc, reverse(r))

    # ## Teardown
    teardown = CircuitGate[]

    if N > 1
        push!(teardown, circuit_gate(G(1), X, (a(0), b(0))))
    end
    push!(teardown, circuit_gate(b(0), X))
    for i in 1:N-2
        push!(teardown, circuit_gate(b(i), X, a(i)))
        push!(teardown, circuit_gate(G(i+1), X, (a(i), b(i))))
        push!(teardown, circuit_gate(b(i), X))
    end

    prepend!(cgc, setup)
    append!(cgc, teardown)

    return cgc
end


function local_update!(model::Ising, T::Real, Js::AbstractArray)
    rng = model.rng
    nsites = numsites(model)
    nbonds = numbonds(model)
    mbeta = -1.0/T

    @inbounds for site in 1:nsites
        center = model.spins[site]
        de = 0.0
        for (n,b) in neighbors(model, site)
            de += 2center * model.spins[n] * Js[bondtype(model,b)]
        end
        if rand(rng) < exp(mbeta*de)
            model.spins[site] *= -1
        end
    end

    return nothing
end

function SW_update!(model::Ising, T::Real, Js::AbstractArray)
    rng = model.rng
    ps = -expm1.((-2.0/T).*Js)
    nsites = numsites(model)
    nbonds = numbonds(model)
    nbt = numbondtypes(model)
    activated_bonds = zeros(Int,nbt)
    uf = UnionFind(nsites)
    @inbounds for bond in bonds(model)
        s1,s2 = source(bond), target(bond)
        bt = bondtype(bond)
        if model.spins[s1] == model.spins[s2] && rand(rng) < ps[bt]
            activated_bonds[bt] += 1
            unify!(uf, s1,s2)
        end
    end
    nc = clusterize!(uf)
    clustersize = zeros(Int, nc)
    clusterspin = rand(rng,[1,-1], nc)

    @inbounds for site in 1:nsites
        id = clusterid(uf, site)
        model.spins[site] = clusterspin[id]
        clustersize[id] += 1
    end
    return SWInfo(activated_bonds, clustersize, clusterspin)
end

function Wolff_update!(model::Ising, T::Real, Js::AbstractArray)
    rng = model.rng
    ps = -expm1.((-2.0/T).*Js)
    nsites = numsites(model)

    clustersize = 0
    st = Stack(Deque{Int}())
    center = rand(rng, 1:nsites)
    sp = model.spins[center]
    model.spins[center] *= -1
    push!(st, center)
    @inbounds while !isempty(st)
        clustersize += 1
        s = pop!(st)
        for (n,b) in neighbors(model, s)
            bt = bondtype(model,b)
            if model.spins[n] == sp && rand(rng) < ps[bt]
                model.spins[n] *= -1
                push!(st, n)
            end
        end
    end

    return nothing
end

@gen_convert_parameter(Ising, ("T", 1, 1.0), ("J", numbondtypes, 1.0))


include("../../src/ZZDiffusionBridge.jl")
using LinearAlgebra
"""
    OUSDE <: AbstractModel

dX_t = ν(μ - X_t)dt + dB_t
ν := intensity
μ := mean reversion
M::Array{Float64, 2} := matrix whose element i,j is equal to ∫_0^T ϕ_i ϕ_j dt
V::Vector{Float64} :=  vector whose element i is equal to ∫_0^T ϕ_i dt
bound1::Vector{Float64} :=  vector whose element i is equal to ∫_0^T upbar{ϕ}_i ϕ_j dt
bound2::Vector{Float64} :=  vector whose element i is equal to ∫_0^T downbar{ϕ}_i ϕ_j dt
"""
struct OUSDE <: AbstractModel
    μ::Float64
    ν::Float64
    M::Array{Float64, 2}
    V::Vector{Float64}
    bound1::Vector{Float64}
    bound2::Vector{Float64}
    function OUSDE(μ, ν, L, T)
        new(μ, ν, generate_matrix(L, T), generate_vector(L, T), generate_bound1(L,T), generate_bound2(L,T))
    end
end

dependence_strucute(::OUSDE) = PartialIndependence()
sampling_scheme(::OUSDE) = Regular()

#Poisson rates
"""
    wait_gengaus(a,b,u)
obtaining waiting time for Inhomogeneous Poisson Process
with rate of the form λ(t) = (a + b*t)^+, `a`,`b` ∈ R, `u` random variable
"""
function wait_gengaus(a,b,u)
    if b > 0
        if a < 0
            τ = sqrt(-log(u)*2.0/b) - a/b
        else #a[i]>0
            τ = sqrt((a/b)^2 - log(u)*2.0/b) - a/b
        end
    elseif  b == 0
        if a > 0
            τ = -log(u)/a
        else #a[i] <= 0
            τ = Inf
        end
    else #b[i] < 0
        if a <= 0
            τ = Inf
        elseif -log(u) <= -a^2/b + a^2/(2*b)
            τ = - sqrt((a/b)^2 - log(u)*2.0/b) - a/b
        else
            τ = Inf
        end
    end
end

"""
    λbar(n, S::System, X::OUSDE , u, v)

Poisson time (real bound) for the coefficient `n`
of model `OUSDE` starting at `u` and ending at `v`

λbar = λ No upper bound and no accept reject step
"""
function λbar(n, S::System, X::OUSDE , u::Float64, v::Float64, t::Float64)
    a = S.θ[n]*(X.ν*X.ν*(dot(X.M[n,:], S.ξ) + X.bound1[n]*v + X.bound2[n]*u - X.μ*X.V[n]) + S.ξ[n])
    b = S.θ[n]*((dot(X.M[:,n], S.θ))*X.ν*X.ν + S.θ[n])
    return wait_gengaus(a, b, rand())
end

####
function runall(SHORT = false)
    Random.seed!(1)
    T = 10.0
    clock = 1000.0
    L = 6
    ν = 1.0
    μ = -5
    u = -1.0
    v = 2.0
    X = OUSDE(μ, ν, L, T)
    XX = zz_sampler(X, T, L, u, v, clock)
    if SHORT == false
        burning = 10.0    #burning
        f = clock - 1.0; n = 100
        db = (f-burning)/n
        b =  burning:db:f
        p = plotmixing(XX, b, T, L, u, v)
        hline!(p, [-5.0], color = :blue)
        xaxis!(p, "t")
        yaxis!(p, "X_t")
        display(p)
        #plot the mean of the process
    end
    return XX
end

error("STOP HERE")

runall()
savefig("../../output/ou_1_m5.pdf")


using Gadfly

plot([sin, cos], color=["sin", "cos"], 0, 25)


n = 17
ijs = [(3,1), (3,2), (4,1), (4,2), (4,3), (5,1), (5,3), (5,4), (8,7),
       (9,5), (9,6), (9,7), (9,8), (11,10), (13,10), (13,11), (13,12),
       (14,10), (14,11), (14,12), (14,13),
       (15, 1), (15,3), (15,4), (15,5), (15,7), (15,8), (15,9),
       (16,5), (16,6), (16,9), (16,12), (16,13), (16,14), (16,15),
       (17,10), (17,11), (17,12), (17,13), (17,14), (17,15), (17,16)]
append!(ijs, [(i,i) for i in 1:n])
I, J = Chordal.unzip(ijs)
A = sparse(I, J, ones(length(I)))
A = A + A'
L = qdldl(A; perm=nothing, logical=true).L

@testset "Test elimination tree" begin
    # VA survey, Figure 4.2
    par_true = vec([3 3 4 5 9 9 8 9 15 11 13 13 14 16 16 17 0])
    ch_true =  [Set{Int}() for _ in 1:n]
    ch_true[3] = Set([1, 2])
    ch_true[4] = Set([3])
    ch_true[5] = Set([4])
    ch_true[8] = Set([7])
    ch_true[9] = Set([5, 6, 8])
    ch_true[11] = Set([10])
    ch_true[13] = Set([11, 12])
    ch_true[14] = Set([13])
    ch_true[15] = Set([9])
    ch_true[16] = Set([15, 14])
    ch_true[17] = Set([16])

    et_par = CD.get_etree(L)
    @test all(et_par .== par_true)
    et_ch = CD.get_children_from_par(et_par)
    @test all([issetequal(et_ch[i], ch_true[i]) for i in 1:n])
    perm = CD.get_postordering(et_par, et_ch)
    iperm = invperm(perm)
    @test all([iperm[i] .< et_par[iperm[i]] for i in 1:17 if et_par[iperm[i]] ≠ 0])
end

@testset "Test Clique Tree" begin
    deg⁺_true = vec([4 2 3 2 3 2 3 2 2 4 3 4 3 2 2 1 0])

    et_par = CD.get_etree(L)
    et_ch = CD.get_children_from_par(et_par)
    deg⁺ = CD.get_higher_deg(L)
    @test all(deg⁺_true .== deg⁺)

    vreps, snd_par, snd_mem = CD.max_supernode_etree(L, et_par)
    # Test vrep (Thm 4.2)
    for v in vreps
        @test all([deg⁺[w] < deg⁺[v] + 1 for w in et_ch[v]])
    end
    # Test supernodes
    for (ind, v) in enumerate(vreps)
        nv = count(x->x==ind, snd_mem) - 1
        pv = v
        for k = 1:nv
            pv = et_par[pv]
            @test deg⁺[v] == deg⁺[pv] + k
            @test snd_mem[pv] == ind
        end
    end
    # Test clique tree (Thm 4.3)
    # Test root
    root_et = findfirst(x->x==0, et_par)
    root_ind = findfirst(x->x==0, snd_par)
    @test snd_mem[root_et] == root_ind
    for (ind, v) in enumerate(vreps)
        ind == root_ind && continue
        col_v = vcat([v], rowvals(L)[nzrange(L, v)])
        qv = vreps[snd_par[ind]]
        col_qv = vcat([qv], rowvals(L)[nzrange(L, qv)])
        col_minus_snd = filter(x->(snd_mem[x] != ind), col_v)
        # println(" --- v = $v --- ")
        # @show col_minus_snd
        # @show col_qv
        @test issubset(col_minus_snd, col_qv) && length(col_minus_snd) < length(col_qv)
    end

end


@testset "Solve ODEs" begin
    "sample ode function"
    f_(u,p,t) = -1.01*u

    ExpDecay = DEsteady(func=f_, p=1.0, u0=3.0)

    # Ordinary solving
    prob = SteadyStateProblem(ExpDecay.func, ExpDecay.u0, ExpDecay.p)
    sol = solve(prob, ExpDecay.method)

    # With finc, u, p  wrapper
    sol_ = solveSS(ExpDecay.func, ExpDecay.u0, ExpDecay.p; method=ExpDecay.method )

    # With ode prameters
    sol_2 = solve(ExpDecay)

    # Solve Time-series
    de = ODEtime(func=f_, u0=3.0, p=1.0, tspan=(0.0,100.0), method= AutoTsit5(Rosenbrock23()))
    sol_time = solve(de)

    @test sol.u == sol_.u
    @test sol_2.u == sol.u
end

@testset "Threading Grid" begin
    f_(u,p,t) = -1.01*u
    us = ParameterGrid([(1.,3.,3)])

    de = DEsteady(func = f_, p=1.0, u0=us[1], method= Tsit5() )

    solve(de)

    @test length(solve(de, us)) == length(us) # mulit-threading mehod
end


# This example is described in more detail in the documentation of Trixi

using Trixi, LinearAlgebra, ForwardDiff

equations = CompressibleEulerEquations2D(1.4)

mesh = TreeMesh((-1.0, -1.0), (1.0, 1.0),
                initial_refinement_level=2, n_cells_max=10^5)

solver = DGSEM(polydeg=3, surface_flux=flux_lax_friedrichs,
               volume_integral=VolumeIntegralFluxDifferencing(flux_ranocha))

"""
    initial_condition_isentropic_vortex(x, t, equations::CompressibleEulerEquations2D)

The classical isentropic vortex test case of
- Chi-Wang Shu (1997)
  Essentially Non-Oscillatory and Weighted Essentially Non-Oscillatory
  Schemes for Hyperbolic Conservation Laws
  [NASA/CR-97-206253](https://ntrs.nasa.gov/citations/19980007543)
"""
function initial_condition_isentropic_vortex(x, t, equations::CompressibleEulerEquations2D)
  # needs appropriate mesh size, e.g. [-10,-10]x[10,10]
  # make sure that the inicenter does not exit the domain, e.g. T=10.0
  # initial center of the vortex
  inicenter = SVector(0.0, 0.0)
  # size and strength of the vortex
  iniamplitude = 0.2
  # base flow
  rho = 1.0
  v1 = 1.0
  v2 = 1.0
  vel = SVector(v1, v2)
  p = 10.0
  rt = p / rho                  # ideal gas equation
  cent = inicenter + vel*t      # advection of center
  cent = x - cent               # distance to centerpoint
  cent = SVector(-cent[2], cent[1])
  r2 = cent[1]^2 + cent[2]^2
  du = iniamplitude/(2*π)*exp(0.5*(1-r2)) # vel. perturbation
  dtemp = -(equations.gamma-1)/(2*equations.gamma*rt)*du^2 # isentrop
  rho = rho * (1+dtemp)^(1\(equations.gamma-1))
  vel = vel + du*cent
  v1, v2 = vel
  p = p * (1+dtemp)^(equations.gamma/(equations.gamma-1))
  prim = SVector(rho, v1, v2, p)
  return prim2cons(prim, equations)
end
semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_isentropic_vortex, solver)

u0_ode = compute_coefficients(0.0, semi)

J = ForwardDiff.jacobian((du_ode, γ) -> begin
    equations_inner = CompressibleEulerEquations2D(first(γ))
    semi_inner = Trixi.remake(semi, equations=equations_inner, uEltype=eltype(γ))
    Trixi.rhs!(du_ode, u0_ode, semi_inner, 0.0)
  end, similar(u0_ode), [1.4]); # γ needs to be an `AbstractArray`


# Utility functions for motif finding algorithms

DNA = Dict("A" => 1, "T" => 2, "C" => 3, "G" => 4)

function align(sequences, s, l)
    """
    Return the length(*s*) x *l* alignment matrix for *sequences*
    """
    align = []
    for i = 1:length(s)
        # add the motif substring to align array
        push!(align, sequences[i][s[i]:s[i]+l-1])
    end
    return align
end

function getprofile(alignment)
    """
    Return the profile matrix of an alignment
    """
    profile = zeros(Int32, 4, length(alignment[1]))
    for i = 1:length(alignment)
        for j = 1:length(alignment[1])
            # get the character at position [i,j] in alignment
            nucl = alignment[i][[j]]
            profile[DNA[nucl], j] += 1
        end
    end
    return profile
end

score(profile) = sum(maximum(profile, 1))

function score(sequences, s, l)
    """
    Find the alignment and calculate the score of the corresponding profile
    matrix.
    """
    alignment = align(sequences, s, l)
    profile = getprofile(alignment)
    return score(profile)
end

function probscore(motif, profile)
    """
    Calculate the probability of seeing *motif* given the matrix *profile*.
    Return a probability in [0, 1]
    """

    # normalize profile
    P = profile ./ sum(profile, 1)

    return prod(P[DNA[motif[[i]]], i] for i = 1:length(motif))
end


module Juliaset

using Images
using Compat
using JuliaWebAPI

export juliaset

function genjuliaset(c1::Float64=0.0, c2::Float64=0.65, w::Int=200, h::Int=200)
    a = Array(UInt8, w, h)

    c = c1 + c2 * im

    idx = 1
    for r_val in linspace(-2.0, 2.0, w), im_val in linspace(-2.0, 2.0, h)
        z = r_val + im_val*im
        n = 255
        while abs(z) < 10 && n >= 5
            z = z*z + c
            n -= 5
        end
        (idx < length(a)) && (a[idx] = n)
        idx += 1
    end
    grayim(a)
#
#    img = grayim(a)
#    imwrite(img, "out.png")
end

function pngbytes(img)
    iob = IOBuffer()
    writemime(iob, MIME"image/png"(), img)
    takebuf_array(iob)
end

juliaset(c1::Float64, c2::Float64) = pngbytes(genjuliaset(c1, c2))
juliaset(s1::AbstractString, s2::AbstractString) = juliaset(parse(Float64, s1), parse(Float64, s2))
juliaset() = juliaset(rand(), rand())

const PNG_HEADER = @compat Dict{AbstractString,AbstractString}("Content-Type" => "image/png")
const JULIABOX_APIS = [
        (juliaset, false, PNG_HEADER)
    ]

serve_juliabox() = process(JULIABOX_APIS)
serve_local() = process(JULIABOX_APIS, "tcp://127.0.0.1:9999"; bind=true)

end # module


## function to perform imaging opetation
##

b3spline1d = [1/16 , 1/4 , 3/8 , 1/4 ,  1/16]

## Wavelet A Trous Algorithm (See e.g. Leon et al. 2001)
## im: image array
## levels : number of Wavelet planes
## kernel1d : 1D kernel to create the 2D kernel. Default is B3-spline but  could be other (e.g. [0,1,0])
##
function atrous(im, levels , kernel1d = b3spline1d, verbose = true)
    let
        w      = []
        ker    = Base.copy(kernel1d)
        approx = Base.copy(im)

        for lev in 1:levels
            if verbose
                println("##WT--A Trous--Plane: $lev")
            end

            kernel2d = zeros(length(ker),length(ker))
            for i in 1:(length(ker))
                kernel2d[i,:] = ker[i] .* ker
            end
            ker2 = ImageFiltering.centered(kernel2d)

            smooth = ImageFiltering.imfilter(approx, ker2)
            wi = approx .- smooth
            approx = smooth

            ker = upscale(ker)
            push!(w, wi)
        end

        ## LSP ###
        if verbose
            println("##WT--A Trous--LSP ($(levels+1))")
        end

        push!(w, approx)

        return(w)
    end
end

## upscale by a factor 2 the 1D kernel
function upscale(kerr)
    xx   = 0:length(kerr)-1
    itp  = Interpolations.interpolate(kerr, Interpolations.BSpline(Interpolations.Linear()))
    sitp = Interpolations.scale(itp, xx)
    xp   = 0:0.5:length(kerr)-1
    ker2 = 0.5 .* sitp(xp)

    return(ker2)
end

### Add the WT coefficients to get the (partial) image
##
function addWav(wt,p1,p2)
    let

        nplane = size(wt)
        if p2 == -1 p2 = nplane[1] end
        nxy = size(wt[1])
        println(nxy)
        im = zeros(nxy[1], nxy[2])

        for i in p1:p2
            im = im .+ wt[i]
        end
    return(im)
    end
end

## Hard thresholding on the WT coefficients.
## wt: wavelet transform
## d : random distribution
## σ : threshold in std from the simulated noise, if an array of length== nw[1] then used for thresholding
##
function thresholdingWav(wt, d::UnivariateDistribution , σ = 3 , kernel1d=b3spline1d)
    wtfilt = Base.copy(wt)
    nw = size(wt)
    nxy = size(wt[1])

    println("##WT--Thresholding the WT")
    if length(σ) == nw[1]
        threshold = σ
    else
        println("##WT--Simulating the noise in the WT.")
        threshold = zeros(nw[1])
        noise = rand(d , nxy[1] , nxy[2])
        wsimul = atrous(noise , nw[1]-1 , kernel1d , false)
        threshold = σ .* noiseWav(wsimul)
    end

    for p in 1:nw[1]
        for i in 1:nxy[1]
            for j in 1:nxy[2]
                if abs(wt[p][i,j]) < threshold[p]
                   wtfilt[p][i,j] = 0
                end
            end
        end
    end
    return(wtfilt)
end

## Compute std per plane for the wt
##
function noiseWav(wt)
    noise = []
    for plane in wt
        σ = std(plane)
        push!(noise,σ)
    end
    return(noise)
end


function dmetric!(dest, x, y, ::SqEuclidean, weights)
    for j in 1:size(x,2)
        w = weights[j]
        for i in 1:size(x,1)
            dest[i,j] = 2 * (x[i,j]-y[i,j]) * w
        end
    end
end


for T in (Float32, Float64)
    a1 = rand(Complex{T}, 10)
    b1 = fft(a1)
    af1 = AFArray(a1)
    bf1 = fft(af1)
    @test bf1 == fft1(af1)
    cf1 = ifft(bf1)
    @test cf1 == ifft1(bf1)
    @test isa(bf1, AFArray)
    @test isa(cf1, AFArray)
    @test eltype(bf1) == eltype(cf1) == Complex{T}
    @test Array(bf1) ≈ b1
    @test Array(cf1) ≈ a1

    fft!(af1)
    @test eltype(af1) == Complex{T}
    @test Array(af1) ≈ b1

    ifft!(af1)
    @test eltype(af1) == Complex{T}
    @test Array(af1) ≈ a1

    a2 = rand(Complex{T}, 10, 11)
    b2 = fft(a2)
    af2 = AFArray(a2)
    bf2 = fft(af2)
    @test bf2 == fft2(af2)
    cf2 = ifft(bf2)
    @test cf2 == ifft2(bf2)
    @test isa(bf2, AFArray)
    @test isa(cf2, AFArray)
    @test eltype(bf2) == eltype(cf2) == Complex{T}
    @test Array(bf2) ≈ b2
    @test Array(cf2) ≈ a2

    fft2!(af2)
    @test eltype(af2) == Complex{T}
    @test Array(af2) ≈ b2

    ifft2!(af2)
    @test eltype(af2) == Complex{T}
    @test Array(af2) ≈ a2

    a3 = rand(Complex{T}, 10, 11, 12)
    b3 = fft(a3)
    af3 = AFArray(a3)
    bf3 = fft(af3)
    @test bf3 == fft3(af3)
    cf3 = ifft(bf3)
    @test cf3 == ifft3(bf3)
    @test isa(bf3, AFArray)
    @test isa(cf3, AFArray)
    @test eltype(bf3) == eltype(cf3) == Complex{T}
    @test Array(bf3) ≈ b3
    @test Array(cf3) ≈ a3

    fft3!(af3)
    @test eltype(af3) == Complex{T}
    @test Array(af3) ≈ b3

    ifft3!(af3, 1 / length(af3))
    @test eltype(af3) == Complex{T}
    @test Array(af3) ≈ a3

    ar1 = rand(T, 10)
    br1 = rfft(ar1)
    arf1 = AFArray(ar1)
    brf1 = rfft(arf1)
    crf1 = irfft(brf1, 1 / length(arf1), isodd(size(arf1, 1)))
    @test isa(brf1, AFArray)
    @test isa(crf1, AFArray)
    @test eltype(brf1) == Complex{T}
    @test eltype(crf1) == T
    @test Array(brf1) ≈ br1
    @test Array(crf1) ≈ ar1

    ar2 = rand(T, 10, 11)
    arf2 = AFArray(ar2)
    brf2 = rfft2(arf2)
    crf2 = irfft2(brf2, 1 / length(arf2), isodd(size(arf2, 1)))
    @test isa(brf2, AFArray)
    @test isa(crf2, AFArray)
    @test isa(bf1, AFArray)
    @test isa(cf1, AFArray)
    @test eltype(brf2) == Complex{T}
    @test eltype(crf2) == T
    @test Array(brf2) ≈ rfft(ar2)
    @test Array(crf2) ≈ ar2

    ar3 = rand(T, 11, 12, 10)
    arf3 = AFArray(ar3)
    brf3 = rfft3(arf3)
    crf3 = irfft3(brf3, 1 / length(arf3), isodd(size(arf3, 1)))
    @test isa(brf3, AFArray)
    @test isa(crf3, AFArray)
    @test eltype(brf3) == Complex{T}
    @test eltype(crf3) == T
    @test Array(brf3) ≈ rfft(ar3)
    @test Array(crf3) ≈ ar3

    for sz in ((10,), (10, 10), (10, 10, 10))
        a = rand(Complex{T}, sz...)
        ad = AFArray(a)
        b = fft(a)
        bd = fft(ad)
        @test b ≈ Array(bd)
        @test fft!(a) ≈ Array(fft!(ad))
        @test ifft!(a) ≈ Array(ifft!(ad))

        a = rand(T, sz...)
        ad = AFArray(a)
        b = fft(a)
        bd = fft(ad)
        @test b ≈ Array(bd)
        @test rfft(a) ≈ Array(rfft(ad))
    end
end


import QuadGK: quadgk

### NIntegrate

@mkapprule NIntegrate

@sjdoc NIntegrate """
    NIntegrate( expr, [x,x0,xf])

Integrate the Symata expression `expr` numerically  between `x0` and `xf`.

    NIntegrate( expr , [x,x0,x1,...,xf])

Specify singular points between the limitx `x0` and `xf`.

    NIntegrate( f, [x0,xf])
    NIntegrate( f , [x0,x1,...,xf])

Integrate the compiled function `f` numerically.

NIntegrate returns `[result, err]`.

!!! note
    Giving an uncompiled expression is slower than giving a compiled function. But, for
    many purposes, the difference in efficiency is not important.
    `f = Compile(expr)` compiles `expr` to the function `f`.
"""

quadgklist(f,args) = mxprcf(:List, quadgk(f, map( x -> float(doeval(x)), args)...)...)

@doap function NIntegrate(expr, range::Mxpr{:List})
    r = margs(range)
    (sym,x0,x1) = length(r) == 3 ? (r...,) : length(r) == 2 ? (:x,r...,) : symerror(range, " is not a valid integration ranges")
    f = doeval(expr)
    isa(f,Function) && return quadgklist(f,(x0,x1))
    quadgklist(wrap_symata(expr,sym), (x0,x1))
end


### Functions to Compute Pose Graph Error ###
### Ted Steiner ###
### September 2014 ###

# Compute pose position RMS error
function posegraphError( posesA::Array{Pose2D,1}, posesB::Array{Pose2D,1} )
    alignPoses!(posesA,posesB)

    err = poseError(posesA,posesB)

    rmse = sqrt( sum(err.*err) / length(posesA) )
    return rmse
end

function poseError( posesA::Array{Pose2D,1}, posesB::Array{Pose2D,1} )
    err = zeros(length(posesA),3)

    for k = 1:length(posesA)
        err[k,:] = poseError(posesA[k],posesB[k])
    end

    return err
end

function poseError( poseA::Pose2D, poseB::Pose2D )
    err = zeros(3)
    err[1] = poseA.x - poseB.x
    err[2] = poseA.y - poseB.y
    err[3] = 0.0 #poseA.theta - poseB.theta
    return err
end

function posegraphError( truth_filename::String, result_filename::String )
    posesTruth = readiSAMOutput(truth_filename)
    posesResult = readiSAMOutput(result_filename)
    return posegraphError(posesTruth,posesResult)
end

function poseErrorCovTrace( cov::Array{Float64,2} )
    err = zeros(size(cov,1))

    for k = 1:length(err)
        err[k] = sqrt(cov[k,1]) + sqrt(cov[k,2]) + sqrt(cov[k,3])
    end

    return err
end

# Pose Sigma (Pose position error standard deviation)
function poseSigma( cov::Array{Float64,2} )
    return sqrt( cov[1,1] + cov[1,2] )
end

function poseSigmas( covs::Array{Float64,2} )
    err = zeros(size(covs,1))

    for k = 1:length(err)
        err[k] = poseSigma( covs[k,:] )
    end

    return err
end

function epsilonTrajectory( covs::Array{Float64,2} )
    return mean(poseSigmas(covs))
end

# Trajectory Error Epsilon
function errorRatio( mean_sigmas::Array, rounds::Int=Inf )
    if rounds > size(mean_sigmas,1)
        rounds = size(mean_sigmas,1)
    end

    return mean(errorRatios(mean_sigmas,rounds),1)
end

function errorRatios( mean_sigmas::Array, rounds::Int=Inf )
    if rounds > size(mean_sigmas,1)
        rounds = size(mean_sigmas,1)
    end

    anchors = size(mean_sigmas,2)-1
    e_ratios = zeros(rounds,anchors)

    for r = 1:rounds
        e_full = mean_sigmas[r,end]

        for a = 1:anchors
            e_ratios[r,a] = (mean_sigmas[r,a] - e_full) / e_full
        end
    end

    return e_ratios
end

######################
### Pose Alignment ###
######################

function alignPoses!(posesA::Array{Pose2D,1}, posesB::Array{Pose2D,1})
    R, t = kabsch(posesA,posesB)
    transformPoses!(posesB,R',-t)
    return nothing
end

# Kabsch Algorithm: Computes rotation and translation between two sets of points
function kabsch( posesA::Array{Pose2D,1}, posesB::Array{Pose2D,1} )
    A = poses2matrix(posesA)
    B = poses2matrix(posesB)

    return kabsch(A,B)
end

function kabsch( A::Array{Float64,2}, B::Array{Float64,2} )
    assert(size(A,1) == size(B,1))

    centroid_A = mean(A,1)
    centroid_B = mean(B,1)

    H = (A - repmat(centroid_A, size(A,1), 1))' * (B - repmat(centroid_B, size(B,1), 1))
    U,S,V = svd(H)
    R = V*U';

    if det(R) < 0
        println("Reflection detected.")
        V[:,3] *= -1
        R = V*U'
    end

    t = -R*centroid_A' + centroid_B'

    return R, t
end

function transformPoses!( poses::Array{Pose2D,1}, R, t )
    R_theta = atan2(R[2,1],R[1,1])

    for pose in poses
        pose.x += t[1]
        pose.y += t[2]
        pose.theta += R_theta
    end

    return nothing
end

function poses2matrix( poses::Array{Pose2D,1} )
    A = zeros(length(poses),3)

    for k = 1:length(poses)
        A[k,1] = poses[k].x
        A[k,2] = poses[k].y
    end

    return A
end


# # [Converting Models to Computation Graphs](@id bilayernetwork_example)
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/examples/covid/bilayerconversion.ipynb)

# using Pkg
# Pkg.status()
# Pkg.activate(".")

# Pkg.instantiate()

# Pkg.develop(path="../../")
# Pkg.instantiate()

using AlgebraicPetri
using AlgebraicPetri.Epidemiology
using AlgebraicPetri.BilayerNetworks

using Catlab
using Catlab.CategoricalAlgebra
using Catlab.Graphics
import Catlab.CategoricalAlgebra: migrate!
using Catlab.WiringDiagrams
using Catlab.Programs.RelationalPrograms
display_uwd(ex) = to_graphviz(ex, box_labels=:name, junction_labels=:variable, edge_attrs=Dict(:len=>".75"));

using PrettyTables
function printsoln(bn::AbstractLabelledBilayerNetwork, soln::Vector)
    pretty_table(soln)
end

# #### SIR Model:

# define model
sir = @relation (s,i,r) begin
    infection(s,i)
    recovery(i,r)
end
display_uwd(sir)
#-
#

# Extract the Petri network form of the model.
psir = apex(oapply_epi(sir))
psir
Graph(psir)
#-

# Convert the Petri network into a bilayer network and draw it.
# This model uses a computation graph to express the computation of the vector field of the Petri net.
bnsir = BilayerNetwork()
migrate!(bnsir, psir)
bnsir
to_graphviz(bnsir)
#-

# We can hand code a Bilayer netowork using the `@acset` macro provided by Catlab. As you can see from the code,
# there is a lot of typing to specify the incidence of all these wires. The Petri Net is more compact.
# This notion of Bilayer network comes from the definition of mass action kinetics for reaction networks.

# hand coded Bilayer network
bnsir_test = @acset BilayerNetwork begin
    Qin = 3
    Qout = 3
    Win = 3
    Wa = 3
    Wn = 3
    Box = 2
    arg = [1,2,2]
    call = [1,1,2]
    efflux = [1,1,2]
    effusion = [1,2,2]
    influx = [1,1,2]
    infusion = [2,2,3]
end

@assert bnsir == bnsir_test

function roundtrip(pn::AbstractPetriNet, bn::AbstractBilayerNetwork)
    roundtrippetri = PetriNet()
    migrate!(roundtrippetri, bn)
    pn_structure = PetriNet()
    copy_parts!(pn_structure, pn)
    return roundtrippetri, pn_structure
end
#-

# We can specify the SEIR model with the relation macro and functorial semantics as usual.

seir = @relation (s,e,i,r) begin
    exposure(s,i,e)
    illness(e,i)
    recovery(i,r)
end

pseir = apex(oapply_epi(seir))

# By converting this to a Bilayer network, we are able to visualize differences in the computational pattern
# of data flow between different reaction network models.

#
bnseir = BilayerNetwork()
migrate!(bnseir, pseir)

bnrt,pnstr = roundtrip(pseir, bnseir)

display_uwd(sir)

Graph(psir)
#-

to_graphviz(bnsir)
#-

to_graphviz(bnseir)
#-

Graph(bnrt)
#-

qm = @relation (s,q) begin
    exposure(s,i,e)
    illness(e,i)
    recovery(i,r)
    quarexp(e,q)
    quarinf(i,q)
end
display_uwd(qm)
#-

import AlgebraicPetri.Epidemiology: exposure_petri, spontaneous_petri
semantics = Dict(
    :infection => exposure_petri(:S, :I, :I, :inf),
    :exposure  => exposure_petri(:S, :I, :E, :exp),
    :illness   => spontaneous_petri(:E,:I,:ill),
    :recovery  => spontaneous_petri(:I,:R,:rec),
    :death     => spontaneous_petri(:I,:D,:death),
    :quarexp   => spontaneous_petri(:E, :Q, :qe),
    :quarinf   => spontaneous_petri(:I, :Q, :qi),
    :quarrec   => spontaneous_petri(:Q, :R, :qr)
)
pn_quar = oapply(qm, semantics)  |> apex
Graph(pn_quar)
#-

bnquar = LabelledBilayerNetwork()
migrate!(bnquar, pn_quar)
to_graphviz(bnquar)

qm = @relation (s,i,q) begin
    exposure(s,i,e)
    illness(e,i)
    recovery(i,r)
    quarexp(e,q)
    quarinf(i,q)
    quarrec(q,r)
end
display_uwd(qm)
#-

pn_quar = oapply(qm, semantics)  |> apex
Graph(pn_quar)
#-

bnquar = LabelledBilayerNetwork()
migrate!(bnquar, pn_quar)
to_graphviz(bnquar)
#-

quarrt = LabelledPetriNet()
migrate!(quarrt, bnquar) |> Graph

bnquar
#-

import AlgebraicPetri.BilayerNetworks: evaluate!, evaluate

bnsir = LabelledBilayerNetwork()
migrate!(bnsir, psir)

evaluate(bnsir, [10.0, 1, 0.0], inf=0.1, rec=0.3)

evaluate(bnquar, [10.0, 1, 0, 0,0,0], exp=0.1, rec=0.3, qi=0.2, ill=0.7, qe=0.23, qr=0.3)

function euler(bn::AbstractLabelledBilayerNetwork, state, nsteps::Integer, stepsize::Real; params...)
    #preallocate storage so that each step is nonallocating
    #create a storage space for steps of euler's method store intermediate
    #states as named tuples so that you can integrate with julia Tables.jl data  analysis ecosystem
    du = zeros(length(state))
    ϕ = ones(nparts(bn, :Box))
    u = tuple(state...)
    results = Vector{NamedTuple{tuple(bn[:,:variable]...)}}()
    for i in 1:nsteps
        u = u .+ stepsize.*evaluate!(du, ϕ, bn, u; params...)
        push!(results, NamedTuple{tuple(bn[:,:variable]...)}(u))
    end
    return results
end

soln = euler(bnquar, (S=10.0, I=1, E=0, R=0, Q=0), 30, 0.15, exp=0.1, rec=0.03, qi=0.37, ill=0.7, qe=0.23, qr=0.03)
printsoln(bnquar, soln)
#-

import AlgebraicPetri.BilayerNetworks: compile
compile(bnquar, :du, :ϕ, :u, :p)
#-
#
compile(bnquar, :du, :ϕ, :u, exp=0.1, rec=0.03, qi=0.37, ill=0.7, qe=0.23, qr=0.03)
#-

function eulers(bn::AbstractLabelledBilayerNetwork, funcname::Symbol; params...)
    f = compile(bn, :du, :ϕ, :u; params...)
    varnames = tuple(bn[:,:variable]...)
    nϕ = nparts(bn, :Box)
    quote
    function $funcname(state, nsteps::Integer, stepsize::Real)
        $f
        #preallocate storage so that each step is nonallocating
        du = zeros(length(state))
        ϕ = ones($nϕ)
        u = tuple(state...)
        #create a storage space for steps of euler's method
        #store intermediate states as named tuples so that you can integrate
        #with julia Tables.jl data analysis ecosystem
        results = Vector{NamedTuple{$varnames}}()
        for i in 1:nsteps
            Δ = f!(du, ϕ, u, 0)
            u = u .+ stepsize.*Δ
            push!(results, NamedTuple{$varnames}(u))
        end
        return results
    end
    end
end

eulers(bnsir, :eulsir, inf=0.3, rec=0.2)
#-

eulseirqexp = eulers(bnquar, :eulseirq, exp=0.1, rec=0.03, qi=0.37, ill=0.7, qe=0.23, qr=0.03)
#-

eulsirexp = eulers(bnsir, :eulsir, inf=0.3, rec=0.2)
#-

eval(eulsirexp)
soln_codegen = eulsir((S=10.0, I=1, R=0), 30, 0.15)
pretty_table(soln_codegen)
#-

eval(eulseirqexp)
soln_codegen = eulseirq((S=10.0, I=1, E=0, R=0, Q=0), 30, 0.15)
pretty_table(soln_codegen)
#-


@testset "logpvalue(A & B)" begin
    using OptEnrichedSetCover: logpvalue

    @test_throws ArgumentError logpvalue(1, -1, 0, 1)
    @test_throws ArgumentError logpvalue(1, 1, -1, 1, :right)
    @test_throws ArgumentError logpvalue(1, 1, 1, -1)

    @test_throws ArgumentError logpvalue(1, 2, 0, 1)
    @test_throws ArgumentError logpvalue(1, 1, 3, 2, :right)
    @test_throws ArgumentError logpvalue(1, 1, 1, 0)

    @test logpvalue(1, 1, 0, 1, :right) == -Inf
    @test logpvalue(1, 1, 0, 1) == -Inf
    @test logpvalue(1, 1, 0, 1, :left) == 0.0
    @test logpvalue(1, 1, 1, 1, :left) == 0.0
    @test logpvalue(1, 1, 0, 1, :both) == -Inf
    @test logpvalue(1, 1, 1, 1, :both) == 0.0
    @test logpvalue(1, 10, 5, 12, :left) == -Inf
    @test logpvalue(1, 10, 5, 12, :right) == 0.0
    @test logpvalue(1, 10, 5, 12, :both) == -Inf
    @test_throws ArgumentError logpvalue(1, 1, 0, 1, :center)

    @test logpvalue(1, 1, 2, 3, :right) ≈ log(2/3)
    @test logpvalue(1, 1, 2, 3, :left) == 0.0
    @test logpvalue(1, 2, 2, 4, :left) ≈ log(5/6)
    @test logpvalue(1, 2, 2, 4, :right) ≈ log(5/6)

    # no overlap
    @test logpvalue(-1, 10, 10, 100, :left) == -Inf
    @test logpvalue( 0, 10, 10, 100, :right) == 0.0

    # full overlap
    @test logpvalue(10, 10, 10, 100, :left) == 0.0
    @test logpvalue(10, 10, 10, 100, :right) ≈ -log(binomial(100, 10))
    @test logpvalue(11, 10, 10, 100, :right) == -Inf

    # https://bugs.r-project.org/bugzilla3/show_bug.cgi?id=17271
    @test logpvalue(4, 22837, 4, 22837, :left) == 0.0
    @test logpvalue(4, 22837, 4, 22837, :right) == 0.0
    @test logpvalue(22837-4, 22837, 22837-4, 22837, :left) == 0.0
    @test logpvalue(22837-4, 22837, 22837-4, 22837, :right) == 0.0

    @test logpvalue(42, 43, 46, 47, :right) == 0.0
    @test logpvalue(41, 43, 46, 47, :left) == -Inf

    @test logpvalue(5, 10, 10, 20, :left) ≈ log(0.6718591)
    @test logpvalue(5, 10, 10, 20, :right) ≈ log(0.6718591)
    @test logpvalue(5, 10, 10, 20, :both) ≈ 0.0

    @test logpvalue(10, 15, 15, 30, :right) ≈ log(0.071555489)
    @test logpvalue(10, 15, 15, 30, :left) ≈ log(0.9865811354)
    @test logpvalue(10, 15, 15, 30, :both) ≈ log(2*0.071555489)

    @test logpvalue(5, 15, 15, 30, :right) ≈ log(0.9865811354)
    @test logpvalue(5, 15, 15, 30, :left) ≈ log(0.071555489)
end


using Revise, InfiniteLinearAlgebra, BlockBandedMatrices, BandedMatrices, BlockArrays, InfiniteArrays, FillArrays, LazyArrays, Test

A = BlockTridiagonal(Vcat([fill(1.0,2,1),Matrix(1.0I,2,2),Matrix(1.0I,2,2),Matrix(1.0I,2,2)],Fill(Matrix(1.0I,2,2), ∞)),
                       Vcat([zeros(1,1)], Fill(zeros(2,2), ∞)),
                       Vcat([fill(1.0,1,2),Matrix(1.0I,2,2)], Fill(Matrix(1.0I,2,2), ∞)))

BlockBandedMatrix(A)
isb


function angle_of_incidence(lat,lon,jd,time)
# angle_of_incidence calculates the cosine of the sun incident angle
#  y = ANGLE_OF_INCIDENCE(lat,lon,jd,time) calculates the cosine of the sun incident angle for
#   a given location (lat,lon), julian day (jd) and time in hours and fraction of hours.

# inclination angle
delta=(asin(-sin(23.45/180*pi)*cos(360/365.25*(jd+10)/180*pi)));
# lat to radians
lamda=lat/180*pi;
# time to radians
omega=time/24*2*pi+lon/360*2*pi+pi;

y=cos(delta)*cos(lamda)*cos(omega)+sin(delta)*sin(lamda);

return y

end


@defcomp radiativeforcing begin
    FORC      = Variable(index=[time])   # Increase in radiative forcing (watts per m2 from 1900)

    forcoth   = Parameter(index=[time])  # Exogenous forcing for other greenhouse gases
    MAT       = Parameter(index=[time])  # Carbon concentration increase in atmosphere (GtC from 1750)
    MAT_final = Parameter()              # MAT calculation one timestep further than the model's index
    fco22x    = Parameter()              # Forcings of equilibrium CO2 doubling (Wm-2)

    function run_timestep(p, v, d, t)
        # Define function for FORC
        if ! is_last(t)
            v.FORC[t] = p.fco22x * (log((((p.MAT[t] + p.MAT[t + 1]) / 2) + 0.000001) / 596.4) / log(2)) + p.forcoth[t]
        else # final timestep
            # need to use MAT_final, calculated one step further
            v.FORC[t] = p.fco22x * (log((((p.MAT[t] + p.MAT_final) / 2) + 0.000001) / 596.4) / log(2)) + p.forcoth[t]
        end
    end
end


# Try to select a reasonable set of sigma values that should be used
# for the adaptive shrinkage model based on the values of x (the noisy
# observations) and s (the standard error in the observations). The
# return value is a vector of sigma values.
function autoselectmixsd(x::Array{Float64,1},
                         s::Array{Float64,1} = ones(size(x));
                         gridmult::Float64 = 1.4,
                         nv::Int = 0)

  # Get the number of samples.
  n = length(x);

  # Check input "s"---it should be the same length.
  if length(s) != n
    throw(ArgumentError("Arguments \"x\" and \"s\" should have the same" *
                        "length"))
  elseif any(s .<= 0)
    throw(ArgumentError("All elements of \"s\" should be positive"))
  end

  # Check input "gridmult".
  if gridmult < 0
    throw(ArgumentError("Input \"gridmult\" should be non-negative"))
  end

  # Check input "nv".
  if !(nv == 0 || nv > 1)
    throw(ArgumentError("Input \"nv\" should be 0, or greater than 1"))
  end

  # Determine the minimum and maximum sigma settings.
  smin = sqrt(minimum(s.^2))/10;
  if all(x.^2 .<= s.^2)

    # This deals with the occasional boundary case.
    smax = 10 * smin;
  else

    # This is, roughly, the largest value you'd want to use.
    smax = 2 * sqrt(maximum(x.^2 - s.^2));
  end

  # Choose the grid of sigmas.
  if nv > 0
    return vcat(0,logspace(log10(smin),log10(smax),nv - 1));
  elseif gridmult == 0
    return vcat(0,smax)
  else
    m = ceil(Int,log2(smax/smin)/log2(gridmult));
    return vcat(0,logspace(log10(smin),log10(smax),m + 1));
  end
end

# Compute the n x k conditional likelihood matrix, where n is the
# number of samples and k is the number of mixture components, for the
# case when the likelihood is univariate normal and prior is a mixture
# of univariate normals.
#
# Entry (i,j) of the conditional likelihood matrix is equal to
# N(0,s[i]^2 + sd[j]^2), the normal density with zero mean and
# variance s[i]^2 + sd[j]^2.
#
# If normalize = true, each row of the likelihood matrix is divided by
# the largest value in the row. After normalization, the largest value
# in each row is 1.
function normlikmatrix(x::Array{Float64,1},
                       s::Array{Float64,1} = ones(size(x));
                       sd::Array{Float64,1} = autoselectmixsd(x,s),
                       normalizerows::Bool = true)

  # Get the number of samples (n) and the number of prior mixture
  # components (k).
  n = length(x);
  k = length(sd);

  # Check input "s"---it should be the same length at x.
  if length(s) != n
    throw(ArgumentError("Arguments \"x\" and \"s\" should have the same" *
                        "length"))
  elseif any(s .<= 0)
    throw(ArgumentError("All elements of \"s\" should be positive"))
  end

  # Check input "sd".
  if any(sd .< 0)
    throw(ArgumentError("All elements of \"sd\" should be non-negative"))
  end

  # Compute the n x k matrix of standard deviations.
  S = sqrt.((s.^2) .+ (sd.^2)');

  # Compute the log-densities, and normalize the rows, if requested.
  L = -(x./S).^2/2 - log.(S) .- log(2*pi)/2;
  if normalizerows

    # This is the same as
    #
    #   L = L - repmat(maximum(L,2),1,k);
    #
    # but uses memory more efficiently to complete the operation.
    L = broadcast(-,L,maximum(L,dims = 2));
  end
  return exp.(L)
end

# Implements the "logspace" function that was available in an
# earlier version of Julia.
function logspace(a, b, n)
  return exp10.(range(a,stop = b,length = n))
end


using ITensors,
      Test

@testset "ITensor broadcast syntax" begin

  i = Index(2,"i")
  A = randomITensor(i,i')
  B = randomITensor(i',i)
  α = 2
  β = 3

  @testset "Copy" begin
    Bc = copy(B)
    Bc .= A
    @test Bc[1,1] == A[1,1]
    @test Bc[2,1] == A[1,2]
    @test Bc[1,2] == A[2,1]
    @test Bc[2,2] == A[2,2]
  end

  @testset "Fill" begin
    Bc = copy(B)
    Bc .= α
    @test Bc[1,1] == α
    @test Bc[2,1] == α
    @test Bc[1,2] == α
    @test Bc[2,2] == α
  end

  @testset "Scaling" begin
    Bc = copy(B)
    Bc .*= α
    @test Bc[1,1] == α * B[1,1]
    @test Bc[2,1] == α * B[2,1]
    @test Bc[1,2] == α * B[1,2]
    @test Bc[2,2] == α * B[2,2]
  end

  @testset "Dividing" begin
    Bc = copy(B)
    Bc ./= α
    @test Bc[1,1] == B[1,1] / α
    @test Bc[2,1] == B[2,1] / α
    @test Bc[1,2] == B[1,2] / α
    @test Bc[2,2] == B[2,2] / α
  end

  @testset "Scalar multiplication (in-place)" begin
    Bc = copy(B)
    Bc .= α .* A
    @test Bc[1,1] == α * A[1,1]
    @test Bc[2,1] == α * A[1,2]
    @test Bc[1,2] == α * A[2,1]
    @test Bc[2,2] == α * A[2,2]
  end

  @testset "Dividing (in-place)" begin
    Bc = copy(B)
    Bc .= A ./ α
    @test Bc[1,1] == A[1,1] / α
    @test Bc[2,1] == A[1,2] / α
    @test Bc[1,2] == A[2,1] / α
    @test Bc[2,2] == A[2,2] / α
    @test_throws ErrorException Bc .= α ./ A
  end

  @testset "Scalar multiplication (out-of-place)" begin
    Bc = α .* A
    @test Bc[1,1] == α * A[1,1]
    @test Bc[2,1] == α * A[2,1]
    @test Bc[1,2] == α * A[1,2]
    @test Bc[2,2] == α * A[2,2]
  end

  @testset "Addition" begin
    Bc = copy(B)
    Bc .= A .+ Bc
    @test Bc[1,1] == A[1,1] + B[1,1]
    @test Bc[2,1] == A[1,2] + B[2,1]
    @test Bc[1,2] == A[2,1] + B[1,2]
    @test Bc[2,2] == A[2,2] + B[2,2]
  end

  @testset "Addition (with α)" begin
    Bc = copy(B)
    Bc .+= A .* α

    @test Bc[1,1] == α * A[1,1] + B[1,1]
    @test Bc[2,1] == α * A[1,2] + B[2,1]
    @test Bc[1,2] == α * A[2,1] + B[1,2]
    @test Bc[2,2] == α * A[2,2] + B[2,2]
  end

  @testset "Addition (with α and β)" begin
    Bc = copy(B)
    Bc .= α .* A .+ β .* Bc

    @test Bc[1,1] == α * A[1,1] + β * B[1,1]
    @test Bc[2,1] == α * A[1,2] + β * B[2,1]
    @test Bc[1,2] == α * A[2,1] + β * B[1,2]
    @test Bc[2,2] == α * A[2,2] + β * B[2,2]
  end

  @testset "Addition errors" begin
    C = randomITensor(i,i')
    @test_throws ErrorException C .= A .+ B
    @test_throws ErrorException C = A .+ B
    @test_throws ErrorException C .= A .* B
  end

  @testset "Contraction" begin
    ii = Index(2; tags="ii")
    jj = Index(2; tags="jj")
    kk = Index(2; tags="kk")

    AA = randomITensor(ii, jj)
    BB = randomITensor(kk, jj)
    CC = randomITensor(kk, ii)

    R = copy(CC)
    R .= AA .* BB
    @test AA*BB ≈ R

    R = copy(CC)
    R .= α .* AA .* BB
    @test α*AA*BB ≈ R

    R = copy(CC)
    R .= AA .* α .* BB
    @test α*AA*BB ≈ R

    R = copy(CC)
    R .= AA .* BB .* α
    @test α*AA*BB ≈ R

    R = copy(CC)
    R .+= α .* AA .* BB
    @test α*AA*BB+CC ≈ R

    R = copy(CC)
    R .= β .* R .+ AA .* BB .* α
    @test α*AA*BB+β*CC ≈ R
  end

  @testset "General functions" begin
    absA = abs.(A)

    @test absA[1,1] == abs(A[1,1])
    @test absA[2,1] == abs(A[2,1])

    Bc = copy(B)
    Bc .= sqrt.(absA)

    @test Bc[1,1] == sqrt(absA[1,1])
    @test Bc[2,1] == sqrt(absA[1,2])

    Bc2 = copy(B)
    Bc2 .+= sqrt.(absA)

    @test Bc2[1,1] == B[1,1]+sqrt(absA[1,1])
    @test Bc2[2,1] == B[2,1]+sqrt(absA[1,2])
  end

  @testset "Some other operations" begin
    i = Index(2)
    A = randomITensor(i)
    B = randomITensor(i)

    absA = abs.(A)

    @test absA[1] == abs(A[1])
    @test absA[2] == abs(A[2])

    Bc = copy(B)
    Bc .= sqrt.(absA)

    @test Bc[1] == sqrt(absA[1])
    @test Bc[2] == sqrt(absA[2])

    Bc2 = copy(B)
    Bc2 .+= sqrt.(absA)

    @test Bc2[1] == B[1]+sqrt(absA[1])
    @test Bc2[2] == B[2]+sqrt(absA[2])

    Bc3 = copy(B)
    Bc3 .= sqrt.(absA) .+ sin.(Bc3)

    @test Bc3[1] == sin(B[1])+sqrt(absA[1])
    @test Bc3[2] == sin(B[2])+sqrt(absA[2])

    sqrtabsA = sqrt.(abs.(A))

    @test sqrtabsA[1] == sqrt(abs(A[1]))
    @test sqrtabsA[2] == sqrt(abs(A[2]))

    sqrtabsA = cos.(sin.(sqrt.(abs.(A))))

    @test sqrtabsA[1] == cos(sin(sqrt(abs(A[1]))))
    @test sqrtabsA[2] == cos(sin(sqrt(abs(A[2]))))

    Ap = A .+ 3

    @test Ap[1] == A[1] + 3
    @test Ap[2] == A[2] + 3

    Apow1 = A .^ 2.0

    @test Apow1[1] == A[1]^2
    @test Apow1[2] == A[2]^2

    Apow2 = A .^ 3

    @test Apow2[1] == A[1]^3
    @test Apow2[2] == A[2]^3
  end

end


using Base, OrthogonalPolynomialsQuasi, ContinuumArrays, QuasiArrays, FillArrays,
        LazyArrays, BandedMatrices, LinearAlgebra, FastTransforms, IntervalSets,
        InfiniteLinearAlgebra, Test
using ForwardDiff, SemiseparableMatrices, SpecialFunctions, LazyBandedMatrices
import ContinuumArrays: BasisLayout, MappedBasisLayout
import OrthogonalPolynomialsQuasi: jacobimatrix, ∞, ChebyshevInterval, LegendreWeight,
            Clenshaw, forwardrecurrence!, singularities
import LazyArrays: ApplyStyle, colsupport, MemoryLayout, arguments
import SemiseparableMatrices: VcatAlmostBandedLayout
import QuasiArrays: MulQuasiMatrix
import Base: OneTo
import InfiniteLinearAlgebra: KronTrav, Block
import FastTransforms: clenshaw!

@testset "singularities" begin
    x = Inclusion(ChebyshevInterval())
    @test singularities(x) == singularities(exp.(x)) == singularities(x.^2) ==
        singularities(x .+ 1) == singularities(1 .+ x) == singularities(x .+ x) ==
        LegendreWeight()
    @test singularities(exp.(x) .* JacobiWeight(0.1,0.2)) ==
        singularities(JacobiWeight(0.1,0.2) .* exp.(x)) ==
        JacobiWeight(0.1,0.2)

    x = Inclusion(0..1)
    @test singularities(x) == singularities(exp.(x)) == singularities(x.^2) ==
        singularities(x .+ 1) == singularities(1 .+ x) == singularities(x .+ x) ==
        legendreweight(0..1)
end

include("test_chebyshev.jl")
include("test_legendre.jl")
include("test_ultraspherical.jl")
include("test_jacobi.jl")
include("test_fourier.jl")
include("test_odes.jl")
include("test_normalized.jl")
include("test_lanczos.jl")

@testset "Auto-diff" begin
    U = Ultraspherical(1)
    C = Ultraspherical(2)

    f = x -> ChebyshevT{eltype(x)}()[x,5]
    @test ForwardDiff.derivative(f,0.1) ≈ 4*U[0.1,4]
    f = x -> ChebyshevT{eltype(x)}()[x,5][1]
    @test ForwardDiff.gradient(f,[0.1]) ≈ [4*U[0.1,4]]
    @test ForwardDiff.hessian(f,[0.1]) ≈ [8*C[0.1,3]]

    f = x -> ChebyshevT{eltype(x)}()[x,1:5]
    @test ForwardDiff.derivative(f,0.1) ≈ [0;(1:4).*U[0.1,1:4]]
end


# ---
# title: 1277. Count Square Submatrices with All Ones
# id: problem1277
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array, Dynamic Programming
# link: <https://leetcode.com/problems/count-square-submatrices-with-all-ones/description/>
# hidden: true
# ---
#
# Given a `m * n` matrix of ones and zeros, return how many **square**
# submatrices have all ones.
#
#
#
# **Example 1:**
#
#
#
#     Input: matrix =
#     [
#       [0,1,1,1],
#       [1,1,1,1],
#       [0,1,1,1]
#     ]
#     Output: 15
#     Explanation:
#     There are **10** squares of side 1.
#     There are **4** squares of side 2.
#     There is  **1** square of side 3.
#     Total number of squares = 10 + 4 + 1 = **15**.
#
#
# **Example 2:**
#
#
#
#     Input: matrix =
#     [
#       [1,0,1],
#       [1,1,0],
#       [1,1,0]
#     ]
#     Output: 7
#     Explanation:
#     There are **6** squares of side 1.
#     There is **1** square of side 2.
#     Total number of squares = 6 + 1 = **7**.
#
#
#
#
# **Constraints:**
#
#   * `1 <= arr.length <= 300`
#   * `1 <= arr[0].length <= 300`
#   * `0 <= arr[i][j] <= 1`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


"TruthTables exports the macro @truth_table"
module TruthTables

## For pretty printing
using NamedArrays: NamedArray

export @truth_table

¬(p::Bool) = !p
∧(p::Bool, q::Bool) = p && q
∨(p::Bool, q::Bool) = p || q
⟹(p::Bool, q::Bool) = if p; q else true end
⟺(p::Bool, q::Bool) = p === q

"""
    @truth_table predicate

Return the truth table of `predicate`.

The `predicate` can contain built-in boolean operators, as well as user-defined boolean-valued functions.
The truth table is a NamedArray whose columns are variables and sub-expressions of the input predicate.
"""
macro truth_table(expr)
    sub_exprs, vars = get_sub_exprs(expr)
    col_names = string.([vars..., sub_exprs...])
    c = gensym(:c)
    vals = gensym(:vals)

    ## escape user-defined functions
    sub_exprs = escape_calls.(sub_exprs)

    ## Form expression
    ex = :($vals[$c+=1, :] = [$(vars...), $(sub_exprs...)])
    for i in reverse(vars)
        ex = :(for $i in [true, false]; $ex end)
    end

    quote
        $c = 0
        $vals = Matrix{Bool}(undef, 2^length($vars), length($vars) + length($sub_exprs))
        $ex
        NamedArray($vals, (1:2^length($vars), $col_names), ("", ""))
    end
end

function get_sub_exprs(expr)
    sub_exprs = Expr[]
    vars = Symbol[]

    ## Traverse expression using multiple dispatch
    function walk!(ex::Expr)
        n = ex.head == :call ? 2 : 1    # Ignore function names
        walk!.(ex.args[n:end])          # Mind the dot
        ex ∉ sub_exprs && push!(sub_exprs, ex)
    end
    walk!(var::Symbol) = var ∉ vars && push!(vars, var)
    walk!(::Bool) = nothing

    walk!(expr)
    return sub_exprs, vars
end

escape_calls(s) = s
function escape_calls(expr::Expr)
    ex = expr
    if ex.head == :call && ex.args[1] isa Symbol && !isdefined(@__MODULE__, ex.args[1])
        ex.args[1] = esc(ex.args[1])
    end
    ex.args[2:end] = escape_calls.(ex.args[2:end])
    ex
end

end #module


### A Pluto.jl notebook ###
# v0.14.5

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

# ╔═╡ 623b2240-c520-11eb-0b1e-4b03cb67ac24
begin
	using PlutoUI,Plots, ForwardDiff
	plotlyjs()
end

# ╔═╡ c4ec78ff-b9c7-4867-81df-9384b965881e
md"## Polinomio de Taylor"

# ╔═╡ bc296b17-b6ca-476d-b9f4-1743559f0a6d
f(x) = sin(x)

# p₀(x=0) = f(x=0) = cos(0) = 1
# f'(x) = -sen(x) // f'(x=0) = -sen(0) = 0

# ╔═╡ 916c7595-d8ee-4ce1-9083-f403613131f7
md"""
a₀ = $(@bind a₀ Slider(-2:0.1:2,default=1,show_value=true))

a₁ = $(@bind a₁ Slider(-2:0.1:2,default=0,show_value=true))

a₂ = $(@bind a₂ Slider(-2:0.1:2,default=-0.5,show_value=true))

a₃ = $(@bind a₃ Slider(-2:0.1:2,default=0,show_value=true))

a₄ = $(@bind a₄ Slider(-2:0.1:2,default=0,show_value=true))
"""

# ╔═╡ 0831c272-bd4a-459a-9d8b-75bc4f926022
md"""
p₀(x) = $(a₀)

p₁(x) = $(a₀) + $(a₁) x

p₂(x) = $(a₀) + $(a₁) x + $(a₂) x²

p₃(x) = $(a₀) + $(a₁) x + $(a₂) x² + $(a₃) x³

p₄(x) = $(a₀) + $(a₁) x + $(a₂) x² + $(a₃) x³ + + $(a₄) x⁴
"""

# ╔═╡ 21815a23-2d33-4813-baca-21adcfa4a85f
begin
	p₀(x) = a₀
	p₁(x) = a₀ + a₁ * x
	p₂(x) = a₀ + a₁ * x + a₂ * x^2
	p₃(x) = a₀ + a₁ * x + a₂ * x^2 +  a₃ * x^3
	p₄(x) = a₀ + a₁ * x + a₂ * x^2 +  a₃ * x^3 + a₄ * x^4
	# p₂(x) = a₀ + a₁ * x + a₂ * x^2
end

# ╔═╡ 9f161f71-89db-44d3-84f3-be08176ff802
begin
	plot(f,xlim=[-2,2],label=false,width=2)
	plot!(p₀,label="p₀",width=2)
	plot!(p₁,label="p₁",width=2)
	plot!(p₂,label="p₂",width=2)
	plot!(p₃,label="p₃",width=2)
	plot!(p₄,label="p₄",width=2)
end

# ╔═╡ Cell order:
# ╟─623b2240-c520-11eb-0b1e-4b03cb67ac24
# ╟─c4ec78ff-b9c7-4867-81df-9384b965881e
# ╠═bc296b17-b6ca-476d-b9f4-1743559f0a6d
# ╟─916c7595-d8ee-4ce1-9083-f403613131f7
# ╟─0831c272-bd4a-459a-9d8b-75bc4f926022
# ╟─21815a23-2d33-4813-baca-21adcfa4a85f
# ╟─9f161f71-89db-44d3-84f3-be08176ff802


module Histogram

# simple histogram functions copied from Julia v0.4.7 to avoid having to
# pull in all of Statistics.jl

export hist, hist!

function histrange(v::AbstractArray{T}, n::Integer) where {T<:AbstractFloat}
    nv = length(v)
    if nv == 0 && n < 0
        throw(ArgumentError("number of bins must be ≥ 0 for an empty array, got $n"))
    elseif nv > 0 && n < 1
        throw(ArgumentError("number of bins must be ≥ 1 for a non-empty array, got $n"))
    end
    if nv == 0
        return 0.0:1.0:0.0
    end
    lo, hi = extrema(v)
    if hi == lo
        step = 1.0
    else
        bw = (hi - lo) / n
        e = 10.0^floor(log10(bw))
        r = bw / e
        if r <= 2
            step = 2*e
        elseif r <= 5
            step = 5*e
        else
            step = 10*e
        end
    end
    start = step*(ceil(lo/step)-1)
    nm1 = ceil(Int,(hi - start)/step)
    return start:step:(start + nm1*step)
end

function histrange(v::AbstractArray{T}, n::Integer) where {T<:Integer}
    nv = length(v)
    if nv == 0 && n < 0
        throw(ArgumentError("number of bins must be ≥ 0 for an empty array, got $n"))
    elseif nv > 0 && n < 1
        throw(ArgumentError("number of bins must be ≥ 1 for a non-empty array, got $n"))
    end
    if nv == 0
        return 0:1:0
    end
    if n <= 0
        throw(ArgumentError("number of bins n = $n must be positive"))
    end
    lo, hi = extrema(v)
    if hi == lo
        step = 1
    else
        bw = (hi - lo) / n
        e = 10^max(0,floor(Int,log10(bw)))
        r = bw / e
        if r <= 1
            step = e
        elseif r <= 2
            step = 2*e
        elseif r <= 5
            step = 5*e
        else
            step = 10*e
        end
    end
    start = step*(ceil(lo/step)-1)
    nm1 = ceil(Int,(hi - start)/step)
    return start:step:(start + nm1*step)
end

# Sturges' formula
function sturges(n::Integer)
    n==0 && return one(n)
    return ceil(Int, log2(n)) + 1
end

function hist!(h::AbstractArray{HT}, v::AbstractVector, edg::AbstractVector) where HT
    n = length(edg) - 1
    length(h) == n || throw(DimensionMismatch("length(histogram) must equal length(edges) - 1"))
    fill!(h, zero(HT))
    for x in v
        i = searchsortedfirst(edg, x)-1
        if 1 <= i <= n
            h[i] += 1
        end
    end
    return edg, h
end

hist(v::AbstractVector, edg::AbstractVector) = hist!(Array{Int}(undef, length(edg)-1), v, edg)
hist(v::AbstractVector, n::Integer) = hist(v,histrange(v,n))
hist(v::AbstractVector) = hist(v,sturges(length(v)))

end


"""
The `Romberg` module provides the functions `romberg` and `romberg!` that
perform a Romberg integration of discrete 1-dimensional data `y` sampled at
equally spaced points over `x`.

Romberg integration combines the trapezoid method of integration with Richardson
extrapolation. Significantly lower error estimates of definite integrals can be
obtained using Romberg integration over the trapezoid method alone, at only
slightly higher computational cost.
"""
module Romberg

using Trapz

export romberg, romberg!

maxsteps(N) = trunc(Int, log2(prevpow(2, N)))

"""
    romberg(x::AbstractRange, y::AbstractVector)

Integrate `y` sampled at `x` by Romberg integration applying the maximum possible
number of Richardson extrapolation steps.

Romberg integration requires equally spaced points. This is enforced by
requiring `x` to be an `AbstractRange`, rather than a dense vector.

The algorithm is most efficient if `length(x) == 2ⁿ + 1` for any positive integer
`n`. However, the function will still work for any `length(x)`.

# Examples

```jldoctest
julia> x = range(0, π, length=2^8+1);

julia> romberg(x, sin.(x))
1.9999999999999996
```
"""
function romberg(x::AbstractRange{Tx}, y::AbstractVector{Ty}) where {Tx,Ty}
    N = length(x)

    # Integral over nothing
    N <= 1 && return zero(promote_type(Tx,Ty))

    @boundscheck begin
        N == length(y) || throw(DimensionMismatch("length of `y` not equal to length of `x`"))

        # NOTE: by requiring x::AbstractRange, a fixed step size is guaranteed
    end

    pN = prevpow(2, N-1) + 1
    if pN == N
        max_steps = maxsteps(N)
        return integrate(x, y, max_steps)
    else
        return recursive_integration(x, y)
    end
end

"""
    romberg(x::AbstractRange, y::AbstractVector, max_steps::Integer)

Integrate `y` sampled at `x` by Romberg integration applying `max_steps` of
Richardson extrapolation.

`max_steps` is the number of Richardson extrapolation steps applied (also equal
to the number of trapezoid integrations - 1). In general, the larger it is, the
lower error in the integration. The largest `max_steps` can be is
`log2(prevpow(2, length(x)))`.

When specifying `max_steps`, `length(x) - 1` must be a power of 2.

# Examples

```jldoctest
julia> x = range(0, π, length=2^8+1);

julia> romberg(x, sin.(x), 8)
1.9999999999999996
```

```jldoctest
julia> romberg(x, sin.(x), 9)
ERROR: DomainError with 9:
`max_steps` cannot exceed `log2(prevpow(2, length(x)))` = 8
[...]
```
"""
function romberg(x::AbstractRange{Tx}, y::AbstractVector{Ty}, max_steps::Integer) where {Tx,Ty}
    N = length(x)

    # Integral over nothing
    N <= 1 && return zero(promote_type(Tx,Ty))

    @boundscheck begin
        max_steps <= maxsteps(N) || throw(DomainError(max_steps, "`max_steps` cannot exceed `log2(prevpow(2, length(x)))` = $(maxsteps(length(x)))"))
        ispow2(N-1) || throw(DomainError(length(x), "`length(x) - 1` must be a power of 2"))
        N == length(y) || throw(DimensionMismatch("length of `y` not equal to length of `x`"))

        # NOTE: by requiring x::AbstractRange, a fixed step size is guaranteed
    end

    return integrate(x, y, max_steps)
end

"""
    romberg!(R::AbstractMatrix, x::AbstractRange, y::AbstractVector)

Integrate `y` sampled at `x` by Romberg integration, filling in the square
matrix `R` in place.

The size of `R` determines the number of Richardson extrapolation steps. If `R`
has size 6×6, 5 extrapolation steps will be performed. In general, for a matrix
with dimensions `L×L`, `max_steps` corresponds to `L - 1`.

This function is useful to see how the Romberg integration progressed.

# Examples

```jldoctest
julia> x = range(0, π, length=2^8+1);

julia> R = zeros(4, 4);

julia> romberg!(R, x, sin.(x))
4×4 Array{Float64,2}:
 1.92367e-16  0.0      0.0      0.0
 1.5708       2.0944   0.0      0.0
 1.89612      2.00456  1.99857  0.0
 1.97423      2.00027  1.99998  2.00001
```
"""
function romberg!(R::AbstractMatrix, x::AbstractRange, y::AbstractVector)
    N = length(x)

    N <= 1 && return fill!(R, 0)

    @boundscheck begin
        # Assume `size(R, 1)` is L (`max_steps + 1`)
        size(R,1) <= maxsteps(N)+1 || throw(DomainError(size(R,1), "dimenensions of `R` cannot be greater than `log2(prevpow(2, length(x))) + 1`"))
        isequal(size(R)...) || throw(DimensionMismatch("`R` must be square"))
        N == length(y) || throw(DimensionMismatch("length of `y` not equal to length of `x`"))

        # NOTE: by requiring x::AbstractRange, a fixed step size is guaranteed
    end

    integrate!(R, x, y)
end

function recursive_integration(x, y)
    # Needed if `length(x) - 1` is not a power of 2

    N = length(x)

    N <= 1 && return zero(promote_type(x,y))

    pN = prevpow(2, N-1) + 1

    if pN == N
        # `ispow2(N - 1) == true`
        return integrate(x, y, maxsteps(N))
    elseif pN < N
        # `ispow2(N - 1) == false`
        return integrate(x[1:pN], view(y,1:pN), maxsteps(pN)) + recursive_integration(x[pN:N], view(y,pN:N))
    else
        return trapz(x, y)
    end
end

@inline function integrate(x::AbstractRange{Tx}, y::AbstractVector{Ty}, max_steps::Integer) where {Tx,Ty}
    # `max_steps` are extrapolation steps, size of `R` is `max_steps + 1`
    L = max_steps + 1

    # `Rp` is "Rprevious", `Rc` is "Rcurrent"
    Rp = Array{promote_type(Tx,Ty)}(undef, L)
    Rc = similar(Rp)

    trapezoid!(Rp, x, y)
    extrapolate!(Rp, Rc)

    # best estimate
    return Rc[end]
end

@inline function integrate!(R::AbstractMatrix, x::AbstractRange, y::AbstractVector)
    L = size(R, 1)

    trapezoid!(view(R,:,1), x, y)
    extrapolate!(R)

    return R
end

@inline function trapezoid!(R, x, y)
    N = length(x)
    @inbounds for i in eachindex(R)
        step_size = div(N-1, 2^(i-1))

        if step_size == 0
            # only triggered when N == 1?
            idxs = 1:N-1:N
        else
            idxs = 1:step_size:N
        end

        R[i] = trapz(x[idxs], view(y,idxs))
    end
end

@inline function extrapolate!(Rp, Rc)
    L = length(Rp)
    @inbounds for j = 2:L
        # Precompute reused values
        pow_prev = 4^(j-1)
        den = inv(pow_prev - 1)

        @inbounds for i = j:L
            Rc[i] = (pow_prev*Rp[i] - Rp[i-1]) * den
        end

        # Only copying the relevant values! (not necessary on last loop)
        if j < L
            copyto!(Rp, j, Rc, j, L-j+1)
        end
    end
end

@inline function extrapolate!(R)
    L = size(R,1)
    @inbounds for j = 2:L
        # Precompute reused values
        pow_prev = 4^(j-1)
        den = inv(pow_prev - 1)

        @inbounds for i = j:L
            R[i,j] = (pow_prev*R[i,j-1] - R[i-1,j-1]) * den
        end
    end
end

function error_test(R, max_steps)
    # See http://www.sci.utah.edu/~beiwang/teaching/cs6210-fall-2016/lecture24.pdf

    # The ratio of the difference between successive entries in column `j`
    # should be approximately 4^j.

    ratio = zeros(max_steps-2, max_steps-2)
    for j = 1:max_steps-2
        ratio[j:end,j] .= (R[1+j:end-1,j] - R[j:end-2,j]) ./ (R[2+j:end,j] - R[1+j:end-1,j])
    end

    return ratio
end

end # module


img = SquareImage(rows=4, cols=4, l=2, calibration_value=1, background=0)
@testset "Testing Radon transform" begin
    @test img ≈ [
        [0 0 0 0];
        [0 1 1 0];
        [0 1 1 0];
        [0 0 0 0]
    ]
    @testset "Testing diag algorithm" begin
        sinog = radon(img, Radon())
        @test sinog ≈ [
            [0.0       0.0       0.0       0.0       0.0];
            [0.178452  0.145332  0.164798  0.164798  0.145332];
            [0.785413  0.763186  0.745526  0.745525  0.763186];
            [0.785413  0.763186  0.745526  0.745525  0.763186];
            [0.178452  0.145332  0.164798  0.164798  0.145332];
            [0.0       0.0       0.0       0.0       0.0]
        ]
    end
    @testset "Testing square algorithm" begin
        sinog = radon(img, RadonSquare())
        @test sinog ≈ [
            [0.0       0.05596    0.044909  0.064763   0.0647629  0.044909  0.05596];
            [0.888889  0.764706   0.846043  0.804828   0.804828   0.846043  0.764706];
            [0.888889  0.764706   0.846043  0.804828   0.804828   0.846043  0.764706];
            [0.0       0.0559601  0.044909  0.0647629  0.0647629  0.044909  0.05596]
        ]
    end
end


include("header.jl")
using LinearAlgebra

@testset "LinearAlgebra" begin

    diagm1(x,y,z) = diagm(1=>x,0=>y,-1=>z)

    function krontest(w1,w2)
        sum(kron(w1,w2))
    end

    function qrtest(w)
        Q,R = qr(w)
        sum(Q) + 2sum(R)
    end

    function lqtest(w)
        Q,L = lq(w)
        sum(Q) + 2sum(L)
    end

    function svdtest(w)
        U,S,V = svd(w)
        sum(U) + 2sum(S) + 3sum(V)
    end

    for T in (Float32, Float64)
        w = randn(T,3,2)
        wt = randn(T,2,3)
        wsquare = I + 0.2*randn(T,3,3)
        wposdef = wsquare' * wsquare
        udg = T[1.,2.,3.]
        dg = T[4.,5.,6.]
        ldg  = T[7.,8.,9.]

        @test gradcheck(adjoint,w)
        @test gradcheck(det,wsquare)
        @test gradcheck(diag,wsquare)
        @test_skip gradcheck(diag,w)             #TODO no support for non-square matrices yet
        @test_skip gradcheck(diagm1,udg,dg,ldg)   #TODO diagm not implemented yet
        @test gradcheck(dot,udg,ldg)
        @test gradcheck(dot,w,copy(w))
        @test gradcheck(inv,wsquare)
        @test gradcheck(krontest,w,copy(w); delta=0.01)
        @test gradcheck(logabsdet,wsquare)
        @test gradcheck(logdet,wposdef)
        @test gradcheck(norm,w,1; args=1)
        @test gradcheck(norm,w,2; args=1)
        @test gradcheck(norm,w,Inf; args=1)
        @test_skip gradcheck(qrtest,w) #TODO iterator error
        @test_skip gradcheck(lqtest,w) #TODO iterator error
        @test_skip gradcheck(svdtest,w) #TODO iterator error
        @test gradcheck(tr,wsquare)
        @test gradcheck(transpose,w)
        @test gradcheck(tril,w)
        @test gradcheck(triu,w)
        @test gradcheck(*,w,wt)
        @test_skip gradcheck(^,wsquare,randn()) #TODO: real matrix powers
        @test_skip gradcheck(^,wsquare,rand(-3:3)) #TODO: integer matrix powers
    end
end


function tracem(x::Array{T, 2}) where T <: AbstractFloat

  # We require the number of rows to be greater than the number of columns, so that m is greater than one.

  trans = false

  if size(x,1) < size(x,2)
    x = x'
    trans = true
  end

  n = size(x,2)
  m = round(Int,size(x,1)/n)

  y = zeros(m,1)  # We want this to be a 2-d array for subsequent matrix multiplication

  for i = 1:m
    y[i,1] = tr(x[(n*(i-1)+1):i*n,1:n])
  end

  if trans == true
    y = y'
  end

  return y

end


#!/usr/bin/env julia

using LKJ
using Test

import Pkg
Pkg.add("ForwardDiff")

using ForwardDiff
using LinearAlgebra

is_corr_matrix(R) = issymmetric(R) & all(diag(R) .== 1.0) & all(R .<= 1.0) & all(R .>= -1.0)
is_corr_cholfac(L) = isapprox(sum(L.^2, dims=2), ones(eltype(L), size(L,1))) && all(diag(L) .> 0.0)

function chol_from_vec(x, K)
    w = vec_to_lt(x, K)
    for i in 1:K
        if sum(w[i,:].^2) > 1
            w[i,i] = 0.0
        else
            w[i,i] = sqrt(1 - sum(w[i,:].^2))
        end
    end
    LowerTriangular(w)
end

function vec_to_lt(y, K)
    w = zeros(eltype(y), K, K)
    k = 1
    for i in 2:K
        for j in 1:(i-1)
            w[i,j] = y[k]
            k += 1
        end
    end
    LowerTriangular(w)
end

function lt_to_vec(w::LowerTriangular)
    K = size(w,1)
    x = Vector{eltype(w)}(undef, Int(K*(K-1)/2))
    k = 1
    for i in 2:K
        for j in 1:(i-1)
            x[k] = w[i,j]
            k += 1
        end
    end
    x
end

## log jacobian factors using automatic differentiation
function log_jacobian_det_free_to_chol_ad(y, K)
    h(x) = lt_to_vec(transform_free_to_chol(x, K))
    log(abs(det(ForwardDiff.jacobian(h, y))))
end

function log_jacobian_det_chol_to_free_ad(w::LowerTriangular, K)
    h(x) = transform_chol_to_free(chol_from_vec(x, K))
    log(abs(det(ForwardDiff.jacobian(h, lt_to_vec(w)))))
end

d = 5
eta = 2.36

distr = LKJcorr(d, eta)
@test params(distr) == (d, eta)
@test length(distr) == d

R = rand(distr)
@test is_corr_matrix(R)

distr = LKJcorrChol(d, eta)
@test params(distr) == (d, eta)
@test length(distr) == d

L = rand(distr)
@test is_corr_cholfac(L)

y = transform_chol_to_free(L)

## check that transform_free_to_chol is inverse of transform_chol_to_free
@test isapprox(transform_free_to_chol(y, d), L)

## check that log jacobian factor from automatic differentation is the same
@test isapprox(log_jacobian_det_free_to_chol(y, d), log_jacobian_det_free_to_chol_ad(y, d))


# # Heights_2 problem with restricted prior on mu.

using DynamicHMCModels

ProjDir = @__DIR__
cd(ProjDir)

# Import the dataset.

data = DataFrame(CSV.read(joinpath("..", "..", "data", "Howell1.csv"), delim=';'));

# Use only adults and standardize

df = filter(row -> row[:age] >= 18, data);

# Flat `σ`, see below.

Base.@kwdef mutable struct Heights_2{Ty <: AbstractVector}
    "Observations."
    y::Ty
end;

# Write a function to return properly dimensioned transformation.

function make_transformation(model::Heights_2)
    as((σ = asℝ₊, μ  = as(Real, 100, 250)), )
end

model = Heights_2(;y = df[:, :height])

# Then make the type callable with the parameters *as a single argument*. Very constraint prior on μ. Flat σ.

function (model::Heights_2)(θ)
    @unpack y = model   # extract the data
    @unpack μ, σ = θ
    loglikelihood(Normal(μ, σ), y) + logpdf(Normal(178, 0.1), μ) +
    logpdf(Uniform(0, 50), σ)
end;

# Wrap the problem with a transformation, then use Flux for the gradient.

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P);

# Tune and sample.

results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000)
posterior = P.transformation.(results.chain)

println()
DynamicHMC.Diagnostics.EBFMI(results.tree_statistics) |> display
println()

DynamicHMC.Diagnostics.summarize_tree_statistics(results.tree_statistics) |> display
println()

a3d = Array{Float64, 3}(undef, 1000, 2, 1);
for j in 1:1
  for i in 1:1000
    a3d[i, 1, j] = values(posterior[i].μ)
    a3d[i, 2, j] = values(posterior[i].σ)
  end
end

pnames = ["μ", "σ"]
sections =   Dict(
  :parameters => pnames,
)
chns = create_mcmcchains(a3d, pnames, sections, start=1);

# cmdstan result

stan_result = "
Iterations = 1:1000
Thinning interval = 1
Chains = 1,2,3,4
Samples per chain = 1000

Empirical Posterior Estimates:
         Mean         SD       Naive SE       MCSE      ESS
sigma  24.604616 0.946911707 0.0149719887 0.0162406632 1000
   mu 177.864069 0.102284043 0.0016172527 0.0013514459 1000

Quantiles:
         2.5%       25.0%     50.0%     75.0%     97.5%
sigma  22.826377  23.942275  24.56935  25.2294  26.528368
   mu 177.665000 177.797000 177.86400 177.9310 178.066000
";

describe(chns)

# end of m4.2d.jl


Base.isapprox(x::Tuple, y::Tuple; kws...) = isapprox(collect(x), collect(y); kws...)

@testset "Analysis struct" begin
@testset "Uniformly distributed values" begin
    n = 20
    data = rand(Uniform(0,10), n)
    curve = CalculatedABC.ABCcurve(data)

    analysis1 = CalculatedABC.ABCanalysis(curve)
    @test isapprox(analysis1.pareto,      (0.4, 0.6192318028485929))
    @test isapprox(analysis1.break_even,  (0.5, 0.7210686469145638))
    @test isapprox(analysis1.demark_AB,   (0.4, 0.6192318028485929))
    @test isapprox(analysis1.submarginal, (0.6, 0.8125961014234772))

    analysis2 = CalculatedABC.ABCanalysis(data)
    @test isapprox(analysis2.pareto,      (0.4, 0.6192318028485929))
    @test isapprox(analysis2.break_even,  (0.5, 0.7210686469145638))
    @test isapprox(analysis2.demark_AB,   (0.4, 0.6192318028485929))
    @test isapprox(analysis2.submarginal, (0.6, 0.8125961014234772))
    @test analysis2.a_indices == [2, 13, 1, 3, 6, 14, 7, 19]
    @test analysis2.b_indices == [12, 16, 18, 4]
    @test analysis2.c_indices == [5, 15, 17, 9, 20, 11, 10, 8]
end

@testset "Exponentially distributed values" begin
    n = 20
    data = rand(Exponential(1), n)
    curve = CalculatedABC.ABCcurve(data)

    analysis1 = CalculatedABC.ABCanalysis(curve)
    @test isapprox(analysis1.pareto,      (0.35, 0.6503414525690906))
    @test isapprox(analysis1.break_even,  (0.35, 0.6503414525690906))
    @test isapprox(analysis1.demark_AB,   (0.35, 0.6503414525690906))
    @test isapprox(analysis1.submarginal, (0.5, 0.7704643555492489))

    analysis2 = CalculatedABC.ABCanalysis(data)
    @test isapprox(analysis2.pareto,      (0.35, 0.6503414525690906))
    @test isapprox(analysis2.break_even,  (0.35, 0.6503414525690906))
    @test isapprox(analysis2.demark_AB,   (0.35, 0.6503414525690906))
    @test isapprox(analysis2.submarginal, (0.5, 0.7704643555492489))
    @test analysis2.a_indices == [2, 19, 13, 20, 6, 4, 5]
    @test analysis2.b_indices == [15, 3, 12]
    @test analysis2.c_indices == [16, 7, 17, 9, 11, 14, 18, 8, 1, 10]
end

@testset "SwissInhabitants data" begin
    data = readdlm("swissinhabitants")[:,1]
    curve = CalculatedABC.ABCcurve(data)

    analysis1 = CalculatedABC.ABCanalysis(curve)
    @test isapprox(analysis1.pareto,      (0.25379834254143646, 0.7219586074342286))
    @test isapprox(analysis1.break_even,  (0.22168508287292818, 0.6913764432243525))
    @test isapprox(analysis1.demark_AB,   (0.22168508287292818, 0.6913764432243525))
    @test isapprox(analysis1.submarginal, (0.3715469613259669, 0.810278191608252))

    analysis2 = CalculatedABC.ABCanalysis(data)
    @test isapprox(analysis2.pareto,      (0.25379834254143646, 0.7219586074342286))
    @test isapprox(analysis2.break_even,  (0.22168508287292818, 0.6913764432243525))
    @test isapprox(analysis2.demark_AB,   (0.22168508287292818, 0.6913764432243525))
    @test isapprox(analysis2.submarginal, (0.3715469613259669, 0.810278191608252))
    # too many to test directly
    @test (analysis2.a_indices |> length) == 642
    @test (analysis2.b_indices |> length) == 434
    @test (analysis2.c_indices |> length) == 1820
    @test minimum(data[analysis2.a_indices]) >= maximum(data[analysis2.b_indices])
    @test minimum(data[analysis2.b_indices]) >= maximum(data[analysis2.c_indices])
end
end


# Note: All the estimates and inequalities below only apply
# for our iterated integrals because the Levy area and the
# corresponding error terms are skew-symmetric!

"""
    abstract type AbstractErrorNorm end

Abstract type for different kind of errors one might consider.
The most important are `MaxL2` and `FrobeniusL2`. These are actually
special cases of the `MaxLp{p}` and `lqLp{p,q}` norms:
```julia
const MaxL2 = MaxLp{2}
const FrobeniusL2 = lqLp{2,2}
```

All currently defined norms:
```jldoctest; setup=:(using InteractiveUtils; using LevyArea)
julia> subtypes(LevyArea.AbstractErrorNorm)
3-element Vector{Any}:
 LevyArea.LpMax
 LevyArea.MaxLp
 LevyArea.lqLp
```
"""
abstract type AbstractErrorNorm end

# maximum of entry-wise Lp(Ω) norms
# this requires new Lp convergence proofs
# and should be defined for each new algorithm
struct MaxLp{p} <: AbstractErrorNorm end

function errcoeff(m, q_12, stepsize, alg::AbstractIteratedIntegralAlgorithm, ::MaxLp{p}) where {p}
    maxqq = maximum(q_12[i]*q_12[j] for i=1:m for j=1:i-1)
    return maxqq * errcoeff(m, stepsize, alg, MaxLp{p}())
end

# lq norm of all entry-wise Lp(Ω) norms
# = (Σᵢⱼ||Mᵢⱼ||_Lp(Ω)^q)^(1/q)
# we can get this from the MaxLp norms
# p=q=2 is the FrobeniusL2 norm
struct lqLp{p,q} <: AbstractErrorNorm end

function errcoeff(m, stepsize, alg::AbstractIteratedIntegralAlgorithm, ::lqLp{p,q}) where {p,q}
    return (m^2-m)^(1/q) * errcoeff(m, stepsize, alg, MaxLp{p}())
end
function errcoeff(m, q_12, stepsize, alg::AbstractIteratedIntegralAlgorithm, ::lqLp{p,q}) where {p,q}
    trQq_sq = abs2(sum(x->x^q, q_12))
    tr_Qqsq = sum(x->x^2q, q_12)
    return (trQq_sq-tr_Qqsq)^(1/q) * errcoeff(m, stepsize, alg, MaxLp{p}())
end

# Lp norm of maximum of matrix
# can be bounded by the lqLp norm with q=p
# (this relies on the fact that lp-Lp == Lp-lp)
# the factor 1/2^1/p comes from the antisymmetry
struct LpMax{p} <: AbstractErrorNorm end

function errcoeff(m, stepsize, alg::AbstractIteratedIntegralAlgorithm, ::LpMax{p}) where {p}
    return 1/(2^(1/p)) * errcoeff(m, stepsize, alg, lqLp{p,p}())
end
function errcoeff(m, q_12, stepsize, alg::AbstractIteratedIntegralAlgorithm, ::LpMax{p}) where {p}
    return 1/(2^(1/p)) * errcoeff(m, q_12, stepsize, alg, lqLp{p,p}())
end

# most important/common norms
const MaxL2 = MaxLp{2}
const FrobeniusL2 = lqLp{2,2}


using LinearAlgebra, PyPlot, Random, FFTW, Distributions, CircularArrays, Printf, Statistics

close(:all)

ci, co, N, ps, b = 16, 16, 256, 64, 50
plot_I = false

eb = randn(Float32, co*N, ps)

chunk(arr, n) = [arr[i:min(i + n - 1, end)] for i in 1:n:length(arr)]
simil(x, y) = 100*dot(x[:], y[:])/(norm(x)*norm(y))
error(x, y) = norm(x - y)/(norm(x)+norm(y))

function draw_e(ps::Integer, co::Integer, N::Integer)
    n = ps ÷ co
    if n < 8
        n = 8
        inds = chunk(randperm(n*co) .% ps .+ 1, n)
        overlap = 1 #n / (ps ÷ co)
    else
        inds = chunk(randperm(ps), n)
        overlap = 1
    end
    e = zeros(Float32, N*co, ps)
    for i=1:co
        iloc = inds[i]
        ni = length(iloc)
        e[(i-1)*N+1:i*N, iloc] .= Float32(sqrt(ps/(overlap*ni)))*randn(Float32, N, ni)
    end

    return e
end

eortho = draw_e(ps, co, N)

plot_e_scale = ps/50
figure()
subplot(221)
imshow(eb, vmin=-.1, vmax=.1, cmap="seismic", aspect="auto")
title("Z, Random")
subplot(222)
imshow(eortho, vmin=-.1, vmax=.1, cmap="seismic", aspect="auto")
title("Z, Random Orthogonalized")
subplot(223)
imshow(eb*eb', vmin=-plot_e_scale, vmax=plot_e_scale, cmap="seismic", aspect="auto")
title(L"$ZZ^T$ Random")
subplot(224)
imshow(eortho*eortho', vmin=-plot_e_scale, vmax=plot_e_scale, cmap="seismic", aspect="auto")
title(L"$ZZ^T$ Random Orthogonalized")
tight_layout()
savefig("./figures/zortho.png", bbox_inches="tight")


import LsqFit
import Distributions
import PyPlot.plt

#Fit result type
immutable JLabFitresult
    #Fit inputs
    xdata::Array{Float64,1}     #independent variable
    ydata::Array{Float64,1}     #dependent variable
    yerrors::Array{Float64,1}   #uncertainties in dependent variable
    model::Function             #model to be fitted against

    #Fit results
    residue::Array{Float64,1}   #weighted residue
    params::Array{Float64,1}    #best fit parameters for model
    perrors::Array{Float64,1}   #uncertainties in best fit parameters
    dof::Int64                  #number of degrees of freedom in fit
    chisq::Float64              #chi^2 value of fit
    probability::Float64        #probability of chi^2 value
end

#Fitting procedure
function getfit{T1<:Real,T2<:Real,T3<:Real}(model::Function,xdata::Array{T1,1},ydata::Array{T2,1},initparams::Array{T3,1};kwargs...)
    return getfit(model,xdata,ydata,ones(Float64,length(ydata)),initparams;kwargs...)
end

function getfit{T1<:Real,T2<:Real,T3<:Real,T4<:Real}(model::Function,xdata::Array{T1,1},ydata::Array{T2,1},yerrors::Array{T3,1},initparams::Array{T4,1};alpha=0.6827)
    #Takes a model function with experimental data and returns a fit result.
    #model is a function f(x,p) with x an array of values for the independent
    #variable, and p an array of parameters for the model, and returns values for the
    #dependent variable.
    #
    #Usage:
    #   fit = getfit(model,xdata,ydata,yerrors,initparams)
    #       model           Function on (x,params)
    #       xdata           independent variable values
    #       ydata           dependent variable values
    #       yerrors         uncertainty in dependent variable
    #       initparams      initial guess for model parameters

    #Conduct fit
    weights = 1./yerrors
    fit = LsqFit.curve_fit(model,float64(xdata),float64(ydata),float64(weights),float64(initparams))
    perrors = LsqFit.estimate_errors(fit,alpha)

    chisq = sum(fit.resid.^2)
    chisqdist = Distributions.Chisq(fit.dof)
    probability = Distributions.ccdf(chisqdist,chisq)

    return JLabFitresult(float64(xdata),float64(ydata),float64(yerrors),model,fit.resid,fit.param,perrors,fit.dof,chisq,probability)
end

#Plot fit result
function plotfit(fit::JLabFitresult;datakwargs=(),modelkwargs=())
    #Makes a plot of data points and fit curve with annotations on fit parameters

    #Plot experimental data
    datapts = plt.errorbar(fit.xdata,fit.ydata,fit.yerrors,fmt=".",label="Expt data";datakwargs...)

    #Plot Model
    x = linspace(minimum(fit.xdata),maximum(fit.xdata),200)
    y = fit.model(x,fit.params)
    modelline = plt.plot(x,y,"r",lw=2,label="Model";modelkwargs...)

    #Annotations
    println("@sprintf(\"\\\$\\\\chi^2\\\$ / ndf = %.2f / %d = %.2f\\np(\\\$\\\\chi^2\\\$) = %.3f\\n\",fit.chisq,fit.dof,fit.chisq/fit.dof,fit.probability)")
    @printf("\$\\chi^2\$ / ndf = %.2f / %d = %.2f\np(\$\\chi^2\$) = %.3f\n\n",fit.chisq,fit.dof,fit.chisq/fit.dof,fit.probability)
    annot = plt.text(0.5*(minimum(fit.xdata)+maximum(fit.xdata)),0.5*(minimum(fit.ydata)+maximum(fit.ydata)),@sprintf("\$\\chi^2\$ / ndf = %.2f / %d = %.2f\np(\$\\chi^2\$) = %.3f\n\n",fit.chisq,fit.dof,fit.chisq/fit.dof,fit.probability))

    #Plot
    plt.xlabel("\$x\$")
    plt.ylabel("\$y\$")
    plt.title("\$y\$ vs \$x\$")
    plt.legend()

    return datapts,modelline,annot
end

# #############################################################################
#Common fitting functions

#Linear
linear{T<:Real,N}(x::Array{T,N},p::Array{Float64,1}) = p[1]+p[2]*x

#Polynomial
polynomial{T<:Real,N}(x::Array{T,N},p::Array{Float64,1}) = broadcast(^,x,transpose([0:length(p)-1]))*p

#Gaussian
gaussian{T<:Real,N}(x::Array{T,N},p::Array{Float64,1}=[1/sqrt(2*pi),1.,0.]) = p[1]*exp(-0.5*(x-p[3]).^2/p[2]^2)

#Lorentzian
lorentzian{T<:Real,N}(x::Array{T,N},p::Array{Float64,1}=[2/pi,1.,0.]) = p[1]/(1+(x-p[3]).^2/p[2]^2)

#Exponential decay
decay{T<:Real,N}(x::Array{T,N},p::Array{Float64,1}=[1.,1.]) = p[1]*exp(-p[2]*x)


function gradient(f, xs...)
  xs = param.(xs)
  back!(f(xs...))
  grad.(xs)
end

function ngradient(f, xs::AbstractArray...)
  grads = zeros.(xs)
  for (x, Δ) in zip(xs, grads), i in 1:length(x)
    δ = sqrt(eps())
    tmp = x[i]
    x[i] = tmp - δ/2
    y1 = f(xs...)
    x[i] = tmp + δ/2
    y2 = f(xs...)
    x[i] = tmp
    Δ[i] = (y2-y1)/δ
  end
  return grads
end

gradcheck(f, xs...) =
  all(isapprox.(ngradient(f, xs...),
                gradient(f, xs...), rtol = 1e-5, atol = 1e-5))