module StatisticalTests

    # packages
    using Printf
    using DataAPI, DataFrames
    using Distributions, Statistics, StatsBase, StatsModels, GLM
    using Reexport
    @reexport using StatsModels
    const D = Distributions
    # end packages

    # imports
    import Base: show
    # import StatsBase: PValue, StatisticalModel, fit, coeftable, tiedrank, countmap
    # import DataAPI: levels
    # import DataFrames: DataFrame
    # import GLM: LinearModel
    import SpecialFunctions: loggamma
    # end imports

    # exports
    export coeftable, fit
    export
        # functions
        t_test,
        f_test,
        ks_test,
        wilcox_test
    # end export

    # types
    const NumVector{T<:Number} = AbstractVector{T}
    const IntVector{Int} = AbstractVector{Int}
    const FloatVector{T<:AbstractFloat} = AbstractVector{T}
    const NumMatrix{T<:Number} = AbstractMatrix{T}
    const IntMatrix{Int} = AbstractMatrix{Int}
    const FloatMatrix{T<:AbstractFloat} = AbstractMatrix{T}
    abstract type TestModel <: StatisticalModel end
    abstract type TestResult end
    # end types

    # includes
    include("utilities.jl")
    include("ttest.jl")
    include("ftest.jl")
    include("kstest.jl")
    include("wilcoxtest.jl")
    # end includes

end # module


function taxicabdistanceXY(d::DisplacementPositions)
    return TaxiCabDistanceXY(
        abs(d.p.x - d.q.x),
        abs(d.p.y - d.q.y))
end

function taxicabdistance(d::DisplacementPositions)
    TaxiCabDistance(abs(d.p.x - d.q.x) + abs(d.p.y - d.q.y))
end

function periodicdomainshift(distance,domain)
     return (distance > domain/2 ? domain - distance : distance)
end

function taxicabdistance(d::DisplacementPositions,domain)
     TaxiCabDistance(
          periodicdomainshift(abs(d.p.x - d.q.x),domain) +
               periodicdomainshift(abs(d.p.y - d.q.y),domain)
     )
end

function taxicabdistanceXY(d::TaxiCabDistanceXY,domain)
   dx = abs(periodicdomainshift(d.dₓ,domain))
   dy = abs(periodicdomainshift(d.dy,domain))
   return dx + dy
end

walker_distance(w) = taxicabdistance(DisplacementPositions(w[1],w[end]),domain)

function dₜ(w,domain)
     d = DisplacementPositions(
          Walker2D(w[1][1],w[1][2]),
          Walker2D(w[end][1],w[end][2])
     )
     D = taxicabdistanceXY(
          taxicabdistanceXY(d),domain
          )
     return D
end


using FourierCore, FourierCore.Grid, FourierCore.Domain
using FFTW, LinearAlgebra, BenchmarkTools, Random, JLD2
using GLMakie
rng = MersenneTwister(1234)
# Random.seed!(123456789)
Random.seed!(12)
# jld_name = "high_order_timestep_spatial_tracer_"
jld_name = "blocky"
include("transform.jl")
include("random_phase_kernel.jl")
# using GLMakie
using CUDA
arraytype = CuArray
Ω = S¹(2π)^2
N = 2^9 # number of gridpoints
grid = FourierGrid(N, Ω, arraytype=arraytype)
nodes, wavenumbers = grid.nodes, grid.wavenumbers

x = nodes[1]
y = nodes[2]
kˣ = wavenumbers[1]
kʸ = wavenumbers[2]
# construct filter
kxmax = maximum(kˣ)
kymax = maximum(kˣ)
filter = @. (kˣ)^2 + (kʸ)^2 ≤ ((kxmax / 2)^2 + (kymax / 2)^2)
filter = @. abs(kˣ) .+ 0 * abs(kʸ) ≤ 2/3 * kxmax
@. filter = filter * ( 0 * abs(kˣ) .+ 1 * abs(kʸ) ≤ 2/3 * kxmax )

# now define the random field
wavemax = 5
𝓀 = arraytype([-wavemax, 0, wavemax]) # arraytype(1.0 .* [-wavemax, -wavemax + 1, wavemax - 1, wavemax])# arraytype(collect(-wavemax:1:wavemax))
𝓀ˣ = reshape(𝓀, (length(𝓀), 1))
𝓀ʸ = reshape(𝓀, (1, length(𝓀)))
# A = @. 0.1 * (𝓀ˣ * 𝓀ˣ + 𝓀ʸ * 𝓀ʸ)^(-11 / 12)
A = @. 0.5 * (𝓀ˣ * 𝓀ˣ + 𝓀ʸ * 𝓀ʸ)^(0.01) # @. 1e-1 / (1 * 2 * wavemax^2) .* (𝓀ˣ * 𝓀ˣ + 𝓀ʸ * 𝓀ʸ)^(0.0) # ( 1 .+ (0 .* 𝓀ˣ) .* 𝓀ʸ)
A[A.==Inf] .= 0.0
φ = arraytype(2π * rand(size(A)...))
field = arraytype(zeros(N, N))

##
# Fields
# velocity
ψ = arraytype(zeros(ComplexF64, N, N))
u = similar(ψ)
v = similar(ψ)

# theta
θ = similar(ψ)
∂ˣθ = similar(ψ)
∂ʸθ = similar(ψ)
κΔθ = similar(ψ)
θ̇ = similar(ψ)
s = similar(ψ)
θ̅ = similar(ψ)
k₁ = similar(ψ)
k₂ = similar(ψ)
k₃ = similar(ψ)
k₄ = similar(ψ)
θ̃ = similar(ψ)

# source
s = similar(ψ)
@. s = cos(kˣ[5] * x)

# phase
φ̇ = similar(A)

# operators
∂x = im * kˣ
∂y = im * kʸ
Δ = @. ∂x^2 + ∂y^2
Δ⁻¹ = 1 ./ Δ
bools = (!).(isnan.(Δ⁻¹))
Δ⁻¹ .*= bools # hack in the fact that false * NaN = 0

# plan ffts
P = plan_fft!(ψ)
P⁻¹ = plan_ifft!(ψ)

# number of gridpoints in transition is about λ * N / 2
bump(x; λ=10 / N, width=π / 2) = 0.5 * (tanh((x + width / 2) / λ) - tanh((x - width / 2) / λ))
bumps(x; λ=20 / N, width=1.0) = 0.25 * (bump(x, λ=λ, width=width) + bump(x, λ=λ, width=2.0 * width) + bump(x, λ=λ, width=3.0 * width) + bump(x, λ=λ, width=4.0 * width))

##
Δx = x[2] - x[1]

Δt = Δx / (2π) * 1

κ = 1.0 * Δx^2 # /Δt  # 1.0 / N * 0.1 # * 2^(2)  # roughly 1/N for this flow

# take the initial condition as negative of the source
tic = Base.time()

# save some snapshots
ψ_save = typeof(real.(Array(ψ)))[]
θ_save = typeof(real.(Array(ψ)))[]

r_A = Array(@. sqrt((x - π)^2 + (y - π)^2))
θ_A = [bumps(r_A[i, j]) - 3.0 for i in 1:N, j in 1:N]
θ .= CuArray(θ_A)
θ .= @. 0.1 * sin(3 * x) * sin(3 * y) + 0im
# @. θ = bump(sqrt(x^2 + y^2)) # scaling so that source is order 1
θclims = extrema(Array(real.(θ))[:])
P * θ # in place fft
@. κΔθ = κ * Δ * θ
P⁻¹ * κΔθ # in place fft
s .= -κΔθ * 0.0
P⁻¹ * θ # in place fft
θ̅ .= 0.0

t = [0.0]
tend = 40 # 5000

phase_speed = 1.0

iend = ceil(Int, tend / Δt)

#=
function θ_rhs_moisture!(θ̇, θ, params)
    #(; ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter) = params
    ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter = params
    τ = 0.01
    e = 0.01
    event = stream_function!(ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ)
    wait(event)
    P * ψ # in place fft
    P * θ # in place fft
    # ∇ᵖψ
    @. u = filter * -1.0 * (∂y * ψ)
    @. v = filter * (∂x * ψ)
    # ∇θ
    @. ∂ˣθ = filter * ∂x * θ
    @. ∂ʸθ = filter * ∂y * θ
    @. κΔθ = κ * Δ * θ # (κ * Δ  - (κ * Δ)^2 ) * θ
    # go back to real space
    P⁻¹ * ψ
    P⁻¹ * θ
    P⁻¹ * u
    P⁻¹ * v
    P⁻¹ * ∂ˣθ
    P⁻¹ * ∂ʸθ
    P⁻¹ * κΔθ
    # construct source
    @. s = v + e - 1 / τ * θ * (real(θ) > 0)
    # Assemble RHS
    @. θ̇ = -u * ∂ˣθ - v * ∂ʸθ + κΔθ + s
    return nothing
end
=#

function θ_rhs_moisture!(θ̇, θ, params)
    #(; ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter) = params
    ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter, Δ⁻¹ = params
    τ = 0.01
    e = 0.01
    event = stream_function!(ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ)
    wait(event)
    # @. ψ = cos(4 * x)
    @. s = ψ
    P * ψ # in place fft
    P * θ # in place fft
    # quick hack
    @. ψ = Δ⁻¹ * θ
    # ∇ᵖψ
    @. u = filter *  (∂y * ψ)
    @. v = filter * -1.0 * (∂x * ψ)
    # ∇θ
    @. ∂ˣθ = filter * ∂x * θ
    @. ∂ʸθ = filter * ∂y * θ
    @. κΔθ = (Δ⁻¹ + κ * Δ - 1e-1 * (κ * Δ)^2 + 1e-3 * (κ * Δ)^3 - 1e-5 * (κ * Δ)^4) * θ # κ * Δ - 0.1 * (κ * Δ)^2 + 0.01 * (κ * Δ)^3) * θ # κ * Δ * θ #
    # go back to real space
    P⁻¹ * ψ
    P⁻¹ * θ
    P⁻¹ * u
    P⁻¹ * v
    P⁻¹ * ∂ˣθ
    P⁻¹ * ∂ʸθ
    P⁻¹ * κΔθ
    # construct source
    # @. s = ψ# v + e - 1 / τ * θ * (real(θ) > 0)
    # Assemble RHS
    @. θ̇ = real(-u * ∂ˣθ - v * ∂ʸθ + κΔθ + s)
    @. θ = real(θ)
    #
    # P * θ̇    # in place fft
    # @. θ̇ = 0 * θ̇
    # P⁻¹ * θ̇
    return nothing
end

params = (; ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter, Δ⁻¹)

size_of_A = size(A)

realizations = 1
for j in 1:realizations
    # θ .= CuArray(θ_A)
    for i = 1:iend
        # fourth order runge kutta on deterministic part
        # keep ψ frozen is the correct way to do it here

        # the below assumes that φ is just a function of time
        θ_rhs_moisture!(k₁, θ, params)
        @. θ̃ = θ + Δt * k₁ * 0.5

        φ_rhs_normal!(φ̇, φ, rng)
        @. φ += phase_speed * sqrt(Δt / 2) * φ̇

        θ_rhs_moisture!(k₂, θ̃, params)
        @. θ̃ = θ + Δt * k₂ * 0.5
        θ_rhs_moisture!(k₃, θ̃, params)
        @. θ̃ = θ + Δt * k₃

        φ_rhs_normal!(φ̇, φ, rng)
        @. φ += phase_speed * sqrt(Δt / 2) * φ̇

        θ_rhs_moisture!(k₄, θ̃, params)
        @. θ += Δt / 6 * (k₁ + 2 * k₂ + 2 * k₃ + k₄)

        # update stochastic part
        # φ_rhs_normal!(φ̇, φ, rng)
        # @. φ += sqrt(Δt) * φ̇

        t[1] += Δt
        # save output

        #=
        if i % div(iend, 10) == 0
            println("Saving at i=", i)
            push!(ψ_save, Array(real.(ψ)))
            push!(θ_save, Array(real.(θ)))
            println("extrema are ", extrema(θ_save[end]))
            println("time is t = ", t[1])
        end

        if t[1] >= tstart
            θ̅ .+= Δt * θ
        end

        if i % div(iend, 100) == 0
            println("time is t = ", t[1])
            local toc = Base.time()
            println("the time for the simulation is ", toc - tic, " seconds")
            println("extrema are ", extrema(real.(θ)))
            println("on wavenumber index ", index_choice)
        end
        =#
        if i % div(iend, 100) == 0
            println("time is t = ", t[1])
            local toc = Base.time()
            println("the time for the simulation is ", toc - tic, " seconds")
            println("extrema are ", extrema(real.(θ)))
            # println("on wavenumber index ", index_choice)
        end

    end
    println("finished realization ", j)
    @. θ̅ += θ / realizations
end

toc = Base.time()
println("the time for the simulation was ", toc - tic, " seconds")

# calculate vorticity
P * ψ # in place fft
@. κΔθ = Δ * ψ
P⁻¹ * κΔθ # in place fft

x_A = Array(x)[:] .- 2π
θ_F = Array(real.(θ))
ψ_F = Array(real.(ψ))
ω_F = Array(real.(κΔθ))
θ̅_F = Array(real.(θ̅))

fig = Figure(resolution=(2048, 512))
ax1 = Axis(fig[1, 1], title="moisture t = 0")
ax2 = Axis(fig[1, 2], title="moisture t = " * string(tend))
ax3 = Axis(fig[1, 4], title="vorticity t = " * string(tend))
println("the extrema of the end field is ", extrema(θ_F))
println("the extrema of the ensemble average is ", extrema(θ̅_F))
colormap = :bone_1
# colormap = :nipy_spectral
colorrange = (-1, 0)
colorrange2 = extrema(θ_F)
colorrange3 = (-maximum(ω_F), maximum(ω_F))
heatmap!(ax1, x_A, x_A, θ_A, colormap=colormap, colorrange=colorrange2, interpolate=true)
hm = heatmap!(ax2, x_A, x_A, θ_F, colormap=colormap, colorrange=colorrange2, interpolate=true)
hm_e = heatmap!(ax3, x_A, x_A, ω_F, colormap=:balance, colorrange=colorrange3, interpolate=true)
Colorbar(fig[1, 3], hm, height=Relative(3 / 4), width=25, ticklabelsize=30, labelsize=30, ticksize=25, tickalign=1,)
Colorbar(fig[1, 5], hm_e, height=Relative(3 / 4), width=25, ticklabelsize=30, labelsize=30, ticksize=25, tickalign=1,)
display(fig)


using jInv.Mesh
using EikonalInv
using jInv.ForwardShare
using jInv.Utils

using Base.Test

function testSensMatVec(pFor,m0)

	println("\t\tchecking derivatives")

	dobs, = getData(vec(m0),pFor)

	function testFun(m,v=[])
		dobs, = getData(vec(m),pFor)
		if isempty(v)
			return vec(dobs)
		else
			JTv = getSensMatVec(vec(v),vec(m),pFor)
			return vec(dobs), JTv
		end
	end

	chkDer, = jInv.Utils.checkDerivative(testFun,vec(m0))
	@test chkDer

	(nd,nm) = getSensMatSize(pFor)
	v0 = randn(nm)
	w0 = randn(nd)

	J  = getSensMat(vec(m0),pFor)
	J1 = J*v0;
	J2 = J'*w0;

	J1mf = getSensMatVec(v0,vec(m0),pFor)
	J2mf = getSensTMatVec(w0,vec(m0),pFor)

	println("\t\ttest matrix free sensitivities")
	@test norm(J1mf-J1)/norm(J1) < 1e-10
	@test norm(J2mf-J2)/norm(J2) < 1e-10

	println("\t\tadjoint test")
	t1 = dot(w0,J1mf)
	t2 = dot(v0,J2mf)

	p = abs(t1-t2)/abs(t1);
	@test p < 1e-10
	if p > 1e-10
		println("Inconsistency in sensitivity: ",p);
	end

end
############################################################################################
# get 2D mesh with uneven number of cells
domain = [1. 4.0 .4 4.2]
n      = [17;16;]
M      = getRegularMesh(domain,n)
# sources / receivers on top edge
idx    = reshape(collect(1:prod(n+1)),tuple(n+1...))
ib     = idx[1:2:end,1];
I      = speye(prod(n+1))
Q      = I[:,vec(ib)]
R      = copy(Q)

println("2D - test low order scheme")
pFor = getEikonalInvParam(M,Q,R,false)
m0   = rand(prod(n+1))+10.0

try
	testSensMatVec(pFor,vec(m0))
catch
	warn("Sensitivity for low-order scheme and mesh size n=$n FAILED!")
end

############################################################################################
println("2D - test high order scheme")
pFor = getEikonalInvParam(M,Q,R,true)
m0   = rand(prod(n+1))+10.0
try
	testSensMatVec(pFor,vec(m0))
catch
	warn("Sensitivity for higher-order scheme and mesh size n=$n FAILED!")
end

n      = [16;32;]
M      = getRegularMesh(domain,n)
# sources / receivers on top edge
idx    = reshape(collect(1:prod(n+1)),tuple(n+1...))
ib     = idx[1:2:end,1];
I      = speye(prod(n+1))
Q      = I[:,vec(ib)]
R      = copy(Q)
println("2D - test low order scheme")
pFor = getEikonalInvParam(M,Q,R,false)
m0   = rand(prod(n+1))+10.0
try
	testSensMatVec(pFor,vec(m0))
catch
	warn("Sensitivity for lower-order scheme and mesh size n=$n FAILED!")
end

println("2D - test high order scheme")
pFor = getEikonalInvParam(M,Q,R,true)
m0   = rand(prod(n+1))+10.0
try
	testSensMatVec(pFor,vec(m0))
catch
	warn("Sensitivity for higher-order scheme and mesh size n=$n FAILED!")
end
############################################################################################

# get 3D mesh
domain = [1. 4.0 .4 4.2 0 1.2]
n      = [16;16;8]
M      = getRegularMesh(domain,n)
# sources / receivers on top edge
idx    = reshape(collect(1:prod(n+1)),tuple(n+1...))
ib     = idx[1:6:end,1,1:4:end];
I      = speye(prod(n+1))
Q      = I[:,vec(ib)]
R      = copy(Q)
println("3D - test low order scheme")
pFor = getEikonalInvParam(M,Q,R,false)
m0   = rand(prod(n+1))+10.0
try
	testSensMatVec(pFor,vec(m0))
catch
	warn("Sensitivity for low-order scheme and mesh size n=$n FAILED!")
end

println("3D - test high order scheme")
pFor = getEikonalInvParam(M,Q,R,true)
m0   = rand(prod(n+1))+10.0
try
	testSensMatVec(pFor,vec(m0))
catch
	warn("Sensitivity for higher-order scheme and mesh size n=$n FAILED!")
end


#some schemes
#Emb 4/3 AK p
embedded_scheme_43 = EmbeddedScheme(
          ( 0.125962888700250514,  0.333588446797901933,
            0.751193431379145450, -0.338296598434303506,
            0.127551831557005609,  0.127551831557005609,
           -0.338296598434303506,  0.751193431379145450,
            0.333588446797901933,  0.125962888700250514 ),

          ( 0.125962888700250514,  0.333588446797901933,
            0.751193431379145450, -0.338296598434303506,
            0.0                 ,  0.261153550449697153,
           -0.242703571757396124,  0.596114052266110425,
            0.365547251678000160,  0.147440548920593995 ),

            4 )
#PP 3/4 A
palindromic_scheme_34 = PalindromicScheme(
          ( 0.268330095781759925,  0.919661523017399857,
           -0.187991618799159782, -0.187991618799159782,
            0.919661523017399857,  0.268330095781759925 ),
            3 )
#PP 5/6 A
palindromic_scheme_56 = PalindromicScheme(
          ( 0.201651044312324230,   0.578800656272664932,
            0.562615975356569200,   0.273128836056524479,
            0.253874038247554845,  -0.102733803148432142,
           -0.835351693190370636,   0.068014946093165092,
            0.068014946093165092,  -0.835351693190370636,
           -0.102733803148432142,   0.253874038247554845,
            0.273128836056524479,   0.562615975356569200,
            0.578800656272664932,   0.201651044312324230 ),
            5 )
#defectbased
embedded_scheme_43_D = DefectBasedScheme(
            embedded_scheme_43.scheme2,
            embedded_scheme_43.order-1 )
palindromic_scheme_34_D = DefectBasedScheme(
            palindromic_scheme_34.scheme,
            palindromic_scheme_34.order )
palindromic_scheme_56_D = DefectBasedScheme(
            palindromic_scheme_56.scheme,
            palindromic_scheme_56.order )
schemes=[embedded_scheme_43  ,palindromic_scheme_34  ,palindromic_scheme_56,
         embedded_scheme_43_D,palindromic_scheme_34_D,palindromic_scheme_56_D];


using PDMats
using Test

_randPDMat(T, n) = (X = randn(T, n, n); PDMat(X * X'))
_randPDiagMat(T, n) = PDiagMat(rand(T, n))
_randScalMat(T, n) = ScalMat(n, rand(T))
function _pd_compare(A::AbstractPDMat, B::AbstractPDMat)
    @test dim(A) == dim(B)
    @test Matrix(A) ≈ Matrix(B)
    @test cholesky(A).L ≈ cholesky(B).L
    @test cholesky(A).U ≈ cholesky(B).U
    nothing
end
function _pd_kron_compare(A::AbstractPDMat, B::AbstractPDMat)
    PDAkB1 = kron(A, B)
    PDAkB2 = PDMat( kron( Matrix(A), Matrix(B) ) )
    _pd_compare(PDAkB1, PDAkB2)
    nothing
end

n = 4
m = 7

for T in [Float32, Float64]
    _pd_kron_compare( _randPDMat(T, n),    _randPDMat(T, m) )
    _pd_kron_compare( _randPDiagMat(T, n), _randPDiagMat(T, m) )
    _pd_kron_compare( _randScalMat(T, n),  _randScalMat(T, m) )
    _pd_kron_compare( _randPDMat(T, n),    _randPDiagMat(T, m) )
    _pd_kron_compare( _randPDiagMat(T, m), _randPDMat(T, n) )
    _pd_kron_compare( _randPDMat(T, n),    _randScalMat(T, m) )
    _pd_kron_compare( _randScalMat(T, m),  _randPDMat(T, n) )
    _pd_kron_compare( _randPDiagMat(T, n), _randScalMat(T, m) )
    _pd_kron_compare( _randScalMat(T, m),  _randPDiagMat(T, n) )
end


using OnlineStats, Test, LinearAlgebra, Random, StatsBase, Statistics, Dates
using LossFunctions, PenaltyFunctions

#-----------------------------------------------------------------------# utils
n = 1000
x,  y,  z  = rand(Bool, n), randn(n), rand(1:10, n)
x2, y2, z2 = rand(Bool, n), randn(n), rand(1:10, n)
xs, ys, zs = vcat(x, x2), vcat(y, y2), vcat(z, z2)

p = 5
xmat,  ymat,  zmat  = rand(Bool, n, p), randn(n, p), rand(1:10, n, p)
xmat2, ymat2, zmat2 = rand(Bool, n, p), randn(n, p), rand(1:10, n, p)

function mergestats(a::OnlineStat, y1, y2)
    b = copy(a)
    fit!(a, y1)             # fit a on y1
    fit!(b, y2)             # fit b on y2
    merge!(a, b)            # merge b into a
    fit!(b, y1)             # fit b on y1
    @test nobs(a) == nobs(b) == length(y1) + length(y2)
    a, b
end
mergevals(o1::OnlineStat, y1, y2) = map(value, mergestats(o1, y1, y2))

#-----------------------------------------------------------------------# AutoCov
@testset "AutoCov" begin
    o = fit!(AutoCov(10), y)
    @test autocov(o) ≈ autocov(y, 0:10)
    @test autocor(o) ≈ autocor(y, 0:10)
    @test nobs(o) == n
end
#-----------------------------------------------------------------------# Bootstrap
@testset "Bootstrap" begin
    o = fit!(Bootstrap(Mean(), 100, [1]), y)
    @test all(value.(o.replicates) .== value(o.stat))
    c = confint(o)
    @test length(c) == 2
    @test c[1] ≤ c[2]
end
#-----------------------------------------------------------------------# CallFun
@testset "CallFun" begin
    i = 0
    o = fit!(CallFun(Mean(), x -> i+=1), y)
    @test value(o) ≈ mean(y)
    @test i == n
    @test ≈(mergevals(CallFun(Mean(), x->nothing), y, y2)...)
end
#-----------------------------------------------------------------------# CovMatrix
@testset "CovMatrix" begin
    o = fit!(CovMatrix(), eachrow(ymat))
    @test value(o) ≈ cov(ymat)
    @test cov(o) ≈ cov(ymat)
    @test cor(o) ≈ cor(ymat)
    @test all(x -> ≈(x...), zip(var(o), var(ymat; dims=1)))
    @test all(x -> ≈(x...), zip(std(o), std(ymat; dims=1)))
    @test all(x -> ≈(x...), zip(mean(o), mean(ymat; dims=1)))

    @test ≈(mergevals(CovMatrix(), OnlineStatsBase.eachrow(ymat), OnlineStatsBase.eachrow(ymat2))...)
    @test ≈(mergevals(CovMatrix(), OnlineStatsBase.eachcol(ymat'), OnlineStatsBase.eachcol(ymat2'))...)
    @test ≈(mergevals(CovMatrix(Complex{Float64}), OnlineStatsBase.eachrow(ymat * im), OnlineStatsBase.eachrow(ymat2))...)
    @test ≈(mergevals(CovMatrix(Complex{Float64}), OnlineStatsBase.eachrow(ymat * im), OnlineStatsBase.eachrow(ymat2 * im))...)
end
#-----------------------------------------------------------------------# Diff
@testset "Diff" begin
    @test value(fit!(Diff(), y)) == y[end] - y[end-1]
    o = fit!(Diff(Int), 1:10)
    @test diff(o) == 1
    @test last(o) == 10
end
#-----------------------------------------------------------------------# Fit[Dist]
@testset "Fit[Dist]" begin
    @testset "FitBeta" begin
        @test value(FitBeta()) == (1.0, 1.0)
        a, b = mergevals(FitBeta(), rand(100), rand(100))
        @test a[1] ≈ b[1]
        @test a[2] ≈ b[2]
    end
    @testset "FitCauchy" begin
        @test value(FitCauchy()) == (0.0, 1.0)
        a, b = mergevals(FitCauchy(), y, y2)
        @test ≈(a[1], b[1]; atol=.5)
        @test ≈(a[2], b[2]; atol=.5)
    end
    @testset "FitGamma" begin
        @test value(FitGamma()) == (1.0, 1.0)
        a, b = mergevals(FitGamma(), rand(100), rand(100))
        @test a[1] ≈ b[1]
        @test a[2] ≈ b[2]
    end
    @testset "FitLogNormal" begin
        @test value(FitLogNormal()) == (0.0, 1.0)
        a, b = mergevals(FitLogNormal(), exp.(y), exp.(y2))
        @test a[1] ≈ b[1]
        @test a[2] ≈ b[2]
    end
    @testset "FitNormal" begin
        @test value(FitNormal()) == (0.0, 1.0)
        a, b = mergestats(FitNormal(), y, y2)
        @test value(a)[1] ≈ value(b)[1]
        @test value(a)[2] ≈ value(b)[2]

        @test mean(a) ≈ mean(ys)
        @test var(a) ≈ var(ys)
        @test std(a) ≈ std(ys)

        # pdf and cdf
        o = fit!(FitNormal(), [-1, 0, 1])
        @test OnlineStats.pdf(o, 0.0) ≈ 0.3989422804014327
        @test OnlineStats.pdf(o, -1.0) ≈ 0.24197072451914337
        @test OnlineStats.cdf(o, 0.0) ≈ 0.5
        @test ≈(OnlineStats.cdf(o, -1.0), 0.15865525393145702; atol=.001)
    end
    @testset "FitMultinomial" begin
        o = FitMultinomial(5)
        @test value(o)[2] == ones(5) / 5
        data = [1 2 3 4 5; 1 2 3 4 5]
        @test value(fit!(o, eachrow(data)))[2] == collect(2:2:10) ./ sum(data)

        data1 = OnlineStatsBase.eachrow(rand(1:4, 10, 3))
        data2 = OnlineStatsBase.eachrow(rand(2:7, 11, 3))
        a, b = mergevals(FitMultinomial(3), data1, data2)
        @test ≈(a[2], b[2])
    end
    @testset "FitMvNormal" begin
        @test value(FitMvNormal(2)) == (zeros(2), Matrix(I, 2, 2))
        a, b = mergevals(FitMvNormal(2), OnlineStatsBase.eachrow([y y2]), OnlineStatsBase.eachrow([y2 y]))
        @test a[1] ≈ b[1]
        @test a[2] ≈ b[2]
    end
end
#-----------------------------------------------------------------------# FastNode
@testset "FastNode" begin
    X, Y = ymat, rand(1:3, 1000)
    X2, Y2 = ymat, rand(1:3, 1000)
    data  = zip(eachrow(X), Y)
    data2 = zip(eachrow(X2), Y2)
    o  = fit!(FastNode(5, 3), data)
    o2 = fit!(FastNode(5, 3), data2)
    merge!(o, o2)
    fit!(o2, data)
    for k in 1:3, j in 1:5
        @test value(o.stats[k][j])[1] ≈ value(o2.stats[k][j])[1]
        @test value(o.stats[k][j])[2] ≈ value(o2.stats[k][j])[2]
    end
    @test length(o[1]) == 3

    pvec = [mean(Y .== 1), mean(Y .== 2), mean(Y .== 3)]
    o = fit!(FastNode(5, 3), data)
    @test probs(o) == pvec
    @test nobs(o) == 1000
    @test OnlineStats.nkeys(o) == 3
    @test OnlineStats.nvars(o) == 5
    @test classify(o) ∈ [1, 2, 3]
end
#-----------------------------------------------------------------------# FastTree
@testset "FastTree" begin
    X, Y = OnlineStats.fakedata(FastNode, 10^4, 10)
    o = fit!(FastTree(10; splitsize=100), zip(eachrow(X),Y))
    @test classify(o, X[1,:]) ∈ [1, 2]
    @test all(0 .< classify(o, X) .< 3)
    @test OnlineStats.nkeys(o) == 2
    @test OnlineStats.nvars(o) == 10
    @test mean(classify(o, X) .== Y) > .5

    # Issue 116
    Random.seed!(218)
    X,Y = OnlineStats.fakedata(FastNode, 10^4, 1)
    fit!(FastTree(1, splitsize=100), zip(eachrow(X),Y))
end
#-----------------------------------------------------------------------# FastForest
@testset "FastForest" begin
    X, Y = OnlineStats.fakedata(FastNode, 10^4, 10)
    o = fit!(FastForest(10; splitsize=500, λ = .7), zip(eachrow(X), Y))
    @test classify(o, randn(10)) in 1:2
    @test mean(classify(o, X) .== Y) > .5
end

#-----------------------------------------------------------------------# HeatMap
@testset "HeatMap" begin
    data1 = OnlineStatsBase.eachrow(ymat[:, 1:2])
    data2 = OnlineStatsBase.eachrow(ymat2[:, 1:2])
    @test ==(mergevals(HeatMap(-5:.1:5, -5:.1:5), data1, data2)...)
end
#-----------------------------------------------------------------------# Hist
@testset "Hist" begin
    @test ==(mergevals(Hist(-5:.1:5), y, y2)...)
    @testset "Hist compared to StatsBase.Histogram" begin
        for edges in (-5:5, collect(-5:5), [-5, -3.5, 0, 1, 4, 5.5])
            for data in (y, -6:.75:6)
                w  = fit(Histogram, data, edges, closed = :left).weights
                w2 = fit(Histogram, data, edges, closed = :right).weights
                @test fit!(Hist(edges, Number; closed=false, left=true),  data).counts == w
                @test fit!(Hist(edges, Number; closed=false, left=false), data).counts == w2
            end
        end
    end
    o = fit!(Hist(-5:.1:5), y)
    for (v1, v2) in zip(extrema(o), extrema(y))
        @test ≈(v1, v2; atol=.1)
    end
    @test ≈(mean(o), mean(y); atol=.1)
    @test ≈(var(o), var(y); atol=.2)

    # merge unequal bins
    r1, r2 = -5:.2:5, -5:.1:5
    @test merge!(fit!(Hist(r1), y), fit!(Hist(r2), y2)) == fit!(Hist(r1), vcat(y, y2))
    @test OnlineStats.pdf(fit!(Hist(-5:.1:5), y), 0) > 0
    @test OnlineStats.pdf(fit!(Hist(-5:.1:5), y), 100) == 0
end
#-----------------------------------------------------------------------# KHist
@testset "KHist" begin
    @test KHist(10) == KHist(10)

    o = fit!(KHist(1000), y)
    @test mean(o) ≈ mean(y)
    @test var(o) ≈ var(y)
    @test median(o) ≈ median(y)
    @test quantile(o) ≈ quantile(y, [0, .25, .5, .75, 1])
    @test std(o) ≈ std(y)
    @test extrema(o) == extrema(y)

    ==(mergevals(KHist(2000), y, y2)...)
    ==(mergevals(KHist(1), y, y2)...)
    ==(mergevals(KHist(2000, Float32), Float32.(y), Float32.(y2))...)

    data = randn(10_000)
    o = fit!(KHist(50), data)
    @test OnlineStats.pdf(o, -10) == 0.0
    @test ≈(OnlineStats.pdf(o, 0.0), 0.3989422804014327, atol=.5)
    @test OnlineStats.pdf(o, 10) == 0.0
    @test OnlineStats.cdf(o, -10) == 0.0
    @test ≈(OnlineStats.cdf(o, 0.0), .5; atol=.1)
    @test OnlineStats.cdf(o, 10) == 1.0
end
#-----------------------------------------------------------------------# HyperLogLog
@testset "HyperLogLog" begin
    @test ≈(value(fit!(HyperLogLog(), y)), n; atol=20)
    @test ==(mergevals(HyperLogLog(), x, x2)...)
    @test ==(mergevals(HyperLogLog(), y, y2)...)
    @test ==(mergevals(HyperLogLog(), z, z2)...)
end
#-----------------------------------------------------------------------# IndexedPartition
@testset "IndexedPartition" begin
    o = IndexedPartition(Float64, Mean())
    fit!(o, zip(y, y2))
end
#-----------------------------------------------------------------------# KahanSum
@testset "KahanSum" begin
    @test value(fit!(KahanSum(), y)) ≈ sum(y)
    @test value(fit!(KahanSum(Int), x)) == sum(x)
    @test value(fit!(KahanSum(Int), z)) == sum(z)
    @test ≈(mergevals(KahanSum(), y, y2)...)
end
#-----------------------------------------------------------------------# KahanMean
@testset "KahanMean" begin
    @test value(fit!(KahanMean(), y)) ≈ mean(y)
    @test ≈(mergevals(KahanMean(), y, y2)...)
end
#-----------------------------------------------------------------------# KahanVariance
@testset "KahanVariance" begin
    o = fit!(KahanVariance(), y)
    @test mean(o) ≈ mean(y)
    @test var(o) ≈ var(y)
    @test std(o) ≈ std(y)
    @test ≈(mergevals(KahanVariance(), y, y2)...)

    # Issue 116
    @test std(KahanVariance()) == 1
    @test std(fit!(KahanVariance(), 1)) == 1
    @test std(fit!(KahanVariance(), [1, 2])) == sqrt(.5)
end
#-----------------------------------------------------------------------# KMeans
@testset "KMeans" begin
    o = fit!(KMeans(5,2), eachrow(ymat))
    sort!(o, rev=true)
    @test o.value[1].n ≥ o.value[2].n
end
#-----------------------------------------------------------------------# LinReg
@testset "LinReg" begin
    ≈(mergevals(LinReg(), zip(eachrow(ymat), y), zip(eachrow(ymat2), y2))...)

    o = fit!(LinReg(), zip(eachrow(ymat), y))
    @test coef(o) ≈ ymat \ y
    @test coef(o, .1) ≈ (ymat'ymat ./ n + .1I) \ ymat'y ./ n
    @test coef(o, .1:.1:.5) ≈ (ymat'ymat ./ n + Diagonal(.1:.1:.5)) \ ymat'y ./ n
    @test predict(o, ymat) == ymat * o.β
    @test predict(o, ymat[1,:]) == dot(ymat[1,:], o.β)
end
#-----------------------------------------------------------------------# LinRegBuilder
@testset "LinRegBuilder" begin
    @test ≈(mergevals(LinRegBuilder(), OnlineStatsBase.eachrow(ymat), OnlineStatsBase.eachrow(ymat2))...)

    o = fit!(LinRegBuilder(), eachrow(ymat))
    for i in 1:5
        data = ymat[:, setdiff(1:5, i)]
        @test coef(o; y=i) ≈ [data ones(n)] \ ymat[:,i]
        @test coef(o, .1; y=i, bias=false) ≈ (data'data ./ n + .1*I) \ data'ymat[:,i] ./ n
    end

    o2 = fit!(LinReg(), zip(eachrow(ymat[:,[4,1]]), ymat[:,3]))
    @test coef(o, [.2,.4]; y=3, x = [4,1], bias=false) ≈ coef(o2, [.2, .4])
end
#-----------------------------------------------------------------------# ML
@testset "ML" begin
    o = OnlineStats.preprocess(OnlineStatsBase.eachrow(ymat))
    for i in 1:5
        @test o.group[i] isa OnlineStats.Numerical
    end
    o = OnlineStats.preprocess(OnlineStatsBase.eachrow(xmat))
    for i in 1:2
        @test o.group[i] isa OnlineStats.Categorical
    end
    o = OnlineStats.preprocess(OnlineStatsBase.eachrow(zmat), 3 => OnlineStats.Categorical(Int))
    @test o.group[3] isa OnlineStats.Categorical
    @test o.group[1] isa OnlineStats.Numerical
end

#-----------------------------------------------------------------------# Mosaic
@testset "Mosaic" begin
    @test ==(mergevals(Mosaic(Int,Int), zip(z, z2), zip(z2, z))...)
end
#-----------------------------------------------------------------------# MovingTimeWindow
@testset "MovingTimeWindow" begin
    dates = Date(2010):Day(1):Date(2011)
    data = Int.(1:length(dates))
    o = fit!(MovingTimeWindow(Day(4); timetype=Date, valtype=Int), zip(dates, data))
    @test value(o) == collect(Pair(a,b) for (a,b) in zip(dates[end-4:end], data[end-4:end]))

    d1 = zip(dates[1:2], data[1:2])
    d2 = zip(dates[3:4], data[3:4])
    @test ==(mergevals(MovingTimeWindow(Day(4); timetype=Date, valtype=Int), d1, d2)...)
end
#-----------------------------------------------------------------------# MovingWindow
@testset "MovingWindow" begin
    o = fit!(MovingWindow(10, Int), 1:12)
    for i in 1:10
        @test o[i] == (1:12)[i + 2]
    end
end
#-----------------------------------------------------------------------# NBClassifier
@testset "NBClassifier" begin
    o = fit!(NBClassifier(5, Bool), zip(eachrow(ymat),x))
    merge!(o, fit!(NBClassifier(5, Bool), zip(eachrow(ymat2), x2)))
    @test nobs(o) == 2000
    @test length(probs(o)) == 2
    @test sum(predict(o, ymat[1,:])) ≈ 1
    @test classify(o, ymat[1, :]) || !classify(o, ymat[1, :])
    @test OnlineStats.nvars(o) == 5
    @test OnlineStats.nkeys(o) == 2
    @test length(o[2]) == 2
end
#-----------------------------------------------------------------------# OrderStats
@testset "OrderStats" begin
    @test ≈(mergevals(OrderStats(100), y, y2)...)
    o = fit!(OrderStats(n), y)
    @test value(o) == sort(y)
    @test quantile(o, 0:.25:1) == quantile(y, 0:.25:1)
end
#-----------------------------------------------------------------------# Partition
@testset "Partition" begin
    # merging
    o = fit!(Partition(Mean(), 1000), y)
    o2 = fit!(Partition(Mean(), 1000), y2)
    merge!(o, o2)
    fit!(o2, y)
    @test nobs(o) == nobs(o2)
    @test all(nobs.(o.parts) .== nobs.(o2.parts))
    for i in 1:5
        @test value(o.parts[i]) ≈ value(o2.parts[500 + i])
    end
end
#-----------------------------------------------------------------------# ProbMap
@testset "ProbMap" begin
    @test sort(collect(keys(fit!(ProbMap(Float64), y).value))) == sort(y)
    # merge
    data, data2 = rand(1:4, 100), rand(1:4, 100)
    o = fit!(ProbMap(Int), data)
    o2 = fit!(ProbMap(Int), data2)
    merge!(o, o2)
    fit!(o2, data)
    @test sort(collect(keys(o.value))) == sort(collect(keys(o2.value)))
    @test probs(fit!(ProbMap(Int), [1,1,2,2,3,3,4,4])) ≈ fill(.25, 4)
    @test probs(fit!(ProbMap(Int), [1,1,2,2,3,3,4,4]), [1,2,9]) ≈ [.5, .5, 0]
end
#-----------------------------------------------------------------------# Quantile
@testset "Quantile/P2Quantile" begin
    data = randn(10_000)
    data2 = randn(10_000)
    τ = .1:.1:.9
    for o in [
            Quantile(τ; alg=SGD()),
            Quantile(τ; alg=MSPI()),
            Quantile(τ; alg=OMAS()),
            Quantile(τ; alg=ADAGRAD())
            ]
        @test ≈(value(fit!(copy(o), data)), quantile(data, τ), atol=.5)
        @test ≈(mergevals(o, data, data2)...; atol=.5)

    end
    for τi in τ
        @test ≈(value(fit!(P2Quantile(τi),data)), quantile(data, τi), atol=.2)
    end
end
#-----------------------------------------------------------------------# ReservoirSample
@testset "ReservoirSample" begin
    @test value(fit!(ReservoirSample(n), y)) == y
    a, b = mergestats(ReservoirSample(9), y, y2)
    for yi in value(a)
        @test (yi ∈ y) || (yi ∈ y2)
    end
end
#-----------------------------------------------------------------------# StatLag
@testset "StatLag" begin
    o = fit!(StatLag(Mean(), 10), 1:20)
    @test length(o.lag.value) == 10
    for (i, m) in enumerate(reverse(o.lag.value))
        @test nobs(m) == 10 + i
    end
    @test nobs(value(o)[1]) == 20
end
#-----------------------------------------------------------------------# StatLearn
@testset "StatLearn" begin
    X = randn(10_000, 5)
    β = collect(-1:.5:1)
    Y = X * β + randn(10_000)
    Y2 = 2.0 .* [rand()< 1 /(1 + exp(-η)) for η in X*β] .- 1.0
    for A in [SGD(),ADAGRAD(),ADAM(),ADAMAX(),ADADELTA(),RMSPROP(),OMAS(),OMAP(),MSPI()]
        print("  > $A")
        print(": ")
        for L in [.5 * L2DistLoss()]
            print(" | $L")
            # sanity checks
            o = fit!(StatLearn(5, A, L; rate=LearningRate(.7)), zip(eachrow(X),Y))
            @test o.loss isa typeof(L)
            @test o.alg isa typeof(A)
            any(isnan.(o.β)) && @info((L, A))
            merge!(o, copy(o))
            @test coef(o) == o.β
            @test predict(o, X) == X * o.β
            @test ≈(coef(o), β; atol=1.5)
            @test OnlineStats.objective(o, X, Y) ≈ value(o.loss, Y, predict(o, X), AggMode.Mean()) + value(o.penalty, o.β, o.λ)
        end
        for L in [LogitMarginLoss(), DWDMarginLoss(1.0)]
            print(" | $L")
            o = fit!(StatLearn(5, A, L), zip(eachrow(X),Y2))
            @test mean(Y2 .== classify(o, X)) > .5
        end
        println()
    end
end
#-----------------------------------------------------------------------# Kahan

include("test_kahan.jl")

#-----------------------------------------------------------------------# Show methods
@testset "Show methods" begin
    for stat in [BiasVec([1,2,3]), Bootstrap(Mean()), CallFun(Mean(), println), FastNode(5),
                 FastTree(5), FastForest(5), FTSeries(Variance()), Group(Mean(), Mean()),
                 HyperLogLog{10}(), LinRegBuilder(4), NBClassifier(5, Float64), ProbMap(Int),
                 P2Quantile(.5), Series(Mean()), StatLearn(5)]
        println("  > ", stat)
    end
end


#=Copyright (c) 2017 Gabriel Goh

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.=#

using Plots

# A = ones(18,67)

# A[2:end-1,2:end-1] = 0
# A[1:7,32-5:34+5] = 1
# A[end-7:end,32-5:34+5] = 1

Box = zeros(16,16)
Cavity = ones(16,16)
Cavity[8:9,:] = 0
R = [Box Cavity Box Cavity Box Cavity Box]
A = ones(size(R,1)+2, size(R,2) + 2)
A[2:end-1,2:end-1] = R

n = size(A,1) - 2
m = size(A,2) - 2

R = reshape(1:n*m, n,m)
Id(i,j) = R[i-1,j-1]

Q = zeros(n*m,n*m)
for i = 2:(size(A,1)-1), j = 2:(size(A,2)-1)
  k = Id(i,j)
  if A[i,j] == 0
    A[i,j+1] == 0 ? Q[Id(i,j+1),k] = 1 : nothing
    A[i+1,j] == 0 ? Q[Id(i+1,j),k] = 1 : nothing
    A[i,j-1] == 0 ? Q[Id(i,j-1),k] = 1 : nothing
    A[i-1,j] == 0 ? Q[Id(i-1,j),k] = 1 : nothing
  end
end

Q = (Q + Q')/2

v = sum(Q,1)[:]
v[v .== 0] = 1
L = diagm(v) - Q
L[1,1] = 2.1

Σ,U = eig(L)

pyplot(leg=false, ticks=nothing)
@gif for i in 2:500
  plot(size = (880,140)); heatmap!(reshape(U[:,224*3+i],n,m))
end

e1 = zeros(n*m)
e1[8:9] = 1

α = 1/max(eig(L)[1]...)
iters = 5000
X = zeros(n*m,iters)
x = 10*zeros(n*m)
for i = 1:iters
  X[:,i] = x
  for i = 1:100
    x = x - (1*α)*(L*x - e1)
  end
  println(norm(L*x - e1))
end

plot(size = (880,140)); heatmap!(reshape(X[:,100],n,m))
pyplot(leg=false, ticks=nothing)

@gif for i in 10:10:2000
  println(i)
  xi = reshape(X[:,i],n,m)
  plot(size = (880,140)); heatmap!(xi)
end

lambdak(k,λ) = α*(1 - (1 - α*λ).^k)./(α*λ)
c = U'e1

@gif for i in 10:500:300000
  println(i)
  z = U*(lambdak(i,Σ).*(c)
  plot(size = (880,140)); heatmap!(reshape(z, n, m))
end

function getiterk(Q,b)
  Σ,U = eig(Q)
  c = U'b
  lambdak(k,λ) = α*(1 - (1 - α*λ).^k)./(α*λ)
  return U*(lambdak(k,Σ).*c)
end


using DelimitedFiles

"Load numbers from file."
function loaddata(file::String)
    return readdlm(file, '\t', Int, '\n')
end

"Find the two entries that sum to 2020."
function findtwo(data::Array{Int64,2})
    return [x * y for x in data, y in data if x + y == 2020][1]
end

"Find the three entries that sum to 2020."
function findthree(data::Array{Int64,2})
    return [x * y * z for x in data, y in data, z in data if x + y + z == 2020][1]
end

function main()
    data = loaddata("../assets/data/01.txt")
    println(typeof(data))

    # part one
    println("The product of the two numbers is ", findtwo(data), ".")

    # part two
    println("The product of the three numbers is ", findthree(data), ".")
end

main()


using SIMD

function transpose_blocks!(data::AbstractVector{T}, ::Val{V}, ::Val{J}) where {T,V,J}
    block_size = V*J
    num_blocks = div(length(data), block_size)

    for k in 1:num_blocks
        block = ntuple(l->vload(Vec{V, T}, data, 1 + (k-1)*block_size + (l-1)*V), Val(J))
        transposed_block = transpose_vecs(block...)
        for v in 1:V
            vstore(transposed_block[v], data, 1 + (k-1)*block_size + (v-1)*J)
        end
    end
    data
end

"""In-place transpose of a 3 dimensional array VKJ into VJK."""
@generated function transpose!(data::AbstractVector{T}, ::Val{V}, ::Val{J}, ::Val{K}) where {T,V,J,K}
    seeds = find_transpose_cycles(Val(K), Val(J))
    :(transpose!(data, Val(V), Val(J), Val(K), $seeds))
end

@inline function transpose!(data::AbstractVector{T}, ::Val{V}, ::Val{J}, ::Val{K}, cycle_seeds) where {T,V,J,K}
    kjm = K*J-1
    @inbounds for c in cycle_seeds
        a = c
        base = pointer(data)
        vlen = V * sizeof(T)
        v1 = vloada(Vec{V, T}, base + vlen * a)
        while true
            pa = (a * K) % kjm
            if pa == c break end
            vpa = vload(Vec{V, T}, base + vlen * pa)
            vstore(vpa, base + vlen * a)
            a = pa
        end
        vstore(v1, base + vlen * a)
    end
    nothing
end

"""Find the cycle seeds to transpose a matrix with K rows and J columns into J rows and K columns."""
function find_transpose_cycles(::Val{J}, ::Val{K}) where {J,K}
    cycles = BitSet(1:K*J-2)
    cycle_seeds = Int32[]
    while length(cycles) > 0
        a = popfirst!(cycles)
        cmin = a
        while true
            pa = (a * K) % (K*J-1)
            if pa ∉ cycles break end
            cmin = min(cmin, pa)
            setdiff!(cycles, Set([pa]))
            a = pa
        end
        push!(cycle_seeds, cmin)
    end
    (cycle_seeds...,)
end

"""
    transpose_vecs(input::Vararg{Vec{N,T}, L}) where {L,N,T}

Transposes a matrix of L vectors of size N into N vectors of size L. Sizes should be powers of 2.
"""
@generated function transpose_vecs(input::Vararg{Vec{N,T}, L})::Vararg{N, Vec{L,T}} where {L,N,T}

    ex = Expr[Expr(:meta, :inline)]

    pa = Val(ntuple(a->((a-1)*N)%(2*N-1), Val(N)))
    pb = Val(ntuple(a->div(N,2)+((a-1)*N)%(2*N-1), Val(N)))

    for t in 1:L
        a1 = Symbol("input_", 0, "_", t)
        push!(ex, :($a1 = input[$t]))
    end

    nsteps = mylog(L)

    L_2 = div(L,2)

    for st in 1:nsteps
        for t in 1:L_2
            a1 = Symbol("input_", st-1, "_", t)
            a2 = Symbol("input_", st-1, "_", t+L_2)
            b1 = Symbol("input_", st, "_", t*2-1)
            b2 = Symbol("input_", st, "_", t*2)
            append!(ex, [:($b1 = shufflevector($a1, $a2, $pa)),
                         :($b2 = shufflevector($a1, $a2, $pb))])
        end
    end

    ## Reshape matrix if not square
    if N < L

        outsteps = mylog(div(L,N))
        for st in 1:outsteps
            pcat = Val(ntuple(a->a-1, N*2^st))
            for t in 1:N*2^(outsteps-st)
                a1 = Symbol("input_", nsteps+st-1, "_", t*2-1)
                a2 = Symbol("input_", nsteps+st-1, "_", t*2)
                b1 = Symbol("input_", nsteps+st, "_", t)
                push!(ex, :($b1 = shufflevector($a1, $a2, $pcat)))
            end
        end

        push!(ex, Expr(:tuple, ntuple(t->Symbol("input_", nsteps+outsteps, "_", t), Val(N))...))

    elseif N > L

        outsteps = mylog(div(N,L))
        for st in 1:outsteps
            pleft = Val(ntuple(a->a-1, div(N,2^st)))
            pright = Val(ntuple(a->a-1+div(N,2^st), div(N,2^st)))
            for t in 1:L*2^(st-1)
                a1 = Symbol("input_", nsteps+st-1, "_", t)
                b1 = Symbol("input_", nsteps+st, "_", t*2-1)
                b2 = Symbol("input_", nsteps+st, "_", t*2)
                append!(ex, [:($b1 = shufflevector($a1, $pleft)),
                             :($b2 = shufflevector($a1, $pright))])
            end
        end

        push!(ex, Expr(:tuple, ntuple(t->Symbol("input_", nsteps+outsteps, "_", t), Val(N))...))

    else

        push!(ex, Expr(:tuple, ntuple(t->Symbol("input_", nsteps, "_", t), Val(L))...))

    end

    quote $(ex...) end
end


export is_braket
export Ket, Bra

"""
    is_braket( ψ :: Vector; atol=ATOL :: Float64) :: Bool

    is_braket(
        ψ :: Adjoint{T, Vector{T}};
        atol=ATOL :: Float64
    ) where T <: Number :: Bool

Returns `true` if vector `ψ` is a valid bra (or ket):

* `ψ` is a real or complex-valued vector.
* `ψ` is normalized with respect to the bra-ket inner prodcut (`ψ' * ψ == 1`).
"""
function is_braket(ψ :: Vector{<:Number}; atol=ATOL :: Float64)
    isapprox( norm(ψ,2), 1, atol=atol )
end
function is_braket(ψ :: Adjoint{T, Vector{T}}; atol=ATOL :: Float64) where T <: Number
    isapprox( norm(ψ,2), 1, atol=atol )
end

"""
    _braket_error(ψ)

Throws the `DomainError`, `"ψ must be a vector that satisfies \`norm(ψ,2) != 1\`."`.
"""
_braket_error(ψ) = throw(DomainError(ψ, "ψ must be a vector that satisfies `norm(ψ,2) != 1`."))

"""
    Ket(ψ :: Vector{T}; atol=ATOL :: Float64) where T <: Number

A normalized column vector on a complex-valued Hilbert space. If called with
an adjoint vector `ψ'`, the adjoint will be taken on construction.

```julia
Ket(ψ :: Adjoint{Vector{T}}; atol=ATOL :: Float64) where T <: Number
```

Similarly a `Bra` can be converted into a `Ket` via the adjoint.

```julia
Ket(ψ :: Bra{T}; atol=ATOL :: Float64) where T <: Number
```

`Ket(ψ)`, throws a `DomainError` if `ψ` is not normalized.
"""
struct Ket{T<:Number} <: AbstractVector{Number}
    ψ :: Vector{T}
    atol :: Float64
    Ket(
        ψ :: Vector{T}; atol=ATOL :: Float64
    ) where T <: Number = is_braket(ψ, atol=atol) ? new{T}(ψ, atol) : _braket_error(ψ)
    Ket(
        ψ :: Adjoint{T,Vector{T}}; atol=ATOL :: Float64
    ) where T <: Number = is_braket(ψ, atol=atol) ? new{T}(
        ψ', atol
    ) : _braket_error(ψ)
end
Base.size(ket::Ket) = size(ket.ψ)
Base.getindex(ket::Ket, I::Int) = getindex(ket.ψ, I...)
Base.setindex!(ket::Ket, val, I::Int) = (ket.ψ[I...] = val)
is_braket(::Ket) = true

"""
    Bra(ψ :: Adjoint{T,Vector{T}}; atol=ATOL :: Float64) where T <: Number

A row vector on a complex valued HIlbet space. Bras are dual to Kets via the adjoint
If called with an `Vector{<:Number}` type, the adjoint will automatically be taken `ψ'`.

```julia
Bra(ψ :: Vector{<:Number}; atol=ATOL :: Float64)
```

Similarly a `Ket` can be converted into a `Bra` via the adjoint.

```julia
Bra(ψ :: Ket; atol=ATOL :: Float64)
```

`Bra(ψ')`, throws a `DomainError` if `ψ'` is not normalized.
"""
struct Bra{T<:Number} <: AbstractMatrix{Number}
    ψ :: Adjoint{T,Vector{T}}
    atol :: Float64
    Bra(
        ψ :: Vector{T}; atol=ATOL :: Float64
    ) where T <: Number = is_braket(ψ, atol=atol) ? new{T}(ψ', atol) : _braket_error(ψ)
    Bra(
        ψ :: Adjoint{T,Vector{T}}; atol=ATOL :: Float64
    ) where T <: Number = is_braket(ψ, atol=atol) ? new{T}(ψ, atol) : _braket_error(ψ)
end
Base.size(bra::Bra) = size(bra.ψ)
Base.getindex(bra::Bra, I::Vararg{Int,2}) = getindex(bra.ψ, I...)
Base.setindex!(bra::Bra, val, I::Vararg{Int,2}) = (bra.ψ[I...] = val)
is_braket(::Bra) = true

# Bra-Ket conversions
Ket(bra :: Bra) = Ket(bra.ψ', atol=bra.atol)
Bra(ket :: Ket) = Bra(ket.ψ', atol=ket.atol)

"""
Outer product ``|\\psi \\rangle\\langle \\psi|``:

```julia
*(ket :: Ket, bra :: Bra) :: Matrix
```

Inner product ``\\langle \\psi | \\psi \\rangle``:

```julia
*(bra :: Bra, ket :: Ket) :: Number
```
"""
*(ket :: Ket, bra :: Bra) :: Matrix = ket.ψ * bra.ψ
*(bra :: Bra, ket :: Ket) :: Number = (bra.ψ * ket.ψ)[1]

"""
The adjoint is the complex conjugate transpose. This operation converts a `Ket` to
a `Bra`, ``\\langle \\psi | = |\\psi \\rangle^{\\dagger}``:

```julia
adjoint(ket :: Ket) :: Bra
```

Or, convert a `Bra` to a `Ket`, ``|\\psi\\rangle = \\langle \\psi |^{\\dagger}``:

```julia
adjoint(bra :: Bra) :: Ket
```
"""
adjoint(ket :: Ket) = Bra(ket)
adjoint(bra :: Bra) = Ket(bra)

"""
Performs the kronecker product of a set of `Ket`s (`Bra`s) to produce a new `Ket` (`Bra`).

```julia
kron( kets :: Vararg{Ket}; atol=ATOL ) :: Ket
kron( bras :: Vararg{Bra}; atol=ATOL ) :: Bra
```
"""
kron(kets :: Vararg{Ket}; atol=ATOL) = Ket(kron(map(ket -> ket.ψ, kets)...), atol=atol)
kron(bras :: Vararg{Bra}; atol=ATOL) = Bra(kron(map(bra -> bra.ψ', bras)...), atol=atol)


"""
    cSIL.jl

This is a Julia port of a MATLAB implementation of batch and incremental
Centroid-based Silhouette (cSIL) Cluster Validity Index.

Authors:
MATLAB implementation: Leonardo Enzo Brito da Silva
Julia port: Sasha Petrenko <sap625@mst.edu>

REFERENCES
[1] L. E. Brito da Silva, N. M. Melton, and D. C. Wunsch II, "Incremental
Cluster Validity Indices for Hard Partitions: Extensions  and  Comparative
Study," ArXiv  e-prints, Feb 2019, arXiv:1902.06711v1 [cs.LG].
[2] P. J. Rousseeuw, "Silhouettes: A graphical aid to the interpretation and
validation of cluster analysis," Journal of Computational and Applied
Mathematics, vol. 20, pp. 53–65, 1987.
[3] M. Rawashdeh and A. Ralescu, "Center-wise intra-inter silhouettes," in
Scalable Uncertainty Management, E. Hüllermeier, S. Link, T. Fober et al.,
Eds. Berlin, Heidelberg: Springer, 2012, pp. 406–419.
"""

"""
    cSIL

The stateful information of the Centroid-based Silhouette (cSIL) Cluster Validity Index.

# References
1. L. E. Brito da Silva, N. M. Melton, and D. C. Wunsch II, "Incremental Cluster Validity Indices for Hard Partitions: Extensions  and  Comparative Study," ArXiv  e-prints, Feb 2019, arXiv:1902.06711v1 [cs.LG].
2. P. J. Rousseeuw, "Silhouettes: A graphical aid to the interpretation and validation of cluster analysis," Journal of Computational and Applied Mathematics, vol. 20, pp. 53–65, 1987.
3. M. Rawashdeh and A. Ralescu, "Center-wise intra-inter silhouettes," in Scalable Uncertainty Management, E. Hüllermeier, S. Link, T. Fober et al., Eds. Berlin, Heidelberg: Springer, 2012, pp. 406–419.
"""
mutable struct cSIL <: CVI
    label_map::LabelMap
    dim::Integer
    n_samples::Integer
    n::IntegerVector        # dim
    v::RealMatrix           # dim x n_clusters
    CP::RealVector          # dim
    G::RealMatrix           # dim x n_clusters
    S::RealMatrix           # n_clusters x n_clusters
    sil_coefs::RealVector   # dim
    n_clusters::Integer
    criterion_value::Float
end # cSIL <: CVI

"""
    cSIL()

Default constructor for the Centroid-based Silhouette (cSIL) Cluster Validity Index.
"""
function cSIL()
    cSIL(
        LabelMap(),                     # label_map
        0,                              # dim
        0,                              # n_samples
        Array{Integer, 1}(undef, 0),    # n
        Array{Float, 2}(undef, 0, 0),   # v
        Array{Float, 1}(undef, 0),      # CP
        Array{Float, 2}(undef, 0, 0),   # G
        Array{Float, 2}(undef, 0, 0),   # S
        Array{Float, 1}(undef, 0),      # sil_coefs
        0,                              # n_clusters
        0.0                             # criterion_value
    )
end # cSIL()

"""
    setup!(cvi::cSIL, sample::Vector{T}) where {T<:RealFP}
"""
function setup!(cvi::cSIL, sample::Vector{T}) where {T<:RealFP}
    # Get the feature dimension
    cvi.dim = length(sample)
    # Initialize the augmenting 2-D arrays with the correct feature dimension
    # NOTE: R is emptied and calculated in evaluate!, so it is not defined here
    cvi.v = Array{T, 2}(undef, cvi.dim, 0)
    cvi.G = Array{T, 2}(undef, cvi.dim, 0)
end # setup!(cvi::cSIL, sample::Vector{T}) where {T<:RealFP}

function param_inc!(cvi::cSIL, sample::RealVector, label::Integer)
    # Get the internal label
    i_label = get_internal_label!(cvi.label_map, label)

    n_samples_new = cvi.n_samples + 1
    if cvi.n_samples == 0
        setup!(cvi, sample)
    end

    if i_label > cvi.n_clusters
        n_new = 1
        v_new = sample
        CP_new = transpose(sample) * sample
        G_new = sample
        # Compute S_new
        if cvi.n_clusters == 0
            # S_new = 0.0
            S_new = zeros(1,1)
        else
            S_new = zeros(cvi.n_clusters + 1, cvi.n_clusters + 1)
            S_new[1:cvi.n_clusters, 1:cvi.n_clusters] = cvi.S
            S_row_new = zeros(cvi.n_clusters + 1)
            S_col_new = zeros(cvi.n_clusters + 1)
            for cl = 1:cvi.n_clusters
                # Column "bmu_temp" - D_new
                C = CP_new + (transpose(cvi.v[:, cl]) * cvi.v[:, cl]) - 2 * (transpose(G_new) * cvi.v[:, cl])
                S_col_new[cl] = C
                # Row "bmu_temp" - E
                C = cvi.CP[cl] + cvi.n[cl] * (transpose(v_new)*v_new) - 2*(transpose(cvi.G[:, cl]) * v_new)
                S_row_new[cl] = C / cvi.n[cl]
            end
            # Column "ind_minus" - F
            S_col_new[i_label] = 0
            S_row_new[i_label] = S_col_new[i_label]
            S_new[:, i_label] = S_col_new
            S_new[i_label, :] = S_row_new
        end
        # Update 1-D parameters with a push
        cvi.n_clusters += 1
        push!(cvi.CP, CP_new)
        push!(cvi.n, n_new)
        # Update 2-D parameters with appending and reassignment
        cvi.v = [cvi.v v_new]
        cvi.G = [cvi.G G_new]
        cvi.S = S_new
    else
        n_new = cvi.n[i_label] + 1
        v_new = (1 - 1/n_new) .* cvi.v[:, i_label] + (1/n_new) .* sample
        CP_new = cvi.CP[i_label] + (transpose(sample) * sample)
        G_new = cvi.G[:, i_label] + sample
        # Compute S_new
        S_row_new = zeros(cvi.n_clusters)
        S_col_new = zeros(cvi.n_clusters)
        for cl = 1:cvi.n_clusters
            # Skip the i_label iteration
            if cl == i_label
                continue
            end
            # Column "bmu_temp" - D_new
            diff_x_v = sample - cvi.v[:, cl]
            C = cvi.CP[i_label] + (transpose(diff_x_v)*diff_x_v) + cvi.n[i_label]*(transpose(cvi.v[:, cl]) * cvi.v[:, cl]) - 2 * (transpose(G_new) * cvi.v[:, cl])
            S_col_new[cl] = C / n_new
            # Row "bmu_temp" - E
            C = cvi.CP[cl] + cvi.n[cl] * (transpose(v_new)*v_new) - 2*(transpose(cvi.G[:, cl]) * v_new)
            S_row_new[cl] = C / cvi.n[cl]
        end
        # Column "ind_minus" - F
        diff_x_v = sample - v_new
        C = cvi.CP[i_label] + (transpose(diff_x_v) * diff_x_v) + cvi.n[i_label] * (transpose(v_new) * v_new) - 2*(transpose(cvi.G[:, i_label])*v_new)
        S_col_new[i_label] = C / n_new
        S_row_new[i_label] = S_col_new[i_label]
        # Update parameters
        cvi.n[i_label] = n_new
        cvi.v[:, i_label] = v_new
        cvi.CP[i_label] = CP_new
        cvi.G[:, i_label] = G_new
        cvi.S[:, i_label] = S_col_new
        cvi.S[i_label, :] = S_row_new
    end
    cvi.n_samples = n_samples_new
end # param_inc!(cvi::cSIL, sample::RealVector, label::Integer)

function param_batch!(cvi::cSIL, data::RealMatrix, labels::IntegerVector)
    cvi.dim, cvi.n_samples = size(data)
    # u = findfirst.(isequal.(unique(labels)), [labels])
    u = unique(labels)
    cvi.n_clusters = length(u)
    cvi.n = zeros(Integer, cvi.n_clusters)
    cvi.v = zeros(cvi.dim, cvi.n_clusters)
    cvi.CP = zeros(cvi.n_clusters)
    cvi.S = zeros(cvi.n_clusters, cvi.n_clusters)
    D = zeros(cvi.n_samples, cvi.n_clusters)
    for ix = 1:cvi.n_clusters
        subset = data[:, findall(x->x==u[ix], labels)]
        cvi.n[ix] = size(subset, 2)
        cvi.v[:, ix] = mean(subset, dims=2)
        # Compute CP in case of switching back to incremental mode
        d_temp = (data - cvi.v[:, ix]*ones(1, cvi.n_samples)).^2
        D[:, ix] = transpose(sum(d_temp, dims=1))
    end
    for ix = 1:cvi.n_clusters
        for jx = 1:cvi.n_clusters
            subset_ind = findall(x->x==u[jx], labels)
            cvi.S[ix, jx] = sum(D[subset_ind, ix]) / cvi.n[jx]
        end
    end
end # param_batch!(cvi::cSIL, data::RealMatrix, labels::IntegerVector)

function evaluate!(cvi::cSIL)
    cvi.sil_coefs = zeros(cvi.n_clusters)
    if !isempty(cvi.S) && cvi.n_clusters > 1
        for ix = 1:cvi.n_clusters
            # Same cluster
            a = cvi.S[ix, ix]
            # Other clusters
            b = minimum(cvi.S[ix, 1:end .!= ix])
            cvi.sil_coefs[ix] = (b-a)/max(a, b)
        end
        # cSIL index value
        cvi.criterion_value = sum(cvi.sil_coefs) / cvi.n_clusters
    else
        cvi.criterion_value = 0.0
    end
end # evaluate(cvi::cSIL)


export chamfer_distance

"""
    chamfer_distance(A::PointCloud, B::PointCloud; w1=1.0, w2=1.0)

Computes the chamfer distance between PointCloud `A` and `B`.

`w1` and `w2` are optional arguments for specifying the weighted mean of between two
pcloud for computing metrics.
"""
chamfer_distance(A::PointCloud, B::PointCloud; w1::Number = 1.0, w2::Number = 1.0) =
    _chamfer_distance(A.points, B.points, Float32(w1), Float32(w2))

chamfer_distance(
    A::AbstractArray{<:Number},
    B::AbstractArray{<:Number};
    w1::Number = 1.0,
    w2::Number = 1.0,
) = _chamfer_distance(Float32.(A), Float32.(B), Float32(w1), Float32(w2))

chamfer_distance(
    A::AbstractArray{Float32},
    B::AbstractArray{Float32};
    w1::Number = 1.0,
    w2::Number = 1.0,
) = _chamfer_distance(A, B, Float32(w1), Float32(w2))

function _chamfer_distance(
    A::AbstractArray{Float32,2},
    B::AbstractArray{Float32,2},
    w1::Float32 = 1.0f0,
    w2::Float32 = 1.0f0,
)
    A = reshape(A, size(A)..., 1)
    B = reshape(B, size(B)..., 1)
    return _chamfer_distance(A, B, w1, w2)
end

function _chamfer_distance(
    A::AbstractArray{Float32,3},
    B::AbstractArray{Float32,3},
    w1::Float32 = 1.0f0,
    w2::Float32 = 1.0f0,
)
    nn_for_A, nn_for_B = @ignore _nearest_neighbors(A, B)

    dist_A_to_B = mean((A .- B[:, nn_for_A]) .^ 2) * 3.0f0
    dist_B_to_A = mean((B .- A[:, nn_for_B]) .^ 2) * 3.0f0

    distance = (w1 * dist_A_to_B) + (w2 * dist_B_to_A)
    return distance
end

function _nearest_neighbors(x::Array{Float32,3}, y::Array{Float32,3})
    nn_for_x = cat(
        [
            CartesianIndex.(reduce(vcat, knn(KDTree(y[:, :, i]), x[:, :, i], 1)[1]), i)
            for i = 1:size(x, 3)
        ]...,
        dims = 2,
    )
    nn_for_y = cat(
        [
            CartesianIndex.(reduce(vcat, knn(KDTree(x[:, :, i]), y[:, :, i], 1)[1]), i)
            for i = 1:size(x, 3)
        ]...,
        dims = 2,
    )
    return nn_for_x, nn_for_y
end

function _nearest_neighbors(x::CuArray{Float32,3}, y::CuArray{Float32,3})
    xx = sum(x .^ 2, dims = 1)
    yy = sum(y .^ 2, dims = 1)
    zz = Flux.batched_mul(permutedims(x, (2, 1, 3)), y)
    rx = reshape(xx, size(xx, 2), 1, :)
    ry = reshape(yy, 1, size(yy, 2), :)
    P = (rx .+ ry) .- (2 .* zz)

    nn_for_x = argmin(P; dims = 2) |> Array
    nn_for_y = argmin(P; dims = 1) |> Array

    nn_for_x = reshape(map(x -> CartesianIndex(x[2], x[3]), nn_for_x), :, size(x, 3))
    nn_for_y = reshape(map(y -> CartesianIndex(y[1], y[3]), nn_for_y), :, size(y, 3))
    return nn_for_x, nn_for_y
end


module AutomatonTests

using Test
using WeaklyHard
using WeaklyHard: nbits, shift, shift_rowhit, _childexists, consword

##########################
#### Common Variables ####
##########################

k_max           = 20
n_constraints   = 6
n_tests         = 1_000
types           = [Int8, Int16, Int32, Int64, BigInt]

#####################
#### Shift tests ####
#####################

function test_shift_correctness(w::T, c) where {T <: Integer}
    newWord = (w << 1) | T(c)
    if c == M
        return newWord == shift(w, c)
    else
        return newWord - T(2^(nbits(newWord)-1)) == shift(w, c)
    end # if
end # function

function test_shift_rowhit_correctness(w::T, c, n::T) where {T <: Integer}
    newWord = (w << 1) | T(c)
    if c == M
        return newWord == shift_rowhit(w, c, n)
    else
        return ((newWord & n == n) ? n : newWord) == shift_rowhit(w, c, n)
    end # if
end # function

@testset "Shift Tests" begin
    chars = [H, M]
    for T in types
        for _ in 1:n_tests
            @test test_shift_correctness(rand(T(1):T(8*sizeof(T)-1)), rand(chars))
        end # for
    end # for
end # testset

@testset "Shift_RowHit Tests" begin
    chars = [H, M]
    for T in types
        n = consword( rand( T(1) : max( T(sizeof(T) - 1), T(1) ) ) )
        for _ in 1:n_tests
            @test test_shift_rowhit_correctness(rand(T(1):T(8*sizeof(T)-1)), rand(chars), n)
        end # for
    end # for
end # testset

####################
### Vertex tests ###
####################

function test_vertex_correctness(Lambda::Set{T}) where {T <: Constraint}
    if length(Lambda) == 1
        L = pop!(Lambda)
    else
        L = dominant_set(Lambda)
    end
    automaton = build_automaton(L)

    seq = random_sequence(automaton, 5000)
    c = is_satisfied(L, seq)
    if !c
        return false
    end

    for n in vertices(automaton)
        if (!is_satisfied(L, Int64(n.w))
                || (_childexists(n, :miss) && !is_satisfied(L, Int64(n.miss.w)))
                || !is_satisfied(L, Int64(n.hit.w)))
            return false
        end
    end

    true
end # function

@testset "Vertices Tests" begin
    constraint_names = ["RowMissConstraint", "RowHitConstraint", "AnyMissConstraint", "AnyHitConstraint"]

    tested = Set{Set{Constraint}}()
    for _ in 1:n_tests

        Lambda = Set{Constraint}()
        nbr_c  = rand(1:n_constraints)
        for _ in 1:nbr_c
            k = rand(1:k_max)
            x = rand(1:max(k-1, 1))
            lambda_type = rand(constraint_names)
            if lambda_type == "RowHitConstraint"
                x = x #floor(Int, x/2)
            end
            push!(Lambda, getfield(WeaklyHard, Symbol(lambda_type))(x, k))
        end # for

        L = dominant_set(Lambda)

        if !(L in tested)
            push!(tested, L)

            c = test_vertex_correctness(Lambda)
            @test c
            if !c

                @show "###############"
                @show "Lambda:"
                for l in Lambda
                    @show l
                end
                @show "---------------"
                @show "Lambda_star:"
                Lambda_star = dominant_set(Lambda)
                for l in Lambda_star
                    @show l
                end
                @show "###############"

                return
            end
        end
    end # for

    println(" ")
    println("Executed tests on $(length(tested)) different automata")
end # testset

############################
##### Transitions tests ####
############################

function test_transition_correctness(Lambda::Set{T}) where {T <: Constraint}

    if length(Lambda) == 1
        L = pop!(Lambda)
    else
        L = Lambda
    end

    automaton = build_automaton(L)
    trans     = transitions(automaton)

    for v in vertices(automaton)
        v1 = v.hit
        if !((v.w, v1.w, H) in trans)
            return false
        end

        if _childexists(v, :miss)
            v2 = v.miss
            if !((v.w, v2.w, M) in trans)
                return false
            end
        end
    end # for

    return true
end # function

@testset "Transitions Tests" begin
    constraint_names = ["RowMissConstraint", "RowHitConstraint", "AnyMissConstraint", "AnyHitConstraint"]

    tested = Set{Set{Constraint}}()
    for _ in 1:n_tests

        Lambda = Set{Constraint}()
        nbr_c  = rand(1:n_constraints)
        for _ in 1:nbr_c
            k = rand(1:k_max)
            x = rand(1:max(k-1, 1))
            lambda_type = rand(constraint_names)
            if lambda_type == "RowHitConstraint"
                x = x #floor(Int, x/2)
            end
            push!(Lambda, getfield(WeaklyHard, Symbol(lambda_type))(x, k))
        end # for

        L = dominant_set(Lambda)

        if !(L in tested)
            push!(tested, L)

            c = test_transition_correctness(Lambda)
            @test c
            if !c

                @show "###############"
                @show "Lambda:"
                @show Lambda
                @show "---------------"
                @show "Lambda_star:"
                Lambda_star = dominant_set(Lambda)
                @show Lambda_star
                @show "###############"

                return
            end
        end
    end # for

    println(" ")
    println("Executed tests on $(length(tested)) different automata")
end # testset

@testset "Set-construction Tests" begin
    constraint_names = ["RowMissConstraint", "RowHitConstraint", "AnyMissConstraint", "AnyHitConstraint"]
    n_constraints    = 3
    k_max            = 15
    n_tests          = 1_000

    for _ in 1:n_tests
        Lambda = Set{Constraint}()
        n = rand(1:n_constraints)
        for _ in 1:n
            k = rand(1:k_max)
            x = rand(0:k)
            lambda_type = rand(constraint_names)
            if lambda_type == "RowHitConstraint"
                x = floor(Int, x/2)
            end
            push!(Lambda, getfield(WeaklyHard, Symbol(lambda_type))(x, k))
        end # for

        Lambda_star = dominant_set(Lambda)

        G = build_automaton(Lambda_star)
        minimize_automaton!(G)
        all_seq = all_sequences(G, k_max+1) # +1 to account for RowMissConstraints
        for seq in all_seq
            if !is_satisfied(Lambda_star, seq)
                @show Lambda
                @show Lambda_star
                @show seq
                @show bitstring(seq)
            end
            @test is_satisfied(Lambda_star, seq)
        end # for
    end # for
end # testset

end # module


# Rohan McLure, Australian National University (2019)
# Distributed implementation of the (Star) Stencil Parallel Research Kernel
using Distributed

function main()
    if length(ARGS) != 3 || nprocs() == 1
        println("Usage: julia -p <# workers> stencil.jl <# iterations> <# order> <# radius>")
        exit(1)
    end

    argv = map(x -> parse(Int, x), ARGS)
    iterations, order, r = argv
    ranks = nworkers()
    procs_x, procs_y = factor(ranks)
    out_channels = Matrix{Vector{Union{RemoteChannel,Nothing}}}(undef, procs_y, procs_x)
    dims = Matrix{Tuple{Int,Int}}(undef, procs_y, procs_x)
    std_width = Int(ceil(order / procs_x))
    std_height = Int(ceil(order / procs_y))
    last_width = order % std_width == 0 ? std_width : order % std_width
    last_height = order % std_height == 0 ? std_height : order % std_height
    for i in 1:procs_y
        for j in 1:procs_x
            v = Vector{Union{RemoteChannel, Nothing}}(undef, 4)
            v[1] = i != 1 ? RemoteChannel(() -> Channel{Matrix{Float64}}(0), (i-2)*procs_x+j) : nothing
            v[2] = j != procs_x ? RemoteChannel(() -> Channel{Matrix{Float64}}(0), (i-1)*procs_x+j+1) : nothing
            v[3] = i != procs_y ? RemoteChannel(() -> Channel{Matrix{Float64}}(0), i*procs_x+j) : nothing
            v[4] = j != 1 ? RemoteChannel(() -> Channel{Matrix{Float64}}(0), (i-1)*procs_x+j-1) : nothing
            out_channels[i,j] = v
            # Assigning coordinates
            height = i == procs_y ? last_height : std_height
            width = j == procs_x ? last_width : std_width
            dims[i,j] = (height, width)
        end
    end
    in_channels = Matrix{Vector{Union{RemoteChannel,Nothing}}}(undef, procs_y, procs_x)
    for i in 1:procs_y
        for j in 1:procs_x
            v = Vector{Union{RemoteChannel, Nothing}}(undef, 4)
            v[1] = i != 1 ? out_channels[i-1,j][3] : nothing
            v[2] = j != procs_x ? out_channels[i,j+1][4] : nothing
            v[3] = i != procs_y ? out_channels[i+1,j][1] : nothing
            v[4] = j != 1 ? out_channels[i,j-1][2] : nothing
            in_channels[i,j] = v
        end
    end

    futures = Matrix{Future}(undef, procs_y, procs_x)
    @sync for i in 1:procs_y
        for j in 1:procs_x
            wid = (i-1)*procs_x+j
            id = workers()[wid]
            futures[i,j] = @spawnat id do_stencil(iterations, order, dims[i,j], (i,j), (procs_y, procs_x), r, out_channels[i,j], in_channels[i,j])
        end
    end

    results = map(fetch, futures)

    active_points = (order - 2*r)^2
    norm = sum(map(x->x[1], results))
    norm /= active_points
    time = maximum(map(x->x[2], results))
    reference_norm = 2 * (iterations + 1)
    if abs(norm - reference_norm) > 1.e-8
        println("ERROR: L1 norm = $norm Reference L1 norm = $reference_norm")
    else
        println("Solution validates")
    end
    stencil_size = 4*r+1
    flops = (2*stencil_size+1) * active_points
    avgtime = time/iterations
    println("Rate (MFlops/s): ",1.e-6*flops/avgtime, " Avg time (s): ",avgtime)
end

@everywhere function do_stencil(
    iterations::Int,
    order::Int,
    local_dims::Tuple{Int,Int},
    proc_dims::Tuple{Int,Int},
    proc_setup::Tuple{Int,Int},
    r::Int,
    out_channels::Vector{Union{RemoteChannel,Nothing}},
    in_channels::Vector{Union{RemoteChannel,Nothing}}
)
    ub, rb, db, lb = out_channels .!== nothing
    h, w = local_dims
    id_y, id_x = proc_dims
    procs_y, procs_x = proc_setup
    std_height, std_width = Int(ceil(order/procs_y)), Int(ceil(order/procs_x))
    i_start, j_start = (id_y-1)*std_height, (id_x-1)*std_width
    i_end, j_end = i_start + h - 1, j_start + w - 1
    i_o_start, j_o_start = (ub ? 1 : r+1), (lb ? 1 : r+1)
    i_o_end, j_o_end = (db ? h : h-r), (rb ? w : w-r)
    A = zeros(Float64, (local_dims .+ 2*r)...)
    for i in 1:h
        for j in 1:w
            A[i+r,j+r] = (i_start + i - 1) + (j_start + j - 1)
        end
    end
    B = zeros(Float64, local_dims...)
    W = zeros(Float64, 2*r+1, 2*r+1)
    for i in 1:r
        W[r+1, r+i+1] = 1.0/(2*i*r)
        W[r+i+1, r+1] = 1.0/(2*i*r)
        W[r+1, r-i+1] = -1.0/(2*i*r)
        W[r-i+1, r+1] = -1.0/(2*i*r)
    end

    u_out, r_out, d_out, l_out = [Matrix{Float64}(undef, (i % 2 == 0 ? (h,r) : (r,w))...) for i in 1:4]
    u_out_v, r_out_v, d_out_v, l_out_v = [
        view(A, r+1:2*r, r+1:r+w),
        view(A, r+1:r+h, w+1:r+w),
        view(A, h+1:r+h, r+1:r+w),
        view(A, r+1:r+h, r+1:2*r)
    ]
    u_in_v, r_in_v, d_in_v, l_in_v = [
        view(A, 1:r, r+1:r+w),
        view(A, r+1:r+h, r+w+1:2*r+w),
        view(A, r+h+1:2*r+h, r+1:r+w),
        view(A, r+1:r+h, 1:r)
    ]

    # Channels
    u_out_c, r_out_c, d_out_c, l_out_c = out_channels
    u_in_c, r_in_c, d_in_c, l_in_c = in_channels

    local t0, t1

    for k in 0 : iterations
        if k == 1
            t0 = time_ns()
        end
        # Send
        @sync begin
            if ub
                @async begin
                    copy!(u_out, u_out_v)
                    put!(u_out_c, u_out)
                end
                @async begin
                    u_in = take!(u_in_c)
                    copy!(u_in_v, u_in)
                end
            end
            if rb
                @async begin
                    copy!(r_out, r_out_v)
                    put!(r_out_c, r_out)
                end
                @async begin
                    r_in = take!(r_in_c)
                    copy!(r_in_v, r_in)
                end
            end
            if db
                @async begin
                    copy!(d_out, d_out_v)
                    put!(d_out_c, d_out)
                end
                @async begin
                    d_in = take!(d_in_c)
                    copy!(d_in_v, d_in)
                end
            end
            if lb
                @async begin
                    copy!(l_out, l_out_v)
                    put!(l_out_c, l_out)
                end
                @async begin
                    l_in = take!(l_in_c)
                    copy!(l_in_v, l_in)
                end
            end
        end

        for j in j_o_start : j_o_end
            for i in i_o_start : i_o_end
                tmp = 0.0
                for ii in -r:r
                    @inbounds tmp += W[r+ii+1, r+1] * A[i+ii+r, j+r]
                end
                for jj in -r:-1
                    @inbounds tmp += W[r+1, r+jj+1] * A[i+r, j+jj+r]
                end
                for jj in 1:r
                    @inbounds tmp += W[r+1, r+jj+1] * A[i+r, j+jj+r]
                end
                @inbounds B[i,j] += tmp
            end
        end

        for j in 1:w
            for i in 1:h
                @inbounds A[i+r, j+r] += 1.0
            end
        end
    end

    t1 = time_ns()

    local_norm = 0.0
    for i in i_o_start : i_o_end
        for j in j_o_start : j_o_end
            local_norm += abs(B[i,j])
        end
    end

    return (local_norm, (t1-t0)*1.e-9)
end

function factor(x::Int)
    s = Int(floor(sqrt(x)))
    for p in s:-1:1
        if x % p == 0
            return (p, x ÷ p)
        end
    end
end

main()


"""
    mirror(N::T) where {T <: AbstractUnipartiteNetwork}

Returns a mirrored version of the adjacency matrix, required for spectral
radius.
"""
function mirror(N::T) where {T <: AbstractUnipartiteNetwork}
    M = copy(N)
    for i in N
        M[i.to, i.from] = M[i.from, i.to]
    end
    return M
end

ρ_phillips(N::T, v::Float64) where {T <: UnipartiteNetwork}  = v/((links(N)*(richness(N)-1))/richness(N))^0.5
ρ_ska(N::T, v::Float64) where {T <: UnipartiteNetwork} = v/sqrt(links(N))
ρ_raw(N::T, v::Float64) where {T <: UnipartiteNetwork} = v

"""
    ρ(N::T; range=EcologicalNetworks.ρ_ska) where {T <: UnipartiteNetwork}

Returns the spectral radius (the absolute value of the largest real part of the
eigenvalues of the adjacency matrix) of any unipartite network whose
interactions are positive or null. The spectral radius is measured on the
mirrored version of the matrix, so that an interaction from i to j also implies
an interaction from j to i.

Note that the spectral radius has been suggested as a measure of nestedness by
Staniczenko *et al.* (2013). Phillips (2011) uses it as a measure of the ability
of a system to dampen or absorb perturbations.

#### Maximal values

The spectral radius is sensitive to network size, and to a certain extent to the
number of links. The keyword argument `range` will divide return a ranged
version of the spectral radius, so that it is expressed relatively to its
maximal value. The `range` argument takes a *function*, which requires two
arguments: the network (which must be unipartite), and the value of the spectral
radius.

Options that come with `EcologicalNetworks.jl` (where L is the number of links
and S the number of nodes) are:

1. `EcologicalNetworks.ρ_phillips`: divides by the square root of (2L(S-1))/S,
   as in Phillips (20110)
1. `EcologicalNetworks.ρ_ska`: divides by the square root of L, as in
   Staniczenko *et al.* (2013) - this is the **default**
1. `EcologicalNetworks.ρ_raw`: returns the raw value

#### References

- Phillips, J.D., 2011. The structure of ecological state transitions:
  Amplification, synchronization, and constraints in responses to environmental
  change. Ecological Complexity, 8, 336–346.
  https://doi.org/10.1016/j.ecocom.2011.07.004

- Staniczenko, P.P.A., Kopp, J.C., Allesina, S., 2013. The ghost of nestedness
  in ecological networks. Nat Commun 4, 1391. https://doi.org/10.1038/ncomms2422
"""
function ρ(N::T; range=EcologicalNetworks.ρ_ska) where {T <: UnipartiteNetwork}
    links(N) == 0 && return NaN
    M = mirror(N)
    @assert minimum(M.edges) ≥ zero(_interaction_type(M))
    absolute_real_part = abs.(real.(eigvals(Array(M.edges))))
    return range(M, maximum(absolute_real_part))
end

"""
    ρ(N::T; varargs...) where {T <: BipartiteNetwork}

Bipartite version of the spectral radius. In practice, this casts the network
into its unipartite representation, since the spectral radius only makes sense
for square matrices.
"""
function ρ(N::T; varargs...) where {T <: BipartiteNetwork}
    return ρ(convert(UnipartiteNetwork, N); varargs...)
end


@testset "Solver $(solver)" for solver in [RAFF.lmlovo, RAFF.gnlslovo]

    model = (x, θ) -> θ[1] * x[1] + θ[2]
    gmodel! = (g, x, θ) -> begin
        g[1] = x[1]
        g[2] = 1.0
    end

    @testset "Simple cases ($(np), $p)" for (np, p) in [(10, 9), (100, 90)]

        data = Array{Float64, 2}(undef, np, 2)
        n = 2
        θsol = [-p, 0.5 * np]

        for (i, x) in enumerate(range(-10, stop=10, length=np))
            data[i, 1] = x
            data[i, 2] = model(x, θsol)
        end

        θ = [1.0, 1.0]

        data[1:np - p, 2] .+= 100

        r = solver(model, gmodel!, θ, data, n, p)

        @test(r.outliers == [1:np - p;])
        @test(r.solution ≈ θsol, atol=1.0e-4)
        @test(r.status == 1)
        @test(r.p == p)
        @test(r.nf >= r.iter)
        @test(r.nj >= r.iter)

        θ = [1.0, 1.0]

        r = solver(model, θ, data, n, p)

        @test(r.outliers == [1:np - p;])
        @test(r.solution ≈ θsol, atol=1.0e-4)
        @test(r.status == 1)
        @test(r.p == p)
        @test(r.nf >= r.iter)
        @test(r.nj >= r.iter)

    end

    @testset "Wrong arguments" begin

        np = 10
        p  = 9

        data = Array{Float64, 2}(undef, np, 2)
        n = 2
        θsol = [-p, 0.5 * np]

        for (i, x) in enumerate(range(-10, stop=10, length=np))
            data[i, 1] = x
            data[i, 2] = model(x, θsol)
        end

        θ = [1.0, 1.0]

        data[1:np - p, 2] .+= 100

        r = solver(model, gmodel!, θ, data, n, p)

        @test_throws(AssertionError, solver(model, θ, data, 0, 1))
        @test_throws(AssertionError, solver(model, θ, data, 2, -1))

    end

    # Test to check Issue #1

    @testset "Error in printing" begin

        m(x, θ) = θ[1] * x[1]^2 + θ[2]

        A = [ -2.0  5.00;
              -1.5  3.25;
              -1.0  2.00;
              -0.5  1.25;
              0.0  1.00;
              0.5  1.25;
              1.0  2.00;
              1.5  3.25;
              2.0  5.00 ]

        θ = [0.0, 0.0]

        # Changes log just for this test
        rout = with_logger(Logging.NullLogger()) do

            lmlovo(m, θ, A, 2, 4)

        end

        @test rout.status == 1
        @test rout.p == 4

    end

end


immutable GammaRmathSampler <: Sampleable{Univariate,Continuous}
    α::Float64
    scale::Float64
end

rand(s::GammaRmathSampler) =
    ccall((:rgamma, "libRmath-julia"), Float64, (Float64, Float64), s.α, s.scale)

#  A simple method for generating gamma variables - Marsaglia and Tsang (2000)
#  http://www.cparity.com/projects/AcmClassification/samples/358414.pdf
#  Page 369
#  basic simulation loop for pre-computed d and c
#

immutable GammaMTSampler <: Sampleable{Univariate,Continuous}
    iα::Float64
    d::Float64
    c::Float64
    κ::Float64
end

function GammaMTSampler(α::Float64, scale::Float64)
    if α >= 1.0
        iα = 1.0
        d = α - 1/3
    else
        iα = 1.0 / α
        d = α + 2/3
    end
    c = 1.0 / sqrt(9.0 * d)
    κ = d * scale
    GammaMTSampler(iα, d, c, κ)
end

GammaMTSampler(α::Float64) = GammaMTSampler(α, 1.0)

function rand(s::GammaMTSampler)
    d = s.d
    c = s.c
    iα = s.iα

    v = 0.0
    while true
        x = randn()
        v = 1.0 + c * x
        while v <= 0.0
            x = randn()
            v = 1.0 + c * x
        end
        v *= (v * v)
        u = rand()
        x2 = x * x
        if u < 1.0 - 0.331 * abs2(x2)
            break
        end
        if log(u) < 0.5 * x2 + d * (1.0 - v + log(v))
            break
        end
    end
    v *= s.κ
    if iα > 1.0
        v *= (rand() ^ iα)
    end
    return v
end


# * Norms

_norm(B::FiniteDifferencesOrRestricted, c::AbstractArray, p::Real=2) =
    norm(c, p)*(step(B)^(inv(p)))

LinearAlgebra.norm(v::FDVecOrMat, p::Real=2) = _norm(v.args..., p)
LinearAlgebra.norm(v::Mul{<:Any, <:Tuple{<:FiniteDifferencesOrRestricted, <:AbstractArray}},
                   p::Real=2) = _norm(v.args..., p)

function LinearAlgebra.normalize!(v::FDVecOrMat, p::Real=2)
    v.args[2][:] /= norm(v,p)
    v
end

function LinearAlgebra.normalize!(v::Mul{<:Any, <:Tuple{<:FiniteDifferencesOrRestricted, <:AbstractArray}},
                                  p::Real=2)
    v.args[2][:] /= norm(v,p)
    v
end

# * Inner products

function _inner_product(a::Adjoint{<:Any,<:AbstractVector}, A::QuasiAdjoint{<:Any,<:FD},
                        B::FD, b::AbstractVector) where {FD<:FiniteDifferencesOrRestricted}
    axes(A.parent) == axes(B) || throw(ArgumentError("Incompatible axes"))
    a*b*step(B)
end

LazyArrays.materialize(inner_product::FDInnerProduct{T,U,FD}) where {T,U,FD<:FiniteDifferencesOrRestricted{U}} =
    _inner_product(inner_product.args...)

function LazyArrays.materialize(inner_product::LazyFDInnerProduct{FD}) where {FD<:FiniteDifferencesOrRestricted}
    aA,Bb = inner_product.args
    _inner_product(aA.args..., Bb.args...)
end

function LazyArrays.materialize(s::Mul{<:Any, <:Tuple{
    <:Mul{<:Any, <:Tuple{
        <:Adjoint{<:Any,<:AbstractVector},
        <:ContinuumArrays.QuasiArrays.QuasiAdjoint{<:Any, <:FD}}},
    <:Mul{<:Any, <:Tuple{
        <:FD,
        <:Diagonal,
        <:ContinuumArrays.QuasiArrays.QuasiAdjoint{<:Any, <:FD}}},
    <:Mul{<:Any, <:Tuple{
        <:FD,
        <:AbstractVector}}}}) where {FD<:FiniteDifferencesOrRestricted}
    a,o,b = s.args
    axes(a.args[2].parent) == axes(o.args[1]) &&
        axes(o.args[3].parent) == axes(b.args[1]) ||
        throw(ArgumentError("Incompatible axes"))
    av = first(a.args)
    ov = o.args[2].diag
    bv = last(b.args)
    v = zero(promote_type(eltype(av),eltype(ov),eltype(bv)))
    @inbounds for i in eachindex(bv)
        v += av[i]*ov[i]*bv[i]
    end
    v*step(first(b.args))
end


using Test, Distributions
@testset "Normal" begin
    @test isa(convert(Normal{Float64}, Float16(0), Float16(1)),
              Normal{Float64})
    @test Inf === logpdf(Normal(0, 0), 0)
    @test -Inf === logpdf(Normal(), Inf)
    @test iszero(logcdf(Normal(0, 0), 0))
    @test iszero(logcdf(Normal(), Inf))
    @test -Inf === logccdf(Normal(0, 0), 0)
    @test iszero(logccdf(Normal(eps(), 0), 0))
    @test -Inf === quantile(Normal(), 0)
    @test iszero(quantile(Normal(), 0.5))
    @test Inf === quantile(Normal(), 1)
    @test -Inf === quantile(Normal(0, 0), 0)
    @test iszero(quantile(Normal(0, 0), 0.75))
    @test Inf === quantile(Normal(0, 0), 1)
    @test -Inf === quantile(Normal(0.25, 0), 0)
    @test 0.25 == quantile(Normal(0.25, 0), 0.95)
    @test Inf === quantile(Normal(0.25, 0), 1)
    @test Inf === cquantile(Normal(), 0)
    @test iszero(cquantile(Normal(), 0.5))
    @test -Inf === cquantile(Normal(), 1)
    @test Inf === cquantile(Normal(0, 0), 0)
    @test iszero(cquantile(Normal(0, 0), 0.75))
    @test -Inf === cquantile(Normal(0, 0), 1)
    @test Inf === cquantile(Normal(0.25, 0), 0)
    @test 0.25 == cquantile(Normal(0.25, 0), 0.95)
    @test -Inf === cquantile(Normal(0.25, 0), 1)
    @test -Inf === invlogcdf(Normal(), -Inf)
    @test isnan(invlogcdf(Normal(), NaN))
    @test Inf === invlogccdf(Normal(), -Inf)
    @test isnan(invlogccdf(Normal(), NaN))
end


include("znccUtils.jl")

#=
	2D IMPLEMENTATION
=#

#  Lewis, J.P.: Fast normalized cross-correlation. Ind. Light Magic10(2001)
function crossCorrelation!( ::ZNCC, ccr::A{T,2}, num::A{T,2},
                                    padF::A{C{T},2}, padG::A{C{T},2}, plan, iplan,
                                    sumF::A{U,2}, sumF2::A{U,2}, meanF, meanF2,
                                    stdG, sizeF, sizeG, shifts
                          ) where {T<:Real,U<:Real}

    # Computing numerator with FFT.
    crossCorrelation!( FFT(), padF, padG, ccr, plan, iplan );
    Base.circshift!( num, ccr, shifts )

    # Computing denominator through summed-area tables.
    N = prod(sizeG);
    N2F, stdI = N*meanF2, stdG*sqrt(N)
    @inbounds for col in 1:size(ccr,2)
        for row in 1:size(ccr,1)
            F, F2 = znccSumValues( sumF, sumF2, row, col, sizeF, sizeG )
            ccr[ row, col ] = num[ row, col ]/( sqrt( F2 + N2F - 2*F*meanF )*stdG );
        end
    end
end

function PIV_2D( ::ZNCC, img1::A{T,2}, img2::A{T,2},
                         IA::III, SM::III, overlap::III, mpass::I, width::I,
                         peak::S, sigNoise::S, filtFun::Function, threshold::F,
                         mpFun::Function, reps::Array{I,1};
                         corrType=Float32
			   ) where {T<:Real}

	# Calculating vector field size.
    h, w   = size( img1 );
    step   = IA .- overlap;
    VFsize = 1 .+ floor.( Int64, ((h,w,0).-IA)./step );
    U  = zeros( corrType, VFsize[1:2] );
    V  = zeros( corrType, VFsize[1:2] );
    SN = zeros( corrType, VFsize[1:2] );

	ignoreSN = sigNoise == "";

    @inbounds for mp in mpass:-1:1

        # Scaling IA, SM and step (overlap) to multi-pass iteration.
        IA_mp   = mpFun.(  IA , mp );
        SM_mp   = mpFun.(  SM , mp );
        step_mp = mpFun.( step, mp );
        last_mp = ( mp == 1 );

        ssize   = IA_mp .+ 2 .* SM_mp
        search  = zeros( T, ssize[1:2] );
        sumS    = zeros( T, ssize[1:2] .+ 1 );
        sumS2   = zeros( T, ssize[1:2] .+ 1 );

        csize   = 2 .* ( IA_mp .+ SM_mp  );
        cmatrix = zeros( corrType, csize[1:2] );
        num     = zeros( corrType, csize[1:2] );
		shifts  = div.( csize[1:2], 2 ) .+ SM_mp[1:2] .- 1;
        pads    = zeros( Complex{corrType}, csize[1:2] );
        padi    = zeros( Complex{corrType}, csize[1:2] );
        plan    =  FFTW.plan_fft!( padi );
        iplan   = FFTW.plan_ifft!( padi );

        # Constructing grid coordinates for each interrogation volume,
        grid = constructGrid( (h, w, IA_mp[3]), IA_mp, step_mp, mp )

        for g in grid
            #  g[1],  g[2] are start and stop rows of interrogation area
            #  g[3],  g[4] are start and stop columns of interrogation area
            #  g[7],  g[8] are start and stop rows inside UV arrays
            #  g[9], g[10] are start and stop columns inside UV arrays

            interr = view( img1, g[1]:g[2], g[3]:g[4] )

            if threshold > 0 && filtFun( interr ) < threshold
                continue;
            end

            # 1-. Shifting the searchArea ranges by previously calculated displacements
            offH = round(Int64, U[ g[7], g[9] ]);
            offW = round(Int64, V[ g[7], g[9] ]);

            # 1.2-. Search volume coordinates after shifting.
            scoords, so = searchGrid( g, SM, offH, offW, 0, h, w, 0 )

            setTo0!( search )
            putWithinSearch!( search, img2, so, scoords )

            setTo0!( sumS  );
            setTo0!( sumS2 );
            znccSumArrays!( search, sumS, sumS2 );

            meanI = Statistics.mean( interr );
            meanS = Statistics.mean( search );

            setTo0!( pads )
            setTo0!( padi )
            putWithinPadded!( padi, interr, meanI );
            putWithinPadded!( pads, search, meanS );

            meanS2 = meanS*meanS;
            meanI2 = meanI*meanI;
            stdI   = Statistics.std( interr, mean=meanI );

            # 2-. Cross-Correlation
            crossCorrelation!( ZNCC(), cmatrix, num, pads, padi, plan, iplan,
                               sumS, sumS2, meanS, meanS2, stdI, ssize, IA_mp, shifts );

            # 3-. Calculation of displacement
            ( r, c ) = approxTranslation( cmatrix, peak, Val(last_mp) )
			if ( !ignoreSN && last_mp )
            	SN[ g[7], g[9] ] = SNRatio( cmatrix, sigNoise, width )
			end

            # 4-. Updating U, V matrices
            t1 =  g[7]:g[8];
            t2 =  g[9]:g[10];
            U[ t1, t2 ] .= offH - r;
            V[ t1, t2 ] .= offW - c;
        end
    end # Multi-pass loop

    return U, V, SN
end

#=
	3D IMPLEMENTATION
=#

function crossCorrelation!( ::ZNCC, ccr::A{T,3}, num::A{T,3},
                                    padF::A{C{T},3}, padG::A{C{T},3}, plan, iplan,
                                    sumF::A{U,3}, sumF2::A{U,3}, meanF, meanF2,
                                    stdG, sizeF, sizeG, shifts
						  ) where {T<:Real, U<:Real}

    # Computing numerator with FFT.
    crossCorrelation!( FFT(), padF, padG, ccr, plan, iplan );
    Base.circshift!( num, ccr, shifts )

    # Computing denominator through summed-area tables.
    N = prod(sizeG);
    N2F, stdG  = N*meanF2, stdG*sqrt(N);
    @inbounds for z in 1:size(ccr,3)
        for c in 1:size(ccr,2)
            for r in 1:size(ccr,1)
                F, F2 = znccSumValues( sumF, sumF2, r, c, z, sizeF, sizeG )
                ccr[r,c,z] = num[r,c,z]/( sqrt( F2 + N2F - 2*F*meanF )*stdG );
            end
        end
    end
end

function PIV_3D( ::ZNCC, vol1::A{T,3}, vol2::A{T,3},
                         IA::III, SM::III, overlap::III, mpass::I, width::I,
                         peak::S, sigNoise::S, filtFun::Function, threshold::F,
                         mpFun::Function, reps::Array{I,1};
                         corrType=Float32
			   ) where {T<:Real}

	# Calculating vector field size.
    h, w, d = size( vol1 );
    step   = IA .- overlap;
    VFsize = 1 .+ floor.(Int64,(size(vol1).-IA)./step);
    U  = zeros( corrType, VFsize );
    V  = zeros( corrType, VFsize );
    W  = zeros( corrType, VFsize );
    SN = zeros( corrType, VFsize );

	ignoreSN = sigNoise == "";

    @inbounds for mp in mpass:-1:1

        # Scaling IA, SM and step (overlap) to multi-pass iteration.
        IA_mp   = mpFun.(  IA , mp );
        SM_mp   = mpFun.(  SM , mp );
        step_mp = mpFun.( step, mp );
        last_mp = ( mp == 1 );

        ssize   = IA_mp .+ 2 .* SM_mp
        search  = zeros( T, ssize );
        sumS    = zeros( T, ssize .+ 1 );
        sumS2   = zeros( T, ssize .+ 1 );

        # Initialize cross-correlation variables once per multi-pass iteration
        csize   = 2 .* ( IA_mp .+ SM_mp  );
        cmatrix = zeros( corrType, csize );
        num     = zeros( corrType, csize );
        shifts  = div.(csize,2) .+ SM_mp .- 1;
        pads    = zeros( Complex{corrType}, csize );
        padi    = zeros( Complex{corrType}, csize );
        plan    =  FFTW.plan_fft!( pads );
        iplan   = FFTW.plan_ifft!( pads );

        # Constructing grid coordinates for each interrogation volume.
        grid = constructGrid( size(vol1), IA_mp, step_mp, mp )

        for g in grid
            #  g[1],  g[2] are start and stop rows of interrogation volume
            #  g[3],  g[4] are start and stop columns of interrogation volume
            #  g[5],  g[6] are start and stop depths of interrogation volume
            #  g[7],  g[8] are start and stop rows inside UVW arrays
            #  g[9], g[10] are start and stop columns inside UVW arrays
            # g[11], g[12] are start and stop depths insie UVW arrays

            interr = view( vol1, g[1]:g[2], g[3]:g[4], g[5]:g[6] )

            if threshold > 0 && filtFun( interr ) < threshold
                continue;
            end

            # 1-. Shifting the searchArea ranges by previously calculated displacements
            offH = round(Int64, U[ g[7], g[9], g[11] ]);
            offW = round(Int64, V[ g[7], g[9], g[11] ]);
            offD = round(Int64, W[ g[7], g[9], g[11] ]);

            # 1.2-. Search volume coordinates after shifting.
            scoords, so = searchGrid( g, SM, offH, offW, offD, h, w, d )

            setTo0!( search )
            putWithinSearch!( search, vol2, so, scoords )

            setTo0!( sumS  );
            setTo0!( sumS2 );
            znccSumArrays!( search, sumS, sumS2 );

            meanI = Statistics.mean( interr );
            meanS = Statistics.mean( search );

            setTo0!( pads )
            setTo0!( padi )
            putWithinPadded!( padi, interr, meanI );
            putWithinPadded!( pads, search, meanS );

            meanS2 = meanS*meanS;
            meanI2 = meanI*meanI;
            stdI   = Statistics.std( interr, mean=meanI );

            # 2-. Cross-Correlation
            crossCorrelation!( ZNCC(), cmatrix, num, pads, padi, plan, iplan,
                               sumS, sumS2, meanS, meanS2, stdI, ssize, IA_mp, shifts );

            # 3-. Calculation of displacement
            ( r, c, z ) = approxTranslation( cmatrix, peak, Val(last_mp) )
			if ( !ignoreSN && last_mp )
            	SN[ g[7], g[9], g[11] ] = SNRatio( cmatrix, sigNoise, width )
			end

            # 4-. Updating U, V, W matrices
            t1 =  g[7]:g[8];
            t2 =  g[9]:g[10];
            t3 = g[11]:g[12];
            U[ t1, t2, t3 ] .= offH - r;
            V[ t1, t2, t3 ] .= offW - c;
            W[ t1, t2, t3 ] .= offD - z;
        end
    end # Multi-pass loop

    return U, V, W, SN
end

""" ZNCC out of place CROSS CORRELATION """
function crossCorrelation( ::ZNCC, f::A{T,3}, g::A{T,3}; typ=Float32 ) where {T<:Real, U<:Real}

    sizef  = size( f );
    sizeg  = size( g );

    sumf   = zeros( typ, sizef .+ 1 );
    sumf2  = zeros( typ, sizef .+ 1 );
    znccSumArrays!( f, sumf, sumf2 );

    csize = sizef .+ sizeg;
    corr  = zeros( typ, csize );
    num   = zeros( typ, csize );
    padf  = zeros( Complex{typ}, csize );
    padg  = zeros( Complex{typ}, csize );
    plan  =  plan_fft!( padf );
    iplan = plan_ifft!( padf );
	shifts = div.( csize, 2 ) .- 1;

    meanf = Statistics.mean( f );
    meang = Statistics.mean( g );

    putWithinPadded!( padf, f, meanf );
    putWithinPadded!( padg, g, meang );

    meanf2 = meanf^2;
    meang2 = meang^2;
    stdg   = Statistics.std( g, mean=meang );

    crossCorrelation!( ZNCC(), corr, num, padf, padg, plan, iplan,
                       sumf, sumf2, meanf, meanf2, stdg, sizef, sizeg, shifts );

    GC.gc();
    return corr;
end

function crossCorrelation( ::ZNCC, f::A{T,2}, g::A{T,2}; typ=Float32 ) where {T<:Real, U<:Real}

    sizef  = size( f );
    sizeg  = size( g );

    sumf   = zeros( typ, sizef .+ 1 );
    sumf2  = zeros( typ, sizef .+ 1 );
    znccSumArrays!( f, sumf, sumf2, (0,0) );

    csize = sizef .+ sizeg;
    corr  = zeros( typ, csize );
    num   = zeros( typ, csize );
    padf  = zeros( Complex{typ}, csize );
    padg  = zeros( Complex{typ}, csize );
    plan  =  plan_fft!( padf );
    iplan = plan_ifft!( padf );
	shifts = div.( csize[1:2], 1 ) .- 1;

    meanf = Statistics.mean( f );
    meang = Statistics.mean( g );

    putWithinPadded!( padf, f, meanf );
    putWithinPadded!( padg, g, meang );

    meanf2 = meanf^2;
    meang2 = meang^2;
    stdg   = Statistics.std( g, mean=meang );

    crossCorrelation!( ZNCC(), corr, num, padf, padg, plan, iplan,
                       sumf, sumf2, meanf, meanf2, stdg, sizef, sizeg, shifts);

    GC.gc();
    return corr;
end


const DEFAULT_JACOBIAN_ϵ = 8sqrt(eps())

abstract type JacobianParameters{T} end

compute_jacobian!(j::AbstractMatrix, x::AbstractVector, jacobian::JacobianParameters) = jacobian(j,x)

function check_jacobian(J::AbstractMatrix)
    println("Condition Number of Jacobian: ", cond(J))
    println("Determinant of Jacobian:      ", det(J))
    println("minimum(|Jacobian|):          ", minimum(abs.(J)))
    println("maximum(|Jacobian|):          ", maximum(abs.(J)))
    println()
end

function print_jacobian(J::AbstractMatrix)
    display(J)
    println()
end

struct JacobianParametersUser{T, JT} <: JacobianParameters{T}
    J!::JT
end

JacobianParametersUser(J!::JT, ::AbstractArray{T}) where {T,JT} = JacobianParametersUser{T,JT}(J!)

function (jac::JacobianParametersUser{T})(J::AbstractMatrix{T}, x::AbstractVector{T}) where {T}
    jac.J!(J, x)
end

struct JacobianParametersAD{T, FT, JT <: ForwardDiff.JacobianConfig, YT <: AbstractVector{T}} <: JacobianParameters{T}
    F!::FT
    Jconfig::JT
    ty::YT

    function JacobianParametersAD{T}(F!::FT, Jconfig::JT, y::YT) where {T, FT, JT, YT}
        new{T, FT, JT, YT}(F!, Jconfig, y)
    end

end

function JacobianParametersAD(F!::FT, x::AbstractVector{T}, y::AbstractVector{T}) where {T, FT}
    Jconfig = ForwardDiff.JacobianConfig(F!, y, x)
    JacobianParametersAD{T}(F!, Jconfig, zero(y))
end

function JacobianParametersAD{T}(F!::FT, nx::Int, ny::Int) where {T, FT}
    tx = zeros(T, nx)
    ty = zeros(T, ny)
    JacobianParametersAD(F!, tx, ty)
end

JacobianParametersAD{T}(F!, n) where {T} = JacobianParametersAD{T}(F!, n, n)

function (jac::JacobianParametersAD{T})(J::AbstractMatrix{T}, x::AbstractVector{T}) where {T}
    ForwardDiff.jacobian!(J, jac.F!, jac.ty, x, jac.Jconfig)
end

function compute_jacobian_ad!(J::AbstractMatrix{T}, x::AbstractVector{T}, F!::FT) where {T, FT}
    jac = JacobianParametersAD{T}(F!, length(x))
    jac(J,x)
end

struct JacobianParametersFD{T, FT} <: JacobianParameters{T}
    ϵ::T
    F!::FT
    f1::Vector{T}
    f2::Vector{T}
    e::Vector{T}
    tx::Vector{T}
end

function JacobianParametersFD{T}(F!::FT, nx::Int, ny::Int; ϵ=DEFAULT_JACOBIAN_ϵ) where {T, FT}
    f1 = zeros(T, ny)
    f2 = zeros(T, ny)
    e  = zeros(T, nx)
    tx = zeros(T, nx)
    JacobianParametersFD{T,FT}(ϵ, F!, f1, f2, e, tx)
end

JacobianParametersFD{T}(F!, n; kwargs...) where {T} = JacobianParametersFD{T}(F!, n, n; kwargs...)

function (jac::JacobianParametersFD{T})(J::AbstractMatrix{T}, x::AbstractVector{T}) where {T}
    local ϵⱼ::T

    for j in eachindex(x)
        ϵⱼ = jac.ϵ * x[j] + jac.ϵ
        fill!(jac.e, 0)
        jac.e[j] = 1
        jac.tx .= x .- ϵⱼ .* jac.e
        jac.F!(jac.f1, jac.tx)
        jac.tx .= x .+ ϵⱼ .* jac.e
        jac.F!(jac.f2, jac.tx)
        for i in eachindex(x)
            J[i,j] = (jac.f2[i] - jac.f1[i]) / (2ϵⱼ)
        end
    end
end

function compute_jacobian_fd!(J::AbstractMatrix{T}, x::AbstractVector{T}, F!::FT; kwargs...) where {T, FT}
    jac = JacobianParametersFD{T}(F!, length(x); kwargs...)
    jac(J,x)
end

function JacobianParameters{T}(ForJ, nx::Int, ny::Int; mode = :autodiff, diff_type = :forward, kwargs...) where {T}
    if mode == :autodiff
        if diff_type == :forward
            Jparams = JacobianParametersAD{T}(ForJ, nx, ny)
        else
            Jparams = JacobianParametersFD{T}(ForJ, nx, ny; kwargs...)
        end
    else
        Jparams = JacobianParametersUser{T, typeof(ForJ)}(ForJ)
    end
    return Jparams
end

JacobianParameters{T}(ForJ, n::Int; kwargs...) where {T} = JacobianParameters{T}(ForJ, n, n; kwargs...)

JacobianParameters{T}(J!, F!, nx, ny; kwargs...) where {T} = JacobianParameters{T}(J!, nx, ny; mode = :user, kwargs...)

JacobianParameters{T}(J!::Nothing, F!, nx, ny; kwargs...) where {T} = JacobianParameters{T}(F!, nx, ny;  mode = :autodiff, kwargs...)

JacobianParameters{T}(J!, F!, n; kwargs...) where {T} = JacobianParameters{T}(J!, F!, n, n; kwargs...)

function compute_jacobian!(j::AbstractMatrix{T}, x::AbstractVector{T}, ForJ; kwargs...) where {T}
    jacobian = JacobianParameters{T}(ForJ, size(j,1), size(j,2); kwargs...)
    jacobian(j,x)
end


using Revise, ApproxOperator, LaTeXStrings, GLMakie

elements, nodes = importmsh("./msh/patchtest.msh")
nₚ = length(nodes[:x])

type = (Node,:Quadratic2D,:□,:CubicSpline)
elements["Ω"] = ReproducingKernel{type...,:Tri3}(elements["Ω"])
s = 2.5*0.25*ones(nₚ)
push!(nodes,:s₁=>s,:s₂=>s,:s₃=>s)

set𝓖!(elements["Ω"],:TriGI3,:∂1)

# n = 12
n = 21
inte = 100
# as = collect(0.0:1/inte:1.0)
as = LinRange(0.0,1.0,inte)

𝗠 = elements["Ω"][1].𝗠
𝝭 = elements["Ω"][1].𝝭
ap = ReproducingKernel{type...,:Node}([Node(i,nodes) for i in 1:nₚ],Node[],𝗠,𝝭)

# x = [x for x in as for y in as]
# y = [y for x in as for y in as]
z = [get𝝭(ap,(x,y,0.0),[n])[1] for x in as, y in as]

# plot(x,y,z,
#     st=:surface,
#     c=:jet1,
#     xlims=(0.0,1.0),
#     ylims=(0.0,1.0),
#     zlims=(-0.1,1.0),
#     xlabel=L"x",
#     ylabel=L"y",
#     zlabel=L"\Psi_I",
#     ztickfontrotation=10.0,
#     framestyle=:box,
#     legend=:none,
#     aspect_ratio=0.8,
#     dpi=300
#     )
f = Figure(fontsize=14)
ax1 = Axis3(f[1,1],
      xlabel="",
      ylabel="",
      zlabel="",
      limits=(0.0,1.0,0.0,1.0,-0.1,1.0),
      aspect=(1,1,0.7)
)
# ax2 = Axis3(f[1,1],limits=(0.0,1.0,0.0,1.0,-0.1,1.0))
surface!(as,as,z,colormap = :jet1)

meshscatter!([nodes[:x][n]], [nodes[:y][n]], [nodes[:z][n]],markersize = 0.05, color=:red)

save("shape1.png",f)
# ax = Axis3(figure[1,1],xlabel="x")