### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ 09637070-a4f8-11eb-3e6d-ab2106c891d3
using TestImages, FourierTools, ImageShow, Colors, IndexFunArrays

# ╔═╡ b0b0700c-a8f0-4f13-b9d2-799db5904b69
md"## More complex example using FourierTools"

# ╔═╡ 2888c718-8300-47be-8716-ca67908e3f32
img = Float32.(testimage("fabio_gray_512"));

# ╔═╡ c77dab69-ea8b-4cd6-8169-f7572e600126
Gray.(img)

# ╔═╡ 2e879e20-bb7b-41f7-8c72-833e959587fd
md"## Put a low pass filter in Fourier space"

# ╔═╡ e1239869-34fb-4984-b72d-0b1eb8d765c8
begin
	low_pass = IndexFunArrays.rr(size(img)) .< 30;
	low_pass_psf = real(iffts(low_pass));
end

# ╔═╡ 0d190942-5888-42b2-85cf-1522b8a2b8c4
Gray.(low_pass)

# ╔═╡ e299a827-50c2-4fee-89b2-96ac95d4f081
md"## Now take only frequencies passing the low pass filter"

# ╔═╡ 56445c4f-69b4-43ce-81a1-809216e79ff5
img_filtered = real(iffts(ffts(img) .* low_pass))

# ╔═╡ e3b63443-5234-4e5b-a074-b91957cffc98
Gray.(img_filtered)

# ╔═╡ 4e82292b-39c2-4cae-a052-2bc2c7be4100
md"## Since the image is clearly bandlimited, we can downsample without loosing data"

# ╔═╡ 72626a88-b95f-45ac-85cc-622195b136ed
img_downsample = resample(img_filtered, (150, 150));

# ╔═╡ 197d57ec-d89e-4d7d-a209-53535f6e8045
Gray.(img_downsample)

# ╔═╡ 1f06fbca-39a2-4c08-84ca-c5a00a524b0e
md"## Downsampling preserves all the information!"

# ╔═╡ d5bae0f3-55d1-41c2-846e-03b4e5c47529
resample(img_downsample, size(img_filtered)) ≈ img_filtered

# ╔═╡ 8eaac2d0-a698-4c16-809f-4c8f61a99675
md"## Use convolution instead of handcrafted convolution"

# ╔═╡ de770699-474d-42f1-a9aa-d5bfdeb2261e
img_filtered2 = conv(img, low_pass_psf)

# ╔═╡ e83ab89a-ea25-4623-a198-503676f445ef
Gray.(img_filtered2)

# ╔═╡ 53546c5d-f718-4618-bcf1-2f6e9cd2c118
img_filtered2 ≈ img_filtered

# ╔═╡ a49ab07b-59e4-488d-ad35-522fb298dd28
md"## Rotate image"

# ╔═╡ 6b9e4923-4104-41ff-959a-2b5423f14542
begin
	img_pad = FourierTools.center_set!(zeros(Float32, (800, 800)), img_filtered2)
	Gray.(rotate(img_pad, 30))
end

# ╔═╡ 63e7924a-9f11-4088-8e9d-35219dfc7caa
md"## Shift image"

# ╔═╡ 4863b1af-16e7-44d0-a868-13220b768bf9
begin
	Gray.(shift(img_filtered, (200.23, 0)))
end

# ╔═╡ Cell order:
# ╠═b0b0700c-a8f0-4f13-b9d2-799db5904b69
# ╠═09637070-a4f8-11eb-3e6d-ab2106c891d3
# ╠═2888c718-8300-47be-8716-ca67908e3f32
# ╠═c77dab69-ea8b-4cd6-8169-f7572e600126
# ╠═2e879e20-bb7b-41f7-8c72-833e959587fd
# ╠═e1239869-34fb-4984-b72d-0b1eb8d765c8
# ╠═0d190942-5888-42b2-85cf-1522b8a2b8c4
# ╠═e299a827-50c2-4fee-89b2-96ac95d4f081
# ╠═56445c4f-69b4-43ce-81a1-809216e79ff5
# ╠═e3b63443-5234-4e5b-a074-b91957cffc98
# ╠═4e82292b-39c2-4cae-a052-2bc2c7be4100
# ╠═72626a88-b95f-45ac-85cc-622195b136ed
# ╠═197d57ec-d89e-4d7d-a209-53535f6e8045
# ╠═1f06fbca-39a2-4c08-84ca-c5a00a524b0e
# ╠═d5bae0f3-55d1-41c2-846e-03b4e5c47529
# ╠═8eaac2d0-a698-4c16-809f-4c8f61a99675
# ╠═de770699-474d-42f1-a9aa-d5bfdeb2261e
# ╠═e83ab89a-ea25-4623-a198-503676f445ef
# ╠═53546c5d-f718-4618-bcf1-2f6e9cd2c118
# ╠═a49ab07b-59e4-488d-ad35-522fb298dd28
# ╠═6b9e4923-4104-41ff-959a-2b5423f14542
# ╠═63e7924a-9f11-4088-8e9d-35219dfc7caa
# ╠═4863b1af-16e7-44d0-a868-13220b768bf9


module FemAdjoint

using LinearAlgebra, StatsBase, GroupSlices
using ForwardDiff, ReverseDiff, SparseDiffTools, SparseArrays, Symbolics
using FileIO, JLD2, UnicodePlots

include("assembly.jl")
include("plot.jl")
# different test cases
include("costproduv.jl")
#-------------------------------------------------------------------------------
function costnormU(pin, t)
    p = if typeof(pin) <: Vector
        reshape(pin, length(pin) ÷ 2, 2)
    elseif typeof(pin) <: Matrix
        pin
    end
    np = size(p, 1)
    # assemble matrices
    I, J = indKM_sparse(t)
    SK, SM = FemAdjoint.assembKM_P12D(p, t)
    K = sparse(I, J, SK, np, np)
    Ib = unique(FemAdjoint.btri(t)[:])
    K += sparse(Ib, Ib, fill(1e8, length(Ib)), np, np) # Dirichlet conditions
    M = sparse(I, J, SM, np, np)
    F = M * ones(np)
    #F = ones(np)
    U = K \ F
    # L² cost
    val = sum(U[k] ^ 2 for k = 1:np) / np
    return val
end
#-------------------------------------------------------------------------------
function ∇costnormU(pin, t;
                    jacK::SparseMatrixCSC{Float64, Int64} = spzeros(0, 0),
                    jacM::SparseMatrixCSC{Float64, Int64} = spzeros(0, 0))
    p = if typeof(pin) <: Vector
        reshape(pin, length(pin) ÷ 2, 2)
    elseif typeof(pin) <: Matrix
        pin
    end
    np = size(p, 1)
    # assemble matrices
    I, J = indKM_sparse(t)
    SK, SM = FemAdjoint.assembKM_P12D(p, t)
    K = sparse(I, J, SK, np, np)
    Ib = unique(FemAdjoint.btri(t)[:])
    K += sparse(Ib, Ib, fill(1e8, length(Ib)), np, np) # Dirichlet conditions
    M = sparse(I, J, SM, np, np)
    # compute state function
    F = M * ones(np)
    #F = ones(np)
    U = K \ F
    # plot
    plotunicode(U,  title = "U")
    # jacobian structures
    fK(y, x) = FemAdjoint.assembK_P12D_inplace!(x, t, y)
    fM(y, x) = FemAdjoint.assembM_P12D_inplace!(x, t, y)
    if length(jacK) == 0
        y = similar(FemAdjoint.assembKM_P12D(p, t)[1])
        sparsity_pattern = Symbolics.jacobian_sparsity(fK, y, p[:])
        jacK = Float64.(sparse(sparsity_pattern))
    end
    if length(jacM) == 0
        y = similar(FemAdjoint.assembKM_P12D(p, t)[2])
        sparsity_pattern = Symbolics.jacobian_sparsity(fM, y, p[:])
        jacM = Float64.(sparse(sparsity_pattern))
    end
    # evaluate jacobians
    colorsK = matrix_colors(jacK)
    JK = forwarddiff_color_jacobian!(jacK, fK, p[:], colorvec = colorsK)
    colorsM = matrix_colors(jacM)
    JM = forwarddiff_color_jacobian!(jacM, fM, p[:], colorvec = colorsM)
    # cost gradient and adjoint function
    JFU = - 2 * U / np
    λ = K' \ JFU
    plotunicode(λ, colormap = :plasma, title = "λ")
    # compute full gradient
    g = zeros(np, 2)
    for (i, j, s) ∈ zip(findnz(JK)...)
        g[j] += λ[J[i]] * s * U[I[i]]
    end
    for (i, j, s) ∈ zip(findnz(JM)...)
        g[j] -= λ[J[i]] * s
    end
    return g
end
#-------------------------------------------------------------------------------
function ∇costnormU_direct(pin, t;
                           plotsolutions::Bool = false)
    p = if typeof(pin) <: Vector
        reshape(pin, length(pin) ÷ 2, 2)
    elseif typeof(pin) <: Matrix
        pin
    end
    np = size(p, 1)
    # assemble matrices
    I, J = indKM_sparse(t)
    SK, SM = FemAdjoint.assembKM_P12D(p, t)
    K = sparse(I, J, SK, np, np)
    Ib = unique(FemAdjoint.btri(t)[:])
    K += sparse(Ib, Ib, fill(1e8, length(Ib)), np, np) # Dirichlet conditions
    M = sparse(I, J, SM, np, np)
    # compute state function
    F = M * ones(np)
    #F = ones(np)
    U = K \ F
    # plot
    plotsolutions && plotunicode(U,  title = "U")
    # cost gradient and adjoint function
    JFU = - 2 * U / np
    λ = K' \ JFU
    plotsolutions && plotunicode(λ, colormap = :plasma, title = "λ")
    # compute full gradient by automatic differentiation
    #∇prodgrad = ForwardDiff.gradient(x -> prodgrad(x, t, U, λ), p)
    ∇prodgrad = ReverseDiff.gradient(x -> prodgrad(x, t, U, λ), p)
    return ∇prodgrad
end
#-------------------------------------------------------------------------------
function prodgrad(pin, t, U, λ)
    p = reshape(pin, length(pin) ÷ 2, 2)
    np = size(p, 1)
    # assemble matrices
    I, J = indKM_sparse(t)
    SK, SM = FemAdjoint.assembKM_P12D(p, t)
    val = 0.
    for (i, j, s) ∈ zip(I, J, SK)
        val += λ[i] * s * U[j]
    end
    for (i, j, s) ∈ zip(I, J, SM)
        val -= λ[i] * s
    end
    return val
end

end # module


# ----------------------------------------------------
# --- Import modules
# ----------------------------------------------------
using DigitalComm
using Test
# ----------------------------------------------------
# --- Tests
# ----------------------------------------------------
println("Tests for Hard constellation");

for alphabet = [4;16;64;256]
	@testset  "$alphabet-QAM" begin
		# Minimal Tx
		mcs	  = alphabet;		  # Constellation size
		n	  = Int(log2(mcs));		  # Bit per symbol
		N	  = 4096;	  # Numbver of test symbols
		bitSeq	= genBitSequence(N*n);
		qamSeq	= bitMappingQAM(mcs,bitSeq);
		# Minimal Rx
		qamDec	= hardConstellation(mcs,qamSeq);
		# Checking size of decoding vector
		@test length(qamDec) == N;
		# Checking that bang and classic methods are the same
		buff  = zeros(Complex{Float64},N);
		hardConstellation!(buff,mcs,qamSeq);
		@test all(qamDec .== buff);
		# Checking that decoded method is initial method
		@test all(qamDec .≈ qamSeq)
		# Checking that introduces a small biais does not change voronoi region
		biais	= 1/sqrt(2/3*(mcs-1)) * 0.1 * (1+1im)
		@test all( hardConstellation(mcs, qamSeq .+ biais) .≈ qamSeq)
	end
end


using ADCME
using PyPlot
using ProgressMeter
using Statistics

function f(x, σ)
    ε = randn(size(x)...) * σ
    return 10 * sin.(2π*x) + ε
end

batch_size = 32
noise = 1.0

X = reshape(LinRange(-0.5, 0.5, batch_size)|>Array, :, 1)
y = f(X, noise)
y_true = f(X, 0.0)

close("all")
scatter(X, y, marker="+", label="Training Data")
plot(X, y_true, label="Truth")
legend()
savefig("bnn_training_data.png")

mutable struct VariationalLayer
    units
    activation
    prior_σ1
    prior_σ2
    prior_π1
    prior_π2
    Wμ
    bμ
    Wρ
    bρ
    init_σ
end

function VariationalLayer(units; activation=relu, prior_σ1=1.5, prior_σ2=0.1,
        prior_π1=0.5)
    init_σ = sqrt(
        prior_π1 * prior_σ1^2 + (1-prior_π1)*prior_σ2^2
    )
    VariationalLayer(units, activation, prior_σ1, prior_σ2, prior_π1, 1-prior_π1,
                        missing, missing, missing, missing, init_σ)
end

function kl_loss(vl, w, μ, σ)
    dist = ADCME.Normal(μ,σ)
    return sum(logpdf(dist, w)-logprior(vl, w))
end

function logprior(vl, w)
    dist1 = ADCME.Normal(constant(0.0), vl.prior_σ1)
    dist2 = ADCME.Normal(constant(0.0), vl.prior_σ2)
    log(vl.prior_π1*exp(logpdf(dist1, w)) + vl.prior_π2*exp(logpdf(dist2, w)))
end

function (vl::VariationalLayer)(x)
    x = constant(x)
    if ismissing(vl.bμ)
        vl.Wμ = get_variable(vl.init_σ*randn(size(x,2), vl.units))
        vl.Wρ = get_variable(zeros(size(x,2), vl.units))
        vl.bμ = get_variable(vl.init_σ*randn(1, vl.units))
        vl.bρ = get_variable(zeros(1, vl.units))
    end
    Wσ = softplus(vl.Wρ)
    W = vl.Wμ + Wσ.*normal(size(vl.Wμ)...)
    bσ = softplus(vl.bρ)
    b = vl.bμ + bσ.*normal(size(vl.bμ)...)
    loss = kl_loss(vl, W, vl.Wμ, Wσ) + kl_loss(vl, b, vl.bμ, bσ)
    out = vl.activation(x * W + b)
    return out, loss
end

function neg_log_likelihood(y_obs, y_pred, σ)
    y_obs = constant(y_obs)
    dist = ADCME.Normal(y_pred, σ)
    sum(-logpdf(dist, y_obs))
end

ipt = placeholder(X)
x, loss1 = VariationalLayer(20, activation=relu)(ipt)
x, loss2 = VariationalLayer(20, activation=relu)(x)
x, loss3 = VariationalLayer(1, activation=x->x)(x)

loss_lf = neg_log_likelihood(y, x, noise)
loss = loss1 + loss2 + loss3 + loss_lf

opt = AdamOptimizer(0.08).minimize(loss)
sess = Session(); init(sess)
@showprogress for i = 1:5000
    run(sess, opt)
end

X_test = reshape(LinRange(-1.5,1.5,32)|>Array, :, 1)
y_pred_list = []
@showprogress for i = 1:10000
    y_pred = run(sess, x, ipt=>X_test)
    push!(y_pred_list, y_pred)
end

y_preds = hcat(y_pred_list...)

y_mean = mean(y_preds, dims=2)[:]
y_std = std(y_preds, dims=2)[:]

close("all")
plot(X_test, y_mean)
scatter(X[:], y[:], marker="+")
fill_between(X_test[:], y_mean-2y_std, y_mean+2y_std, alpha=0.5)
savefig("bnn_prediction.png")


using QuadGK

@testset "Integration and Coupling generation" begin
	# ----------------------------------------
	# Test the numerical integration capabilities
	antiDerX = ChiralWaveguide.__integrate(x -> x, (0.0, 1.0))

	@test antiDerX(1.0) ≈ 1/2	# local accuracy
	@test sum(x -> abs(antiDerX(x) - x^2 / 2), [0.0:0.01:1.0;]) ≤ 1e-12 # global accuracy
end

@testset "Modes with offset" begin
	μ, σ = 10.0, 0.1

	# ----------------------------------------
	# Test normalisation
	compressions = [Algebraic(σ = σ, μ = μ), Exponential(σ = σ, μ = μ), Trigonometric(σ = σ, μ = μ)]
	for compression ∈ compressions
		testMode = Mode(
			t -> abs(t - μ) < σ/2 ? √(1/σ) : 0.0, compression = compression,
			abstol = 1e-8, reltol = 1e-10
		)
		@test abs(testMode.norm(2μ) - 1.0) < 1e-7
	end
end

@testset "Numerical Mode Couplings" begin
	modes = [
		FourierMode(1/√(5), [1/√(5), 1/√(5)], [1/√(5), 1/√(5)], σ = 3.0),
		HardBoxMode(σ = 3.0),    # give a width to help the integrator
		SoftBoxMode(σ = 3.0),    # A little bit of fine tuning is involved here, but
		SoftBoxExpMode(σ = 2.5), # I simple want to test the normalisation conditions.
		GaussMode(),
		ExpMode(γ =  4.0),
		ExpMode(γ = -4.0)
	]
	# ----------------------------------------
	# Test Numerical Derivations
	ts = range(-3.0, 3.0, length = 100)
	for mode ∈ modes
		numericalMode = Mode(t -> mode.modeFunction(t), reltol = 1e-8, abstol = 1e-12)
		@test sum(abs, numericalMode.gₒ.(ts) .- mode.gₒ.(ts)) / length(ts) < 1e-7
	end

	# ----------------------------------------
	# Integration test
	for mode ∈ modes
		@test quadgk(
			t -> mode.modeFunction(t) * mode.gₒ(t),
			-10.0, 10.0
		)[1] ≈ -2.0
	end
end

@testset "FourierMode and HardBoxMode equivalence" begin
	hb = HardBoxMode()
	f  = FourierMode(1., [], [])

	ts = range(0.0, 1.0, length = 101)

	@test sum(abs, hb.modeFunction.(ts) .- f.modeFunction.(ts)) / length(ts) < 1e-9
	@test sum(abs, hb.gᵢ.(ts) .- f.gᵢ.(ts)) / length(ts) < 1e-9
	@test sum(abs, hb.gₒ.(ts) .- f.gₒ.(ts)) / length(ts) < 1e-9
	@test sum(abs, hb.norm.(ts) .- f.norm.(ts)) / length(ts) < 1e-9
end

using QuantumOptics
@testset "Cavity Hopping" begin
	# Idea of this test is, that a cavity for a given mode
	# should have 100% into a cavity for the same mode.
	modes = [
		FourierMode(rand(), rand(2), rand(2), t₀ = -5.0),
		HardBoxMode(t₀ = -5.0),
		SoftBoxMode(),
		SoftBoxExpMode(),
		GaussMode(),
		ExpMode(γ =  3.5), # Needed to fit in the time window
		ExpMode(γ = -3.5)
	]

	for mode ∈ modes
		for n ∈ [1, 2, 3] # Fock state scattering
		  basis = FockBasis(n)
		  a = destroy(basis)

		  gᵢ, gₒ = mode.gᵢ, mode.gₒ

		  H(t) = 0.5im * gᵢ(t) * gₒ(t) * (a' ⊗ a - a ⊗ a')
		  L(t) = gᵢ(t) * a ⊗ one(basis) + gₒ(t) * one(basis) ⊗ a

		  ψ₀ = fockstate(basis, n) ⊗ fockstate(basis, 0)

		  local ts, ρs
		  ts, ρs = timeevolution.master_dynamic(
				[-5:0.01:5.0;], ψ₀, (t,ψ) -> (H(t), [L(t)], [L(t)']), reltol = 1e-7, abstol = 1e-9
			)

		  @test abs(expect(2, a'a, ρs[end]) - n) < 1e-7
		end
	end
end


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Compute the Sun velocity vector.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] The Astronomical Almanac for the year 2000 (p. C24).
#
#   [2] http://aa.usno.navy.mil/faq/docs/SunApprox.php
#
#   [3] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorne, CA.
#
#   [4] The Astronomical Almanac for the year 2006.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export sun_velocity_i

"""
    function sun_velocity_i(JD::Number)

Compute the Sun velocity represented in the Mean Equinox of Date (MOD) at the
Julian Day `JD`. The algorithm was obtained by computing the time derivative of
the Sun position in [3, p. 277-279].

"""
function sun_velocity_i(JD::Number)
    # Constants
    cen2s = 1/(36525.0*86400.0)

    # Sun position
    # ==========================================================================

    # Number of Julian centuries from J2000 epoch.
    T_UT1 = (JD-JD_J2000)/36525.0

    # Mean longitude of the Sun [deg].
    λ_m = mod(280.460 + 36000.771T_UT1, 360)

    # Mean anomaly of the Sun [deg].
    #
    # Here, we should use T_TBD (Barycentric Dynamical Time). However, it is
    # sufficient to use T_UT1 because this is a low precision computation [1].
    Ms = mod(357.5291092 + 35999.05034T_UT1, 360)

    # Convert Ms to [rad].
    Ms = deg2rad(Ms)

    # Ecliptic latitude of the Sun [deg].
    λ_e = mod(λ_m + 1.914666471sin(Ms) + 0.019994643sin(2Ms), 360)

    # Obliquity of the ecliptic [deg].
    #
    # Here, we are also approximating T_TBD ≈ T_UT1.
    ϵ = mod(23.439291 - 0.0130042T_UT1, 360)

    # Auxiliary variables.
    sin_Ms  , cos_Ms  = sincos(Ms)
    sin_2Ms , cos_2Ms = sincos(2Ms)
    sin_ϵ   , cos_ϵ   = sincos(deg2rad(ϵ))
    sin_λ_e , cos_λ_e = sincos(deg2rad(λ_e))

    # Distance of the Sun from Earth [m].
    r = ( 1.000140612 - 0.016708617cos_Ms - 0.000139589cos_2Ms )*au2m

    # Sun velocity
    # ==========================================================================

    # Compute the required time derivatives [rad/s].
    ∂λ_m = deg2rad(36000.771)   * cen2s
    ∂Ms  = deg2rad(35999.05034) * cen2s
    ∂ϵ   = deg2rad(-0.0130042)  * cen2s
    ∂r   = (+0.016708617sin_Ms*∂Ms + 2*0.000139589cos_2Ms*∂Ms)*au2m
    ∂λ_e = ∂λ_m + deg2rad(1.914666471cos_Ms*∂Ms + 2*0.019994643cos_2Ms*∂Ms)

    # Compute the Sun velocity vector represented in the Mean Equinox of Date
    # (MOD).
    vs_MOD = SVector{3,Float64}(
               ∂r*cos_λ_e -                            r*sin_λ_e*∂λ_e,
               ∂r*cos_ϵ*sin_λ_e - r*sin_ϵ*sin_λ_e*∂ϵ + r*cos_ϵ*cos_λ_e*∂λ_e,
               ∂r*sin_ϵ*sin_λ_e + r*cos_ϵ*sin_λ_e*∂ϵ + r*sin_ϵ*cos_λ_e*∂λ_e
              )
end


let
    setprecision(BigFloat, 256)

    tol = 1e-20
    p = 3
    n = 20
    t = map(BigFloat, (1:n)//n)
    Ut,Vt = egrss.generators(t,p);
    d = BigFloat("1e-8")*ones(BigFloat,n)
    Wt,c = egrss.potrf(Ut,Vt,d)
    Yt,Zt = egrss.trtri(Ut,Wt,c)

    K = egrss.full(Ut,Vt,d)
    Lref = cholesky(K).L
    L = egrss.full_tril(Ut,Wt,c)
    L_inv = egrss.full_tril(Yt,Zt,1.0./c)

    @test ( norm(I-L*L_inv) < n*tol )
    @test ( norm(L-Lref) < norm(Lref)*tol )
end


struct GaussianCopula
    cov::Array{Float64, 2}
    _chol::Array{Float64, 2}
    # GaussianCopula(cov::Array{Float64,2}) = new(cov, cholesky(cov, Val{:L}))
    # two argument constructor is allowed for easy copy-paste for debugging
end
GaussianCopula(cov::Array{Float64,2}) = GaussianCopula(cov, cholesky(cov).factors')
GaussianCopula(ndim::Int, rho::Float64) = GaussianCopula(rho*ones(ndim, ndim) + (1.0-rho)*diagm(ones(ndim)))

ndim(c::GaussianCopula) = size(c.cov,1)

normal_pdf(x) = exp.(-(x .^ 2) ./ 2.0) ./ sqrt(2*pi) # assumes zero mean, unit variance
normal_cdf(x) = 0.5 .* (1.0 .+ erf.(x ./ sqrt(2.0))) # assumes zero mean, unit variance

function sample_mvn(rng::AbstractRNG, cov_chol::Array{Float64})
    return cov_chol*randn(rng, size(cov_chol,1))
end

function rand(rng::AbstractRNG, c::GaussianCopula, sample::Vector{Float64}=Vector{Float64}(undef, ndim(c)))
    return normal_cdf(sample_mvn(rng, c._chol))
end
rand(c::GaussianCopula) = rand(Base.GLOBAL_RNG, c)


# spherical bessel functions
sbesselj(nu, x) = √(π./2x)*besselj.(nu+1/2, x)
sbessely(nu, x) = √(π./2x)*bessely.(nu+1/2, x)
shankelh1(nu,x) = sbesselj.(nu,x) + 1im.*sbessely.(nu,x)

function cart2d2sph2d(xyz,ax_horz,ax_vert)
# generate 2D meshgrid in spherical coordinates system
# corresponding to Cartesian 2D profile (xz, yz, xy)
# for calculation of field profiles
# Input:
# xyz: 1->xz, 2->yz, 3->xy
# ax_horz, ax_vert: 1D array of axes
    mat1 = ax_horz' .* ones(length(ax_vert))
    mat2 = ones(length(ax_horz))' .* ax_vert

    if xyz == 1
        # xz plot range
        mat_r,mat_theta,mat_phi = cart2sph(mat1,0,mat2)
    elseif xyz == 2
        # yz plot range
        mat_r,mat_theta,mat_phi = cart2sph(0,mat1,mat2)
    elseif xyz == 3
        # xy plot range
        mat_r,mat_theta,mat_phi = cart2sph(mat1,mat2,0)
    end
    return mat_r,mat_theta,mat_phi
end

function cart2sph(x,y,z)
# convert Cartesian coordinates(x,y,z) to spherical coordinates(r,θ,φ)
# θ is a polar angle (θ=0 is z-axis), φ is an azimuthal angle (φ=0 is x-axis)
    r = sqrt.(x.^2 .+ y.^2 .+ z.^2)
    φ = mod.(atan.(y,x) .+ 2*π, 2*π)
    θ = acos.(z./r)
    return r,θ,φ
end


using MomentClosure
using MomentClosure: expected_coeff
using ModelingToolkit: value, simplify
using SymbolicUtils: expand
using Catalyst
using Test

@parameters b
b = b.val
@test expected_coeff(2, 3) == 8
@test isequal(expected_coeff(b, 3), b*(1 + (6b) + (6(b^2))))
expr = ((b*((1 + b)^-1)) + ((b^6)*((1 + b)^-6)) + (57(b^2)*((1 + b)^-2)) + (302(b^3)*((1 + b)^-3)) + (302(b^4)*((1 + b)^-4)) + (57(b^5)*((1 + b)^-5)))*((1 + b)^6)
@test isequal(expected_coeff(b, 6), expr)

@parameters t, c₁, c₂, c₃, c₄, Ω
@variables X(t), Y(t)

S_mat = [ 1 -1  1 -1;
         -1  1  0  0]
a = [c₁*X*Y*(X-1)/Ω^2, c₂*X, c₃*Ω, c₄*X]
rn = ReactionSystemMod(t, [X, Y], [c₁, c₂, c₃, c₄, Ω], a, S_mat)

sys = generate_central_moment_eqs(rn, 2, 4)
expr1 = sys.odes.eqs[2].rhs
@test isequal(MomentClosure.Differential(t)(sys.μ[1,0]), sys.odes.eqs[1].lhs)
μ = sys.μ
M = sys.M
# simplify(value()) is needed due to diffs between ModelingToolkit Num type and SymbolicUtils types
expr2 = c₂*μ[1,0] + c₁*M[1,1]*(Ω^-2) + c₁*μ[0,1]*μ[1,0]*Ω^-2 - c₁*M[2,1]*Ω^-2 -
        2*c₁*M[1,1]*μ[1,0]*Ω^-2 - c₁*M[2,0]*μ[0,1]*Ω^-2 - c₁*μ[0,1]*Ω^-2*μ[1,0]^2
expr2 = simplify(value.(expr2))
@test isequal(expand(expr1), expr2)

sys = generate_central_moment_eqs(rn, 2)
expr1 = sys.odes.eqs[2].rhs
@test isequal(MomentClosure.Differential(t)(sys.μ[1,0]), sys.odes.eqs[1].lhs)
@test isequal(expand(expr1), expr2)

sys = generate_raw_moment_eqs(rn, 2)
μ = sys.μ
expr1 = sys.odes.eqs[4].rhs
expr2 = c₂*μ[2,0] + c₁*μ[1,1]/Ω^2 - c₁*μ[3,1]/Ω^2 -c₂*(μ[1,0] + μ[1,1]) +
        c₃*μ[0,1]*Ω - c₄*μ[1,1] - c₁*μ[1,2]/Ω^2 + c₁*μ[2,2]/Ω^2
expr2 = simplify(value.(expr2))
@test isequal(expand(expr1), expand(expr2))

# corner case - a linear propensity with rate coefficient 1
rn = @reaction_network begin
    1, X → 0
end
sys = generate_raw_moment_eqs(rn, 2)
μ = sys.μ
@test isequal(MomentClosure.Differential(t)(μ[(1,)]) ~ -μ[(1,)], sys.odes.eqs[1])

# time-dependent propensity tests

rn = @reaction_network begin
        (X^1.5), X ⇒ Y
end c₁
@test_throws ErrorException generate_raw_moment_eqs(rn, 2)

rn = @reaction_network begin
        (c₁*Y^3*X^1.5), X ⇒ Y
end c₁
@test_throws ErrorException generate_raw_moment_eqs(rn, 2)

rn = @reaction_network begin
        (c₁*sin(exp(X)+Y)), X ⇒ Y
end c₁
@test_throws ErrorException generate_raw_moment_eqs(rn, 2)

rn = @reaction_network begin
        (c₁*sin(exp(t))+Y^(3+1)), X ⇒ Y
end c₁
eqs = generate_raw_moment_eqs(rn, 2)
@test eqs.q_order == 5


#########################################################################
##Define the function that extracts the first r factors of
##the data in the X matrix from time 1 to time t
#########################################################################
function get_factors(vec_î, X_lag, r)
    x_j = X_lag[1:vec_î, :]
    x_j_c = x_j .- mean(x_j, 1)
    # computing the weight matrix
    covx = cov(x_j_c)
    # W is the weight matrix
    F = eig(covx)[2]
    W = flipdim(F, 2)::Array{Float64, 2}
    res = (x_j_c * W)[:, 1:r]
    return res
end

#########################################################################
#Forecast using a dynamic factor model with lagged aggregate variables
#and lagged values of first "r" principal components (factors).
#########################################################################
function dfm_oos_forc(J, h, h_lag, Y, Y_lag, X_lag, r)
    ## Vector Index for oos forecast goes from J - (h_lag + h) to length(Y)
    ŷ = zeros(Float64, size(J - (h_lag + h):length(Y)))
    for (n, vec_î) in enumerate(J - (h_lag + h):length(Y))
        LHS = Y[1:(vec_î - h)]
        sample_size = length(LHS)
        factors = get_factors(vec_î, X_lag, r)
        RHS = [ones(sample_size) Y_lag[1:sample_size] factors[1:sample_size, :]]
        β̂ = OLSestimator(LHS, RHS)
        # Compute Forecasts
        RHS_crnt = [1; Y_lag[vec_î]; factors[vec_î, :]]
        ŷ[n] = dot(β̂, RHS_crnt)
    end
    RMSE = sqrt(mean((Y[J - (h_lag + h):length(Y)] - ŷ) .^ 2))
    return ŷ, RMSE
end

#########################################################################
#Forecast using a dynamic factor model with lagged aggregate variables
#and lagged values of first 1 principal components from
#1) the first order terms and
#2) the second order interactions terms.
#########################################################################
function dfm_oos_forc(J, h, h_lag, Y, Y_lag, X_lag, X_lag2, r1, r2)
    ## Vector Index for oos forecast goes from J - (h_lag + h) to length(Y)
    ŷ = zeros(Float64, size(J - (h_lag + h):length(Y)))
    for (n, vec_î) in enumerate(J - (h_lag + h):length(Y))
        println(n)
        LHS = Y[1:(vec_î - h)]
        sample_size = length(LHS)
        factors1 = get_factors(vec_î, X_lag, r1)
        factors2 = get_factors(vec_î, X_lag2, r2)
        RHS = [ ones(sample_size) Y_lag[1:sample_size] factors1[1:sample_size, :] factors2[1:sample_size, :] ]
        β̂ = OLSestimator(LHS, RHS)
        # Compute Forecasts
        RHS_crnt = [1; Y_lag[vec_î]; factors1[vec_î, :]; factors2[vec_î, :]]
        ŷ[n] = dot(β̂, RHS_crnt)
    end
    RMSE = sqrt(mean((Y[J - (h_lag + h):length(Y)] - ŷ) .^ 2))
    return ŷ, RMSE
end


@polly function kernel_correlation(float_n, data, corr, mean, stddev)
    n,m = size(data)
    eps = 0.1;

    for j = 1:m
        mean[j] = zero(eltype(mean))
        for i = 1:n
            mean[j] += data[i,j]
        end
        mean[j] /= float_n
    end

    for j = 1:m
        stddev[j] = zero(eltype(stddev))
        for i = 1:n
            stddev[j] += (data[i,j] - mean[j]) * (data[i,j] - mean[j])
        end
        stddev[j] /= float_n
        @fastmath stddev[j] = sqrt(stddev[j])
        stddev[j] = stddev[j] <= eps ? one(eltype(stddev)) : stddev[j]
    end

    # Center and reduce the column vectors.
    for i = 1:n, j = 1:m
        data[i,j] -= mean[j]
        @fastmath data[i,j] /= sqrt(float_n) * stddev[j]
    end

    # Calculate the m * m correlation matrix.
    for i = 1:(m-1)
        corr[i,i] = one(eltype(corr))
        for j = (i+1):m
            corr[i,j] = one(eltype(corr))
            for k = 1:n
                corr[i,j] += (data[k,i] * data[k,j])
            end
            corr[j,i] = corr[i,j]
        end
    end

    corr[m-1,m-1] = one(eltype(corr))
end

let
    m = 1200
    n = 1400

    data = zeros(Float32,n,m)
    corr = zeros(Float32,m,m)
    mean = zeros(Float32,m)
    stddev = zeros(Float32,m)

    float_n = Float32(n);

    for i = 1:n, j = 1:m
        data[i,j] = (i*j)/m + i
    end

    SUITE["correlation"] = @benchmarkable kernel_correlation($float_n, data, corr, mean, stddev) setup=(data=copy($data); corr=copy($corr); mean=copy($mean); stddev=copy($stddev))
end


import Base: size, ndims
import LinearAlgebra: diag

export ndoms,
       domainType,
       codomainType,
       is_linear,
       is_eye,
       is_null,
       is_diagonal,
       is_AcA_diagonal,
       is_AAc_diagonal,
       is_orthogonal,
       is_invertible,
       is_full_row_rank,
       is_full_column_rank,
       is_sliced,
       diag_AcA,
       diag_AAc,
       displacement,
       remove_displacement

"""
`domainType(A::AbstractOperator)`

Returns the type of the domain.

```julia
julia> domainType(DFT(10))
Float64

julia> domainType(hcat(Eye(Complex{Float64},(10,)),DFT(Complex{Float64},10)))
(Complex{Float64}, Complex{Float64})
```
"""
domainType

"""
`codomainType(A::AbstractOperator)`

Returns the type of the codomain.

```julia
julia> codomainType(DFT(10))
Complex{Float64}

julia> codomainType(vcat(Eye(Complex{Float64},(10,)),DFT(Complex{Float64},10)))
(Complex{Float64}, Complex{Float64})
```
"""
codomainType

"""
`size(A::AbstractOperator, [dom,])`

Returns the size of an `AbstractOperator`. Type `size(A,1)` for the size of the codomain and `size(A,2)` for the size of the codomain.
"""
size(L::AbstractOperator, i::Int) = size(L)[i]

"""
`ndims(A::AbstractOperator, [dom,])`

Returns a `Tuple` with the number of dimensions of the codomain and domain of an `AbstractOperator`.  Type `ndims(A,1)` for the number of dimensions of the codomain and `ndims(A,2)` for the number of dimensions of the codomain.

"""
ndims(L::AbstractOperator) = count_dims(size(L,1)), count_dims(size(L,2))
ndims(L::AbstractOperator, i::Int) = ndims(L)[i]

count_dims(dims::NTuple{N,Int}) where N = N
count_dims(dims::NTuple{N,Tuple}) where N = count_dims.(dims)

"""
`ndoms(L::AbstractOperator, [dom::Int]) -> (number of codomains, number of domains)`

Returns the number of codomains and domains  of a `AbstractOperator`. Optionally you can specify the codomain (with `dom = 1`) or the domain (with `dom = 2`)

```julia
julia > ndoms(DFT(10,10))
(1,1)

julia> ndoms(hcat(DFT(10,10),DFT(10,10)))
(1, 2)

julia> ndoms(hcat(DFT(10,10),DFT(10,10)),2)
2

julia> ndoms(DCAT(DFT(10,10),DFT(10,10)))
(2,2)
```
"""
ndoms(L::AbstractOperator) = length.(ndims(L))
ndoms(L::AbstractOperator, i::Int) = ndoms(L)[i]

diag_AcA(L::AbstractOperator) = error("is_AAc_diagonal($L) == false")
diag_AAc(L::AbstractOperator) = error("is_AcA_diagonal($L) == false")

"""
`is_linear(A::AbstractOperator)`

Test whether `A` is a `LinearOperator`

```julia
julia> is_linear(Eye(2))
true

julia> is_linear(Sigmoid(Float64,(2,)))
false
```

"""
is_linear(          L::LinearOperator  ) = true
is_linear(          L::AbstractOperator) = false

"""
`is_null(A::AbstractOperator)`

Test whether `A` is null.

```julia
julia> is_null(Zeros(Float64,(10,),(10,)))
true

julia> is_null(Eye(10))
false

```
"""
is_null(            L::AbstractOperator) = false

"""
`is_eye(A::AbstractOperator)`

Test whether `A` is an Identity operator

```julia
julia> is_eye(Eye(10))
true

julia> is_eye(Zeros(Float64,(10,),(10,)))
false

```
"""
is_eye(             L::AbstractOperator) = false

"""
`is_diagonal(A::AbstractOperator)`

Test whether `A` is diagonal.

```julia
julia> is_diagonal(Eye(10))
true

julia> is_diagonal(FiniteDiff((10,)))
false

```
"""
is_diagonal(        L::AbstractOperator) = false

"""
`is_AcA_diagonal(A::AbstractOperator)`

Test whether `A'*A` is diagonal.

```julia
julia> is_AcA_diagonal(Eye(10))
true

julia> is_AcA_diagonal(GetIndex((10,),1:3))
false

```
"""
is_AcA_diagonal(    L::AbstractOperator) = is_diagonal(L)

"""
`is_AAc_diagonal(A::AbstractOperator)`

Test whether `A*A'` is diagonal.

```julia
julia> is_AAc_diagonal(Eye(10))
true

julia> is_AAc_diagonal(GetIndex((10,),1:3))
false

```
"""
is_AAc_diagonal(    L::AbstractOperator) = is_diagonal(L)

"""
`is_orthogonal(A::AbstractOperator)`

Test whether `A` is orthogonal.

```julia
julia> is_orthogonal(DCT(10))
true

julia> is_orthogonal(MatrixOp(randn(3,4)))
false

```
"""
is_orthogonal(      L::AbstractOperator) = false

"""
`is_invertable(A::AbstractOperator)`

Test whether `A` is easily invertable.

```julia
julia> is_invertible(DFT(10))
true

julia> is_invertable(MatrixOp(randn(3,4)))
false

```
"""
is_invertible(      L::AbstractOperator) = false

"""
`is_full_row_rank(A::AbstractOperator)`

Test whether `A` is easily invertable.

```julia
julia> is_full_row_rank(MatrixOp(randn(3,4)))
true

julia> is_full_row_rank(MatrixOp(randn(4,3)))
false
```
"""
is_full_row_rank(   L::AbstractOperator) = false

"""
`is_full_row_rank(A::AbstractOperator)`

Test whether `A` is easily invertable.

```julia
julia> is_full_column_rank(MatrixOp(randn(4,3)))
true

julia> is_full_column_rank(MatrixOp(randn(3,4)))
false
```
"""
is_full_column_rank(L::AbstractOperator) = false

"""
`is_sliced(A::AbstractOperator)`

Test whether `A` is a sliced operator.

```julia
julia> is_sliced(GetIndex((10,), 1:5))
true

```
"""
is_sliced(L::AbstractOperator) = false

"""
`displacement(A::AbstractOperator)`

Returns the displacement of the operator.

```julia
julia> A = AffineAdd(Eye(4),[1.;2.;3.;4.])
I+d  ℝ^4 -> ℝ^4

julia> displacement(A)
4-element Array{Float64,1}:
 1.0
 2.0
 3.0
 4.0

```
"""
function displacement(S::AbstractOperator)
  D = domainType(S)
  if typeof(D) <: Tuple
    x = ArrayPartition(zeros.(D, size(S, 2))...)
  else
	  x = zeros(D, size(S, 2))
  end
  d = S*x
  if all(y -> y == d[1], d )
    return d[1]
  else
    return d
  end
end

"""
`remove_displacement(A::AbstractOperator)`

Removes the displacement of the operator.

"""
remove_displacement(A::AbstractOperator) = A

import Base: convert
function convert(::Type{T}, dom::Type, dim_in::Tuple, L::T) where {T <: AbstractOperator}
	domainType(L) != dom && error("cannot convert operator with domain $(domainType(L)) to operator with domain $dom ")
	size(L,1) != dim_in && error("cannot convert operator with size $(size(L,1)) to operator with domain $dim_in ")
	return L
end

#printing
function Base.show(io::IO, L::AbstractOperator)
	print(io, fun_name(L)*" "*fun_space(L) )
end

function fun_space(L::AbstractOperator)
	dom = fun_dom(L,2)
	codom = fun_dom(L,1)
	return dom*"->"*codom
end

function fun_dom(L::AbstractOperator,n::Int)
	dm = n == 2 ? domainType(L) : codomainType(L)
	sz = size(L,n)
	return string_dom(dm,sz)
end

function string_dom(dm::Type,sz::Tuple)
	dm_st = dm <: Complex ? " ℂ" : " ℝ"
	sz_st = length(sz) == 1 ? "$(sz[1]) " : "$sz "
	return dm_st*"^"*sz_st
end

function string_dom(dm::Tuple,sz::Tuple)
	s = string_dom.(dm,sz)
	length(s) > 3 ? s[1]*"..."*s[end]  : *(s...)
end


# ---
# title: 795. Number of Subarrays with Bounded Maximum
# id: problem795
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array
# link: <https://leetcode.com/problems/number-of-subarrays-with-bounded-maximum/description/>
# hidden: true
# ---
#
# We are given an array `A` of positive integers, and two positive integers `L`
# and `R` (`L <= R`).
#
# Return the number of (contiguous, non-empty) subarrays such that the value of
# the maximum array element in that subarray is at least `L` and at most `R`.
#
#
#
#     **Example :**
#     Input:
#     A = [2, 1, 4, 3]
#     L = 2
#     R = 3
#     Output: 3
#     Explanation: There are three subarrays that meet the requirements: [2], [2, 1], [3].
#
#
# **Note:**
#
#   * L, R  and `A[i]` will be an integer in the range `[0, 10^9]`.
#   * The length of `A` will be in the range of `[1, 50000]`.
#
#
## @lc code=start
using LeetCode

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


@testset "mixture sampler" begin

	cc = dummyChoiceContext()

	@testset "getparams, numparams, paramranges, sample" begin

		subsA = DataGenerators.BernoulliSampler()
		subsB = DataGenerators.DiscreteUniformSampler([7.0,9.0])
		s = DataGenerators.MixtureSampler(subsA,subsB)

	    @testset "numparams and paramranges" begin
			@test DataGenerators.numparams(s) == 2 + DataGenerators.numparams(subsA) + DataGenerators.numparams(subsB)
			prs = DataGenerators.paramranges(s)
	        @test typeof(prs) <: Vector{Tuple{Float64,Float64}}
			@test prs == [(0.0,1.0); (0.0,1.0); DataGenerators.paramranges(subsA); DataGenerators.paramranges(subsB)]
	    end

	    @testset "default params" begin
			@test DataGenerators.getparams(s) == [0.5; 0.5; DataGenerators.getparams(subsA); DataGenerators.getparams(subsB)]
	    end

	    @mtestset "default sampling" reps=Main.REPS alpha=Main.ALPHA begin
	        x, trace = DataGenerators.sample(s, (0,1), cc)
	        @test typeof(x) <: Int
 			@mtest_values_are [0,1,7,8,9] x
			if x <= 1
				@mtest_distributed_as Bernoulli(0.5) x
			else
				@mtest_distributed_as DiscreteUniform(7,9) x
			end
	    end

	end

	@testset "setparams" begin

		subsA = DataGenerators.BernoulliSampler()
		subsB = DataGenerators.DiscreteUniformSampler()
		subsC = DataGenerators.GeometricSampler()
		s = DataGenerators.MixtureSampler(subsA,subsB,subsC)
		prs = DataGenerators.paramranges(s)
		midparams = map(pr->robustmidpoint(pr[1],pr[2]), prs)

		@testset "parameters of subsamplers" begin

			params = [0.1, 0.2, 0.7, 0.4, 1.0, 6.0, 0.3]
			DataGenerators.setparams!(s, params)
			@test DataGenerators.getparams(subsA) == [0.4]
			@test DataGenerators.getparams(subsB) == [1.0, 6.0]
			@test DataGenerators.getparams(subsC) == [0.3]

			@mtestset "consistent with mixture" reps=Main.REPS alpha=Main.ALPHA begin
            	x, trace = DataGenerators.sample(s, (0,1), cc)
		        @mtest_distributed_as MixtureModel([Bernoulli(0.4),DiscreteUniform(1,6),Geometric(0.3)],[0.1,0.2,0.7]) x
			end

		end

		@testset "boundary values choice params $choiceparams" for choiceparams in [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0],]

            DataGenerators.setparams!(s, [choiceparams; [0.4, 1.0, 6.0, 0.3]])

			@mtestset "consistent with mixture" reps=Main.REPS alpha=Main.ALPHA begin
            	x, trace = DataGenerators.sample(s, (0,1), cc)
				if choiceparams[1] == 1.0
					@mtest_distributed_as Bernoulli(0.4) x
				elseif choiceparams[2] == 1.0
					@mtest_distributed_as DiscreteUniform(1,6) x
				elseif choiceparams[3] == 1.0
					@mtest_distributed_as Geometric(0.3) x
				end
			end

		end

		@testset "range check exception choice params $pidx bound $bidx" for pidx in 1:3, bidx in 1:2
			choiceparams = [0.5, 0.4, 0.1]
            choiceparams[pidx] = bidx == 1 ? prevfloat(prs[pidx][bidx]) : nextfloat(prs[pidx][bidx])
            @test_throws ErrorException DataGenerators.setparams!(s, [choiceparams, [0.4, 1.0, 6.0, 0.3]])
		end

	    @testset "wrong number of parameters" begin
	        @test_throws ErrorException DataGenerators.setparams!(s, midparams[1:end-1])
	        @test_throws ErrorException DataGenerators.setparams!(s, [midparams; 0.5])
	    end

	end

	@testset "estimateparams" begin

		subs1A = DataGenerators.BernoulliSampler()
		subs1B = DataGenerators.DiscreteUniformSampler()
		s1 = DataGenerators.MixtureSampler(subs1A,subs1B)

		@testset "from parameters $params" for params in [[0.6, 0.4, 0.7, 10.0, 13.0],]

			DataGenerators.setparams!(s1, params)

			subs2A = DataGenerators.BernoulliSampler()
			subs2B = DataGenerators.DiscreteUniformSampler()
			s2 = DataGenerators.MixtureSampler(subs2A,subs2B)
			otherparams = [0.3, 0.7, 0.3, -40.0, -27.0]
			DataGenerators.setparams!(s2, otherparams)

			traces = map(1:100) do i
				x, trace = DataGenerators.sample(s1, (0,1), cc)
				trace
			end
			estimateparams!(s2, traces)

			@mtestset "consistent with mixture" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s2, (0,1), cc)
		        @mtest_distributed_as MixtureModel([Bernoulli(params[3]),DiscreteUniform(Int(params[4]),Int(params[5]))],params[1:2]) x
			end

		end

	    @testset "too few traces" begin

	        params = [0.2, 0.8, 0.5, 1.0, 6.0]
			DataGenerators.setparams!(s1, params)

			subs2A = DataGenerators.BernoulliSampler()
			subs2B = DataGenerators.DiscreteUniformSampler()
			s2 = DataGenerators.MixtureSampler(subs2A,subs2B)
			otherparams = [0.3, 0.7, 0.3, -40.0, -27.0]
			DataGenerators.setparams!(s2, otherparams)

	        traces = map(1:0) do i
	            x, trace = DataGenerators.sample(s1, (0,1), cc)
	            trace
	        end
	        estimateparams!(s2, traces)

			@mtestset "consistent with mixture" reps=Main.REPS alpha=Main.ALPHA begin
	        	x, trace = DataGenerators.sample(s2, (0,1), cc)
		        @mtest_distributed_as MixtureModel([Bernoulli(otherparams[3]),DiscreteUniform(Int(otherparams[4]),Int(otherparams[5]))],otherparams[1:2]) x
			end

	    end

	end

end


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

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

# ╔═╡ a0b3813e-adab-11eb-2983-616cf2bb6f5e
using LinearAlgebra, DifferentialEquations, PlutoUI,Plots

# ╔═╡ d84dbb08-b1a4-11eb-2d3b-0711fddd1347
html"""

Supplementary Materials

 Forecasting virus outbreakes with Neural ODEs.

"""

# ╔═╡ ef8d6690-720d-4772-a41f-b260d306b5b2
TableOfContents(title="📚 Table of Contents", indent=true, depth=4, aside=true)

# ╔═╡ 30969341-9079-4732-bf55-d6bba2c2c16c
md"""
`sign(x)` returns 0 if x is 0, or ±1 for positive/negative x.
"""

# ╔═╡ 6139554e-c6c9-4252-9d64-042074f68391
begin
	f(y,a) =   sign(y)  + a+1 -  y
	f(y,a,t) = f(y,a) # just for the difeq solver
end

# ╔═╡ 61960e99-2932-4a8f-9e87-f64a7a043489
md"""
a = $(@bind a Slider(-5:.1:5, show_value=true, default=0) )

"""

# ╔═╡ 6f61180d-6900-48fa-998d-36110e79d2dc
gr()

# ╔═╡ 5027e1f8-8c50-4538-949e-6c95c550016e
md"""
### Flux argument

"""

# ╔═╡ 465f637c-0555-498b-a881-a2f6e5714cbb
md"""
y₀ = $(@bind y₀ Slider(-6:.1:6, show_value=true, default=2.0) )
"""

# ╔═╡ a09a0064-c6ab-4912-bc9f-96ab72b8bbca
sol = solve(  ODEProblem( f, y₀, (0,10.0), a ));

# ╔═╡ fb6e34b0-b4eb-11eb-3687-257bc18c3e52
md"""
### More

"""

# ╔═╡ bc96abb0-b4ea-11eb-21a1-fdd01945f9c6
md"""
### Auxiliary Code

"""

# ╔═╡ f8ee2373-6af0-4d81-98fb-23bde10198ef
function plotit(y₀, a)

	sol = solve(  ODEProblem( f, y₀, (0,10.0), a ));
	p = plot( sol , legend=false, background_color_inside=:black , ylims=(-7,7), lw=3, c=:red)

	# plot direction field:
		xs = Float64[]
		ys = Float64[]

	    lrx = LinRange( xlims(p)..., 30)
		for x in  lrx
			for y in LinRange( ylims(p)..., 30)
				v = [1, f(y, a ,x) ]
				v ./=  (100* (lrx[2]-lrx[1]))

				push!(xs, x - v[1], x + v[1], NaN)
				push!(ys, y - v[2], y + v[2], NaN)

			end
		end

	if a<1 hline!( [-1+a],c=:white,ls=:dash); end
	if a>-1  hline!( [1+a],c=:white,ls=:dash);end
	    hline!( [0 ],c=:pink,ls=:dash)
		plot!(xs, ys, alpha=0.7, c=:yellow)
	    ylabel!("y")
	    annotate!(-.5,y₀,text("y₀",color=:red))
	if a>-1
	   annotate!(5,1+a,text(round(1+a,digits=3),color=:white))
	end
	if a<1
	   annotate!(5,-1+a-.4,text(round(-1+a,digits=3),color=:white))
	end
    title!("Solution to y'=f(y,a)")
return(p)
end

# ╔═╡ a4686bca-90d6-4e02-961c-59f08fc37553
plotit(y₀, a)

# ╔═╡ Cell order:
# ╟─d84dbb08-b1a4-11eb-2d3b-0711fddd1347
# ╟─a0b3813e-adab-11eb-2983-616cf2bb6f5e
# ╟─ef8d6690-720d-4772-a41f-b260d306b5b2
# ╟─30969341-9079-4732-bf55-d6bba2c2c16c
# ╟─6139554e-c6c9-4252-9d64-042074f68391
# ╟─61960e99-2932-4a8f-9e87-f64a7a043489
# ╟─6f61180d-6900-48fa-998d-36110e79d2dc
# ╟─a09a0064-c6ab-4912-bc9f-96ab72b8bbca
# ╟─5027e1f8-8c50-4538-949e-6c95c550016e
# ╟─465f637c-0555-498b-a881-a2f6e5714cbb
# ╟─a4686bca-90d6-4e02-961c-59f08fc37553
# ╟─fb6e34b0-b4eb-11eb-3687-257bc18c3e52
# ╟─bc96abb0-b4ea-11eb-21a1-fdd01945f9c6
# ╟─f8ee2373-6af0-4d81-98fb-23bde10198ef


module LinRegOutliers

# using StatsModels
# using DataFrames
# using Distributions
# using Clustering
# using StatsBase
# using LinearAlgebra
# using Plots
# using Optim

# import functions from corresponding packages
import StatsModels: @formula, FormulaTerm, ModelFrame, ModelMatrix
import DataFrames: DataFrame
import Distributions: TDist, Chisq, std, cov, median
import Clustering: Hclust, hclust, cutree
import StatsBase: quantile, standardize, ZScoreTransform, mean, sample
import LinearAlgebra: inv, pinv, det, diag, eigen, norm, rank, qr, Diagonal, UpperTriangular
import Plots: scatter, title!, xlabel!, ylabel!, hline!, vline!, RGB, plot, font, text
import Optim: optimize, NelderMead
import Combinatorics: combinations

# Basis
include("basis.jl")

# Predefined datasets used in outlier detection literature
include("data.jl")

# Regression diagnostics
include("diagnostics.jl")

# Ordinary least squares type and functions
# for fast regression tasks in outlier detection algorithms
include("ols.jl")

# Hadi & Simonoff (1993) algorithm
include("hs93.jl")

# Kianifard & Swallow (1989) algorithm
include("ks89.jl")

# Sebert et. al. (1998) algorithm
include("smr98.jl")

# asm (2000) algorithm
include("asm2000.jl")

# Rousseeuw's Least Median of Squares estimator
include("lms.jl")

# Rousseeuw's Least Trimmed Squares estimator
include("lts.jl")

# Minimum Volume Ellipsoid estimator
# for robust covariance matrix
include("mve.jl")

# MVE - LTS plot for visual detection of regression outliers
include("mveltsplot.jl")

# Billor & Chatterjee & Hadi Algorithm for detecting outliers
include("bch.jl")

# Pena & Yohai (1995) algorithm
include("py95.jl")

# Satman (2013) algorithm
include("satman2013.jl")

# Satman (2015) algorithm
include("satman2015.jl")

# Least Absolute Deviations estimator
include("lad.jl")

# Least Trimmed Absolute Deviations estimator
include("lta.jl")

# Hadi (1992) detecting outliers in multivariate data
# This algorithm is not directly related with the regression
# however, detecting bad leverage points is same as detecting
# outliers in X-space
include("hadi1992.jl")

# Hadi (1994) algorithm
include("hadi1994.jl")

# Gray-scale images of distance matrices
include("dataimage.jl")

# Compact genetic algorithm
include("cga.jl")

# Genetic Algorithm
include("ga.jl")

# Modified and original Satman (2012) algorithms
include("gwlts.jl")

# RANSAC Algorithm
include("ransac.jl")

# CCF formulation and heuristic; Barratt, Angeris, and Boyd (2020)
include("ccf.jl")

# Atkinson94 Algorithm
include("atkinson94.jl")
include("atkinsonstalactiteplot.jl")

# BACON 2000 Algorithm
include("bacon.jl")

# Imon 2005 Algorithm
include("imon2005.jl")

#  Chatterjee & Mächler (1997) Algorithm
include("cm97.jl")

# All-in-one
include("summary.jl")

# Essentials from other packages
export @formula, DataFrame
export mean, quantile

# Basics
export RegressionSetting
export createRegressionSetting
export designMatrix
export responseVector
export applyColumns
export find_minimum_nonzero
export @extractRegressionSetting

# Data
export phones, hbk, stackloss
export weightloss, hs93randomdata, woodgravity
export hills, softdrinkdelivery

# Diagnostics
export dffit, dfbeta
export hatmatrix
export studentizedResiduals
export adjustedResiduals
export jacknifedS
export cooks
export mahalanobisSquaredMatrix
export covratio
export hadimeasure

# Ordinary least squares
export OLS, ols, wls, residuals, predict, coef

# Algorithms
export hs93, hs93initialset, hs93basicsubset
export ks89
export smr98, asm2000
export lms
export lts
export mve, mcd
export mveltsplot
export bch, bchplot
export py95, py95SuspectedObservations
export satman2013
export satman2015, dominates
export lad
export lta
export hadi1992
export hadi1994
export dataimage
export gwcga, galts, ga, cga, RealChromosome
export detectOutliers
export ransac
export ccf
export imon2005
export atkinson94, atkinsonstalactiteplot, generate_stalactite_plot
export bacon
export cm97

end # module


function mean_hausdorff(u, v)
    "
    This function extracts the edges of two images and then finds the average Hausdorff distance
    Both arrays u and v are required to be the same size and they must be binary or boolean images.
    "
    d = Euclidean()

    # Find the coordinates of the pixels that are edges in the images
    edges_1 = find_edges(u)
    edges_2 = find_edges(v)

    # Turn the array of CartesianIndices into an array of integers corresponding to the CartesianIndices
    edges_1 = (transpose ∘ reshape)(reinterpret(Int, edges_1), 2, :)
    edges_2 = (transpose ∘ reshape)(reinterpret(Int, edges_2), 2, :)

    min_euc_list_u = []
    min_euc_list_v = []

    # Loop through every edge point on `edge_1` and find its corresponding closest point to `edge_2`
    for i in 1:size(edges_1, 1)
        euc_list_1 = []
        for j in 1:size(edges_2, 1)
            euc = evaluate(d, edges_1[i, :], edges_2[j, :])
            append!(euc_list_1 , euc)
        end
        append!(min_euc_list_u, minimum(euc_list_1))
    end

    # Loop through every edge point on `edge_2` and find its corresponding closest point to `edge_1`
    for i in 1:size(edges_2, 1)
        euc_list_1 = []
        for j in 1:size(edges_1, 1)
            euc = evaluate(d, edges_2[i, :], edges_1[j, :])
            append!(euc_list_1 , euc)
        end
        append!(min_euc_list_v, minimum(euc_list_1))
    end

    # Take the average of each of these points to return the average Hausdorff distance
    return mean([mean(min_euc_list_u), mean(min_euc_list_v)])
end

function mean_dice(score, target)
    smooth = 1e-5
    intersect = sum(score .* target)
    y_sum = sum(target .* target)
    z_sum = sum(score .* score)
    metric = (2 * intersect + smooth) / (z_sum + y_sum + smooth)
    return metric
end


using DynamicalSystemsBase
using DelayEmbeddings
using Test
using DelimitedFiles

println("\nTesting mdop_embedding.jl...")
@testset "Nichkawde method MDOP" begin

# For comparison reasons using Travis CI we carry out the integration on a UNIX
# OS and save the resulting time series
# # solve Mackey-Glass-Delay Diff.Eq. as in the Paper
# function mackey_glass(du,u,h,p,t)
#   beta,n,gamma,tau = p
#   hist = h(p, t-tau)[1]
#   du[1] = (beta*hist)/(1+hist^n) - gamma * u[1]
# end
# # set parameters
# h(p,t) = 0
# tau_d = 44
# n = 10
# β = 0.2
# γ = 0.1
# δt = 0.5
# p = (β,n,γ,tau_d)
#
# # time span
# tspan = (0.0, 12000.0)
# u0 = [1.0]
#
# prob = DDEProblem(mackey_glass,u0,h,tspan,p; constant_lags=tau_d)
# alg = MethodOfSteps(Tsit5())
# sol = solve(prob,alg; adaptive=false, dt=δt)
#
# s = [u[1] for u in sol.u]
# s = s[4001:end]

s = readdlm(joinpath(tsfolder, "1.csv"))
s = vec(s)
Y = Dataset(s)

theiler = 57

@testset "beta statistic" begin
    ## Test beta_statistic (core algorithm of mdop_embedding)

    taus = 0:100
    β = @inferred DelayEmbeddings.beta_statistic(Y, s, taus, theiler)
    maxi, max_idx = findmax(β)

    @test maxi>4.1
    @test taus[max_idx]>=50

    # # display results as in Fig. 3 of the paper
    # using Plots
    # plot(taus, β, linewidth = 3, label = "1st embedding cycle")
    # plot!(title = "β-statistic for Mackey Glass System as in Fig. 3 in the Paper")
    # xlabel!("τ")
    # ylabel!("β-Statistic")

    # # display results as in Fig. 3 of the paper
    # using PyPlot
    # pygui(true)
    # figure()
    # plot(taus, β, linewidth = 3, label = "1st embedding cycle")
    # scatter(max_idx-1,maxi, c="red")
    # title("β-statistic for Mackey Glass System as in Fig. 3 in the Paper")
    # xlabel("τ")
    # ylabel("β-Statistic")
    # xticks(0:10:100)
    # grid()

    # test different tau range
    taus2 = 1:4:100
    β2 = @inferred beta_statistic(Y, s, taus2, theiler)
    maxi2, max_idx2 = findmax(β2)

    @test maxi2>4.1
    @test taus2[max_idx2]>=40

    # # display results as in Fig. 3 of the paper
    # using Plots
    # plot(taus2, β2, linewidth = 3, label = "1st embedding cycle")
    # plot!(title = "coarse β-statistic for Mackey Glass System as in Fig. 3 in the Paper")
    # xlabel!("τ")
    # ylabel!("β-Statistic")

end

@testset "mdop_embedding univariate" begin

    taus = 0:100
    β = DelayEmbeddings.beta_statistic(Y, s, taus, theiler)
    Y, τ_vals, ts_vals, FNNs, betas = mdop_embedding(s; τs = taus, w = theiler)
    # for different τs
    taus2 = 1:4:100
    Y2, τ_vals2, ts_vals2, FNNs2, betas2 = mdop_embedding(s; τs = taus2, w = theiler)

    @test round.(β, digits=6) == round.(betas[:,1], digits=6)
    @test size(Y,2) == 5
    @test size(Y,2) == size(Y2,2)
    @test sum(findall(x -> x != 1, ts_vals))==0
    @test sum(abs.(diff(τ_vals)) .< 10) == 0

    # # display results as in Fig. 3 of the paper
    # using Plots
    # # Figure as in Fig.3 in the paper
    # plot(taus, betas[:,1], linewidth = 3, label = "embedding cycle 1")
    # plot!([taus[τ_vals[2]+1]],[betas[τ_vals[2]+1,1]], seriestype = :scatter, color="red", label = "")
    # for i = 2:size(betas,2)
    #   plot!(taus, betas[:,i], linewidth = 3, label = "embedding cycle $i")
    #   plot!([taus[τ_vals[i+1]+1]],[betas[τ_vals[i+1]+1,i]], seriestype = :scatter, color="red", label = "")
    # end
    # plot!(title = "β-statistic's for each embedding cycle of Mackey Glass System as in Fig. 3")
    # xlabel!("delay τ")
    # ylabel!("log10 β(τ)")
    #
    # # Figure of coarse grained analysis
    # taus22 = zeros(length(taus2))
    # [taus22[i]=taus2[i] for i = 1: length(taus2)]
    # plot(taus22, betas2[:,1], linewidth = 3, label = "embedding cycle 1")
    # trueind = findall(x -> x == τ_vals2[2], taus22)
    # plot!(taus22[trueind],[betas2[trueind,1]], seriestype = :scatter, color="red", label = "")
    # for i = 2:size(betas2,2)
    #   plot!(taus22, betas2[:,i], linewidth = 3, label = "embedding cycle $i")
    #   trueind = findall(x -> x == τ_vals2[i+1], taus22)
    #   plot!(taus22[trueind],[betas2[trueind,i]], seriestype = :scatter, color="red", label = "")
    # end
    # plot!(title = "β-statistic's for each embedding cycle of Mackey Glass System as in Fig. 3")
    # xlabel!("delay τ")
    # ylabel!("log10 β(τ)")

end

@testset "estimate τ max (Roessler)" begin
    # For comparison reasons using Travis CI we carry out the integration on a UNIX
    # OS and save the resulting time series
    # roe = Systems.roessler([1.0, 0, 0]; a=0.2, b=0.2, c=5.7)
    # sroe = trajectory(roe, 500; dt = 0.05, Ttr = 100.0)
    # writedlm("2.csv", sroe)

    sroe = readdlm(joinpath(tsfolder, "2.csv"))
    tws = 25:32

    τ_m, L = @inferred mdop_maximum_delay(sroe[:, 2], tws)
    @test τ_m == 26
    τ_m, Ls = @inferred mdop_maximum_delay(Dataset(sroe[:, 1:2]), tws)
    @test τ_m == 26

    # # reproduce Fig.2 of the paper
    # tws = 1:2:101
    # τ_m, L = DelayEmbeddings.mdop_maximum_delay(s[:,2]; tw = tws, samplesize=1.0)
    #
    # using Plots
    # twss = zeros(length(tws))
    # [twss[cnt] = i for (cnt,i) in enumerate(tws)]
    # plot(twss,L, label="")
    # xlabel!("time window")
    # ylabel!("L")

    # tw=1:4:200
    # tau_max, LL = DelayEmbeddings.mdop_maximum_delay(sroe; tw=tw)
    #
    # using Plots
    # gui()
    # twss = zeros(length(tw))
    # [twss[cnt] = i for (cnt,i) in enumerate(tw)]
    # plot(twss,LL, label="")
    # xlabel!("time window")
    # ylabel!("L")
end

@testset "mdop_embedding multivariate" begin
    # For comparison reasons using Travis CI we carry out the integration on a UNIX
    # OS and save the resulting time series
    # roe = Systems.roessler([1.0, 0, 0]; a=0.2, b=0.2, c=5.7)
    # sroe = trajectory(roe, 500; dt = 0.05, Ttr = 100.0)
    # writedlm("2.csv", sroe)

    sroe = readdlm(joinpath(tsfolder, "2.csv"))
    tra = Dataset(sroe)
    w1 = estimate_delay(sroe[:,1], "mi_min")
    w2 = estimate_delay(sroe[:,2], "mi_min")

    theiler = w2
    taus = 0:26
    mc = 10

    Y, τ_vals, ts_vals, FNNs, betas =  mdop_embedding(sroe[:,1]; τs = taus, w = theiler, max_num_of_cycles = mc)

    max_idx, ts_number = @inferred DelayEmbeddings.choose_optimal_tau2(betas)
    @test ts_number == 1
    @test taus[max_idx] == τ_vals[2]

    Y2, τ_vals2, ts_vals2, FNNs2, betas2 = mdop_embedding(tra; τs = taus, w = theiler, max_num_of_cycles = mc)
    ttra = standardize(tra)
    b1 = DelayEmbeddings.beta_statistic(Dataset(ttra[:,ts_vals2[1]]), ttra[:,1], taus, theiler)
    b2 = DelayEmbeddings.beta_statistic(Dataset(ttra[:,ts_vals2[1]]), ttra[:,2], taus, theiler)
    b3 = DelayEmbeddings.beta_statistic(Dataset(ttra[:,ts_vals2[1]]), ttra[:,3], taus, theiler)

    @test betas2[1][:,1] == b1
    @test betas2[1][:,2] == b2
    @test betas2[1][:,3] == b3

    @test size(Y2,2) == 3
    @test τ_vals2[1] == τ_vals2[2] == 0
    @test τ_vals2[3] == maximum(taus)

    @test ts_vals2[2] == ts_vals2[3] == 2
    @test ts_vals2[1] == 3

end

end


Random.seed!(0)

# Dynamics
dt = 0.1
m = 50
A = [1 dt 0 0 ; 0 1 0 0 ; 0 0 1 dt; 0 0 0 1]
B = [0 0; dt/m 0; 0 0; 0 dt/m]
W = 0.1*Matrix{Float64}(I,4,4)

dmodel = LinearDynamicsModel(A,B,W);

# Measurement
C = [0 1.0 0 0; 0 0 0 1.0]
V = 0.5*Matrix{Float64}(I,2,2)

omodel = LinearObservationModel(C,V)

# Build KF
kf = KalmanFilter(dmodel,omodel)

# Simulation
b0 = GaussianBelief([0.0,0.0,0.0,0.0], 2.0*Matrix{Float64}(I,4,4))

times = 0:dt:10
Fmag = 1000
action_sequence = [[Fmag*cos(t), Fmag*sin(t)] for t in times]

sim_states, sim_measurements = simulation(kf, b0,action_sequence);

# Filter
filtered_beliefs = run_filter(kf, b0, action_sequence, sim_measurements)
μ, Σ = unpack(filtered_beliefs);

# Tests
@test kf.d isa LinearDynamicsModel
@test kf.o isa LinearObservationModel
@test length(sim_measurements) == length(action_sequence)

first_update = update(kf, b0, action_sequence[1], sim_measurements[1])
first_predict_measure = measure(kf, predict(kf, b0, action_sequence[1]),
    sim_measurements[1]; u = action_sequence[1])
array_isapprox(first_update.μ, first_predict_measure.μ)
array_isapprox(first_update.Σ, first_predict_measure.Σ)

@test length(filtered_beliefs[end].μ) == length(b0.μ)
@test size(filtered_beliefs[end].Σ) == size(b0.Σ)
@test size(μ) == (length(filtered_beliefs), length(filtered_beliefs[1].μ))
@test size(Σ) == (length(filtered_beliefs), size(filtered_beliefs[1].Σ)...)

# change with algorithm / rng changes
array_isapprox(sim_states[end], [59.9223, -7.17527, 178.149, 30.4653];atol=0.001)
array_isapprox(sim_measurements[end], [-7.01552, 30.4259];atol=0.001)
array_isapprox(filtered_beliefs[end].μ, [60.1627, -7.13376, 183.227, 30.3895];
                atol=0.001)
array_isapprox(filtered_beliefs[end].Σ,   [12.6069 0.0320871 0.0 0.0;
                                            0.0320871 0.179129 0.0 0.0;
                                            0.0 0.0 12.6069 0.0320871;
                                            0.0 0.0 0.0320871 0.179129]; atol=0.001)


# Shannon entropy of discrete distribution with probabilities given by p
entropy(p::Real) = iszero(p) ? p : -p*log2(p)
entropy(p::AbstractVector) = sum(entropy, p)
perplexity(p::AbstractVector) = 2^entropy(p)

# Kullback-Leibler divergence
kl(p::Real, q::Real) = iszero(p) ? zero(p) : p*log(p/q) # what about q = 0?
kl(pq::NTuple{2}) = kl(pq...) # unpack tuple argument
# Z is assumed to be normalization constant of Q
function kl(P::AbstractMatOrFac, Q::AbstractMatOrFac, Z::Real = 1)
    val = zero(promote_type(eltype(P), eltype(Q)))
    for i in 1:size(P, 1), j in 1:size(Q, 2)
        val += kl(P[i, j], Q[i, j] / Z)
    end
    return val
end
# specialization for sparse matrix P
# Z is assumed to be normalization constant of Q
function kl(P::SparseMatrixCSC, Q::AbstractMatOrFac, Z::Real = 1)
    I, J, _ = findnz(P)
    val = zero(promote_type(eltype(P), eltype(Q)))
    for (i, j) in zip(I, J) # IDEA: parallelize, simd?
        val += kl(P[i, j], Q[i, j] / Z)
    end
    return val
end

# euclidean distance
_d2(x::Real, y::Real) = (x-y)^2
_d2(x::NTuple{2}) = _d2(x...)
_euclidean(x, y) = sqrt(sum(_d2, zip(x, y)))
euclidean(x::AbstractVector, y::AbstractVector) = _euclidean(x, y) # avoids temporary array
euclidean(x::NTuple, y::NTuple) = _euclidean(x, y) # avoids temporary array

# kernels
gaussian(x::Real) = exp(-x^2/2)
cauchy(x::Real) = inv(1+x^2)

# linear algebra
function hadamard_product(P::AbstractMatrix, Q::AbstractMatOrFac)
    n, m = size(P)
    H = zeros(eltype(P), n, m)
    for j in 1:n, i in 1:m
        H[i, j] = P[i, j] * Q[i, j]
    end
    return H
end

function hadamard_product(P::SparseMatrixCSC, Q::AbstractMatOrFac)
    PQ = similar(P) # this copies the sparsity pattern of P, therefore Q is nevery instantiated densely
    I, J, _ = findnz(P)
     for (i, j) in zip(I, J)
          PQ[i, j] = P[i, j] * Q[i, j]
     end
     return PQ
end

# synthetic data generators
function two_bump_data(d::Int, n::Int)
    X = randn(d, n)
    r = rem(n, 2)
    n = div(n, 2)
    @. X[:, 1:n] = .1X[:, 1:n] + 1
    @. X[:, n+1:end] = .1X[:, n+1:end] - 1
    return X
end

############################# optimization #####################################
function newtons_method(objective, x; tol::Real = 1e-6, max_iter::Int = 32)
    g(x) = ForwardDiff.derivative(objective, x)
    for t in 1:max_iter # Newton's method
        ft, gt = value_derivative(objective, x)
        gt = sign(gt) * maximum(abs, (ft, gt)) # results in abs(ft/gt) ≤ 1 which leads to more robust convergence
        abs(ft) > tol || break
        x -= ft / gt
    end
    return x
end

# auto-diff helper functions to calculate value and gradient with one call
function value_derivative(f, x::Real, diffresults::Val{true} = Val(true))
    r = DiffResults.DiffResult(zero(x), zero(x))
    r = ForwardDiff.derivative!(r, f, x)
    r.value, r.derivs[1]
end
function value_derivative(f, x::Real, diffresults::Val{false})
    f(x), ForwardDiff.derivative(f, x)
end

############################### neighborhoods ##################################
# defaults to euclidean distance
function k_nearest_neighbors(X::AbstractMatrix, k::Int)
    metric = Euclidean()
    tree = BallTree(X, metric; leafsize = 32, reorder = true)
    neighbors, distances = knn(tree, X, k + 1) # plus one because we are putting in the original points as well
    for i in 1:size(X, 2) # delete self-references
         filter!(!=(i), neighbors[i])
    end
    return neighbors
end

# ε = sqrt(-2log(δ)) # based on Gaussian kernel, minimum distance for kernel to be smaller than δ
function in_range_neighbors(X::AbstractMatrix, r::Real)
    reorder = true # reorder copies the data and puts nearby points close in memory, think about using Distance.jl for the rest of the code too
    leafsize = 8 # to optimize
    metric = Euclidean()
    tree = BallTree(X, metric; leafsize = 8, reorder = true)
    neighbors = inrange(tree, X, r) # plus one because we are putting in the original points as well
    for i in 1:size(X, 2) # delete self-references
         filter!(!=(i), neighbors[i])
    end
    return neighbors
end


using Revise
using PyPlot
using Infiltrator
using LinearAlgebra
using Bem2d

"""
    gravityparticularfunctions()

From Pape and Bannerjee 1987
"""
function gravityparticularfunctions(x, y, g, rho, lambda, mu)
    U = zeros(length(x), 2)
    S = zeros(length(x), 3)
    U[:, 1] = @. -lambda * rho * g / (4 * mu * (lambda + mu)) * x * y  # Pape and Banerjee (1987) equation (6a)
    U[:, 2] = @. (rho * g) / (8 * mu * (lambda + mu)) * (lambda * x^2 + (lambda + 2 * mu) * y^2) # Pape and Banerjee (1987) equation (6b)
    S[:, 1] .= 0 # Pape and Banerjee (1987) equation (6c)
    S[:, 2] = @. rho * g * y  # Pape and Banerjee (1987) equation (6d)
    S[:, 3] .= 0 # Pape and Banerjee (1987) equation (6e)
    return U, S
end

"""
    gravitysquareparticular()

Experiments with gravity body force.
"""
function gravitysquareparticular()
    # TODO: Move particular solution to Bem2d.jl
    # TODO: It's strange that the top of the model has to be at zero.  Can we generalize this?
    # TODO: Rule of thumb for choosing precondtioner value (alpha)?
    # TODO: Make a version of gravityparticularfunctions() for BC generation

    close("all")
    alpha = 7e-8 # scalar preconditioner for traction terms
    fontsize = 20
    mu = 3e10
    lambda = 3e10
    nu = 0.25
    rho = 2700
    g = 9.81
    nels = 20
    npts = 100
    L = 1e4
    offset = 10
    x, y = obsgrid(-L+offset, -2*L+offset, L-offset, 0-offset, npts)

    # Define BEM geometry
    els = Elements(Int(1e5))
    x1, y1, x2, y2 = discretizedline(-L, -2*L, L, -2*L, nels) # Bottom
    addelsez!(els, x1, y1, x2, y2, "B")
    x1, y1, x2, y2 = discretizedline(L, -2*L, L, 0, nels) # Right hand side
    addelsez!(els, x1, y1, x2, y2, "R")
    x1, y1, x2, y2 = discretizedline(L, 0, -L, 0, nels) # Top
    addelsez!(els, x1, y1, x2, y2, "T")
    x1, y1, x2, y2 = discretizedline(-L, 0, -L, -2*L, nels) # Left hand side
    addelsez!(els, x1, y1, x2, y2, "L")

    # Common indexing
    idx = getidxdict(els) # Should this return "all" - YES TODO
    bcidxU = idx["B"] # Boundaries with *displacement* BCs
    bcidxT = [idx["R"] ; idx["T"]; idx["L"]] # Boundaries with *traction* BCs
    bcidxall = collect(1:1:els.endidx) # All boundaries

    # Gravity square problem with quadratic elements
    T_pU_qall, _ = PUTQ(slip2dispstress, els, bcidxU, bcidxall, mu, nu)
    _, H_pT_qall = PUTQ(slip2dispstress, els, bcidxT, bcidxall, mu, nu)

    TH = [T_pU_qall ; alpha .* H_pT_qall] # Assemble combined linear operator

    # Particular solution and effective boundary conditions
    xnodes = transpose(els.xnodes[idx["B"], :])[:]
    ynodes = transpose(els.ynodes[idx["B"], :])[:]
    UB, _ = gravityparticularfunctions(xnodes, ynodes, g, rho, lambda, mu)
    bcs = zeros(6 * els.endidx)
    bcs[1:2:6*nels] = UB[:, 1] # Bottom boundary (x-component)
    bcs[2:2:6*nels] = UB[:, 2] # Bottom boundary (y-component)
    bcs *= -1 # This is neccesary for the right answer and is consistent with derivation

    # BEM solve to get particular solution
    Ueffparticular = inv(TH) * bcs

    # Evaluate and plot interior solution
    Uinteriorcomplementary, Sinteriorcomplementary = quaddispstress(slip2dispstress, x, y, els, bcidxall, quadstack(Ueffparticular[1:2:end]), quadstack(Ueffparticular[2:2:end]), mu, nu)
    Uinteriorparticular, Sinteriorparticular = gravityparticularfunctions(x, y, g, rho, lambda, mu)
    U = @. Uinteriorcomplementary + Uinteriorparticular
    S = @. Sinteriorcomplementary + Sinteriorparticular
    plotfields(els, reshape(x, npts, npts), reshape(y, npts, npts), Uinteriorcomplementary, Sinteriorcomplementary, "Complementary solution")
    plotfields(els, reshape(x, npts, npts), reshape(y, npts, npts), Uinteriorparticular, Sinteriorparticular, "Particular solution")
    plotfields(els, reshape(x, npts, npts), reshape(y, npts, npts), U, S, "Complementary + Particular solutions")

    # Interior point solution for nice vector plot
    npts = 10
    offset = 500
    x, y = obsgrid(-L+offset, -2*L+offset, L-offset, 0-offset, npts)
    Uinteriorcomplementary, Sinteriorcomplementary = quaddispstress(slip2dispstress, x, y, els, bcidxall, quadstack(Ueffparticular[1:2:end]), quadstack(Ueffparticular[2:2:end]), mu, nu)
    Uinteriorparticular, Sinteriorparticular = gravityparticularfunctions(x, y, g, rho, lambda, mu)
    U = @. Uinteriorcomplementary + Uinteriorparticular
    S = @. Sinteriorcomplementary + Sinteriorparticular

    # Plot of interior points
    figure()
    fill([-10000, 10000, 10000, -10000], [-20000, -20000, 0, 0] .+ 20000, color="lightgray", zorder=1)
    quiver(x[:], y[:].+20000, U[:, 1], U[:, 2], units="width", color="cyan", edgecolor="k", linewidths=0.25, zorder=10)
    xlabel("x (m)")
    ylabel("y (m)")
    gca().set_aspect("equal")
end
gravitysquareparticular()


module day14

struct Part1 end
struct Part2 end

include("inputs.jl")

ex1 = inputs.ex1_d14()
ex2 = inputs.ex2_d14()
input = inputs.input_d14()

function parse_input(input)
    lines = split(input, "\n")
    memrx = r"mem\[(\d+)\]"
    instructions = []
    for line in lines
        first, last= split(line, " = ")
        if startswith(line, "mem")
            idx = parse(Int, match(memrx, first)[1])
            val = parse(Int, last)
            msk = ""
        elseif startswith(line, "mask")
            msk = lpad(last, 64, '.') # make mask 64 bits for easy comparison
            idx, val = missing, missing
        end
        push!(instructions, (idx=idx, val=val, msk=msk))
    end
    return instructions
end

test_mask = "............................XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X"
@assert parse_input(ex1)[1].msk == test_mask

# Part 1
function apply_mask(p::Part1, mask, val)
    result = Array{Char,1}()
    valbit = bitstring(UInt64(val))
    for bit = 29:64
        if mask[bit] == 'X'
            push!(result, valbit[bit])
        elseif mask[bit] == '1'
            push!(result, '1')
        elseif mask[bit] == '0'
            push!(result, '0')
        end
    end
    parse(Int, String(result), base=2)
end

@assert apply_mask(Part1(), test_mask, 11) == 73
@assert apply_mask(Part1(), test_mask, 101) == 101
@assert apply_mask(Part1(), test_mask, 0) == 64

write!(p::Part1, memory, mask, i) = memory[i.idx] = apply_mask(p, mask, i.val)

function initialize_program(p::Union{Part1,Part2}, input)
    insts = parse_input(input)
    memory = Dict{Int, Int}()
    mask = String(['X' for _ in 1:64])

    for i in insts
        if !isempty(i.msk)
            mask = i.msk
        else
            write!(p, memory, mask, i)
        end
    end
    return sum(values(memory))
end

part1(input) = initialize_program(Part1(), input)
@assert part1(ex1) == 165
@info "Day 14, Part 1 answer: $(part1(input))"

# Part 2
function apply_mask(p::Part2, mask, addr)
    result = Array{Char,1}()
    addrbit = bitstring(UInt64(addr))
    for bit = 29:64
        if mask[bit] == 'X'
            push!(result, 'X')
        elseif mask[bit] == '1'
            push!(result, '1')
        elseif mask[bit] == '0'
            push!(result, addrbit[bit])
        end
    end

    addrs = [result]
    while 'X' ∈ addrs[1]
        x = popfirst!(addrs)
        push!(addrs, replace(x, 'X'=>'0', count=1))
        push!(addrs, replace(x, 'X'=>'1', count=1))
    end
    [parse(Int, String(addr), base=2) for addr in addrs]
end

test_mask2 = lpad("000000000000000000000000000000X1001X", 64, '.')
@assert apply_mask(Part2(), test_mask2, 42) == [26, 27, 58, 59]

function write!(p::Part2, memory, mask, i)
    addrs = apply_mask(p, mask, i.idx)
    [memory[addr] = i.val for addr in addrs]
end

part2(input) = initialize_program(Part2(), input)
@assert part2(ex2) == 208
@info "Day 14, Part 2 answer: $(part2(input))"

end


@testset "product" begin

    # Test the multiplication of a GP by a constant.
    let
        rng, N, N′, D = MersenneTwister(123456), 10, 11, 2
        X, X′ = ColsAreObs(randn(rng, D, N)), ColsAreObs(randn(rng, D, N′))
        g1, c, c′ = GP(ConstantMean(1.0), EQ(), GPC()), -4.3, 2.1
        g2, g2′ = c * g1, g1 * c′
        g3, g3′ = c * g2, g2′ * c′
        g4, g4′ = c * g3, g3′ * c′

        @test mean_vec(g2(X)) == c .* mean_vec(g1(X))
        @test mean_vec(g3(X)) == c .* mean_vec(g2(X))
        @test mean_vec(g4(X)) == c .* mean_vec(g3(X))
        @test mean_vec(g2′(X)) == mean_vec(g1(X)) .* c′
        @test mean_vec(g3′(X)) == mean_vec(g2′(X)) .* c′
        @test mean_vec(g4′(X)) == mean_vec(g3′(X)) .* c′

        @test cov(g2(X)) ≈ c^2 .* cov(g1(X))
        @test cov(g3(X)) ≈ c^2 .* cov(g2(X))
        @test cov(g4(X)) ≈ c^2 .* cov(g3(X))
        @test cov(g2′(X)) ≈ c′^2 .* cov(g1(X))
        @test cov(g3′(X)) ≈ c′^2 .* cov(g2′(X))
        @test cov(g4′(X)) ≈ c′^2 .* cov(g3′(X))

        @test xcov(g2(X), g1(X′)) ≈ c .* xcov(g1(X), g1(X′))
        @test xcov(g3(X), g1(X′)) ≈ c^2 .* xcov(g1(X), g1(X′))
        @test xcov(g4(X), g1(X′)) ≈ c^3 .* xcov(g1(X), g1(X′))
        @test xcov(g2′(X), g1(X′)) ≈ c′ .* xcov(g1(X), g1(X′))
        @test xcov(g3′(X), g1(X′)) ≈ c′^2 .* xcov(g1(X), g1(X′))
        @test xcov(g4′(X), g1(X′)) ≈ c′^3 .* xcov(g1(X), g1(X′))

        @test xcov(g2(X), g2′(X′)) ≈ (c * c′) .* xcov(g1(X), g1(X′))
        @test xcov(g3(X), g3′(X′)) ≈ (c * c′)^2 .* xcov(g1(X), g1(X′))
        @test xcov(g4(X), g4′(X′)) ≈ (c * c′)^3 .* xcov(g1(X), g1(X′))

        @test xcov(g2(X), g3′(X′)) ≈ (c * c′^2) .* xcov(g1(X), g1(X′))
        @test xcov(g3(X), g2′(X′)) ≈ (c^2 * c′) .* xcov(g1(X), g1(X′))
        @test xcov(g2(X), g4′(X′)) ≈ (c * c′^3) .* xcov(g1(X), g1(X′))
        @test xcov(g4(X), g2′(X′)) ≈ (c^3 * c′) .* xcov(g1(X), g1(X′))
    end

    # Test the multiplication of a GP by a known function.
    let
        rng, N, N′, D = MersenneTwister(123456), 10, 11, 2
        X, X′ = ColsAreObs(randn(rng, D, N)), ColsAreObs(randn(rng, D, N′))
        g1, f, f′ = GP(ConstantMean(1.0), EQ(), GPC()), x->sum(sin, x), x->sum(cos, x)
        g2, g2′ = f * g1, g1 * f′
        g3, g3′ = f * g2, g2′ * f′
        g4, g4′ = f * g3, g3′ * f′

        @test mean_vec(g2(X)) == map(f, X) .* mean_vec(g1(X))
        @test mean_vec(g3(X)) == map(f, X) .* mean_vec(g2(X))
        @test mean_vec(g4(X)) == map(f, X) .* mean_vec(g3(X))
        @test mean_vec(g2′(X)) == mean_vec(g1(X)) .* map(f′, X)
        @test mean_vec(g3′(X)) == mean_vec(g2′(X)) .* map(f′, X)
        @test mean_vec(g4′(X)) == mean_vec(g3′(X)) .* map(f′, X)

        fX, f′X = map(f, X), map(f′, X)
        @test cov(g2(X)) ≈ fX .* cov(g1(X)) .* fX'
        @test cov(g3(X)) ≈ fX .* cov(g2(X)) .* fX'
        @test cov(g4(X)) ≈ fX .* cov(g3(X)) .* fX'
        @test cov(g2′(X)) ≈ f′X .* cov(g1(X)) .* f′X'
        @test cov(g3′(X)) ≈ f′X .* cov(g2′(X)) .* f′X'
        @test cov(g4′(X)) ≈ f′X .* cov(g3′(X)) .* f′X'

        fX′, f′X′ = map(f, X′), map(f′, X′)
        @test xcov(g2(X), g1(X′)) ≈ fX .* xcov(g1(X), g1(X′))
        @test xcov(g3(X), g1(X′)) ≈ fX.^2 .* xcov(g1(X), g1(X′))
        @test xcov(g4(X), g1(X′)) ≈ fX.^3 .* xcov(g1(X), g1(X′))
        @test xcov(g2′(X), g1(X′)) ≈ f′X .* xcov(g1(X), g1(X′))
        @test xcov(g3′(X), g1(X′)) ≈ (f′X.^2) .* xcov(g1(X), g1(X′))
        @test xcov(g4′(X), g1(X′)) ≈ (f′X.^3) .* xcov(g1(X), g1(X′))

        @test xcov(g1(X′), g2(X)) == xcov(g2(X), g1(X′))'
        @test xcov(g1(X′), g3(X)) == xcov(g3(X), g1(X′))'
        @test xcov(g1(X′), g4(X)) == xcov(g4(X), g1(X′))'

        @test xcov(g2(X), g2′(X′)) ≈ fX .* xcov(g1(X), g1(X′)) .* f′X′'
        @test xcov(g3(X), g3′(X′)) ≈ fX.^2 .* xcov(g1(X), g1(X′)) .* (f′X′.^2)'
        @test xcov(g4(X), g4′(X′)) ≈ fX.^3 .* xcov(g1(X), g1(X′)) .* (f′X′.^3)'

        @test xcov(g2(X), g3′(X′)) ≈ fX .* xcov(g1(X), g1(X′)) .* (f′X′').^2
        @test xcov(g3(X), g2′(X′)) ≈ fX.^2 .* xcov(g1(X), g1(X′)) .* (f′X′')
        @test xcov(g2(X), g4′(X′)) ≈ fX .* xcov(g1(X), g1(X′)) .* (f′X′').^3
        @test xcov(g4(X), g2′(X′)) ≈ fX.^3 .* xcov(g1(X), g1(X′)) .* (f′X′')
    end
end


# exported hyperbolic functions

over_sch(::Type{Float64}) = 710.0
over_sch(::Type{Float32}) = 89f0

"""
    sinh(x)

Compute hyperbolic sine of `x`.
"""
function sinh(x::T) where {T<:Union{Float32,Float64}}
    u = abs(x)
    d = expk2(Double(u))
    d = dsub(d, drec(d))
    u = T(d) * T(0.5)
    u = abs(x) > over_sch(T) ? T(Inf) : u
    u = isnan(u) ? T(Inf) : u
    u = flipsign(u, x)
    u = isnan(x) ? T(NaN) : u
    return u
end

"""
    cosh(x)

Compute hyperbolic cosine of `x`.
"""
function cosh(x::T) where {T<:Union{Float32,Float64}}
    u = abs(x)
    d = expk2(Double(u))
    d = dadd(d, drec(d))
    u = T(d) * T(0.5)
    u = abs(x) > over_sch(T) ? T(Inf) : u
    u = isnan(u) ? T(Inf) : u
    u = isnan(x) ? T(NaN) : u
    return u
end

over_th(::Type{Float64}) = 18.714973875
over_th(::Type{Float32}) = 18.714973875f0

"""
    tanh(x)

Compute hyperbolic tangent of `x`.
"""
function tanh(x::T) where {T<:Union{Float32,Float64}}
    u = abs(x)
    d = expk2(Double(u))
    e = drec(d)
    d = ddiv(dsub(d, e), dadd(d, e))
    u = T(d)
    u = abs(x) > over_th(T) ? T(1.0) : u
    u = isnan(u) ? T(1) : u
    u = flipsign(u, x)
    u = isnan(x) ? T(NaN) : u
    return u
end

"""
    asinh(x)

Compute the inverse hyperbolic sine of `x`.
"""
function asinh(x::T) where {T<:Union{Float32,Float64}}
    y = abs(x)

    d = y > 1 ? drec(x) : Double(y, T(0.0))
    d = dsqrt(dadd2(dsqu(d), T(1.0)))
    d = y > 1 ? dmul(d, y) : d

    d = logk2(dnormalize(dadd(d, x)))
    y = T(d)

    y = (abs(x) > SQRT_MAX(T) || isnan(y)) ? flipsign(T(Inf), x) : y
    y = isnan(x) ? T(NaN) : y
    y = isnegzero(x) ? T(-0.0) : y

    return y
end

"""
    acosh(x)

Compute the inverse hyperbolic cosine of `x`.
"""
function acosh(x::T) where {T<:Union{Float32,Float64}}
    d = logk2(dadd2(dmul(dsqrt(dadd2(x, T(1.0))), dsqrt(dsub2(x, T(1.0)))), x))
    y = T(d)

    y = (x > SQRT_MAX(T) || isnan(y)) ? T(Inf) : y
    y = x == T(1.0) ? T(0.0) : y
    y = x < T(1.0) ? T(NaN) : y
    y = isnan(x) ? T(NaN) : y

    return y
end

"""
    atanh(x)

Compute the inverse hyperbolic tangent of `x`.
"""
function atanh(x::T) where {T<:Union{Float32,Float64}}
    u = abs(x)
    d = logk2(ddiv(dadd2(T(1.0), u), dsub2(T(1.0), u)))
    u = u > T(1.0) ? T(NaN) : (u == T(1.0) ? T(Inf) : T(d) * T(0.5))

    u = isinf(x) || isnan(u) ? T(NaN) : u
    u = flipsign(u, x)
    u = isnan(x) ? T(NaN) : u

    return u
end


mutable struct P_filt{T<:Real}
	fs::Float64
	freq::Vector{Float64}
	p_conv::Conv.P_conv{T,1,1,1}
end

function P_filt(nt::Int=128, filt_func=x->butterworth(x, order=2, fmin=0.0, fmax=1.0); fs=1)
	# create a plan for conv
	p_conv=Conv.P_conv(gsize=[nt], dsize=[nt], ssize=[nt],)

	freq=rfftfreq(p_conv.np2, fs)
	filt=filt_func(freq)
	for i in eachindex(p_conv.sfreq)
		p_conv.sfreq[i] = filt[i]
	end

	A_mul_B!(p_conv.spad, p_conv.sifftp, p_conv.sfreq)
	Conv.pad_truncate!(p_conv.s, p_conv.spad, p_conv.slags[1],
		    p_conv.slags[2], p_conv.np2, -1)

	fs=Float64(fs)
	return P_filt(fs, Vector(freq), p_conv)

end

function filter!(xx, x, pa::P_filt)

	for ii in 1:size(xx,2)
		xv=view(x, :, ii)
		xxv=view(xx, :, ii)
		Conv.mod!(pa.p_conv, :d, d=xxv, g=xv)
	end
end

"""
A bandpass butterworth filter using `fmin` and `fmax` in the frequency domain.
Return either a zero-phase or minimum-phase filter using `attrib`.
More info on minimum-phase filtering: http://www.katjaas.nl/minimumphase/minimumphase.html.
! then it is converted to minimum phase filter in the cepstral domain
! more info: http://www.katjaas.nl/minimumphase/minimumphase.html
! positive quefrencies correspond to minimum phase component of the signal
! negetive quefrencies correspond to maximum phase component of the signal
! signal = minimum phase [convolved with] maximum phase

# Arguments

* ``

# Keyword Arguments

* `order::Int64` : order of the butterworth filter
* `attrib::Symbol`
  * `=:zp` means zero phase
  * `=:mp` means zero phase
* `fmin::Float64`
  * `=0.` means a lowpass filter at fmax
* `fmax::Float64`
  * `=0.` means a highpass filter at fmin

"""
function butterworth(fvec; order::Int64=2, attrib::Symbol=:zp, fmin=0.0, fmax=0.0)
# Author : Pawan Bharadwaj
#          p.b.pisupati@tudelft.nl
# August 2017, imported to Julia from FORTRAN90
	(fmin < 0.0) && error("fmin cannot be .lt. zero")
	(fmax < 0.0) && error("fmax cannot be .lt. zero")
	(fmax ≠ 0.0) && (fmax <= fmin) && error("fmax .le. fmin")

	if(fmin == 0.0)
		# lowpass filter
		F  = complex.((1.0 + (fvec./fmax).^(2.0*order)).^(-1.0), 0.0)
	elseif(fmax == 0.0)
		# highpass filter
		F  = complex.((1.0 + (fmin./fvec).^(2.0*order)).^(-1.0), 0.0)
	elseif((fmin ≠ 0.0) & (fmax ≠ 0.0))
		# bandpass filter
		x1 = sqrt(fmin*fmax) / (fmax-fmin) * (fvec./sqrt(fmin*fmax) + sqrt(fmax*fmin)./fvec);
		F  = complex.((1.0 + (x1).^(2.0*order)).^(-1.0), 0.0)
	end
	# DC component is always forced zero
	F[1] = complex(0.0, 0.0);

	# conversion to minimum phase
	if(attrib == :mp)
	# to prevent log(0)
		damp = 1e-20 * maximum(abs.(F))
		# logarithm
		X = log.(complex.(abs.(F) + damp, 0.0))
		# to cepstral domain - IFFT
		ifft!(X)
		# only real part output
		X = complex.(real.(X), 0.0);
		# scaling
		# X = X / complex(real(npow2), 0.0)

		# positive cepstrum x 2
		X[2 : npow2/2 + 1] *= complex(2.0, 0.0)
		# remove negative quefrencies
		X[npow2/2 + 2 : npow2] = complex(0.0, 0.0)

		# FFT
		fft!(X)

		# exponential
		F = exp.(X)

		F /= complex(maximum(abs.(F)), 0.0)
	end

	return F

end # butterworth


using Combinatorics
using GraphModularDecomposition
using GraphModularDecomposition: overlap
using IterTools
using LinearAlgebra
using LinearAlgebra: checksquare
using Random
using SparseArrays
using Test

const \ = setdiff

function is_module(G::AbstractMatrix, S::Vector{Int})
    n = checksquare(G)
    all(G[i,k] == G[j,k] && G[k,i] == G[k,j]
        for i in S for j in S for k in 1:n if k ∉ S)
end

function all_modules(G::AbstractMatrix)
    n = checksquare(G)
    filter!(collect(subsets(1:n))) do S
        1 < length(S) && is_module(G, S)
    end
end

function maximal_independent_sets(G::AbstractMatrix)
    n = checksquare(G)
    C(i, j) = G[i,j] == G[j,i] == 0
    filter!(collect(subsets(1:n))) do S
        all(C(i, j) for i in S for j in S) &&
        !any(all(C(i, j) for i in S) for j in 1:n if j ∉ S)
    end
end

function all_strong_modules(G::AbstractMatrix)
    modules = all_modules(G)
    filter(modules) do A
        all(modules) do B
            !overlap(A, B)
        end
    end
end

function is_modular_permutation(G::AbstractMatrix, p::Vector{Int};
    modules = all_strong_modules(G))
    isempty(modules) && return true
    diffs = map(modules) do M
        diff(findall(in(M), p))
    end
    maximum(maximum, diffs) == 1
end

findin_partition(P, S) = sort!(map(x->findfirst(X->x in X, P), S))

function is_modular_partition(G::AbstractMatrix, P::Vector{Vector{Int}};
    modules = all_strong_modules(G))
    sort!(vcat(P...)) == collect(1:size(G,2)) || error("not a partition")
    diffs = map(modules) do M
        diff(findin_partition(P, M))
    end
    maximum(maximum, diffs) == 1
end

is_tournament(G::AbstractMatrix) = all(G + G' + I .== 1)


using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(n::Int)::Int
    perimeters = Dict{Int, Int}()
    for a ∈ 3:n÷3
        for b ∈ a:n÷2
            k = a * a + b * b
            if √k % 1 == 0
                p = a + b + trunc(Int, √k)
                if haskey(perimeters, p)
                    perimeters[p] += 1
                else
                    push!(perimeters, p => 1)
                end
            end
        end
    end
    return findmax(perimeters)[2]
end

compute(1_000)

@benchmark compute(1_000)


using Images, ImageView, ImageFiltering
br = load("bridge.jpg")
imshow(br)
#= You can find misty.jpg and bridge.jpg in the resources folder,
make sure to put it in the same directory you run the code =#

misty = load("misty.jpg")
imshow(misty)
channel_view = channelview(misty); #converts misty to an array that can be modified
channel_view = permuteddimsview(channel_view, (2, 3, 1)); #changes channel locations for convenience: height, width, rgb

x_coords = 1:599 #you can specify any range of x coords here, since width is 1200 this is right half of image now
x_coords_2 = 600:1200

y_coords = 1:799
y_coords_2 = 800:1600

channel_view[y_coords, x_coords, 1] = min.(channel_view[y_coords, x_coords, 1] .* 1.8, 1);
# multiplies red pixels (top left) by 1.8 but restricts their maximum value to 1, using the min function.

channel_view[y_coords, x_coords_2 , 2] = min.(channel_view[y_coords, x_coords_2, 2] .* 1.2, 1);
# multiplies green pixels (top right) by 1.2 but restricts their maximum value to 1, using the min function

channel_view[y_coords_2, x_coords_2, 3] = min.(channel_view[y_coords_2, x_coords_2, 3] .* 1.4, 1);
# multiplies blue pixels (bottom right) in  by 1.4 but restricts their maximum value to 1, using the min function

# bottom left of image is left untouched

imshow(misty)

# individual pixel values can be modified like this to incrase saturation.
# now let's do something even cooler!
# we can create a color filter using the same principle

img = load("bridge.jpg")
gray_img = RGB.(Gray.(img)) # creates a gray img

channel_view = channelview(img);
channel_view_gray = channelview(gray_img);

img_array = permuteddimsview(channel_view, (2, 3, 1)); # creates array of img with height, width, rgb
gray_img_array = permuteddimsview(channel_view_gray, (2, 3, 1)); # creates array of gray img with height, width, grayvalue

mask = fill(true, size(img)); # creates a grid of true values with the size of img

retained_colors_x = 500:508; # the pixels that are in this range are going to be
retained_colors_y = 400:470; # the colors that will stay in the final image (as a range, not specific values)

# img[retained_colors_x, retained_colors_y] # uncomment this to see the selected color range

for channel_id = 1:3

    # views the image array only in specific channel (for loop)
    current_channel = view(img_array, :, :, channel_id)
    # takes the retained color part only in specified channel
    retained_color_area = current_channel[retained_colors_y, retained_colors_x, :]

    #gets minimum and maximum channel values from the retained color area
    channel_min = minimum(retained_color_area)
    channel_max = maximum(retained_color_area)

    # only the pixels from the full image that are in the retained color area return true to the image mask
    channel_mask = channel_min .< current_channel .< channel_max
    mask = mask .& channel_mask
end
# only the pixels that correspond to true in the general mask retain their pixel values
# the rest is replaced with the gray image
img_masked = img_array .* mask .+ gray_img_array .* .~(mask);
# converts the array to an image that we can view (also puts the channels back into place)
final_image = colorview(RGB, permutedims(img_masked, (3, 1, 2)))

imshow(final_image)


module BenchmarkAStarSearchCollatz

using BenchmarkTools
using AStarSearch

neighbours(n) = n % 2 == 0 ? [n ÷ 2, 3n + 1] : [3n + 1]

# The best case is if the number is a power of 2, then the moves are just the log2, otherwise this is a very optimistic heuristic (it doesn't overestimate)
heuristic(n, goal) = floor(Int64, log2(n))

astar_collatz(x; kwargs...) = astar(neighbours, x, 1; heuristic, kwargs...)

suite = BenchmarkGroup()

numbers = [1, 34, 523, 1345, 8765]
for number in numbers
  suite[[string(number)]] = @benchmarkable astar_collatz($number)
end

end

BenchmarkAStarSearchCollatz.suite


"""
    InverseGaussian(μ,λ)

The *inverse Gaussian distribution* with mean `μ` and shape `λ` has probability density function

```math
f(x; \\mu, \\lambda) = \\sqrt{\\frac{\\lambda}{2\\pi x^3}}
\\exp\\!\\left(\\frac{-\\lambda(x-\\mu)^2}{2\\mu^2x}\\right), \\quad x > 0
```

```julia
InverseGaussian()              # Inverse Gaussian distribution with unit mean and unit shape, i.e. InverseGaussian(1, 1)
InverseGaussian(μ),            # Inverse Gaussian distribution with mean μ and unit shape, i.e. InverseGaussian(μ, 1)
InverseGaussian(μ, λ)          # Inverse Gaussian distribution with mean μ and shape λ

params(d)           # Get the parameters, i.e. (μ, λ)
mean(d)             # Get the mean parameter, i.e. μ
shape(d)            # Get the shape parameter, i.e. λ
```

External links

* [Inverse Gaussian distribution on Wikipedia](http://en.wikipedia.org/wiki/Inverse_Gaussian_distribution)

"""
struct InverseGaussian{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    λ::T
    InverseGaussian{T}(μ::T, λ::T) where {T<:Real} = new{T}(μ, λ)
end

function InverseGaussian(μ::T, λ::T; check_args::Bool=true) where {T<:Real}
    @check_args InverseGaussian (μ, μ > zero(μ)) (λ, λ > zero(λ))
    return InverseGaussian{T}(μ, λ)
end

InverseGaussian(μ::Real, λ::Real; check_args::Bool=true) = InverseGaussian(promote(μ, λ)...; check_args=check_args)
InverseGaussian(μ::Integer, λ::Integer; check_args::Bool=true) = InverseGaussian(float(μ), float(λ); check_args=check_args)
InverseGaussian(μ::Real; check_args::Bool=true) = InverseGaussian(μ, one(μ); check_args=check_args)
InverseGaussian() = InverseGaussian{Float64}(1.0, 1.0)

@distr_support InverseGaussian 0.0 Inf

#### Conversions

function convert(::Type{InverseGaussian{T}}, μ::S, λ::S) where {T <: Real, S <: Real}
    InverseGaussian(T(μ), T(λ))
end
function convert(::Type{InverseGaussian{T}}, d::InverseGaussian{S}) where {T <: Real, S <: Real}
    InverseGaussian(T(d.μ), T(d.λ), check_args=false)
end

#### Parameters

shape(d::InverseGaussian) = d.λ
params(d::InverseGaussian) = (d.μ, d.λ)
partype(::InverseGaussian{T}) where {T} = T

#### Statistics

mean(d::InverseGaussian) = d.μ

var(d::InverseGaussian) = d.μ^3 / d.λ

skewness(d::InverseGaussian) = 3sqrt(d.μ / d.λ)

kurtosis(d::InverseGaussian) = 15d.μ / d.λ

function mode(d::InverseGaussian)
    μ, λ = params(d)
    r = μ / λ
    μ * (sqrt(1 + (3r/2)^2) - (3r/2))
end

#### Evaluation

function pdf(d::InverseGaussian{T}, x::Real) where T<:Real
    if x > 0
        μ, λ = params(d)
        return sqrt(λ / (twoπ * x^3)) * exp(-λ * (x - μ)^2 / (2μ^2 * x))
    else
        return zero(T)
    end
end

function logpdf(d::InverseGaussian{T}, x::Real) where T<:Real
    if x > 0
        μ, λ = params(d)
        return (log(λ) - (log2π + 3log(x)) - λ * (x - μ)^2 / (μ^2 * x))/2
    else
        return -T(Inf)
    end
end

function cdf(d::InverseGaussian, x::Real)
    μ, λ = params(d)
    y = max(x, 0)
    u = sqrt(λ / y)
    v = y / μ
    z = normcdf(u * (v - 1)) + exp(2λ / μ) * normcdf(-u * (v + 1))

    # otherwise `NaN` is returned for `+Inf`
    return isinf(x) && x > 0 ? one(z) : z
end

function ccdf(d::InverseGaussian, x::Real)
    μ, λ = params(d)
    y = max(x, 0)
    u = sqrt(λ / y)
    v = y / μ
    z = normccdf(u * (v - 1)) - exp(2λ / μ) * normcdf(-u * (v + 1))

    # otherwise `NaN` is returned for `+Inf`
    return isinf(x) && x > 0 ? zero(z) : z
end

function logcdf(d::InverseGaussian, x::Real)
    μ, λ = params(d)
    y = max(x, 0)
    u = sqrt(λ / y)
    v = y / μ

    a = normlogcdf(u * (v - 1))
    b = 2λ / μ + normlogcdf(-u * (v + 1))
    z = logaddexp(a, b)

    # otherwise `NaN` is returned for `+Inf`
    return isinf(x) && x > 0 ? zero(z) : z
end

function logccdf(d::InverseGaussian, x::Real)
    μ, λ = params(d)
    y = max(x, 0)
    u = sqrt(λ / y)
    v = y / μ

    a = normlogccdf(u * (v - 1))
    b = 2λ / μ + normlogcdf(-u * (v + 1))
    z = logsubexp(a, b)

    # otherwise `NaN` is returned for `+Inf`
    return isinf(x) && x > 0 ? oftype(z, -Inf) : z
end

@quantile_newton InverseGaussian

#### Sampling

# rand method from:
#   John R. Michael, William R. Schucany and Roy W. Haas (1976)
#   Generating Random Variates Using Transformations with Multiple Roots
#   The American Statistician , Vol. 30, No. 2, pp. 88-90
function rand(rng::AbstractRNG, d::InverseGaussian)
    μ, λ = params(d)
    z = randn(rng)
    v = z * z
    w = μ * v
    x1 = μ + μ / (2λ) * (w - sqrt(w * (4λ + w)))
    p1 = μ / (μ + x1)
    u = rand(rng)
    u >= p1 ? μ^2 / x1 : x1
end

#### Fit model

"""
Sufficient statistics for `InverseGaussian`, containing the weighted
sum of observations, the weighted sum of inverse points and sum of weights.
"""
struct InverseGaussianStats <: SufficientStats
    sx::Float64      # (weighted) sum of x
    sinvx::Float64   # (weighted) sum of 1/x
    sw::Float64      # sum of sample weight
end

function suffstats(::Type{<:InverseGaussian}, x::AbstractVector{<:Real})
    sx = sum(x)
    sinvx = sum(inv, x)
    InverseGaussianStats(sx, sinvx, length(x))
end

function suffstats(::Type{<:InverseGaussian}, x::AbstractVector{<:Real}, w::AbstractVector{<:Real})
    n = length(x)
    if length(w) != n
        throw(DimensionMismatch("Inconsistent argument dimensions."))
    end
    T = promote_type(eltype(x), eltype(w))
    sx = zero(T)
    sinvx = zero(T)
    sw = zero(T)
    @inbounds @simd for i in eachindex(x)
        sx += w[i]*x[i]
        sinvx += w[i]/x[i]
        sw += w[i]
    end
    InverseGaussianStats(sx, sinvx, sw)
end

function fit_mle(::Type{<:InverseGaussian}, ss::InverseGaussianStats)
    mu = ss.sx / ss.sw
    invlambda = ss.sinvx / ss.sw  -  inv(mu)
    InverseGaussian(mu, inv(invlambda))
end


using LinearAlgebra
using Test
using ProximalOperators

@testset "LogisticLoss" for T in [Float32, Float64]

y = T[1.0, -1.0, 1.0, -1.0, 1.0]
mu = T(1.5)

f = LogisticLoss(y, mu)

x = T[-1.0, -2.0, 3.0, 2.0, 1.0]

f_x_1 = f(x)

@test typeof(f_x_1) == T

grad_f_x, f_x_2 = gradient(f, x)

f_x_ref = 5.893450123044199
grad_f_x_ref = [-1.0965878679450072, 0.17880438303317633, -0.07113880976635019, 1.3211956169668235, -0.4034121320549927]

@test f_x_1 ≈ f_x_ref
@test f_x_2 ≈ f_x_ref
@test all(grad_f_x .≈ grad_f_x_ref)

z1, f_z1 = prox(f, x)
grad_f_z1, = gradient(f, z1)

@test typeof(f_z1) == T
@test norm((x - z1)./1.0 - grad_f_z1, Inf)/norm(grad_f_z1, Inf) <= 1e-4

z2, f_z2 = prox(f, x, T(2.0))
grad_f_z2, = gradient(f, z2)

@test typeof(f_z2) == T
@test norm((x - z2)./2.0 - grad_f_z2, Inf)/norm(grad_f_z2, Inf) <= 1e-4

end


function test_jump_lp(solver)
    @testset "Linear Programming" begin
        @testset "LP1" begin
            # simple 2 variable, 1 constraint problem
            # min -x
            # st   x + y <= 1   (x + y - 1 ∈ Nonpositives)
            #       x, y >= 0   (x, y ∈ Nonnegatives)

            m = Model(solver=solver)
            @variable(m, x >= 0.0)
            @variable(m, y >= 0.0)
            @objective(m, Min, -x)

            @constraint(m, x + y <= 1)

            JuMP.solve(m)

            @test JuMP.isattached(m)
            @test JuMP.hasvariableresult(m)

            @test JuMP.terminationstatus(m) == MOI.Success
            @test JuMP.primalstatus(m) == MOI.FeasiblePoint

            @test JuMP.resultvalue(x) ≈ 1.0 atol=1e-6
            @test JuMP.resultvalue(y) ≈ 0.0 atol=1e-6
            @test JuMP.objectivevalue(m) ≈ -1.0 atol=1e-6

        end
    end
end


using GAIO

# This example demonstrates how to get a ~2x speedup
# in your code using your CPU's SIMD capabilities.

N = 3
const σ, ρ, β = 10.0, 28.0, 0.4

# By default, GAIO is set up to accept functions of the form
function f(x)
    # Some map, here we use the Lorenz equation
    dx = (
           σ * x[2] -    σ * x[1],
           ρ * x[1] - x[1] * x[3] - x[2],
        x[1] * x[2] -    β * x[3]
    )
    return dx
end

F(x) = rk4_flow_map(f, x)

# Internally, GAIO calls this function on a set of test points
# within the domain. This means many function calls have to be made.
# If your function only uses "basic" instructions, then it is
# possible to simultaneously apply
# Single Instructions to Multiple Data (SIMD).
# This way less total function calls have to be made, saving time.

# To see which instructions are supported, refer to
# https://github.com/eschnett/SIMD.jl.git

# All we need to do is pass :cpu to the BoxMap command.
center, radius = (0,0,25), (30,30,30)
P = BoxPartition(Box(center, radius), (128,128,128))
G = BoxMap(F, P, :cpu)

x = (sqrt(β*(ρ-1)), sqrt(β*(ρ-1)), ρ-1)
W = unstable_set!(G, P[x])

plot(W)


@snn_kw struct PoissonParameter{FT=Float32}
    rate::FT = 1Hz
end

@snn_kw mutable struct Poisson{VFT=Vector{Float32},VBT=Vector{Bool}}
    param::PoissonParameter = PoissonParameter()
    N::Int32 = 100
    randcache::VFT = rand(N)
    fire::VBT = zeros(Bool, N)
    records::Dict = Dict()
end

"""
[Poisson Neuron](https://www.cns.nyu.edu/~david/handouts/poisson.pdf)
"""
Poisson

function integrate!(p::Poisson, param::PoissonParameter, dt::Float32)
    @unpack N, randcache, fire = p
    @unpack rate = param
    prob = rate * dt
    rand!(randcache)
    @inbounds for i = 1:N
        fire[i] = randcache[i] < prob
    end
end


# Helper code

"""
  spdiagm_nonsquare(m, n, args...)

Construct a sparse diagonal matrix from Pairs of vectors and diagonals. Each
vector arg.second will be placed on the arg.first diagonal. By default (if
size=nothing), the matrix is square and its size is inferred from kv, but a
non-square size m×n (padded with zeros as needed) can be specified by passing
m,n as the first arguments.

# Arguments
  - `m::Int64`: First dimension of the output matrix
  - `n::Int64`: Second dimension of the output matrix
  - `args::Tuple{T} where T<:Pair{<:Integer,<:AbstractVector}`

# Outputs

  - sparse matrix of size mxn containing the values in args

"""
function spdiagm_nonsquare(m, n, args...)
    I, J, V = SparseArrays.spdiagm_internal(args...)
    return sparse(I, J, V, m, n)
end

# One-dimensional codes

"""
    nablasq_grid(n, h)

Laplacian matrix on a 1D grid

# Arguments:
    - `n`: Number of points
    - `h`: Aspect ratio

# Outputs:
    - discrete Laplacian matrix
"""
function nablasq_grid(n::Int64, h::Float64 = 1.0)
  o = ones(n) / h
  del = spdiagm_nonsquare(n + 1, n, -1 => -o, 0 => o)
  A1D = del' * del
  A1D[1, 1] = 1.0 / h ^ 2
  A1D[n, n] = 1.0 / h ^ 2
  return A1D
end

"""

   matern_1d_grid(y, idx, m, eps, h)

Matern Interpolation in one dimension

# Arguments:
  - `y`: the vector of y's for which the values are known
  - `idx`: the vector of indices for which values are to be interpolated
  - `m::Int64 = 1`: Matern parameter m
  - `eps = 0.0`: Matern parameter eps
  - `h = 1.0`: aspect ratio

# Outputs:
  - vector of interpolated data

# Example:
```<julia-repl>
x = 1:100
h = x[2] - x[1]
y = sin.(2 * pi * x * 0.2)
discard = randperm(100)[1:50]
# Laplace interpolation
y_lap = matern_1d_grid(y, discard, 1, 0.0, h)
# Matern interpolation
y_mat = matern_1d_grid(y, discard, 2, 0.1, h)
```

"""
function matern_1d_grid(y::Vector{T}, idx::Vector{Int64},
                        m::Int64 = 1, eps = 0.0, h::Float64 = 1.0) where{T<:Number}
  n = length(y)
  A1D = nablasq_grid(n, h)
  C = sparse(I, n, n)
  for i in idx
    C[i, i] = 0.0
  end
  if ((eps == 0)||(eps == 0.0)) && (m == 1)
    # Laplace Interpolation
    return ((C - (sparse(I, n, n) - C) * A1D)) \ (C * y)
  else
    # Matern Interpolation
    for i = 1:size(A1D,1)
      A1D[i, i] = A1D[i, i] + eps^2
    end
    A1DM = A1D ^ m
    return ((C - (sparse(I, n, n) - C) * A1DM)) \ (C * y)
  end
end

# Two dimensional codes

"""
  bdy_nodes(Nx, Ny)

...
# Arguments

  - `Nx::Int64`: the number of points in the first dimension
  - `Ny::Int64`: the number of points in the second dimension
...

...
# Outputs
  - vector containing the indices of coordinates on the boundary of the 2D rectangle
...

"""
function bdy_nodes(Nx, Ny)
  bdy = []
  xnb = []
  ynb = []
  counter = 0
      for j = 1:Ny
          for i = 1:Nx
              counter += 1
              if (j == 1 || j == Ny || i == 1 || i == Nx)
                  bdy = push!(bdy, counter)
                  if (j == 1 || j == Ny)
                      push!(ynb, 1)
                  else
                      push!(ynb, 2)
                  end
                  if (i == 1 || i == Nx)
                      push!(xnb, 1)
                  else
                      push!(xnb, 2)
                  end
              end
          end
      end
  return bdy, xnb, ynb
end

"""
    nablasq_grid(Nx, Ny, h, k)

Laplacian matrix on a 2D grid

# Arguments:
    - `Nx::Int64`: Number of points in first dimension
    - `Ny::Int64`: Number of points in second dimension
    - `h::Float64`: Aspect ratio in first dimension
    - `k::Float64`: Aspect ratio in second dimension

# Outputs:
    - discrete Laplacian matrix in 2D
"""
function nablasq_grid(Nx::Int64, Ny::Int64, h::Float64, k::Float64)
  o1 = ones(Nx) / h
  del1 = spdiagm_nonsquare(Nx + 1, Nx, -1 => -o1, 0 => o1)
  o2 = ones(Ny) / k
  del2 = spdiagm_nonsquare(Ny + 1, Ny, -1 => -o2,0 => o2)
  A2D = (kron(sparse(I, Ny, Ny), del1' * del1) +
          kron(del2' * del2, sparse(I, Nx, Nx)))
  bdy, xnb, ynb = bdy_nodes(Nx, Ny)
  count = 1
  for i in bdy
      A2D[i, i] = 0.0
      A2D[i, i] = A2D[i, i] + xnb[count] / h ^ 2 + ynb[count] / k ^ 2
      count += 1
  end
  return A2D
end

"""

  matern_2d_grid(mat, discard, m, eps, h, k)

...
# Arguments
  - `mat`: the matrix containing the image
  - `idx`: the linear indices of the nodes to be discarded
  - `eps`: Matern parameter eps
  - `m`: The Matern exponent (integer)
  - `h`: The aspect ratio in the first dimension
  - `k`: The aspect ratio in the second dimension

# Outputs
  - matrix containing the interpolated image

# Example:

```<julia-repl>
x = y = 1:30
h = k = x[2] - x[1]
y = sin.(2 * pi * x * 0.2) * cos.(2 * pi * y * 0.3)
discard = randperm(900)[1:450]
# Laplace interpolation
y_lap = matern_2d_grid(y, discard, 1, 0.0, h, k)
# Matern interpolation
y_mat = matern_2d_grid(y, discard, 2, 0.1, h, k)
...

"""
function matern_2d_grid(mat::Matrix, discard::Vector, m::Int64 = 1, eps = 0.0,
                        h::Float64 = 1.0, k::Float64 = 1.0)
    rows, columns = size(mat)
    A2D = nablasq_grid(rows, columns, h, k)
    sizeA = size(A2D, 1)
    C = sparse(I, sizeA, sizeA)
    for i in discard
        C[i, i] = 0.0
    end
    if ((m == 1)||(m == 1.0)) && (eps == 0.0)
        # Laplace interpolation
        u = ((C - (sparse(I, sizeA, sizeA) - C) * A2D)) \ (C * mat[:])
        return reshape(u, rows, columns)
    else
        # Matern interpolation
        for i = 1:sizeA
            A2D[i,i] = A2D[i,i] + eps^2
        end
        A2D = A2D^m
        u = ((C - (sparse(I, sizeA, sizeA) - C) * A2D)) \ (C * mat[:])
        return reshape(u, rows, columns)
    end
end