test = (4, 8)
puzzle = (1, 10)

rot(v::Integer, m::Integer = 10) = mod(v - 1, m) + 1

function q1(p1::Integer, p2::Integer)::Int
    roll(i) = sum(rot.([i, i + 1, i + 2], 100))

    rolls = [[roll(i), roll(i + 3)] for i in 1:6:1000]
    fields = [rot.(fs) for fs in cumsum([[p1, p2], rolls...])]
    scores = cumsum(fields[2:end])
    last_round = findfirst(v -> any(>=(1000), v), scores)
    winner = argmax(scores[last_round])
    loser_score = winner == 1 ? scores[last_round - 1][2] : scores[last_round][1]
    total_rolls = 3 * (2 * (last_round - 1) + winner)
    return total_rolls * loser_score
end

function q2(p1::Integer, p2::Integer)::Int
    (s1, s2) = quantum_scores.((p1, p2))
    return foldl(2:size(s1, 1); init = (0, 0)) do (w1, w2), r
        # p1 reaches 21 exactly after r rounds this many times
        p1w = sum(s1[r, 21:end]) - 27 * sum(s1[r - 1, 21:end])
        # p2 stays below 21 after r rounds this many times
        p2l = sum(s2[r - 1, 1:20])
        # vice versa
        p2w = sum(s2[r, 21:end]) - 27 * sum(s2[r - 1, 21:end])
        p1l = sum(s1[r, 1:20])
        # count wins
        return (w1 + p1w * p2l, w2 + p2w * p1l)
    end |> maximum
end

function quantum_scores(start::Integer; rounds = 10, max_score = 90)
    rolls = [3 => 1, 4 => 3, 5 => 6, 6 => 7, 7 => 6, 8 => 3, 9 => 1]

    # states[r, f, s] = number of parallel states after a solo game of `r - 1`
    #   rounds ((r - 1) * 3 rolls) end on field `f`, with a score of `s`.
    states = fill(0, rounds + 1, 10, max_score + 1)
    states[1, start, 1] = 1
    for round in 1:rounds, score in 1:max_score, field in 1:min(score, 10)
        states[round + 1, field, score + 1] =
            sum([count * states[round, rot(field - roll), score + 1 - field]
                 for (roll, count) in rolls])
    end

    # clean, check and trim
    r = reshape(sum(states[2:end, :, 2:end]; dims = 2), rounds, max_score)
    @assert all(iszero, r[end, 1:20]) "Round count too low, not all games end!"
    @assert sum(r[end, :]) == 27 ^ rounds "Max score too low!"
    return r[CartesianIndex(1, 1):findlast(!iszero, r)]
end

@assert q1(test...) == 739785
@assert q2(test...) == 444356092776315

println("Q1: ", q1(puzzle...))
println("Q2: ", q2(puzzle...))


# dtft.jl

import MIRT: dtft_init, dtft, dtft_adj
import MIRT: dtft_loop_n
import MIRT: dtft_loop_m
import MIRT: dtft_dist_m
import MIRT: dtft_pmap_m
import MIRT: dtft_matvec

using FFTW: fft
#using LinearAlgebra: norm
using Test: @test, @testset
using LinearMapsAA: LinearMapAA, LinearMapAO

# 1D small test data to verify correctness
function dtft_test1( ; N::Int=2^10)
	x = randn(ComplexF64, N)
	w = (2*pi) * (0:N-1) / N

	o0 = fft(x)

	o1 = dtft_loop_n(w, x)
	o2 = dtft_loop_m(w, x)
	o3 = dtft_dist_m(w, x)
	o4 = dtft_pmap_m(w, x)
	o5 = dtft_matvec(w, x)
#	@show norm(o1 - o0) / norm(o0)

	@test isapprox(o1, o0)
	@test isapprox(o2, o0)
	@test isapprox(o3, o0)
	@test isapprox(o4, o0)
	@test isapprox(o5, o0)

	true
end

# 2D small test
function dtft_test2( ; N::Dims = (2^3,2^2))
	x = randn(ComplexF64, N)
	w1 = (2*pi) * (0:N[1]-1) / N[1]
	w2 = (2*pi) * (0:N[2]-1) / N[2]
	w1 = repeat(w1, 1, N[2])
	w2 = repeat(w2', N[1], 1)
	w = [vec(w1) vec(w2)]
	d = dtft_init(w, N)

	o0 = fft(x)
#	o1 = dtft_loop_n(w, x)
	o1 = d.dtft(x)
	o1 = reshape(o1, N)
	@test isapprox(o1, o0)
	true
end

# 1D verify consistency
function dtft_test1c( ; N::Int=2^10, M::Int=2^11, n_shift::Real=7)
	x = randn(ComplexF64, N)
	w = randn(M)

	o1 = dtft_loop_n(w, x ; n_shift=n_shift)
	o2 = dtft_loop_m(w, x ; n_shift=n_shift)
	o3 = dtft_dist_m(w, x ; n_shift=n_shift)
	o4 = dtft_pmap_m(w, x ; n_shift=n_shift)
	o5 = dtft_matvec(w, x ; n_shift=n_shift)

	@test isapprox(o2, o1)
	@test isapprox(o3, o1)
	@test isapprox(o4, o1)
	@test isapprox(o5, o1)

	d = dtft_init(w, N ; n_shift=n_shift)
	o6 = d.dtft(x)
	@test isapprox(o6, o1)

	o7 = d.A * x
	@test isapprox(o7, o1)

	b1 = dtft_adj(w, o1, N ; n_shift=n_shift)
	b2 = d.adjoint(o1)
	@test isapprox(b2, b1)
	b3 = d.A' * o1
	@test isapprox(b3, b1)

	A = d.A
	@test A.name == "dtft1"
	@test A.N == (N,)

	true
end

# 2D verify consistency
function dtft_test2c( ;
	M::Int=2^9,
	N::Dims=(2^6,2^4),
	n_shift::AbstractVector{<:Real} = zeros(Int, length(N)),
)

	x = randn(ComplexF64, N)
	w = (rand(M,2) .- 0.5) * 2 * pi

	o1 = dtft_loop_n(w, x ; n_shift=n_shift)
	sd = dtft_init(w, N ; n_shift=n_shift)
	o2 = sd.dtft(x)
	@test isequal(o2, o1)
	o3 = sd.A * x
	@test isequal(o3, o1)

	b1 = dtft_adj(w, o1, N ; n_shift=n_shift)
	b2 = sd.adjoint(o1)
	@test isequal(b2, b1)

	b3 = sd.A' * o1
	@test isequal(b3, b1)

	A = sd.A
	@test A.name == "dtft2"
	@test A.N == N
	@test A isa LinearMapAO
	true
end

"""
dtft_test1_adj()
test adjoint
"""
function dtft_test1_adj( ; N::Int=20, M::Int=30, n_shift::Int=5)
	w = randn(M)
	_, _, A = dtft_init(w, N ; n_shift = n_shift)
	@test Matrix(A)' ≈ Matrix(A')
	true
end

"""
dtft_test2_adj()
test adjoint
"""
function dtft_test2_adj( ;
	M::Int=2^4,
	N::Dims=(2^3,2^2),
	n_shift::AbstractVector{<:Real} = [6,1.7],
)
	w = (rand(M,2) .- 0.5) * 2 * pi
	sd = dtft_init(w, N ; n_shift=n_shift)
	A = sd.A
	@test A isa LinearMapAO
	@test Matrix(A)' ≈ Matrix(A')
	true
end

@testset "1d" begin
	@test dtft_test1()
	@test dtft_test1c()
	@test dtft_test1_adj()
end

@testset "2d" begin
	@test dtft_test2()
	@test dtft_test2c()
	@test dtft_test2_adj()
end


log_temperature(t::Real) = 1 / log(t)

constant_temperature(t::Real) = 1.0

function default_neighbor!(x::Array, x_proposal::Array)
    @assert size(x) == size(x_proposal)
    for i in 1:length(x)
        @inbounds x_proposal[i] = x[i] + randn()
    end
    return
end

immutable SimulatedAnnealing <: Optimizer
    neighbor!::Function
    temperature::Function
    keep_best::Bool # not used!?
end

SimulatedAnnealing(; neighbor!::Function = default_neighbor!,
                     temperature::Function = log_temperature,
                     keep_best::Bool = true) =
  SimulatedAnnealing(neighbor!, temperature, keep_best)

type SimulatedAnnealingState{T}
    @add_generic_fields()
    iteration::Int64
    x_current::Array{T}
    x_proposal
    f_x_current::T
    f_proposal::T
end

initial_state(method::SimulatedAnnealing, options, d, initial_x::Array) = initial_state(method, options, d.f, initial_x)

function initial_state{T}(method::SimulatedAnnealing, options, f::Function, initial_x::Array{T})
    # Count number of parameters
    n = length(initial_x)

    # Store f(x) in f_x
    f_x = f(initial_x)
    f_calls = 1

    # Store the best state ever visited
    best_x = copy(initial_x)
    best_f_x = f_x
    SimulatedAnnealingState("Simulated Annealing", n, copy(initial_x), f_x, f_calls, 0, 0, 1, copy(initial_x), best_x, best_f_x, f_x)
end

update_state!(d, state::SimulatedAnnealingState, method::SimulatedAnnealing) = update_state!(d.f, state, method)
function update_state!{T}(f::Function, state::SimulatedAnnealingState{T}, method::SimulatedAnnealing)

    # Determine the temperature for current iteration
    t = method.temperature(state.iteration)

    # Randomly generate a neighbor of our current state
    method.neighbor!(state.x_current, state.x_proposal)

    # Evaluate the cost function at the proposed state
    state.f_proposal = f(state.x_proposal)
    state.f_calls += 1

    if state.f_proposal <= state.f_x_current
        # If proposal is superior, we always move to it
        copy!(state.x_current, state.x_proposal)
        state.f_x_current = state.f_proposal

        # If the new state is the best state yet, keep a record of it
        if state.f_proposal < state.f_x
            state.f_x = state.f_proposal
            copy!(state.x, state.x_proposal)
        end
    else
        # If proposal is inferior, we move to it with probability p
        p = exp(-(state.f_proposal - state.f_x_current) / t)
        if rand() <= p
            copy!(state.x_current, state.x_proposal)
            state.f_x_current = state.f_proposal
        end
    end

    state.iteration += 1
    false
end


# do some profiling
using Plots, LaTeXStrings
gr(colorbar=false,size=(600,150),legend=false,grid=false,xticks=false,yticks=false,axis=false)
msg =:mediumseagreen

using Revise, FourierGPE

# ==== Initialize simulation
L = (60.0,)
N = (512,)
sim = Sim(L,N)
@unpack_Sim sim

# ==== set simulation parameters
μ = 25.0
g = -0.01
γ = 0.0
Ns = 200
ξs = 2/abs(g)/Ns
us = 20
tf = π |> Float64
Nt = 150
ti = 0.0
t = LinRange(ti,tf,Nt)

# ==== soliton wavefunction
x = X[1]
ψs(x) = sqrt(Ns/2ξs)*sech(x/ξs)*exp(im*us*x)
ψi = ψs.(x)
ϕi = kspace(ψi,sim)

# ==== Set all fields
@pack_Sim! sim

# ====== Evolve in k space ==========
sol = runsim(sim)
# ===================================

import FourierGPE:Lgp!, nlin!, V
V(x,t) = 0.0

function nlin!(dϕ,ϕ,sim::Sim{1},t)
    @unpack g,X,V0 = sim; x = X[1]
    dϕ .= ϕ
    xspace!(dϕ,sim)
    @. dϕ *= V0 + V(x,[t]) + g*abs2(dϕ)
    kspace!(dϕ,sim)
    return nothing
end

function Lgp!(dϕ,ϕ,sim,t)
    @unpack γ,μ,espec = sim
    nlin!(dϕ,ϕ,sim,t)
    @. dϕ = -im*(1.0 - im*γ)*(dϕ + (espec - μ)*ϕ)
    return nothing
end

using BenchmarkTools

dϕi = copy(ϕi)
@btime Lgp!(dϕi,ϕi,sim,t)
# factor of 10 speed up in MKL in 1D
# if there isn't a speed up, there must be another bottleneck

@btime kspace!(dϕi,sim)
@btime xspace!(dϕi,sim)
@btime @. dϕi *= g*abs2(dϕ)

# ===============================
# 2d profiling: vortex precession example

using Revise, FourierGPE

# ==== Initialize simulation
L = (18.0,18.0)
N = (128,128)
sim = Sim(L,N)
@unpack_Sim sim

# ==== set simulation parameters
μ = 12.0

# Time dependent potential function (here trivial t dep)
import FourierGPE.V
V(x,y,t) = 0.5*(x^2 + y^2)
ψ0(x,y,μ,g) = sqrt(μ/g)*sqrt(max(1.0-V(x,y,0.0)/μ,0.0)+im*0.0)

# ==== make initial state
x,y = X
ψi = ψ0.(x,y',μ,g)
ϕi = kspace(ψi,sim)
@pack_Sim! sim

# ==== evolve
@time sol = runsim(sim)

# ==== ground state

ϕg = sol[end]
ψg = xspace(ϕg,sim)
showpsi(x,y,ψg)

# ==== set simulation parameters
γ = 0.0
t = LinRange(ti,tf,Nt)
ϕi = kspace(ψg,sim)
# reltol = 1e-7
# alg = DP5()ß

#---
using VortexDistributions
R(w) = sqrt(2*μ/w^2)
R(1)
rv = 3.
healinglength(x,y,μ,g) = 1/sqrt(g*abs2(ψ0(x,y,μ,g)))
ξ0 = healinglength.(0.,0.,μ,g)
ξ = healinglength(rv,0.,μ,g)
#---

pv = PointVortex(rv,0.,1)
vi = ScalarVortex(ξ,pv)
psi = Torus(copy(ψg),x,y)
vortex!(psi,vi)
showpsi(x,y,psi.ψ)
ψi .= psi.ψ
ϕi = kspace(ψi,sim)

# ===== compare with Fetter JLTP 2010
ξ = 1/sqrt(μ)
Rtf = R(1)
Ωm = 3*log(Rtf/ξ/sqrt(2))/2/Rtf^2
Ωv = Ωm/(1-rv^2/Rtf^2)

# or a precession period
Tv = 2*pi/Ωv

ti = 0.; tf = 0.1*Tv
t = LinRange(ti,tf,Nt)

@pack_Sim! sim
#---
# ==== evolve
@time solv = runsim(sim)

# FFTW 1.8s
# MKL 1.1s

import FourierGPE:Lgp!, nlin!, V
V(x,t) = 0.0

function nlin!(dϕ,ϕ,sim::Sim{1},t)
    @unpack g,X,V0 = sim; x = X[1]
    dϕ .= ϕ
    xspace!(dϕ,sim)
    @. dϕ *= V0 + V(x,[t]) + g*abs2(dϕ)
    kspace!(dϕ,sim)
    return nothing
end

function Lgp!(dϕ,ϕ,sim,t)
    @unpack γ,μ,espec = sim
    nlin!(dϕ,ϕ,sim,t)
    @. dϕ = -im*(1.0 - im*γ)*(dϕ + (espec - μ)*ϕ)
    return nothing
end

using BenchmarkTools

dϕi = copy(ϕi)
@btime Lgp!(dϕi,ϕi,sim,t)
# factor of ?? speed up in MKL in 2D
# if there isn't a speed up, there must be another bottleneck

@btime kspace!(dϕi,sim)
@btime xspace!(dϕi,sim)
@btime @. dϕi *= g*abs2(dϕ)


using StatsBase
using Test, Random, Statistics, LinearAlgebra
using Missings

const ≅ = isequal

Random.seed!(1)

# to avoid using specialized method
arbitrary_fun(x, y) = cor(x, y)

@testset "pairwise and pairwise! with $f" for f in (arbitrary_fun, cor, cov)
    @testset "basic interface" begin
        x = [rand(10) for _ in 1:4]
        y = [rand(Float32, 10) for _ in 1:5]
        # to test case where inference of returned eltype fails
        z = [Vector{Any}(rand(Float32, 10)) for _ in 1:5]

        res = @inferred pairwise(f, x, y)
        @test res isa Matrix{Float64}
        res2 = zeros(Float64, size(res))
        @test pairwise!(f, res2, x, y) === res2
        @test res == res2 == [f(xi, yi) for xi in x, yi in y]

        res = pairwise(f, y, z)
        @test res isa Matrix{Float32}
        res2 = zeros(Float32, size(res))
        @test pairwise!(f, res2, y, z) === res2
        @test res == res2 == [f(yi, zi) for yi in y, zi in z]

        res = pairwise(f, Any[[1.0, 2.0, 3.0], [1.0f0, 3.0f0, 10.5f0]])
        @test res isa Matrix{Float64}
        res2 = zeros(AbstractFloat, size(res))
        @test pairwise!(f, res2, Any[[1.0, 2.0, 3.0], [1.0f0, 3.0f0, 10.5f0]]) === res2
        @test res == res2 ==
            [f(xi, yi) for xi in ([1.0, 2.0, 3.0], [1.0f0, 3.0f0, 10.5f0]),
                           yi in ([1.0, 2.0, 3.0], [1.0f0, 3.0f0, 10.5f0])]
        @test res isa Matrix{Float64}

        @inferred pairwise(f, x, y)

        @test_throws ArgumentError pairwise(f, [Int[]], [Int[]])
        @test_throws ArgumentError pairwise!(f, zeros(1, 1), [Int[]], [Int[]])

        res = pairwise(f, [], [])
        @test size(res) == (0, 0)
        @test res isa Matrix{Any}
        res2 = zeros(0, 0)
        @test pairwise!(f, res2, [], []) === res2

        res = pairwise(f, Vector{Int}[], Vector{Int}[])
        @test size(res) == (0, 0)
        @test res isa Matrix{Float64}
        res2 = zeros(0, 0)
        @test pairwise!(f, res2, Vector{Int}[], Vector{Int}[]) === res2

        res = pairwise(f, [[1, 2]], Vector{Int}[])
        @test size(res) == (1, 0)
        @test res isa Matrix{Float64}
        res2 = zeros(1, 0)
        @test pairwise!(f, res2, [[1, 2]], Vector{Int}[]) === res2

        res = pairwise(f, Vector{Int}[], [[1, 2], [2, 3]])
        @test size(res) == (0, 2)
        @test res isa Matrix{Float64}
        res2 = zeros(0, 2)
        @test pairwise!(f, res2, [], [[1, 2], [2, 3]]) === res2

        @test_throws DimensionMismatch pairwise!(f, zeros(1, 2), x, y)
        @test_throws DimensionMismatch pairwise!(f, zeros(1, 2), [], [])
        @test_throws DimensionMismatch pairwise!(f, zeros(0, 0),
                                                 [], [[1, 2], [2, 3]])
    end

    @testset "missing values handling interface" begin
        xm = [ifelse.(rand(100) .> 0.9, missing, rand(100)) for _ in 1:4]
        ym = [ifelse.(rand(100) .> 0.9, missing, rand(Float32, 100)) for _ in 1:4]
        zm = [ifelse.(rand(100) .> 0.9, missing, rand(Float32, 100)) for _ in 1:4]

        res = pairwise(f, xm, ym)
        @test res isa Matrix{Missing}
        res2 = zeros(Union{Float64, Missing}, size(res))
        @test pairwise!(f, res2, xm, ym) === res2
        @test res ≅ res2 ≅ [missing for xi in xm, yi in ym]

        res = pairwise(f, xm, ym, skipmissing=:pairwise)
        @test res isa Matrix{Float64}
        res2 = zeros(Union{Float64, Missing}, size(res))
        @test pairwise!(f, res2, xm, ym, skipmissing=:pairwise) === res2
        @test res ≅ res2
        @test isapprox(res, [f(collect.(skipmissings(xi, yi))...) for xi in xm, yi in ym],
                       rtol=1e-6)

        res = pairwise(f, ym, zm, skipmissing=:pairwise)
        @test res isa Matrix{Float32}
        res2 = zeros(Union{Float32, Missing}, size(res))
        @test pairwise!(f, res2, ym, zm, skipmissing=:pairwise) === res2
        @test res ≅ res2
        @test isapprox(res, [f(collect.(skipmissings(yi, zi))...) for yi in ym, zi in zm],
                       rtol=1e-6)

        nminds = mapreduce(x -> .!ismissing.(x),
                        (x, y) -> x .& y,
                        [xm; ym])
        res = pairwise(f, xm, ym, skipmissing=:listwise)
        @test res isa Matrix{Float64}
        res2 = zeros(Union{Float64, Missing}, size(res))
        @test pairwise!(f, res2, xm, ym, skipmissing=:listwise) === res2
        @test res ≅ res2
        @test isapprox(res, [f(view(xi, nminds), view(yi, nminds)) for xi in xm, yi in ym],
                       rtol=1e-6)

        if VERSION >= v"1.6.0-DEV"
            # inference of cor fails so use an inferrable function
            # to check that pairwise itself is inferrable
            for skipmissing in (:none, :pairwise, :listwise)
                g(x, y=x) = pairwise((x, y) -> x[1] * y[1], x, y, skipmissing=skipmissing)
                @test Core.Compiler.return_type(g, Tuple{Vector{Vector{Union{Float64, Missing}}}}) ==
                    Core.Compiler.return_type(g, Tuple{Vector{Vector{Union{Float64, Missing}}},
                                                    Vector{Vector{Union{Float64, Missing}}}}) ==
                    Matrix{<: Union{Float64, Missing}}
                if skipmissing in (:pairwise, :listwise)
                    @test_broken Core.Compiler.return_type(g, Tuple{Vector{Vector{Union{Float64, Missing}}}}) ==
                        Core.Compiler.return_type(g, Tuple{Vector{Vector{Union{Float64, Missing}}},
                                                        Vector{Vector{Union{Float64, Missing}}}}) ==
                        Matrix{Float64}
                end
            end
        end

        @test_throws ArgumentError pairwise(f, xm, ym, skipmissing=:something)
        @test_throws ArgumentError pairwise!(f, zeros(Union{Float64, Missing},
                                                      length(xm), length(ym)), xm, ym,
                                             skipmissing=:something)

        # variable with only missings
        xm = [fill(missing, 10), rand(10)]
        ym = [rand(10), rand(10)]

        res = pairwise(f, xm, ym)
        @test res isa Matrix{Union{Float64, Missing}}
        res2 = zeros(Union{Float64, Missing}, size(res))
        @test pairwise!(f, res2, xm, ym) === res2
        @test res ≅ res2 ≅ [f(xi, yi) for xi in xm, yi in ym]

        if VERSION >= v"1.5" # Fails with UndefVarError on Julia 1.0
            @test_throws ArgumentError pairwise(f, xm, ym, skipmissing=:pairwise)
            @test_throws ArgumentError pairwise(f, xm, ym, skipmissing=:listwise)

            res = zeros(Union{Float64, Missing}, length(xm), length(ym))
            @test_throws ArgumentError pairwise!(f, res, xm, ym, skipmissing=:pairwise)
            @test_throws ArgumentError pairwise!(f, res, xm, ym, skipmissing=:listwise)
        end

        for sm in (:pairwise, :listwise)
            @test_throws ArgumentError pairwise(f, [[1, 2]], [1], skipmissing=sm)
            @test_throws ArgumentError pairwise(f, [1], [[1, 2]], skipmissing=sm)
            @test_throws ArgumentError pairwise(f, [1], [1], skipmissing=sm)
        end
    end

    @testset "iterators" begin
        x = (v for v in [rand(10) for _ in 1:4])
        y = (v for v in [rand(10) for _ in 1:4])

        res = @inferred pairwise(f, x, y)
        res2 = zeros(size(res))
        @test pairwise!(f, res2, x, y) === res2
        @test res == res2 == pairwise(f, collect(x), collect(y))

        res = @inferred(pairwise(f, x))
        res2 = zeros(size(res))
        @test pairwise!(f, res2, x) === res2
        @test res == res2 == pairwise(f, collect(x))
    end

    @testset "non-vector entries" begin
        x = (Iterators.drop(v, 1) for v in [rand(10) for _ in 1:4])
        y = (Iterators.drop(v, 1) for v in [rand(10) for _ in 1:4])

        @test pairwise((x, y) -> f(collect(x), collect(y)), x, y) ==
            [f(collect(xi), collect(yi)) for xi in x, yi in y]
        @test pairwise((x, y) -> f(collect(x), collect(y)), x) ==
            [f(collect(xi1), collect(xi2)) for xi1 in x, xi2 in x]
        @test_throws ArgumentError pairwise((x, y) -> f(collect(x), collect(y)), x, y,
                                            skipmissing=:pairwise)
        @test_throws ArgumentError pairwise((x, y) -> f(collect(x), collect(y)), x, y,
                                            skipmissing=:listwise)
    end

    @testset "two-argument method" begin
        x = [rand(10) for _ in 1:4]
        res = pairwise(f, x)
        res2 = zeros(size(res))
        @test pairwise!(f, res2, x) === res2
        @test res == res2 == pairwise(f, x, x)
    end

    @testset "symmetric" begin
        x = [rand(10) for _ in 1:4]
        y = [rand(10) for _ in 1:4]

        @test pairwise(f, x, x, symmetric=true) ==
            pairwise(f, x, symmetric=true) ==
            Symmetric(pairwise(f, x, x), :U)

        res = zeros(4, 4)
        res2 = zeros(4, 4)
        @test pairwise!(f, res, x, x, symmetric=true) === res
        @test pairwise!(f, res2, x, symmetric=true) === res2
        @test res == res2 == Symmetric(pairwise(f, x, x), :U)

        @test_throws ArgumentError pairwise(f, x, y, symmetric=true)
        @test_throws ArgumentError pairwise!(f, res, x, y, symmetric=true)
    end

    @testset "cor corner cases" begin
        # Integer inputs must give a Float64 output
        res = pairwise(cor, [[1, 2, 3], [1, 5, 2]])
        @test res isa Matrix{Float64}
        @test res == [cor(xi, yi) for xi in ([1, 2, 3], [1, 5, 2]),
                                      yi in ([1, 2, 3], [1, 5, 2])]

        # NaNs are ignored for the diagonal
        res = pairwise(cor, [[1, 2, NaN], [1, 5, 2]])
        @test res isa Matrix{Float64}
        @test res ≅ [1.0 NaN
                     NaN 1.0]

        # missings are ignored for the diagonal
        res = pairwise(cor, [[1, 2, 7], [1, 5, missing]])
        @test res isa Matrix{Union{Float64, Missing}}
        @test res ≅ [1.0 missing
                     missing 1.0]
        res = pairwise(cor, Vector{Union{Int, Missing}}[[missing, missing, missing],
                                                        [missing, missing, missing]])
        @test res isa Matrix{Union{Float64, Missing}}
        @test res ≅ [1.0 missing
                     missing 1.0]
        if VERSION >= v"1.5"
            # except when eltype is Missing
            res = pairwise(cor, [[missing, missing, missing],
                                 [missing, missing, missing]])
            @test res isa Matrix{Missing}
            @test res ≅ [missing missing
                         missing missing]
        end

        for sm in (:pairwise, :listwise)
            res = pairwise(cor, [[1, 2, NaN, 4], [1, 5, 5, missing]], skipmissing=sm)
            @test res isa Matrix{Float64}
            @test res ≅ [1.0 NaN
                         NaN 1.0]
            if VERSION >= v"1.5"
                @test_throws ArgumentError pairwise(cor, [[missing, missing, missing],
                                                          [missing, missing, missing]],
                                                    skipmissing=sm)
            end
        end
    end
end


@testset "simplify" begin

    @test simplify(:(10 - 5)) == 5
    @test simplify(:(u - 2)) == :(u - 2)
    @test simplify(:(1 * (u - 2))) == :(u - 2)
    @test simplify(:(1.0 * cos(x1))) == :(cos(x1))
    @test simplify(:(2 * 3x)) == :(6x)
    @test simplify(:(x ^ 0 / 2)) == 1/2

end


module  ChernovA1P2

#this is useful if we have mdoules
if pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

using DataFrames, Distributions, StatsBase, GZip, JLD,
  Gadfly, CTModCopy, JuMP, NLopt, ForwardDiff

const DATA_PATH = pwd() * "\\data" #path where we store the data
const OUTPUT_PATH = pwd() * "\\output"
const DATE_FORMAT_STR = "m/d/y"
const SP500_NAME = "GSPC"

const DEF_NUM_MOMENTS = 6

function preProcessSP500()::Void
  dateFormat::String = DATE_FORMAT_STR
  sp500DF::DataFrame = readtable("$DATA_PATH\\$SP500_NAME.csv")

  #fix the dates
  rename!(sp500DF, :Date, :DATE)
  sp500DF[:,:DATE] = ((s::String)->Date(s, dateFormat)::Date).(sp500DF[:,:DATE])::DataVector{Date}
  sort!(sp500DF, cols = [:DATE])
  sp500DF[:,:returns] = similar(sp500DF[:,:Adj_Close])
  sp500DF[2:end,:returns] = log.(sp500DF[2:end,:Adj_Close] ./  sp500DF[1:(end-1),:Adj_Close])
  #showcols(sp500DF)

  stream::IOStream = open("$DATA_PATH\\$SP500_NAME.jls", "w")
  serialize(stream, sp500DF)
  close(stream)

  return nothing
end

#this function gets the sample moments of a given vector of data
function getSampleMoments!(x::Vector{Float64}, N::Int, moments::Vector{Float64} = Vector{Float64}(N))

  #for each moment we are getting
  for i::Int ∈ 1:N
    moments[i] = sum(x.^Float64(i))/Float64(length(x))
  end

  println(moments)
  println("Var: ")

  return moments
end

function getCumulantsOfJumpNorm(α::Real, σ²::Real, θ::Real, δ²::Real, ω::Real)::Vector{Real}
  κ::Vector{Real} = Vector{Real}(DEF_NUM_MOMENTS)

  #δ²::Real = δ*δ #helpers
  θ²::Real = θ*θ
  θ³::Real = θ*θ²

  κ[1] = α + ω*θ
  κ[2] = σ²+ω*(δ²+θ²) #σ^2 +ω*(δ^2+θ^2)
  κ[3] = 3.*ω*δ²*θ+ω*θ³ #3*ω*δ^2*θ+ω*θ^3
  κ[4] = ω*(6.*δ²*θ²+θ²*θ²+3.*δ²*δ²) #6*ω*δ^2*θ^2+ω*θ^4 +3*ω*δ^4
  κ[5] = ω*(15.*θ*δ²*δ²+10.*δ²*θ³+θ³*θ²) #θ*ω*(15*δ^4+10*δ^2*θ^2+θ^4)
  κ[6] = ω*(15.*δ²*δ²*δ²+45.*δ²*δ²*θ²+15.*δ²*θ²*θ²+θ³*θ³) #ω*(15*δ^6+45*δ^4*θ^2+15*δ^2*θ^4+θ^6)

  return κ
end

function get6CumulantMoments(κ...)::Vector{Real}
  μ::Vector{Real} = Vector{Real}(DEF_NUM_MOMENTS)

  κ1²::Real = κ[1] * κ[1]
  κ1³::Real = κ[1] * κ1²
  κ2²::Real = κ[2] * κ[2]

  μ[1] = κ[1]
  μ[2] = κ[2] + κ1² #κ[2]+κ[1]^2
  μ[3] = κ[3] + 3.*κ[2]*κ[1] + κ1³
  μ[4] = κ[4] + 4.*κ[3]*κ[1] + 3.*κ2² + 6.*κ[2]*κ1²+κ1²*κ1²
  μ[5] = κ[5] + 5.*κ[4]*κ[1] + 10.*κ[3]*κ[2] + 10.*κ[3]*κ1² + 15.*κ2² *κ[1] + 10.*κ[2]*κ1³+κ1³*κ1²
  μ[6] = κ[6] + 6.*κ[5]*κ[1] + 15.*κ[4]*κ[2] + 15.*κ[4]*κ1² + 10.*κ[3]*κ[3]+60*κ[3]*κ[2]*κ[1] +
    20.*κ[3]*κ1³+15.*κ2²*κ[2] + 45.*κ2²*κ1² + 15.*κ[2]*κ1²*κ1²+κ1³*κ1³

  return μ
end

function getParametersOfJumpNorm(x::Vector{Float64})::Void

  numMoments::Int = DEF_NUM_MOMENTS

  #get the sample moments
  M::Vector{Float64} = getSampleMoments!(x, numMoments)

  #this is our objective function
  function obj(α::Real, σ²::Real, θ::Real, δ²::Real, ω::Real)::Real

    κ::Vector{Real} = getCumulantsOfJumpNorm(α, σ², θ, δ², ω)
    μ::Vector{Real} = get6CumulantMoments(κ...)
    d::Real = sum((M .- μ) .* (M .- μ))

    return d
  end

  αBest::Float64=0.0; σ²Best::Float64=0.0; θBest::Float64=0.0
  δ²Best::Float64=0.0; ωBest::Float64=0.0; dBest::Float64=100.0

  solveFails::Int = 0
  solveAttempts::Int = 0
  failed::Bool = false

  #encapsulate the routing for searching a grid
  #Algorithms: :LD_LBFGS, :LN_COBYLA, :LN_BOBYQA
  #  :LD_SLSQP, :LD_TNEWTON, :LD_VAR2/:LD_VAR1,
  #  :LD_MMA, :GN_ESCH, :GN_DIRECT, GN_DIRECT_L, :GD_STOGO
  # GN_ISRES, :LN_NELDERMEAD, :LN_SBPLX,
  #GN_CRS2_LM
  function solveProblem(αMin::Float64, σ²Min::Float64, θMin::Float64, δ²Min::Float64, ωMin::Float64, Δ::Float64)
    solveAttempts += 1
    mod = Model(solver=NLoptSolver(algorithm=:LD_SLSQP, maxtime=1))

    #JuMP.build(mod)
    #maxtime!(internalmodel(mod), 1.0)

    @variable(mod, αMin <= α <= αMin + Δ)
    @variable(mod, σ²Min <= σ² <= σ²Min + Δ)
    @variable(mod, θMin <= θ <= θMin + Δ)
    @variable(mod, δ²Min <= δ² <= δ²Min + Δ)
    @variable(mod, ωMin <= ω <= ωMin + Δ)

    #register the outside functions
    JuMP.register(mod, :getCumulantsOfJumpNorm, 5, getCumulantsOfJumpNorm, autodiff=true)
    JuMP.register(mod, :get6CumulantMoments, 6, get6CumulantMoments, autodiff=true)
    JuMP.register(mod, :obj, 5, obj, autodiff=true)

    @NLobjective(mod, Min, obj(α, σ², θ, δ², ω))

    #set the starting point at the midpoint of the grid
    setvalue(α , (αMin + Δ/2.0))
    setvalue(σ², (σ²Min + Δ/2.0))
    setvalue(θ , (θMin + Δ/2.0))
    setvalue(δ², (δ²Min + Δ/2.0))
    setvalue(ω , (ωMin + Δ/2.0))

    #indicator
    failed = false
    try #solve the model
      output = solve(mod)
    catch
      solveFails+=1
      failed = true
    end

    objValue::Float64 = obj(getvalue(α), getvalue(σ²), getvalue(θ), getvalue(δ²), getvalue(ω))
    if !failed && objValue<dBest #if both statements are true, update the best parameters
      dBest = objValue
      αBest, σ²Best, θBest, δ²Best, ωBest = getvalue(α), getvalue(σ²), getvalue(θ), getvalue(δ²), getvalue(ω)
    end
  end

  #iterate the parameter space
  Δ::Float64=0.1
  @fastmath for αTry::Float64 ∈ 0.0:Δ:(0.1-Δ),
     σ²Try::Float64 ∈ 0.0:Δ:(0.2-Δ),
      θTry::Float64 ∈ -1.0:Δ:(0.0-Δ),
      δ²Try::Float64 ∈ 0.0:Δ:(0.2-Δ),
      ωTry::Float64 ∈ 0.0:Δ:(1.0-Δ)
        solveProblem(αTry, σ²Try, θTry, δ²Try, ωTry, Δ)
  end

  #= Output code here
  println("Trys: $solveAttempts")
  println("Fails: $solveFails")
  println("α=$(αBest), σ²=$(σ²Best), θ=$(θBest), δ²=$(δ²Best), ω=$(ωBest)")
  println("Also: σ=$(σ²Best^0.5) annualized: $((σ²Best)^0.5*255.0^0.5)")
  println("Objective: $dBest")
  println("For Reference: x̄=$(mean(x)), std(x)= $(std(x))")
  =#

  return nothing
end

#the main method for running the optimziation routines
function problemScript(;refreshData::Bool = true)::Void
  if !isfile("$DATA_PATH\\$SP500_NAME.jls") || refreshData
    preProcessSP500()
  end

  #read the binary file
  stream::IOStream = open("$DATA_PATH\\$SP500_NAME.jls")
  sp500DF::DataFrame = deserialize(stream)
  close(stream)

  x::Vector{Float64} = sp500DF[2:end, :returns]
  getParametersOfJumpNorm(x)

  return nothing
end

@time problemScript()
end


# ---
# title: 675. Cut Off Trees for Golf Event
# id: problem675
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Breadth-first Search
# link: <https://leetcode.com/problems/cut-off-trees-for-golf-event/description/>
# hidden: true
# ---
#
# You are asked to cut off all the trees in a forest for a golf event. The
# forest is represented as an `m x n` matrix. In this matrix:
#
#   * `0` means the cell cannot be walked through.
#   * `1` represents an empty cell that can be walked through.
#   * A number greater than `1` represents a tree in a cell that can be walked through, and this number is the tree's height.
#
# In one step, you can walk in any of the four directions: north, east, south,
# and west. If you are standing in a cell with a tree, you can choose whether to
# cut it off.
#
# You must cut off the trees in order from shortest to tallest. When you cut off
# a tree, the value at its cell becomes `1` (an empty cell).
#
# Starting from the point `(0, 0)`, return _the minimum steps you need to walk
# to cut off all the trees_. If you cannot cut off all the trees, return `-1`.
#
# You are guaranteed that no two trees have the same height, and there is at
# least one tree needs to be cut off.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/11/26/trees1.jpg)
#
#
#
#     Input: forest = [[1,2,3],[0,0,4],[7,6,5]]
#     Output: 6
#     Explanation: Following the path above allows you to cut off the trees from shortest to tallest in 6 steps.
#
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2020/11/26/trees2.jpg)
#
#
#
#     Input: forest = [[1,2,3],[0,0,0],[7,6,5]]
#     Output: -1
#     Explanation: The trees in the bottom row cannot be accessed as the middle row is blocked.
#
#
# **Example 3:**
#
#
#
#     Input: forest = [[2,3,4],[0,0,5],[8,7,6]]
#     Output: 6
#     Explanation: You can follow the same path as Example 1 to cut off all the trees.
#     Note that you can cut off the first tree at (0, 0) before making any steps.
#
#
#
#
# **Constraints:**
#
#   * `m == forest.length`
#   * `n == forest[i].length`
#   * `1 <= m, n <= 50`
#   * `0 <= forest[i][j] <= 109`
#
#
## @lc code=start
using LeetCode

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


using TikzGraphs
using TikzPictures
using GraphRecipes
using LightGraphs
using Plots

default(size=(1000, 1000))

f = :(x * y + sin(x))

vertices = Dict([
1 => (name = "x", 			x = 1, y = 0, color = :lightblue, ),
2 => (name = "y", 			x = 0, y = 1, color = :lightblue, ),
3 => (name = "h₁ = sin", 	x = 1, y = 2, color = :lightblue, ),
4 => (name = "h₂ = *", 		x = 0, y = 3, color = :lightblue, ),
5 => (name = "h₃ = +", 		x = 1, y = 4, color = :lightblue, ),
6 => (name = "z", 			x = 1, y = 5, color = :lightblue, ),
7 => (name = "∂z/∂h₃", 		x = 2.5, y = 4, color = :orange, ),
8 => (name = "∂h₃/∂h₂", 	x = 2.5, y = 3, color = :orange, ),
9 => (name = "∂z/∂h₂", 		x = 3.5, y = 3, color = :orange, ),
10 => (name = "∂h₃/∂h₁", 	x = 2.5, y = 2, color = :orange, ),
11 => (name = "∂z/∂h₁", 	x = 3.5, y = 2, color = :orange, ),
12 => (name = "∂h₂/∂y", 	x = 4.5, y = 1, color = :orange, ),
13 => (name = "∂h₂/∂x + ∂h₁/∂x", x = 4.5, y = 0, color = :orange, ),
14 => (name = "∂z/∂x", x = 6, y = 0, color = :orange, ),
])

n = length(vertices)
g = LightGraphs.DiGraph(n)
add_edge!(g, 1, 3)  # x -> sin
add_edge!(g, 1, 4)	# x ->  *
add_edge!(g, 2, 4)	# y ->  *
add_edge!(g, 3, 5)	# sin(x) -> sin(x) + x*y
add_edge!(g, 4, 5)  # x*y -> sin(x) + x*y
add_edge!(g, 5, 6)  # sin(x) + x*y -> z

#add_edge!(g, 5, 7,)	# ∂z/∂h₃ -> z
#add_edge!(g, 4, 8,)	# ∂h₃/∂h₂ -> h₂ = *
#add_edge!(g, 8, 9,)	# ∂z/∂h₂ -> ∂h₃/∂h₂
#add_edge!(g, 7, 9,)	# ∂z/∂h₂ -> ∂h₃/∂h₂
#add_edge!(g, 3, 10,)	# ∂h₃/∂h₁ -> h₁
#add_edge!(g, 10, 11,)# ∂z/∂h₁  -> ∂h₃/∂h₁
#add_edge!(g, 7, 11,) # ∂z/∂h₁  -> ∂z/∂h₃
#add_edge!(g, 2, 12,) # ∂h₂/∂y  -> y
#add_edge!(g, 9, 12,) # ∂h₂/∂y  -> y
#add_edge!(g, 1, 13,) # ∂h₂/∂x + ∂h₁/∂x  -> y
#add_edge!(g, 9, 14,) # "∂z/∂x"  -> ∂z/∂h₂
#add_edge!(g, 11, 14,) # "∂z/∂x"  -> ∂z/∂h₁
#add_edge!(g, 13, 14,) # "∂z/∂x"  -> ∂h₂/∂x + ∂h₁/∂x
# graphplot(adjacency_matrix(g),
# 	names = [vertices[i].name for i in 1:n],
# 	x = 0.25 .* [vertices[i].x for i in 1:n],
# 	y = 0.25 .* [vertices[i].y falseor i in 1:n],
# 	curves=false,
# 	markercolor = [vertices[i].color for i in 1:n],
# )

for n in [6, 7, 9, 11, 13, 14]
    #graphplot(adjacency_matrix(g)[1:n, 1:n],
	#	names = [vertices[i].name for i in 1:n],
	#	x = 0.25 .* [vertices[i].x for i in 1:n],
	#	y = 0.25 .* [vertices[i].y for i in 1:n],
	#	curves=false,
	#	markercolor = [vertices[i].color for i in 1:n],
	#)
	names = [vertices[i].name for i in 1:length(vertices)]
    t = TikzGraphs.plot(g, names)
    TikzPictures.save(SVG("graphdiff_$(n).svg"),t)
end


@testset "218.the-skyline-problem.jl" begin
    @test get_skyline([[2, 9, 10], [3, 7, 15], [5, 12, 12], [15, 20, 10], [19, 24, 8]]) ==
          [(2, 10), (3, 15), (7, 12), (12, 0), (15, 10), (20, 8), (24, 0)]
    @test get_skyline([[0, 2, 3], [2, 5, 3]]) == [(0, 3), (5, 0)]
end


using Test, OrderedCollections, HDF5, BenchmarkTools
include(joinpath(dirname(@__FILE__), "../../src/includeall.jl"))

time_results = false

@testset "Pseudo-Transient Relaxation" begin
    rp = joinpath(dirname(@__FILE__), "../reference/solve/pseudo_transient_relaxation.h5")
    η = h5read(rp, "eta")
    vₑ = h5read(rp, "ve")
    vₑ₁ = h5read(rp, "ve1")
    vₕ = h5read(rp, "vh")
    vₕ₁ = h5read(rp, "vh1")
    Mvₑ = h5read(rp, "Mve")
    Mvₕ = h5read(rp, "Mvh")
    MU = h5read(rp, "MU")
    S = h5read(rp, "S")
    Δ = h5read(rp, "Delta")
    G = h5read(rp, "G")
    stategrid = StateGrid(OrderedDict(:η => η))

    # Test copy from Princeton Initiative script
    myvₑ₁ = upwind_parabolic_pde(η, Mvₑ, MU, S.^2, G, vₑ, Δ)
    myvₕ₁ = upwind_parabolic_pde(η, Mvₕ, MU, S.^2, G, vₕ, Δ)
    @test myvₑ₁ ≈ vₑ₁
    @test myvₕ₁ ≈ vₕ₁

    # Test against an implementation where the operators are constructed rather than the raw LHS matrix
    N = length(η)
    DU1 = zeros(N)
    DU2 = zeros(N)
    DD1 = zeros(N)
    DD2 = zeros(N)
    D01 = zeros(N)
    D02 = zeros(N)
    Σ²0 = zeros(N)
    dX = diff(η)
    Σ²0[2:N-1] .= S[2:N-1].^2 ./ (dX[1:N-2] + dX[2:N-1]) # approx Σ² / (2 * dx): this term is the Σ²/2 coefficient
    DU1[2:N] = max.(MU[1:N - 1], 0.) ./ dX
    DU2[2:N] = Σ²0[1:N - 1] ./ dX
    DD1[1:N - 1] = max.(-MU[2:N], 0.) ./ dX
    DD2[1:N - 1] = Σ²0[2:N] ./ dX
    D01[1:N - 1] .= -DU1[2:N]
    D01[2:N] .-= DD1[1:N - 1]
    D02[1:N - 1] .= -DU2[2:N]
    D02[2:N] .-= DD2[1:N - 1]
    L1 = spdiagm(0 => D01, -1 => DD1[1:N - 1], 1 => DU1[2:N])
    L2 = spdiagm(0 => D02, -1 => DD2[1:N - 1], 1 => DU2[2:N])
    A = spdiagm(0 => (1 - Δ) .* ones(N) + Δ .* Mvₑ) - Δ .* (L1 + L2)
    directvₑ₁ = A \ ((1 - Δ) .* vₑ)
    @test myvₑ₁ ≈ directvₑ₁

    # Test new wrapper scripts utilizing DiffEqOperator for more flexible upwinding (e.g. higher-order)
    L₁ = UpwindDifference(1, 1, diff(η)[1], N)
    L₂ = CenteredDifference(2, 2, diff(η)[1], N)
    Qbc = RobinBC((0., 1., 0.), (0., 1., 0.), diff(η)[1])
    myvₑ₁ = similar(vₑ)
    myvₕ₁ = similar(vₕ)
    pseudo_transient_relaxation!(stategrid, (myvₑ₁, myvₕ₁), (vₑ, vₕ), (Mvₑ, Mvₕ),
                                 MU ./ η, (S ./ η).^2, Δ; uniform = true,
                                 L₁ = L₁, L₂ = L₂, Q = Qbc)

    @test myvₑ₁ ≈ vₑ₁
    @test myvₕ₁ ≈ vₕ₁

    L₁ = UpwindDifference(1, 1, diff(η)[1], N)
    L₂ = CenteredDifference(2, 2, diff(η)[1], N)
    Qbc = RobinBC((0., 1., 0.), (0., 1., 0.), diff(η)[1])
    myvₑ₁ = similar(vₑ)
    myvₕ₁ = similar(vₕ)
    pseudo_transient_relaxation!(stategrid, (myvₑ₁, myvₕ₁), (vₑ, vₕ), (Mvₑ, Mvₕ),
                                 MU ./ η, (S ./ η).^2, Δ; uniform = false,
                                 L₁ = L₁, L₂ = L₂, Q = Qbc)

    @test myvₑ₁ ≈ vₑ₁
    @test myvₕ₁ ≈ vₕ₁

    # Test fast default using BandedMatrices
    myvₑ₁ = similar(vₑ)
    myvₕ₁ = similar(vₕ)
    pseudo_transient_relaxation!(stategrid, (myvₑ₁, myvₕ₁), (vₑ, vₕ), (Mvₑ, Mvₕ),
                                 MU ./ η, (S ./ η).^2, Δ; uniform = true)

    @test @test_matrix_approx_eq myvₑ₁[2:end - 1] vₑ₁[2:end - 1]
    @test @test_matrix_approx_eq myvₕ₁ vₕ₁

    myvₑ₁ = similar(vₑ)
    myvₕ₁ = similar(vₕ)
    pseudo_transient_relaxation!(stategrid, (myvₑ₁, myvₕ₁), (vₑ, vₕ), (Mvₑ, Mvₕ),
                                 MU ./ η, (S ./ η).^2, Δ; uniform = false)

    @test @test_matrix_approx_eq myvₑ₁ vₑ₁
    @test @test_matrix_approx_eq myvₕ₁ vₕ₁
end

if time_results
    println("\nTiming different methods for calculating the implicit time step.\n")
    rp = joinpath(dirname(@__FILE__), "../reference/solve/pseudo_transient_relaxation.h5")
    η = h5read(rp, "eta")
    vₑ = h5read(rp, "ve")
    vₑ₁ = h5read(rp, "ve1")
    vₕ = h5read(rp, "vh")
    vₕ₁ = h5read(rp, "vh1")
    Mvₑ = h5read(rp, "Mve")
    Mvₕ = h5read(rp, "Mvh")
    MU = h5read(rp, "MU")
    S = h5read(rp, "S")
    Δ = h5read(rp, "Delta")
    G = h5read(rp, "G")

    # Test copy from Princeton Initiative script
    println("Princeton Initiative script")
    @btime begin
        upwind_parabolic_pde(η, Mvₑ, MU, S.^2, G, vₑ, Δ)
        upwind_parabolic_pde(η, Mvₕ, MU, S.^2, G, vₕ, Δ)
    end

    # Test new wrapper scripts utilizing DiffEqOperator for more flexible upwinding (e.g. higher-order)
    L₁ = UpwindDifference(1, 1, diff(η)[1], N)
    L₂ = CenteredDifference(2, 2, diff(η)[1], N)
    Qbc = RobinBC((0., 1., 0.), (0., 1., 0.), diff(η)[1])
    myvₑ₁ = similar(vₑ)
    myvₕ₁ = similar(vₕ)
    println("Using DiffEqOperator finite differencing, uniform")
    @btime begin
        pseudo_transient_relaxation!(stategrid, (myvₑ₁, myvₕ₁), (vₑ, vₕ), (Mvₑ, Mvₕ),
                                     MU ./ η, (S ./ η).^2, Δ; uniform = true,
                                     L₁ = L₁, L₂ = L₂, Q = Qbc)
    end

    L₁ = UpwindDifference(1, 1, diff(η)[1], N)
    L₂ = CenteredDifference(2, 2, diff(η)[1], N)
    Qbc = RobinBC((0., 1., 0.), (0., 1., 0.), diff(η)[1])
    myvₑ₁ = similar(vₑ)
    myvₕ₁ = similar(vₕ)
    println("Using DiffEqOperator finite differencing, non-uniform")
    @btime begin
        pseudo_transient_relaxation!(stategrid, (myvₑ₁, myvₕ₁), (vₑ, vₕ), (Mvₑ, Mvₕ),
                                     MU ./ η, (S ./ η).^2, Δ; uniform = false,
                                     L₁ = L₁, L₂ = L₂, Q = Qbc)
    end

    # Test fast default using BandedMatrices
    myvₑ₁ = similar(vₑ)
    myvₕ₁ = similar(vₕ)
    println("Using BandedMatrices, uniform")
    @btime begin
        pseudo_transient_relaxation!(stategrid, (myvₑ₁, myvₕ₁), (vₑ, vₕ), (Mvₑ, Mvₕ),
                                     MU ./ η, (S ./ η).^2, Δ; uniform = true)
    end

    myvₑ₁ = similar(vₑ)
    myvₕ₁ = similar(vₕ)
    println("Using BandedMatrices, non-uniform")
    @btime begin
        pseudo_transient_relaxation!(stategrid, (myvₑ₁, myvₕ₁), (vₑ, vₕ), (Mvₑ, Mvₕ),
                                     MU ./ η, (S ./ η).^2, Δ; uniform = false)
    end
end

nothing


# ---
# title: 1106. Parsing A Boolean Expression
# id: problem1106
# author: Indigo
# date: 2022-03-08
# difficulty: Hard
# categories: String
# link: <https://leetcode.com/problems/parsing-a-boolean-expression/description/>
# hidden: true
# ---
#
# Return the result of evaluating a given boolean `expression`, represented as a
# string.
#
# An expression can either be:
#
#   * `"t"`, evaluating to `True`;
#   * `"f"`, evaluating to `False`;
#   * `"!(expr)"`, evaluating to the logical NOT of the inner expression `expr`;
#   * `"&(expr1,expr2,...)"`, evaluating to the logical AND of 2 or more inner expressions `expr1, expr2, ...`;
#   * `"|(expr1,expr2,...)"`, evaluating to the logical OR of 2 or more inner expressions `expr1, expr2, ...`
#
#
#
# **Example 1:**
#
#
#
#     Input: expression = "!(f)"
#     Output: true
#
#
# **Example 2:**
#
#
#
#     Input: expression = "|(f,t)"
#     Output: true
#
#
# **Example 3:**
#
#
#
#     Input: expression = "&(t,f)"
#     Output: false
#
#
# **Example 4:**
#
#
#
#     Input: expression = "|(&(t,f,t),!(t))"
#     Output: false
#
#
#
#
# **Constraints:**
#
#   * `1 <= expression.length <= 20000`
#   * `expression[i]` consists of characters in `{'(', ')', '&', '|', '!', 't', 'f', ','}`.
#   * `expression` is a valid expression representing a boolean, as given in the description.
#
#
## @lc code=start
using LeetCode

function parse_bool_expr(expr::String)
    operations = Char[]
    operands = Char[]
    for ch in expr
        if ch == '!' || ch == '&' || ch == '|'
            push!(operations, ch)
        elseif ch == ')'
            cur_oprs = Char[]
            while operands[end] != '('
                push!(cur_oprs, pop!(operands))
            end
            pop!(operands)
            op = pop!(operations)
            rs = true
            if op == '|'
                rs = any(==('t'), cur_oprs)
            elseif op == '&'
                rs = all(==('t'), cur_oprs)
            else
                rs = cur_oprs[1] != 't'
            end
            push!(operands, rs ? 't' : 'f')
        elseif ch != ','
            push!(operands, ch)
        end
    end
    return operands[1] == 't' ? true : false
end
## @lc code=end


function parseInput(file)
    allNumbers::Vector{Int} = []
    open(file) do file
        for line in eachline(file)
            push!(allNumbers, parse(Int,line))
        end
    end
    return allNumbers
end

function testNumberAtIndex(allNumbers, numberIndex, preambleLength)
    result = 0
    numberToTest = allNumbers[numberIndex]
    lastIx = numberIndex-1
    startIx = numberIndex-preambleLength
    for i in startIx:1:lastIx
        for j in startIx:1:lastIx
            a = allNumbers[i]
            b = allNumbers[j]
            # println("a+b=", a+b)
            if a != b && a + b == numberToTest
                result += 1
            end
        end
    end
    return result > 0
end

function findWeaknessA(file, preambleLength, allNumbers)
    for n in preambleLength+1:length(allNumbers)
        # println(allNumbers[n])
        if !testNumberAtIndex(allNumbers,n,preambleLength)
            println("Found weaknes in ", allNumbers[n], " at index ", n)
            return(n,allNumbers[n])
        end
    end
end

function findWeaknessB(file, preambleLength)
    allNumbers = parseInput(file)
    ixA, numberToTest = findWeaknessA(file, preambleLength, allNumbers)
    contiguousNumbers = []
    for i in 1:ixA-1
        for j in i+1:ixA-1
            view = @view allNumbers[i:j]
            sum = reduce(+, view)
            if sum == numberToTest
                println("SUCCESS! ", view, "=", sum )
                return extremeSum(view)
            end
        end
    end
end

function extremeSum(array)
    min = findmin(array)[1]
    max = findmax(array)[1]
    return min+max
end

println("ExtremeSum=",findWeaknessB("data.txt", 25))


using ChainRulesCore: include
using SparseADRules
using Test, Random, LinearAlgebra, NiLang, SparseArrays
using NiLang.AD, ForwardDiff
using ChainRulesCore, ChainRulesTestUtils
import FiniteDifferences

using SparseADRules: Normal_QR
function ChainRulesTestUtils.test_approx(actual::Normal_QR, expected::Normal_QR, msg=""; kwargs...)
    ChainRulesTestUtils.test_approx(actual.Q, expected.Q, msg; kwargs...)
    ChainRulesTestUtils.test_approx(actual.R, expected.R, msg; kwargs...)
end

include("testutils.jl")

@testset "SparseADRules.jl" begin
    include("linalg.jl")
    include("chainrules.jl")
    include("lowranksvd.jl")
end


d2 = readdlm("./tsp/distances_usa.txt")
c2 = readdlm("./tsp/cities_usa.txt")

# Some basic unit tests
# If a Test fails, there is an error in your implementation
# Passing the tests does not mean your implementation is correct!
@testset begin
    _s1 = find_subtours(
        [
            [0.0 1.0];
            [1.0 0.0]
        ]
    )  # no sub-tour should be returned here

    _s2 = find_subtours(
        [
            [0.0 1.0 0.0 0.0];
            [1.0 0.0 0.0 0.0];
            [0.0 0.0 0.0 1.0];
            [0.0 0.0 1.0 0.0]
        ]
    )  # two sub-tours should be returned here

    @test length(_s1) == 0
    @test length(_s2) == 2
end;

sol2 = solve_tsp2(50, d2, GLPK.Optimizer)
RoadTrip(50, c2, sol2)
![](tsp/solution_usa.png)


using Gadfly
using CSV
using DataFrames

data = CSV.read("C:/Users/ztjsw/Desktop/data.txt", delim = ',', header = true)

data[:chromosome] = string.(data[:chromosome])
data[:neg_log_p] = -log10.(data[:pvalue])
data[:adj_pos] = copy(data[:position])

function cal_adj_pos(data)
    #calculate adjusted position
    start = 0
    for i in 2:nrow(data)
        if data[:chromosome][i] != data[:chromosome][i-1]
            start += data[:position][i-1]
        end
        data[:adj_pos][i] += start
    end
    data
end

data_new = cal_adj_pos(data)

## Plot
#calculate postion for xticks
x_chrom_pos = by(data_new, :chromosome, mid_pos = :adj_pos => d -> (maximum(d) + minimum(d))/2)[2]
x_chrom_label = unique(data_new[:chromosome])

function pos_to_chrom(iPos)
    # return the name of its chromosome for a position
    return x_chrom_label[x_chrom_pos .== iPos][1]
end

p=plot(data_new, x=:adj_pos, y=:neg_log_p, color=:chromosome, Geom.point,
        Guide.xlabel("Chromosome"),
        Guide.ylabel("-log10(p)"),
        Guide.xticks(ticks=x_chrom_pos),
        Scale.x_continuous(labels = pos_to_chrom),  #labels is a function which maps a coordinate value to a string label
        Theme(key_position = :none, grid_line_width=0mm, point_size=0.1cm))

display(p)


import CLBLAS
import OpenCL

    const clblas = CLBLAS
    clblas.setup()

    const cl = OpenCL

    device, ctx, queue = clblas.get_next_compute_context()

    N = unsigned(5)
    data = Array(Float32, 5)
    data[1] = 1.1
    data[2] = 2.22
    data[3] = 3.333
    data[4] = 4.4444
    data[5] = 5.55555

    bufX = cl.Buffer(Float32, ctx, (:r, :copy),hostbuf=data)

    bufNRM2 = cl.Buffer(Float32, ctx, :w, 1)
    scratchBuff = cl.Buffer(Float32, ctx, :rw, 2 * length(data))

    offNRM2 = offx = unsigned(0)
    incx = cl.cl_int(1)
    ncq = cl.cl_uint(1)
    clq = queue.id
    newl = cl.cl_uint(0)

    event = cl.UserEvent(ctx, retain=true)
    ptrEvent = [event.id]

    clblas.clblasSnrm2( N, bufNRM2.id, offNRM2, bufX.id, offx, incx, scratchBuff.id, ncq, [clq], newl, C_NULL, ptrEvent);

    cl.api.clWaitForEvents(cl.cl_uint(1), ptrEvent)

    result = Array(Float32, 1)
    cl.enqueue_read_buffer(queue, bufNRM2, result, unsigned(0), nothing, true)

    expected = BLAS.nrm2(5, data, 1)
    println(result)
    println("------------")
    println(expected)

    try
        @assert(isapprox(norm(expected - result[1]), zero(Float32)))
        info("success!")
    finally
        clblas.teardown()
    end


@inline function C_se(Cell,s,z,Def,cse_tf,D,res0)
"""
Concentration Solid-Electrolyte Transfer Function

C_se(Cell,s,z,Def)

"""

if Def == "Pos"
    Electrode = Cell.Pos #Electrode Length
else
    Electrode = Cell.Neg #Electrode Length
end

κ_eff = Cell.Const.κ*Electrode.ϵ_e^Electrode.κ_brug #Effective Electrolyte Conductivity
σ_eff = Electrode.σ*Electrode.ϵ_s^Electrode.σ_brug #Effective Electrode Conductivity

#Defining SOC
θ = Cell.Const.SOC * (Electrode.θ_100-Electrode.θ_0) + Electrode.θ_0

#Prepare for j0
cs0 = Electrode.cs_max * θ

#Current Flux Density
if Cell.Const.CellTyp == "Doyle_94"
    κ = Electrode.k_norm/Electrode.cs_max/Cell.Const.ce0^(1-Electrode.α)
    j0 = κ*(Cell.Const.ce0*(Electrode.cs_max-cs0))^(1-Electrode.α)*cs0^Electrode.α
else
    j0 = Electrode.k_norm*(Cell.Const.ce0*cs0*(Electrode.cs_max-cs0))^(1-Electrode.α)
end

#Resistances
Rtot = R*Cell.Const.T/(j0*F^2) + Electrode.RFilm
#Rtot = R*Cell.Const.T/(j0*Cell.Const.CC_A*F) + Electrode.RFilm

#∂Uocp_Def
∂Uocp_elc = Cell.Const.∂Uocp(Def,θ)/Electrode.cs_max

res0 .= -3/(Electrode.as*F*Electrode.L*Cell.Const.CC_A*Electrode.Rs)*ones(length(z)) #Residual Variable for Pole Removal - eq. 4.45
ν = @. Electrode.L*sqrt((Electrode.as/σ_eff+Electrode.as/κ_eff)/(Rtot+∂Uocp_elc*(Electrode.Rs/(F*Electrode.Ds))*(tanh(Electrode.β)/(tanh(Electrode.β)-Electrode.β)))) #Condensing Variable - eq. 4.13

cse_tf .= @. ν*Electrode.Rs*(σ_eff*cosh(ν*z)+κ_eff*cosh(ν*(z-1)))*tanh(Electrode.β)/(Electrode.as*F*Electrode.L*Cell.Const.CC_A*Electrode.Ds*(κ_eff+σ_eff)*sinh(ν)*(tanh(Electrode.β)-Electrode.β))-res0/s #Transfer Function - eq. 4.17
#cse_tf .= @. cse_tf-res0t/s  #Pole removal - eq. 4.43
zero_tf = @. (5*Electrode.as*Electrode.Ds*F*Electrode.L^2*(κ_eff*(2-6*z+3*z^2)+(3*z^2-1)*σ_eff)-6*∂Uocp_elc*Electrode.Rs*κ_eff*σ_eff)/(30*Cell.Const.CC_A*Electrode.as*Electrode.Ds*∂Uocp_elc*F*Electrode.L*κ_eff*σ_eff) #For s = 0 / Wolfram Alpha
cse_tf[:,findall(s.==0)] .= zero_tf[:,findall(s.==0)]
D .= zeros(length(z))

if Def == "Pos"
    cse_tf .= -cse_tf
    res0 .= -res0
end

if abs.(cse_tf[:,1]) > abs.(cse_tf[:,2])*10
    cse_tf[:,1] = cse_tf[:,2]*10
end

end


using OnlinePCA
using DelimitedFiles: writedlm
writecsv(filename::AbstractString, data) = writedlm(filename, data, ',')

# Setting
chunksize = parse(Int64, ARGS[1])
outdir = ARGS[2]

# Randomized SVD
out = OnlinePCA.tenxpca(
	tenxfile="../Data/Brain/1M_neurons_filtered_gene_bc_matrices_h5.h5",
	rowmeanlist="../Data/Brain/Feature_LogMeans.csv",
	scale="log",
	dim=20,
	niter=3,
	chunksize=chunksize,
	group="mm10")

# Output
writecsv(joinpath(outdir, "Eigen_vectors.csv"), out[1])
writecsv(joinpath(outdir, "Eigen_values.csv"), out[2])
writecsv(joinpath(outdir, "Loadings.csv"), out[3])


module Day02
export day02

function parseLine(line)
    cmd, distance = split(line, " ")
    return (Symbol(cmd), parse(Int, distance))
end

function readInput(inputfile="inputs/input02.txt")
    return parseLine.(eachline(inputfile))
end

function day02(inputdata = readInput("inputs/input02.txt"))
    return (day02Part1(inputdata), day02Part2(inputdata))
end

function day02Part1(cmdlist)
    pos = 0+0im
    factors = (; forward=1, down=im, up=-im)
    for (cmd, dist) in cmdlist
        pos += factors[cmd]*dist
    end
    return real(pos)*imag(pos)
end

function day02Part2(cmdlist)
    pos = 0+0im
    aim = 0im
    ## simple "if" statements are probably cleaner...
    moves = (; forward = x -> (x*(1+aim), 0),
               down = x -> (0, im*x),
               up = x -> (0, -im*x))
    for (cmd, dist) in cmdlist
        Δpos, Δaim = moves[cmd](dist)
        pos += Δpos
        aim += Δaim
    end
    return real(pos)*imag(pos)
end

end #module


# # Histogram of the most significant digit (MSD)

using LiterateTest                                                     #src
using Test                                                             #src

# ## An allocation-free function to compute MSD

function msd(x::Real)
    x = abs(x)
    d = x
    while true
        x < 1 && return floor(Int, d)
        d = x
        x ÷= 10
    end
end
nothing  # hide
#-

@test begin
    msd(34513)
end == 3
#-

@test begin
    msd(-51334)
end == 5
#-

@assert 2.76e19 > typemax(Int64)  #src
@test begin
    msd(2.76e19)
end == 2
#-

# ## MSD of random numbers

using CUDA
using FLoops
using FoldsCUDA
using Setfield

function histogram_msd(xs, ex = nothing)
    zs = ntuple(_ -> 0, 9)  # a tuple of 9 zeros
    @floop ex for x in xs
        d = msd(x)
        1 <= d <= 9 || continue  # skip it if `msd` returns 0
        h2 = @set zs[d] = 1      # set `d`-th position of the tuple to 1
        @reduce(h1 = zs .+ h2)   # point-wise addition merges the histogram
    end
    return h1
end
nothing  # hide

# Generate some random numbers

xs = let randn = has_cuda_gpu() ? CUDA.randn : randn
    exp.(10.0 .* (randn(10^8) .+ 6))
end
@assert all(isfinite, xs)
collect(view(xs, 1:10))  # preview

# Pass an array of (real) numbers to `histogram_msd` to compute the
# histogram of MSD:

hist1 = histogram_msd(xs)

# Frequency in percentage:

aspercentage(h) = pairs(round.((collect(h) ./ length(xs) .* 100); digits = 1))
aspercentage(hist1)
#-

# ## MSD of lazily computed numbers

# FoldsCUDA supports running reduction over non-`CuArray` containers
# such as `UnitRange` and iterator transformation wrapping it.
# However, you need to explicitly specify to use CUDA by, e.g.,
# passing `CUDAEx` to `@floop`:

hist2 = histogram_msd(x^2 for x in 1:10^8)

# Frequency in percentage:
aspercentage(hist2)
#-

hist3 = histogram_msd(exp(x) for x in range(1, 35, length=10^8))

# Frequency in percentage:
aspercentage(hist3)


println(sqrt(1im))
println(6//9)
println(numerator(6//9))
println(denominator(6//9))


struct QGramIterator{S <: Union{AbstractString, AbstractVector}}
	s::S   # Collection
	q::Int # Length of Qgram
	function QGramIterator{S}(s, q) where {S <: Union{AbstractString, AbstractVector}}
		q > 0 || throw(ArgumentError("The qgram length must be higher than zero"))
		new(s, q)
	end
end
function QGramIterator(s::Union{AbstractString, AbstractVector}, q::Integer)
	QGramIterator{typeof(s)}(s, q)
end
Base.length(qgram::QGramIterator) = max(length(qgram.s) - qgram.q + 1, 0)

# q-grams of AbstractString
function Base.iterate(qgram::QGramIterator{<: AbstractString},
	state = (1, nextind(qgram.s, 0, qgram.q)))
	istart, iend = state
	iend > ncodeunits(qgram.s) && return nothing
	element = SubString(qgram.s, istart, iend)
	nextstate = nextind(qgram.s, istart), nextind(qgram.s, iend)
	element, nextstate
end
Base.eltype(qgram::QGramIterator{SubString{S}}) where {S} = SubString{S}
Base.eltype(qgram::QGramIterator{S}) where {S <: AbstractString} = SubString{S}
qgrams(s::AbstractString, q::Integer) = QGramIterator(s, q)

#q-grams of AbstractVector
# Alternatively, I could also use partition in IterTools but it creates a vector for each iteration
# so it does not seem to be worth it.
function Base.iterate(qgram::QGramIterator{<: AbstractVector}, state = firstindex(qgram.s))
	state + qgram.q - 1 > lastindex(qgram.s) && return nothing
	view(qgram.s, state:(state + qgram.q - 1)), state + 1
end
Base.eltype(qgram::QGramIterator{<: AbstractVector}) = typeof(first(qgram))
qgrams(s::AbstractVector, q::Integer) = QGramIterator(s, q)
qgrams(s, q::Integer) = QGramIterator(collect(s), q)

@doc """
Return an iterator corresponding to the the q-gram of an iterator.
When the iterator is a String, qgrams are SubStrings.

### Arguments
* `s` iterator
* `q::Integer`: length of q-gram

## Examples
```julia
for x in qgrams("hello", 2)
	println(x)
end
```
"""
qgrams

# For two iterators s1 and s2, that define a length and eltype method,
# this returns an iterator that,
# for each element in s1 ∪ s2, returns (numbers of times it appears in s1, numbers of times it appears in s2)
function _count(qgrams1, qgrams2)
	K = promote_type(eltype(qgrams1), eltype(qgrams2))
	d = Dict{K, Tuple{Int, Int}}()
	sizehint!(d, length(qgrams1) + length(qgrams2))
	# I use a faster way to change a dictionary key
	# see setindex! in https://github.com/JuliaLang/julia/blob/master/base/dict.jl#L380
	for x1 in qgrams1
		index = Base.ht_keyindex2!(d, x1)
		if index > 0
			d.age += 1
			@inbounds d.keys[index] = x1
			@inbounds d.vals[index] = (d.vals[index][1] + 1, 0)
		else
			@inbounds Base._setindex!(d, (1, 0), x1, -index)
		end
	end
	for x2 in qgrams2
		index = Base.ht_keyindex2!(d, x2)
		if index > 0
			d.age += 1
			@inbounds d.keys[index] = x2
			@inbounds d.vals[index] = (d.vals[index][1], d.vals[index][2] + 1)
		else
			@inbounds Base._setindex!(d, (0, 1), x2, -index)
		end
	end
	return values(d)
end

abstract type AbstractQGramMatchCounter end

abstract type AbstractQGramDistance <: SemiMetric end

function (dist::AbstractQGramDistance)(s1, s2)
	((s1 === missing) | (s2 === missing)) && return missing
	counter = newcounter(dist)
	for (n1, n2) in _count(qgrams(s1, dist.q), qgrams(s2, dist.q))
		count!(dist, counter, n1, n2)
	end
	calculate(dist, counter)
end

"""
	QGram(q::Int)

Creates a QGram distance.

The distance corresponds to

``||v(s1, q) - v(s2, q)||``

where ``v(s, q)`` denotes the vector on the space of q-grams of length q,
that contains the number of times a q-gram appears for the string s
"""
struct QGram <: AbstractQGramDistance
	q::Int
end

mutable struct SingleCounter <: AbstractQGramMatchCounter
	shared::Int
end

newcounter(::QGram) = SingleCounter(0)
@inline function count!(::QGram, c::SingleCounter, n1::Integer, n2::Integer)
	c.shared += abs(n1 - n2)
end
calculate(::QGram, c::SingleCounter) = c.shared

"""
	Cosine(q::Int)

Creates a Cosine distance.

The distance corresponds to

`` 1 - v(s1, q).v(s2, q)  / ||v(s1, q)|| * ||v(s2, q)||``

where ``v(s, q)`` denotes the vector on the space of q-grams of length q,
that contains the  number of times a q-gram appears for the string s
"""
struct Cosine <: AbstractQGramDistance
	q::Int
end

mutable struct ThreeCounters <: AbstractQGramMatchCounter
	left::Int
	right::Int
	shared::Int
end

newcounter(::Cosine) = ThreeCounters(0, 0, 0)
@inline function count!(::Cosine, c::ThreeCounters, n1::Integer, n2::Integer)
	c.left += n1^2
	c.right += n2^2
	c.shared += n1 * n2
end
calculate(::Cosine, c::ThreeCounters) =
	1.0 - c.shared / (sqrt(c.left) * sqrt(c.right))

"""
	Jaccard(q::Int)

Creates a Jaccard distance.

The distance corresponds to

``1 - |Q(s1, q) ∩ Q(s2, q)| / |Q(s1, q) ∪ Q(s2, q))|``

where ``Q(s, q)``  denotes the set of q-grams of length n for the string s
"""
struct Jaccard <: AbstractQGramDistance
	q::Int
end
newcounter(::Jaccard) = ThreeCounters(0, 0, 0)
@inline function count!(::Jaccard, c::ThreeCounters, n1::Integer, n2::Integer)
	c.left += n1 > 0
	c.right += n2 > 0
	c.shared += (n1 > 0) & (n2 > 0)
end
calculate(::Jaccard, c::ThreeCounters) =
	1.0 - c.shared / (c.left + c.right - c.shared)

"""
	SorensenDice(q::Int)

Creates a SorensenDice distance.

The distance corresponds to

``1 - 2 * |Q(s1, q) ∩ Q(s2, q)|  / (|Q(s1, q)| + |Q(s2, q))|)``

where ``Q(s, q)``  denotes the set of q-grams of length n for the string s
"""
struct SorensenDice <: AbstractQGramDistance
	q::Int
end
newcounter(::SorensenDice) = ThreeCounters(0, 0, 0)
@inline function count!(::SorensenDice, c::ThreeCounters, n1::Integer, n2::Integer)
	c.left += n1 > 0
	c.right += n2 > 0
	c.shared += (n1 > 0) & (n2 > 0)
end
calculate(::SorensenDice, c::ThreeCounters) =
	1.0 - 2.0 * c.shared / (c.left + c.right)

"""
	Overlap(q::Int)

Creates a Overlap distance.

The distance corresponds to

``1 - |Q(s1, q) ∩ Q(s2, q)|  / min(|Q(s1, q)|, |Q(s2, q)|)``

where ``Q(s, q)``  denotes the set of q-grams of length n for the string s
"""
struct Overlap <: AbstractQGramDistance
	q::Int
end
newcounter(::Overlap) = ThreeCounters(0, 0, 0)
@inline function count!(::Overlap, c::ThreeCounters, n1::Integer, n2::Integer)
	c.left += n1 > 0
	c.right += n2 > 0
	c.shared += (n1 > 0) & (n2 > 0)
end
calculate(::Overlap, c::ThreeCounters) =
	1.0 - c.shared / min(c.left, c.right)

"""
	NMD(q::Int)
	NMD(q::Int)

Creates a NMD (Normalized Multiset Distance) as introduced by Besiris and
Zigouris 2013. The goal with this distance is to behave similarly to a normalized
compression distance without having to do any actual compression (and thus being
faster to compute).

The distance corresponds to

``(sum(max.(m(s1), m(s2)) - min(M(s1), M(s2))) / max(M(s1), M(s2))``

where ``m(s)`` is the vector of q-gram counts for string ``s`` and ``M(s)`` is the
sum of those counts.

For details see:
https://www.sciencedirect.com/science/article/pii/S1047320313001417
"""
struct NMD <: AbstractQGramDistance
	q::Int
end

newcounter(::NMD) = ThreeCounters(0, 0, 0)
@inline function count!(::NMD, c::ThreeCounters, n1::Integer, n2::Integer)
	c.left += n1
	c.right += n2
	c.shared += max(n1, n2)
end
calculate(::NMD, c::ThreeCounters) =
	(c.shared - min(c.left, c.right)) / max(c.left, c.right)

"""
	MorisitaOverlap(q::Int)

Creates a MorisitaOverlap distance, a general, statistical measure of
dispersion which can also be used on dictionaries such as created
from q-grams. See https://en.wikipedia.org/wiki/Morisita%27s_overlap_index
This is more fine-grained than many of the other QGramDistances since
it is based on the counts per q-gram rather than only which q-grams are
in the strings.

The distance corresponds to

``(2 * sum(m(s1) .* m(s2)) / (sum(m(s1).^2)*M(s2)/M(s1) + sum(m(s2).^2)*M(s1)/M(s2))``

where ``m(s)`` is the vector of q-gram counts for string ``s`` and ``M(s)`` is the
sum of those counts.
"""
struct MorisitaOverlap <: AbstractQGramDistance
	q::Int
end

mutable struct FiveCounters <: AbstractQGramMatchCounter
	leftsum::Int    # sum(m(s1))
	rightsum::Int   # sum(m(s2))
	leftsq::Int     # sum(m(s1).^2)
	rightsq::Int    # sum(m(s2).^2)
	shared::Int     # sum(m(s1) .* m(s2))
end

newcounter(::MorisitaOverlap) = FiveCounters(0, 0, 0, 0, 0)
@inline function count!(::MorisitaOverlap, c::FiveCounters, n1::Integer, n2::Integer)
	c.leftsum += n1
	c.rightsum += n2
	c.leftsq += n1^2
	c.rightsq += n2^2
	c.shared += n1 * n2
end
calculate(::MorisitaOverlap, c::FiveCounters) =
	1.0 - ((2 * c.shared) / (c.leftsq*c.rightsum/c.leftsum + c.rightsq*c.leftsum/c.rightsum))


abstract type AbstractNormalSample{T} <: ContinuousEBayesSample{T} end

"""
    NormalSample(Z,σ)

An observed sample ``Z`` drawn from a Normal distribution with known variance ``\\sigma^2 > 0``.

```math
Z \\sim \\mathcal{N}(\\mu, \\sigma^2)
```

``\\mu`` is assumed unknown. The type above is used when the sample ``Z`` is to be used for estimation or inference of ``\\mu``.

```jldoctest
julia> NormalSample(0.5, 1.0)          #Z=0.5, σ=1
Z=     0.5 | σ=1.0
```
"""
struct NormalSample{T,S} <: AbstractNormalSample{T}
    Z::T
    σ::S
end

# TODO: Should not need this eventually.
function NormalSample(Z::P, σ::S) where {T,P<:EBInterval{T},S}
    NormalSample{EBInterval{T},S}(Z, σ)
end

function NormalSample(σ::S) where {S}
    NormalSample(missing, σ)
end

"""
    StandardNormalSample(Z)

An observed sample ``Z`` drawn from a Normal distribution with known variance ``\\sigma^2 =1``.

```math
Z \\sim \\mathcal{N}(\\mu, 1)
```

``\\mu`` is assumed unknown. The type above is used when the sample ``Z`` is to be used for estimation or inference of ``\\mu``.

```jldoctest
julia> StandardNormalSample(0.5)          #Z=0.5
Z=     0.5 | σ=1.0
```
"""
struct StandardNormalSample{T} <: AbstractNormalSample{T}
    Z::T
end

StandardNormalSample() = StandardNormalSample(missing)

eltype(Z::AbstractNormalSample{T}) where {T} = T
support(Z::AbstractNormalSample) = RealInterval(-Inf, +Inf)

response(Z::AbstractNormalSample) = Z.Z
var(Z::AbstractNormalSample) = Z.σ^2
var(Z::StandardNormalSample) = one(eltype(response(Z)))
var(Z::StandardNormalSample{Missing}) = 1.0

std(Z::AbstractNormalSample) = Z.σ
std(Z::StandardNormalSample) = one(eltype(response(Z)))
std(Z::StandardNormalSample{Missing}) = 1.0

nuisance_parameter(Z::AbstractNormalSample) = std(Z)

likelihood_distribution(Z::AbstractNormalSample, μ) = Normal(μ, std(Z))

function Base.show(io::IO, Z::AbstractNormalSample{<:Real})
    Zz = response(Z)
    print(io, "Z=")
    print(io, lpad(round(Zz, digits = 4),8))
    print(io, " | ", "σ=")
    print(io, rpad(round(std(Z), digits = 3),5))
end

function Base.show(io::IO, Z::AbstractNormalSample{<:Interval})
    Zz = response(Z)
    print(io, "Z ∈ ")
    show(IOContext(io, :compact => true), Zz)
    print(io, " | ", "σ=")
    print(io, rpad(round(std(Z), digits = 3),5))
end

# Targets

# TODO: Note this is not correct for intervals.
function cf(target::MarginalDensity{<:AbstractNormalSample}, t)
    error_dbn = likelihood_distribution(location(target))
    cf(error_dbn, t)
end

# Conjugate computations
function default_target_computation(::BasicPosteriorTarget,
    ::AbstractNormalSample,
    ::Normal
)
    Conjugate()
end

function marginalize(Z::AbstractNormalSample, prior::Normal)
    prior_var = var(prior)
    prior_μ = mean(prior)
    likelihood_var = var(Z)
    marginal_σ = sqrt(likelihood_var + prior_var)
    Normal(prior_μ, marginal_σ)
end

function posterior(Z::AbstractNormalSample, prior::Normal)
    z = response(Z)
    sigma_squared = var(Z)
    prior_mu = mean(prior)
    prior_A = var(prior)

    post_mean =
        (prior_A) / (prior_A + sigma_squared) * z +
        sigma_squared / (prior_A + sigma_squared) * prior_mu
    post_var = prior_A * sigma_squared / (prior_A + sigma_squared)
    Normal(post_mean, sqrt(post_var))
end

function _set_defaults(
    convexclass::DiscretePriorClass,
    Zs::VectorOrSummary{<:AbstractNormalSample};  #TODO for MultinomialSummary
    hints,
)
    eps = get(hints, :eps, 1e-4)
    prior_grid_length = get(hints, :prior_grid_length, 300)::Integer
    _sample_min, _sample_max = extrema(response.(Zs))
    # This won't handle infinity correctly. TODOO
    # Also TODO once switch from Intervals.jl to IntervalSets.jl
    _sample_min = isa(_sample_min, Interval) ? first(_sample_min) : _sample_min
    _sample_max = isa(_sample_max, Interval) ? last(_sample_max) : _sample_max
    _grid = range(_sample_min - eps; stop = _sample_max + eps, length = prior_grid_length)
    DiscretePriorClass(_grid)
end

function _set_defaults(
    convexclass::GaussianScaleMixtureClass,
    Zs::AbstractVector{<:AbstractNormalSample};  #TODO for MultinomialSummary
    hints,
)
    grid_scaling = get(hints, :grid_scaling, sqrt(2))

    σ_min =  minimum(std.(Zs))./ 10
    _max = maximum(response.(Zs).^2 .-  var.(Zs))
    _σ_max = _max > 0.0 ? 2*sqrt(_max) : 8*σ_min
    σ_max = get(hints, :σ_max, _σ_max) #somewhat redundant computations above.

    npoint = ceil(Int, log2(σ_max/σ_min)/log2(grid_scaling))
    σ_grid = σ_min*grid_scaling.^(0:npoint)

    GaussianScaleMixtureClass(σ_grid)
end

# Target specifics
function Base.extrema(density::MarginalDensity{<:AbstractNormalSample{<:Real}})
    (0.0, 1 / sqrt(2π * var(location(density))))
end


using CompScienceMeshes, BEAST

trc = X->ntrace(X,γ)
dvg = divergence

h1 = h2 = h3 = 1/15; h = max(h1, h2, h3)
width, height = 1.0, 0.5
Γ1 = meshrectangle(width, height, h1)
Γ2 = meshrectangle(width, height, h2); CompScienceMeshes.rotate!(Γ2, 0.5π*x̂)
Γ3 = meshrectangle(width, height, h3); CompScienceMeshes.rotate!(Γ3, 1.0π*x̂)
γ = meshsegment(width, width, 3)

κ = 1.0
S, T, I = SingleLayerTrace(κ*im), MWSingleLayer3D(κ), Identity()
St = transpose(S)
E = Maxwell3D.planewave(direction=ẑ, polarization=x̂, wavenumber=κ)
e = (n×E)×n

in_interior1 = CompScienceMeshes.interior_tpredicate(Γ1)
in_interior2 = CompScienceMeshes.interior_tpredicate(Γ2)
in_interior3 = CompScienceMeshes.interior_tpredicate(Γ3)

on_junction = CompScienceMeshes.overlap_gpredicate(γ)
edges1 = skeleton(e -> in_interior1(e) || on_junction(chart(Γ1,e)), Γ1,1)
edges2 = skeleton(e -> in_interior2(e) || on_junction(chart(Γ2,e)), Γ2,1)
edges3 = skeleton(e -> in_interior3(e) || on_junction(chart(Γ3,e)), Γ3,1)

X1 = raviartthomas(Γ1, edges1)
X2 = raviartthomas(Γ2, edges2)
X3 = raviartthomas(Γ3, edges3)
# X1, X2, X3 = raviartthomas(Γ1, γ), raviartthomas(Γ2, γ), raviartthomas(Γ3, γ)
X = X1 × X2 × X3

@hilbertspace j
@hilbertspace k
α, β = 1/(im*κ), log(abs(h))/(im*κ)
Eq = @varform T[k,j] + α*S[trc(k), dvg(j)] + α*St[dvg(k), trc(j)] + β*I[trc(k), trc(j)] == e[k]
eq = @discretise Eq  j∈X k∈X
u = solve(eq)

Θ, Φ = range(0.0,stop=π,length=100), 0.0
ffpoints = [point(cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)) for θ in Θ for ϕ in Φ]
farfield = potential(MWFarField3D(κ*im), ffpoints, u, X)

using Plots
using LinearAlgebra
plot(Θ,norm.(farfield))


include("../../ApproximateCOSMO.jl/src/COSMO.jl")
using Main.COSMO, Random
include("../src/closest_correlation.jl"); include("../utils/plot_utils.jl")
using RandomCorrelationMatrices

n = 300
C = RandomCorrelationMatrices.randcormatrix(n, 0.1) + 0.1*Symmetric(randn(n, n))
all(eigvals(C) .>= 0) && @warn "The perturbed correlation matrix is still positive definite!"

settings = COSMO.Settings(verbose=true, verbose_timing=true,
    eps_abs=1e-7, eps_rel=1e-7, max_iter=5000,
    rho=.1, adaptive_rho=false,
    check_termination=40, check_infeasibility=1000)
    # kkt_solver=COSMO.with_options(COSMO.IndirectReducedKKTSolver, solver=:CG)
model = create_closest_correlation_model(C, settings; psdcone_type=COSMO.PsdConeTriangleLanczos)
COSMO.optimize!(model)


@with_kw struct FastTree
    tolerance::Float64 = 1.0
end

# This is the main function

function solve(solver::FastTree, problem::Problem)
    center = problem.input.center
    radius = problem.input.radius[1]
    (W, b) = (problem.network.layers[1].weights, problem.network.layers[1].bias)

    input = forward_affine_map(solver, W, b, problem.input)
    lower, upper = low(input), high(input)
    local_lower, local_upper = similar(lower), similar(lower)

    np = pyimport("numpy")
    q, e, p = np.linalg.svd(W)

    k_1 = size(W, 1)
    k_0 = size(W, 2)

    C = vcat(W, -W)
    d = zeros(2k_1)

    hps = Array{Hyperplane, 1}(undef, k_0-k_1)

    for i in k_1+1:k_0
        hps[i-k_1] = Hyperplane(p[i,:], dot(p[i,:], center))
    end

    S = hps[1]
    if k_0-k_1 > 1
        for i in 2:k_0-k_1
            S = intersection(S, hps[i])
        end
    end

    stack = Vector{Hyperrectangle}(undef, 0)
    push!(stack, input)
    #count = 1
    while !isempty(stack)
        interval = popfirst!(stack)
        reach = forward_network(solver, problem.network, interval)
        if issubset(reach, problem.output)
            continue
        else
            if get_largest_width(interval) > solver.tolerance
                sections = bisect(interval)
                for i in 1:2
                    local_lower = low(sections[i])
                    local_upper = high(sections[i])

                    d = vcat(local_upper - b, b - local_lower)
                    P_i = HPolyhedron(C, d)

                    inter = intersection(problem.input, P_i)

                    if isempty(inter) == false
                        O_i = box_approximation(intersection(P_i, S))
                        inner = true
                        for j in 1:k_0
                            if low(O_i)[j] ≤ low(problem.input)[j]
                                inner = false
                                break
                            end

                            if high(problem.input)[j] ≤ high(O_i)[j]
                                inner = false
                                break
                            end
                        end
                        if inner == false
                            push!(stack, sections[i])
                            #count += 1
                        end
                    end
                end
            else
                return BasicResult(:violated)
            end
        end
    end
    #print("\n$(count)\n")
    return BasicResult(:holds)
end

function forward_network(solver::FastTree, nnet::Network, input::Hyperrectangle)
    layers = nnet.layers
    act = layers[1].activation
    reach = Hyperrectangle(low = act.(low(input)), high = act.(high(input)))

    for i in 2:length(layers)
        reach = forward_layer(solver, layers[i], reach)
    end
    return reach
end

function forward_layer(solver::FastTree, L::Layer, input::Hyperrectangle)
    (W, b, act) = (L.weights, L.bias, L.activation)
    center = zeros(size(W, 1))
    gamma  = zeros(size(W, 1))
    for j in 1:size(W, 1)
        node = Node(W[j,:], b[j], act)
        center[j], gamma[j] = forward_node(solver, node, input)
    end
    return Hyperrectangle(center, gamma)
end

function forward_node(solver::FastTree, node::Node, input::Hyperrectangle)
    output    = node.w' * input.center + node.b
    deviation = sum(abs.(node.w) .* input.radius)
    βmax = node.act(output + deviation)
    βmin = node.act(output - deviation)
    return ((βmax + βmin)/2, (βmax - βmin)/2)
end

function forward_affine_map(solver::FastTree, W::Matrix, b::Vector, input::Hyperrectangle)
    center = W * input.center + b
    radius = abs.(W) * input.radius
    return Hyperrectangle(center, radius)
end

#to determine whether
function ishull(x::Hyperrectangle, y::Hyperrectangle, W::Matrix, b::Vector)
    center = y.center
    radius = y.radius[1]
    k_1 = size(W, 1) #length(lower)
    #I = zeros(k_1, k_1) #Identity matrix
    #IN = zeros(k_1, k_1) #negative Identity matrix
    #for k = 1:k_1
        #I[k, k] = 1.0
        #IN[k, k] = -1.0
    #end
    #Hi
    #C_1 = vcat(I, IN)
    C_1 = vcat(Diagonal(ones(k_1)), -Diagonal(ones(k_1)))
    d_1 = vcat(high(x), -low(x))
    #Pi
    C = C_1 * W #size(2k_1, k_0)
    d = d_1 - C_1 * b #size(2k_1)

    #print("C: ")
    #println(C)
    #print("d: ")
    #println(d)

    inter = true #if the intersection of Pi and the input set is empty
    for l = 1:length(d)
        sum = 0.0
        for m = 1:size(C, 2)
            if C[l,m] > 0
                sum += low(y, m) * C[l,m]
            elseif C[l,m] < 0.0
                sum += high(y, m) * C[l,m]
            end
        end
        if sum > d[l]
            inter = false
            break
        end
    end

    if inter
        hull = false #if Hi is a hull of Z, or if Pi and center of the input set are close
        for l = 1:length(d)
            #print(l)
            #print(" distance: ")
            #println(distance(center, C[l,:], d[l]))
            #print("radius: ")
            #println(radius)
            if distance(center, C[l,:], d[l]) > radius
                return true
            end
        end
    end
    return false
end


using CairoMakie
using CubedSphere

w(z) = (1 - cn(z, 1/2)) / sn(z, 1/2)

CubedSphere.visualize_conformal_mapping(w, x_min=-1.85, y_min=-1.85, x_max=1.85, y_max=1.85, n_lines=51, n_samples=101, filepath="square_to_disk_conformal_mapping.png")


using CompScienceMeshes
using BEAST
using LinearAlgebra
using FastBEAST
using IterativeSolvers

CM = CompScienceMeshes

function test_beast_laplace_singlelayer(h; threading=:single,
    farquaddata=nothing, svdrecompress=false)

    Γ = CM.meshsphere(1, h)

    X = lagrangecxd0(Γ)

    𝒱 = Helmholtz3D.singlelayer(wavenumber=0.0)

    hmat = hassemble(𝒱,X,X, treeoptions=BoxTreeOptions(nmin=50), threading=threading, svdrecompress=svdrecompress)

    mat = assemble(𝒱,X,X)
    return mat, hmat
end

mat, hmat_single = test_beast_laplace_singlelayer(0.1)

@test nnz(hmat_single) == 3916760

@test compressionrate(hmat_single) > 0.3
@test estimate_reldifference(hmat_single, mat) ≈ 0 atol=1e-4

mat, hmat_multi = test_beast_laplace_singlelayer(0.1, threading=:multi)

@test nnz(hmat_multi) == 3916760

@test compressionrate(hmat_multi) > 0.3
@test estimate_reldifference(hmat_multi, mat) ≈ 0 atol=1e-4
@test compressionrate(hmat) > 0.3

mat, hmat_single = test_beast_laplace_singlelayer(0.1, farquaddata=quaddata)

@test nnz(hmat_single) == 3916760

@test compressionrate(hmat_single) > 0.3
@test estimate_reldifference(hmat_single, mat) ≈ 0 atol=1e-3

mat, hmat_multi = test_beast_laplace_singlelayer(0.1, threading=:multi, farquaddata=quaddata)

@test nnz(hmat_multi) == 3916760

@test compressionrate(hmat_multi) > 0.3
@test estimate_reldifference(hmat_multi, mat) ≈ 0 atol=1e-3
@test compressionrate(hmat) > 0.3

mat, hmat_svdmulti = test_beast_laplace_singlelayer(0.1, threading=:multi, farquaddata=quaddata, svdrecompress=true)

@test nnz(hmat_svdmulti) == 3356811

@test compressionrate(hmat_svdmulti) > 0.3
@test estimate_reldifference(hmat_svdmulti, mat) ≈ 0 atol=1e-3
@test compressionrate(hmat) > 0.3

#function test_beast_laplace_singlelayer_manufactured(h)

# ##

# mutable struct Monopole{T,P} <: BEAST.Functional
#     location::P
#     wavenumber::T
#     amplitude::T
# end

# function Monopole(p,k,a = 1)
#     T = promote_type(eltype(p), typeof(k), typeof(a))
#     P = similar_type(typeof(p), T)
#     Monopole{T,P}(p,k,a)
# end

# monopolehh3d(;
#     location    = error("missing arguement `location`"),
#     wavenumber   = error("missing arguement `wavenumber`"),
#     ) = Monopole(location, wavenumber)

# function (f::Monopole)(r)
#     p = f.location
#     k = f.wavenumber
#     a = f.amplitude
#     return a * 1/norm(r-p)
# end

#     h=0.2
#     Γ = CM.meshsphere(1, h)

#     X = lagrangecxd0(Γ)

#     κ = 0.0
#     𝒱 = Helmholtz3D.singlelayer(wavenumber=κ)
#     q = Monopole(SVector(0.35,0.35,0.25), 0.0) #monopolehh3d(location = SVector(0.35,0.35,0.25), wavenumber = 0.0)
#     f = assemble(strace(q,Γ), X)
#     hmat = hassemble(𝒱,X,X, nmin=50)
#     mat = assemble(𝒱,X,X)

# ##
#     hx, hch = cg(hmat,f, log=true,reltol=1e-4)

#     x, ch = cg(mat,f, log=true,reltol=1e-4)
# @test norm(hx - x)/norm(x) ≈ 0 atol=1e-3
# ##
#     Φ, Θ = [0.0], range(0,stop=π,length=100)
#     pts = [point(cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)) for ϕ in Φ for θ in Θ]
#     ffd = potential(𝒱, pts, x, X)

# #end

# ##
# h=0.02
# Γ = CM.meshsphere(1, h)

# X = lagrangecxd0(Γ)

# κ = 0.0
# 𝒱 = Helmholtz3D.singlelayer(wavenumber=κ)
# hmat = hassemble(𝒱,X,X, nmin=100)
# println(compressionrate(hmat))


import Distributions: truncated

function validate_bounds(lower, upper, uv, constraint)
    if lower >= upper
        msg = "lower_bound < upper_bound required, got $lower <= $upper"
        msg2 = "cannot truncate $uv with $constraint\n"
        throw(ArgumentError(msg2*msg))
    end
end
################################################################
# Truncating uncertain values based on theoretical distributions
# Operating on the union of both TheoreticalFittedUncertainScalar
# and TheoreticalDistributionScalarValue as the type of the
# uncertain value is ok, because Truncated is defined for
# FittedDistribution, which is the .distribution fild for
# fitted scalars.
################################################################
"""
    truncate(uv::TheoreticalDistributionScalarValue, constraint::SamplingConstraint)

Truncate an uncertain value `uv` represented by a theoretical distribution
according to the sampling `constraint`.
"""
Base.truncate(uv::TheoreticalDistributionScalarValue, constraint::SamplingConstraint)

"""
    truncate(uv::TheoreticalDistributionScalarValue,
        constraint::NoConstraint)

Truncate the theoretical distribution furnishing `uv` using a
`NoConstraint` sampling constraint.
"""
function Base.truncate(uv::TheoreticalDistributionScalarValue,
        constraint::NoConstraint)
    s = support(uv.distribution)
    lower_bound, upper_bound = s.lb, s.ub

    return truncated(uv.distribution, lower_bound, upper_bound)
end

"""
    truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateQuantiles)

Truncate the theoretical distribution furnishing `uv` using a
`TruncateQuantiles` sampling constraint.
"""
function Base.truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateQuantiles)
    lower_bound = quantile(uv.distribution, constraint.lower_quantile)
    upper_bound = quantile(uv.distribution, constraint.upper_quantile)

    validate_bounds(lower_bound, upper_bound, uv, constraint)

    truncated(uv.distribution, lower_bound, upper_bound)
end

"""
    truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateLowerQuantile)

Truncate the theoretical distribution furnishing `uv` using a
`TruncateLowerQuantile` sampling constraint.
"""
function Base.truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateLowerQuantile)
    lower_bound = quantile(uv.distribution, constraint.lower_quantile)
    upper_bound = maximum(uv)

    validate_bounds(lower_bound, upper_bound, uv, constraint)

    truncated(uv.distribution, lower_bound, upper_bound)
end

"""
    truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateUpperQuantile)

Truncate the theoretical distribution furnishing `uv` using a
`TruncateUpperQuantile` sampling constraint.
"""
function Base.truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateUpperQuantile)
    lower_bound = minimum(uv)
    upper_bound = quantile(uv.distribution, constraint.upper_quantile)

    validate_bounds(lower_bound, upper_bound, uv, constraint)

    truncated(uv.distribution, lower_bound, upper_bound)
end

"""
    truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateMinumum)

Truncate the theoretical distribution furnishing `uv` using a
`TruncateMinimum` sampling constraint.
"""
function Base.truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateMinimum)

    lower_bound = constraint.min
    upper_bound = maximum(uv)

    validate_bounds(lower_bound, upper_bound, uv, constraint)

    truncated(uv.distribution, lower_bound, upper_bound)
end

"""
    truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateMaximum)

Truncate the theoretical distribution furnishing `uv` using a
`TruncateMaximum` sampling constraint.
"""
function Base.truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateMaximum)

    lower_bound = minimum(uv)
    upper_bound = constraint.max

    validate_bounds(lower_bound, upper_bound, uv, constraint)

    truncated(uv.distribution, lower_bound, upper_bound)
end

"""
    truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateRange)

Truncate the theoretical distribution furnishing `uv` using a
`TruncateRange` sampling constraint.
"""
function Base.truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateRange)

    lower_bound = constraint.min
    upper_bound = constraint.max

    validate_bounds(lower_bound, upper_bound, uv, constraint)

    truncated(uv.distribution, lower_bound, upper_bound)
end

"""
    truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateStd, n::Int = 10000)

Truncate the theoretical distribution furnishing `uv` using a
`TruncateStd` sampling constraint.

This functions needs to compute the mean and standard deviation
of a truncated distribution, so takes an extra optional argument `n_draws` to allow
this.
"""
function Base.truncate(uv::TheoreticalDistributionScalarValue,
        constraint::TruncateStd)

    m = mean(uv.distribution)
    s = std(uv.distribution)
    lower_bound = m - s
    upper_bound = m + s

    validate_bounds(lower_bound, upper_bound, uv, constraint)

    truncated(uv.distribution, lower_bound, upper_bound)
end

export truncate


# Following code can be plugged into cellbox.jl

begin
    function  loss_all(p)
        loss_epoch = zeros(Float32, n_exp);
        for i_exp in 1:n_exp
            loss_epoch[i_exp] = loss_neuralode(p, i_exp)
        end
        loss_total = mean(loss_epoch);
    end
    p_grad = Zygote.gradient(x -> loss_neuralode(x), p)[1]

    backend(:pyplot)

    xs = vcat(["α"], [string("W i=", i) for i = 1:ns])
    ys = [string("W j=", i) for i = 1:ns]
    plt = heatmap(xs, ys, abs.(p_grad) ./ maximum(abs.(p_grad)), aspect_ratio = 1);
    title!(plt, "Scaled sensitivity map p(α | W)");

    plt_p = heatmap(xs, ys, abs.(p), aspect_ratio = 1);
    title!(plt_p, "map p");

    plt_p_gold = heatmap(xs, ys, abs.(p_gold), aspect_ratio = 1);
    title!(plt_p_gold, "map p_gold");

    plt_p_diff = heatmap(xs, ys, abs.(p_gold - p), aspect_ratio = 1);
    title!(plt_p_diff, "map p_gold - p");

    plt_all = plot(plt, plt_p, plt_p_gold, plt_p_diff, layout=(2, 2), size=(1600, 1600));
    png(plt_all, string(fig_path, "/sens"))

    backend(:gr)
end


module TwoVarianceComponentTest

using VarianceComponentModels, MathProgBase, Ipopt, Random, Test, LinearAlgebra

Random.seed!(123)

# generate data from a d-variate response variane component model
n = 100   # no. observations
d = 2     # no. categories
m = 2     # no. variance components
p = 2     # no. covariates
X = randn(n, p)
B = ones(p, d)
Σ = ntuple(x -> zeros(d, d), m)
for i in 1:m
  Σi = randn(d, d)
  copyto!(Σ[i], Σi' * Σi)
end
## make the first variance component 0 matrix
#fill!(Σ[1], 0.0)
V = ntuple(x -> zeros(n, n), m)
for i = 1:m-1
  Vi = randn(n, 50)
  copyto!(V[i], Vi * Vi')
end
copyto!(V[m], Matrix(I, n, n))
# form Ω
Ω = zeros(n*d, n*d)
for i = 1:m
  global Ω += kron(Σ[i], V[i])
end
Ωchol = cholesky(Ω)
Y = X * B + reshape(Ωchol.L * randn(n*d), n, d)

@info "Forming VarianceComponentModel from data"
@inferred VarianceComponentVariate(Y, X, V)
vcdata  = VarianceComponentVariate(Y, X, V)
vcmodel = VarianceComponentModel(vcdata)

@info "Pre-compute eigen-decomposition and rotate data"
vcdatarot  = TwoVarCompVariateRotate(vcdata)
vcmodelrot = TwoVarCompModelRotate(vcmodel)

@info "Evaluate log-pdf"
#@code_warntype logpdf(vcmodelrot, vcdatarot)
@inferred logpdf(vcmodelrot, vcdatarot)
@test logpdf(vcmodel, vcdata) == logpdf(vcmodelrot, vcdatarot)
@test (logpdf(vcmodelrot, [vcdatarot vcdatarot; vcdatarot vcdatarot]) -
    logpdf(vcmodel, [vcdata vcdata; vcdata vcdata])) ≈ 0.0

@info "Evaluate gradient"
∇ = zeros(2d^2)
#@code_warntype gradient!(∇, vcmodelrot, vcdatarot)
@inferred VarianceComponentModels.gradient!(∇, vcmodelrot, vcdatarot)
@inferred VarianceComponentModels.gradient!(∇, vcmodel, vcdatarot)
@inferred VarianceComponentModels.gradient!(∇, vcmodel, vcdata)
@test norm(VarianceComponentModels.gradient(vcmodel, vcdata) -
    VarianceComponentModels.gradient(vcmodelrot, vcdatarot)) ≈ 0.0
@test norm(VarianceComponentModels.gradient(vcmodel, vcdata) -
    VarianceComponentModels.gradient(vcmodel, vcdatarot)) ≈ 0.0
@test norm(VarianceComponentModels.gradient(vcmodel, [vcdata vcdata]) -
  2.0VarianceComponentModels.gradient(vcmodel, vcdata)) ≈ 0.0
@test norm(VarianceComponentModels.gradient(vcmodel, [vcdata vcdata]) -
  VarianceComponentModels.gradient(vcmodelrot, [vcdatarot vcdatarot])) ≈ 0.0

@info "Evaluate Fisher information matrix of Σ"
H = zeros(2d^2, 2d^2)
#@code_warntype fisher!(H, vcmodelrot, vcdatarot)
@inferred fisher_Σ!(H, vcmodelrot, vcdatarot)
@inferred fisher_Σ!(H, vcmodel, vcdata)
@test norm(fisher_Σ(vcmodel, vcdata) - fisher_Σ(vcmodelrot, vcdatarot)) ≈ 0.0
@test norm(fisher_Σ(vcmodel, vcdata) - fisher_Σ(vcmodel, vcdatarot)) ≈ 0.0
@test norm(fisher_Σ(vcmodel, [vcdata vcdata]) -
    2fisher_Σ(vcmodel, vcdata)) ≈ 0.0
@test norm(fisher_Σ(vcmodel, [vcdata vcdata]) -
    fisher_Σ(vcmodelrot, [vcdatarot vcdatarot])) ≈ 0.0

@info "Evaluate Fisher information matrix of B"
H = zeros(p * d, p * d)
#@code_warntype fisher_B!(H, vcmodelrot, vcdatarot)
@inferred fisher_B!(H, vcmodelrot, vcdatarot)
@inferred fisher_B!(H, vcmodel, vcdatarot)
@inferred fisher_B!(H, vcmodel, vcdata)
@test norm(fisher_B(vcmodel, vcdata) - fisher_B(vcmodelrot, vcdatarot)) ≈ 0.0
@test norm(fisher_B(vcmodel, vcdata) - fisher_B(vcmodel, vcdatarot)) ≈ 0.0
@test norm(fisher_B(vcmodel, [vcdata vcdata]) -
    2.0fisher_B(vcmodel, vcdata)) ≈ 0.0
@test norm(fisher_B(vcmodel, [vcdata vcdata]) -
    fisher_B(vcmodelrot, [vcdatarot vcdatarot])) ≈ 0.0

@info "Find MLE using Fisher scoring"
vcmfs = deepcopy(vcmodel)
#@code_warntype mle_fs!(vcmfs, vcdatarot; solver = :Ipopt)
#@inferred mle_fs!(vcmfs, vcdatarot; solver = :Ipopt)
logl_fs, _, _, Σcov_fs, Bse_fs, = mle_fs!(vcmfs, vcdatarot; solver = :Ipopt)
logl_fs_array, = mle_fs!(vcmfs, [vcdatarot vcdatarot]; solver = :Ipopt)
@show vcmfs.B
@show Bse_fs
@show B

@info "Find MLE using MM algorithm"
vcmmm = deepcopy(vcmodel)
#@code_warntype mle_mm!(vcmm, vcdatarot)
#@inferred mle_mm!(vcmmm, vcdatarot)
logl_mm, _, _, _, Bse_mm, = mle_mm!(vcmmm, vcdatarot)
@test abs(logl_fs - logl_mm) / (abs(logl_fs) + 1.0) < 1.0e-4
@show vcmmm.B
@show Bse_mm
@show B
vcmmm = deepcopy(vcmodel)
logl_mm_array, = mle_mm!(vcmmm, [vcdatarot vcdatarot])
@test abs(logl_fs_array - logl_mm_array) / (abs(logl_fs_array) + 1.0) < 1.0e-4

@info "Find MLE using Fisher scoring (linear equality + box constraints)"
vcmfs       = deepcopy(vcmodel)
vcmfs.A     = [1.0 -1.0 zeros(1, p*d-2)]
vcmfs.sense = '='
vcmfs.b     = 0.0
vcmfs.lb    = 0.0
vcmfs.ub    = 1.0
logl_fs, _, _, Σcov_fs, Bse_fs, = mle_fs!(vcmfs, vcdatarot;
    solver = :Ipopt, qpsolver = :Ipopt)
@show vcmfs.B
@test vcmfs.B[1] ≈ vcmfs.B[2]
@test all(vcmfs.B .≥ 0.0)
@test all(vcmfs.B .≤ 1.0)

@info "Find MLE using MM algorithm (linear equality + box constraints)"
vcmm = deepcopy(vcmodel)
vcmm.A = [1.0 -1.0 zeros(1, p*d-2)]
vcmm.sense = '='
vcmm.b = 0.0
vcmm.lb = 0.0
vcmm.ub = 1.0
logl_mm, _, _, Σcov_mm = mle_mm!(vcmm, vcdatarot; qpsolver = :Ipopt)
@show vcmm.B
@test vcmm.B[1] ≈ vcmm.B[2]
@test all(vcmm.B .≥ 0.0)
@test all(vcmm.B .≤ 1.0)
@test abs(logl_fs - logl_mm) / (abs(logl_fs) + 1.0) < 1.0e-4

@info "Heritability estimation"
h, h_se = heritability(vcmfs.Σ, Σcov_fs)
@show h, h_se
@test all(0.0 .≤ h .≤ 1.0)

@info "test fit_mle (FS)"
vcmmle = deepcopy(vcmodel)
logl_mle, _, _, Σcov_mle, Bse_mle, = fit_mle!(vcmmle, vcdata; algo = :FS)
@show vcmmle.B, Bse_mle, B
vcmle = deepcopy(vcmodel)
fit_mle!(vcmmle, [vcdata vcdata]; algo = :FS)

@info "test fit_mle (MM)"
vcmmle = deepcopy(vcmodel)
logl_mle, _, _, Σcov_mle, Bse_mle, = fit_mle!(vcmmle, vcdata; algo = :MM)
@show vcmmle.B, Bse_mle, B

@info "test fit_reml (FS)"
vcmreml = deepcopy(vcmodel)
logl_reml, _, _, Σcov_reml, Bse_reml, = fit_reml!(vcmreml, vcdata; algo = :FS)
@show vcmreml.B, Bse_reml, B

@info "test fit_reml (MM)"
vcmreml = deepcopy(vcmodel)
logl_reml, _, _, Σcov_reml, Bse_reml, = fit_reml!(vcmreml, vcdata; algo = :MM)
@show vcmreml.B, Bse_reml, B

end # module VarianceComponentTypeTest


module IPR

import ..ED
using ..SimLib
using ..SimLib: FArray
using SpinSymmetry: basissize
using SharedArrays: sdata

import Statistics

export ipr, inverse_participation_ratio, IPRData, IPRDataDescriptor, load_ipr, InverseParticipationRatio

## Data structure

"""
    struct IPRDataDescriptor <: ED.EDDerivedDataDescriptor

See [`EDDataDescriptor`](!ref) and [`EDDerivedDataDescriptor`](!ref).
"""
struct IPRDataDescriptor <: ED.EDDerivedDataDescriptor
    derivedfrom::ED.EDDataDescriptor
end

IPRDataDescriptor(args...; kwargs...) = IPRDataDescriptor(EDDataDescriptor(args...; kwargs...))

"""
    struct IPRData

Stores the actual inverse participation ratios (IPR) specified by the descriptor [`LSRDataDescriptor`](@ref).
Inverse participation ratio is defined as IPR(|ψ⟩) = 1/(∑ᵢ |⟨i|ψ⟩|⁴ ) relative to a basis |i⟩.
Usually |i⟩ is just the z-basis and then the IPR tells a story of localization.
IPR = 2^N for maximally entangled states
IPR = 1 for maximally localized states

The first dimension hold the participation ratio of the i-th eigenstate IPR(|ϕᵢ⟩).

The default save directory is "ipr".

`Statistics.mean` and `Statistics.std` are overloaded to act on the first dimension to conveniently compute
mean LSR and its variance.
"""
struct IPRData{N} <: SimLib.AbstractSimpleData
    descriptor::IPRDataDescriptor
    data::FArray{N}
end

IPRData(iprdd::IPRDataDescriptor) = IPRData(iprdd, FArray{4}(undef, ED.ed_size(iprdd), iprdd.shots, length(iprdd.fields), length(iprdd.ρs)))

## Saving/Loading

ED._default_folder(::IPRDataDescriptor) = "ipr"

"""
    load_ipr(edd)
    load_ipr(model[, diagtype][, location])
"""
load_ipr(args...; kwargs...) = load(IPRDataDescriptor(args...; kwargs...))

## Functions

ipr(mat) = ipr!(Vector{Float64}(undef, size(mat,1)), mat)

const inverse_participation_ratio = ipr

function ipr!(v, mat)
    I, J = size(mat)
    J == length(v) || throw(ArgumentError("Sizes not compatible: $(size(v)) and $(size(mat))"))
    @inbounds for j in 1:J
        tmp = zero(Float64)
        for i in 1:I
            tmp += abs2(mat[i,j])^2
        end
        v[j] = 1/tmp
    end
    v
end

### Task

mutable struct IPRTask <: ED.EDTask
    data
end

InverseParticipationRatio() = IPRTask(nothing)

function ED.initialize!(task::IPRTask, arrayconstructor, spectral_size)
    task.data = arrayconstructor(Float64, spectral_size)
end

function ED.compute_task!(task::IPRTask, evals, evecs, inds...)
    n = min(size(task.data,1), size(evecs, 2))
    @views ipr!(task.data[1:n, inds...], evecs[:, 1:n])
end

function ED.failed_task!(task::IPRTask, inds...)
    task.data[:, inds...] .= NaN64
end

function ED.assemble(task::IPRTask, edd)
    IPRData(IPRDataDescriptor(edd), sdata(task.data))
end

end #module