using ApproxFun,SingularIntegralEquations,Base.Test
    import SingularIntegralEquations:stieltjesmoment,stieltjesjacobimoment

c = [0.9731840665678853,0.11644664868790366,0.8961305368364185,0.30663942299763747,0.4564158422422153,0.7262396331533589,0.6448725310785985,0.05623500638232981,0.8452582204677404,0.25385833878392283]

for z in [1.+1.0im,0.1+0.1im]
    o=Fun(one)
    x=Fun(identity)

    @test_approx_eq stieltjes(o)(z) sum(1/(z-x))
    @test_approx_eq stieltjes(x)(z) sum(x/(z-x))

    for w in (sqrt(1-x),sqrt(1+x))
        @test_approx_eq stieltjes(w*o)(z) sum(w/(z-x))
        @test_approx_eq stieltjes(w*x)(z) sum(w*x/(z-x))
    end

    for a in -0.5:0.5:0.5, b in -0.5:0.5:0.5
        f = Fun(c,WeightedJacobi(a,b))
        @test_approx_eq stieltjes(f)(z) sum(f/(z-x))
    end

    f = Fun(c,WeightedJacobi(0.123,0.456))
    @test_approx_eq stieltjes(f)(z) sum(f/(z-x))

    f = Fun(identity,[-1.,0.,1.])
    @test_approx_eq cauchy(sqrt(Fun(one,space(f))-f^2))(z) cauchy(sqrt(1-Fun()^2),z)
end

x=Fun(identity,[im,0.,1.])
w=2/(sqrt(1-x)*sqrt(1+im*x))

for x in (0.9im,0.4im,0.4,0.9)
    #@test_approx_eq cauchy(w,x,true)-cauchy(w,x,false) w(x)
end


@testset "LinearAlgebra" begin
    @testset "Constructor + basics" begin
        A = rand(ComplexF32, 12, 12)
        WS = HC2.MatrixWorkspace(A)
        @test WS isa HC2.MatrixWorkspace

        WS.factorized[] = true
        B = randn(ComplexF64, 12, 12)
        HC2.update!(WS, B)
        @test WS.A == B
        @test WS.A !== B
        @test WS.factorized[] == false

        # test indexing etc
        @test WS[2, 1] == B[2, 1]
        @test WS[2] == B[2]
        WS[2, 3] = 5.0
        @test WS[2, 3] == 5.0
        WS[5] = 2.32
        @test WS[5] == 2.32
        @test size(WS) == (12, 12)
        C = randn(ComplexF64, 12, 12)
        WS .= C
        @test WS.A == C

        #test struct array switch
        A = rand(ComplexF32, 30, 30)
        WS = HC2.MatrixWorkspace(A)
    end

    @testset "QR" begin
        A = randn(ComplexF64, 12, 7)
        MW = HC2.MatrixWorkspace(A)
        B = randn(ComplexF64, 12, 7)
        true_qr = LinearAlgebra.qrfactUnblocked!(copy(B))
        MW .= B

        HC2.updated!(MW)

        b = randn(ComplexF64, 12)
        b2 = copy(b)
        x = zeros(ComplexF64, 7)

        ldiv!(x, MW, b)

        @test norm(MW.qr.factors - true_qr.factors) / norm(MW.qr.factors) < 1e-14
        @test norm(b2 - b) == 0
        @test norm(x - (B \ b2)) / norm(x) < 1e-14
    end

    @testset "ldiv" begin
        for n in [3, 13, 31] # test that struct array also works
            A = randn(ComplexF64, n, n)
            b = randn(ComplexF64, n)
            x = zeros(ComplexF64, n)
            WS = HC2.MatrixWorkspace(A)
            HC2.update!(WS, A)
            ldiv!(x, WS, b)
            WS.factorized[] = false
            @test (@allocated ldiv!(x, WS, b)) == 0
            @test (lu(A) \ b) ≈ x rtol = cond(A) * 10
        end
    end

    @testset "Inf-norm estimator / cond" begin
        d_r = rand() * exp10.(range(-6, stop = 6, length = 6))
        D_R = diagm(d_r)
        d_l = rand() * exp10.(range(6, stop = -6, length = 6))
        D_L = diagm(d_l)

        A = randn(6, 6)

        B = A * inv(D_R)
        WB = HC2.MatrixWorkspace(B)
        @test 0.1 ≤ opnorm(inv(B), Inf) / HC2.inverse_inf_norm_est(WB) ≤ 10
        @test 0.1 ≤
              opnorm(inv(B * D_R), Inf) / HC2.inverse_inf_norm_est(WB, nothing, d_r) ≤
              10
        @test 0.1 ≤ opnorm(inv(D_R * B), Inf) / HC2.inverse_inf_norm_est(WB, d_r) ≤ 10

        @test 0.1 ≤ cond(B, Inf) / cond(WB) ≤ 10
        @test 0.1 ≤ cond(B * D_R, Inf) / cond(WB, nothing, d_r) ≤ 10
        @test 0.1 ≤ cond(D_R * B, Inf) / cond(WB, d_r, nothing) ≤ 10

        C = inv(D_L) * A
        WC = HC2.MatrixWorkspace(C)
        @test 0.1 ≤ opnorm(inv(C), Inf) / HC2.inverse_inf_norm_est(WC) ≤ 10
        @test 0.1 ≤ opnorm(inv(D_L * C), Inf) / HC2.inverse_inf_norm_est(WC, d_l) ≤ 10
        @test 0.1 ≤
              opnorm(inv(C * D_L), Inf) / HC2.inverse_inf_norm_est(WC, nothing, d_l) ≤
              10
        @test 0.1 ≤ cond(C, Inf) / cond(WC) ≤ 10
        @test 0.1 ≤ cond(C * D_L, Inf) / cond(WC, nothing, d_l) ≤ 10
        @test 0.1 ≤ cond(D_L * C, Inf) / cond(WC, d_l, nothing) ≤ 10

        D = inv(D_L) * A * inv(D_R)
        WD = HC2.MatrixWorkspace(D)
        @test 0.1 ≤ opnorm(inv(D), Inf) / HC2.inverse_inf_norm_est(WD) ≤ 10
        @test 0.1 ≤
              opnorm(inv(D_L * D * D_R), Inf) / HC2.inverse_inf_norm_est(WD, d_l, d_r) ≤
              10
        @test HC2.inverse_inf_norm_est(WD, d_r, d_l) >
              100 * HC2.inverse_inf_norm_est(WD, d_l, d_r)
        @test 0.1 ≤ cond(D, Inf) / cond(WD) ≤ 10
        @test 0.1 ≤ cond(D_L * D * D_R, Inf) / cond(WD, d_l, d_r) ≤ 10
    end

    @testset "Jacobian" begin
        A = randn(ComplexF64, 6, 6)
        b = randn(ComplexF64, 6)
        x̂ = similar(b)
        x = similar(b)
        JM = HC2.Jacobian(zeros(ComplexF64, 6, 6))
        @test JM isa HC2.Jacobian
        HC2.matrix(JM) .= A
        HC2.updated!(JM)
        ldiv!(x̂, JM, b)
        ldiv!(x, HC2.workspace(JM), b)
        @test x == x̂
    end
end


n=100
nr=3
d=randn(n,nr)
g=randn(n,nr)
s=randn(n,nr)
pa=Conv.P_conv(d=d,g=g,s=s, gsize=[n,nr], dsize=[n,nr], ssize=[n,nr]);

function adjtest()
	x=randn(size(F,2))
	y=randn(size(F,1))
	a=LinearAlgebra.dot(y,F*x)
	b=LinearAlgebra.dot(x,adjoint(F)*y)
	c=LinearAlgebra.dot(x, transpose(F)*F*x)
	println("adjoint test: ", a, "\t", b)
	@test isapprox(a,b,rtol=1e-6)
	println("must be positive: ", c)
	@test c>0.0
end

function do_tests(m=3)

	@testset "S" begin
		global pa
		global F=Conv.operator(pa, Conv.S());
		adjtest()
	end

	@testset "G" begin
		global pa
		global F=Conv.operator(pa, Conv.G());
		adjtest()
	end

	@testset "D" begin
		global pa
		global F=Conv.operator(pa, Conv.D());
		adjtest()
	end
end
do_tests()

n=100
nr=3
d=randn(n,nr)
g=randn(n,nr)
s=randn(n)
pa=Conv.P_conv(d=d,g=g,s=s, gsize=[n,nr], dsize=[n,nr], ssize=[n]);
@testset "S" begin
	global pa
	global F=Conv.operator(pa, Conv.S());
	adjtest()
end

#=
@testset "BD" begin
	global p=pa.plsbd
	for attrib in [FBD.S(), FBD.G()]
		global F=FBD.operator(p, attrib);
		adjtest()
	end
end

# adj tests of filter
pa=FBD.P_bandpass(Float64, fmin=0.1, fmax=0.4, nt=101)
global F=FBD.create_operator(pa)
@testset "bandpass" begin
	adjtest()
end
=#


struct UnionFind
    par::Vector{Int}
    size::Vector{Int}
    UnionFind(n) = new(collect(1:n), ones(Int, n))
end

root!(uf, x) = uf.par[x] == x ? x : uf.par[x] = root!(uf, uf.par[x])

same(uf, x, y) = root!(uf, x) == root!(uf, y)

function unite!(uf, x, y)
    x = root!(uf, x)
    y = root!(uf, y)
    x == y && return false

    if uf.size[x] > uf.size[y]
        uf.par[y] = x
        uf.size[x] += uf.size[y]
    else
        uf.par[x] = y
        uf.size[y] += uf.size[x]
    end
    return true
end
size!(uf, x) = uf.size[root!(uf, x)]

function main()
    N, K, L = parse.(Int, split(readline()))

    uf₁ = UnionFind(N)
    uf₂ = UnionFind(N)

    for _ in 1:K
        p, q = parse.(Int, split(readline()))
        unite!(uf₁, p, q)
    end

    for _ in 1:L
        r, s = parse.(Int, split(readline()))
        unite!(uf₂, r, s)
    end

    d = Dict{Tuple{Int, Int}, Int}()
    for i in 1:N
        key = (root!(uf₁, i), root!(uf₂, i))
        if haskey(d, key)
            d[key] += 1
        else
            d[key] = 1
        end
    end

    for i in 1:N-1
        print(d[(root!(uf₁, i), root!(uf₂, i))], " ")
    end

    println(d[(root!(uf₁, N), root!(uf₂, N))])
end

main()


const N_A = 6.02214076e23 # mol-1
const k_B = 1.380649e-23 # m2 kg s-2 K-1
const R̄   = N_A*k_B # m2 kg s-2 K-1 mol-1
const ħ = 1.054571817e-34 # J s
const h = 1.054571817e-34*2*π # J s


##用于实现一些有用的函数

export  add_gaussian_noise,
        hilbert,
        find_fft_peak

#高斯白噪音注入,信噪比越大越好
function add_gaussian_noise(single::AbstractArray{T, 1}, SNR::Float64) where{T<:Number}
    ps = sum(single.^2) / length(single)
    pn_set = sqrt(ps / (10^(SNR / 10)))

    return single + pn_set * randn(Float64, size(single))
end

function add_gaussian_noise(single::AbstractArray{T, 2}, SNR::Float64; dims=1) where{T<:Number}
    buff = similar(single)
    if dims == 1
        @views Threads.@threads for i in 1:size(single, 2)
            buff[:, i] = add_gaussian_noise(single[:, i], SNR)
        end
    elseif dims == 2
        @views Threads.@threads for i in 1:size(single, 1)
            buff[i, :] = add_gaussian_noise(single[i, :], SNR)
        end
    else
        error("dims should be 1 or 2")
    end
    return @view buff[:,:]
end
#希尔伯特变换实现
function hilbert(x::AbstractArray{T, 1}) where{T<:Real}
    buffer = similar(x, Complex)
    buffer = fft(x, 1)
    len = length(x)
    if len %2 ==0
        buffer[2: div(len, 2)] .= 2 .*buffer[2: div(len, 2)]
        buffer[div(len, 2)+1: end] .= 0
    elseif len % 2 == 1
         buffer[2: div(len+1, 2)] .= 2 .*buffer[2: div(len+1, 2)]
         buffer[div(len+1, 2)+1:end] .= 0
    end
    buffer = ifft(buffer, 1)
    buffer1 = similar(x, Complex)
    map!((x, y) -> complex(x, imag(y)), buffer1, x, buffer)
    return view(buffer1,:)
end

function hilbert(x::AbstractArray{T, 2}; dims = 1) where{T <: Real}
    buff =  similar(x, Complex)
    if dims == 1
        @views Threads.@threads for i in 1:size(x, 2)
             buff[:, i] = hilbert(x[:, i])
        end
    elseif  dims == 2
        @views Threads.@threads for i in 1:size(x, 1)
             buff[i, :] = hilbert(x[i, :])
        end
    else
        error("dims should be 1 or 2")
    end

    return @view buff[:,:]
end

##进行傅里叶变换并寻找波峰值,找到前5个波峰
##如果调用上面那种阈值需要调整
function find_fft_peak(x::Vector{T}, fs, Threshold::Float64) where{T}
    xf = abs.(fft(x))
    index = Int[]
    L = length(xf)
    P = (xf ./ L)[1:Int(L/2)]
    @. P[2: end] = 2 * P[2: end]

    f = view(fs*(0:L-1) / L, 1:Int(L/2))

    ##直接for循环遍历
    Threads.@threads for i in 1:length(P)
        P[i] >= Threshold && push!(index, i)
    end

    return view(f, index)
end
function find_fft_peak(x::Vector,fs, n::Int)
    xf = abs.(fft(x))
    L = length(x)
    index = partialsortperm(view(xf, 1:Int(L/2)), 1:n, rev = true)
    f = (fs*(0:L-1)/L)[1:Int(L/2)]

    return view(f, index)
end


## JL-ML SPGL Comparison
export jl_cs

"""
An example of using implicit JOLI operators to solve a system using GenSPGL

Use: x, r, g, info, SNR = jl_cs()
"""
function jl_cs()
    n = 1024
    k = 200

    # Create a system
    p = randperm(n)
    x0 = zeros(n)
    x0[p[1:k]] = sign.(randn(k))

    # Create a modelling op
    F = joDCT(n)

    # Create a restriction op
    ind = randperm(n)
    ind = ind[1:Int(floor(0.6*n))]
    R = joRestriction(n,ind, DDT = Float64, RDT = Float64)

    # Create data
    b = R*F*x0

    # Solve
    opts = spgOptions(optTol = 1e-4,
                         verbosity = 1)

    @time x, r, g, info = spgl1(R*F, vec(b), tau = 0., sigma = 1e-3, options = opts)

    # Calc SNR of recovery
    SNR = -20*log10(norm(x0-x)/norm(x0));

    return x, r, g, info, SNR
end


immutable MixtureModel <: Distribution
    components::Vector # Vector should be able to contain any type of
                       # distribution with comparable support
    probs::Vector{Float64}
    drawtable::DiscreteDistributionTable
    function MixtureModel(c::Vector, p::Vector{Float64})
        if length(c) != length(p)
            error("components and probs must have the same number of elements")
        end
        sump = 0.0
        for i in 1:length(p)
            if p[i] < 0.0
                error("MixtureModel: probabilities must be non-negative")
            end
            sump += p[i]
        end
        table = DiscreteDistributionTable(p)
        new(c, p ./ sump, table)
    end
end

function mean(d::MixtureModel)
    m = 0.0
    for i in 1:length(d.components)
        m += mean(d.components[i]) * d.probs[i]
    end
    return m
end

function pdf(d::MixtureModel, x::Any)
    p = 0.0
    for i in 1:length(d.components)
        p += pdf(d.components[i], x) * d.probs[i]
    end
    return p
end

function rand(d::MixtureModel)
    i = draw(d.drawtable)
    return rand(d.components[i])
end

# TODO: Correct this definition
function var(d::MixtureModel)
    m = 0.0
    for i in 1:length(d.components)
        m += var(d.components[i]) * d.probs[i]^2
    end
    return m
end


using COSMOAccelerators
using ProximalOperators
using ProximalAlgorithms
using LinearAlgebra
using Test
# problem taken from
# https://github.com/kul-forbes/ProximalAlgorithms.jl/blob/master/test/problems/test_lasso_small.jl
T = Float64
A = T[
    1.0 -2.0 3.0 -4.0 5.0
    2.0 -1.0 0.0 -1.0 3.0
    -1.0 0.0 4.0 -3.0 2.0
    -1.0 -1.0 -1.0 1.0 3.0
]

b = T[1.0, 2.0, 3.0, 4.0]
m, n = size(A)
R = real(T)
lam = R(0.1) * norm(A' * b, Inf)

# functions for proximal operator
f = LeastSquares(A, b)
g = NormL1(lam)

# correct solution
x_star = T[-3.877278911564627e-01, 0, 0, 2.174149659863943e-02, 6.168435374149660e-01]

# Douglas-Rachford splitting
function solve(iter::ProximalAlgorithms.DRS_iterable, maxiter, gamma, tol)
    tol_stop(state::ProximalAlgorithms.DRS_state) = norm(state.res, Inf) / gamma <= tol

    iter = ProximalAlgorithms.halt(iter, tol_stop) # tolerance check
    iter = Base.Iterators.take(iter, maxiter) # max iter
    iter = enumerate(iter) # state = (k, state)

    state_final = nothing
    for state in iter
        state_final = state
    end
    return state_final[2].y, state_final[2].x, state_final[1]
end

# Douglas-Rachford splitting and Anderson Acceleration
function solve_and_accelerate(iter::ProximalAlgorithms.DRS_iterable, aa::AndersonAccelerator, maxiter, gamma, tol)
    tol_stop(state::ProximalAlgorithms.DRS_state) = norm(state.res, Inf) / gamma <= tol
    x_prev = copy(iter.x)

    iter = ProximalAlgorithms.halt(iter, tol_stop)
    iter = Base.Iterators.take(iter, maxiter)

    state_final = nothing
    k_final = nothing
    for (k, state) in enumerate(iter)
        state_final = state
        k_final = k

        # accelerate
        # if k > 6
            COSMOAccelerators.update!(aa, state.x, x_prev, k)
            COSMOAccelerators.accelerate!(state.x, x_prev, aa, k)
        # end
        @. x_prev = state.x

    end
    return state_final.y, state_final.z, k_final
end

# test
@testset "Accelerator and Douglas-Rachford (ProximalOperators.jl)" begin
    tol = 1e-6
    maxiter = 200
    x0 = zeros(5)
    gamma = R(10.0) / opnorm(A)^2

    # vanilla Douglas-Rachford
    iter = ProximalAlgorithms.DRS_iterable(f, g, x0, gamma)
    y, z, it = solve(iter, maxiter, gamma, tol)
    @test norm(y - x_star, Inf) <= tol

    # Douglas-Rachford + Anderson Acceleration
    iter = ProximalAlgorithms.DRS_iterable(f, g, x0, gamma)
    aa = AndersonAccelerator(length(x0))
    y_aa, z_aa, it_aa = solve_and_accelerate(iter, aa, maxiter, gamma, tol)
    @test norm(y_aa - x_star, Inf) <= tol
    @test it_aa < it

end


# Functions for the Vector Autoregressive model
import Pkg;
using ProgressMeter;
using LinearAlgebra, Statistics;
using Gadfly, Colors;

eye(n) = Matrix{Float64}(I, n, n)
chol(mat) = convert(Array{Float64}, cholesky(Hermitian(mat)).U');

# (1) Function to calculate impulse response
function func_VAR(data, p)
    # x: partial equilibrium data
    # p: lag order
    T, k = size(data)
    # Set up L.H.S.: set t-by-k array
    y = data';
    Y = y[:, p:T];

    for i = 1:(p-1)
        Y = [Y; y[:, (p-i):(T-i)]];
    end
    # Set up R.H.S.
    X = [ones(1, T-p); Y[:, 1:(T-p)]];
    Y = y[:, ((p+1):T)];

    # Least square estimation
    𝚩 = (Y*X')/(X*X');
    𝞄 = (Y - 𝚩 * X)';
    𝝨 = 𝞄' * 𝞄 ./ (T-p-p*k-1);
    return 𝚩, 𝞄, 𝝨;
end

function func_IRFvar(data, p, h, ℏ)
    T, k = size(data);

    𝚩_ols, 𝞄_ols, 𝝨_ols = func_VAR(data, p);
    V_ols = 𝚩_ols[:, 1];
    𝚩_ols = 𝚩_ols[:, 2:size(𝚩_ols,2)];
    # For standardization of the impulse response,
    ρ = convert(Array{Float64}, cholesky(𝝨_ols).U');

    𝝖 = [𝚩_ols; [kron(eye(k), eye(p - 1)) zeros(k*(p-1), k)]];
    J = [eye(k) zeros(k, k*(p-1))];

    # Compute impulse response
    ψ = zeros(k, k, (h + 1));

    for i_irf = 0:1:h
        ψ[:,:, (i_irf + 1)] = J * 𝝖^(i_irf) * J';
        # If needed to be standardized, multiply ρ
    end

    # Include recursive-design wild bootstrap calculation of confidence band
    #   Follow Kilian (2009, AER)
    #srand(8301993);
    n_repli = 2000; # Number of iterations

    ψ_boot = zeros(k, k, (h + 1), n_repli);

    y = data';
    Y = y[:, p:T];
    for i = 1:(p-1)
        Y = [Y; y[:, (p-i):(T-i)]];
    end

    Ur = zeros(k*p, T-p+1);
    Yr = zeros(k*p, T-p+1);

    @showprogress 1 "Wild-bootstrap confidence band " for i_iter = 1:n_repli
        pos = Int64(round(rand(1)[]*(T-p+1)));
        if pos == 0 #randomization hits zero, default to 1
            pos = 1
        end
        Yr[:, 1] = Y[:, pos];

        # Recrusive design bootstrap
        η = randn(1, size(𝞄_ols, 1));
        η = kron(η, ones(k, 1));

        Ur[1:k, 2:(T-p+1)] = 𝞄_ols' .* η;

        for ii_iter = 2:(T-p+1)
            Yr[:, ii_iter] = [V_ols; zeros(k*(p-1),1)] +
                                    𝝖 * Yr[:, (ii_iter - 1)] +
                                    Ur[:, ii_iter];
        end

        yr = Yr[1:k, :];
        for i=2:p
    		yr = [Yr[(((i-1)*k)+1):(i*k), 1] yr];
        end;
        yr = yr';

        # Iterated partial equilibrium estimation.;
        𝚩, 𝞄, 𝝨 = func_VAR(yr, p);

        𝚩 = 𝚩[:, 2:size(𝚩,2)];
        # For standardization of the impulse response,
        ρ = cholesky(𝝨).L[:,:]; # = cholesky()';

        𝝖 = [𝚩; [kron(eye(k), eye(p - 1)) zeros(k*(p-1), k)]];
        J = [eye(k) zeros(k, k*(p-1))];
        # Computaiton of impulse response
        temp_ψ_prev = [];
        for i_irf = 0:1:h
            temp_ψ = J * 𝝖^(i_irf) * J';
            ψ_boot[:,:, (i_irf + 1), i_iter] = temp_ψ ;
            # If needed to be standardized, multiply ρ
        end
    end
    # Need to calculate percentile
    α = 1; # One standard deviaiton confidence band of bootstrap iterations
    # Thus, this measure of confidence band does not carry literal definition
    #   as in classical model.

    # Calculate the standard deviation, across differen iteration (4th dim)
    σ_ψ_boot = std(ψ_boot, dims = 4);

    ψ_ub_1sd = ψ + σ_ψ_boot
    ψ_lb_1sd = ψ - σ_ψ_boot
    ψ_ub_2sd = ψ + σ_ψ_boot * 2;
    ψ_lb_2sd = ψ - σ_ψ_boot * 2;

    ############################################################################
    # FORECAST ERROR VARIANCE DECOMPOSITION
    T, k = size(data);

    𝚩_ols, 𝞄_ols, 𝝨_ols = func_VAR(data, p);
    V_ols = 𝚩_ols[:, 1];
    𝚩_ols = 𝚩_ols[:, 2:size(𝚩_ols,2)];
    # For standardization of the impulse response,
    ρ = convert(Array{Float64}, cholesky(𝝨_ols).U');

    𝝖 = [𝚩_ols; [kron(eye(k), eye(p - 1)) zeros(k*(p-1), k)]];
    J = [eye(k) zeros(k, k*(p-1))];

    # forecast error variance decomposition
    TH1 = J * 𝝖^(0) *J';
    TH = TH1 * ρ;
    TH = TH';
    TH2 = (TH .* TH);
    TH3 = TH2;

    for i=2:ℏ
        TH = J * 𝝖^(i-1) * J' * ρ;
        TH = TH'; TH2 = (TH.*TH);
        TH3 = TH3 + TH2;
    end;
    TH4 = sum(TH3, dims = 1); VC = zeros(k, k);
    for j=1:k
        VC[j,:] = TH3[j, :] ./ TH4';
    end;

    FEVDC = VC'*100;

    return ψ,
           ψ_lb_1sd, ψ_ub_1sd, ψ_lb_2sd, ψ_ub_2sd,
           FEVDC;
end

function func_contemp(bmat, 𝝨_ols, k, p)
    Ik = eye(k);
    i_block = 1:k; # initial, block of coefficients for same lag.
    for i = 1:p
        Ik = Ik .- bmat[:, i_block];
        # update the block index.
        i_block = i_block .+ k;
    end
    Ik = inv(Ik);

    Bk = eye(k);
    i_block = 1:k; # initial, block of coefficients for same lag.
    for i = 1:p
        Bk = Bk .- bmat[:, i_block]';
        # update the block index.
        i_block = i_block .+ k;
    end
    Bk = inv(Bk);

    # Long run impact matrix
    ψ0 = Ik * 𝝨_ols * Bk;
    #return Ik, Bk, ψ0
    ψ0 = chol(ψ0);

    # Contemporaneous shock identificaiton.
    ρ = Ik\ψ0';
    return ρ
end

function func_IRFvar_LR(data, p, h)
    # Implementation of Blanchard and Quah (1989)
    T, k = size(data);

    𝚩_ols, 𝞄_ols, 𝝨_ols = func_VAR(data, p);
    V_ols = 𝚩_ols[:, 1];
    𝚩_ols = 𝚩_ols[:, 2:size(𝚩_ols,2)];

    𝝖 = [𝚩_ols; [kron(eye(k), eye(p - 1)) zeros(k*(p-1), k)]];
    J = [eye(k) zeros(k, k*(p-1))];

    # For contemporaneous impact matrix - long run identificaiton
    # Contemporaneous shock identificaiton.
    ρ = func_contemp(𝚩_ols, 𝝨_ols, k, p);

    # Compute impulse response
    ψ = zeros(k, k, (h + 1));

    for i_irf = 0:1:h
        ψ[:,:, (i_irf + 1)] = J * 𝝖^(i_irf) * J' * ρ;
        # If needed to be standardized, multiply ρ
    end

    # Include recursive-design wild bootstrap calculation of confidence band
    #   Follow Kilian (2009, AER)
    #srand(8301993);
    n_repli = 2000; # Number of iterations

    ψ_boot = zeros(k, k, (h + 1), n_repli);

    y = data';
    Y = y[:, p:T];
    for i = 1:(p-1)
        Y = [Y; y[:, (p-i):(T-i)]];
    end

    Ur = zeros(k*p, T-p+1);
    Yr = zeros(k*p, T-p+1);

    @showprogress 1 "Wild-bootstrap confidence band " for i_iter = 1:n_repli
        pos = Int64(round(rand(1)[]*(T-p+1)));
        if pos == 0 #randomization hits zero, default to 1
            pos = 1
        end
        Yr[:, 1] = Y[:, pos];

        # Recrusive design bootstrap
        η = randn(1, size(𝞄_ols, 1));
        η = kron(η, ones(k, 1));

        Ur[1:k, 2:(T-p+1)] = 𝞄_ols' .* η;

        for ii_iter = 2:(T-p+1)
            Yr[:, ii_iter] = [V_ols; zeros(k*(p-1),1)] +
                                    𝝖 * Yr[:, (ii_iter - 1)] +
                                    Ur[:, ii_iter];
        end

        yr = Yr[1:k, :];
        for i=2:p
    		yr = [Yr[(((i-1)*k)+1):(i*k), 1] yr];
        end;
        yr = yr';

        # Iterated partial equilibrium estimation.;
        𝚩, 𝞄, 𝝨 = func_VAR(yr, p);

        𝚩 = 𝚩[:, 2:size(𝚩,2)];
        # For standardization of the impulse response,
        # For contemporaneous impact matrix - long run identificaiton
        # Contemporaneous shock identificaiton.
        ρ = func_contemp(𝚩, 𝝨, k, p);

        𝝖 = [𝚩; [kron(eye(k), eye(p - 1)) zeros(k*(p-1), k)]];
        J = [eye(k) zeros(k, k*(p-1))];
        # Computaiton of impulse response
        temp_ψ_prev = [];
        for i_irf = 0:1:h
            temp_ψ = J * 𝝖^(i_irf) * J' * ρ;
            ψ_boot[:,:, (i_irf + 1), i_iter] = temp_ψ ;
            # If needed to be standardized, multiply ρ
        end
    end
    # Need to calculate percentile
    α = 1; # One standard deviaiton confidence band of bootstrap iterations
    # Thus, this measure of confidence band does not carry literal definition
    #   as in classical model.

    # Calculate the standard deviation, across differen iteration (4th dim)
    σ_ψ_boot = std(ψ_boot, dims = 4);

    ψ_ub_1sd = ψ + σ_ψ_boot
    ψ_lb_1sd = ψ - σ_ψ_boot
    ψ_ub_2sd = ψ + σ_ψ_boot * 2;
    ψ_lb_2sd = ψ - σ_ψ_boot * 2;

    return ψ,
           ψ_lb_1sd, ψ_ub_1sd, ψ_lb_2sd, ψ_ub_2sd
end


module Regularization

using Flux.Losses: logitcrossentropy
using Statistics

"""
	spectral_decoupling(ŷ, y, λ)

Logitcrossentropy extended with spectral decoupling regularization as proposed by
[Pezeshki et al.](https://arxiv.org/abs/2011.09468). The implementation here somewhat differs from the formula presented
in the paper but the methodology is consistent with the python code provided by the authors.
"""
function spectral_decoupling(ŷ, y, λ)
	return logitcrossentropy(ŷ, y) + λ * mean(ŷ .^ 2)
end

end


# approximation algorithm from ``Algorithm 917: Complex Double-Precision
# Evaluation of the Wright ω Function'' by Piers W. Laurence, Robert M. Corless
# and David J. Jeffrey
# Series approximations can be found in the book chapter ``The Wright ω Function''
# by Corless, R. M. and Jeffrey, D. J.
#
# see also https://github.com/emsr/tr29124_test

module WrightOmega

export wrightomega

# TODO handle infs
# TODO don't use two methods
# FIXME test failures around zero and regularization

function wrightomega(z::T) where {T <: Real}
    if z < -2
        t = exp(z)
        w = t * (1 + t * (-1 + t * (T(3) / 2 + t * (-T(8) / 3 + T(125) / 24 * t))))
    elseif z < 1 + T(π)
        z1 = z - 1
        w = 1 + z1 / 2 * (1 + z1 / 8 * (1 + z1 / 12 * (-1 + z1 / 16 *
            (-1 + z1 * T(13) / 20))))
    else
        l = log(z)
        w = z + l * (-1 + (1 + (l / 2 - 1 + (l * (l / 3 - T(3) / 2) + 1) / z) / z) / z)
    end
    # refinement steps
    for _ in 1:5
        r = z - w - log(w)
        w1 = w + 1
        t = w1 * (w1 + T(2) / 3 * r)
        w *= 1 + r / w1 * (t - r / 2) / (t - r)
        fscn = abs(r^4 * (2 * w * (w - 4) - 1))
        if fscn < eps(float(T)) * 72 * w1^6
            break
        end
    end
    return w::float(T)
end

function wrightomega(z::Complex{T}) where {T <: Real}
    x = real(z)
    y = imag(z)
    Tπ = float(T)(π)
    if x == -1 && abs(y) == Tπ
        return zero(Complex{T})
    end
    # region 1
    if -2 < x <= 1 && 1 < y < 2Tπ
        imp = im * conj(sqrt(2 * conj(z + 1 - im * Tπ)))
        w = -1 + imp * (1 + imp / 3 * (-1 + imp / 6 * (inv(T(2)) + imp / 15 *
            (1 + imp / 16))))
    # region 2
    elseif -2 < x <= 1 && -2Tπ < y < -1
        imp = im * conj(sqrt(2 * conj(z + 1 - im * Tπ)))
        w = -1 + imp * (-1 + imp / 3 * (-1 + imp / 6 * (-inv(T(2)) + imp / 15 *
            (1 - imp / 16))))
    # region 3
    elseif x < -2 && -Tπ < y <= Tπ
        t = exp(z)
        w = t * (1 + t * (-1 + t * (T(3) / 2 + t * (-T(8) / 3 + T(125) / 24 * t))))
    # region 4
    elseif (-2 < x <= 1 && -1 <= y <= 1) || (-1 < x && (x - 1)^2 + y^2 <= Tπ^2)
        z1 = z - 1
        w = 1 + z1 / 2 * (1 + z1 / 8 * (1 + z1 / 12 * (-1 + z1 / 16 *
            (-1 + z1 * T(13) / 20))))
    # region 5
    elseif x <= -2 && Tπ < y && y - Tπ <= -3 / T(4) * (x + 1)
        t = z - im * Tπ
        l = log(-t)
        w = t - l + l * (1 + (l / 2 - 1 + (l^2 / 3 - 3 * l / 2 + 1) / t) / t) / t
    # region 6
    elseif x <= -2 && y <= -Tπ && y + Tπ >= T(3) / 4 * (x + 1)
        t = z + im * Tπ
        l = log(-t)
        w = t - l + l * (1 + (l / 2 - 1 + (l^2 / 3 - 3 * l / 2 + 1) / t) / t) / t
    # region 7
    else
        l = log(z)
        w = z + l * (-1 + (1 + (l / 2 - 1 + (l * (l / 3 - T(3) / 2) + 1) / z) / z) / z)
    end
    if x <= -0.99 && (abs(y - Tπ) <= 0.01 || abs(y + Tπ) <= 0.01)
        s = -1
        if abs(y - Tπ) <= 0.01
            # TODO setrounding doesn't work for non-BigFloat, decide what to do
            # T == BigFloat && setrounding(T, RoundUp)
            z = x + im * (y - Tπ)
        else
            z = x + im * (y + Tπ)
        end
    else
        s = 1
    end
    w *= s
    # refinement steps
    for _ in 1:5
        r = z - s * w - log(Complex(w))
        w1 = s * w + 1
        t = w1 * (w1 + T(2) / 3 * r)
        w *= 1 + r / w1 * (t - r / 2) / (t - r)
        fscn = abs(r^4 * (2 * w * (w - 4) - 1))
        if fscn < eps(float(T)) * 72 * abs(w1)^6
            break
        end
    end
    w *= s
    return w::Complex{float(T)}
end

end


using Tensors: ∇, ∇∇, Δ
import ForwardDiff

function Ψ(C, μ, Kb)
    detC = det(C)
    J = sqrt(detC)
    Ĉ = detC^(-1/3)*C
    return 1/2*(μ * (tr(Ĉ)- 3) + Kb*(J-1)^2)
end

function S(C, μ, Kb)
    I = one(C)
    J = sqrt(det(C))
    invC = inv(C)
    return μ * det(C)^(-1/3)*(I - 1/3*tr(C)*invC) + Kb*(J-1)*J*invC
end

const μ = 1e10;
const Kb = 1.66e11;
Ψ(C) = Ψ(C, μ, Kb)
S(C) = S(C, μ, Kb)

@testsection "AD" begin
    for dim in 1:3
        @testsection "dim $dim" begin
        F = one(Tensor{2,dim}) + rand(Tensor{2,dim});
        C = tdot(F);
        C2 = F' ⋅ F;

        @test 2(@inferred ∇(Ψ, C))::typeof(C) ≈ S(C)
        @test 2(@inferred ∇(Ψ, C2))::typeof(C2) ≈ S(C2)

        b = rand(SymmetricTensor{2, dim})
        @test 2 * (@inferred ∇∇(Ψ, C))::SymmetricTensor{4, dim} ⊡ b ≈ ∇(S, C) ⊡ b
        @test 2 * (@inferred ∇∇(Ψ, C2))::Tensor{4, dim} ⊡ b ≈ ∇(S, C2) ⊡ b

        @test ∇(Ψ, C) ≈ ∇(Ψ, C2)
        @test ∇(S, C) ⊡ b ≈ ∇(S, C2) ⊡ b
        @test ∇∇(Ψ, C) ⊡ b ≈ ∇∇(Ψ, C2) ⊡ b

        for T in (Float32, Float64)
            Random.seed!(1234) # needed for getting "good" tensors for calculating det and friends
            A = rand(Tensor{2, dim, T})
            B = rand(Tensor{2, dim, T})
            A_sym = rand(SymmetricTensor{2, dim, T})
            B_sym = rand(SymmetricTensor{2, dim, T})
            v = rand(Vec{dim, T})
            II = one(Tensor{4, dim, T})
            II_sym = one(SymmetricTensor{4, dim, T})

            # Gradient of scalars
            @testsection "scalar grad" begin
                @test (@inferred ∇(norm, v))::typeof(v) ≈ ((@inferred ∇(norm, v, :all))[1])::typeof(v) ≈ v / norm(v)
                @test (∇(norm, v, :all)[2])::T ≈ norm(v)
                @test (@inferred ∇(norm, A))::typeof(A) ≈ ((@inferred ∇(norm, A, :all))[1])::typeof(A) ≈ A / norm(A)
                @test (∇(norm, A, :all)[2])::T ≈ norm(A)
                @test (@inferred ∇(norm, A_sym))::typeof(A_sym) ≈ ((@inferred ∇(norm, A_sym, :all))[1])::typeof(A_sym) ≈ A_sym / norm(A_sym)
                @test (∇(norm, A_sym, :all)[2])::T ≈ norm(A_sym)

                @test ∇(v -> 3*v, v) ≈ ∇(v -> 3*v, v, :all)[1] ≈ diagm(Tensor{2, dim}, 3.0)
                @test ∇(v -> 3*v, v, :all)[2] ≈ 3*v
                # function does not return dual
                @test ∇(A -> T(1), A)::typeof(A) ≈ ∇(A -> T(1), A, :all)[1] ≈ 0*A
                @test ∇(A -> T(1), A, :all)[2] == 1
                @test ∇(A -> T(1), A_sym)::typeof(A_sym) ≈ ∇(A -> T(1), A, :all)[1] ≈ 0*A_sym
                @test ∇(A -> T(1), A_sym, :all)[2] == 1
            end

            @testsection "2nd tensor grad" begin
            # Gradient of second order tensors
            # https://en.wikipedia.org/wiki/Tensor_derivative_(continuum_mechanics)#Derivatives_of_the_invariants_of_a_second-order_tensor
                I1, DI1 = A -> tr(A), A -> one(A)
                I2, DI2 = A -> (tr(A)^2 - tr(A⋅A)) / 2, A -> I1(A) * one(A) - A'
                I3, DI3 = A -> det(A), A -> det(A) * inv(A)'

                @test (@inferred ∇(I1, A))::typeof(A) ≈ ((@inferred ∇(I1, A, :all))[1])::typeof(A) ≈ DI1(A)
                @test (∇(I1, A, :all)[2])::T ≈ I1(A)
                @test (@inferred ∇(I2, A))::typeof(A) ≈ ((@inferred ∇(I2, A, :all))[1])::typeof(A) ≈ DI2(A)
                @test (∇(I2, A, :all)[2])::T ≈ I2(A)
                @test (@inferred ∇(I3, A))::typeof(A) ≈ ((@inferred ∇(I3, A, :all))[1])::typeof(A) ≈ DI3(A)
                @test (∇(I3, A, :all)[2])::T ≈ I3(A)
                @test (@inferred ∇(I1, A_sym))::typeof(A_sym) ≈ ((@inferred ∇(I1, A_sym, :all))[1])::typeof(A_sym) ≈ DI1(A_sym)
                @test (∇(I1, A_sym, :all)[2])::T ≈ I1(A_sym)
                @test (@inferred ∇(I2, A_sym))::typeof(A_sym) ≈ ((@inferred ∇(I2, A_sym, :all))[1])::typeof(A_sym) ≈ DI2(A_sym)
                @test (∇(I2, A_sym, :all)[2])::T ≈ I2(A_sym)
                @test (@inferred ∇(I3, A_sym))::typeof(A_sym) ≈ ((@inferred ∇(I3, A_sym, :all))[1])::typeof(A_sym) ≈ DI3(A_sym)
                @test (∇(I3, A_sym, :all)[2])::T ≈ I3(A_sym)

                @test (@inferred ∇(identity, A))::Tensor{4, dim, T} ≈ ((@inferred ∇(identity, A, :all))[1])::Tensor{4, dim, T} ≈ II
                @test (∇(identity, A, :all)[2])::typeof(A) ≈ A
                @test (@inferred ∇(identity, A_sym))::SymmetricTensor{4, dim, T} ≈ ((@inferred ∇(identity, A_sym, :all))[1])::SymmetricTensor{4, dim, T} ≈ II_sym
                @test (∇(identity, A_sym, :all)[2])::typeof(A_sym) ≈ A_sym
                @test (@inferred ∇(transpose, A))::Tensor{4, dim, T} ⊡ B ≈ ((@inferred ∇(transpose, A, :all))[1])::Tensor{4, dim, T} ⊡ B ≈ B'
                @test (∇(transpose, A, :all)[2])::typeof(A) ≈ A'
                @test (@inferred ∇(transpose, A_sym))::SymmetricTensor{4, dim, T} ⊡ B_sym ≈ ((@inferred ∇(transpose, A_sym, :all))[1])::SymmetricTensor{4, dim, T} ⊡ B_sym ≈ B_sym'
                @test (∇(transpose, A_sym, :all)[2])::typeof(A_sym) ≈ A_sym'
                @test (@inferred ∇(inv, A))::Tensor{4, dim, T} ⊡ B ≈ ((@inferred ∇(inv, A, :all))[1])::Tensor{4, dim, T} ⊡ B ≈ - inv(A) ⋅ B ⋅ inv(A)
                @test (∇(inv, A, :all)[2])::typeof(A) ≈ inv(A)
                @test (@inferred ∇(inv, A_sym))::SymmetricTensor{4, dim, T} ⊡ B_sym ≈ ((@inferred ∇(inv, A_sym, :all))[1])::SymmetricTensor{4, dim, T} ⊡ B_sym ≈ - inv(A_sym) ⋅ B_sym ⋅ inv(A_sym)
                @test (∇(inv, A_sym, :all)[2])::typeof(A_sym) ≈ inv(A_sym)
                # function does not return dual
                @test ∇(A -> B, A) ≈ ∇(A -> B, A, :all)[1] ≈ zero(Tensor{4, dim, T})
                @test ∇(A -> B, A, :all)[2] ≈ B
                @test isa(∇(A -> B, A), Tensor{4, dim, T})
            end

            # Hessians of scalars
            @testsection "hessian" begin
                @test (@inferred ∇∇(norm, A))::Tensor{4, dim, T} ≈ ((@inferred ∇∇(norm, A, :all))[1])::Tensor{4, dim, T} ≈ reshape(ForwardDiff.hessian(x -> sqrt(sum(abs2, x)), A), (dim, dim, dim, dim))
                @test (∇∇(norm, A, :all)[2])::typeof(A) ≈ reshape(ForwardDiff.gradient(x -> sqrt(sum(abs2, x)), A), (dim, dim))
                @test (∇∇(norm, A, :all)[3])::T ≈ norm(A)
                @test (@inferred ∇∇(norm, A_sym))::SymmetricTensor{4, dim, T} ≈ ((@inferred ∇∇(norm, A_sym, :all))[1])::SymmetricTensor{4, dim, T}
                @test (∇∇(norm, A_sym, :all)[2])::typeof(A_sym) ≈ ∇(norm, A_sym)
                @test (∇∇(norm, A_sym, :all)[3])::T ≈ norm(A_sym)
                # function does not return dual
                @test ∇∇(A -> T(1), A)::Tensor{4, dim, T} ≈ ∇∇(A -> T(1), A, :all)[1] ≈ 0*II
                @test ∇∇(A -> T(1), A, :all)[2] ≈ 0*A
                @test ∇∇(A -> T(1), A, :all)[3] == T(1)
            end
            end # loop T
        end # testsection
    end # loop dim

    @testsection "vector calculus identities" begin
        Random.seed!(1234)
        φ(x) = norm(x)^4
        ϕ(x) = sum(x)
        A(x) = Vec{3}((x[1]*x[2]^3*x[3], x[1]*x[2]*x[3]^3, x[1]^3*x[2]*x[3]))
        B(x) = Vec{3}((x[1]*x[1], x[1]*x[2], x[1]*x[3]))
        for T in (Float32, Float64)
            x = rand(Vec{3, T})
            # gradient
            @test gradient(x -> φ(x) + ϕ(x), x) ≈ gradient(φ, x) + gradient(ϕ, x)
            @test gradient(x -> φ(x) * ϕ(x), x) ≈ φ(x) * gradient(ϕ, x) + gradient(φ, x) * ϕ(x)
            @test gradient(x -> A(x) ⋅ B(x), x) ≈ gradient(B, x)⋅A(x) + gradient(A, x)⋅B(x) + A(x)×curl(B, x) + B(x)×curl(A, x)
            # divergence
            @test divergence(x -> A(x) + B(x), x) ≈ divergence(A, x) + divergence(B, x)
            @test divergence(x -> φ(x) * A(x), x) ≈ φ(x)*divergence(A, x) + gradient(φ, x)⋅A(x)
            @test divergence(x -> A(x) × B(x), x) ≈ B(x)⋅curl(A, x) - A(x)⋅curl(B, x)
            # curl
            @test curl(x -> A(x) + B(x), x) ≈ curl(A, x) + curl(B, x)
            @test curl(x -> φ(x) * A(x), x) ≈ φ(x)*curl(A, x) + gradient(φ, x)×A(x)
            @test curl(x -> A(x) × B(x), x) ≈ A(x)*divergence(B, x) - B(x)*divergence(A, x) + gradient(A, x)⋅B(x) - gradient(B, x)⋅A(x)
            # second derivatives
            @test divergence(x -> curl(A, x), x) ≈ 0 atol = eps(T)
            @test curl(x -> gradient(φ, x), x) ≈ zero(Vec{3}) atol = 10eps(T)
            @test divergence(x -> gradient(φ, x), x) ≈ laplace(φ, x)
            @test gradient(x -> divergence(A, x), x) ≈ curl(x -> curl(A, x), x) + laplace.(A, x)
            @test divergence(x -> ϕ(x)*gradient(φ, x), x) ≈ ϕ(x)*laplace(φ, x) + gradient(ϕ, x)⋅gradient(φ, x)
            @test divergence(x -> φ(x)*gradient(ϕ, x) - gradient(φ, x)*ϕ(x), x) ≈ φ(x)*laplace(ϕ, x) - laplace(φ, x)*ϕ(x)
            @test laplace(x -> ϕ(x)*φ(x), x) ≈ laplace(ϕ, x)*φ(x) + 2*gradient(ϕ, x)⋅gradient(φ, x) + ϕ(x)*laplace(φ, x)
            @test laplace.(x -> φ(x)*A(x), x) ≈ A(x)*laplace(φ, x) + 2*(gradient(A, x)⋅gradient(φ, x)) + φ(x)*laplace.(A, x)
            @test laplace(x -> A(x)⋅B(x), x) ≈ A(x)⋅laplace.(B, x) - B(x)⋅laplace.(A, x) + 2*divergence(x -> gradient(A, x)⋅B(x) + B(x)×curl(A, x), x)
            # third derivatives
            @test laplace.(x -> gradient(φ, x), x) ≈ gradient(x -> divergence(x -> gradient(φ, x), x), x) ≈ gradient(x -> laplace(φ, x), x)
            @test laplace(x -> divergence(A, x), x) ≈ divergence(x -> gradient(x->divergence(A, x), x), x) ≈ divergence(x -> laplace.(A, x), x)
            @test laplace.(x -> curl(A, x), x) ≈ -curl(x -> curl(x -> curl(A, x), x), x) ≈ curl(x -> laplace.(A, x), x)
        end # loop T
    end # testsection

    @testsection "f: scalar -> scalar" begin
        f(x) = 2 * x^3
        x = 3.0
        @test gradient(f, x) ≈ gradient(f, x, :all)[1] ≈ 6 * x^2
        @test gradient(f, x, :all)[2] ≈ f(x) ≈ 2 * x^3
        @test hessian(f, x) ≈ hessian(f, x, :all)[1] ≈ 12 * x
        @test hessian(f, x, :all)[2] ≈ gradient(f, x) ≈ 6 * x^2
        @test hessian(f, x, :all)[3] ≈ gradient(f, x, :all)[2] ≈ 2 * x^3
    end
    @testsection "f: scalar -> AbstractTensor" begin
        for dim in 1:3, S in (rand(Vec{dim}), rand(Tensor{2,dim}), rand(SymmetricTensor{2,dim}))
            f(x) = x^3 * S
            x = 3.0
            @test gradient(f, x) ≈ gradient(f, x, :all)[1] ≈ 3 * x^2 * S
            @test gradient(f, x, :all)[2] ≈ f(x) ≈ x^3 * S
            @test hessian(f, x) ≈ hessian(f, x, :all)[1] ≈ 6 * x * S
            @test hessian(f, x, :all)[2] ≈ gradient(f, x) ≈ 3 * x^2 * S
            @test hessian(f, x, :all)[3] ≈ f(x) ≈ x^3 * S
        end
    end
    @testsection "mixed scalar/vec" begin
        f(v::Vec, s::Number) = s * v[1] * v[2]
        v = Vec((2.0, 3.0))
        s = 4.0
        @test gradient(x -> f(v, x), s)::Float64 ≈ v[1] * v[2]
        @test gradient(x -> f(x, s), v)::Vec{2} ≈ Vec((s * v[2], s * v[1]))
        @test gradient(y -> gradient(x -> f(y, x), s), v)::Vec{2} ≈ Vec((v[2], v[1]))
        @test gradient(y -> gradient(x -> f(x, y), v), s)::Vec{2} ≈ Vec((v[2], v[1]))
    end
end # testsection


"""
    S(r, t, laser)

Compute the [Poynting vector](https://en.wikipedia.org/wiki/Poynting_vector) defined as
```math
    S(r, t) = \\frac{1}{μ₀} \\vec{E} × \\vec{B}
```
"""
function S(r, t, laser)
    ElectricField, MagneticField = EB(r, t, laser)
    μ₀ = fundamental_constants(laser, :μ₀)

    1 / μ₀ * (ElectricField × MagneticField)
end


mutable struct QMRCGSTABIterable{
    opT, solT, realT<:Real, scalarT<:Number}

    A::opT
    x::solT

    rt::solT
    r::solT
    p::solT
    v::solT
    s::solT
    d::solT
    dt::solT
    xt::solT
    t::solT

    rho::scalarT
    rho_prev::scalarT
    alpha::scalarT
    omega::scalarT
    tau::scalarT
    theta::scalarT
    eta::scalarT
    beta::scalarT
    thetat::scalarT
    etat::scalarT

    maxiter::Int
    niter::Int
    tol::realT
    residual::realT
end

function init_qmrcgstab(
    x, A, b;
    abstol::Real=zero(real(eltype(b))),
    reltol::Real=sqrt(eps(real(eltype(b)))),
    maxiter::Int=2000)

    T = typeof(x)
    N = size(x)
    niter = 0

    rt = T(undef, N)
    r = T(undef, N)
    p = T(undef, N)
    v = T(undef, N)
    s = T(undef, N)
    d = T(undef, N)
    dt = T(undef, N)
    xt = T(undef, N)
    t = T(undef, N)

    fill!(p, zero(eltype(x)))
    fill!(v, zero(eltype(x)))
    fill!(d, zero(eltype(x)))

    mul!(r, A, x)
    sub!(r, b, r)
    niter += 1

    copyto!(rt, r) # rand!(rt)
    nrm = norm(r)
    tolerance = max(reltol * nrm, abstol)

    rho = one(eltype(x))
    rho_prev = one(eltype(x))
    alpha = one(eltype(x))
    omega = one(eltype(x))
    tau = nrm * one(eltype(x))
    theta = zero(eltype(x))
    eta = zero(eltype(x))
    beta = zero(eltype(x))
    thetat = zero(eltype(x))
    etat = zero(eltype(x))

    QMRCGSTABIterable(
        A, x,
        rt, r, p, v, s, d, dt, xt, t,
        rho, rho_prev, alpha, omega, tau, theta, eta,
        beta, thetat, etat,
        maxiter, niter, tolerance, nrm)
end

@inline converged(it::QMRCGSTABIterable) = it.residual ≤ it.tol
@inline start(::QMRCGSTABIterable) = 0
@inline done(it::QMRCGSTABIterable, iteration::Int) = it.niter ≥ it.maxiter || converged(it)

function iterate(it::QMRCGSTABIterable, iteration::Int=start(it))
    if done(it, iteration) return nothing end

    it.rho_prev = it.rho
    it.rho = dot(it.rt, it.r)

    it.beta = it.rho * it.alpha / (it.rho_prev * it.omega)
    sub_scmul!(it.p, it.omega, it.v)
    mul!(it.p, it.beta, it.p)
    add!(it.p, it.r, it.p)

    mul!(it.v, it.A, it.p)
    it.alpha = it.rho / dot(it.rt, it.v)
    mul!(it.s, it.alpha, it.v)
    sub!(it.s, it.r, it.s)

    # First quasi-minimization and update
    it.thetat = norm(it.s) / it.tau
    c = one(it.thetat) / sqrt(one(it.thetat) + it.thetat*it.thetat) ####
    taut = it.tau * it.thetat * c ####

    it.etat = c*c * it.alpha
    mul!(it.dt, (it.theta * it.theta * it.eta / it.alpha), it.d)
    add!(it.dt, it.p, it.dt)
    mul!(it.xt, it.etat, it.dt)
    add!(it.xt, it.x, it.xt)

    # Update r
    mul!(it.t, it.A, it.s)
    it.omega = (it.thetat * it.tau)^2 / dot(it.s, it.t) # dot(s, t) / dot(t, t)
    mul!(it.r, it.omega, it.t)
    sub!(it.r, it.s, it.r)

    # Second quasi-minimization and update
    it.theta = norm(it.r) / taut
    c = one(it.theta) / sqrt(one(it.theta) + it.theta*it.theta)
    it.tau = taut * it.theta * c

    it.eta = c*c * it.omega
    mul!(it.d, (it.thetat * it.thetat * it.etat / it.omega), it.dt)
    add!(it.d, it.s, it.d)
    mul!(it.x, it.eta, it.d)
    add!(it.x, it.xt, it.x)

    it.niter += 2
    it.residual = sqrt(iteration+1) * abs(it.tau)

    it.residual, iteration + 1
end

"""
- Saad, Y. (2003). Interactive method for sparse linear system.
- Freund, W. R., & Nachtigal, N. M. (1990). QMR : for a Quasi-Minimal
  Residual Linear Method Systems. (December).
- New implementation of QMR-type algorithms (2005)
- A Quasi-Minimal Residual Variant of the Bi-CGSTAB Algorithm for Nonsymmetric Systems
"""
function qmrcgstab!(
    x, A, b;
    abstol::Real=zero(real(eltype(b))),
    reltol::Real=sqrt(eps(real(eltype(b)))),
    maxiter::Int=2000,
    verbose::Bool = false)

    history = zeros(0)

    iterable = init_qmrcgstab(
        x, A, b;
        abstol=abstol,
        reltol=reltol,
        maxiter=maxiter)

    for (iteration, item) = enumerate(iterable)
        push!(history, iterable.residual)
        verbose && @printf("%3d\t%1.2e\n", iteration, iterable.residual)
    end

    if !converged(iterable)
        push!(history, -1)
        verbose && println("QMRCGStab: not converged...")
    end
    # println()

    return history
end


"""
    ExactReach

ExactReach performs exact reachability analysis to compute the output reachable set for a network.

# Problem requirement
1. Network: any depth, ReLU activation
2. Input: HPolytope
3. Output: AbstractPolytope

# Return
`ReachabilityResult`

# Method
Exact reachability analysis.

# Property
Sound and complete.

# Reference
[W. Xiang, H.-D. Tran, and T. T. Johnson,
"Reachable Set Computation and Safety Verification for Neural Networks with ReLU Activations,"
*ArXiv Preprint ArXiv:1712.08163*, 2017.](https://arxiv.org/abs/1712.08163)
"""
struct ExactReach <: Solver end

function solve(solver::ExactReach, problem::Problem)
    reach = forward_network(solver, problem.network, problem.input)
    return check_inclusion(reach, problem.output)
end

# Forward_linear that coverts the input to HPolytope
forward_linear(solver::ExactReach, L::Layer, input) = forward_linear(solver, L, convert(HPolytope, input))

# Forward linear for a single polytope and a list of polytopes
forward_linear(solver::ExactReach, L::Layer, input::HPolytope) = affine_map(L, input)
forward_linear(solver::ExactReach, L::Layer, input::Vector{<:HPolytope}) = [affine_map(L, set) for set in input]

# Forward act for a single polytope and a list of polytopes
forward_act(solver::ExactReach, L::Layer, input::HPolytope) = forward_partition(L.activation, input)
function forward_act(solver::ExactReach, L::Layer, input::Vector{<:HPolytope})
    output = Vector{HPolytope}(undef, 0)
    for i in 1:length(input)
        append!(output, forward_partition(L.activation, input[i]))
    end
    return output
end


function mul_analytic!(τ::Array{R,1}, dx1::R, ie1::Int, xs1::Int) where {R <: AbstractFloat}
    for i1 = 1:ie1
        @inbounds τ[i1] *= sqrt(sum(((i1-xs1)*dx1)^2))
    end
end

function mul_analytic!(τ::Array{R,2}, dx1::R, dx2::R, ie1::Int, ie2::Int, xs1::Int, xs2::Int) where {R <: AbstractFloat}
    for i2 = 1:ie2
        for i1 = 1:ie1
            @inbounds τ[i1,i2] *= sqrt(sum(((i1-xs1)*dx1)^2+((i2-xs2)*dx2)^2))
        end
    end
end

function mul_analytic!(τ::Array{R,3}, dx1::R, dx2::R, dx3::R, ie1::Int, ie2::Int, ie3::Int, xs1::Int, xs2::Int, xs3::Int) where {R <: AbstractFloat}
    for i3 = 1:ie3
        for i2 = 1:ie2
            for i1 = 1:ie1
                @inbounds τ[i1,i2,i3] *= sqrt(sum(((i1-xs1)*dx1)^2+((i2-xs2)*dx2)^2+((i3-xs3)*dx3)^2))
            end
        end
    end
end

function grad_analytic(τ0::Array{R,1}, dx1::R, ie1::Int, xs1::Int) where {R <: AbstractFloat}
    ∇τ0 = [ones(size(τ0))]
    for i1 = 1:ie1
        @inbounds ∇τ0[1][i1] = dx1*(i1-xs1)/τ0[i1]
    end
    #set terms at source
    ∇τ0[1][xs1] = one(R)
    ∇τ0
end

function grad_analytic(τ0::Array{R,2}, dx1::R, dx2::R, ie1::Int, ie2::Int, xs1::Int, xs2::Int) where {R <: AbstractFloat}
    ∇τ0 = [ones(size(τ0)), ones(size(τ0))]
    for i2 = 1:ie2
        for i1 = 1:ie1
            @inbounds ∇τ0[1][i1,i2] = dx1*(i1-xs1)/τ0[i1,i2]
            @inbounds ∇τ0[2][i1,i2] = dx2*(i2-xs2)/τ0[i1,i2]
        end
    end
    ∇τ0[1][xs1, xs2] = convert(R, 1/sqrt(2))
    ∇τ0[2][xs1, xs2] = convert(R, 1/sqrt(2))
    ∇τ0
end

function grad_analytic(τ0::Array{R,3}, dx1::R, dx2::R, dx3::R, ie1::Int, ie2::Int, ie3::Int, xs1::Int, xs2::Int, xs3::Int) where {R <: AbstractFloat}
    ∇τ0 = [ones(size(τ0)), ones(size(τ0)), ones(size(τ0))]
    for i3 = 1:ie3
        for i2 = 1:ie2
            for i1 = 1:ie1
                @inbounds ∇τ0[1][i1,i2,i3] = dx1*(i1-xs1)/τ0[i1,i2,i3]
                @inbounds ∇τ0[2][i1,i2,i3] = dx2*(i2-xs2)/τ0[i1,i2,i3]
                @inbounds ∇τ0[3][i1,i2,i3] = dx3*(i3-xs3)/τ0[i1,i2,i3]
            end
        end
    end
    ∇τ0[1][xs1, xs2, xs3] = convert(R, 1/sqrt(3))
    ∇τ0[2][xs1, xs2, xs3] = convert(R, 1/sqrt(3))
    ∇τ0[3][xs1, xs2, xs3] = convert(R, 1/sqrt(3))
    ∇τ0
end


    using LowRankIntegrators, LaTeXStrings, GLMakie
    # Reaction Network
    # ∅ -> A -> ∅, [k1/(1+(B/θ)^3), k2*A]
    # ∅ -> B -> ∅, [k3/(1+A), k4*B]

    # computing the propensities
    function props(k1, k2, k3, k4, θ, N)
        ax = [zeros(N) for i in 1:4]
        ay = [zeros(N) for i in 1:4]
        for x in 1:N
            ax[1][x] = 1
            ax[2][x] = k2*(x-1)
            ax[3][x] = k3/(1+(x-1))
            ax[4][x] = 1
        end
        for y in 1:N
            ay[1][y] = k1/(1+((y-1)/θ)^3)
            ay[2][y] = 1
            ay[3][y] = 1
            ay[4][y] = k4*(y-1)
        end
        return ax, ay
    end

    # shift operators:
    # let A be an n×n matrix. Then,
    # shift(x,n)*A shifts rows by x rows down and padds the first x rows with 0s
    # A*shift(x,n)' shifts cols by x cols right and padds the x leftmost cols with 0s
    function shift(x,n)
        S = zeros(n,n)
        if x > 0
            S[1+x:end,1:end-x] = Matrix(I,n-x, n-x)
        end
        if x < 0
            S[1:end+x,1-x:end] = Matrix(I,n+x,n+x)
        end
        if x == 0
            S = I
        end
        return S
    end

    ## simulation parameters
    const N = 600
    const Nrxn = 4
    const ν = [(1,0), (-1,0), (0,1), (0,-1)]
    k1 = 30.0
    k2 = 1
    k3 = 10
    k4 = 1
    θ = 1

    ax, ay = props(k1,k2,k3,k4,θ,N) # propensity factors
    Srows = [shift(ν[r][1],N) for r in 1:Nrxn] # shift operators
    Scols = [shift(ν[r][2],N)' for r in 1:Nrxn] # shift operators
    A = [Srows[r]*TwoFactorRepresentation(ax[r],ay[r])*Scols[r] for r in 1:Nrxn] # coefficient matrices
    Asum = TwoFactorRepresentation(hcat(ax...), hcat(ay...)) # coefficient matrix

    # chemical master equation
    function cme(P, (A,Asum,Srows,Scols), t)
        return sum(A[r] .* (Srows[r]*P*Scols[r]) for r in 1:Nrxn) - Asum .* P
    end

    # Gaussian initial condition
    x0 = 20
    y0 = 20
    Q = qr(randn(2,2)).Q
    D = Q*Diagonal([rand()/20, rand()/20])*Q'
    P0(x,y) = exp(-[x-x0; y-y0]'*D*[x-x0; y-y0])
    P_init = [P0(x,y) for x in 1:N, y in 1:N]
    P_init ./= sum(P_init)
    Tf = 10.0

    dt = 2e-2
    prob = ODEProblem(cme, P_init, (0, Tf), (Matrix.(A), Matrix(Asum), Srows, Scols))
    t_direct = @elapsed sol = DifferentialEquations.solve(prob, Tsit5(), saveat = dt)

    P_lr = truncated_svd(P_init, tol = 1e-6)
    @show rank(P_lr)
    lr_prob = MatrixDEProblem((x,t) -> cme(x, (A, Asum, Srows, Scols), t), P_lr, (0.0, 10.0))
    t_lr = @elapsed lr_sol = LowRankIntegrators.solve(lr_prob, UnconventionalAlgorithm(), dt)

    @show t_direct/t_lr

    # comparison movie
    fig = Figure()
    ax_lra = Axis(fig[1,1], xlabel = "A", ylabel = "B", title = "full")
    ax_full = Axis(fig[1,2], xlabel = "A", ylabel = "B", title = "low rank")
    ax_error = Axis(fig[2,1:2], yscale = log10, xlabel="time", ylabel = "error")
    limits!(ax_error, 0, Tf, 1e-8, 1e-0)
    display(fig)

    N_relevant = 40
    sol_lra = Observable(Matrix(lr_sol.Y[1][1:N_relevant, 1:N_relevant]))
    sol_full = Observable(sol.u[1][1:N_relevant, 1:N_relevant])
    error = Observable([norm(sol.u[1] - Matrix(lr_sol.Y[1]))])
    ts = Observable([0.0])
    Makie.heatmap!(ax_lra, sol_lra)
    Makie.heatmap!(ax_full, sol_full)
    Makie.lines!(ax_error, ts, error, linewidth = 2, color = :black)

    record(fig, "markov_chain.mp4", enumerate(lr_sol.t); framerate = 24) do (i,t)
        sol_lra[] = Matrix(lr_sol.Y[i][1:N_relevant,1:N_relevant])
        sol_full[] = sol(t)[1:N_relevant,1:N_relevant]
        push!(error.val, maximum(abs.(Matrix(lr_sol.Y[i][1:N_relevant, 1:N_relevant]) - sol(t)[1:N_relevant, 1:N_relevant])))
        push!(ts.val, t)
        ts[] = ts[]
        error[] = error[]
    end


@testset "Rangefinder" begin
    Random.seed!(0)
    m, n = 30, 20
    l = 5
    tol = 1e-8
    B = randn(m, l) * rand(l, n)

    Q = RP.rangefinder(B, 2l; q=0)
    @test norm(B - Q*Q'*B) <= tol

    Q = RP.rangefinder(B, 2l; q=5)
    @test norm(B - Q*Q'*B) <= tol

end


@testset verbose = true "HMM" begin

    @testset "Discrete" begin
        tr = DiscreteMarkovChain([0.3, 0.7], [0.9 0.1; 0.2 0.8])
        em1 = Dists.Normal(1, 0.3)
        em2 = Dists.Normal(2, 0.3)
        hmm = HiddenMarkovModel(tr, [em1, em2])

        states, observations = rand(hmm, 1000)

        tr_init = DiscreteMarkovChain(randprobvec(2), randtransmat(2))
        em1_init = Dists.Normal(rand(), 1)
        em2_init = Dists.Normal(rand(), 1)
        hmm_init = HiddenMarkovModel(tr_init, [em1_init, em2_init])

        hmm_est, logL_evolution = baum_welch(hmm_init, observations, iterations = 100)

        error = transition_matrix(hmm_est) - transition_matrix(hmm)

        @test maximum(error) < 0.1
    end

    @testset "Continuous" begin
        tr = ContinuousMarkovChain([0.3, 0.7], [-1.0 1.0; 2.0 -2.0])
        em1 = PoissonProcess(1.0, Dists.Normal(1, 1))
        em2 = PoissonProcess(2.0, Dists.Normal(-1, 1))
        mmpp = MarkovModulatedPoissonProcess(tr, [em1, em2])

        state_history, observations = rand(mmpp, 0.0, 1000.0)
        nb_events(observations)

    end
end


using StatsPlots

function pk_qualify(pk::Vector{Float64})
    pk_good = sum(pk .<= 0.5)
    pk_ok = length(pk[pk .<= 0.7]) - pk_good
    pk_bad = length(pk[pk .<= 1]) - pk_good - pk_ok
    (good=pk_good, ok=pk_ok, bad=pk_bad, very_bad=sum(pk .> 1))
end

function pk_plot(pk::Vector{Float64}; title="PSIS diagnostic plot.",
    leg=:topleft, kwargs...)
    scatter(pk, xlab="Datapoint", ylab="Pareto shape k",
        marker=2.5, lab="Pk points", leg=leg)
  hline!([0.5], lab="pk = 0.5");hline!([0.7], lab="pk = 0.7")
  hline!([1], lab="pk = 1.0")
  title!(title)
end

export
    pk_qualify,
    pk_plot


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
Random.seed!(233)

function spatial_varying_tangent_elastic(mu,m,n,h,type=1)
    spatial_varying_tangent_elastic_ = load_op_and_grad("./build/libSpatialVaryingTangentElastic","spatial_varying_tangent_elastic")
    mu,m_,n_,h,type = convert_to_tensor([mu,m,n,h,type], [Float64,Int64,Int64,Float64,Int64])
    H = spatial_varying_tangent_elastic_(mu,m_,n_,h,type)
    set_shape(H, (4*m*n, 2, 2))
end

# TODO: specify your input parameters
m = 10
n = 5
h = 0.1
mu = rand(4*m*n)
H = zeros(4*m*n, 2, 2)

for i = 1:4*m*n
    H[i,:,:] = mu[i] * diagm(0=>ones(2))
end
type = 1
u = spatial_varying_tangent_elastic(mu,m,n,h,type)
sess = Session(); init(sess)
# @show run(sess, u)-H

m = 10
n = 5
h = 0.1
mu = [rand(4*m*n); rand(4*m*n)]
H = zeros(4*m*n, 2, 2)

for i = 1:4*m*n
    H[i,:,:] = [mu[i] 0.0; 0.0 mu[i+4*m*n]]
end
type = 2
u = spatial_varying_tangent_elastic(mu,m,n,h,type)
sess = Session(); init(sess)
@show run(sess, u)-H

m = 10
n = 5
h = 0.1
mu = rand(4*m*n*3)
H = zeros(4*m*n, 2, 2)

for i = 1:4*m*n
    H[i,:,:] = [mu[i] mu[i+8*m*n]; mu[i+8*m*n] mu[i+4*m*n]]
end
type = 3
u = spatial_varying_tangent_elastic(mu,m,n,h,type)
sess = Session(); init(sess)
@show run(sess, u)-H

# uncomment it for testing gradients
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(m_)
    return sum(spatial_varying_tangent_elastic(m_,m,n,h,type)^2)
end

# TODO: change `m_` and `v_` to appropriate values
m_ = constant(rand(4*m*n*3))
v_ = rand(4*m*n*3)
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


include("../src/pressurecorrelations.jl")
include("../src/tempcorrelations.jl")

@testset "Temperature correlations" begin

q_o = 2219
q_w = 11
APIoil = 24
sg_water = 1
GLR = 1.762 * 1000^2 / (2219 + 11)
sg_gas = 0.7
id = 2.992
whp = 150 #psig
A = Shiu_Beggs_relaxationfactor(q_o, q_w, GLR, APIoil, sg_water, sg_gas, id, whp)

@test A ≈ 2089 atol = 1

g_g = 0.0106 * 100
bht = 173
z = 5985 - 0

@test Ramey_temp(z, bht, A, g_g) ≈ 130 atol = 2
@test Ramey_temp(1, bht, A, g_g) ≈ 173 atol = 0.1

#= visual test
depths = range(1, stop = z, length = 100) ;

test_temps = Ramey_wellboretemp.(depths, 0, bht, A, g_g)
depths_plot = z .- depths |> collect

using Gadfly

plot(y = depths_plot, x = test_temps, Geom.line, Coord.cartesian(yflip = true))
=#

end #Temperature correlations


@testset "test delong" begin
    target = [0,0,1,1,1]
    modela = [0.1, 0.2, 0.6, 0.7, 0.8]
    modelb = [0.3,0.6,0.2,0.7,0.9]
    nontara, tara = modela[1:2], modela[3:end]
    nontarb, tarb = modelb[1:2], modelb[3:end]

    @test round(ROCAnalysis.delong_test(tara, nontara, tarb, nontarb), digits=4) == 0.3173

    target = [0,0,0,0,0,0,1,1,1,1,1,1,1]
    modela = [0.1,0.2,0.05,0.3,0.1,0.6,0.6,0.7,0.8,0.99,0.8,0.67,0.5]
    modelb = [0.3,0.6,0.2,0.1,0.1,0.9,0.23,0.7,0.9,0.4,0.77,0.3,0.89]
    nontara, tara = modela[1:6], modela[7:end]
    nontarb, tarb = modelb[1:6], modelb[7:end]

    @test round(ROCAnalysis.delong_test(tara, nontara, tarb, nontarb), digits=5) == 0.09453
end


#------------------------------
# Linear ODE system with two unkown parameters
#------------------------------

#------------------------------
# Setup
#------------------------------

using Pkg;

# all paths are relative to the repository main folder

Pkg.activate(".")
Pkg.instantiate()
Pkg.status()

using Distributions
using Random
using Flux
using DiffEqFlux
using OrdinaryDiffEq
using Plots
using LaTeXStrings
gr()

include("../src/simulation.jl")
include("../src/model.jl")
include("../src/plotting.jl")

#------------------------------
# Define and obtain ground-truth developments
#------------------------------

# define initial condition
true_u0 = Float32[2, 1]
# define time span on which to solve the ODE
tspan = (0.0f0, 10.0f0)
# define parameters for the two distinct groups
true_odeparams_group1 = Float32[-0.2, 0.00, 0.00, -0.2]
true_odeparams_group2 = Float32[-0.2, 0.00, 0.00, 0.2]

# define corresponding ODE problems for the two groups
prob1 = ODEProblem(linear_2d_system,true_u0,tspan,true_odeparams_group1)
prob2 = ODEProblem(linear_2d_system,true_u0,tspan,true_odeparams_group2)

# solve ODE systems to obtain "true" underlying trajectory in each group
dt=0.1
sol_group1 = solve(prob1, Tsit5(), saveat = dt);
sol_group2 = solve(prob2, Tsit5(), saveat = dt);

#------------------------------
# Simulate data based on true trajectories
#------------------------------

# define numbers of individuals and variables
n = 100
p = 10
q, q_info = 50, 10

# choose simulation setting for baseline variables
baseline_simulation = "groupsonly" # alternatives: "trueparams" or "groupsonly"

# set seed for reproducibility
Random.seed!(12);

# generate time dependent variables
xs, tvals, group1, group2 = generate_xs(n, p, true_u0, sol_group1, sol_group2, maxntps=1, σ_var=0.1, σ_ind=0.1); # default vals: t_start=1.5, t_end=10, maxntps=10, dt=0.1, σ_var=0.1, σ_ind=0.5
# generate baseline variables
if baseline_simulation == "trueparams"
    x_baseline = generate_baseline(n, q, q_info, group1, true_odeparams_group1, true_odeparams_group2); # defaul vals: σ_info=0.1, σ_noise=0.1
elseif baseline_simulation == "groupsonly"
    x_baseline = generate_baseline(n, q, q_info, group1,  σ_info=0.5, σ_noise=0.5); # default vals: σ_info=1, σ_noise=1
else
    error("Please select either 'trueparams' or 'groupsonly' as mode for simulation of baseline variables")
end

# look at data:
data = simdata(xs, x_baseline, tvals, group1, group2);

plot(plot_truesolution(2, data, sol_group1, sol_group2, showdata=false),
    plot_truesolution(1,  data, sol_group1, sol_group2, showdata=false),
    plot_truesolution(2, data, sol_group1, sol_group2, showdata=true),
    plot_truesolution(1,  data, sol_group1, sol_group2, showdata=true),
    layout = (2,2),
    legend = false
)

#------------------------------
# Define and train model
#------------------------------

zdim = nODEparams = 2
m = init_vae(p, q, zdim, nODEparams, prob1)
L = loss_wrapper(m)
ps = getparams(m)
opt = ADAM(0.0005)
trainingdata = zip(xs, x_baseline, tvals);
evalcb() = @show(mean(L.(xs, x_baseline, tvals)))
evalcb_zs() = eval_z_trajectories(xs, x_baseline, tvals, group1, sol_group1, sol_group2, m, dt, swapcolorcoding=true)

for epoch in 1:35
    Flux.train!(L, ps, trainingdata, opt)
    evalcb()
    evalcb_zs()
end

#------------------------------
# Look at results
#------------------------------

# individual plots
p1 = plot_individual_solutions(20,xs, x_baseline, tvals, group1, sol_group1, sol_group2, m, dt,
                                swapcolorcoding=true, showlegend=true)
p2 = plot_individual_solutions(28,xs, x_baseline, tvals, group1, sol_group1, sol_group2, m, dt,
                                swapcolorcoding=true, showlegend=false)
p3 = allindsplot(2, data, m, sol_group1, sol_group2, swapcolorcoding=true)
p4 = allindsplot(1, data, m, sol_group1, sol_group2, swapcolorcoding=true, showlegend=false)

# solutions across the entire dataset
plot(p3,p4, layout = (1,2))


using Distributions: Normal, logpdf
using LinearAlgebra: ldiv!, diag, inv, tr

########################
# Leave-one-out        #
########################

function predict_LOO(Σ::AbstractPDMat, alpha::AbstractVector{<:Real}, y::AbstractVector{<:Real})
    invΣ = inv(Σ)
    σi2 = 1 ./ diag(invΣ)
    μi = -alpha .* σi2 .+ y
    return μi, σi2
end
"""
    predict_LOO(gp::GPE)

Leave-one-out cross-validated predictions.
Returns predictions of yᵢ given all other observations y₋ᵢ,
as a vector of means and variances.
Using the notation from Rasmussen & Williams, see e.g. equation 5.12:
    σᵢ = 𝕍 (yᵢ | y₋ᵢ)^(1/2)
    μᵢ = 𝔼 (yᵢ | y₋ᵢ)

The output is the same as fitting the GP on x₋ᵢ,y₋ᵢ,  and calling `predict_f`
on xᵢ to obtain the LOO predictive mean and variance, repeated for each observation i.
With GPs, this can thankfully be done analytically with a bit of linear algebra,
which is what this function implements.

See also: [`predict_CVfold`](@ref), [`logp_LOO`](@ref)
"""
function predict_LOO(gp::GPE)
    # extract relevant bits from GPE object
    Σ = gp.cK
    alpha = gp.alpha
    y = gp.y
    return predict_LOO(Σ, alpha, y)
end

"""
    logp_LOO(gp::GPE)

Leave-one-out log probability CV criterion.
This is implemented by  summing the
normal log-pdf of each observation
with predictive LOO mean and variance parameters
obtained by the `predict_LOO` function.

See also: [`logp_CVfold`](@ref), [`update_mll!`](@ref)
"""
function logp_LOO(gp::GPE)
    y = gp.y
    μ, σ2 = predict_LOO(gp)
    return sum(logpdf(Normal(μi,√σi2), yi)
               for (μi,σi2,yi)
               in zip(μ,σ2,y)
               )
end

"""
    dlogpdθ_LOO_kern!(gp::GPE)

Gradient of leave-one-out CV criterion with respect to the kernel hyperparameters.
See Rasmussen & Williams equations 5.13.

See also: [`logp_LOO`](@ref), [`dlogpdσ2_LOO`](@ref), [`dlogpdθ_LOO`](@ref)
"""
function dlogpdθ_LOO_kern!(∂logp∂θ::AbstractVector{<:Real}, invΣ::PDMat, kernel::Kernel, x::AbstractMatrix, y::AbstractVector, data::KernelData, alpha::AbstractVector)
    dim = num_params(kernel)
    nobs = length(y)
    @assert length(∂logp∂θ) == dim

    σi2 = 1 ./ diag(invΣ)
    μi = -alpha .* σi2 .+ y

    # Note: if useful, the derivatives of μ and σ could be moved to a separate function.
    # ∂μ∂θ = Matrix{Float64}(undef, nobs, dim)
    # ∂σ∂θ = Matrix{Float64}(undef, nobs, dim)
    Zj = Matrix{Float64}(undef, nobs, nobs)
    for j in 1:dim
        grad_slice!(Zj, kernel, x, x, data, j)
        Zj = invΣ.mat * Zj
        # ldiv!(Σ, Zj)

        ZjΣinv = diag(Zj*Matrix(invΣ))
        ∂σ2∂θj = ZjΣinv.*(σi2.^2)
        ∂μ∂θj = (Zj*alpha).*σi2 .- alpha .* ∂σ2∂θj
        # ∂μ∂θ[:,j] = ∂μ∂θj
        # ∂σ∂θ[:,j] = ∂σ∂θj

        ∂logp∂θj = 0.0
        for i in 1:nobs
            # exponentiated quadratic component:
            ∂logp∂θj -= 2*(y[i]-μi[i]) / σi2[i] * ∂μ∂θj[i]
            ∂logp∂θj -= (y[i]-μi[i])^2 * ZjΣinv[i]
            # log determinant component:
            @assert ZjΣinv[i] * σi2[i] ≈ ∂σ2∂θj[i] / σi2[i]
            ∂logp∂θj += ZjΣinv[i] * σi2[i]
        end
        ∂logp∂θ[j] = ∂logp∂θj
    end
    ∂logp∂θ .*= -1/2
    return ∂logp∂θ
end

"""
    dlogpdσ2_LOO(invΣ::PDMat, x::AbstractMatrix, y::AbstractVector, data::KernelData, alpha::AbstractVector)

Derivative of leave-one-out CV criterion with respect to the logNoise parameter.

See also: [`logp_LOO`](@ref), [`dlogpdθ_LOO_kern!`](@ref), [`dlogpdθ_LOO`](@ref)
"""
function dlogpdσ2_LOO(invΣ::PDMat, x::AbstractMatrix, y::AbstractVector, data::KernelData, alpha::AbstractVector)
    nobs = length(y)

    σi2 = 1 ./ diag(invΣ)
    μi = -alpha .* σi2 .+ y

    Zj = invΣ.mat
    ZjΣinv = diag(Zj^2)
    ∂σ2∂σ2 = ZjΣinv.*(σi2.^2)
    ∂μ∂σ2 = (Zj*alpha).*σi2 .- alpha .* ∂σ2∂σ2

    ∂logp∂σ2 = 0.0
    for i in 1:nobs
        # exponentiated quadratic component:
        ∂logp∂σ2 -= 2*(y[i]-μi[i]) / σi2[i] * ∂μ∂σ2[i]
        ∂logp∂σ2 -= (y[i]-μi[i])^2 * ZjΣinv[i]
        # log determinant component:
        @assert ZjΣinv[i] * σi2[i] ≈ ∂σ2∂σ2[i] / σi2[i]
        ∂logp∂σ2 += ZjΣinv[i] * σi2[i]
    end
    return -∂logp∂σ2 ./ 2
end

"""
    dlogpdθ_LOO(gp::GPE; noise::Bool, domean::Bool, kern::Bool)

Derivatives of the leave-one-out CV criterion.

See also: [`logp_LOO`](@ref), [`dlogpdθ_LOO_kern!`](@ref), [`dlogpdσ2_LOO`](@ref)
"""
function dlogpdθ_LOO(gp::GPE; noise::Bool, domean::Bool, kern::Bool)
    Σ = gp.cK
    x, y = gp.x, gp.y
    data = gp.data
    kernel = gp.kernel
    alpha = gp.alpha

    invΣ = inv(Σ)

    n_mean_params = num_params(gp.mean)
    n_kern_params = num_params(gp.kernel)
    ∂logp∂θ = Vector{Float64}(undef, noise + domean*n_mean_params + kern*n_kern_params)
    i = 1
    if noise
        ∂logp∂θ[i] = dlogpdσ2_LOO(invΣ, x, y, data, alpha)*2*exp(2 * gp.logNoise.value)
        i += 1
    end
    if domean && n_mean_params>0
        throw("I don't know how to do means yet")
        Mgrads = grad_stack(gp.mean, gp.x)
        for j in 1:n_mean_params
            gp.dmll[i] = dot(Mgrads[:,j], gp.alpha)
            i += 1
        end
    end
    if kern
        ∂logp∂θ_k = @view(∂logp∂θ[i:end])
        dlogpdθ_LOO_kern!(∂logp∂θ_k, invΣ, kernel, x, y, data, alpha)
    end
    return ∂logp∂θ
end

########################
# Arbitrary fold       #
########################

const Folds = AbstractVector{<:AbstractVector{Int}}

function predict_CVfold(Σ::AbstractPDMat, alpha::AbstractVector{<:Real}, y::AbstractVector{<:Real}, folds::Folds)
    invΣ = inv(Σ)
    μ = Vector{Float64}[]
    Σ = Matrix{Float64}[]
    for V in folds
        ΣVT = inv(Matrix(invΣ)[V,V])
        μVT = y[V]-ΣVT*alpha[V]
        push!(μ, μVT)
        push!(Σ, ΣVT)
    end
    return μ, Σ
end

"""
    predict_CVfold(gp::GPE, folds::Folds)

Cross-validated predictions for arbitrary folds.
A fold is a set of indices of the validation set.
Returns predictions of y_V (validation set) given all other observations y_T (training set),
as a vector of means and covariances.
Using the notation from Rasmussen & Williams, see e.g. equation 5.12:
    σᵢ = 𝕍 (yᵢ | y₋ᵢ)^(1/2)
    μᵢ = 𝔼 (yᵢ | y₋ᵢ)

The output is the same as fitting the GP on x_T,y_T,  and calling `predict_f`
on x_V to obtain the LOO predictive mean and covariance, repeated for each fold V.
With GPs, this can thankfully be done analytically with a bit of linear algebra,
which is what this function implements.

See also: [`predict_LOO`](@ref)
"""
function predict_CVfold(gp::GPE, folds::Folds)
    # extract relevant bits from GPE object
    Σ = gp.cK
    alpha = gp.alpha
    y = gp.y
    return predict_CVfold(Σ, alpha, y, folds)
end

"""
    logp_CVfold(gp::GPE, folds::Folds)

CV criterion for arbitrary fold.  A fold is a set of indices of the validation set.

See also [`predict_CVfold`](@ref), [`logp_LOO`](@ref)
"""
function logp_CVfold(gp::GPE, folds::Folds)
    y = gp.y
    μ, Σ = predict_CVfold(gp, folds)
    CV = 0.0
    for (μVT,ΣVT,V) in zip(μ,Σ,folds)
        chol = similar(ΣVT)
        ΣVT, chol = make_posdef!(ΣVT, chol; nugget=1e-10)
        ΣPD = PDMat(ΣVT, chol)
        CV += logpdf(MvNormal(μVT, ΣPD), y[V])
    end
    return CV
end

"""
    gradient_fold(invΣ, alpha, ZjΣinv, Zjα, V::AbstractVector{Int})

Gradient with respect to the kernel hyperparameters of the CV criterion
component for one validation fold:
```math
    \\nabla_{\\theta}\\left(\\log p (Y_V \\mid Y_T, \\theta) \\right)
```
where `Y_V` is the validation set and `Y_T` is the training set (all other observations)
for this fold.
"""
@inline function gradient_fold(invΣ, alpha, ZjΣinv, Zjα, V::AbstractVector{Int})
    ∂logp∂θj = 0.0
    ΣVTinv = PDMats.PDMat(mat(invΣ)[V,V])
    ΣVTα = ΣVTinv\alpha[V]
    # exponentiated quadratic component:
    ZjΣinvVV = ZjΣinv[V,V]
    ∂logp∂θj -= 2*dot(ΣVTα,Zjα[V])
    ∂logp∂θj += dot(ΣVTα, ZjΣinvVV*ΣVTα)
    # log determinant component:
    ∂logp∂θj += tr(ΣVTinv\ZjΣinvVV)
    return ∂logp∂θj
end

function dlogpdθ_CVfold_kern!(∂logp∂θ::AbstractVector{<:Real}, invΣ::PDMat, kernel::Kernel, x::AbstractMatrix, y::AbstractVector, data::KernelData, alpha::AbstractVector, folds::Folds)
    nobs = length(y)
    dim = num_params(kernel)

    @assert length(∂logp∂θ) == dim
    buffer1 = Matrix{Float64}(undef, nobs, nobs)
    buffer2 = Matrix{Float64}(undef, nobs, nobs)
    for j in 1:dim
        grad_slice!(buffer2, kernel, x, x, data, j)
        mul!(buffer1, mat(invΣ), buffer2)
        Zj = buffer1
        Zjα = Zj*alpha

        mul!(buffer2, Zj, mat(invΣ))
        ZjΣinv = buffer2

        ∂logp∂θj = 0.0
        for V in folds
            ∂logp∂θj += gradient_fold(invΣ, alpha, ZjΣinv, Zjα, V)
        end
        ∂logp∂θ[j] = ∂logp∂θj
    end
    ∂logp∂θ .*= -1/2
    return ∂logp∂θ
end

function dlogpdσ2_CVfold(invΣ::PDMat, x::AbstractMatrix, y::AbstractVector, data::KernelData, alpha::AbstractVector, folds::Folds)
    nobs = length(y)

    Zj = mat(invΣ)
    Zjα = Zj*alpha
    ZjΣinv = Zj^2

    ∂logp∂σ2 = 0.0
    for V in folds
        ∂logp∂σ2 += gradient_fold(invΣ, alpha, ZjΣinv, Zjα, V)
    end
    return -∂logp∂σ2 / 2
end

"""
    dlogpdθ_CVfold_kern!(∂logp∂θ::AbstractVector{<:Real}, gp::GPE, folds::Folds;
                         noise::Bool, domean::Bool, kern::Bool)

Derivative of leave-one-out CV criterion with respect to the noise, mean and kernel hyperparameters.
See Rasmussen & Williams equations 5.13.
"""
function dlogpdθ_CVfold(gp::GPE, folds::Folds; noise::Bool, domean::Bool, kern::Bool)
    Σ = gp.cK
    x, y = gp.x, gp.y
    data = gp.data
    kernel = gp.kernel
    alpha = gp.alpha

    invΣ = inv(Σ)

    n_mean_params = num_params(gp.mean)
    n_kern_params = num_params(gp.kernel)
    ∂logp∂θ = Vector{Float64}(undef, noise + domean*n_mean_params + kern*n_kern_params)
    i = 1
    if noise
        ∂logp∂θ[i] = dlogpdσ2_CVfold(invΣ, x, y, data, alpha, folds)*2*exp(2 * gp.logNoise.value)
        i += 1
    end
    if domean && n_mean_params>0
        throw("I don't know how to do means yet")
        Mgrads = grad_stack(gp.mean, gp.x)
        for j in 1:n_mean_params
            gp.dmll[i] = dot(Mgrads[:,j], gp.alpha)
            i += 1
        end
    end
    if kern
        ∂logp∂θ_k = @view(∂logp∂θ[i:end])
        dlogpdθ_CVfold_kern!(∂logp∂θ_k, invΣ, kernel, x, y, data, alpha, folds)
    end
    return ∂logp∂θ
end


module BernoulliTest

using Test
using ReactiveMP
using Distributions
using Random

@testset "Bernoulli" begin

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

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

        @test typeof(d) <: Bernoulli
        @test mean(d) === 0.5
        @test succprob(d) === 0.5
        @test failprob(d) === 0.5
    end

    @testset "prod" begin
        @test prod(ProdAnalytical(), Bernoulli(0.5), Bernoulli(0.5))   ≈ Bernoulli(0.5)
        @test prod(ProdAnalytical(), Bernoulli(0.1), Bernoulli(0.6))   ≈ Bernoulli(0.14285714285714285)
        @test prod(ProdAnalytical(), Bernoulli(0.78), Bernoulli(0.05)) ≈ Bernoulli(0.1572580645161291)
    end

    @testset "probvec" begin
        @test probvec(Bernoulli(0.5)) === (0.5, 0.5)
        @test probvec(Bernoulli(0.3)) === (0.3, 0.7)
        @test probvec(Bernoulli(0.6)) === (0.6, 0.4)
    end

end

end


function check(i, j)
    id, im = div(i, 9), mod(i, 9)
    jd, jm = div(j, 9), mod(j, 9)

    jd == id && return true
    jm == im && return true

    div(id, 3) == div(jd, 3) &&
    div(jm, 3) == div(im, 3)
end

const lookup = zeros(Bool, 81, 81)

for i in 1:81
    for j in 1:81
        lookup[i,j] = check(i-1, j-1)
    end
end

function solve_sudoku(callback::Function, grid::Array{Int64})
    (function solve()
        for i in 1:81
            if grid[i] == 0
                t = Dict{Int64, Void}()

                for j in 1:81
                    if lookup[i,j]
                        t[grid[j]] = nothing
                    end
                end

                for k in 1:9
                    if !haskey(t, k)
                        grid[i] = k
                        solve()
                    end
                end

                grid[i] = 0
                return
            end
        end

        callback(grid)
    end)()
end

function display(grid)
    for i in 1:length(grid)
        print(grid[i], " ")
        i %  3 == 0 && print(" ")
        i %  9 == 0 && print("\n")
        i % 27 == 0 && print("\n")
    end
end

grid = Int64[5, 3, 0, 0, 2, 4, 7, 0, 0,
             0, 0, 2, 0, 0, 0, 8, 0, 0,
             1, 0, 0, 7, 0, 3, 9, 0, 2,
             0, 0, 8, 0, 7, 2, 0, 4, 9,
             0, 2, 0, 9, 8, 0, 0, 7, 0,
             7, 9, 0, 0, 0, 0, 0, 8, 0,
             0, 0, 0, 0, 3, 0, 5, 0, 6,
             9, 6, 0, 0, 1, 0, 3, 0, 0,
             0, 5, 0, 6, 9, 0, 0, 1, 0]

solve_sudoku(display, grid)


ϕK(t) = -1 < t < 1 ? (1.0-t^2)^3 : 0.0

# sin(tx)/t
function sinc(t, x)
    if abs(t) <= 1e-4
        return x - t^2*x^3/6.0
    else
        return sin(t*x) / t
    end
end
"""
    very literal implementation of equation 8.5
    in Wang&Wang 2011
"""
function Fhatx_brute(x::Real, W::AbstractVector, num_t::Int, h::Real,
        U::AbstractVector{D} where D <: UnivariateDistribution)
    quad_nodes, quad_weight = FastGaussQuadrature.gausslegendre(num_t)
    htt = quad_nodes
    tt = htt ./ h
    t_weight = quad_weight ./ h

    ϕKht = ϕK.(htt)
    n = length(W)
    s = 0.0
    sum_abs_ϕUkt = abs2.(cf.(U[1], tt))
    for k in 2:n
        sum_abs_ϕUkt .+= abs2.(cf.(U[k], tt))
    end
    mean_abs_ϕUkt = sum_abs_ϕUkt/n
    @inbounds for j in 1:n
        sinctxW = sinc.(tt, x-W[j])
        ϕUj = cf.(U[j], -tt)
        inv_ψUj = ϕUj ./ mean_abs_ϕUkt
        s += dot(t_weight, sinctxW .* ϕKht  .* inv_ψUj)
        isfinite(s) || throw(AssertionError("s=$s is not finite, j=$j"))
    end
    Fx = 0.5 + 1/(2π*n) * s # 2h is the Jacobian
    @assert isfinite(Fx)
    (imag(Fx) ≈ 0.0) || AssertionError("Fx has non-negligible imaginary part $Fx")
    return real(Fx)
end

@inline function real_cf_Uj_t(Uj::UnivariateDistribution, t::Real)
    return @fastmath real(cf(Uj, t))
end
@inline function real_cf_Uj_t(Uj::Normal, t::Real)
    # assumes Uj has mean zero!
    return exp(-var(Uj)*t^2/2)
end

function Fhat(xx::Vector{Float64}, W::Vector{Float64}, num_t::Int, h::Real,
        U::Vector{D} where D <: UnivariateDistribution)
    quad_nodes, t_weight = FastGaussQuadrature.gausslegendre(num_t)
    # FastGaussQuadrature outputs nodes from -1 to 1
    # we need nodes from 0 to 1
    htt = (quad_nodes .+ 1.0) ./ 2.0
    tt = htt ./ h
    ϕKht = ϕK.(htt)
    n = length(W)
    sum_abs_ϕUkt = abs2.(cf.(U[1], tt))
    for k in 2:n
        sum_abs_ϕUkt .+= abs2.(cf.(U[k], tt))
    end
    inv_mean_abs_ϕUkt = n ./ sum_abs_ϕUkt
    s_cache = Vector{Float64}[]
    for j in 1:n
        Uj = U[j]
        sUj = real_cf_Uj_t.(Uj, -tt) .* inv_mean_abs_ϕUkt .* t_weight .* ϕKht
        push!(s_cache, sUj)
    end
    F_xx = zeros(Float64, length(xx))
    @inbounds for (ix, x) in enumerate(xx)
        s = zero(Float64)
        for j in 1:n
            Wj, Uj = W[j], U[j]
            scache_j = s_cache[j]
            for (it, t) in enumerate(tt)
                # sinctxW = @fastmath sin(t*(x-Wj))/t
                sinctxW = sinc(t, x-Wj)
                s += scache_j[it] * sinctxW
            end
        end
        F_xx[ix] = 0.5 + 1/(π*n) * s / (2*h) # 2h is the Jacobian
    end
    return F_xx
end

"""
    Estimation using Fourier transforms leads
    to estimates of the CDF that are not strictly valid:
    they don't go from 0 to 1 and aren't monotonic.
    This is a function to "fix them up", intended
    as a post-processing step.
"""
function fix_CDF!(Fhat_xx)
    discretised_pdf = diff(Fhat_xx)
    while minimum(discretised_pdf) < 0.0
        imin = argmin(discretised_pdf)
        pmin = discretised_pdf[imin]
        discretised_pdf[imin] = 0.0
        discretised_left = @view(discretised_pdf[imin-1:-1:1])
        push_negative_prob!(discretised_left, pmin / 2)
        discretised_right = @view(discretised_pdf[imin+1:end])
        push_negative_prob!(discretised_right, pmin / 2)
    end
    corrected_cdf = cumsum(discretised_pdf)
    corrected_cdf .-= minimum(corrected_cdf)
    corrected_cdf ./= maximum(corrected_cdf)
    @assert maximum(corrected_cdf) ≈ 1.0
    Fhat_xx[1] = 0.0
    Fhat_xx[2:end] .= corrected_cdf
    return Fhat_xx
end
function push_negative_prob!(prob_vec, p::Real)
    i = 0
    while p < 0
        i += 1
        if i == length(prob_vec)
            p = 0.0
            break
        end
        p += prob_vec[i] # pick this bit of probability up
        prob_vec[i] = 0.0 # new probability is zero
    end
    prob_vec[i] = p
end


Random.seed!(0)

#sig = [1 2; 2 2]
#sig = [1 1; 2 3]
sig = [2 3]
#sig = [2 2]
#sig = [3 2; 2 3]

S = random_S0Graph(sig)

w = random_bounded(S.n)

U = random_S1_unitary(sig);
SU = S0Graph(S.sig, U * S.S * U')
@time opt1 = dsw(S, w, eps=solver_eps)[1]
@time opt2 = dsw(SU, U*w*U', eps=solver_eps)[1]
@test opt1 ≈ opt2  atol=tol


#=

  Euler #37 in Julia.

  Problem 37:
  """
  The number 3797 has an interesting property. Being prime itself, it is possible to
  continuously remove digits from left to right, and remain prime at each stage:
  3797, 797, 97, and 7. Similarly we can work from right to left: 3797, 379, 37, and 3.

  Find the sum of the only eleven primes that are both truncatable from left to right
  and right to left.

  NOTE: 2, 3, 5, and 7 are not considered to be truncatable primes.
  """

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

=#

include("Euler.jl")

# check is n is a truncated prime
function truncated_prime(n)
    len = nlen(n)
    for i in 1:len-1
        ii = 10^i
        if !isPrime(n % ii) || !isPrime(div(n,ii))
            return false
        end
    end
    return true
end

# 0.02413909s
function euler37a()
    p = 11
    sum = 0
    c = 0
    while c < 11
        if truncated_prime(p) && isPrime(p)
            c += 1
            sum += p
        end
        p += 2
    end
    return sum
end

run_euler(euler37a)


# QuObj Tests
using Schrodinger
using Test, LinearAlgebra, SparseArrays, Statistics

println("Testing Quantum Info Functions...")

@testset "Quantum Information" begin
# Build a few different variable for testing
N = 13
α = complex(1.1,-0.33)

g = basis(2,0)
e1 = basis(2,1)
coherentα = coherent(N,α)
coherentαhalf = coherent(N,0.5α)
numop = numberop(N)

N! = normalize!

@testset "Expectation Value" begin
    # Test calculation of expectation valu
    @test expect(σz,g) == 1
    @test expect(e1,σz) == -1
    @test expect(σx,N!(g+e1)) ≈ 1
    @test expect(N!(g-e1),σx) ≈ -1
    @test expect(N!(g+1im*e1),σy) ≈ 1
    @test expect(σy,N!(g-1im*e1)) ≈ -1
    @test expect(coherentα,destroy(N)) ≈ α
    @test expect(destroy(N),coherent(N,0.5α,true)) ≈ 0.5α
    @test [expect(basis(N,i),numop) for i=0:N-1] == collect(0:N-1)
    @test expect(N!(basis(N,2)+2*basis(N,3)),numop) == (2+2^2*3)/5
    @test expect(N!(Operator(basis(N,2))+2*Operator(basis(N,3))),numop) == (1*2+2*3)/3
    @test expect(numop,thermal(N,0.21)) ≈ 0.21
    @test expect(projectorop(N,[2,4]),maxmixed(N)) == 2/N
end

ONEQUBIT = [g,e1,N!(g+e1),N!(g-e1),N!(g+1im*e1),N!(g-1im*e1)]
include(joinpath(@__DIR__,"twoqubit.jl"))

@testset "Fidelity" begin # Test calculation of various fidelities
    # fidelity
    @test fidelity(g,g) == 1
    @test fidelity(g,e1) == 0
    @test fidelity(g,N!(g+e1)) == 1/√(2)
    @test fidelity(coherentα,basis(N,3)) ≈ abs(coherentα[4])
    @test fidelity(Operator(coherentα),basis(N,3)) ≈ abs(coherentα[4])
    @test fidelity(basis(N,3),Operator(coherentαhalf)) ≈ abs(coherentαhalf[4])
    @test fidelity(Operator(coherentα),Operator(basis(N,3))) ≈ abs(coherentα[4])
    @test fidelity(Operator(basis(N,2)),Operator(coherent(N,1))) ≈ abs(coherent(N,1)[3])
    for i = 1:20
        P = 10*rand()*randn(N,N) + 10*rand()*im*randn(N,N); P = normalize!(Operator(P'*P))
        Q = 10*rand()*randn(N,N) + 10*rand()*im*randn(N,N); Q = normalize!(Operator(Q'*Q))
        @test fidelity(P,Q) ≈ (rtQ=√(Q); real(trace(√(rtQ*P*rtQ))))
    end
    # fidelity2
    @test fidelity2(g,g) == 1
    @test fidelity2(g,e1) == 0
    @test fidelity2(g,N!(g+e1)) ≈ 1/2
    @test fidelity2(coherentαhalf,basis(N,3)) ≈ abs2(coherentαhalf[4])
    @test fidelity2(Operator(coherentαhalf),basis(N,3)) ≈ abs2(coherentαhalf[4])
    @test fidelity2(basis(N,3),Operator(coherentα)) ≈ abs2(coherentα[4])
    @test fidelity2(Operator(coherentαhalf),Operator(basis(N,3))) ≈ abs2(coherentαhalf[4])
    @test fidelity2(Operator(coherent(N,α,true)),Operator(basis(N,3))) ≈ abs2(coherent(N,α,true)[4])
    @test fidelity2(Operator(basis(N,2)),Operator(coherent(N,1))) ≈ abs2(coherent(N,1)[3])
    # entanglement and average gate fidelity
    U1, V1 = rand_unitary(5), rand_unitary(5)
    @test entanglement_fidelity(U1,V1) ≈ (ε=qeye(5)⊗(U1'V1);ϕ=maxentangled(2,5);abs2(ϕ⋅(ε*ϕ)))
    U2, V2 = rand_unitary(2), rand_unitary(2)
    @test gate_fidelity(U2,V2) ≈ mean(fidelity2(ψ,U2'*(V2*ψ)) for ψ ∈ ONEQUBIT)
    @test gate_fidelity(U2,V2) ≈ (2*entanglement_fidelity(U2,V2)+1)/(2+1)
    U3, V3 = rand_unitary(4,(2,2)), rand_unitary(4,(2,2))
    @test gate_fidelity(U3,V3) ≈ mean(fidelity2(ψ,U3'*(V3*ψ)) for ψ ∈ TWOQUBIT)
    @test gate_fidelity(U3,V3) ≈ (4*entanglement_fidelity(U3,V3)+1)/(4+1)
end

@testset "Purity" begin
    U = rand_unitary(N)
    ρ = N!(Operator(coherentαhalf)+numop)
    @test purity(g) == purity(e1) == 1
    @test purity(coherentα) == purity(coherentαhalf) == 1
    @test purity(ρ) ≈ inner(ρ,ρ)
    @test purity(U*ρ*U') ≈ purity(ρ)
    @test purity(maxmixed(4)) == 1/4
end

@testset "Level Probabilities" begin
    # Test calculation of level probabilities
    @test levelprobs(g) == [1,0]
    @test levelprobs(e1) == [0,1]
    @test levelprobs(N!(e1+g)) ≈ [0.5,0.5]
    @test levelprobs(e1⊗g) == [0,0,1,0]
    @test levelprobs(N!(e1⊗g+g⊗e1)) ≈ [0,0.5,0.5,0]
    @test levelprobs(N!(e1⊗g+g⊗g),1:2) ≈ [[0.5,0.5],[1,0]]
end

function inner_cs_test(pe=0.1,ϕ=1.23)
    # σx on qubit 1 of 2-qubit system
    # no phase difference
    V1 = ((1-pe)*σx + pe*qeye(2)) .⊗ (Operator(basis(2,0)), Operator(basis(2,1)))
    # with phase difference
    V2 = (V1[1], cis(ϕ)*V1[2])
    # Ideal gate
    U = σx⊗Operator(basis(2,0)) + σx⊗Operator(basis(2,1))
    # inner products
    w = sqrt(sum(abs2,diag(sum(V2)'U)))/2
    x = abs(inner(sum(data.(V1)),data(U)))/4
    y = sum(abs,inner.(data.(V2),(data(U),)))/4
    z = abs(inner(sum(data.(V2)),data(U)))/4
    w,x,y,z
end

function anharm_NOT_RF(steps)
    dragpulse = Schrodinger.dragpulse_cos
    α = -0.4 * 2π # anharmonicity (GHz)
    Hd = α*Operator(basis(3,2)) # drift Hamiltonian
    σX, σY = 1/2*(create(3)+destroy(3)), 1im/2*(create(3)-destroy(3))
    λ = √2
    T = 5 # (ns)
    tspan = (-T/2,T/2)
    paramsx, paramsy = [0,T,0,α,λ,0,π], [0,T,0,α,λ,π/2,π]
    V = Operator(SchrodingerProp(Hd,((σX,dragpulse,paramsx),(σY,dragpulse,paramsy)),tspan,steps))
    #U = Operator([0 1 0; 1 0 0; 0 0 1]) # 3lvl NOT gate
    ## Fidelity
    #ϕ = N!(basis(3,0)⊗basis(3,0) + basis(3,1)⊗basis(3,1))
    #ϕ' * (qeye(3)⊗U'V) * Operator(ϕ) * (qeye(3)⊗V'U) * ϕ
    # Create objective function type
    #O = CoherentSubspaces(Ut,1:2,Hd,Hc,T,n) # care only about computational subspace
    #grape(O,ui)
    return V
end

end


import Distributions

"""
    perfectNormal(n::Int, mean::Number=0, sd::Number=1)

Generate an almost-perfect normal distribution of size `n`.

# Arguments
- `n::Int`: Length of the vector.
- `mean::Number`: Mean of the vector.
- `sd::Number`: SD of the vector.

# Examples
```julia
perfectNormal(10, 0, 1)
```
"""
function perfectNormal(n::Int, mean::Number=0, sd::Number=1)
    x = Distributions.quantile.([Distributions.Normal(mean, sd)],
                                range(1/n, stop=1-1/n, length=n))
    return x
end
perfectNormal(; n::Int, mean::Number=0, sd::Number=1) = perfectNormal(n, mean, sd)


@testset "Elements in algebras" begin
  Qx, x = FlintQQ["x"]
  f = x^2 + 1
  A = AlgAss(f)

  @testset "Is integral" begin
    @test Hecke.isintegral(A[1]) == true
    @test Hecke.isintegral(fmpq(1, 2)*A[1]) == false

    B = group_algebra(FlintQQ, small_group(2, 1))
    @test Hecke.isintegral(B[1]) == true
    @test Hecke.isintegral(fmpq(1, 2)*B[1]) == false

    C = matrix_algebra(FlintQQ, B, 2)
    @test Hecke.isintegral(C[1]) == true
    @test Hecke.isintegral(fmpq(1, 2)*C[1]) == false
  end

  @testset "Characteristic polynomial" begin
    K, a = number_field(f, "a")

    b = rand(K, -10:10)
    c = A(coeffs(b))

    cpb = charpoly(b)
    cpc = charpoly(c)
    rcpc = reduced_charpoly(c)

    # We have to compare the polynomials by hand as they live in "different" polynomial rings...
    for p in [ cpc, rcpc ]
      for i = 0:2
        @test coeff(p, i) == coeff(cpb, i)
      end
    end

    mpb = minpoly(b)
    mpc = minpoly(c)

    @test degree(mpb) == degree(mpc)
    for i = 0:2
      @test coeff(mpb, i) == coeff(mpc, i)
    end
  end

end


using LinearAlgebra
using SparseArrays
using SuiteSparse

struct IndGraphSparse{T <: RealOrComplex, Ti} <: IndGraph
  m::Int
  n::Int
  A::SparseMatrixCSC{T, Ti}
  F::SuiteSparse.CHOLMOD.Factor{T} #LDL factorization

  tmp::Array{T, 1}
  tmpx::SubArray{T, 1, Array{T, 1}, Tuple{UnitRange{Int}}, true}
  res::Array{T, 1}
end

function IndGraphSparse(A::SparseMatrixCSC{T,Ti}) where {T <: RealOrComplex, Ti}

  m, n = size(A)
  K = [SparseMatrixCSC{T}(I, n, n) A'; A -SparseMatrixCSC{T}(I, m, m)]

  F = ldlt(K)

  tmp = Array{T, 1}(undef, m + n)
  tmpx = view(tmp, 1:n)
  tmpy = view(tmp, (n + 1):(n + m)) #second part is always zeros
  fill!(tmpy, 0)

  res = Array{T,1}(undef, m + n)
  return IndGraphSparse(m, n, A, F, tmp, tmpx, res)
end

function prox!(
    x::AbstractVector{T},
    y::AbstractVector{T},
    f::IndGraphSparse,
    c::AbstractVector{T},
    d::AbstractVector{T},
    gamma=1.0
  ) where {T <: RealOrComplex}
  #instead of res = [c + f.A' * d; zeros(f.m)]
  mul!(f.tmpx, adjoint(f.A), d)
  f.tmpx .+= c
  # A_ldiv_B!(f.res, f.F, f.tmp) #is not working
  f.res .= f.F \ f.tmp
  # f.res .= f.F \ f.tmp #note here f.tmp which is m+n array
  copyto!(x, 1, f.res, 1, f.n)
  copyto!(y, 1, f.res, f.n + 1, f.m)
  return 0.0
end

function (f::IndGraphSparse)(x::AbstractVector{T}, y::AbstractVector{T}) where
    {T <: RealOrComplex}
  # the tolerance in the following line should be customizable
  tmpy = view(f.res, 1:f.m) # the res is rewritten in prox!
  mul!(tmpy, f.A, x)
  tmpy .-= y
  if norm(tmpy, Inf) <= 1e-12
    return 0.0
  end
  return +Inf
end

fun_name(f::IndGraphSparse) = "Indicator of an operator graph defined by sparse matrix"
# fun_dom(f::IndGraph) = "AbstractArray{Real,1}, AbstractArray{Complex,1}"
# fun_expr(f::IndGraph) = "x,y ↦ 0 if Ax = y, +∞ otherwise"
# fun_params(f::IndGraph) =
#   string( "A = ", typeof(f.A), " of size ", size(f.A))

function prox_naive(
    f::IndGraphSparse,
    c::AbstractVector{T},
    d::AbstractVector{T},
    gamma=1.0
  ) where {T <: RealOrComplex}

  tmp = f.A'*d
  tmp .+= c
  res = [tmp; zeros(f.m)]
  xy = f.F \ res
  return xy[1:f.n], xy[f.n + 1:f.n + f.m], 0.0
end

function (f::IndGraphSparse)(xy::AbstractVector{T}) where
  {T <: RealOrComplex}
  x, y = splitinput(f, xy)
  return f(x, y)
end

(f::IndGraphSparse)(xy::Tuple{AbstractVector{T}, AbstractVector{T}}) where
  {T <: RealOrComplex} = f(xy[1], xy[2])


# ------------------------------------------------------------------------------
# Schwarz smoother example
# ------------------------------------------------------------------------------

using LFAToolkit
using LinearAlgebra

# setup
mesh = Mesh1D(1.0)
p = 2

# diffusion operator
diffusion = GalleryOperator("diffusion", p + 1, p + 1, mesh)
# Schwarz smoother
schwarz = Schwarz(diffusion)

# compute operator symbols
A = computesymbols(schwarz, [], [π])
eigenvalues = real(eigvals(A))
#println(eigenvalues)

# ------------------------------------------------------------------------------


using Distributions, SpecialFunctions
import Distributions: mean

""" Exact Bayes (Adams and MacKay, 2007) for the Categorical task (section 2.3.2)
"""
struct TBayesFilter{D, A} <: BayesFilter
    ns::Int
    posteriors::Array{D, 1}
    prior::D
    probabilities::Array{Float64, 1}
    stochasticity::Float64
    changeprobability::Float64
    cutoff::Float64
    approximation::A
    Ps1::Array{Float64, 1}
end
function TBayesFilter(;ns=5, stochasticity = .01,
                    changeprobability = 0.01, cutoff = 1e-16)

    prior = MyDirichlet(ones(ns) * stochasticity)
    approximation = CutoffApprox(cutoff)
    TBayesFilter(ns, [deepcopy(prior)], prior, [1.], stochasticity,
            changeprobability, cutoff, approximation,
            Array{Float64, 1}(undef, ns))
end
meanparameter(f::Union{TBayesFilter, TMPN, TSOR}) = sum(mean.(f.posteriors) .* f.probabilities)
function update_posterior!(f::BayesFilter, x)
    πt= dataprobability.(f.posteriors, x)       # P(x_t|νʳ, χʳ)
    f.probabilities .*= πt
    f.probabilities ./= sum(f.probabilities)
    update!(f.posteriors, x)
end
function update_prior!(f::BayesFilter)
    p_switch = sum(f.probabilities) * f.changeprobability      # ∝ P(r_t = 0, x_{1:t})
    # @show sum(f.probabilities)
    f.probabilities .*= (1 - f.changeprobability)              # ∝ P(r_t = 1, x_{1:t})
    push!(f.probabilities, p_switch)
    push!(f.posteriors, deepcopy(f.prior))
end
function update!(f::Union{TBayesFilter, TMPN, TSOR}, x)
    update_posterior!(f, x)
    approximate!(f.approximation, f)
    parameter_estimate = meanparameter(f)       # E[θ | x_{1:t}]
    update_prior!(f)
    @. f.Ps1 .= copy(parameter_estimate)
end
struct MyDirichlet
    alpha::Vector{Float64}
end
function dataprobability(d::MyDirichlet, x)
    alpha0 = sum(d.alpha)
    tmp = lgamma(d.alpha[x] + 1) - lgamma(d.alpha[x]) +
            lgamma(alpha0) - lgamma(alpha0 + 1)
    exp(tmp)
end
mean(d::MyDirichlet) = d.alpha / sum(d.alpha)
function update!(d::Vector{<:MyDirichlet}, x)
    @inbounds for i in 1:length(d)
        d[i].alpha[x] += 1
    end
    d
end