#######################################################################
# functions to makle lags and leads
#######################################################################

"""
    F!(df, v, horizons=1, fun=identity)

Forward operator.
Applies function `fun` to the column `v` of `df` and then makes the
leads specified by horizons and adds them to `df` as the column `F(i)(v)`.
Modifies `df`.
Return the list of new columns created.
"""
function F!(df, v, horizons=1, fun=identity)
    # Forward operator
    newvars = Symbol[]
    for i in horizons
        Fiv = Symbol("F$(i)$(v)")
        df[:, Fiv] = lead( map(fun, df[!, v]), i)
        push!(newvars, Fiv)
    end
    return newvars
end

"""
    L!(df, v, horizons=1, fun=identity)

Lag operator.
Applies function `fun` to the column `v` of `df` and then makes the
laggs specified by horizons and adds them to `df` as the column `L(i)(v)`.
Modifies `df`.
Return the list of new columns created.
"""
function L!(df, v, horizons=1, fun=identity)
    # Lag operator
    newvars = Symbol[]
    for i in horizons
        Fiv = Symbol("L$(i)$(v)")
        df[:, Fiv] = lag(map(fun, df[!, v]), i)
        push!(newvars, Fiv)
    end
    return newvars
end

"""
    F!(df, v, horizons=1, fun=identity, baseperiod=0, relativebase=false)

Cumlative forward difference operator
Applies function `fun` to the column `v` of `df` and then calculates
the growth of this column, saving the results to a new columns in
`df` as the column `FΔ(i)(v)`.
If `x = map(fun, df[!, v])`, `df.FΔ(i)(v) = lead(x, horizons) - lead(x, baseperiod)`
or `x = map(fun, df[!, v])`, `df.FΔ(i)(v) = lead(x, horizons) - lead(x, horizon + baseperiod)`
if relativebase == true
Modifies `df`.
Return the list of new columns created.
"""
function FΔ!(df, v, horizons=1, fun=identity, baseperiod=0, relativebase=false)
    # Forward difference operator
    newvars = Symbol[]
    for i in horizons
        Fiv = Symbol("FΔ$(i)$(v)")
        if relativebase
            df[:, Fiv] = lead( map(fun, df[!, v]), i) .- lead(map(fun, df[!, v]), i+baseperiod)
        else
            df[:, Fiv] = lead( map(fun, df[!, v]), i) .- lead(map(fun, df[!, v]), baseperiod)
        end
        push!(newvars, Fiv)
    end
    return newvars
end

"""
    LΔ!(df, v, horizons=1, fun=identity, baseperiod=0, relativebase=false)
Cumlative lag difference operator

Applies function `fun` to the column `v` of `df` and then calculates
the lagged growth of this column, saving the results to a new columns in
`df` as the column `LΔ(i)(v)`.
If `x = map(fun, df[!, v])`, `df.FΔ(i)(v) = lag(x, baseperiod) - lag(x, horizons)`
or `x = map(fun, df[!, v])`, `df.FΔ(i)(v) = lag(x, i+baseperiod) - lag(x, horizons)`
if relativebase == true
Modifies `df`.
Return the list of new columns created.
"""
function LΔ!(df, v, horizons=1, fun=identity, baseperiod=0, relativebase=false)
    # Lag difference operator
    newvars = Symbol[]
    for i in horizons
        Fiv = Symbol("LΔ$(i)$(v)")
        if relativebase
            df[:, Fiv] = lag(map(fun, df[!, v]), i+baseperiod) .-  lag( map(fun, df[!, v]), i)
        else
            df[:, Fiv] = lag(map(fun, df[!, v]), i+baseperiod) .-  lag( map(fun, df[!, v]), i)
        end
        push!(newvars, Fiv)
    end
    return newvars
end

"""
    Δ(x, horizons=1, fun=identity)

Single period growth function
    `map(fun, x) .-  lag(map(fun, x), horizons)`
"""
function Δ(x, horizons=1, fun=identity)
    return map(fun, x) .-  lag(map(fun, x), horizons)
end


using PkgSRA
using Plots, Random, Distributions, Interpolations
pyplot()
Random.seed!(11)
using BSON: @save
include("../scripts/paper_settings.jl");
include("../../utils/sindy_turing_utils.jl")
# include("../../utils/sindy_statistics_utils.jl")
# include("../../utils/main_algorithm_utils.jl")

################################################################################
#####
##### Rossler example
#####
include(EXAMPLE_FOLDERNAME*"example_rossler.jl")

# Define the multivariate forcing function
num_ctr = 10;
    U_starts = rand(3, num_ctr) .* tspan[2]
    U_widths = [0.0, 0.0, 0.5];
    amplitudes = [1.0, 1.0, 1000.0]
my_U_func_time2(t) = U_func_time_multivariate(t, u0,
                        U_widths, U_starts,
                        F_dim=[1,2,3],
                        amplitudes=amplitudes)

# Get data
sol = solve_rossler_system(U_func_time=my_U_func_time2)
dat = Array(sol)
plot3d(dat[1,:],dat[2,:],dat[3,:],label="Data")
# plot(ts,dat[1,:]);plot!(ts,dat[2,:]);plot!(ts,dat[3,:])
# plot(ts,dat[3,:])
numerical_grad = numerical_derivative(dat, ts)

## Also get baseline true/ideal cases
# Uncontrolled
dat_raw = Array(solve_rossler_system())
numerical_grad_raw = numerical_derivative(dat_raw, ts)
true_grad_raw = zeros(size(dat))
# for i in 1:size(dat,2)
#     true_grad_raw[:,i] = rossler_system(dat_raw[:,i], Float64.(p), [0])
# end

# True control signal
U_true = zeros(size(dat))
for (i, t) in enumerate(ts)
    U_true[:,i] = my_U_func_time2(t)
end
plot(ts,dat[1,:]);plot!(ts,dat[2,:]);plot!(ts,dat[3,:])
    plot!(ts,U_true')

###
### Fit a naive model
###
sindy_library = Dict("cross_terms"=>2,"constant"=>nothing);

# Upgrade: Use AIC
# val_list = calc_permutations(5,3)
# val_list = combinations(1:5, 3)
val_list = Iterators.product(1:3,1:3,1:3) # DEBUG
(sindy_unctr,best_criterion,all_criteria,all_models) =
    sindyc_ensemble(dat, numerical_grad, sindy_library, val_list,
                    selection_criterion=my_aicc,
                    sparsification_mode="num_terms",
                    selection_dist=Normal(0.0,20),
                    use_clustering_minimization=true)
print_equations(sindy_unctr)

# Integrate this (poor) model
# Callback aborts if it blows up
condition(u,t,integrator) = any(abs.(u).>1e4)
cb = DiscreteCallback(condition, terminate!)
prob_unctr = ODEProblem(sindy_unctr, u0, tspan, [0], callback=cb)
sindy_dat_unctr = Array(solve(prob_unctr, Tsit5(), saveat=ts));

let d = sindy_dat_unctr, d2 = dat
    plot(d[1,:]);plot!(d[2,:]);plot!(d[3,:])
    plot!(d2[1,:]);plot!(d2[2,:]);plot!(d2[3,:])
end
dat_unctr = sindy_dat_unctr

###
### Calculate the residuals
###
# NOTE: do not use Turing, for speed
sindy_grad1 = sindy_unctr(dat, 0)
residual = numerical_grad .- sindy_grad1
noise_guess = 10*abs(median(residual[3,:]))

scatter(residual[3,:])
    hline!([-noise_guess noise_guess])

ctr_guess = process_residual(residual, noise_guess)

###
### Subsample based on residual
###
accepted_ind = subsample_using_residual(residual,
            noise_guess, min_length=4, shorten_length=0)

num_pts = 1000
start_ind = 1
subsample_ind = accepted_ind[start_ind:num_pts+start_ind-1]
dat_sub = dat[:,subsample_ind]
grad_sub = numerical_grad[:,subsample_ind]

# Plot the subsampled points
i = 3;
    plot(dat[i,:], label="Data; NOT residual")
    scatter!(subsample_ind, dat_sub[i,:], color=:blue, label="Subsampled points")
    plot!(U_true[i,:], label="Control Signal")

# Better SINDy model
# val_list = calc_permutations(7,3)
# val_list = combinations(1:5, 3)
val_list = Iterators.product(1:3,1:3,1:3) # DEBUG
(sindy_sub,best_criterion,all_criteria,all_models,best_index) =
    sindyc_ensemble(dat_sub, grad_sub, sindy_library, val_list,
                    selection_criterion=my_aicc,
                    sparsification_mode="num_terms",
                    selection_dist=Normal(0.0,0.1*noise_guess),
                    use_clustering_minimization=false)
println("Best index is $best_index")
print_equations(sindy_sub)
print("True equations")
print_equations(core_dyn_true)
scatter(sum.(val_list), all_criteria)
    title!("AIC for various sparsities")
    xlabel!("Number of nonzero terms")

# Generate the attractor using this system
condition(u,t,integrator) = any(abs.(u).>1e3)
cb = DiscreteCallback(condition, terminate!)
prob_ctr = ODEProblem(sindy_sub, u0, tspan, [0], callback=cb)
sindy_dat_ctr = solve(prob_ctr, Tsit5(), saveat=ts);
dat_ctr = Array(sindy_dat_ctr)

plot3d(sindy_dat_ctr[1,:], sindy_dat_ctr[2,:], sindy_dat_ctr[3,:])

# Also generate the derivatives
sindy_grad_ctr = sindy_sub(dat, 0)

# And then the guessed control signal

ctr_guess2 = process_residual(numerical_grad.-sindy_grad_ctr, noise_guess)

#####
##### Save Rossler data
#####
this_dat_name = DAT_FOLDERNAME*"TMP_dat_library_of_examples_rossler"

# Uncontrolled data
fname = this_dat_name*"uncontrolled.bson"
@save fname dat_raw numerical_grad_raw

# Original ODE and controller variables
fname = this_dat_name*"ode_vars.bson";
@save fname dat U_true #dat_grad

# Bayesian variables 1: before control
fname = this_dat_name*"naive_vars_bayes.bson";
@save fname dat_unctr noise_guess residual accepted_ind ctr_guess

# SINDY variables 2: after control
fname = this_dat_name*"ctr_vars_sindy.bson";
@save fname dat_ctr sindy_grad_ctr ctr_guess2

# Bayesian variables 2: after control
# fname = this_dat_name*"ctr_vars_turing.bson";
# # @save fname dat_ctr sample_trajectory_noise_ctr sample_gradients_ctr turing_dat_ctr ctr_guess2
# @save fname dat_ctr


function rel_err(sol, ref)
    maximum(abs.((sol.d[:,end] .- ref) ./ ref))
end


# ---
# title: 1037. Valid Boomerang
# id: problem1037
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Math
# link: <https://leetcode.com/problems/valid-boomerang/description/>
# hidden: true
# ---
#
# A _boomerang_ is a set of 3 points that are all distinct and **not** in a
# straight line.
#
# Given a list of three points in the plane, return whether these points are a
# boomerang.
#
#
#
# **Example 1:**
#
#
#
#     Input: [[1,1],[2,3],[3,2]]
#     Output: true
#
#
# **Example 2:**
#
#
#
#     Input: [[1,1],[2,2],[3,3]]
#     Output: false
#
#
#
# **Note:**
#
#   1. `points.length == 3`
#   2. `points[i].length == 2`
#   3. `0 <= points[i][j] <= 100`
#
#
#
#
## @lc code=start
using LeetCode

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


using Decimals
using Base.Test

# Number/string to decimal conversions
@test decimal("0.01") == decimal(0.01) == Decimal(0, 1, -2)
@test decimal(".001") == decimal(.001) == Decimal(0, 1, -3)
@test decimal("15.23") == decimal(15.23) == Decimal(0, 1523, -2)
@test decimal("543") == decimal(543) == Decimal(0, 543, 0)
@test decimal("-345") == decimal(-345) == Decimal(1, 345, 0)
@test decimal("000123") == decimal(000123) == Decimal(0, 123, 0)
@test decimal("-00032") == decimal(-00032) == Decimal(1, 32, 0)
@test decimal("200100") == decimal(200100) == Decimal(0, 2001, 2)
@test decimal("-.123") == decimal(-.123) == Decimal(1, 123, -3)
@test decimal("1.23000") == decimal(1.23000) == Decimal(0, 123, -2)
@test decimal("4734.612") == decimal(4734.612) == Decimal(0, 4734612, -3)
@test decimal("541724.2") == decimal(541724.2) == Decimal(0,5417242,-1)
@test decimal("2.5e6") == decimal(2.5e6) == Decimal(0, 25, 5)
@test decimal("2.385350e8") == decimal(2.385350e8) == Decimal(0, 238535, 3)
@test decimal("12.3e-4") == decimal(12.3e-4) == Decimal(0, 123, -5)

@test decimal("-12.3e4") == decimal(-12.3e4) == Decimal(1, 123, 3)

@test decimal("-12.3e-4") == decimal(-12.3e-4) == Decimal(1, 123, -5)

@test decimal("0.1234567891") == decimal(0.1234567891) == Decimal(0,1234567891, -10)
@test decimal("0.12345678912") == decimal(0.12345678912) == Decimal(0,12345678912, -11)

# Array-to-decimal conversion
@test decimal([0.1 0.2 0.3]) == [decimal(0.1) decimal(0.2) decimal(0.3)]

# Decimal-to-string conversions
@test string(Decimal(0, 1, -2)) == "0.01"
@test string(Decimal(0, 1, -3)) == "0.001"
@test string(Decimal(0, 1523, -2)) == "15.23"
@test string(Decimal(0, 543, 0)) == "543"
@test string(Decimal(1, 345, 0)) == "-345"
@test string(Decimal(0, 123, 0)) == "123"
@test string(Decimal(1, 32, 0)) == "-32"
@test string(Decimal(0, 2001, 2)) == "200100"
@test string(Decimal(1, 123, -3)) == "-0.123"
@test string(Decimal(0, 123, -2)) == "1.23"

# Decimal-to-number conversions
@test float(Decimal(0, 1, -2)) == 0.01
@test float(Decimal(0, 1, -3)) == 0.001
@test float(Decimal(0, 1523, -2)) == 15.23
@test float(Decimal(0, 543, 0)) == 543
@test float(Decimal(1, 345, 0)) == -345
@test float(Decimal(0, 123, 0)) == 123
@test float(Decimal(1, 32, 0)) == -32
@test float(Decimal(0, 2001, 2)) == 200100
@test float(Decimal(1, 123, -3)) == -0.123
@test float(Decimal(0, 123, -2)) == 1.23
@test number(Decimal(0, 1, -2)) == 0.01
@test number(Decimal(0, 1, -3)) == 0.001
@test number(Decimal(0, 1523, -2)) == 15.23
@test number(Decimal(0, 543, 0)) == 543
@test number(Decimal(1, 345, 0)) == -345
@test number(Decimal(0, 123, 0)) == 123
@test number(Decimal(1, 32, 0)) == -32
@test number(Decimal(0, 2001, 2)) == 200100
@test number(Decimal(1, 123, -3)) == -0.123
@test number(Decimal(0, 123, -2)) == 1.23
@test string(float(Decimal(0, 543, 0))) == "543.0"
@test string(number(Decimal(0, 543, 0))) == "543"
@test string(number(Decimal(0, 543, -1))) == "54.3"


struct AdaGradIterable{T, F, S, R}
    w0::T
    f::F
    stepsize::S
    epsilon::R
end

AdaGradIterable(w0, f; stepsize=1e-3, epsilon=1e-8) = AdaGradIterable(w0, f, to_iterator(stepsize), epsilon)

Base.IteratorSize(::Type{<:AdaGradIterable}) = Base.IsInfinite()

mutable struct AdaGradState{T, G, S}
    w::T
    f_w
    grad_f_w::G
    stepsize_iterator::S
    sum_squared_grad::G
end

function Base.iterate(iter::AdaGradIterable)
    w = copy(iter.w0)
    grad_f_w, f_w = gradient(iter.f, w)
    state = AdaGradState(w, f_w, grad_f_w, Iterators.Stateful(iter.stepsize), zero(w))
    return IterationOutput(state.w, state.f_w, state.grad_f_w), state
end

function Base.iterate(iter::AdaGradIterable, state::AdaGradState)
    state.sum_squared_grad .+= state.grad_f_w .^ 2
    stepsize = popfirst!(state.stepsize_iterator)
    state.w .-= stepsize .* state.grad_f_w ./ (sqrt.(state.sum_squared_grad) .+ iter.epsilon)
    state.grad_f_w, state.f_w = gradient(iter.f, state.w)
    return IterationOutput(state.w, state.f_w, state.grad_f_w), state
end

struct AdaGrad
    kwargs
    AdaGrad(; kwargs...) = new(kwargs)
end

(alg::AdaGrad)(args...) = AdaGradIterable(args...; alg.kwargs...)


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ 8183b318-5ebb-11eb-1cd8-a96e8704a378
using Pkg, DrWatson

# ╔═╡ 686dac30-5ebb-11eb-00f1-434980dba906
begin
	@quickactivate "StatsModelComparisons"
	using StatsModelComparisons
	using StanSample, StatsFuns, StatsPlots
	using DataFrames, CSV, JSON
end

# ╔═╡ 923212a8-630d-11eb-390f-75d21be80011
begin
	ProjDir = joinpath(mc_path, "..", "examples", "arsenic")
	include(joinpath(ProjDir, "cvit.jl"))
end;

# ╔═╡ ca0d916e-5ebe-11eb-3af1-0993bf5b82c8
md" ### Psis_loo."

# ╔═╡ 2ee71b16-5ebd-11eb-3a7d-4553f2f22274
md" ##### Data."

# ╔═╡ 37f06bc2-5ebd-11eb-31f6-65bfa74ea9cf
begin
	data = JSON.parsefile(joinpath(ProjDir, "wells.data.json"))
	y = Float64.(data["switched"])
	x = Float64[data["arsenic"]  data["dist"]]
	n, m = size(x)
end

# ╔═╡ 37f0b87a-5ebd-11eb-092a-3d9b9e8e427b
md" ##### Model."

# ╔═╡ 37f14af6-5ebd-11eb-0aae-d57e8b3d951a
begin
	model_str = read(open(joinpath(ProjDir, "arsenic_logistic.stan")), String)
	tmpdir = joinpath(ProjDir, "tmp")
	sm1 = SampleModel("arsenic_logistic", model_str; tmpdir)
	data1 = (p = m, N = n, y = Int.(y), x = x)
	rc1 = stan_sample(sm1; data=data1)
end

# ╔═╡ 38128be4-5ebd-11eb-36da-a96a68ca1316
md" ##### Fit the model in Stan."

# ╔═╡ 38135bc8-5ebd-11eb-0644-511bb08b3e59
if success(rc1)
	nt1 = read_samples(sm1)

	# Compute LOO and standard error

	log_lik1 = nt1.log_lik'
	loo1, loos1, pk1 = psisloo(log_lik1)
	elpd_loo1 = sum(loos1)
	se_elpd_loo1 = std(loos1) * sqrt(n)
	(elpd_loo = elpd_loo1, se_elpd_loo = se_elpd_loo1)
end

# ╔═╡ 44457a20-5ebd-11eb-00a1-855bfbb45e22
if success(rc1)

	# Check the shape parameter k of the generalized Pareto distribution

	pk_qualify(pk1)
end

# ╔═╡ 5a76e4aa-5ebd-11eb-2e15-6d5808ead825
md" ##### Fit a second model, using log(arsenic) instead of arsenic."

# ╔═╡ 5a771bfa-5ebd-11eb-1431-a5bc45055a5a
begin
	x2 = Float64[log.(data["arsenic"])  data["dist"]]
	data2 = (p = m, N = n, y = Int.(y), x = x2)
	rc2 = stan_sample(sm1; data=data2)
end;

# ╔═╡ 5a77d188-5ebd-11eb-0090-8d80d495f1b1
if success(rc2)
    nt2 = read_samples(sm1)

    # Compute LOO and standard error

    log_lik2 = nt2.log_lik'
    loo2, loos2, pk2 = psisloo(log_lik2)
    elpd_loo2 = sum(loos2)
    se_elpd_loo2 = std(loos2) * sqrt(n)
    (elpd_loo = elpd_loo2, se_elpd_loo = se_elpd_loo2)
end

# ╔═╡ b88657e2-5ebd-11eb-343e-9d666eef86cc
if success(rc2)

	# Check the shape parameter k of the generalized Pareto distribution

	pk_qualify(pk2)
end

# ╔═╡ dd28d430-5ebd-11eb-1854-ab4c13e82c34
if success(rc1) && success(rc2)

    ## Compare the models

    loodiff = loos1 - loos2
    (elpd_diff = sum(loodiff), se_elpd_diff = std(loodiff) * sqrt(n))
end

# ╔═╡ 98d7208e-5ebd-11eb-3a64-e568c1ceebc0
md" ### k-fold-CV."

# ╔═╡ e86ee35c-5ebd-11eb-1932-ffbde638f11e
md" ##### k-fold-CV should be used if several pk>0.5, in this case it is not needed, but provided as an example."

# ╔═╡ e86f0ab2-5ebd-11eb-1a0b-6958db3a7435
begin
	model_str_2 = read(open(joinpath(ProjDir, "arsenic_logistic_t.stan")), String)
	sm3 = SampleModel("arsenic_logistic_t", model_str_2)
end;

# ╔═╡ d83916ba-6309-11eb-309b-474efacd9f85
md" ##### Be patient... 10 cross Validations."

# ╔═╡ e70f73aa-5eb8-11eb-1960-bf6731681898
begin
	cvitr, cvitst = cvit(n, 10, true)
	kfcvs = similar(loos1)
	for cvi in 1:10
		standatacv = (p = m, N = length(cvitr[cvi]), Nt = length(cvitst[cvi]),
			x = x[cvitr[cvi],:], y = Int.(y[cvitr[cvi]]),
			xt = x[cvitst[cvi],:], yt = Int.(y[cvitst[cvi]]))

		# Fit the model in Stan
		rc3 = stan_sample(sm3; data=standatacv)
		if success(rc3)
			nt3 = read_samples(sm3)
			# Compute LOO and standard error
			log_likt = nt3.log_likt'
        	local n_sam, n_obs = size(log_likt)
        	kfcvs[cvitst[cvi]] .=
            	reshape(logsumexp(log_likt .- log(n_sam), dims=1), n_obs)
		end
	end
end

# ╔═╡ 1bf8d5a2-5ebe-11eb-12b3-2501e2f9d288
begin

	# compare PSIS-LOO and k-fold-CV

	plot([-3.5, 0], [-3.5, 0], color=:red)
	scatter!(loos1, kfcvs, xlab = "PSIS-LOO", ylab = "10-fold-CV",
		leg=false, color=:darkblue)
end

# ╔═╡ Cell order:
# ╠═8183b318-5ebb-11eb-1cd8-a96e8704a378
# ╠═686dac30-5ebb-11eb-00f1-434980dba906
# ╠═923212a8-630d-11eb-390f-75d21be80011
# ╟─ca0d916e-5ebe-11eb-3af1-0993bf5b82c8
# ╟─2ee71b16-5ebd-11eb-3a7d-4553f2f22274
# ╠═37f06bc2-5ebd-11eb-31f6-65bfa74ea9cf
# ╟─37f0b87a-5ebd-11eb-092a-3d9b9e8e427b
# ╠═37f14af6-5ebd-11eb-0aae-d57e8b3d951a
# ╟─38128be4-5ebd-11eb-36da-a96a68ca1316
# ╠═38135bc8-5ebd-11eb-0644-511bb08b3e59
# ╠═44457a20-5ebd-11eb-00a1-855bfbb45e22
# ╟─5a76e4aa-5ebd-11eb-2e15-6d5808ead825
# ╠═5a771bfa-5ebd-11eb-1431-a5bc45055a5a
# ╠═5a77d188-5ebd-11eb-0090-8d80d495f1b1
# ╠═b88657e2-5ebd-11eb-343e-9d666eef86cc
# ╠═dd28d430-5ebd-11eb-1854-ab4c13e82c34
# ╟─98d7208e-5ebd-11eb-3a64-e568c1ceebc0
# ╟─e86ee35c-5ebd-11eb-1932-ffbde638f11e
# ╠═e86f0ab2-5ebd-11eb-1a0b-6958db3a7435
# ╟─d83916ba-6309-11eb-309b-474efacd9f85
# ╠═e70f73aa-5eb8-11eb-1960-bf6731681898
# ╠═1bf8d5a2-5ebe-11eb-12b3-2501e2f9d288


@testset "Basis" begin
    @variables u[1:3]
    @parameters w[1:2] t
    h = [u[1]; u[2]; cos(w[1]*u[2]+w[2]*u[3]); u[3]+u[2]]
    h_not_unique = [u[1]; u[1]; u[1]^1; h]
    basis = Basis(h_not_unique, u, parameters = w, iv = t)

    @test isequal(variables(basis), u)
    @test isequal(parameters(basis), w)
    @test isequal(independent_variable(basis), t)
    @test free_parameters(basis) == 6
    @test free_parameters(basis, operations = [+, cos]) == 7
    @test DataDrivenDiffEq.count_operation((ModelingToolkit.Constant(1) + cos(u[2])*sin(u[1]))^3, [+, cos, ^, *]) == 4

    basis_2 = unique(basis)
    @test isequal(basis, basis_2)
    @test size(basis) == size(h)
    @test basis([1.0; 2.0; π], [0. 1.]) ≈ [1.0; 2.0; -1.0; π+2.0]
    @test size(basis) == size(basis_2)
    push!(basis_2, sin(u[2]))
    @test size(basis_2)[1] == length(h)+1

    basis_3 = merge(basis, basis_2)
    @test size(basis_3) == size(basis_2)
    @test isequal(variables(basis_3), variables(basis_2))
    @test isequal(parameters(basis_3), parameters(basis_2))

    merge!(basis_3, basis)
    @test basis_3 == basis_2

    push!(basis, u[3]+u[2])
    unique!(basis)
    @test size(basis) == size(h)

    @variables a
    g = [u[1]; u[3]; a]
    basis = Basis(g, [u; a])
    @test basis([1; 2; 3; 4]) == [1; 3; 4]
    g = [u[1]; u[3]; u[2]]
    basis = Basis(g, u, parameters = [])
    X = ones(Float64, 3, 10)
    X[1, :] .= 3*X[1, :]
    X[3, :] .= 5*X[3, :]
    # Check the array evaluation
    @test basis(X) ≈ [1.0 0.0 0.0; 0.0 0.0 1.0; 0.0 1.0 0.0] * X
    Y = similar(X)
    basis(Y, X)
    @test Y ≈ [1.0 0.0 0.0; 0.0 0.0 1.0; 0.0 1.0 0.0] * X
    f = jacobian(basis)
    @test f([1;1;1], [0.0 0.0], 0.0) ≈ [1.0 0.0 0.0; 0.0 0.0 1.0; 0.0 1.0 0.0]
    @test_nowarn sys = ODESystem(basis)
    @test_nowarn [xi for xi in basis]
    @test_nowarn basis[2:end]; basis[2]; first(basis); last(basis); basis[:]

    @variables u[1:3] t

    g = [u[2]; -sin(u[1])*exp(-t); u[2]+u[3]]
    basis = Basis(g, u, iv = t)
    @test_nowarn ODESystem(basis)

    f(u, p, t) = [u[3]; u[2]*u[1]; p[1]*sin(u[1])*u[2]; p[2]*t]
    b = Basis(f, u, parameters = w, iv = t)
    @test f([1; 2; 3], [2; 0], 3.0) ≈ b([1; 2; 3], [2; 0], 3.0)
    @test_throws AssertionError ODESystem(b)
end


using Base.Test, GLM, DataFrames

## Formaldehyde data from the R Datasets package
form = DataFrame(Carb=[0.1,0.3,0.5,0.6,0.7,0.9],OptDen=[0.086,0.269,0.446,0.538,0.626,0.782])
lm1 = fit(LinearModel, OptDen ~ Carb, form)
@test_approx_eq coef(lm1) linreg(array(form[:Carb]),array(form[:OptDen]))

dobson = DataFrame(Counts=[18.,17,15,20,10,20,25,13,12], Outcome=gl(3,1,9), Treatment=gl(3,3))
gm1 = fit(GeneralizedLinearModel, Counts ~ Outcome + Treatment, dobson, Poisson());
@test_approx_eq deviance(gm1) 5.12914107700115
@test_approx_eq coef(gm1)[1:3] [3.044522437723423,-0.45425527227759555,-0.29298712468147375]

## Example from http://www.ats.ucla.edu/stat/r/dae/logit.htm
df = readtable(Pkg.dir("GLM","data","admit.csv.gz"))
df[:rank] = pool(df[:rank])

gm2 = fit(GeneralizedLinearModel, admit ~ gre + gpa + rank, df, Binomial())
@test_approx_eq deviance(gm2) 458.5174924758994
@test_approx_eq coef(gm2) [-3.9899786606380734,0.0022644256521549043,0.8040374535155766,-0.6754428594116577,-1.3402038117481079,-1.5514636444657492]

gm3 = fit(GeneralizedLinearModel, admit ~ gre + gpa + rank, df, Binomial(), ProbitLink())
@test_approx_eq deviance(gm3) 458.4131713833386
@test_approx_eq coef(gm3) [-2.3867922998680786,0.0013755394922972369,0.47772908362647015,-0.4154125854823675,-0.8121458010130356,-0.9359047862425298]

gm4 = fit(GeneralizedLinearModel, admit ~ gre + gpa + rank, df, Binomial(), CauchitLink())
@test_approx_eq deviance(gm4) 459.3401112751141

gm5 = fit(GeneralizedLinearModel, admit ~ gre + gpa + rank, df, Binomial(), CloglogLink())
@test_approx_eq deviance(gm5) 458.89439629612616

## Example with offsets from Venables & Ripley (2002, p.189)
df = readtable(Pkg.dir("GLM","data","anorexia.csv.gz"))
df[:Treat] = pool(df[:Treat])

gm6 = fit(GeneralizedLinearModel, Postwt ~ Prewt + Treat, df, Normal(), IdentityLink(), offset=array(df[:Prewt]))
@test_approx_eq deviance(gm6) 3311.262619919613
@test_approx_eq coef(gm6) [49.7711090149846,-0.5655388496391,-4.0970655280729,4.5630626529188]
@test_approx_eq scale(gm6.model, true) 48.6950385282296
@test_approx_eq stderr(gm6) [13.3909581420259,0.1611823618518,1.8934926069669,2.1333359226431]

gm7 = fit(GeneralizedLinearModel, Postwt ~ Prewt + Treat, df, Normal(), LogLink(), offset=array(df[:Prewt]),
	      convTol=1e-8)
@test_approx_eq deviance(gm7) 3265.207242977156
@test_approx_eq coef(gm7) [3.992326787835955,-0.994452693131178,-0.050698258703974,0.051494029957641]
@test_approx_eq scale(gm7.model, true) 48.01787789178518
@test_approx_eq stderr(gm7) [0.157167944259695,0.001886285986164,0.022584069426311,0.023882826190166]

## Gamma example from McCullagh & Nelder (1989, pp. 300-2)
clotting = DataFrame(u = log([5,10,15,20,30,40,60,80,100]),
                     lot1 = [118,58,42,35,27,25,21,19,18])
gm8 = fit(GeneralizedLinearModel, lot1 ~ u, clotting, Gamma())
@test_approx_eq deviance(gm8) 0.01672971517848353
@test_approx_eq coef(gm8) [-0.01655438172784895,0.01534311491072141]
@test_approx_eq scale(gm8.model, true) 0.002446059333495581
@test_approx_eq stderr(gm8) [0.0009275466067257,0.0004149596425600]

gm9 = fit(GeneralizedLinearModel, lot1 ~ u, clotting, Gamma(), LogLink(), convTol=1e-8)
@test_approx_eq deviance(gm9) 0.16260829451739
@test_approx_eq coef(gm9) [5.50322528458221,-0.60191617825971]
@test_approx_eq scale(gm9.model, true) 0.02435442293561081
@test_approx_eq stderr(gm9) [0.19030107482720,0.05530784660144]

gm10 = fit(GeneralizedLinearModel, lot1 ~ u, clotting, Gamma(), IdentityLink(), convTol=1e-8)
@test_approx_eq deviance(gm10) 0.60845414895344
@test_approx_eq coef(gm10) [99.250446880986,-18.374324929002]
@test_approx_eq scale(gm10.model, true) 0.1041772704067886
@test_approx_eq stderr(gm10) [17.864388462865,4.297968703823]


# Approximate G_3 for small \sqrt{|\beta|s}:

using ForwardDiff

#const cg3 = 1./[6.0,-120.0,5040.0,-362880.0,39916800.0,-6227020800.0]
#const cH2 = 1./[3.0,-30.0,840.0,-45360.0,3991680.0,-518918400.0]
#const cH1 = 1./[12.0,-180.0,6720.0,-453600.0,47900160.0,-7264857600.0]

# Define a dummy function for automatic differentiation:
function G3(param::Array{T,1}) where {T <: Real}
return G3(param[1],param[2])
end

function G3(s::T,beta::T) where {T <: Real}
sqb = sqrt(abs(beta))
x = sqb*s
if beta >= 0
  return (x-sin(x))/(sqb*beta)
else
  return (x-sinh(x))/(sqb*beta)
end
end

y = zeros(2)
dg3 = y -> ForwardDiff.gradient(g3,y);

function H2(s::T,beta::T) where {T <: Real}
sqb = sqrt(abs(beta))
x=sqb*s
if beta >= 0
  return (sin(x)-x*cos(x))/(sqb*beta)
else
  return (sinh(x)-x*cosh(x))/(sqb*beta)
end
end

function H1(s::T,beta::T) where {T <: Real}
x=sqrt(abs(beta))*s
if beta >= 0
  return (2.0-2cos(x)-x*sin(x))/beta^2
else
  return (2.0-2cosh(x)+x*sinh(x))/beta^2
end
end

function g3_series(s::T,beta::T) where {T <: Real}
epsilon = eps(s)
# Computes G_3(\beta,s) using a series tailored to the precision of s.
x2 = -beta*s^2
term = one(s)
g3 = one(s)
n=0
# Terminate series when required precision reached:
while abs(term) > epsilon*abs(g3)
  n += 1
  term *= x2/((2n+3)*(2n+2))
  g3 += term
end
g3 *= s^3/6
return g3::T
end

dg3_series = y -> ForwardDiff.gradient(g3_series,y);

# Define a dummy function for automatic differentiation:
function g3_series(param::Array{T,1}) where {T <: Real}
return g3_series(param[1],param[2])
end

#function H2_series(s::T,beta::T) where {T <: Real}
#x2 = abs(beta)*s^2
#pm = sign(beta)
#return s^3*(cH2[1]+x2*(pm*cH2[2]+x2*(cH2[3]+
#   x2*(pm*cH2[4]+x2*(cH2[5]+x2*pm*cH2[6])))))
#end

function H2_series(s::T,beta::T) where {T <: Real}
# Computes H_2(\beta,s) using a series tailored to the precision of s.
epsilon = eps(s)
x2 = -beta*s^2
term = one(s)
h2 = one(s)
n=0
# Terminate series when required precision reached:
while abs(term) > epsilon*abs(h2)
  n += 1
  term *= x2
  term /= (4n+6)*n
  h2 += term
end
h2 *= s^3/3
return h2::T
end

#function H1_series(s::T,beta::T) where {T <: Real}
#x2 = abs(beta)*s^2
#pm = sign(beta)
#return x2*x2*(cH1[1]+x2*(pm*cH1[2]+x2*(cH1[3]+
#   x2*(pm*cH1[4]+x2*(cH1[5]+x2*pm*cH1[6])))))/beta^2
#end

function H1_series(s::T,beta::T) where {T <: Real}
# Computes H_1(\beta,s) using a series tailored to the precision of s.
epsilon = eps(s)
x2 = -beta*s^2
term = one(s)
h1 = one(s)
n=0
# Terminate series when required precision reached:
while abs(term) > epsilon*abs(h1)
  n += 1
  term *= x2*(n+1)
  term /= (2n+4)*(2n+3)*n
  h1 += term
end
h1 *= s^4/12
return h1::T
end

nx = 10001
s = linearspace(-3.0,3.0,nx)
#s = linearspace(-0.5,0.5,nx)

beta = -rand(); epsilon = eps(beta)
sqb = sqrt(abs(beta))
x = sqb*s
g31_bigm= g3.(convert(Array{BigFloat,1},s),big(beta))
println(typeof(x))
g31m= g3.(s,beta)
@time g31m= g3.(s,beta)
e31m = convert(Array{Float64,1},1.0-g31m./g31_bigm)
e31m[isnan.(e31m)]=epsilon
g32m = g3_series.(s,beta)
@time g32m = g3_series.(s,beta)
e32m = convert(Array{Float64,1},1.0-g32m./g31_bigm)
e32m[isnan.(e32m)]=epsilon

H21_bigm= H2.(convert(Array{BigFloat,1},s),big(beta))
H21m= H2.(s,beta)
@time H21m= H2.(s,beta)
eH21m = convert(Array{Float64,1},1.0-H21m./H21_bigm)
eH21m[isnan.(eH21m)]=epsilon
H22m = H2_series.(s,beta)
@time H22m = H2_series.(s,beta)
eH22m = convert(Array{Float64,1},1.0-H22m./H21_bigm)
eH22m[isnan.(eH22m)]=epsilon

H11_bigm= H1.(convert(Array{BigFloat,1},s),big(beta))
H11m= H1.(s,beta)
@time H11m= H1.(s,beta)
eH11m = convert(Array{Float64,1},1.0-H11m./H11_bigm)
eH11m[isnan.(eH11m)]=epsilon
H12m = H1_series.(s,beta)
@time H12m = H1_series.(s,beta)
eH12m = convert(Array{Float64,1},1.0-H12m./H11_bigm)
eH12m[isnan.(eH12m)]=epsilon

beta = rand()
sqb = sqrt(abs(beta))
x = sqb*s
g31_bigp= g3.(convert(Array{BigFloat,1},s),big(beta))
println(typeof(x))
@time g31p= g3.(s,beta)
e31p = convert(Array{Float64,1},1.0-g31p./g31_bigp)
e31p[isnan.(e31p)]=epsilon
@time g32p = g3_series.(s,beta)
e32p = convert(Array{Float64,1},1.0-g32p./g31_bigp)
e32p[isnan.(e32p)]=epsilon

H21_bigp= H2.(convert(Array{BigFloat,1},s),big(beta))
@time H21p= H2.(s,beta)
eH21p = convert(Array{Float64,1},1.0-H21p./H21_bigp)
eH21p[isnan.(eH21p)]=epsilon
@time H22p = H2_series.(s,beta)
eH22p = convert(Array{Float64,1},1.0-H22p./H21_bigp)
eH22p[isnan.(eH22p)]=epsilon

H11_bigp= H1.(convert(Array{BigFloat,1},s),big(beta))
@time H11p= H1.(s,beta)
eH11p = convert(Array{Float64,1},1.0-H11p./H11_bigp)
eH11p[isnan.(eH11p)]=epsilon
@time H12p = H1_series.(s,beta)
eH12p = convert(Array{Float64,1},1.0-H12p./H11_bigp)
eH12p[isnan.(eH12p)]=epsilon

# It looks like for x < 0.5, the fractional error on g3_series is <epsilon for
# float64 numbers.

# So, for x > 0.5, we can use x-sin(x), while for x < 0.5, we can use the series.
# Also need to consider what happens for negative values.

using PyPlot
clf()
semilogy(x,abs.(e31m),label="G3, exact, beta<0 ")
println("e31m: ",minimum(abs.(e31m)))
semilogy(x,abs.(e32m),label="G3, series, beta<0 ")
#semilogy([.5,.5],[1e-16,maxabs(e31m)],linestyle="dashed")
println("e32m: ",minimum(abs.(e32m)))
#semilogy(-[.5,.5],[1e-16,maxabs(e31m)],linestyle="dashed")
semilogy(x,0.*x+eps(1.0))
read(STDIN,Char)
clf()
semilogy(x,abs.(e31p),label="G3, exact, beta>0 ")
println("e31p: ",minimum(abs.(e31p)))
#semilogy([.5,.5],[1e-16,1.],linestyle="dashed")
semilogy(x,abs.(e32p),label="G3, series, beta>0 ")
println("e32p: ",minimum(abs.(e32p)))
#semilogy(-[.5,.5],[1e-16,maxabs(e31p)],linestyle="dashed")
semilogy(x,0.*x+eps(1.0))
read(STDIN,Char)
clf()
semilogy(x,abs.(eH21m),label="G1G2-G0G3, exact, beta<0 ")
println("eH21m: ",minimum(abs.(eH21m)))
#semilogy( [.5,.5],[1e-16,1.],linestyle="dashed")
semilogy(x,abs.(eH22m),label="G1G2-G0G3, series, beta<0 ")
println("eH22m: ",minimum(abs.(eH22m)))
#semilogy(-[.5,.5],[1e-16,maxabs(eH21m)],linestyle="dashed")
semilogy(x,0.*x+eps(1.0))
read(STDIN,Char)
clf()
semilogy(x,abs.(eH21p),label="G1G2-G0G3, exact, beta>0 ")
println("eH21p: ",minimum(abs.(eH21p)))
#semilogy( [.5,.5],[1e-16,1.],linestyle="dashed")
semilogy(x,abs.(eH22p),label="G1G2-G0G3, series, beta>0 ")
println("eH22p: ",minimum(abs.(eH22p)))
#semilogy(-[.5,.5],[1e-16,maxabs(abs.(eH21p))],linestyle="dashed")
semilogy(x,0.*x+eps(1.0))
read(STDIN,Char)
clf()
semilogy(x,abs.(eH11m),label="G2^2-G1G3, exact, beta<0 ")
println("eH11m: ",minimum(abs.(eH11m)))
#semilogy( [.5,.5],[1e-16,1.],linestyle="dashed")
semilogy(x,abs.(eH12m),label="G2^2-G1G3, series, beta<0 ")
println("eH12m: ",minimum(abs.(eH11p)))
#semilogy(-[.5,.5],[1e-16,maxabs(eH11m)],linestyle="dashed")
semilogy(x,0.*x+eps(1.0))
read(STDIN,Char)
clf()
semilogy(x,abs.(eH11p),label="G2^2-G1G3, exact, beta>0 ")
println("eH11p: ",minimum(abs.(eH11p)))
#semilogy( [.5,.5],[1e-16,1.],linestyle="dashed")
semilogy(x,abs.(eH12p),label="G2^2-G1G3, series, beta>0 ")
println("eH12p: ",minimum(abs.(eH12p)))
semilogy(x,0.*x+eps(1.0))
#semilogy(-[.5,.5],[1e-16,maxabs(eH11p)],linestyle="dashed")
#legend(loc="lower left")

read(STDIN,Char)
clf()
plot(x,abs.(g31m))
plot(x,abs.(g32m))
plot(x,abs.(g31p))
plot(x,abs.(g32p))
plot(x,abs.(H21m))
plot(x,abs.(H22m))
plot(x,abs.(H21p))
plot(x,abs.(H22p))
plot(x,abs.(H11m))
plot(x,abs.(H12m))
plot(x,abs.(H11p))
plot(x,abs.(H12p))


### A Pluto.jl notebook ###
# v0.12.17

using Markdown
using InteractiveUtils

# ╔═╡ bb427b30-47ae-11eb-28e6-19ba8afc227d
begin
using Flux
images = Flux.Data.MNIST.images();
labels = Flux.Data.MNIST.labels();
end

# ╔═╡ 96433c12-47f1-11eb-1589-bb257b66361c
begin
	using Random
	x=rand(1:60000)
	md"""
	$(images[x])
	# $(labels[x])"""
end

# ╔═╡ 03e5e240-4800-11eb-0dac-254f8ca6e0f8
begin
	using Flux: onehotbatch
	y=onehotbatch(labels, 0:9)
end

# ╔═╡ 13575c8a-49d2-11eb-35d5-678446f5c47e
using Flux:@epochs

# ╔═╡ a8485fde-47fb-11eb-0c26-05f1e3f300e3
foo=images[800]

# ╔═╡ ee91bc14-47fc-11eb-216c-61f1357586eb
reshape(foo,:)

# ╔═╡ c4977514-47fe-11eb-21f1-6da498eb7362
float(reshape(foo,:))

# ╔═╡ a7e7e35a-47ff-11eb-18df-13c13cf2690f
X=hcat(float.(reshape.(images,:))...)

# ╔═╡ a0362f42-49d1-11eb-0b18-cb34ea364cd9
m = Chain(
  Dense(28*28,40, relu),
  Dense(40, 10),
  softmax)

# ╔═╡ ab272c08-49d1-11eb-218e-fbd72b5323fd
begin
	using Flux:crossentropy,onecold,throttle
	loss(X, y) = crossentropy(m(X), y)
	opt = ADAM()
end

# ╔═╡ dae68f32-49d4-11eb-0f5f-b9caf19a35a2
begin
	using Statistics:mean
	accuracy(X, y) = mean(onecold(m(X)) .== onecold(y))
	md"""
	# Training accuracy: $(accuracy(X, y))
	# Test accuracy: $(accuracy(hcat(float.(reshape.(Flux.Data.MNIST.images(:test), :))...), onehotbatch(Flux.Data.MNIST.labels(:test), 0:9)))"""
end

# ╔═╡ d288829a-49d1-11eb-330c-b3ad51d54203
progress = () -> @show(loss(X, y))

# ╔═╡ d9026228-49d1-11eb-31ed-370d7411fa45
@epochs 100 Flux.train!(loss, params(m),[(X,y)], opt, cb = throttle(progress, 10))

# ╔═╡ 86716a10-49d4-11eb-0c5f-c15aa095dd26
begin
	foobar=rand(1:10000)
	md"""
	$(Flux.Data.MNIST.images(:test)[foobar])
	# Prediction: $(onecold(m(hcat(float.(reshape.(Flux.Data.MNIST.images(:test), :))...)[:,foobar])) - 1)
	"""
end

# ╔═╡ Cell order:
# ╠═bb427b30-47ae-11eb-28e6-19ba8afc227d
# ╠═96433c12-47f1-11eb-1589-bb257b66361c
# ╠═a8485fde-47fb-11eb-0c26-05f1e3f300e3
# ╠═ee91bc14-47fc-11eb-216c-61f1357586eb
# ╠═c4977514-47fe-11eb-21f1-6da498eb7362
# ╠═a7e7e35a-47ff-11eb-18df-13c13cf2690f
# ╠═03e5e240-4800-11eb-0dac-254f8ca6e0f8
# ╠═a0362f42-49d1-11eb-0b18-cb34ea364cd9
# ╠═ab272c08-49d1-11eb-218e-fbd72b5323fd
# ╠═d288829a-49d1-11eb-330c-b3ad51d54203
# ╠═13575c8a-49d2-11eb-35d5-678446f5c47e
# ╠═d9026228-49d1-11eb-31ed-370d7411fa45
# ╠═86716a10-49d4-11eb-0c5f-c15aa095dd26
# ╠═dae68f32-49d4-11eb-0f5f-b9caf19a35a2


using Test, NeXLUncertainties

include("uncertainValue.jl")
include("uncertainValues.jl")
include("propagate.jl")


let
    x, y = variables(["x", "y"])

    #@test (x*y).alpha == [1 1]
    #@test (x*y).c == [1]

    #@test (1 + 2*x*y).alpha == [0 0; 1 1]
    #@test (1 + 2*x*y).c == [1; 2]

    #@test ( (x+y)*(x-y) ).alpha == [2 0; 0 2]
    #@test ( (x+y)*(x-y) ).c == [1; -1]

    #@test ( (1+x+y)^2 * y ).deg == 3
    #@test ( (1+x+y)^2 * y ).alpha == [0 1; 1 1; 0 2; 2 1; 1 2; 0 3]
    #@test ( (1+x+y)^2 * y ).c == [1; 2; 2; 1; 2; 1]

    @test ( x*(x^2-y) == x^3 - x*y)
    @test ( (x^2-y)*x == x^3 - x*y)
    @test ( (x^2-y)*(x+y+x*y) == x^3 + x^2*y + x^3*y - x*y - y^2 - x*y^2)
    @test ( (x+y+x*y)*(x^2-y) == x^3 + x^2*y + x^3*y - x*y - y^2 - x*y^2)

    @test ( [x, y]*[x, y]' == [ x^2 x*y; x*y y^2] )

    @test ( [1 0; 0 2]*[x, y] == [x, 2*y])
    @test ( dot([x,y], [x,y]) == 1.0*x^2 + y^2 )

    @test ( dot([1//2*x, y, 1], [2, 1, 1]) == x + y + 1 )

    @test ( x + (Polyopt.Poly(0) + Polyopt.Poly(0)) == x )
end

let
    x, = variables(["x"])
    @test ( (1+x)^2 == 1+2*x+x^2 )
end

let
    x, y, z = variables(["x", "y", "z"])
    @test Polyopt.monomials(3,[x,y,z]) == sort([1,x,y,z,x^2,x*y,x*z,y^2,y*z,z^2,x^3,x^2*y,x^2*z,x*y^2,x*y*z,x*z^2,y^3,y^2*z,y*z^2,z^3])
end


seating = open("seating.txt") do f
    collect.(readlines(f))
end

function print_map(arr)
    for r in arr
        println(join(r))
    end
    println()
end

function count_neighborhood(arr, x::Int64, y::Int64, c::Char)::Int64
    n = 0

    if y > 1
        if arr[y-1][x] == c
            n += 1
        end
        if x > 1 && arr[y-1][x-1] == c
            n += 1
        end
        if x < length(arr[1]) && arr[y-1][x+1] == c
            n += 1
        end
    end

    if y < length(arr)
        if arr[y+1][x] == c
            n += 1
        end
        if x > 1 && arr[y+1][x-1] == c
            n += 1
        end
        if x < length(arr[1]) && arr[y+1][x+1] == c
            n += 1
        end
    end

    if x > 1 && arr[y][x-1] == c
        n += 1
    end

    if x < length(arr[1]) && arr[y][x+1] == c
        n += 1
    end

    return n
end

function count_visible(arr, x::Int64, y::Int64, c::Char)::Int64
    n = 0
    unit_vectors = map(x -> [round(Int,sin(x)),round(Int,cos(x))], pi*(-1:1/4:0.9))
    for uv in unit_vectors
        l = 1
        while true
            ox, oy = l * uv
            if y+oy < 1 || y+oy > length(arr) || x+ox < 1 || x+ox > length(arr[1])
                break
            end
            if arr[y+oy][x+ox] == c
                n += 1
            end
            if arr[y+oy][x+ox] != '.'
                break
            end
            l += 1
        end
    end
    return n
end

function simulate(arr, fn, n)
    step = deepcopy(arr)

    for y in 1:length(arr)
        for x in 1:length(arr[1])
            # print("at $(y),$(x): $(arr[y][x]): ")
            if arr[y][x] == '.'
                # println("do nothing")
            elseif arr[y][x] == 'L'
                occupied = fn(arr, x, y, '#')
                if occupied == 0
                    step[y][x] = '#'
                end
            elseif arr[y][x] == '#'
                occupied = fn(arr, x, y, '#')
                if occupied >= n
                    step[y][x] = 'L'
                end
            end
        end
    end

    return step
end

function run(state, fn, n)
    i = 0
    while true
        i += 1
        next_state = simulate(state, fn, n)
        if state == next_state
            break
        end
        state = deepcopy(next_state)
    end

    return state
end

println(sum(map(r -> count(x -> x == '#', r), run(seating, count_neighborhood, 4))))
println(sum(map(r -> count(x -> x == '#', r), run(seating, count_visible, 5))))


"""
The Mathenize module exists mainly to make math calculations in
strings.
"""
module Mathenize
export calculate
include("computation.jl")

"""
    calculate(math::String, print_info::Bool=false)

Return calculated value(s) in a string.

# Arguments
- `math::String` Math operation.
- `print_info::Bool` Shows log of what happened during the calculation.

# Examples
```julia-repl
julia> calculate("sqrt(complex(-90)) + 10im")
0.0 + 19.486832980505138im
```

"""
function calculate(math::String, print_info::Bool=false)

    r = nothing
    math = Meta.parse(math)
    hasproperty(math, :args) ? tasks = length(math.args) : tasks = 0

    #Adding to log info.
    LOG_INFO = []
    push!(LOG_INFO, "-> $(math) <- \n └Tasks: $(tasks)\n └$(hasproperty(math, :args) ? math.args : "Empty")")
    if math !== nothing
    if tasks >= 1
        push!(LOG_INFO, "   └ Performing subtasks:")
        r = subtasking(math, tasks, math.args, LOG_INFO, print_info)
    elseif math isa Number
        push!(LOG_INFO, "   └$(math) is a number")
        r = math
    elseif math in sym
        push!(LOG_INFO, "   └$(math) is a valid symbol")
        r = Core.eval(Base.Math, math)
    elseif tasks == 1 && ispermitted(math, LOG_INFO)
        push!(LOG_INFO, "   └$(math) is a valid a permitted number or expression")
        r = Core.eval(Base.Math, math)
    else
        r = nothing
        unknownmath(math, LOG_INFO, print_info)
    end
    end

    if print_info == true
        @info join(LOG_INFO, "\n")
    end

    return r #Return result

end
end #End module


using BilliardModels

function run_sinai(num_collisions)
    billiard_table = Sinai_billiard(0.354, true, true)  # periodic in x and y

    # x = Vector2D(0.3, 0.)
    # v = Vector2D(0.99, 0.01)

    x, v = initial_condition(billiard_table, -.5, .5, -.5, .5)

    l = Vector2D(0, 0)
    p = ParticleOnLattice(x, v, l)

    xs, ls, free_paths = billiard_dynamics_on_lattice(p, billiard_table, num_collisions);
end

N = 0
try
    N = int(ARGS[1])
catch
    N = 1_000_000
end

println("# Using N=$N")

run_sinai(1)
@time run_sinai(N)

using BilliardModels
billiard_table = Sinai_billiard(0.354, true, true)  # periodic in x and y
x, v = initial_condition(billiard_table, -.5, .5, -.5, .5)
l = Vector2D(0, 0)
p = ParticleOnLattice(x, v, l)
num_collisions = 10
xs, ls, free_paths = billiard_dynamics_on_lattice(p, billiard_table, num_collisions);
continuous_time(xs, ls, free_paths, 0.25)

billiard_dynamics_on_lattice


# function test_mesh()
img = rand(100,100)
offset = [0,0]
dist = 20
mesh = create_mesh(img, offset, dist)
i1 = 100 % (dist*sin(pi/3)) / 2
@test_approx_eq mesh.src_nodes[1,:] [i1,0]
@test_approx_eq mesh.src_nodes[2,:] [i1,20]
@test_approx_eq mesh.src_nodes[3,:] [i1,40]
@test_approx_eq mesh.src_nodes[4,:] [i1,60]
@test_approx_eq mesh.src_nodes[5,:] [i1,80]
@test_approx_eq mesh.src_nodes[6,:] [i1,100]
@test_approx_eq mesh.src_nodes[7,:] [i1 + 20*(sin(pi/3)), 10]
@test maximum(mesh.src_nodes[:,1]) <= 100
@test maximum(mesh.src_nodes[:,2]) <= 100
@test size(mesh.src_nodes, 1) == 33
# test that src_nodes match dst_nodes, initially
for (a, b) in zip(mesh.src_nodes, mesh.dst_nodes)
	@test a == b
end
@test size(mesh.edges, 2) == 81


@testset "variables macro" begin
    @rafvars x
    @test x == RevisedAffineForm(0, [1], 0)
    @test 1 + x ± 0.1 == RevisedAffineForm(1.0, @SVector([1.0]), 0.1)

    @rafvars x y z
    @test 1 + x + y + z ± 1 == RevisedAffineForm(1.0, @SVector([1.0, 1.0, 1.0]), 1.0)

    @rafvars x[1:4]
    p2 = 1 + x[1] + x[2] ± 2
    @test p2 == RevisedAffineForm(1.0, [1, 1, 0, 0], 2)
    @test string(p2) == "1.0 + 1.0x₁ + 1.0x₂ ± 2.0"
end

@testset "affinize" begin
    @affinize (a=1 .. 2, b=3 .. 4)
    @test a + b == RevisedAffineForm(5.0, @SVector([0.5, 0.5]), 0.0)
end


using SimpleSMC
import SequentialMonteCarlo.SMCModel
using RNGPool
using SMCExamples.LinearGaussian
using Random

include("test.jl")

setRNGs(0)

model, theta, ys, ko = LinearGaussian.defaultLGModel(10)
println(ko.logZhats)

numParticles = 1024
numTrials = 2

## just run the algorithm a few times
testSMC(model, numParticles, numTrials)

testCSMC(model, numParticles, numTrials)

lM = LinearGaussian.makelM(theta)
testCSMCBS(model, lM, numParticles, numTrials)


using ..ShallowWater: TurbulenceClosure, LinearDrag, ConstantViscosity

function null_init_state!(
    ::HomogeneousBox,
    ::TurbulenceClosure,
    state,
    aux,
    local_geometry,
    t,
)
    T = eltype(state.U)
    state.U = @SVector zeros(T, 2)
    state.η = 0
    return nothing
end

η_lsw(x, y, t) = cos(π * x) * cos(π * y) * cos(√2 * π * t)
u_lsw(x, y, t) = 2^(-0.5) * sin(π * x) * cos(π * y) * sin(√2 * π * t)
v_lsw(x, y, t) = 2^(-0.5) * cos(π * x) * sin(π * y) * sin(√2 * π * t)

function lsw_init_state!(
    m::ShallowWaterModel,
    p::HomogeneousBox,
    state,
    aux,
    local_geometry,
    t,
)

    coords = local_geometry.coord

    state.U = @SVector [
        u_lsw(coords[1], coords[2], t),
        v_lsw(coords[1], coords[2], t),
    ]

    state.η = η_lsw(coords[1], coords[2], t)

    return nothing
end

v_lkw(x, y, t) = 0
u_lkw(x, y, t) = exp(-0.5 * y^2) * exp(-0.5 * (x - t + 5)^2)
η_lkw(x, y, t) = 1 + u_lkw(x, y, t)

function lkw_init_state!(
    m::ShallowWaterModel,
    p::HomogeneousBox,
    state,
    aux,
    local_geometry,
    t,
)

    coords = local_geometry.coord

    state.U = @SVector [
        u_lkw(coords[1], coords[2], t),
        v_lkw(coords[1], coords[2], t),
    ]

    state.η = η_lkw(coords[1], coords[2], t)

    return nothing
end

R₋(ϵ) = (-1 - sqrt(1 + (2 * π * ϵ)^2)) / (2ϵ)
R₊(ϵ) = (-1 + sqrt(1 + (2 * π * ϵ)^2)) / (2ϵ)
D(ϵ) =
    (R₊(ϵ) * (exp(R₋(ϵ)) - 1) + R₋(ϵ) * (1 - exp(R₊(ϵ)))) /
    (exp(R₊(ϵ)) - exp(R₋(ϵ)))
R₂(x₁, ϵ) =
    (1 / D(ϵ)) * (
        (
            (R₊(ϵ) * (exp(R₋(ϵ)) - 1)) * exp(R₊(ϵ) * x₁) +
            (R₋(ϵ) * (1 - exp(R₊(ϵ)))) * exp(R₋(ϵ) * x₁)
        ) / (exp(R₊(ϵ)) - exp(R₋(ϵ)))
    )
R₁(x₁, ϵ) =
    (π / D(ϵ)) * (
        1 .+
        (
            (exp(R₋(ϵ)) - 1) * exp(R₊(ϵ) * x₁) .+
            (1 - exp(R₊(ϵ))) * exp(R₋(ϵ) * x₁)
        ) / (exp(R₊(ϵ)) - exp(R₋(ϵ)))
    )

𝒱(x₁, y₁, ϵ) = R₂(x₁, ϵ) * sin.(π * y₁)
𝒰(x₁, y₁, ϵ) = -R₁(x₁, ϵ) * cos.(π * y₁)
ℋ(x₁, y₁, ϵ, βᵖ, fₒ, γ) =
    (R₂(x₁, ϵ) / (π * fₒ)) * γ * cos(π * y₁) +
    (R₁(x₁, ϵ) / π) *
    (sin(π * y₁) * (1.0 + βᵖ * (y₁ - 0.5)) + (βᵖ / π) * cos(π * y₁))

function gyre_init_state!(
    m::SWModel,
    p::HomogeneousBox,
    T::LinearDrag,
    state,
    aux,
    local_geometry,
    t,
)

    coords = local_geometry.coord

    FT = eltype(state)
    τₒ = p.τₒ
    fₒ = m.fₒ
    β = m.β
    Lˣ = p.Lˣ
    Lʸ = p.Lʸ
    H = p.H

    γ = T.λ

    βᵖ = β * Lʸ / fₒ
    ϵ = γ / (Lˣ * β)

    _grav::FT = grav(m.param_set)

    uˢ(ϵ) = (τₒ * D(ϵ)) / (H * γ * π)
    hˢ(ϵ) = (fₒ * Lˣ * uˢ(ϵ)) / _grav

    u = uˢ(ϵ) * 𝒰(coords[1] / Lˣ, coords[2] / Lʸ, ϵ)
    v = uˢ(ϵ) * 𝒱(coords[1] / Lˣ, coords[2] / Lʸ, ϵ)
    h = hˢ(ϵ) * ℋ(coords[1] / Lˣ, coords[2] / Lʸ, ϵ, βᵖ, fₒ, γ)

    state.U = @SVector [H * u, H * v]

    state.η = h

    return nothing
end

t1(x, δᵐ) = cos((√3 * x) / (2 * δᵐ)) + (√3^-1) * sin((√3 * x) / (2 * δᵐ))
t2(x, δᵐ) = 1 - exp((-x) / (2 * δᵐ)) * t1(x, δᵐ)
t3(y, Lʸ) = π * sin(π * y / Lʸ)
t4(x, Lˣ, C) = C * (1 - x / Lˣ)

η_munk(x, y, Lˣ, Lʸ, δᵐ, C) = t4(x, Lˣ, C) * t3(y, Lʸ) * t2(x, δᵐ)

function gyre_init_state!(
    m::SWModel,
    p::HomogeneousBox,
    V::ConstantViscosity,
    state,
    aux,
    local_geometry,
    t,
)

    coords = local_geometry.coord

    FT = eltype(state.U)
    _grav::FT = grav(m.param_set)

    τₒ = p.τₒ
    fₒ = m.fₒ
    β = m.β
    Lˣ = p.Lˣ
    Lʸ = p.Lʸ
    H = p.H

    ν = V.ν

    δᵐ = (ν / β)^(1 / 3)
    C = τₒ / (_grav * H) * (fₒ / β)

    state.η = η_munk(coords[1], coords[2], Lˣ, Lʸ, δᵐ, C)
    state.U = @SVector zeros(FT, 2)

    return nothing
end


# --------10--------20--------30--------40--------50--------60--------70--------80--------90--------100-------110-------120--------130
# vector potential
function getAt(t::TF, pulse::Pulse{At, TF}) where TF
	@unpack tL, tR, T, A, ω, ϕ = pulse
	if t < tL || tR < t
		return 0
	end
	return A*get_env(t, pulse)*sin(ω*t + ϕ)
end
function getAt(t::TVF, pulse::Pulse{At}) where {TVF <: AbstractVector{<:AbstractFloat}}
	At = zeros(eltype(t), size(t))
	for it in 1:length(t)
		At[it] = getAt(t[it], pulse)
	end
	return At
end
function getAt(t::TF, laser::Vector{Pulse{At, TF}}) where {TF <: AbstractFloat}
	At = zero(t)
	for ip in 1:length(laser)
		At += getAt(t, laser[ip])
	end
	return At
end
function getAt(t::TVF, laser::Vector{Pulse{At}}) where {TVF <: AbstractVector{<:AbstractFloat}}
	At = zeros(eltype(t), size(t))
	for it in 1:length(t)
		At[it] = getAt(t[it], laser)
	end
	return At
end
# electric field
# Et = -dA/dt
function getEt(t::TF, pulse::Pulse{At})where TF <: AbstractFloat
	@unpack tL, tR, T = pulse
	if t < tL || tR < t
		return 0
	end
	dt = T*1e-5
	Et =  (1/60dt)*(getAt(t + 3dt, pulse) - getAt(t - 3dt, pulse)) -
			(3/20dt)*(getAt(t + 2dt, pulse) - getAt(t - 2dt, pulse)) +
			(3/4dt)*(getAt(t + dt, pulse) - getAt(t - dt, pulse))
	return -Et
end
function getEt(t::TVF, pulse::Pulse{At}) where {TVF <: AbstractVector{<:AbstractFloat}}
	Et = zeros(eltype(t), size(t))
	for it in 1:length(t)
		Et[it] = getEt(t[it], pulse)
	end
	return Et
end
function getEt(t::TF, laser::Vector{Pulse{At}}) where {TF <: AbstractFloat}
	Et = zero(t)
	for ip in 1:length(laser)
		Et += getEt(t, laser[ip])
	end
	return Et
end
function getEt(t::TVF, laser::Vector{Pulse{At}}) where {TVF <: AbstractVector{<:AbstractFloat}}
	Et = zeros(eltype(t), size(t))
	for it in 1:length(t)
		Et[it] = getEt(t[it], laser)
	end
	return Et
end


using Plots,LaTeXStrings
using LinearAlgebra,FundamentalsNumericalComputation

z = [-1,3];
v = z - norm(z)*[1,0];
P = I - 2*(v*v')/dot(v,v)
Pz = P*z;

gr()
default(legend=nothing,linewidth=1)
plt = plot(aspect_ratio=1,showaxis=false,grid=false,xaxis=:none,yaxis=false)

plot!([0,z[1]],[0,z[2]],aspect_ratio=1)
plot!([Pz[1],z[1]],[Pz[2],z[2]])
plot!([0,Pz[1]],[0,Pz[2]])

plot!([-2,2]*v[2],[2 ,-2]*v[1],l=(:dash,:black))

scatter!([0],[0],m=(5,:black,:black))
xlims!(-1,4)
ylims!(-1, 3)

annotate!(-0.68,1.5,text(L"\mathbf{z}"))
annotate!(2,1,text(L"\mathbf{v}"))
annotate!(1.5,-0.3,text(L"\|\mathbf{z}\|\,\mathbf{e}_1"))
annotate!(-0.2,0,text(L"0"))
annotate!(1.6,1.9,text(L"S"))

savefig(plt,"hhreflect.svg")


using GraphPlot, SimpleHypergraphs
import LightGraphs

using DataFrames
using CSV

#cd(raw"copy_the_path_here_if_your_Julia_cannot_find_files")

b = CSV.read("business.csv",delim="\t")

b.business_id = Symbol.(b.business_id);
for name in names(b)[2:15]
    b[!,name] = Symbol.(b[!, name])
end

revs = CSV.read("review2.csv",delim="\t")

revs.business_id = Symbol.(revs.business_id)

@assert sort(b.business_id) == sort(unique(revs.business_id))

using StatsBase
using Random

using SimpleHypergraphs

function report(revs::DataFrame, b::DataFrame, stars)
    revsg = revs[revs.stars .== stars, :]
    sort!(revsg,:business_id)

    business_ids_set = Set{Symbol}(revsg.business_id)
    #list of businesses filtered to the data
    bf = filter(row -> row[:business_id] in business_ids_set, b)
    sort!(bf,:business_id)
    bf.b_id = 1:nrow(bf)

    revsg.user_id = categorical(revsg.user_id,true)

    # Verices x hyperedges
    h = Hypergraph{Float64}(nrow(bf),length(levels(revsg.user_id)));
    for r in 1:nrow(revsg)
        bix = searchsorted(bf.business_id,revsg.business_id[r]).start
        h[bix,revsg.user_id[r].level] = 1
    end

    println("Constructed a hypergraph having size: " ,size(h))

    for gt in [:city, :state, :Alcohol,:NoiseLevel,:RestaurantsTakeOut,:category]
        partition = Vector{Set{Int}}()
        for bfg in groupby(bf,gt)
            push!(partition,Set{Int}(sort(bfg.b_id)))
        end
        println(gt," ",modularity(h,partition))
    end
end

for stars in 1:5
    println(stars)
    report(revs, b, stars)
end


# include("../../src/SFFM.jl")
using LinearAlgebra, Plots, JLD2, GLM

## define the model(s)
# include("exampleModelDef.jl")

## load sims
@load pwd()*"/examples/baseAdvComparison/dump/sims.jld2" sims

## analytic X distribution for comparison
# construction
Ψₓ = SFFM.PsiFunX(approxModel)
ξₓ = SFFM.MakeXiX(approxModel, Ψₓ)
pₓ, πₓ, Πₓ, Kₓ = SFFM.StationaryDistributionX(approxModel, Ψₓ, ξₓ)

## section 4.5: error for approximation of stationary distribution of X
Basis = "lagrange"
Δs = [1.6; 0.8; 0.4; 0.2; 0.1; 0.05]
NBasesRange = [1; 2; 3; 4]

πnorms = zeros(length(Δs), length(NBasesRange))
πnormsAdv = zeros(length(Δs), length(NBasesRange))
times = Array{Any}(undef, length(Δs), length(NBasesRange))
timesAdv = Array{Any}(undef, length(Δs), length(NBasesRange))
gctimes = Array{Any}(undef, length(Δs), length(NBasesRange))
gctimesAdv = Array{Any}(undef, length(Δs), length(NBasesRange))
mems = Array{Any}(undef, length(Δs), length(NBasesRange))
memsAdv = Array{Any}(undef, length(Δs), length(NBasesRange))
alloc = Array{Any}(undef, length(Δs), length(NBasesRange))
allocAdv = Array{Any}(undef, length(Δs), length(NBasesRange))
approxSpec = Array{Any}(undef, length(Δs), length(NBasesRange))

for d = 1:length(Δs), n = 1:length(NBasesRange)
    if d + 2 * n <= 10
        # define the mesh for each iteration
        nBases = NBasesRange[n]
        Δtemp = Δs[d]
        println("mesh details; h = ", Δtemp, " NBases = ", nBases)

        # collect time and memory stats
        ~, times[d, n], mems[d, n], gctimes[d, n], alloc[d, n] = @timed begin
            Nodes = collect(approxBounds[1, 1]:Δtemp:approxBounds[1, 2])
            mesh = SFFM.DGMesh(
                approxModel,
                Nodes,
                nBases,
                Basis = Basis,
            )
            approxSpec[d, n] = (Δtemp, nBases, SFFM.TotalNBases(mesh) * SFFM.NPhases(approxModel))

            # compute the marginal via DG
            Matrices = SFFM.MakeMatrices(approxModel, mesh,probTransform=false)
            MatricesR = SFFM.MakeMatricesR( approxModel, mesh)
            B = SFFM.MakeB( approxModel, mesh, Matrices,probTransform=false)
            Dr = SFFM.MakeDR(
                Matrices,
                MatricesR,
                approxModel,
                mesh,
                B,
            )
            Ψ = SFFM.PsiFun(Dr.DDict)
            R = SFFM.MakeR( approxModel, mesh, approxType="projection", probTransform=false)

            # the distribution of X when Y first returns to 0
            ξ = SFFM.MakeXi( B.BDict, Ψ, mesh = mesh, model = approxModel, probTransform = false)

            marginalX, p, K = SFFM.MakeLimitDistMatrices(
                B.BDict,
                Dr.DDict,
                R.RDict,
                Ψ,
                ξ,
                mesh,
                approxModel,
                probTransform = false,
            )
        end
        # convert marginalX to a distribution for analysis
        DGStationaryDist = SFFM.Coeffs2Dist(
            approxModel,
            mesh,
            marginalX,
            SFFM.SFFMProbaility,
            probTransform = false,
        )

        ## stationary distirbution stuff
        # evaluate the analytic result given the mesh
        analyticX = SFFM.SFFMProbability(
            [pₓ[:]; 0; 0],
            Πₓ(Matrix(mesh.Nodes[2:end]')) - Πₓ(Matrix(mesh.Nodes[1:end-1]')),
            mesh.Nodes[1:end-1] + SFFM.Δ(mesh) / 2,
        )

        # save them
        πnorms[d, n] = SFFM.starSeminorm(DGStationaryDist, analyticX)
    else
        πnorms[d, n] = NaN
        times[d, n], mems[d, n], gctimes[d, n], alloc[d, n] = NaN, NaN, NaN, NaN
    end
end

let p = plot()
    types = [:solid,:dash,:dashdot,:dashdotdot]
    for n = 1:length(NBasesRange)
        p = plot!(
            Δs,
            πnorms[:, n],
            xaxis = :log,
            yaxis = :log,
            xlabel = "log(h)",
            ylabel = "log error",
            label = "N_k = " * string(NBasesRange[n]),
            legend = :outertopright,
            linestyle = types[n],
            seriestype = :line,
            markershape = :auto,
            title = "Error",
            grid = false,
        )
    end
    display(p)
    savefig(pwd()*"/examples/baseAdvComparison/dump/piErrorVsDelta.png")
end

let p = plot()
    types = [:solid,:dash,:dashdot,:dashdotdot,:solid,:dash]
    for d = 1:length(Δs)
        p = plot!(
            NBasesRange,
            πnorms[d, :],
            yaxis = :log,
            xlabel = "N_k",
            ylabel = "log error",
            label = "h = " * string(Δs[d]),
            legend = :outertopright,
            linestyle = types[d],
            seriestype = :line,
            markershape = :auto,
            title = "Error",
            grid = false,
        )
    end
    display(p)
    savefig(pwd()*"/examples/baseAdvComparison/dump/piErrorVsNBases.png")
end

begin
    n = 1
    data = log.(πnorms[:,n])
    ind = .!isnan.(data)
    data = data[ind]
    println("rate of convergence for ", n, " basis functions is ")
    display(
        lm([ones(length(Δs[ind])) log.(Δs[ind])], data)
    )
    println("")

    n = 2
    data = log.(πnorms[:,n])
    ind = .!isnan.(data)
    data = data[ind]
    println("rate of convergence for ", n, " basis functions is ")
    display(
        lm([ones(length(Δs[ind])) log.(Δs[ind])], data)
    )
    println("")

    n = 3
    data = log.(πnorms[:,n])
    ind = .!isnan.(data)
    ind[4] = false
    data = data[ind]
    data = data
    println("rate of convergence for ", n, " basis functions is ")
    display(
        lm([ones(length(Δs[ind])) log.(Δs[ind])], data)
    )
    println("")
end

for d = 1:min(4,length(Δs))
    data = log.(πnorms[d,:])
    ind = .!isnan.(data)
    data = data[ind]
    println("rate of convergence for h = ", Δs[d])
    display(
        lm([ones(length(NBasesRange[ind])) NBasesRange[ind]], data)
    )
    println("")
end

mp =
    (log.(πnorms[1:end-1, :]) .- log.(πnorms[2:end, :])) ./
    (log.(Δs[1:end-1]) - log.(Δs[2:end]))
np =
    (log.(πnorms[:, 1:end-1]) .- log.(πnorms[:, 2:end])) ./
    ((NBasesRange[1:end-1]) - (NBasesRange[2:end]))'

stats = Array{Any}(undef,length(Δs)+1,length(NBasesRange)+1)
for n in 1:length(NBasesRange), d in 1:length(Δs)
    if d + 2 * n <= 10
        stats[d+1,n+1] = (πnorms[d,n], times[d,n], mems[d,n]/(1024*1024), approxSpec[d,n][3])
    else
        stats[d+1,n+1] = (0,0,0,0)
    end
end
stats[1,2:end] = NBasesRange
stats[2:end,1] = Δs
display(stats)
using NumericIO
innards = stats[2:end,2:end]
asciiexponentfmt = NumericIO.IOFormattingExpNum(
	"\\times10^{", false, '+', '-', NumericIO.ASCII_SUPERSCRIPT_NUMERALS
)
fmt = NumericIO.IOFormattingReal(asciiexponentfmt,
	ndigits=3, decpos=0, decfloating=true, eng=false, minus='-', inf="Inf"
)

begin
    println("\\begin{tabular}{ | c | l | l | l | l | l | }")
    println("\\hline")
    println("\\multirow{2}{*}{\\(\\Delta\\)} & \\multirow{2}{*}{}& \\multicolumn{4}{c|}{Number of basis functions, \\(N_k\\)} \\\\ \\cline{3-6}
          &      & 1            & 2           & 3           & 4           \\\\\\hline"
        )
    for d in 1:length(Δs)
        println(Δs[d], " &")
        println("
        \\(\\begin{array}{l}
            \\mbox{error}
            \\\\ \\mbox{time (sec)}
            \\\\ \\mbox{memory (MB)}
            \\\\ n_\\varphi
        \\end{array}\\)"
        )
        for n in 1:length(NBasesRange)
            println("&")
            if d + 2 * n <= 10
                println("
                \\(\\begin{array}{l}
                    ", formatted(innards[d,n][1],fmt),"}
                    \\\\",round(innards[d,n][2],sigdigits=3),"
                    \\\\",round(innards[d,n][3],sigdigits=3),"
                    \\\\",convert(Int,round(innards[d,n][4],digits=0)),"
                \\end{array}\\)"
                )
            else
                println(" - ")
            end
        end
        println("\\\\ \\hline")
    end
    println("\\end{tabular}")
end


using Base.Test
using QuantumOptics

const ode = QuantumOptics.ode_dopri.ode

const ζ = 0.5
const ω₀ = 10.0

const y₀ = Float64[0., sqrt(1-ζ^2)*ω₀]
const A = 1
const ϕ = 0

function f(t::Float64)
    α = sqrt(1-ζ^2)*ω₀
    x = A*exp(-ζ*ω₀*t)*sin(α*t + ϕ)
    p = A*exp(-ζ*ω₀*t)*(-ζ*ω₀*sin(α*t + ϕ) + α*cos(α*t + ϕ))
    return [x,p]
end

function df(t::Float64, y::Vector{Float64}, dy::Vector{Float64})
    dy[1] = y[2]
    dy[2] = -2*ζ*ω₀*y[2] - ω₀^2*y[1]
    return nothing
end

# Test absolute and relative tolerances
T = [0.:0.5:5.;]

for reltol=[1e-3,1e-4,1e-5,1e-6,1e-7,1e-8,1e-9]
    tout, y = ode(df, T, y₀, reltol=reltol, abstol=0.)
    for i=2:length(T)
        # println("T = $(T[i])")
        # println("Target Reltol: $(reltol*T[i]); Real reltol: ", norm(y[i]-f(T[i]))/norm(y[i]))
        @test T[i]*reltol*20 > norm(y[i]-f(T[i]))/norm(y[i]) > T[i]*reltol/20
    end
end

for abstol=[1e-3,1e-4,1e-5,1e-6,1e-7,1e-8,1e-9]
    tout, y = ode(df, T, y₀, reltol=0., abstol=abstol)
    for i=2:length(T)
        # println("T = $(T[i])")
        # println("Target abstol: $(abstol*T[i]); Real abstol: ", norm(y[i]-f(T[i])))
        @test abstol*10 > norm(y[i]-f(T[i])) > abstol/100
    end
end

# Test output of intermediate steps
T = [0.,10.]
tout, yout = ode(df, T, y₀; display_intermediatesteps=true)
@test length(tout)>2
maxstep = maximum(tout[2:end]-tout[1:end-1])

# Test hmax stepszie control
hmax = maxstep/10
tout, yout = ode(df, T, y₀; display_intermediatesteps=true, hmax=hmax)
maxstep2 = maximum(tout[2:end]-tout[1:end-1])
@test (maxstep2-hmax)<1e-12

# Test hmin
@test_throws ErrorException ode(df, T, y₀; display_intermediatesteps=true, hmin=0.4)

# Test h0
h0 = 1e-5
tout, yout = ode(df, T, y₀; display_intermediatesteps=true, h0=h0)
@test abs(h0-(tout[2]-tout[1]))<1e-14
h0 = 0.4
tout, yout = ode(df, T, y₀; display_intermediatesteps=true, h0=h0)
@test abs(h0-(tout[2]-tout[1]))>0.3

# Test fout
tout_ = Float64[]
yout_ = Vector{Float64}[]
function fout(t, y)
    push!(tout_, t)
    push!(yout_, deepcopy(y))
end
T = [0.,0.5,1.0]
result = ode(df, T, y₀, fout)
@test result==nothing
tout, yout = ode(df, T, y₀)

@test tout==tout_
@test yout==yout_


## Instrumental variables problems

dt = readcsv("iv.csv");

y = dt[:,1];
x = dt[:,2];
z = dt[:,3:7];

f(theta) = z.*(y-x*theta);

gmms2 = GMMEstimator(f, [.0])
initialize!(gmm2s)
estimate!(gmm2s)

srand(1)
using MomentBasedEstimators
using ModelsGenerators
y,x,z = ModelsGenerators.sim_iv_d01(n = 10000, CP = 2000)
xb    = randn(10000,2)
x     = [x xb]
z     = [z xb]

f(theta) = z.*(y-x*theta);

gmm = MomentBasedEstimators.GMMEstimator(f, [.1, 0, 0]);
md = MomentBasedEstimators.MDEstimator(f, [.1, 0, 0]);

MomentBasedEstimators.initialize!(gmm);
MomentBasedEstimators.initialize!(md);

estimate!(gmm);
estimate!(md);

maximum(vcov(gmm)-vcov(md))<1e-07
maximum(vcov(md, false, :Unweighted) - vcov(md, false, :Weighted))<1e-05

MomentBasedEstimators.setparLB!(md, coef(md))
MomentBasedEstimators.setparUB!(md, coef(md))
estimate!(md);

MomentBasedEstimators.setparLB!(md, [0.,0,0])
MomentBasedEstimators.setparUB!(md, [0.,0,0])

## Methods

# 1. Coefficient
coef(gmm)
coef(md)

## 2. Variances
vcov(gmm)
vcov(md)

## 3. Objectiv value

srand(1)
y,x,z=randiv(100)
xb = randn(100,2)
x = [x xb]
z = [z xb]

z=sparse(z)

f(theta) = broadcast(*, z, (y-x*theta))

g = MomentBasedEstimators.GMMEstimator(f, [.1, 0, 0])
m = MomentBasedEstimators.MDEstimator(f, [.1, 0, 0])

MomentBasedEstimators.initialize!(g);
MomentBasedEstimators.initialize!(m);

@time estimate!(g);
@time estimate!(m);

srand(1)
y,x,z=randiv(100)
xb = randn(100,2)
x = [x xb]
z = [z xb]

f(theta) = z.*(y-x*theta);

g = MomentBasedEstimators.GMMEstimator(f, [.1, 0, 0])
m = MomentBasedEstimators.MDEstimator(f, [.1, 0, 0])

MomentBasedEstimators.initialize!(g)
MomentBasedEstimators.initialize!(m)

@time estimate!(g);
@time estimate!(m);

## Before
## elapsed time: 0.129503039 seconds (69486712 bytes allocated, 40.36% gc time)


"""
    LogNormalExpert(μ, σ)

PDF:

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

See also: [Lognormal Distribution](https://en.wikipedia.org/wiki/Log-normal_distribution) (Wikipedia)

"""
struct LogNormalExpert{T<:Real} <: NonZIContinuousExpert
    μ::T
    σ::T
    LogNormalExpert{T}(µ::T, σ::T) where {T<:Real} = new{T}(µ, σ)
end

function LogNormalExpert(μ::T, σ::T; check_args=true) where {T <: Real}
    check_args && @check_args(LogNormalExpert, σ >= zero(σ))
    return LogNormalExpert{T}(μ, σ)
end

## Outer constructors
LogNormalExpert(μ::Real, σ::Real) = LogNormalExpert(promote(μ, σ)...)
LogNormalExpert(μ::Integer, σ::Integer) = LogNormalExpert(float(μ), float(σ))
LogNormalExpert() = LogNormalExpert(0.0, 1.0)

## Conversion
function convert(::Type{LogNormalExpert{T}}, μ::S, σ::S) where {T <: Real, S <: Real}
    LogNormalExpert(T(μ), T(σ))
end
function convert(::Type{LogNormalExpert{T}}, d::LogNormalExpert{S}) where {T <: Real, S <: Real}
    LogNormalExpert(T(d.μ), T(d.σ), check_args=false)
end
copy(d::LogNormalExpert) = LogNormalExpert(d.μ, d.σ, check_args=false)

## Loglikelihood of Expert
logpdf(d::LogNormalExpert, x...) = Distributions.logpdf.(Distributions.LogNormal(d.μ, d.σ), x...)
pdf(d::LogNormalExpert, x...) = Distributions.pdf.(Distributions.LogNormal(d.μ, d.σ), x...)
logcdf(d::LogNormalExpert, x...) = Distributions.logcdf.(Distributions.LogNormal(d.μ, d.σ), x...)
cdf(d::LogNormalExpert, x...) = Distributions.cdf.(Distributions.LogNormal(d.μ, d.σ), x...)

## expert_ll, etc
expert_ll_exact(d::LogNormalExpert, x::Real) = LRMoE.logpdf(d, x)
function expert_ll(d::LogNormalExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_ll = (yl == yu) ? logpdf.(d, yl) : logcdf.(d, yu) + log1mexp.(logcdf.(d, yl) - logcdf.(d, yu))
    expert_ll = (tu == 0.) ? -Inf : expert_ll
    return expert_ll
end
function expert_tn(d::LogNormalExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_tn = (tl == tu) ? logpdf.(d, tl) : logcdf.(d, tu) + log1mexp.(logcdf.(d, tl) - logcdf.(d, tu))
    expert_tn = (tu == 0.) ? -Inf : expert_tn
    return expert_tn
end
function expert_tn_bar(d::LogNormalExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_tn_bar = (tl == tu) ? 0.0 : log1mexp.(logcdf.(d, tu) + log1mexp.(logcdf.(d, tl) - logcdf.(d, tu)))
    return expert_tn_bar
end

exposurize_expert(d::LogNormalExpert; exposure = 1) = d

## Parameters
params(d::LogNormalExpert) = (d.μ, d.σ)
function params_init(y, d::LogNormalExpert)
    pos_idx = (y .> 0.0)
    μ_init, σ_init = mean(log.(y[pos_idx])), sqrt(var(log.(y[pos_idx])))
    μ_init = isnan(μ_init) ? 0.0 : μ_init
    σ_init = isnan(σ_init) ? 1.0 : σ_init
    return LogNormalExpert(μ_init, σ_init)
end

## KS stats for parameter initialization
function ks_distance(y, d::LogNormalExpert)
    p_zero = sum(y .== 0.0) / sum(y .>= 0.0)
    return max(abs(p_zero-0.0), (1-0.0)*HypothesisTests.ksstats(y[y .> 0.0], Distributions.LogNormal(d.μ, d.σ))[2])
end

## Simululation
sim_expert(d::LogNormalExpert) = Distributions.rand(Distributions.LogNormal(d.μ, d.σ), 1)[1]

## penalty
penalty_init(d::LogNormalExpert) = [Inf 1.0 Inf]
no_penalty_init(d::LogNormalExpert) = [Inf 1.0 Inf]
penalize(d::LogNormalExpert, p) = 0 # (d.μ/p[1])^2 + (p[2]-1)*log(d.σ) - d.σ/p[3]

## statistics
mean(d::LogNormalExpert) = mean(Distributions.LogNormal(d.μ, d.σ))
var(d::LogNormalExpert) = var(Distributions.LogNormal(d.μ, d.σ))
quantile(d::LogNormalExpert, p) = quantile(Distributions.LogNormal(d.μ, d.σ), p)
lev(d::LogNormalExpert, u) = exp(d.μ+0.5*d.σ^2) * cdf.(Normal(d.μ+d.σ^2, d.σ), log(u)) + u*(1-cdf.(Normal(d.μ, d.σ), log(u)))
excess(d::LogNormalExpert, u) = mean(d) - lev(d, u)

## Misc functions for E-Step
function _diff_dens_series(d::LogNormalExpert, yl, yu)
    return exp(-0.5*(log(yl)-d.μ)^2/(d.σ^2)) - exp(-0.5*(log(yu)-d.μ)^2/(d.σ^2))
end

function _diff_dist_series(d::LogNormalExpert, yl, yu)
    return ( 0.5 + 0.5*erf((log(yu)-d.μ)/(sqrt2*d.σ)) ) - ( 0.5 + 0.5*erf((log(yl)-d.μ)/(sqrt2*d.σ)) )
end

function _int_obs_logY(d::LogNormalExpert, yl, yu, expert_ll)
    if yl == yu
        return log(yl)
    else
        return exp(-expert_ll) * ( d.σ*invsqrt2π*0.5*_diff_dens_series(d, yl, yu) + d.μ*_diff_dist_series(d, yl, yu) )
    end
end

function _int_lat_logY(d::LogNormalExpert, tl, tu, expert_tn_bar)
    return exp(-expert_tn_bar) * (d.μ - ( d.σ*invsqrt2π*0.5*_diff_dens_series(d, tl, tu) + d.μ*_diff_dist_series(d, tl, tu) ) )
end

function _zdensz_series(z)
    return (isinf(z) ? 0.0 : z*exp(-0.5*z^2))
end

function _diff_zdensz_series(d::LogNormalExpert, yl, yu)
    return _zdensz_series((log(yl)-d.μ)/d.σ) - _zdensz_series((log(yu)-d.μ)/d.σ)
end

function _int_obs_logY_sq(d::LogNormalExpert, yl, yu, expert_ll)
    if yl == yu
        return (log(yl))^2
    else
        return exp(-expert_ll) * ( ((d.σ)^2 + (d.μ)^2)*_diff_dist_series(d, yl, yu) + d.σ*d.μ*invsqrt2π*_diff_dens_series(d, yl, yu) + (d.σ)^2*invsqrt2π*_diff_zdensz_series(d, yl, yu) )
    end
end

function _int_lat_logY_sq(d::LogNormalExpert, tl, tu, expert_tn_bar)
    return exp(-expert_tn_bar) * ( ((d.σ)^2 + (d.μ)^2) - ( ((d.σ)^2 + (d.μ)^2)*_diff_dist_series(d, tl, tu) + d.σ*d.μ*invsqrt2π*_diff_dens_series(d, tl, tu) + (d.σ)^2*invsqrt2π*_diff_zdensz_series(d, tl, tu) ))
end

## EM: M-Step
function EM_M_expert(d::LogNormalExpert,
                     tl, yl, yu, tu,
                     exposure,
                     z_e_obs, z_e_lat, k_e;
                     penalty = true, pen_pararms_jk = [Inf 1.0 Inf])

    if d.σ < 0.000001
        return d
    end

    expert_ll_pos = expert_ll.(d, tl, yl, yu, tu)
    expert_tn_bar_pos = expert_tn_bar.(d, tl, yl, yu, tu)

    # Further E-Step
    logY_e_obs = vec( _int_obs_logY.(d, yl, yu, expert_ll_pos) )
    logY_e_lat = vec( _int_lat_logY.(d, tl, tu, expert_tn_bar_pos) )
    nan2num(logY_e_lat, 0.0) # get rid of NaN

    logY_sq_e_obs = vec( _int_obs_logY_sq.(d, yl, yu, expert_ll_pos) )
    logY_sq_e_lat = vec( _int_lat_logY_sq.(d, tl, tu, expert_tn_bar_pos) )
    nan2num(logY_sq_e_lat, 0.0) # get rid of NaN

    # Update parameters
    pos_idx = (yu .!= 0.0)
    term_zkz = z_e_obs[pos_idx] .+ (z_e_lat[pos_idx] .* k_e[pos_idx])
    term_zkz_logY = (z_e_obs[pos_idx] .* logY_e_obs[pos_idx]) .+ (z_e_lat[pos_idx] .* k_e[pos_idx] .* logY_e_lat[pos_idx])
    term_zkz_logY_sq = (z_e_obs[pos_idx] .* logY_sq_e_obs[pos_idx]) .+ (z_e_lat[pos_idx] .* k_e[pos_idx] .* logY_sq_e_lat[pos_idx])

    μ_new = sum(term_zkz_logY)[1] / sum(term_zkz)[1]
    σ_new = sqrt( 1/sum(term_zkz)[1] * (sum(term_zkz_logY_sq)[1] - 2.0*μ_new*sum(term_zkz_logY)[1] + (μ_new)^2*sum(term_zkz)[1] ) )

    return LogNormalExpert(μ_new, σ_new)
end

## EM: M-Step, exact observations
function EM_M_expert_exact(d::LogNormalExpert,
                            ye, exposure,
                            z_e_obs;
                            penalty = true, pen_pararms_jk = [Inf 1.0 Inf])

    if d.σ < 0.000001
        return d
    end

    # Further E-Step
    logY_e_obs = log.(ye)

    logY_sq_e_obs = (log.(ye)).^2

    # Update parameters
    pos_idx = (ye .!= 0.0)
    term_zkz = z_e_obs[pos_idx]
    term_zkz_logY = (z_e_obs[pos_idx] .* logY_e_obs[pos_idx])
    term_zkz_logY_sq = (z_e_obs[pos_idx] .* logY_sq_e_obs[pos_idx])

    μ_new = sum(term_zkz_logY)[1] / sum(term_zkz)[1]
    σ_new = sqrt( 1/sum(term_zkz)[1] * (sum(term_zkz_logY_sq)[1] - 2.0*μ_new*sum(term_zkz_logY)[1] + (μ_new)^2*sum(term_zkz)[1] ) )

    return LogNormalExpert(μ_new, σ_new)
end


module TestNetworkRoles
   using Base.Test
   using Brim
   #=testdir = dirname(@__FILE__)=#
   A = [1 0 0; 1 1 0; 0 1 1; 0 1 1]
   S = [1 0; 0 1; 0 1; 1 0; 1 0; 0 1; 0 1]
   M = Modular(A, S)
   roles = network_roles(M)
   ## Approximate values, calculated by hand
   @test abs(roles[:z][1] - 1.1547) <= 0.01
   @test abs(roles[:c][2] - 0.4444) <= 0.01
   @test roles[:c][5] == 0.5
end


import MitosisStochasticDiffEq as MSDE
using Mitosis
using Test, Random
using LinearAlgebra

@testset "backward information filtering tests" begin
  # define SDE function
  f(u,p,t) = p[1]*u .+ p[2]
  g(u,p,t) = p[3]

  # time span
  tstart = 0.0
  tend = 1.0
  dt = 0.001
  trange = tstart:dt:tend

  # set of linear parameters Eq.~(2.2)
  B, β, σ̃ = -0.1, 0.2, 1.3
  plin = [B, β, σ̃]

  # initial values for ODE
  mynames = (:logscale, :μ, :Σ);
  c = 0.0
  ν = 0.0
  P = 10.0
  myvalues = [c, ν, P]
  NT = NamedTuple{mynames}(myvalues)

  # covariance filter
  kernel = MSDE.SDEKernel(Mitosis.AffineMap(B, β), Mitosis.ConstantMap(σ̃), trange, plin)
  message, solend = MSDE.backwardfilter(kernel, NT)

  # information filter
  H = inv(P)
  F = H*ν
  myvalues2 = [c, F, H]
  NT2 = NamedTuple{mynames}(myvalues2)
  message2, solend2 = MSDE.backwardfilter(kernel, NT2, filter=MSDE.InformationFilter())

  @test !isapprox(solend, solend2, rtol=1e-1)
  invH = inv(solend2[2])
  invHF = invH*solend2[1]
  @test isapprox(solend[1], invHF, rtol=1e-3)
  @test isapprox(solend[2], invH, rtol=1e-3)
  @test isapprox(solend[3], solend2[3], rtol=1e-10)

  # multivariate tests
  dim = 5
  Random.seed!(123)
  c = randn()
  ν = randn(dim)
  P = randn(dim,dim)
  myvalues = [c, ν, P]
  NT = NamedTuple{mynames}(myvalues)
  H = inv(P)
  F = H*ν
  myvalues2 = [c, F, H]
  NT2 = NamedTuple{mynames}(myvalues2)

  # covariance filter
  kernel = MSDE.SDEKernel(Mitosis.AffineMap(B, β), Mitosis.ConstantMap(σ̃), trange, plin)
  message, solend = MSDE.backwardfilter(kernel, NT)

  # information filter
  message2, solend2 = MSDE.backwardfilter(kernel, NT2,
    filter=MSDE.InformationFilter())

  @test !isapprox(solend, solend2, rtol=1e-1)
  invH = inv(solend2.x[2])
  invHF = invH*solend2.x[1]
  @test isapprox(solend.x[1], invHF, rtol=1e-3)
  @test isapprox(solend.x[2], invH, rtol=1e-2)
  @test isapprox(solend.x[3], solend2.x[3], rtol=1e-10)

  # test inplace version
  message2, solend2 = MSDE.backwardfilter(kernel, NT2,
    filter=MSDE.InformationFilter(), inplace=true)

  @test !isapprox(solend, solend2, rtol=1e-1)
  invH = inv(solend2.x[2])
  invHF = invH*solend2.x[1]
  @test isapprox(solend.x[1], invHF, rtol=1e-3)
  @test isapprox(solend.x[2], invH, rtol=1e-2)
  @test isapprox(solend.x[3], solend2.x[3], rtol=1e-10)
end
#
# # define SDE function
# f(u,p,t) = p[1]*u .+ p[2]
# g(u,p,t) = p[3]
#
# # time span
# tstart = 0.0
# tend = 1.0
# dt = 0.001
# trange = tstart:dt:tend
#
# m = 3 # some number of Brownian processes
# Random.seed!(123)
# plin = [randn(dim,dim), randn(dim), randn(dim,m)] # B, β, σtil
#
# # multivariate tests
# dim = 5
# Random.seed!(123)
# c = randn()
# ν = randn(dim)
# P = randn(dim,dim)
# myvalues = [c, ν, P]
# NT = NamedTuple{mynames}(myvalues)
# H = inv(P)
# F = H*ν
# myvalues2 = [c, F, H]
# NT2 = NamedTuple{mynames}(myvalues2)
#
# # covariance filter
# kernel = MSDE.SDEKernel(Mitosis.AffineMap(B, β), Mitosis.ConstantMap(σ̃), trange, plin)
# message, solend = MSDE.backwardfilter(kernel, NT)
#
# # information filter
# message2, solend2 = MSDE.backwardfilter(kernel, NT2,
#   filter=MSDE.InformationFilter())
#
#
# @test !isapprox(solend, solend2, rtol=1e-1)
# invH = inv(solend2.x[2])
# invHF = invH*solend2.x[1]
# @test isapprox(solend.x[1], invHF, rtol=1e-3)
# @test isapprox(solend.x[2], invH, rtol=1e-2)
# @test isapprox(solend.x[3], solend2.x[3], rtol=1e-10)

# test inplace version
# message2, solend2 = MSDE.backwardfilter(kernel, NT2, inplace=true,
#   filter=MSDE.InformationFilter())
#
# @test !isapprox(solend, solend2, rtol=1e-1)
# invH = inv(solend2.x[2])
# invHF = invH*solend2.x[1]
# @test isapprox(solend.x[1], invHF, rtol=1e-3)
# @test isapprox(solend.x[2], invH, rtol=1e-2)
# @test isapprox(solend.x[3], solend2.x[3], rtol=1e-10)


using SequentialMonteCarlo
using RNGPool
import SMCExamples.LinearGaussian.defaultLGModel
import SMCExamples.Particles.Float64Particle
import Statistics: mean, var

setRNGs(0)

model, theta, ys, ko = defaultLGModel(100)

smcio = SMCIO{model.particle, model.pScratch}(1024*1024, model.maxn,
  Threads.nthreads(), false)
smc!(model, smcio)

predMeanSMC = SequentialMonteCarlo.eta(smcio, p -> p.x, false, smcio.n)
predVarSMC = SequentialMonteCarlo.eta(smcio, p -> p.x^2, false, smcio.n) -
  predMeanSMC^2
println("Predictive mean: $(ko.predictionMeans[smcio.n])")
println("Estimate: $predMeanSMC")
println("Predictive variance: $(ko.predictionVariances[smcio.n])")
println("Estimate: $predVarSMC")

filtMeanSMC = SequentialMonteCarlo.eta(smcio, p -> p.x, true, smcio.n)
filtVarSMC = SequentialMonteCarlo.eta(smcio, p -> p.x^2, true, smcio.n) -
  filtMeanSMC^2
println("Filtering mean: $(ko.filteringMeans[smcio.n])")
println("Estimate: $filtMeanSMC")
println("Filtering variance: $(ko.filteringVariances[smcio.n])")
println("Estimate: $filtVarSMC")

println("Running many particle filters with only 128 particles in parallel...")
m = 10000
lZs = Vector{Float64}(undef, m)
Vs = Vector{Float64}(undef, m)
lZhats = Vector{Float64}(undef, m)
Vhats = Vector{Float64}(undef, m)
f1(p::Float64Particle) = 1.0

nthreads = Threads.nthreads()
smcios = Vector{SMCIO}(undef, nthreads)
Threads.@threads for i = 1:nthreads
  smcios[i] = SMCIO{model.particle, model.pScratch}(128, model.maxn, 1, true)
end
@time Threads.@threads for i = 1:m
  smcio = smcios[Threads.threadid()]
  smc!(model, smcio)
  lZs[i] = smcio.logZhats[smcio.n-1]
  Vs[i] = SequentialMonteCarlo.V(smcio, f1, false, false, smcio.n)
  lZhats[i] = smcio.logZhats[smcio.n]
  Vhats[i] = SequentialMonteCarlo.V(smcio, f1, true, false, smcio.n)
end

print("Empirical relative variance of Z_n^N: ")
println(var(exp.(lZs.-ko.logZhats[smcio.n-1])))
print("Mean of unbiased estimator of relative variance of Z_n^N: ")
println(mean(Vs .* exp.(2*(lZs.-ko.logZhats[smcio.n-1]))))
print("Estimated standard deviation of the value above: ")
println(sqrt(var(Vs .* exp.(2*(lZs.-ko.logZhats[smcio.n-1])))/m))

print("Empirical relative variance of \\hat{Z}_n^N: ")
println(var(exp.(lZhats.-ko.logZhats[smcio.n])))
print("Mean of unbiased estimator of relative variance of \\hat{Z}_n^N: ")
println(mean(Vhats .* exp.(2*(lZhats.-ko.logZhats[smcio.n]))))
print("Estimated standard deviation of the value above: ")
println(sqrt(var(Vhats .* exp.(2*(lZhats.-ko.logZhats[smcio.n])))/m))


# Taken from https://github.com/org-arl/AlphaStableDistributions.jl/blob/master/src/AlphaStableDistributions.jl, which has a more complete implementation.
# Maybe we should just depend on it ?

struct AlphaStable{T} <: Distributions.ContinuousUnivariateDistribution
    α::T
    β::T
    scale::T
    location::T
end
Base.eltype(::Type{<:AlphaStable{T}}) where {T<:AbstractFloat} = T
function Base.rand(rng::Distributions.AbstractRNG, d::AlphaStable{T}) where {T<:AbstractFloat}
    # Shamelessly taken from https://github.com/org-arl/AlphaStableDistributions.jl/blob/master/src/AlphaStableDistributions.jl
    α=d.α; β=d.β; scale=d.scale; loc=d.location
    (α < 0 || α > 2) && throw(DomainError(α, "α must be in the range 0 to 2"))
    abs(β) > 1 && throw(DomainError(β, "β must be in the range -1 to 1"))
    ϕ = (rand(rng, T) - 0.5) * π
    if α == one(T) && β == zero(T)
        return loc + scale * tan(ϕ)
    end
    w = -log(rand(rng, T))
    α == 2 && (return loc + 2*scale*sqrt(w)*sin(ϕ))
    β == zero(T) && (return loc + scale * ((cos((1-α)*ϕ) / w)^(one(T)/α - one(T)) * sin(α * ϕ) / cos(ϕ)^(one(T)/α)))
    cosϕ = cos(ϕ)
    if abs(α - one(T)) > 1e-8
        ζ = β * tan(π * α / 2)
        aϕ = α * ϕ
        a1ϕ = (one(T) - α) * ϕ
        return loc + scale * (( (sin(aϕ) + ζ * cos(aϕ))/cosϕ * ((cos(a1ϕ) + ζ*sin(a1ϕ))) / ((w*cosϕ)^((1-α)/α)) ))
    end
    bϕ = π/2 + β*ϕ
    x = 2/π * (bϕ * tan(ϕ) - β * log(π/2*w*cosϕ/bϕ))
    α == one(T) || (x += β * tan(π*α/2))
    return loc + scale * x
end


@testset verbose=true "wigner_matrices" begin
    include("wigner_matrices/small_d.jl")
    include("wigner_matrices/H.jl")
    include("wigner_matrices/big_D.jl")
    include("wigner_matrices/sYlm.jl")
end


function find_arrest(S₀, S₁; rtol = eps(), artol = sqrt(eps()))
    grid = S₁.grid
    Λᵤ = Interpolator(S₀, S₁)
    ws = ApproximationGrids.weights(grid)
    avars, kvars, D₀ = initialize_asymptotics(S₁)

    K = nodes(grid)
    S = kvars.svars.S
    υ = kvars.svars.υ
    w = kvars.svars.w
    d = dimensionality(S₁.f.liquid)
    Z = similar(w)
    T = eltype(K)

    i  = argmax(Λᵤ.A)
    u₀ = zero(T)
    u₁ = log(floatmax(T)) / abs(Λᵤ.A[i])
    uₘ = u₁ / 2

    g = LiquidsDynamics.FixedPoint(Z, kvars, D₀)

    while true
        interpolate!(S, Λᵤ, uₘ)
        LiquidsDynamics.asymptotics_weights!(w, ws, K, S, d, grid)
        ζ = υ * sum(w)

        asymptotics!(g, avars, ζ; rtol = artol)

        if iszero(avars.ζ∞[])
            u₀ = uₘ
        else
            u₁ = uₘ
        end

        if isapprox(u₀, u₁, rtol = rtol, nans = true)
            if iszero(avars.ζ∞[])
                interpolate!(S, Λᵤ, u₁)
                LiquidsDynamics.asymptotics_weights!(w, ws, K, S, d, grid)
                asymptotics!(g, avars, υ * sum(w); rtol = artol)
            end
            break
        end

        uₘ = (u₀ + u₁) / 2
    end

    return u₀, u₁, avars.ζ∞[]
end


using LinearAlgebra
using Zygote
using ForwardDiff
using ReverseDiff
using DiffEqOperators
using StaticArrays
using Base.Iterators
using Plots

## Stencil generation
neighbours(CI, deg)        = CI .+ CartesianIndices(ntuple(x -> -deg:deg, length(CI)))

stencil_depth(CI, depth)   = [ CI + CartesianIndex(ntuple(x -> ifelse(x == i, depth, 0), length(CI))) for i in 1:length(CI) ]

stencils(CI, deg, poke = false) = ifelse(poke, [ stencil_depth(CI, -(deg + 1)); neighbours(CI, deg)[:]; stencil_depth(CI, deg + 1) ], neighbours(CI, deg)[:])

von_neumann_stencil(CI, n) = stencils(CI, n, true)
moore_stencil(CI, n)       = stencils(CI, n, false)
cs_stencil(φ, index)       = @view φ[von_neumann_stencil(index, 0)]

## Cells and boundaries
abstract type AbstractCell end

# This structure can be encoded generically using the Store comonad; here we define it explicitly for testing.
struct CartesianPoint{M,N,T} <: AbstractCell
    index  :: CartesianIndex{N}
    values :: SVector{M,T}
end

von_neumann_stencil(ci :: CartesianPoint, n) = von_neumann_stencil(ci.index, n)
moore_stencil(ci :: CartesianPoint, n)       = moore_stencil(ci.index, n)

struct CartesianBoundary{M,N,T} <: AbstractCell
    point      :: CartesianPoint{M,N,T}
    neighbours :: Vector{Tuple{CartesianIndex{N},SVector{M,T}}}
end

## Finite differencing operations on values and stencils

# First derivatives
Δ_central(φp, φm, Δx)     = (φp - φm) / 2Δx
Δ_forward(φ1, φ2, Δt)     = (φ2 - φ1) / Δt
central_2nd(φp, φ, φm)    = φm - 2φ + φp

# Stencil
Δx_central(φ, dx)  = Δ_central(φ[4], φ[1], dx)
Δy_central(φ, dy)  = Δ_central(φ[5], φ[2], dy)

Δx_backward(φ, dx) = Δ_backward(φ[3], φ[1], dx)
Δy_backward(φ, dy) = Δ_backward(φ[3], φ[2], dy)

Δx_forward(φ, dx)  = Δ_forward(φ[4], φ[3], dx)
Δy_forward(φ, dy)  = Δ_forward(φ[5], φ[3], dy)

# Second derivatives
Δ²_central(φp, φ, φm, Δx) = central_2nd(φp, φ, φm) / Δx^2
Δ²_central(φp, φ, φm, Δxp, Δxm) = Δ_forward(φp, φ, Δxp) - Δ_forward(φ, φm, Δxm)

# Stencil
Δ²x_central(φ, dx) = Δ²_central(φ[4], φ[3], φ[1], dx)
Δ²y_central(φ, dy) = Δ²_central(φ[5], φ[3], φ[2], dy)

Δ²x_central(φ, dxs) = Δ²_central(φ[4], φ[3], φ[1], dxs[1], dxs[2])
Δ²y_central(φ, dys) = Δ²_central(φ[5], φ[3], φ[2], dys[1], dys[2])

# Boundary (this needs type distinctions between boundary and non-boundary cells to be efficient)
function substitute(φ, index, val)
    if index ∈ CartesianIndices(φ)
        φ[index]
    else
        val
    end
end

## Governing equations and residual setup
temperature_residual(φ, k, q, ds) = -k * (Δ²x_central(φ, ds[1:2]) + Δ²y_central(φ, ds[3:4])) + q

# Ruleset for thermal diffusion with a source field
function ruleset(φs, σs, k, dxs, dys, index, φ_boundaries)
    # Boundary conditions check
    if any(index.I .<= 1) || any(index.I .>= size(φs))
        # There have to be more conditions generally...
        φs_vec = [ substitute(φs, local_index, φ_boundaries) for local_index in von_neumann_stencil(index, 0) ]
        ds_vec = [ substitute(dxs, local_index, 0.1)          for local_index in von_neumann_stencil(index, 0) ]
    else
        # Get neighbours of a cell and corresponding sizes
        φs_vec   = cs_stencil(φs, index)
        dxp, dxn = dxs[index], dxs[index + CartesianIndex(1,0)]
        dyp, dyn = dys[index], dys[index + CartesianIndex(0,1)]
        ds_vec   = [ dxp; dxn; dyp; dyn ]
    end

    σ = σs[index]

    # Evaluate residuals
    temperature_residual(φs_vec, k, σ, ds_vec)
end

# Compute residual equations
compute_residuals(p, σ, k, dxs, dys, φ_boundary) = map(index -> ruleset(p, σ, k, dxs, dys, index, φ_boundary), CartesianIndices(p))

newton_update(p, δp, α = 1.0) = p - α * δp

# Pre-allocated version
newton_update!(p, δp, α = 1.0) = @. p = p - α * δp

## Initial setup
nx, ny  = 50, 40

# Less allocations than broadcasting inside sum
polynomial(x, coeffs) = sum(k -> k[1] * x^k[2], zip(coeffs, 0:length(coeffs)-1))

# Grid generation
xs      = range(0, 1, length = nx)
ys      = polynomial.(xs, Ref(ones(5)))
grid    = reshape([ SVector(x, y) for (x, y_max) in zip(xs, ys) for y in range(y_max, -y_max, length = 2ny) ], 2ny, nx)

# Compute grid spacings
dxs = norm.(grid[:,2:end] .- grid[:,1:end-1])
dys = norm.(grid[2:end,:] .- grid[1:end-1,:])

##
α       = 1.0
k       = 1.
φ_bound = 200

φ0      = φ_bound .* ones(size(grid)...)
σ       = rand(size(grid)...)

p       = φ0
R       = similar(p[:])
# ∂R∂p = zeros(2 .* (prod(size(ωs)), prod((reverse ∘ size)(ωs)))...)

num_iters = 3
compute_residuals!(R, p) = R .= compute_residuals(p, σ, k, dxs, dys, φ_bound)[:]
compute_grad_fwd(R, p) = ForwardDiff.jacobian(compute_residuals!, R, p)

##
# using NLsolve
# prob = nlsolve(compute_residuals!, p)

## NEWTON ITERATIONS
#=================================================================#

function newton_solver!(R, p, num_iters = 3, α = 1.0)
    # Array to store errors
    ε = zeros(num_iters);

    # Newton iteration loop
    for i in 1:num_iters
        # Compute residuals
        R    = compute_residuals!(R, p)

        # Compute Jacobian
        ∂R∂p = compute_grad_fwd(R, p)

        # Compute Newton step
        Δp   = ∂R∂p \ -R

        # Update state with relaxation factor
        p    = newton_update!(p, reshape(Δp, size(p)), α)

        # Error processing
        ε[i] = maximum(abs.(Δp))
        println("Newton step error L² norm: $(ε[i])")
    end
    R, p, ε
end

##
R, p, ε = newton_solver!(R, p, num_iters, α)


# ------------------------------------------------------------------
# It is advantageous to have the coeffs as
# one single vector in order to make use of BLAS
# and related libraries when defining operators on
# our space of functions
#
# In this script, we "vectorize" our functions,
# not in the regular meaning of the word, but in the
# sense that dictionaries will all be recelld by vectors
#
# The main difficulty with this is that, unlike with dictionaries
# there is no easy way to go from (level, cell, mode) to a
# corresponding index in a 1-D vector.
#
# For this, we have the reference functions:
# full/sparse_referenceV2D/D2V
# D2V generates a dict that, upon input of a level, cell, mode
# gives the corresponding index in a vector
# V2D generates a vector, with row i having the three numbers
# level, cell, mode corresponding to index i in the vector
#
# These methods work in all dimensions, and there are ones for both
# full and sparse grids.
#
# The entire script culminates in a final result: a matrix
# representation of the derivative operator
# (both in full and sparse bases)
# ------------------------------------------------------------------

# Calculates the exact dimension of interpolating
# basis functions in the full or sparse grid schemes
# using the degree k-1 Galerkin polynomials
# with multi-resolution up to n
function get_size(::Val{D}, k::Int, n::Int, scheme::Val{Scheme}=Val(:sparse)) where {D, Scheme}
    ls      = ntuple(i->(n+1), D)
    size    = 0
    for level in CartesianIndices(ls)
        cutoff(scheme, level, n) && continue

        ks = ntuple(q -> 1<<max(0, level[q]-2), D)
        size += prod(ks)*k^D
    end
    return size
end

function D2V(d::Int, k::Int, n::Int, coeffs::Dict{CartesianIndex{D}, <:AbstractArray{<:AbstractArray{T, D}, }};
    scheme="sparse") where {D, T <: Real}
    @assert d == D
    D2V(Val{D}(), k, n, coeffs, Val(Symbol(scheme)))
end
function D2V(::Val{D}, k::Int, n::Int, coeffs::Dict{CartesianIndex{D}, <:AbstractArray{<:AbstractArray{T, D}, }},
    scheme::Val{Scheme}) where {D, T <: Real, Scheme}

    size        = get_size(Val{D}(), k, n, scheme)
    vect        = Array{T}(undef, size)
    modes       = ntuple(i-> k, D)
    ls            = ntuple(i->(n+1), D)
    j = 1
    for level in CartesianIndices(ls) #This really goes from 0 to l_i for each i
        cutoff(scheme, level, n) && continue

        ks = ntuple(q -> 1<<max(0, level[q]-2), D) #This sets up a specific k+1 vector
        for cell in CartesianIndices(ks)
            for mode in CartesianIndices(modes)
                vect[j] = coeffs[level][cell][mode]
                j += 1
            end
        end
    end
    return vect
end

function V2D(D::Int, k::Int, n::Int, vect::Array{T}; scheme="sparse") where {T <: Real}
    V2D(Val{D}(), k, n, vect, Val(Symbol(scheme)))
end
function V2D(::Val{D}, k::Int, n::Int, vect::Array{T}, scheme::Val{Scheme}) where {D, T <: Real, Scheme}
    coeffs        = Dict{CartesianIndex{D}, Array{Array{T,D},D}}()
    modes        = ntuple(q-> k, D)
    ls            = ntuple(i->(n+1), D)
    j = 1
    for level in CartesianIndices(ls) #This really goes from 0 to l_i for each i
        cutoff(scheme, level, n) && continue

        ks = ntuple(q -> 1<<max(0, level[q]-2), D)  #This sets up a specific k+1 vector
        level_coeffs = Array{Array{T}}(undef, ks) #all the coefficients at this level
        for cell in CartesianIndices(ks)
            cell_coeffs = Array{T}(undef, modes)
            for mode in CartesianIndices(modes)
                cell_coeffs[mode] = vect[j]
                j += 1
            end
            level_coeffs[cell] = cell_coeffs
        end
        coeffs[level] = level_coeffs
    end
    return coeffs
end

function D2Vref(::Val{D}, k::Int, n::Int, scheme::Val{Scheme}) where {D, Scheme}
    size        = get_size(Val{D}(), k, n, scheme)
    dict        = Dict{NTuple{3,CartesianIndex{D}}, Int}()
    modes        = ntuple(q-> k, D)
    ls            = ntuple(i->(n+1), D)
    j = 1
    for level in CartesianIndices(ls)
        cutoff(scheme, level, n) && continue

        ks = ntuple(q -> 1<<max(0, level[q]-2), D)  #This sets up a specific k+1 vector
        lvl = ntuple(i -> level[i]-1,D)
        for cell in CartesianIndices(ks)
            for mode in CartesianIndices(modes)
                dict[(level, cell, mode)] = j
                j += 1
            end
        end
    end
    return dict
end

function V2Dref(::Val{D}, k::Int, n::Int, scheme::Val{Scheme}) where {D, Scheme}
    size        = get_size(Val{D}(), k, n, scheme)
    vect        = Array{NTuple{3,CartesianIndex{D}}}(undef, size)
    modes       = ntuple(q-> k, D)
    ls            = ntuple(i->(n+1), D)
    j = 1
    for level in CartesianIndices(ls)
        cutoff(scheme, level, n) && continue

        ks = ntuple(q -> 1<<max(0, level[q]-2), D)  #This sets up a specific k+1 vector
        lvl = ntuple(i -> level[i]-1,D)
        for cell in CartesianIndices(ks)
            for mode in CartesianIndices(modes)
                vect[j] = (level, cell, mode)
                j += 1
            end
        end
    end
    return vect
end

# -----------------------------------------------------
# Let's now make the coefficient operators work on
# and return vectors
# -----------------------------------------------------
function vcoeffs_DG(D::Int, k::Int, n::Int, f::Function;
                                rtol=REL_TOL, atol=ABS_TOL,
                                maxevals=MAX_EVALS,
                                scheme="sparse")
    vcoeffs_DG(Val(D), k, n, f, rtol, atol, maxevals, Val(Symbol(scheme)))
end
function vcoeffs_DG(::Val{D}, k::Int, n::Int, f::Function,
                                rtol, atol,
                                maxevals,
                                scheme::Val{Scheme}) where {D, Scheme}
    len            = get_size(Val{D}(), k, n, scheme)
    coeffs        = Array{Float64}(undef, len)
    modes        = ntuple(i-> k, D)
    ls            = ntuple(i-> (n+1), D)
    j = 1
    for level in CartesianIndices(ls)  # This really goes from 0 to l_i for each i,
        cutoff(scheme, level, n) && continue

        ks = ntuple(i -> 1<<max(0, level[i]-2), D)  #This sets up a specific k+1 vector
        lvl = ntuple(i -> level[i]-1,D)
        for cell in CartesianIndices(ks)
            for mode in CartesianIndices(modes)
                coeffs[j] = get_coefficient_DG(k, lvl, cell, mode, f;
                                                rtol=rtol, atol=atol,
                                                maxevals=maxevals)
                j += 1
            end
        end
    end
    return coeffs
end