using Turing
using Distributions
using AdvancedMH
using AdvancedHMC
using LinearAlgebra: I, dot
using Random
using Test

using AnnealedIS

function geomspace(start, stop, length)
    logstart = log10(start)
    logstop = log10(stop)
    points = 10 .^ range(logstart, logstop; length=length)
    points[1] = start
    points[end] = stop
    return points
end

@testset "AnnealedIS.jl" begin
    rng = MersenneTwister(42)
    Random.seed!(42)

    @testset "AnnealedISSampler" begin
        D = 5
        N = 2

        prior_sampling(rng) = rand(rng, MvNormal(D, 1.0))
        prior_density(params) = logpdf(MvNormal(D, 1.0), params)
        joint_density(params) = logpdf(MvNormal(D, 0.5), params)

        ais = AnnealedISSampler(
            prior_sampling, prior_density, joint_density, N)

        @test ais.betas == [0.0, 0.5, 1.0]
        x = ais.prior_sampling(rng)
        y = rand(rng, ais.transition_kernels[1])
        @test typeof(x) == typeof(y)

        samples = [ais.prior_sampling(rng) for _ in 1:N]

        @test AnnealedIS.logdensity(ais, 1, samples[1]) == prior_density(samples[1])
        @test AnnealedIS.logdensity(ais, 3, samples[1]) == joint_density(samples[1])
        @test AnnealedIS.logdensity(ais, 2, samples[1]) == (0.5*prior_density(samples[1]) + 0.5*joint_density(samples[1]))

        for i in 1:N
            dratio = AnnealedIS.logdensity(ais, i+1, samples[1]) - AnnealedIS.logdensity(ais, i, samples[1])
            @test AnnealedIS.logdensity_ratio(ais, i, samples[1]) == dratio
        end
    end

    @testset "Convergence test" begin
        D = 1
        N = 100
        num_samples = 1000

        y_obs = [3.0]

        prior_sampling(rng) = rand(rng, MvNormal(D, 1.0))
        prior_density(params) = logpdf(MvNormal(D, 1.0), params)
        joint_density(params) = logpdf(MvNormal(params, I), y_obs) + prior_density(params)

        ais = AnnealedISSampler(
            prior_sampling, prior_density, joint_density, N)

        samples, diagnostics = ais_sample(rng, ais, num_samples)

        posterior_mean = AnnealedIS.estimate_expectation(samples, x -> x)
        @test isapprox(posterior_mean, [1.5]; atol=3e-2)

        ess = effective_sample_size(samples)
        # TODO: What is a reasonable ESS to expect?
        @test ess > 900
        @test ess == diagnostics[:ess]
    end

    @testset "Logdensity ratio" begin
        D = 1
        N = 100
        num_samples = 1000

        y_obs = [3.0]

        prior_sampling(rng) = rand(rng, MvNormal(D, 1.0))
        prior_density(params) = logpdf(MvNormal(D, 1.0), params)
        joint_density(params) = logpdf(MvNormal(params, I), y_obs) + prior_density(params)

        betas = (geomspace(1, 1001, N+1) .- 1) ./ 1000
        anis_alg =
        ais = AnnealedISSampler(
            AnnealedIS.AnISModel(prior_sampling, prior_density, joint_density),
            AnIS(betas, (x, i) -> Normal(0, 1), SimpleRejection())
        )

        sample = [1.0]

        for i in 1:(length(ais.betas)-1)
            @test isapprox(
                AnnealedIS.logdensity_ratio(ais, i, sample),
                (AnnealedIS.logdensity(ais, i+1, sample) - AnnealedIS.logdensity(ais, i, sample));
                atol=1e-12 # TODO: What is the numerical error we would expect?
            )
        end
    end

    @testset "Sample from Turing" begin
        @model function test_model(y)
            x ~ Normal(0, 1)
            y ~ Normal(y, 1)
        end

        y_obs = 2
        tm = test_model(y_obs)

        named_tuple = sample_from_prior(rng, tm)
        named_tuple2 = sample_from_prior(rng, tm)
        @test typeof(named_tuple[:x]) == Float64
        @test named_tuple != named_tuple2

        # Test model with multiple latents and multivariate latent.
        D = 5
        @model function test_model2(y)
            a ~ Normal(0, 1)
            x ~ MvNormal(D, 1.0)
            y ~ MvNormal(x, I)
        end

        y_obs = ones(D)
        tm2 = test_model2(y_obs)

        named_tuple = sample_from_prior(rng, tm2)
        @test typeof(named_tuple[:a]) == Float64
        @test typeof(named_tuple[:x]) == typeof(rand(MvNormal(D, 1.0)))
    end

    @testset "Prior density from Turing" begin
        D = 5

        @model function test_model(y)
            a ~ MvNormal(D, 1)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
        end

        y_obs = 2
        tm = test_model(y_obs)
        logprior_density = make_log_prior_density(tm)

        xval = 1
        aval = ones(D)
        nt = (x = xval, a = aval)
        prior_hand(nt) = logpdf(Normal(0, 1), nt[:x]) + logpdf(MvNormal(D, 1), nt[:a])
        @test logprior_density(nt) == prior_hand(nt)

        # Test that we can evaluate samples from prior
        nt = sample_from_prior(rng, tm)
        @test logprior_density(nt) == prior_hand(nt)
    end

    @testset "Joint density from Turing" begin
        @model function test_model(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
        end

        y_obs = 2
        tm = test_model(y_obs)
        logjoint_density = make_log_joint_density(tm)

        xval = 1
        nt = (x = xval,)
        @test logjoint_density(nt) == logpdf(Normal(0, 1), xval) + logpdf(Normal(xval, 1), y_obs)
    end

    @testset "Turing AnnealedISSampler" begin
        N = 3
        num_samples = 10

        @model function test_model(y)
            x ~ Normal(0, 1)
            z ~ MvNormal(zeros(2))
            y ~ Normal(x, 1)
        end

        y_obs = 2
        tm = test_model(y_obs)

        ais = AnnealedISSampler(tm, N)
        samples, diagnostics = ais_sample(rng, ais, num_samples)

        chains = sample(rng, tm, AnIS(N), num_samples)

        @test typeof(chains) <: MCMCChains.Chains
        @test typeof(chains.logevidence) <: Real
        @test haskey(chains.info, :ess)
        @test haskey(get(chains, :log_weight), :log_weight)
        @test haskey(get(chains, :x), :x)

        # Test intermediate samples
        chains = sample(
            rng, tm, AnIS(20), num_samples; store_intermediate_samples=true
        )
        @test haskey(chains.info, :intermediate_samples)

        inter_samples = hcat(chains.info[:intermediate_samples]...)
        inter_samples = permutedims(inter_samples, [2, 1])
        @test size(inter_samples) == (num_samples, 11)

        @test haskey(chains.info, :intermediate_log_weights)
        inter_weights = hcat(chains.info[:intermediate_log_weights]...)
        inter_weights = permutedims(inter_weights, [2, 1])
        @test size(inter_weights) == (num_samples, 11)
    end

    @testset "Transition kernels" begin
        prior_sample = (a = 0.0, b = [1; 2])
        expected_transition_kernel = (a = Normal(0, 1), b = MvNormal(2, 1.0))
        transition_kernel = AnnealedIS.get_normal_transition_kernel(prior_sample, 1)
        @test expected_transition_kernel == transition_kernel
    end

    @testset "Intermediate Samples" begin
        N = 20
        num_samples = 10

        @model function test_model(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
        end

        y_obs = 2
        tm = test_model(y_obs)

        ais = AnnealedISSampler(tm, N)
        samples, diagnostics = ais_sample(
            rng, ais, num_samples; store_intermediate_samples=true)

        inter_samples = hcat(diagnostics[:intermediate_samples]...)
        inter_samples = permutedims(inter_samples, [2, 1])
        @test size(inter_samples) == (num_samples, 11)
        @test map(x->x.params, samples) == inter_samples[:,end]

        @test haskey(diagnostics, :intermediate_log_weights)
    end

    @testset "Resampling Tests" begin
        D = 1
        N = 10
        num_samples = 100

        y_obs = [3.0]

        # Roughly half the samples should get rejected.
        prior_sampling(rng) = rand(rng, MvNormal(D, 1.0))
        prior_density(params) = logpdf(MvNormal(D, 1.0), params)
        joint_density(params) = logpdf(MvNormal(params, I), y_obs) + prior_density(params) + log(max(params[1],0))

        ais_simple_rejection = AnnealedISSampler(
            prior_sampling,
            prior_density,
            joint_density,
            N
        )
        samples_no_resampling, no_resampling_diag = ais_sample(
            rng,
            ais_simple_rejection,
            num_samples
        )

        ais_resampling = AnnealedISSampler(
            prior_sampling,
            prior_density,
            joint_density,
            N,
            RejectionResample()
        )
        samples_resampling, resampling_diag = ais_sample(rng, ais_resampling, num_samples)

        non_zero_weights(samples) = sum(samples) do weighted_sample
            Int(weighted_sample.log_weight != -Inf)
        end

        @test non_zero_weights(samples_resampling) == num_samples
        @test sum(resampling_diag[:num_rejected]) > 0

        num_accepted_samples = non_zero_weights(samples_no_resampling)
        @test num_accepted_samples < num_samples
        @test (num_samples - sum(no_resampling_diag[:num_rejected])) == num_accepted_samples
        @test maximum(no_resampling_diag[:num_rejected]) == 1
    end

    @testset "Out of distribution catch" begin
        @model function model()
            x ~ LogNormal(1, 1)
        end

        prior_density = make_log_prior_density(model())
        joint_density = make_log_joint_density(model())

        @test prior_density((x = -1,)) == -Inf
        @test joint_density((x = -1,)) == -Inf
    end

    @testset "AnISHMC" begin
        betas = [0.0, 0.5, 1.0]
        proposal = AdvancedHMC.StaticTrajectory(AdvancedHMC.Leapfrog(0.05), 10)
        anis = AnISHMC(
            betas,
            proposal,
            10,
            SimpleRejection()
        )
        @test anis.betas[1] == 0.0
    end

    @testset "AnnealedISSamplerHMC" begin
        N = 100
        num_samples = 10

        @model function test_model(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
        end

        yval = 3
        tm = test_model(yval)

        #betas = collect(range(0, 1, length=100))
        betas = (geomspace(1, 1001, 101) .- 1) ./ 1000
        proposal = AdvancedHMC.StaticTrajectory(AdvancedHMC.Leapfrog(0.05), 10)
        anis = AnISHMC(
            betas,
            proposal,
            5,
            SimpleRejection()
        )
        spl = Turing.DynamicPPL.Sampler(anis, tm)
        anis_sampler = AnnealedISSamplerHMC(tm, anis, spl)

        # Check that hamiltonians are correct. The first hamiltonian should
        # have the target density specified by the first non-zero beta.
        prior_sample = anis_sampler.prior_sampling()
        for i in 2:(length(betas)-1)
            h = anis_sampler.hamiltonians[i-1]
            d = AnnealedIS.logdensity(anis_sampler, i, prior_sample)
            @test h.ℓπ(prior_sample) == d

            dr = AnnealedIS.logdensity_ratio(anis_sampler, i, prior_sample)
            hp1 = anis_sampler.hamiltonians[i]
            hdr = hp1.ℓπ(prior_sample) - h.ℓπ(prior_sample)
            @test isapprox(dr, hdr, atol=1e-12)
        end

        chain = sample(
            Random.GLOBAL_RNG,
            test_model(yval),
            anis,
            num_samples;
            store_intermediate_samples=true
        )

        acceptance_rates = chain.info[:intermediate_acceptance_rate]
        @test haskey(chain.info, :intermediate_acceptance_rate)
        @test haskey(chain.info, :acceptance_rate)

        #samples = Array(chain[:x])[:,1]
        #log_weights = Array(chain[:log_weight])[:,1]

        #posterior_mean = dot(samples, exp.(log_weights)) / (num_samples * exp(chain.logevidence))
        #@test isapprox(posterior_mean, 1.5; atol=3e-2)
    end

    # @testset "Constrained Variables" begin
    #     @model function test_model(y)
    #         x ~ Beta(1, 1)
    #         y ~ Normal(x, 1)
    #     end
    #     yval = 3

    #     betas = [0.0, 0.5, 1.0]
    #     proposal = AdvancedHMC.StaticTrajectory(AdvancedHMC.Leapfrog(0.05), 10)
    #     anis = AnISHMC(
    #         betas,
    #         proposal,
    #         10,
    #         SimpleRejection()
    #     )

    #     #chain = sample(Random.GLOBAL_RNG, test_model(yval), anis, 10)
    #     #xs = Array(chain[:x])
    #     #@test all(0 .< xs) && all(xs .<= 1)

    #     true_logjoint(x) = logpdf(Beta(1, 1), x) + logpdf(Normal(x, 1), yval)
    #     tm = test_model(yval)
    #     spl = Turing.DynamicPPL.Sampler(anis, tm)
    #     Turing.DynamicPPL.link!(spl.state.vi, spl)
    #     logjoint = AnnealedIS.gen_logjoint(spl.state.vi, tm, spl)

    #     @test logjoint([2.0]) == true_logjoint(2)
    #     @show logjoint([2.0])
    #     @test logjoint([0.5]) == true_logjoint(0.5)
    # end

    @testset "Density closures" begin
        @model function test_model(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
        end

        betas = [0.0, 0.5, 1.0]
        proposal = AdvancedHMC.StaticTrajectory(AdvancedHMC.Leapfrog(0.05), 10)
        anis = AnISHMC{Turing.Core.ReverseDiffAD{false}}(
            betas,
            proposal,
            10,
            SimpleRejection()
        )

        yval = 3
        tm = test_model(yval)
        spl = Turing.DynamicPPL.Sampler(anis, tm)

        vi = spl.state.vi
        logjoint = AnnealedIS.gen_logjoint(vi, tm, spl)
        logprior = AnnealedIS.gen_logprior(vi, tm, spl)

        true_logjoint(x) = logpdf(Normal(0, 1), x[1]) + logpdf(Normal(x[1], 1), yval)
        @test logjoint([2.0]) == true_logjoint([2])

        true_logprior(x) = logpdf(Normal(0, 1), x[1])
        @test logprior([2.0]) == true_logprior([2])

        # Tests to check whether functions are ThreadSafe.
        num_vals = 20
        test_vals = rand(num_vals)
        test_vals = [[x] for x in test_vals]

        vi = Turing.VarInfo(tm)
        logjoint = AnnealedIS.gen_logjoint(vi, tm, spl)
        logjoint_grad = AnnealedIS.gen_logjoint_grad(vi, tm, spl)

        logprior = AnnealedIS.gen_logprior(vi, tm, spl)
        logprior_grad = AnnealedIS.gen_logprior_grad(vi, tm, spl)

        for fn in [logjoint, logjoint_grad, logprior, logprior_grad]
            vals_seq = map(fn, test_vals)
            vals_parallel = similar(vals_seq)
            Threads.@threads for i in 1:num_vals
                vals_parallel[i] = fn(test_vals[i])
            end
            @test vals_seq == vals_parallel

            vals_broadcast = fn.(test_vals)
            @test vals_seq == vals_broadcast
        end
    end

    @testset "AD backends" begin
        @model function test_model(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
        end

        yval = 3
        tm = test_model(yval)

        num_vals = 20
        test_vals = rand(num_vals)
        test_vals = [[x] for x in test_vals]

        # Forward mode AD
        betas = [0.0, 0.5, 1.0]
        proposal = AdvancedHMC.StaticTrajectory(AdvancedHMC.Leapfrog(0.05), 10)
        anis_fw = AnISHMC{Turing.ForwardDiffAD{1}}(
            betas,
            proposal,
            10,
            SimpleRejection()
        )
        spl = Turing.DynamicPPL.Sampler(anis_fw, tm)
        vi = spl.state.vi

        logjoint_grad_fw = AnnealedIS.gen_logjoint_grad(vi, tm, spl)
        logprior_grad_fw = AnnealedIS.gen_logprior_grad(vi, tm, spl)

        joint_grads_fw = logjoint_grad_fw.(test_vals)
        prior_grads_fw = logprior_grad_fw.(test_vals)

        # ReverseDiff AD
        anis_rv = AnISHMC{Turing.Core.ReverseDiffAD{false}}(
            betas,
            proposal,
            10,
            SimpleRejection()
        )
        spl = Turing.DynamicPPL.Sampler(anis_rv, tm)
        vi = spl.state.vi

        logjoint_grad_rv = AnnealedIS.gen_logjoint_grad(vi, tm, spl)
        logprior_grad_rv = AnnealedIS.gen_logprior_grad(vi, tm, spl)

        joint_grads_rv = logjoint_grad_rv.(test_vals)
        prior_grads_rv = logprior_grad_rv.(test_vals)

        @test joint_grads_fw == joint_grads_rv
        @test prior_grads_fw == prior_grads_rv

        # Some sanity check that the joint and prior are different
        @test prior_grads_rv != joint_grads_rv
        @test prior_grads_fw != joint_grads_fw
    end
end


module ManufacturedSolMHDTest

using Gridap
using Gridap.FESpaces: residual
using Test

u(x) = VectorValue(x[1]^2+x[2]^2,-x[1]-x[2]^2,one(x[1])+x[2]^2)
p(x) = x[1]+x[2]+x[3]
j(x) = VectorValue(x[1]^2+x[2],-x[1]-x[2]^2,one(x[1])+x[3]^2)
φ(x) = x[1]+x[2]+x[3]
B = VectorValue(1.0,1.0,1.0)

∇u(x) = ∇(u)(x)
Δu(x) = Δ(u)(x)
∇p(x) = ∇(p)(x)
∇φ(x) = ∇(φ)(x)

f_u(x) = (∇u(x)')⋅u(x) - Δu(x) + ∇p(x) - j(x)×B
f_p(x) = (∇⋅u)(x)
f_j(x) = j(x) + ∇φ(x) - u(x)×B
f_φ(x) = (∇⋅j)(x)

g_u(x) = u(x)
g_j(x) = j(x)

n = 4
domain = (-0.5,0.5,-0.5,0.5,-0.5,0.5)
partition = (n,n,n)
model = CartesianDiscreteModel(domain,partition)

dirichlet_u = collect(1:24)
dirichlet_j = collect(1:24)
neumann_u = collect(25:26)
neumann_j = collect(25:26)

labels = get_face_labeling(model)
add_tag_from_tags!(labels,"dirichlet_u",dirichlet_u)
add_tag_from_tags!(labels,"dirichlet_j",dirichlet_j)
add_tag_from_tags!(labels,"neumann_u",neumann_u)
add_tag_from_tags!(labels,"neumann_j",neumann_j)

order = 2

Vu = FESpace(model, ReferenceFE(lagrangian,VectorValue{3,Float64},order);
        conformity=:H1, dirichlet_tags="dirichlet_u")

if length(neumann_u) == 0
  Vp = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:P);
          conformity=:L2, constraint=:zeromean)
else
  Vp = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:P);
          conformity=:L2)
end

Vj = FESpace(model, ReferenceFE(raviart_thomas,Float64,order-1);
    conformity=:Hdiv, dirichlet_tags="dirichlet_j")

if length(neumann_j) == 0
  Vφ = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:Q);
      conformity=:L2, constraint=:zeromean)
else
  Vφ = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:Q);
      conformity=:L2)
end

U = TrialFESpace(Vu,g_u)
P = TrialFESpace(Vp)
J = TrialFESpace(Vj,g_j)
Φ = TrialFESpace(Vφ)

Y = MultiFieldFESpace([Vu, Vp, Vj, Vφ])
X = MultiFieldFESpace([U, P, J, Φ])

trian = Triangulation(model)
degree = 2*(order)
dΩ = Measure(trian,degree)

uk = interpolate(u,U)

function a(X,Y)
  u  , p  , j  , φ   = X
  v_u, v_p, v_j, v_φ = Y

  ∫((∇(u)'⋅uk)⋅v_u + inner(∇(u),∇(v_u)) - p*(∇⋅v_u) - (j×B)⋅v_u +
    (∇⋅u)*v_p +
    j⋅v_j - φ*(∇⋅v_j) - (u×B)⋅v_j +
    (∇⋅j)*v_φ )*dΩ
end

function l(Y)
  v_u, v_p, v_j, v_φ = Y

  ∫(v_u⋅f_u + v_p*f_p + v_j⋅f_j + v_φ*f_φ)*dΩ
end

if length(neumann_u) > 0
  btrian_u = BoundaryTriangulation(model,tags=["neumann_u"])
  dΓ_u = Measure(btrian_u,degree)
  nb_u = get_normal_vector(btrian_u)

  function l_Γ_u(Y)
    v_u, v_p, v_j, v_φ = Y

    ∫(v_u⋅(nb_u⋅∇u) - (nb_u⋅v_u)*p)*dΓ_u
  end
end

if length(neumann_j) > 0
  btrian_j = BoundaryTriangulation(model,tags=["neumann_j"])
  dΓ_j = Measure(btrian_j,degree)
  nb_j = get_normal_vector(btrian_j)

  function l_Γ_j(Y)
    v_u, v_p, v_j, v_φ = Y

    ∫(-(v_j⋅nb_j)*φ)*dΓ_j
  end
end

if (length(neumann_u) > 0) & (length(neumann_j) > 0)
  b(Y) = l(Y) + l_Γ_u(Y) + l_Γ_j(Y)
elseif length(neumann_u) > 0
  b(Y) = l(Y) + l_Γ_u(Y)
elseif length(neumann_j) > 0
  b(Y) = l(Y) + l_Γ_j(Y)
else
  b(Y) = l(Y)
end

op  = AffineFEOperator(a,b,X,Y)

xh = solve(op)
uh, ph, jh, φh = xh

r = residual(op,xh)

rh = FEFunction(Y,r)
ruh, rph, rjh, rφh = rh

eu = uh - u
ep = ph - p
ej = jh - j
eφ = φh - φ

l2(v) = sqrt(sum(∫(v⋅v)*dΩ ))
h1(v) = sqrt(sum(∫(v⋅v + inner(∇(v),∇(v)) )*dΩ ))
hdiv(v) = sqrt(sum(∫(v⋅v + inner((∇⋅v),(∇⋅v)) )*dΩ ))

eu_l2 = l2(eu)
eu_h1 = h1(eu)
ep_l2 = l2(ep)
ej_l2 = l2(ej)
ej_hdiv = hdiv(ej)
eφ_l2 = l2(eφ)

@test eu_l2 < 1e-12
@test eu_h1 < 1e-12
@test ep_l2 < 1e-12
@test ej_l2 < 1e-12
@test ej_hdiv < 1e-12
@test eφ_l2 < 1e-12

end #module


using StochasticOptimization
using Base.Test

using ObjectiveFunctions
using Transformations.TestTransforms
using CatViews
using MLDataPattern
# using MLDataUtils

# include("tst_iteration.jl")
# Stop the tests
# error()

using Plots; unicodeplots(show=true,leg=false)

@testset "Rosenbrock-2" begin

    srand(1)
    n = 2
    t = tfunc(rosenbrock, n, rosenbrock_gradient)
    @show t
    # t = rosenbrock_transform(n)
    # obj = objective(t, NoLoss())

    # random starting values
    θ = params(t)
    startvals = 8rand(n)-4
    @show θ startvals

    # build a MetaLearner to use RMSProp w/ fixed learning rate,
    # setting max iterations, a custom convergence check, and a
    # custom iteration callback to collect data to plot
    converged = (m,i) -> totalcost(m) < 1e-6
    maxiter = 50000

    # this problem has no input (we're learning the params only),
    # and we know the minimum is zero, so we forever pull from this
    # fixed (inputs,targets) pair
    # data = zeros(0,1),zeros(1,1)

    # test the choices of ParamUpdaters
    for (T, lr) in [
                    (SGD, 5e-4),
                    (Adagrad, 1e-0),
                    (Adadelta, 1e-3),
                    (Adam, 1e-2),
                    (Adamax, 1e-2),
                    (RMSProp, 1e-3),
                    ]
        @show T,lr
        learner = make_learner(
            GradientLearner(lr, T()),
            # TimeLimit(10),
            # ShowStatus(1000),
            maxiter = maxiter,
            converged = converged
        )

        # learn forever (our maxiter and converge sub-learners will stop us)
        θ[:] = startvals
        learn!(t, learner) #, infinite_obs(data))

        tc = totalcost(t)
        @show tc
        @test 0 < tc < 1e-3
    end

    # rerun while tracking x/y
    x,y = zeros(0),zeros(0)
    learner = make_learner(
        GradientLearner(1e-1, Adamax()),
        maxiter = 50000,
        converged = converged,
        oniter = (m,i) -> begin
            θ = params(m)
            push!(x, θ[1])
            push!(y, θ[2])
            if mod1(i,2000)==1
                println("Iter: $i Loss: $(output_value(m)[1]) θ: $θ")
            end
        end
    )

    # learn forever (our maxiter and converge sub-learners will stop us)
    θ[:] = startvals
    learn!(t, learner)

    tc = totalcost(t)
    @show tc
    @test 0 < tc < 1e-3

    # plot our path to solution
    plot(x,y, ann=[(θ..., text("$θ", :left))])
end

@testset "LinReg" begin
    nin, nout = 10, 1

    # build our objective
    t = Affine(nin,nout)
    l = L2DistLoss()
    p = scaled(L1Penalty(), 1e-8)
    obj = objective(t, l, p)

    # create some fake data... affine transform plus noise
    # note: θ is the "true params"
    τ = 1000
    θ = randn(nout*(nin+1))
    w, b = splitview(θ, ((nout,nin),(nout,)))[1]
    inputs = randn(nin, τ)
    targets = w * inputs + repmat(b, 1, τ) + 0.1randn(nout, τ)

    # add norms to a trace vector
    # normvals = zeros(0)
    tracer = Tracer(Float64, (model,i) -> norm(θ - params(model)))

    # the MetaLearner have a bunch of specialized sub-learners
    learner = make_learner(
        GradientLearner(FixedLR(5e-3), Adamax()),
        ShowStatus(40),
        tracer,
        ConvergedTo(params, θ, tol=0.1, every=20),
        maxiter=5000,
    )

    learn!(obj, learner, RandomBatches(ObsView((inputs, targets)), size=20))

    # some summary output:

    println()
    plot(tracer.storage, title = "‖θₜᵣᵤₑ - θ‖²", xguide="Iteration")

    # scatter predicted output vs ground truth... should be diagonal line
    est_w, est_b = t.params.views
    pred = est_w * inputs + repmat(est_b, 1, τ)
    truth = w * inputs + repmat(b, 1, τ)
    @test maximum(pred - truth) < 5e-1

    println()
    plot(pred', truth', t=:scatter, xguide="Predicted Output", yguide="Actual Output")
end


function dDrs(d1, d2)
    n, m = map(min, size(d1), size(d2))
    weights = (m:-1:1) ./ m
    @views vd1 = d1[1:n, 1:m]
    @views vd2 = d2[1:n, 1:m]
    abs2.(vd1 - vd2) * weights ./ sum(weights) .|> sqrt
end

function analyze(rs, ϕs; fs = [normK, normT, normH], isD=true)
    label = ["ϕ$i" for i in 1:length(ϕs)]
    for (ϕ, lbl) in zip(ϕs, label)
        scatter(Tuple.(ϕ); title=lbl, label=nothing, markersize=2) |> display
    end

    for f in fs
        vs = [f(rs, ϕ) for ϕ in ϕs]
        plot(rs, vs; label=permutedims(label), title="$f(r)") |> display
    end

    if isD
        ds = [normD(length(ϕ) - 1, rs, ϕ) for ϕ in ϕs]
        [plot(ds[i], label=nothing, title=label[i]) for i in 1:length(ϕs)] .|> display

        d1 = ds[1]
        d2 = ds[end]

        dd = dDrs(d1, d2)
        plot(dd, title="∑Dk(r)/k") |> display
    end
    return
end

lossK(ϕ, p) = mapreduce((k, k0) -> abs2(k - k0), +, normK(p.rs, ϕ), p.k0) / length(p.rs)
lossH(ϕ, p) = mapreduce((h, h0) -> abs2(h - h0), +, normH(p.rs, ϕ), p.h0) / length(p.rs)

function errD(d1, d2)
    ddrs = PointsReconstruction.dDrs(d1, d2)
    sum(ddrs) / length(ddrs)
end

lossD(ϕ, p) = errD(normD(p.maxk, p.rs, ϕ), p.d0)

function default_loss_params(rs, ϕ; isK=false, isH=false, isD=false, isT=false, maxk=length(ϕ) - 1)
    p = (;rs)

    if isK
        k = (; k0 = normK(rs, ϕ))
        p = merge(p, k)
    end

    if isH
        h = (; h0 = normH(rs, ϕ))
        p = merge(p, h)
    end

    if isD
        d = (; maxk, d0 = normD(maxk, rs, ϕ))
        p = merge(p, d)
    end

    p
end


module StateSpaceSimulation

using Modia
@usingModiaPlot

# Second order system:
#     w = 1.0
#     D = 0.1
#     k = 2.0
#     der(x1) = x2
#     der(x2) = -w^2*x1 - 2*D*w*x2 + w^2*u
#          y  = k*x1
SecondOrder1 = Model(
    w = 20.0,
    D =  0.1,
    k =  2.0,
    sys = Modia.StateSpace | Map(A = :([  0        1;
                                       -w^2  -2*D*w]),
                                 B = :([0; w^2]),
                                 C = :([k 0]),
                                 D = :(zeros(1)),
                                 x = Var(init = zeros(2)) ),
    equations = :[sys.u = 1.0]
)

secondOrder1 = @instantiateModel(SecondOrder1, unitless=true)
simulate!(secondOrder1, merge=Map(D=0.3),
          requiredFinalStates = [0.9974089572681231, 0.011808652820321723])
plot(secondOrder1, [("sys.u", "sys.y"); "sys.x"], figure = 1)

SecondOrder2 = Model(
    sys = Modia.StateSpace | Map(   A = [   0.0    1.0;
                                         -100.0  -12.0],
                                    B = [0; 400.0],
                                    C = [2.0 0],
                                    D = zeros(1,1),
                                    x = Var(init = zeros(2))),
    equations = :[sys.u = 1.0]
)

secondOrder2 = @instantiateModel(SecondOrder2, unitless=true)
simulate!(secondOrder2, Tsit5(), merge=Map(sys=Map(A = [0.0 1.0; -400.0  -12.0])),
          requiredFinalStates = [0.9974089572681231, 0.011808652820321723])
plot(secondOrder2, [("sys.u", "sys.y"); "sys.x"], figure = 2)

simulate!(secondOrder2, Tsit5(), merge=Map(sys=Map(x=[1.0,1.0])),
          logParameters=true, logStates=true,
          requiredFinalStates = [1.0000295203337868, 0.0022367686372974493])
plot(secondOrder2, [("sys.u", "sys.y"); "sys.x"], figure = 3)
end


module TestStateSpace
using CustomTest
using Base.Test
using ControlSystems
# Naming convention:
# ------------------
# {type}[S]_{dims}[_d][_n]
# type: C: Continuous, D: Discrete
# scalar: add "S" if any input is scalar
# dims: "nxnuny"
# feedthrough: append "_d" if `D` is present
# names: append "_n" if some inputs/outputs/states are named

# SCALARS
a_2 = [-5 -3; 2 -9]
CS_111 = ss(-5, 2, 3, [0])
CS_111_d = ss([3], 2, 1, 1)
CS_211 = ss(a_2, [1; 2], [1 0], 0)
CS_221 = ss(a_2, [1 0; 0 2], [1 0], 0)
CS_222 = ss(a_2, [1 0; 0 2], eye(2), 0)

# CONTINUOUS
a_1 = [-5]
C_111 = ss(a_1, [2], [3], [0])
C_211 = ss(a_2, [1; 2], [1 0], [0])
C_212 = ss(a_2, [1; 2], eye(2), [0; 0])
C_221 = ss(a_2, [1 0; 0 2], [1 0], [0 0])
C_222 = ss(a_2, [1 0; 0 2], eye(2), zeros(2,2))
C_222_d = ss(a_2, [1 0; 0 2], eye(2), eye(2))
C_022 = ss(4*eye(2))

# DISCRETE
da_1 = [-0.5]
da_2 = [0.2 -0.8; -0.8 0.07]
D_111 = ss(da_1, [2], [3], [0], 0.005)
D_211 = ss(da_2, [1; 2], [1 0], [0], 0.005)
D_221 = ss(da_2, [1 0; 0 2], [1 0], [0 0], 0.005)
D_222 = ss(da_2, [1 0; 0 2], eye(2), zeros(2,2), 0.005)
D_222_d = ss(da_2, [1 0; 0 2], eye(2), eye(2), 0.005)
D_022 = ss(4*eye(2), 0.005)

# Definition of input, output and state names
C_222_d_n = ss(a_2, [1 0; 0 2], eye(2), eye(2),
  statenames=["i","u"],inputnames=String("e"),outputnames="theta")

# TESTS
# Contstuct with scalars
@test CS_111 == C_111
@test CS_111_d == ss([3],[2],[1],[1])
@test CS_211 == C_211
@test CS_221 == C_221
@test CS_222 == C_222

# Addition
@test C_111 + C_111 == ss([-5 0; 0 -5],[2; 2],[3 3],[0])
@test C_222 + C_222 == ss([-5 -3 0 0; 2 -9 0 0; 0 0 -5 -3;
        0 0 2 -9],[1 0; 0 2; 1 0; 0 2], [1 0 1 0; 0 1 0 1],[0 0; 0 0])
@test C_222 + 1 == ss([-5 -3; 2 -9],[1 0; 0 2],[1 0; 0 1],[1 1; 1 1])
@test D_111 + D_111 == ss([-0.5 0; 0 -0.5],[2; 2],[3 3],[0], 0.005)

# Subtraction
@test C_111 - C_211 == ss([-5 0 0; 0 -5 -3; 0 2 -9],[2; 1; 2],[3 -1 -0],[0])
@test 1 - C_222 == ss([-5 -3; 2 -9],[1 0; 0 2],[-1 -0; -0 -1],[1 1; 1 1])
@test D_111 - D_211 == ss([-0.5 0 0; 0 0.2 -0.8; 0 -0.8 0.07],[2; 1; 2],
        [3 -1 -0],[0], 0.005)

# Multiplication
@test C_111 * C_221 == ss([-5 2 0; 0 -5 -3; 0 2 -9],
        [0 0; 1 0; 0 2],[3 0 0],[0 0])
@test C_212 * C_111 == ss([-5 -3 3; 2 -9 6; 0 0 -5],
        [0; 0; 2],[1 0 0; 0 1 0],[0; 0])
@test 4*C_222 == ss([-5 -3; 2 -9],[1 0; 0 2],[4 0; 0 4],[0 0; 0 0])
@test D_111 * D_221 == ss([-0.5 2 0; 0 0.2 -0.8; 0 -0.8 0.07],
        [0 0; 1 0; 0 2],[3 0 0],[0 0],0.005)

# Division
@test 1/C_222_d == ss([-6 -3; 2 -11],[1 0; 0 2],[-1 0; -0 -1],[1 -0; 0 1])
@test C_221/C_222_d == ss([-5 -3 -1 0; 2 -9 -0 -2; 0 0 -6 -3;
        0 0 2 -11],[1 0; 0 2; 1 0; 0 2],[1 0 0 0],[0 0])
@test 1/D_222_d == ss([-0.8 -0.8; -0.8 -1.93],[1 0; 0 2],[-1 0; -0 -1],
        [1 -0; 0 1],0.005)

# Indexing
@test size(C_222) == (2, 2)
@test size(C_212) == (2, 1)
@test C_222[1,1] == ss([-5 -3; 2 -9],[1; 0],[1 0],[0])
@test C_222[1:1,1] == ss([-5 -3; 2 -9],[1; 0],[1 0],[0])
@test C_222[1,1:2] == C_221
@test size(C_222[1,[]]) == (1,0)

# Naming signals
@test C_222_d_n.statenames == String["i","u"]
@test C_222_d_n.inputnames == String["e1","e2"]
@test C_222_d_n.outputnames == String["theta1","theta2"]

# Printing
res = ("StateSpace:\nA = \n          x1      x2 \n  x1   -5.0    -3.0  \n  x2"*
       "    2.0    -9.0  \nB = \n         u1     u2 \n  x1   1.0    0.0  \n"*
       "  x2   0.0    2.0  \nC = \n         x1     x2 \n  y1   1.0    0.0  \n"*
       "  y2   0.0    1.0  \nD = \n         u1     u2 \n  y1   0.0    0.0  \n"*
       "  y2   0.0    0.0  \n\nContinuous-time state-space model")
@test sprint(show, C_222) == res
res = ("StateSpace:\nA = \n          x1      x2 \n  x1    0.2    -0.8  \n  x2"*
       "   -0.8     0.07 \nB = \n         u1     u2 \n  x1   1.0    0.0  \n"*
       "  x2   0.0    2.0  \nC = \n         x1     x2 \n  y1   1.0    0.0  \n"*
       "  y2   0.0    1.0  \nD = \n         u1     u2 \n  y1   0.0    0.0  \n"*
       "  y2   0.0    0.0  \n\nSample Time: 0.005 (seconds)\n"*
       "Discrete-time state-space model")
@test sprint(show, D_222) == res
res = ("StateSpace:\nD = \n         u1     u2 \n  y1   4.0    0.0  \n  y2  "*
       " 0.0    4.0  \n\nStatic gain")
@test sprint(show, C_022) == res
res = ("StateSpace:\nD = \n         u1     u2 \n  y1   4.0    0.0  \n  y2  "*
       " 0.0    4.0  \n\nSample Time: 0.005 (seconds)\nStatic gain")
@test sprint(show, D_022) == res

# Errors
@test_throws ErrorException C_111 + C_222             # Dimension mismatch
@test_throws ErrorException C_111 - C_222             # Dimension mismatch
@test_throws ErrorException C_111 * C_222             # Dimension mismatch
@test_throws ErrorException D_111 + C_111             # Sampling time mismatch
@test_throws ErrorException D_111 - C_111             # Sampling time mismatch
@test_throws ErrorException D_111 * C_111             # Sampling time mismatch
D_diffTs = ss([1], [2], [3], [4], 0.1)
@test_throws ErrorException D_111 + D_diffTs            # Sampling time mismatch
@test_throws ErrorException D_111 - D_diffTs            # Sampling time mismatch
@test_throws ErrorException D_111 * D_diffTs            # Sampling time mismatch
@test_throws ErrorException 1/C_222                     # Not invertible
@test_throws ErrorException 1/C_212                     # Not invertible
@test_throws ErrorException ss([1 2], [1], [2], [3])      # Not square A
@test_throws ErrorException ss([1], [2 0], [1], [2])      # I/0 dim mismatch
@test_throws ErrorException ss([1], [2], [3 4], [1])      # I/0 dim mismatch
@test_throws ErrorException ss([1], [2], [3], [4], -0.1)  # Negative samping time
@test_throws ErrorException ss(eye(2), eye(2), eye(2), [0]) # Dimension mismatch
end


function _normal_transition_params(model, scheme::ImplicitEulerMaruyama, t0, x0, x1)
  t1 = t0 + scheme.Δt
  μ = x0 + corrected_drift(model, t1, x1, 1.0) * scheme.Δt
  σ = diffusion(model, t1, x1)
  Σ = scheme.Δt * σ * σ'
  μ, Σ
end


function solve_model(M,rates,dens,cons)
    f = (du,u,p,t) -> M.model(du,u,p,t,dens,cons)
    prob = ODEProblem(f,M.u₀,M.tspan,rates)
    sol = solve(prob,AutoVern7(KenCarp4()), abstol=1e-8,reltol=1e-8)
end

function plot_diss_fit!(plt,M,rates,data;cmap=nothing)

    for i in 1:length(M.cons)
        # Setup the ODE problem, then solve
        sol = solve_model(M,rates,M.dens[i],M.cons[i])

        Σ = sum(Array(sol),dims=1)
        Σ = Im2Re(Σ)
        Σ = normalize_lnp_fit(Σ,data.y[i])
        # h = plot!(sol.t,real.(Σ)',ylims=(-5,0))

        if cmap==nothing
            color = i
        else
            color = cmap
        end
        plot!(plt,sol.t,Σ,
            ylims=(0.0,1.0),
            # xlims=(0,15),
            color=color,
            label=string(data.f[i]))
    end
    if cmap==nothing
        colors = collect(1:length(M.cons))'
    else
        colors = cmap
    end
    plot_dissociations!(plt,data,colors)
    # return plt
end

function Im2Re(S)  # ln of the solutions (gets rid of integration instabilities)
    Σ = Complex.(S)
    Σ = real.(log.(Σ)')
end

function normalize_lnp(lnp)
    normlnpv = lnp ./ log(length(lnp)) .+ 1.0
end

function normalize_lnp_fit(lnp_fit,lnp_data)
    normlnpv_fit = lnp_fit ./ log(length(lnp_data)) .+ 1.0
end

function plot_dissociations!(plt,data,colors)
    normlnp = normalize_lnp.(p2lnp.(data.y))
    ε=0.01
    scatter!(plt,data.t,normlnp,
            xlims = (0.0-ε, maximum(maximum.(data.t))+ε),
            ylims = (0.0-ε, 1.0+ε),
            markercolors = colors, #, markerstrokecolor = colors,
            # markersize = 5, markershape=:square,
            # markeralpha = 0., markerstrokealpha = 1.,
            label=false)
end

function show_diss_fit!(plt,M,rates,data;cmap=nothing)

    plot_diss_fit!(plt,M,rates,data;cmap)
    # return plt
end
# function Im2Re(S)
#     Σ = Complex.(S)
#     Σ = real.(Σ)
#     Σ = log.(Σ')
# end
#

# function plot_diss_fit(rates,dens,cons)
#     for (i,F) in enumerate(cons)
#         # Setup the ODE problem, then solve
#         solve_model(rates,dens,cons)
#
#         Σ = Complex.(sum(Array(sol),dims=1))
#         # h = plot!(sol.t,real.(Σ)',ylims=(-5,0))
#         h = plot!(sol.t,real.(log.(Σ)'),
#             ylims=(-6.,0),
#             xlims=(0,15),
#             color=i,
#             label=string(F),
#             palette = :Dark2_5)
#
#         if generate_data
#             push!(dataset,vec(generate_data(sol;logplot=true)))
#         end
#
#         # display(h)
#     end
#     return h
# end
#


using Test
using QuantumOptics
using LinearAlgebra

@testset "nlevel" begin

N = 3
b = NLevelBasis(N)

# Test basis equality
@test b == NLevelBasis(N)
@test b != NLevelBasis(N+1)
@test_throws DimensionMismatch NLevelBasis(0)

# Test transition operator
@test_throws BoundsError transition(b, 0, 1)
@test_throws BoundsError transition(b, N+1, 1)
@test_throws BoundsError transition(b, 1, 0)
@test_throws BoundsError transition(b, 1, N+1)
@test dense(transition(b, 2, 1)) == basisstate(b, 2) ⊗ dagger(basisstate(b, 1))

# Test nlevel states
@test_throws BoundsError nlevelstate(b, 0)
@test_throws BoundsError nlevelstate(b, N+1)
@test norm(nlevelstate(b, 1)) == 1.
@test norm(dagger(nlevelstate(b, 1))*nlevelstate(b, 2)) == 0.
@test norm(dagger(nlevelstate(b, 1))*transition(b, 1, 2)*nlevelstate(b, 2)) == 1.

end # testset


function normal_closure(sys::MomentEquations, binary_vars::Array{Int,1}=Int[])

    closure = OrderedDict()
    closure_exp = OrderedDict()
    N = sys.N

    if !isempty(binary_vars)
        sys = bernoulli_moment_eqs(sys, binary_vars)
    end

    # build symbolic expressions of cumulants up to q_order in terms of central/raw moments
    if sys isa CentralMomentEquations
        moments = sys.M
        K = cumulants_to_central_moments(N, sys.q_order)
    else
        moments = sys.μ
        K = cumulants_to_raw_moments(N, sys.q_order)
    end

    #unique_iter_q = unique(sort(i) for i in sys.iter_q)
    iter_qs = sys.iter_m
    sub = Dict()

    # construct the corresponding truncated expressions of higher order central moments
    for order in sys.m_order+1:sys.q_order

        iter_r = filter(x -> sum(x) == order, sys.iter_q)

        for r in unique(sort(i) for i in iter_r)
            # the last term in the symbolic expression of cumulant κᵣ is Mᵣ (μᵣ)
            # therefore, as we here set κᵣ = 0, only simple manipulation is needed
            closed_moment = -(K[r]-moments[r])
            closed_moment = simplify(closed_moment, expand=true)

            closure[moments[r]] = closed_moment
            closure_exp[moments[r]] = substitute(closed_moment, closure_exp)
            closure_exp[moments[r]] = simplify(closure_exp[moments[r]], expand=true)

            perms = collect(multiset_permutations(r, length(r)))[2:end]

            for iter_perm in perms

                iter_perm_ind = sortperm(sortperm(iter_perm))
                for i in sys.iter_1
                    sub[sys.μ[i]] = sys.μ[i[iter_perm_ind]]
                end
                for i in iter_qs
                    sub[moments[i]] = moments[i[iter_perm_ind]]
                end

                iter_perm = Tuple(iter_perm)
                closure[moments[iter_perm]] = substitute(closure[moments[r]], sub)
                closure_exp[moments[iter_perm]] = substitute(closure_exp[moments[r]], sub)
            end

        end

        iter_qs = vcat(iter_qs, iter_r)

    end

    close_eqs(sys, closure_exp, closure, true)

end


using MovingWeightedLeastSquares

xs = [0.1 * i for i in 0:100]
fs = hcat([sin(x) for x in xs], [cos(x) for x in xs])
println(fs)

m = mwls(xs, fs, 0.5, (d, e) -> exp(-d^2))

println(m(1))


"""
```
    LinearStretching <: AbstractHistogramAdjustmentAlgorithm
    LinearStretching(; [src_minval], [src_maxval],
                       dst_minval=0, dst_maxval=1,
                       no_clamp=false)

    LinearStretching((src_minval, src_maxval) => (dst_minval, dst_maxval))
    LinearStretching((src_minval, src_maxval) => nothing)
    LinearStretching(nothing => (dst_minval, dst_maxval))

    adjust_histogram([T,] img, f::LinearStretching)
    adjust_histogram!([out,] img, f::LinearStretching)
```

Returns an image where the range of the intensities spans the interval [`dst_minval`, `dst_maxval`].

# Details

Linear stretching (also called *normalization*) is a contrast enhancing
transformation that is used to modify the dynamic range of the image. In
particular, suppose that the input image has gray values in the range [A,B] and
one wishes to change the dynamic range to [a,b] using a linear mapping, then the
necessary transformation is given by the relation

```math
f(x) = (x-A) \\frac{b-a}{B-A} + a.
```

# Options

Various options for the parameters of the `adjust_histogram` and
`LinearStretching` type  are described in more detail below.

## Choices for `img`

The function can handle a variety of input types. The returned
image depends on the input type.

For colored images, the input is converted to the
[YIQ](https://en.wikipedia.org/wiki/YIQ)  type and the intensities of the Y
channel are stretched to the specified range. The modified Y channel is then
combined with the I and Q channels and the resulting image converted to the same
type as the input.

## Choices for `dst_minval` and `dst_maxval`

If destination value range `dst_minval` and `dst_maxval` are specified then intensities are
mapped to the range [`dst_minval`, `dst_maxval`]. The default values are 0 and 1.

## Choices for `src_minval` and `src_maxval`

The source value range `src_minval` and `src_maxval` specifies the intensity range of input
image. By default, the values are `extrema(img)` (finite). If custom values are provided,
the output intensity value will be clamped to range `[dst_minval, dst_maxval]` if it exceeds that.

## `no_clamp`

Setting `no_clamp=true` to disable the automatic clamp even if the output intensity value
exceeds the range `[dst_minval, dst_maxval]`. Note that a clamp is still applied for types
that has limited value range, for example, if the input eltype is `N0f8`, then the output will
be clamped to `[0.0N0f8, 1.0N0f8]` even if `no_clamp==true`.

# Example

```julia
using ImageContrastAdjustment, TestImages

img = testimage("mandril_gray")
# Stretches the contrast in `img` so that it spans the unit interval.
imgo = adjust_histogram(img, LinearStretching(dst_minval = 0, dst_maxval = 1))
```

For convenience, Constructing a `LinearStretching` object using `Pair` is also supported

```julia
# these two constructors are equivalent
LinearStretching(src_minval=0.1, src_maxval=0.9, dst_minval=0.05, dst_maxval=0.95)
LinearStretching((0.1, 0.9) => (0.05, 0.95))

# replace the part with `nothing` to use default values, e.g.,
# specify only destination value range
LinearStretching(nothing => (0.05, 0.95))
# specify only source value range and use default destination value range, i.e., (0, 1)
LinearStretching((0.1, 0.9) => nothing)
```

# References
1. W. Burger and M. J. Burge. *Digital Image Processing*. Texts in Computer Science, 2016. [doi:10.1007/978-1-4471-6684-9](https://doi.org/10.1007/978-1-4471-6684-9)

"""
@with_kw struct LinearStretching{T} <: AbstractHistogramAdjustmentAlgorithm
    src_minval::T = nothing
    src_maxval::T = nothing
    dst_minval::T = 0.0f0
    dst_maxval::T = 1.0f0
    minval::T = nothing
    maxval::T = nothing
    no_clamp::Bool = false
    function LinearStretching(src_minval::T1,
                              src_maxval::T2,
                              dst_minval::T3,
                              dst_maxval::T4,
                              minval::T5=nothing,
                              maxval::T6=nothing,
                              no_clamp::Bool=false) where {T1 <: Union{Nothing,Real,AbstractGray},
                                                 T2 <: Union{Nothing,Real,AbstractGray},
                                                 T3 <: Union{Nothing,Real,AbstractGray},
                                                 T4 <: Union{Nothing,Real,AbstractGray},
                                                 T5 <: Union{Nothing,Real,AbstractGray},
                                                 T6 <: Union{Nothing,Real,AbstractGray}}
        # in order to deprecate old fields we have to introduce new fields if we still want to use @with_kw
        # https://github.com/JuliaImages/ImageContrastAdjustment.jl/pull/28#discussion_r395751301
        if !isnothing(minval)
            dst_minval = minval
            Base.depwarn("deprecated: use `dst_minval` for keyword `minval`", :LinearStretching)
        end
        if !isnothing(maxval)
            dst_maxval = maxval
            Base.depwarn("deprecated: use `dst_maxval` for keyword `maxval`", :LinearStretching)
        end

        dst_minval <= dst_maxval || throw(ArgumentError("dst_minval $dst_minval should be less than dst_maxval $dst_maxval"))
        if !(isnothing(src_minval) || isnothing(src_maxval))
            src_minval <= src_maxval || throw(ArgumentError("src_minval $src_minval should be less than src_maxval $src_maxval"))
        end
        T = promote_type(T1, T2, T3, T4, T5, T6)
        new{T}(convert(T, src_minval), convert(T, src_maxval),
               convert(T, dst_minval), convert(T, dst_maxval),
               convert(T, dst_minval), convert(T, dst_maxval),
               no_clamp)
    end
end
function LinearStretching(rangemap::Pair{Tuple{T1, T2}, Tuple{T3, T4}}; no_clamp=false) where {T1, T2, T3, T4}
    LinearStretching(rangemap.first..., rangemap.second..., nothing, nothing, no_clamp)
end
function LinearStretching(rangemap::Pair{Nothing, Tuple{T3, T4}}; no_clamp=false) where {T3, T4}
    LinearStretching(nothing, nothing, rangemap.second..., nothing, nothing, no_clamp)
end
function LinearStretching(rangemap::Pair{Tuple{T1, T2}, Nothing}; no_clamp=false) where {T1, T2}
    LinearStretching(src_minval=rangemap.first[1], src_maxval=rangemap.first[2], no_clamp=no_clamp)
end

function (f::LinearStretching)(out::GenericGrayImage, img::GenericGrayImage)
    T = eltype(out)
    FT = eltype(floattype(T))
    img_min, img_max = minfinite(img), maxfinite(img)
    # explicit annotation is needed because the ?: line mixes three value types:
    # Nothing, T, and typeof(f.src_minval)
    src_minval::FT = isnothing(f.src_minval) ? img_min : f.src_minval
    src_maxval::FT = isnothing(f.src_maxval) ? img_max : f.src_maxval
    dst_minval::FT = f.dst_minval
    dst_maxval::FT = f.dst_maxval

    # the kernel operation `r * x - o` is equivalent to `(x-A) * ((b-a)/(B-A)) + a`
    # precalculate these and make inner loop contains only multiplication and addition
    # to get better performance
    r = (dst_maxval - dst_minval) / (src_maxval - src_minval)
    o = (src_minval*dst_maxval - src_maxval*dst_minval) / (src_maxval - src_minval)

    if 1 ≈ r && 0 ≈ o
        # when image intensity is already adjusted, there's no need to do it again
        # it's a trivial but common case in practice
        out === img || (out .= img)
        return out
    end

    # In most cases, we don't need to clamp the output
    # this is only used when user specifies custom parameters
    out_minval = r * img_min - o
    out_maxval = r * img_max - o
    do_clamp = f.no_clamp ? false : (out_minval < dst_minval) || (out_maxval > dst_maxval)
    # re-adjust clamp option to avoid ArgumentError
    if !do_clamp && (out_minval < typemin(T) || out_maxval > typemax(T))
        do_clamp = true
        dst_minval = convert(typeof(dst_minval), typemin(eltype(T)))
        dst_maxval = convert(typeof(dst_maxval), typemax(eltype(T)))
    else
        do_clamp = do_clamp || out_minval < typemin(T) || out_maxval > typemax(T)
        dst_minval = max(dst_minval, convert(typeof(dst_minval), typemin(eltype(T))))
        dst_maxval = min(dst_maxval, convert(typeof(dst_maxval), typemax(eltype(T))))
    end

    # although slightly faster, clamp before linear stretching has some rounding issue, e.g,
    # it's likely to get results like -9.313226f-10, which can't be directly converted to N0f8
    # thus the line below is commented out and not used
    # do_clamp && (img = clamp.(img, src_minval, src_maxval))

    # tweak the performance of FixedPoint by fusing operations into one broadcast
    # for Float32 the fallback implementation is faster
    if eltype(T) <: FixedPoint
        # ?: is faster than if-else
        @. out = do_clamp ? clamp(r * img - o, dst_minval, dst_maxval) : r * img - o
        return out
    end

    # fallback implementation
    @inbounds @simd for p in eachindex(img)
        val = img[p]
        if isnan(val)
            out[p] = val
        else
            newval = r * val - o
            do_clamp && (newval = clamp(newval, dst_minval, dst_maxval))
            out[p] = T <: Integer ? round(Int, newval) : newval
        end
    end
    out
end

function (f::LinearStretching)(out::AbstractArray{<:Color3}, img::AbstractArray{<:Color3})
    T = eltype(out)
    yiq = convert.(YIQ, img)
    yiq_view = channelview(yiq)
    adjust_histogram!(view(yiq_view,1,:,:), f)
    out .= convert.(T, yiq)
end

(f::LinearStretching)(out::GenericGrayImage, img::AbstractArray{<:Color3}) =
    f(out, of_eltype(Gray, img))


# # Example: Blowfly Model
# Likelihood-free inference for the blow-fly model was introduced by [Simon N.
# Wood](http://dx.doi.org/10.1038/nature09319). We model here the discrete time
# stochastic dynamics of the size ``N`` of an adult blowfly population as given in [section 1.2.3 of the supplementary
# information](https://static-content.springer.com/esm/art%3A10.1038%2Fnature09319/MediaObjects/41586_2010_BFnature09319_MOESM302_ESM.pdf).
# ```math
# N_{t+1} = P N_{t-\tau}\exp(-N_{t-\tau}/N_0)e_t + N_t\exp(-\delta \epsilon_t)
# ```
# where ``eₜ`` and ``ϵₜ`` are independent Gamma random deviates with
# mean 1 and variance ``σp²`` and ``σd²``, respectively.
using Distributions, StatsBase, LikelihoodfreeInference
Base.@kwdef struct BlowFlyModel
    burnin::Int = 50
    T::Int = 1000
end
function (m::BlowFlyModel)(P, N₀, σd, σp, τ, δ)
    p1 = Gamma(1/σp^2, σp^2)
    p2 = Gamma(1/σd^2, σd^2)
    T = m.T + m.burnin + τ
    N = fill(180., T)
    for t in τ+1:T-1
        N[t+1] = P * N[t-τ] * exp(-N[t-τ]/N₀)*rand(p1) + N[t]*exp(-δ*rand(p2))
    end
    N[end-m.T+1:end]
end

# Let us plot four realizations from this model with the same parameters.
using StatsPlots
gr()
m = BlowFlyModel()
plot([plot(m(29, 260, .6, .3, 7, .2),
           xlabel = "t", ylabel = "N", legend = false) for _ in 1:4]...,
     layout = (2, 2))

# To compare different realizations we will use histogram summary statistics.
# In the literature one finds also other summary statistics for this data.
summary_statistics(N) = fit(Histogram, N, 140:16:16140).weights

# We will use a normal prior on log-transformed parameters.
function parameter(logparams)
    lP, lN₀, lσd, lσp, lτ, lδ = logparams
    (P = round(exp(2 + 2lP)),
    N₀ = round(exp(4 + .5lN₀)),
    σd = exp(-.5 + lσd),
    σp = exp(-.5 + lσp),
    τ = round(Int, max(1, min(500, exp(2 + lτ)))),
    δ = exp(-1 + .4lδ))
end
(m::BlowFlyModel)(logparams) = m(parameter(logparams)...)
target(m::BlowFlyModel) = [(log(29) - 2)/2,
                           (log(260) - 4)*2,
                           log(.6) + .5,
                           log(.3) + .5,
                           log(7) - 2,
                           (log(.2) + 1)/.4]
lower(m::BlowFlyModel) = fill(-5., 6)
upper(m::BlowFlyModel) = fill(5., 6)
prior = TruncatedMultivariateNormal(zeros(6), ones(6),
                                    lower = lower(m), upper = upper(m))

# Let us now generate some target data.
model = BlowFlyModel()
x0 = target(model)
data = summary_statistics(model(x0))

# ## Adaptive SMC
smc = AdaptiveSMC(prior = prior)
result = run!(smc, x -> summary_statistics(model(x)), data,
              maxfevals = 2*10^5, verbose = false)
using PrettyTables
pretty_table([[keys(parameter(zeros(6)))...] quantile(smc, .05) median(smc) mean(smc) x0 quantile(smc, .95)],
             ["names", "5%", "median", "mean", "actual", "95%"],
             formatter = ft_printf("%10.3f"))
#-
histogram(smc)
#-
corrplot(smc)

# ## KernelABC
k = KernelABC(prior = prior, delta = 1e-1, K = 10^3, kernel = Kernel())
result = run!(k, x -> summary_statistics(model(x)), data)
pretty_table([[keys(parameter(zeros(6)))...] quantile(k, .05) median(k) mean(k) x0 quantile(k, .95)],
             ["names", "5%", "median", "mean", "actual", "95%"],
             formatter = ft_printf("%10.3f"))
#-
histogram(k)

# ## Kernel Recursive ABC (with callback)
k = KernelRecursiveABC(prior = prior,
                       K = 100,
                       delta = 1e-3,
                       kernel = Kernel(bandwidth = Bandwidth(heuristic = MedianHeuristic(2^3))),
                       kernelx = Kernel());
# We will use a callback here to show how the estimated parameters evolves.
using LinearAlgebra
res_krabc = run!(k, x -> summary_statistics(model(x)), data,
                 maxfevals = 1300,
                 verbose = true,
                 callback = () -> @show norm(k.theta - x0)/norm(x0))


cd(@__DIR__); include("setups/grid23x22.jl")
gr(dpi = 200)

##
frame = sgwt_frame(W; nf = 6)
x = 242
for j = 1:6
    plt = heatmap(reshape(frame[:, x, j], (Nx, Ny))', c = :viridis, ratio = 1,
            frame = :none, xlim = [1, Nx], size = (500, 400))
    savefig(plt, "../figs/Grid$(Nx)x$(Ny)_SGWT_frame_j$(j-1)_x$(x).png")
end


@reexport module LogicBase

export ¬, ∧, ∨, ⊻, ⇒, ⇔, ⋁, ⋀

"""
    ¬(x::Bool) → Bool

A logical "not". Return `false` if the given argument is `true`, and return
`true` if the given argument is `false`.
"""
¬(x::Bool) = ~x

"""
    ∧(x::Bool, y::Bool) → Bool

A logical "and" that does not short-circuit. Return `true` if both arguments are
`true`, and `false` otherwise.
"""
∧(x::Bool, y::Bool) = x & y

"""
    ∨(x::Bool, y::Bool) → Bool

A logical "or" that does not short-circuit. Return `true` if at least one
argument is `true`, and `false` otherwise.
"""
∨(x::Bool, y::Bool) = x | y

"""
    ⊻(x::Bool, y::Bool) → Bool

A logical "exclusive or". Return `true` if exactly one argument is `true`, and
`false` otherwise.
"""
⊻(x::Bool, y::Bool) = x $ y

"""
    ⇒(x::Bool, y::Bool) → Bool

A logical "implies". Return `true` if the first argument is `false`, or if the
second argument is `true`. Return `false` if the first argument is `true` and
the second argument is `false`.
"""
⇒(x::Bool, y::Bool) = ~x | y

"""
    ⇔(x::Bool, y::Bool) → Bool

A logical "if and only if". Return `true` if the first argument is the same as
the second argument.
"""
⇔(x::Bool, y::Bool) = x ≡ y

"""
    ⋁(xs) → Bool

Return `true` if at least one of the `xs` is true.
"""
⋁ = any

"""
    ⋀(xs) → Bool

Return `true` if all of the `xs` are true.
"""
⋀ = all

end  # module


import Interpolations
"""
	interpdf(data,timecol,timvec)

Re-sample dataframe applying linear interpolation to all columns

**Input**
* data: DataFrame where at least one column contains DateTime
* timevec: output time vector (=Vector{DateTime,1})
* timecol: column containing DateTime (default value = :datetime)

**Output**
* dataframe containing all re-sampled columns + timevec

**Example**
```
data = DataFrame(Temp=[10,11,12,14],
			datetime=[DateTime(2010,1,1,0),DateTime(2010,1,1,1),
			DateTime(2010,1,1,2),DateTime(2010,1,1,4)]);
datai = interpdf(data,[DateTime(2010,1,1,12,0,0)]);

```
"""
function interpdf(data::DataFrame,timevec::Vector{DateTime};timecol=:datetime)
	# prepare input time vector for interpolation
	x,xi = preptime(data[timecol],timevec);
	# run for all input columns except for DateTime (x vector) and Types not
	# suitable for interpolation
	dfi = DataFrame(datetime=timevec);
	for i in names(data)
		if (i != timecol) && (eltype(data[i]) <: Real)
			dfi[i] = interp1(x,data[i],xi);
		end
	end
    return dfi
end

"""
	interp1(x,y,xi)
Interpolate 1D data

**Input**
* x: x coordinate vector
* y: y values (vector)
* xi: x coordinate to be interpolated

**Output**
* yi: interpolated value

**Example**
```
yi = interp1([1,2,3,4],[10,20,30,40],1.5);
```
"""
function interp1(x,y,xi)
	# Declare output vector
	out = zeros(length(xi));
	# Create interpolation object
    itp = Interpolations.interpolate((x,),y,
            Interpolations.Gridded(Interpolations.Linear()));
	# exclude (set to NA) extrapolated values
	maxx = maximum(x);minx = minimum(x);
	for i = 1:length(out)
        if (xi[i] <= maxx) && (xi[i] >= minx)
            out[i] = itp(xi[i]);
		else
			out[i] = NaN;
        end
    end
    return out;
end

function interp1(x::Vector{DateTime},y::Vector{Float64},xi::Vector{DateTime})
	t1,t2 = preptime(x,xi);
	return interp1(t1,y,t2)
end
"""
	preptime(dft,timevec)

Auxiliary function to prepare time in DateTime format for interpolation
Will convert DateTime to Number/Int and declare/create output (length) dataframe

"""
function preptime(dft::Vector{DateTime},timevec::Vector{DateTime})
	x = Dates.value.(dft);
	xi = Dates.value.(timevec);
	return x, xi
end

"""
	meshgrid(x,y)

Function to create meshgrid, i.e. returns the same as 2D Matlab meshgrid [function](https://de.mathworks.com/help/matlab/ref/meshgrid.html).

**Input**
* x: vector or range (Int or Float)
* y: vector or range (Int or Float)

**Output**
* xi: matrix with x coordinates  (Int or Float)
* yi: matrix with y coordinates  (Int or Float)

**Example**

```
xi,yi = meshgrid(0:1:10,100:1:110);

```

"""
function meshgrid(x,y)
    xi = [j for i in y, j in x];
    yi = [i for i in y, j in x];
    return xi, yi;
end

"""
	mesh2vec(xi,yi)

Convert meshgrid back to vector

**Input**
* xi: matrix of x coordinates (in mesghrid format)
* yi: matrix of y coordinates (in mesghrid format)

**Output**
* x: vector of x coordinates
* y: vector of x coordinates

**Example**

```
xi = [0 1 2;0 1 2;0 1 2];
yi = [10 10 10;11 11 11;12 12 12];
x,y = mesh2vec(xi,yi);

```
"""
function mesh2vec(x,y)
    return x[1,:], y[:,1]
end

"""
	interp2(x,y,z,xi,yi)

Re-sample 2D data using linear interpolation

**Input**
* x: vector with x coordinates
* y: vector with y coordinates
* z: matrix with (x,y)-dependent values (meshgrid-like)
* xi: new x coordinates
* yi: new y coordinates

**Output**
* zi: interpolated values

**Example**
```
# Prepare input
x = [1.,2.,3.,4.,5.,];
y = [2.,3.,4.,5.,6.,7.,8.];
z = ones(Float64,(length(y),length(x)));
z[2,3] = 2.;
xi = [4.1,3.0,3.0,3.5,5.0,5.01]
yi = [2.1,3.0,3.5,3.5,2.0,3.30]
# compute
zi = interp2(x,y,z,xi,yi)
```
"""
function interp2(x::Vector{Float64},y::Vector{Float64},z::Matrix{Float64},
				 xi::Vector{Float64},yi::Vector{Float64})
	minx,miny,maxx,maxy,itp = prep2interp(x,y,z);
	# declare output
	zi = zeros(Float64,size(xi));
	for i in 1:length(xi[:])
		# do not extrapolate!
		if (xi[i] >= minx && yi[i] >= miny) && (xi[i] <= maxx && yi[i] <= maxy)
		   zi[i] = itp(yi[i],xi[i]) # use yi,xi order as matlab sorts masgrid in different order
		else
		   zi[i] = NaN;
		end
	end
	return zi
end
function interp2(x::Vector{Float64},y::Vector{Float64},z::Matrix{Float64},
				 xi::Matrix{Float64},yi::Matrix{Float64})
	zi = interp2(x,y,z,xi[:],yi[:]);
	return reshape(zi,size(xi));
end
function interp2(x::Matrix{Float64},y::Matrix{Float64},z::Matrix{Float64},
				 xi::Matrix{Float64},yi::Matrix{Float64})
	xm,ym = mesh2vec(x,y);
	return interp2(xm,ym,z,xi,yi);
end
function interp2(x::Matrix{Float64},y::Matrix{Float64},z::Matrix{Float64},
				 xi::Float64,yi::Float64)
	xm,ym = mesh2vec(x,y);
	return interp2(xm,ym,z,xi,yi);
end
function interp2(x::Vector{Float64},y::Vector{Float64},z::Matrix{Float64},
				 xi::Float64,yi::Float64)
	return interp2(x,y,z,[xi],[yi])[1]
end

"""
	Auxilliary function for determination of min and max values of input vectors
"""
function prep2interp(x::Vector{Float64},y::Vector{Float64},z::Matrix{Float64})
    # Need to switch x <--> y as matlab meshgrid uses transposed matrix
    itp = Interpolations.interpolate((y,x),z,
            Interpolations.Gridded(Interpolations.Linear()));
    return minimum(x), minimum(y), maximum(x), maximum(y), itp;
end


function add_rules(from, to, n)
    for α in [0, 1, 2, 3]
        from_rot = rotr90(from, α)
        rules[n][from_rot] = to
        rules[n][flipdim(from_rot, 1)] = to
        rules[n][flipdim(from_rot, 2)] = to
    end
end
parse_array(v, n) = reshape(BitArray([x == '.' ? false : true for x in v]), (n, n))'

d = split.(replace.(readlines("input.txt"), "/", ""), " => ")
rules = Dict(2 => Dict(), 3 => Dict())
for l in d
    n = Int(sqrt(length(l[1])))
    from = parse_array(l[1], n)
    to = parse_array(l[2], n + 1)
    add_rules(from, to, n)
end

function iter(grid)
    s = size(grid)[1]
    t = s % 2 == 0 ? 2 : 3
    squares = Int(s / t)
    new_grid = falses(squares * (t+1), squares * (t+1))
    for x in 1:squares
        for y in 1:squares
            update = rules[t][grid[t*(x-1)+1:t*x, t*(y-1)+1:t*y]]
            new_grid[(t+1)*(x-1)+1:(t+1)*x, (t+1)*(y-1)+1:(t+1)*y] = update
        end
    end
    return new_grid
end

grid = parse_array(".#...####", 3)
for i in 1:18
    grid = iter(grid)
    if i in [5, 18]
        println(sum(grid))
    end
end


""" Sign preserving 1D distance between nodes and a data sample x. """
function distance(nodes::Array{Float64,1},x::Float64)
    return x .- nodes
end

function αexp(s::Int,r::Float64=0.05)
    return exp(-r*s)
end

function αlin(s::Int,smax::Int)
    return 1-(s-1)/smax
end

function maxspread(data::AbstractArray)
    return maximum(data)-minimum(data)
end

function crenel(x::Real, Δ::Real)
    if x > -Δ && x < Δ
        return one(x)
    else
        return zero(x)
    end
end

function average_rounding_error(data::Array{Float64,1},nodes::Array{Float64,1})
    n = length(data)
    S = 0.0
    for i = 1:n
        S += minimum(abs.(distance(nodes,data[i])))/n
    end
    return S
end


# this script tests whether RegressionTables works with GLFixedEffectModels.jl

using GLFixedEffectModels, GLM, RDatasets, Test

# include("src/RegressionTables.jl")
df = dataset("datasets", "iris")
df[!, :isSmall] = df[!, :SepalWidth] .< 2.9
df[!, :isWide] = df[!, :SepalWidth] .> 2.5
df[!, :binary] = @. ifelse(df[!, :PetalWidth] .> 1.0, 1.0, 0.0)

# One FE, Poisson
m = GLFixedEffectModels.@formula SepalLength ~ SepalWidth + GLFixedEffectModels.fe(Species)
rr1 = GLFixedEffectModels.nlreg(df, m, Poisson(), LogLink() , start = [0.2] )
# Two FE, Poisson
m = GLFixedEffectModels.@formula SepalLength ~ SepalWidth + PetalLength + GLFixedEffectModels.fe(Species) +  GLFixedEffectModels.fe(isSmall)
rr2 = GLFixedEffectModels.nlreg(df, m, Poisson(), LogLink() , start = [0.2;0.2] )

m = GLFixedEffectModels.@formula binary ~ SepalWidth + PetalLength + GLFixedEffectModels.fe(Species)
rr3 = GLFixedEffectModels.nlreg(df, m, Binomial(), LogitLink(), GLFixedEffectModels.Vcov.robust() , start = [0.2, 0.2] )

m = GLFixedEffectModels.@formula SepalLength ~ SepalWidth + PetalLength + GLFixedEffectModels.fe(Species) +  GLFixedEffectModels.fe(isSmall)
rr4 = GLFixedEffectModels.nlreg(df, m, Poisson(), LogLink() , start = [0.2;0.2] )

RegressionTables.regtable(rr1,rr2,rr3,rr4; renderSettings = RegressionTables.asciiOutput(joinpath(dirname(@__FILE__), "tables", "glftest1.txt")), regression_statistics = [:nobs, :r2, :adjr2, :r2_within, :f, :p, :f_kp, :p_kp, :dof])
@test checkfilesarethesame(joinpath(dirname(@__FILE__), "tables", "glftest1.txt"), joinpath(dirname(@__FILE__), "tables", "glftest1_reference.txt"))

rm(joinpath(dirname(@__FILE__), "tables", "glftest1.txt"))

# using Econometrics
# data = RDatasets.dataset("Ecdat", "PSID")
# data = data[data.Earnings .> 0 .&
#             data.Kids .< 98,:]


module FuzzySets

name() = return "FuzzySets.jl"

using GLMakie
using LaTeXStrings
import Plots
using FLoops

include("Intervals.jl")
include("FuzzySet.jl")
include("FuzzyNumbers.jl")
include("FuzzyVectors.jl")
include("FuzzyMatrices.jl")
include("FuzzyWeightedAverages.jl")
include("FCM.jl")
include("PCM.jl")
include("IT2FCM.jl")
include("Dampening.jl")
include("LFCM.jl")
include("LPCM.jl")

export
    Interval, .., mid, rad,
    FuzzySet, support, core, height, cut, centroid,
    FuzzyNumber, SingletonFuzzyNumber, TriangularFuzzyNumber, TrapezoidalFuzzyNumber, GaussianFuzzyNumber,
    FuzzyVector,
    peak, draw,
    fuzzy_weighted_average,
    euclidean_distance,
    u_uncertainty
end


# Constants

# Physical parameters

## Materials
const n₁ = 1.49  # acrylate
const n₂ = 1.00  # air

## The sides of the block to be carved
const Block_Side = 0.15 # meters
const Block_Height = Block_Side
const Block_Width = Block_Side

## Depth parameters. If we imagine a reference place, the top face (facing the caustics) is at
## Top_Offset above it. The flat face facing the light source is Bottom below it.
const Top_Offset = 0.01 # 1 centimeters
const Bottom_Offset = 0.02 # 2 centimeters

## Optics
const Focal_Length = 1.0  # meters

const Caustics_Long_Side = 0.20 # m
global Caustics_Height = Caustics_Long_Side
global Caustics_Width = Caustics_Long_Side

# Caustics picture
global N_Pixel_Side = 512
global N_Pixel_Height = N_Pixel_Side
global N_Pixel_Width = N_Pixel_Side

global Meters_Per_Pixel = Caustics_Height / N_Pixel_Height

# calculation
const ω = 1.99
# const ω = 2 / (1 + π / sqrt(N_Pixel_Height * N_Pixel_Width))
const N_Iterations_Convergence = 10_000        # CHECK: What is a reasonable number of iterations?

# # Global allocations to avoid excessive re-allocations. USEFUL ???
# global container = zeros(Float64, 1_024, 1_024)
# global divergence_intensity = zeros(Float64, 1_024, 1_024)
# global target_map = zeros(Float64, 1_024, 1_024)


using Manopt, Manifolds, Test
@testset "Stochastic Gradient Descent" begin
    M = Sphere(2)
    # 5 point mean
    p = [0.0, 0.0, 1.0]
    s = 1.0
    pts = [
        exp(M, p, X) for
        X in [zeros(3), [s, 0.0, 0.0], [-s, 0.0, 0.0], [0.0, s, 0.0], [0.0, -s, 0.0]]
    ]
    F(y) = 1 / 2 * sum([distance(M, y, x)^2 for x in pts])
    ∇F(y) = sum([-log(M, y, x) for x in pts])
    s∇F1(y) = [-log(M, y, x) for x in pts]
    s∇F2 = [(y -> -log(M, y, x)) for x in pts]
    @testset "Constructors" begin
        p1 = StochasticGradientProblem(M, s∇F1)
        p2 = StochasticGradientProblem(M, s∇F2)
        @test get_gradient(p1, 1, p) == zeros(3)
        @test get_gradient(p2, 1, p) == zeros(3)
        @test get_gradients(p1, p) == get_gradients(p2, p)
        o = StochasticGradientDescentOptions(deepcopy(p))
        o.order = collect(1:5)
        o.order_type = :Linear
        initialize_solver!(p1, o)
        o.order_type = :Linear
        step_solver!(p1, o, 1)
        @test o.x == exp(M, p, get_gradient(p1, 1, p))
    end
    @testset "Momentum and Average Processor Constructors" begin
        p1 = StochasticGradientProblem(M, s∇F1)
        p2 = GradientProblem(M, F, ∇F)
        m1 = MomentumGradient(M, p, StochasticGradient())
        m2 = MomentumGradient(p1, p)
        @test m1.direction == m2.direction #both use StochasticGradient
        m3 = MomentumGradient(M, p)
        m4 = MomentumGradient(p2, p)
        @test m3.direction == m4.direction #both use Gradient
        a1 = AverageGradient(M, p, 10, StochasticGradient())
        a2 = AverageGradient(p1, p, 10)
        @test a1.direction == a2.direction #both use StochasticGradient
        a3 = AverageGradient(M, p, 10)
        a4 = AverageGradient(p2, p, 10)
        @test a3.direction == a4.direction #both use Gradient
    end
    @testset "Comparing Stochastic Methods" begin
        x1 = stochastic_gradient_descent(M, s∇F1, p; order_type=:Linear)
        @test norm(x1) ≈ 1
        o1 = stochastic_gradient_descent(
            M, s∇F1, p; order_type=:Linear, return_options=true
        )
        x1a = get_solver_result(o1)
        @test x1 == x1a
        x2 = stochastic_gradient_descent(M, s∇F1, p; order_type=:FixedRandom)
        @test norm(x2) ≈ 1
        x3 = stochastic_gradient_descent(M, s∇F1, p; order_type=:Random)
        @test norm(x3) ≈ 1
        x4 = stochastic_gradient_descent(
            M,
            s∇F1,
            p;
            order_type=:Random,
            direction=AverageGradient(M, p, 10, StochasticGradient()),
        )
        @test norm(x4) ≈ 1
        x5 = stochastic_gradient_descent(
            M,
            s∇F1,
            p;
            order_type=:Random,
            direction=MomentumGradient(M, p, StochasticGradient()),
        )
        @test norm(x5) ≈ 1
    end
end


export ExplicitSolverStep, GeneralizedAlphaSolverStep, ImplicitStaticSolver, LinearStaticSolver,
LinearGeneralizedAlphaSolverStep, LinearGeneralizedAlphaSolverInit!
@doc raw"""
    ExplicitSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64)

Central Difference explicit solver for `M a + fint(u) = fext(u)`. `a`, `v`, `u` are acceleration, velocity and displacement.

```math
\begin{aligned}
u_{n+1} =& u_n + dtv_n + dt^2/2 a_n \\
v_{n+1} =& v_n + dt/2(a_n + a_{n+1}) \\
M a_{n+1} + f^{int}(u_{n+1}) =& f^{ext}_{n+1} \\
M a_{n+1} =& f^{ext}_{n+1} - f^{int}(u_{n+1}) \\
\end{aligned}
```

!!! info
    You need to call SolverInitial! before the first time step, if $f^{ext}_0 \neq 0$.
    Otherwise we assume the initial acceleration `globdat.acce[:] = 0`.
"""
function ExplicitSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64)
    assembleMassMatrix!(globdat, domain)

    u = globdat.state[:]
    ∂u  = globdat.velo[:]
    ∂∂u = globdat.acce[:]

    globdat.time  += Δt
    updateTimeDependentEssentialBoundaryCondition!(domain, globdat)
    fext = getExternalForce(domain, globdat)

    u += Δt*∂u + 0.5*Δt*Δt*∂∂u
    ∂u += 0.5*Δt * ∂∂u

    domain.state[domain.eq_to_dof] = u[:]
    fint  = assembleInternalForce( globdat, domain, Δt)
    ∂∂up = globdat.M\(fext - fint)

    ∂u += 0.5 * Δt * ∂∂up

    globdat.Dstate = globdat.state[:]
    globdat.state = u[:]
    globdat.velo = ∂u[:]
    globdat.acce = ∂∂up[:]

    commitHistory(domain)
    updateStates!(domain, globdat)

    return globdat, domain
end

@doc raw"""
    ImplicitStaticSolver(globdat::GlobalData, domain::Domain;
        N::Int64 = 10, ε::Float64 = 1.e-6, maxiterstep::Int64=100)

Solves
$$K(u) = F$$
using the incremental load method. Specifically, at step $i$, we solve
```math
f^{int}(u_i) = \frac{i}{N} f^{ext}
```

- `globdat`, GlobalData

- `domain`, Domain

- `N`, an integer, load stepping steps

- `ε`, Float64, absolute error for Newton convergence

-  `maxiterstep`, Int64, maximum iteration number for Newton convergence
"""
function ImplicitStaticSolver(globdat::GlobalData, domain::Domain;
        N::Int64 = 10, ε::Float64 = 1.e-6, maxiterstep::Int64=100)
    assembleMassMatrix!(globdat, domain)
    fext = getExternalForce(domain, globdat)

    for iterstep = 1:N

        # Newton's method
        Newtoniterstep, Newtonconverge = 0, false

        while  !Newtonconverge

            Newtoniterstep += 1

            fint, stiff = assembleStiffAndForce( globdat, domain )

            res = fint - iterstep/N * fext

            Δstate = stiff\res

            globdat.state -= Δstate
            # @show Newtoniterstep, norm(res)
            if (norm(res) < ε  || Newtoniterstep > maxiterstep)
                if Newtoniterstep > maxiterstep
                    @error "$Newtoniterstep Newton iteration does not converge"
                end
                Newtonconverge = true
            end
            updateStates!(domain, globdat)
        end
        commitHistory(domain)
        globdat.Dstate = copy(globdat.state)
    end
    globdat, domain
end

@doc raw"""
    LinearStaticSolver(globdat::GlobalData, domain::Domain)

Solves the linear static problem
$$\begin{aligned}\div \sigma &= f\\ \sigma = H \epsilon \end{aligned}$$

"""
function LinearStaticSolver(globaldata::GlobalData, domain::Domain)
    @assert globaldata.time ≈ 0.0
    fext = getExternalForce(domain, globaldata)
    fint, stiff = assembleStiffAndForce( globaldata, domain, 0.0)
    res = fext - fint
    globaldata.state = stiff\res
    updateStates!(domain, globaldata)
    return globaldata, domain
end

@doc raw"""
    GeneralizedAlphaSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64,
    ρ::Float64 = 0.0, ε::Float64 = 1e-8, ε0::Float64 = 1e-8, maxiterstep::Int64=100,
    η::Float64 = 1.0, failsafe::Bool = false, verbose::Bool = false)

Implicit solver for
$$Ma  + f_{int}(u) = fext$$
Here ``a``, ``v``, ``u`` are acceleration, velocity and displacement respectively.

- `ρ`: controls the damping effect of the α-scheme, ρ∈[0,1], ρ=1 corresponds to the maximum damping
- `ε`: Float64, absolute error for Newton convergence
- `ε0`: Float64, relative error for Newton convergence
- `max_iter`: Int64, maximum iteration number for Newton convergence
- `η`: Float64, Newton step size at the first iteration
- `failsafe`: Bool, if failsafe is true, when the Newton fails to converge,
              revert back, and return false

The nonlinear $\alpha$
```math
u_{n+1} = u_n + dtv_n + dt^2/2 ((1 - 2\beta)a_n + 2\beta a_{n+1})
v_{n+1} = v_n + dt((1 - \gamma)a_n + \gamma a_{n+1})
2\beta = 0.5*(1 - αm + αf)^2
\gamma = 0.5 - \alpha_m + \alpha_f
```

```math
a_{n+1-\alpha_m} = (1-\alpha_m)a_{n+1} + \alpha_m a_{n}
v_{n+1-\alpha_f} = (1-\alpha_f)v_{n+1} + \alpha_f v_{n}
u_{n+1-\alpha_f} = (1-\alpha_f)u_{n+1} + \alpha_f u_{n}
M a_{n+1-\alpha_m} + f^{int}(u_{n+1-\alpha_f}) = f^{ext}_{n+1-\alpha_f}
```

'a_{n+1}' is solved by

```math
M ((1-\alpha_m)a_{n+1} + \alpha_m a_{n})
+ f^{int}((1-\alpha_f)(u_n + dtv_n + dt^2/2 ((1 - 2\beta)a_n + 2\beta a_{n+1}))) + \alpha_f u_{n})
= f^{ext}_{n+1-\alpha_f}
```

As for ```\alpha_m``` and ``\alpha_f``
```math
\alpha_m = (2\rho_{\infty} - 1)/(\rho_{\infty} + 1)
\alpha_f = \rho_{\infty}/(\rho_{\infty} + 1)
```

use the current states `a`, `v`, `u`, `time` in globdat, and update these stetes to next time step
update domain history, when failsafe is true, and Newton's solver fails, nothing will be changed.

You need to call SolverInitial! before the first time step, if f^{ext}_0 != 0.
SolverInitial! updates a_0 in the globdat.acce
a_0 = M^{-1}(- f^{int}(u_0) + f^{ext}_0)

We assume globdat.acce[:] = a_0 and so far initialized to 0
We also assume the external force is conservative (it does not depend on the current deformation)
"""
function GeneralizedAlphaSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64;
        ρ::Float64 = 0.0, ε::Float64 = 1e-8, ε0::Float64 = 1e-8, maxiterstep::Int64=100,
        η::Float64 = 1.0, failsafe::Bool = false, verbose::Bool = false)
    assembleMassMatrix!(globdat, domain)
    @assert 0<=ρ<=1
    αm = (2ρ-1)/(1+ρ)
    αf = ρ/(1+ρ)
    local res0

    β2 = 0.5*(1 - αm + αf)^2
    γ = 0.5 - αm + αf

    # compute solution at uⁿ⁺¹
    failsafe_time = copy(globdat.time)
    globdat.time  += (1 - αf)*Δt

    # domain.Dstate = uⁿ
    failsafe_Dstate = copy(domain.Dstate)
    failsafe_state = copy(domain.state)
    domain.Dstate = domain.state[:]

    updateTimeDependentEssentialBoundaryCondition!(domain, globdat)
    M = globdat.M

    ∂∂u = globdat.acce[:] #∂∂uⁿ
    u = globdat.state[:]  #uⁿ
    ∂u  = globdat.velo[:] #∂uⁿ

    fext = getExternalForce(domain, globdat)
    # fext = zeros(domain.neqs)
    ∂∂up = ∂∂u[:]

    Newtoniterstep, Newtonconverge = 0, false

    norm0 = Inf

    while !Newtonconverge

        Newtoniterstep += 1

        domain.state[domain.eq_to_dof] = (1 - αf)*(u + Δt*∂u + 0.5 * Δt * Δt * ((1 - β2)*∂∂u + β2*∂∂up)) + αf*u
        fint, stiff = assembleStiffAndForce( globdat, domain, Δt)
        res = M * (∂∂up *(1 - αm) + αm*∂∂u)  + fint - fext
        if Newtoniterstep==1
            res0 = res
        end
        A = M*(1 - αm) + (1 - αf) * 0.5 * β2 * Δt^2 * stiff
        Δ∂∂u = A\res

        while η * norm(Δ∂∂u) > norm0
            η /= 2.0
            verbose && (@info "η", η)
        end
        ∂∂up -= η*Δ∂∂u

        if maxiterstep==1
            break
        end
        verbose && println("$Newtoniterstep/$maxiterstep, abs = $(norm(res)), rel = $(norm(res)/norm(res0))")
        if (norm(res)< ε || norm(res)< ε0*norm(res0) ||Newtoniterstep > maxiterstep)
            if Newtoniterstep > maxiterstep
                # Newton method does not converge
                if failsafe
                    globdat.time = failsafe_time
                    domain.state = failsafe_state[:]
                    domain.Dstate = failsafe_Dstate[:]
                    return false
                end
                # When failsafe is not on, test the gradient
                function f(∂∂up)
                    domain.state[domain.eq_to_dof] = (1 - αf)*(u + Δt*∂u + 0.5 * Δt * Δt * ((1 - β2)*∂∂u + β2*∂∂up)) + αf*u
                    fint, stiff = assembleStiffAndForce( globdat, domain )
                    fint, (1 - αf) * 0.5 * β2 * Δt^2 * stiff
                end
                gradtest(f, ∂∂up)
                @warn("Newton iteration cannot converge $(norm(res))"); Newtonconverge = true
            else
                Newtonconverge = true
                verbose && printstyled("[Newmark] Newton converged $Newtoniterstep\n", color=:green)
            end
        end

        η = min(1.0, 2η)
        norm0 = norm(Δ∂∂u)
    end

    #update globdat to the next time step
    globdat.Dstate = globdat.state[:]
    globdat.state += Δt * ∂u + Δt^2/2 * ((1 - β2) * ∂∂u + β2 * ∂∂up)
    globdat.velo += Δt * ((1 - γ) * ∂∂u + γ * ∂∂up)
    globdat.acce = ∂∂up[:]
    globdat.time  += αf*Δt

    #commit history in domain
    commitHistory(domain)
    updateStates!(domain, globdat)
    if options.save_history>=2
        fint, stiff = assembleStiffAndForce( globdat, domain, Δt)
        push!(domain.history["fint"], fint)
        push!(domain.history["fext"], fext)
    end
    if options.save_history>=1
        push!(domain.history["time"], [globdat.time])
    end
    return globdat, domain

end

function LinearGeneralizedAlphaSolverInit!(globdat::GlobalData, domain::Domain, Δt::Float64;
    ρ::Float64 = 0.0)
    assembleMassMatrix!(globdat, domain)
    @assert 0<=ρ<=1
    αm = (2ρ-1)/(1+ρ)
    αf = ρ/(1+ρ)
    β2 = 0.5*(1 - αm + αf)^2
    γ = 0.5 - αm + αf
    M = globdat.M

    _, stiff = assembleStiffAndForce(domain, Δt)
    A = M*(1 - αm) + (1 - αf) * 0.5 * β2 * Δt^2 * stiff
    STORAGE["LinearGeneralizedAlphaSolverInit!"] = lu(A)
    nothing
end

"""
    LinearGeneralizedAlphaSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64;
    ρ::Float64 = 0.0)

Solves the dynamic structural equation with time-independent coefficient matrix. This is useful for linear elasticity simulation.
"""
function LinearGeneralizedAlphaSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64;
    ρ::Float64 = 0.0)
    if haskey(STORAGE, "LinearGeneralizedAlphaSolverInit")
        error("You must call LinearGeneralizedAlphaSolverInit! before LinearGeneralizedAlphaSolverStep")
    end
    A = STORAGE["LinearGeneralizedAlphaSolverInit!"]
    M = globdat.M
    @assert 0<=ρ<=1
    αm = (2ρ-1)/(1+ρ)
    αf = ρ/(1+ρ)
    β2 = 0.5*(1 - αm + αf)^2
    γ = 0.5 - αm + αf

    # compute solution at uⁿ⁺¹
    globdat.time  += (1 - αf)*Δt

    # domain.Dstate = uⁿ
    domain.Dstate = domain.state[:]

    updateTimeDependentEssentialBoundaryCondition!(domain, globdat)

    ∂∂u = globdat.acce[:] #∂∂uⁿ
    u = globdat.state[:]  #uⁿ
    ∂u  = globdat.velo[:] #∂uⁿ

    fext = getExternalForce(domain, globdat)
    ∂∂up = ∂∂u[:]

    domain.state[domain.eq_to_dof] = (1 - αf)*(u + Δt*∂u + 0.5 * Δt * Δt * ((1 - β2)*∂∂u + β2*∂∂up)) + αf*u

    fint = assembleInternalForce(domain, Δt)
    res = M * (∂∂up *(1 - αm) + αm*∂∂u)  + fint - fext
    Δ∂∂u = A\res
    ∂∂up -= Δ∂∂u

    #update globdat to the next time step
    globdat.Dstate = globdat.state[:]
    globdat.state += Δt * ∂u + Δt^2/2 * ((1 - β2) * ∂∂u + β2 * ∂∂up)
    globdat.velo += Δt * ((1 - γ) * ∂∂u + γ * ∂∂up)
    globdat.acce = ∂∂up[:]
    globdat.time  += αf*Δt

    #commit history in domain
    commitHistory(domain)
    updateStates!(domain, globdat)
    if options.save_history>=2
        fint, stiff = assembleStiffAndForce( globdat, domain, Δt)
        push!(domain.history["fint"], fint)
        push!(domain.history["fext"], fext)
    end
    if options.save_history>=1
        push!(domain.history["time"], [globdat.time])
    end
    return globdat, domain

end


#=
polylagrange:
- Julia version:
- Author: ymocquar
- Date: 2019-11-25
=#
include("polyexp.jl")
function getpolylagrange(k::Int64, j::Int64, N::DataType)
    @assert k <= j "_getpolylagrange(k=$k,j=$j) k must be less or equal to j"
    @assert N <: Signed "the type $N must be an Integer"
    result = Polynomial([one(Complex{Rational{N}})])
    for l = 0:j
        if l != k
            result *= Polynomial([l // 1, 1 // 1]) / (l - k)
        end
    end
    return result
end
function interpolate(tab, order, value, N::DataType)
    T = (N == BigInt) ? BigFloat : Float64
    res = zeros(Complex{T}, size(tab[1]))
    for i = 0:order
        res .+= getpolylagrange(i, order, N)(-value) * tab[i+1]
    end
    return res
end
function getpolylagrange(t::Vector{T}, k, j) where {T<:Number}
    result = Polynomial([one(T)])
    for l = 0:j
        if l != k
            result *= Polynomial([-t[l+1], one(T)]) / (t[k+1] - t[l+1])
        end
    end
    return result
end
function interpolate(tab_time::Vector{T}, tab, order, value) where {T<:Number}
    res = zeros(Complex{T}, size(tab[1]))
    for i = 0:order
        res .+= getpolylagrange(tab_time, i, order)(value) * tab[i+1]
    end
    return res
end
interpolate(tab, order, value) = interpolate(tab, order, value, BigInt)
struct CoefExpABRational
    tab_coef::Any
    function CoefExpABRational(
        order::Int64,
        epsilon::AbstractFloat,
        list_tau,
        dt::AbstractFloat,
    )
        n_tau = size(list_tau, 1)
        T = typeof(epsilon)
        tab_coef = zeros(Complex{T}, n_tau, order + 1, order + 1)
        N = T == BigFloat ? BigInt : Int64
        epsilon = rationalize(N, epsilon, tol = epsilon * 10 * Base.eps(T))
        dt = rationalize(N, dt, tol = dt * 10 * Base.eps(T))
        list_tau = rationalize.(N, list_tau)
        pol_x = Polynomial([0 // 1, 1 // dt])
        for j = 0:order
            for k = 0:j
                res = view(tab_coef, :, k + 1, j + 1)
                pol = getpolylagrange(k, j, N)
                pol2 = pol(pol_x)
                for ind = 1:n_tau
                    ell = list_tau[ind]
                    pol3 = undef
                    pol_int = if ell == 0
                        # in this case the exponentiel value is always 1
                        Polynomials.integrate(pol2)
                    else
                        pol3 = PolyExp(pol2, im * ell / epsilon, -im * ell * dt / epsilon)
                        Polynomials.integrate(pol3)
                    end
                    res[ind] = pol_int(dt) - pol_int(0)
                end
            end
        end
        return new(tab_coef)
    end
end
struct CoefExpAB
    tab_coef::Any
    tab_coef_neg::Any
    function CoefExpAB(order::Int64, epsilon::AbstractFloat, n_tau, dt)
        T = typeof(epsilon)
        N = T == BigFloat ? BigInt : Int64
        new_prec = precision(BigFloat)
        new_prec += T == BigFloat ? order * 16 : 0
        setprecision(BigFloat, new_prec) do
            list_tau = [collect(0:n_tau/2-1); collect(-n_tau/2:-1)]
            T2 = BigFloat
            epsilon = T2(epsilon)
            dt = T2(dt)
            tab_coef = zeros(Complex{T2}, n_tau, order + 1, order + 1)
            pol_x = Polynomial([0, 1 / dt])
            for j = 0:order
                for k = 0:j
                    res = view(tab_coef, :, k + 1, j + 1)
                    pol = getpolylagrange(k, j, N)
                    pol2 = pol(pol_x)
                    for ind = 1:n_tau
                        ell = list_tau[ind]
                        pol_int = if ell == 0
                            # in this case the exponentiel value is always 1
                            Polynomials.integrate(pol2)
                        else
                            Polynomials.integrate(
                                PolyExp(pol2, im * ell / epsilon, -im * ell * dt / epsilon),
                            )
                        end
                        res[ind] = pol_int(dt) - pol_int(0)
                    end
                end
            end # end of for j=....
        end # end of setprecision(...)
        # conversion to the new precision
        tab_coef = T.(real(tab_coef)) + im * T.(imag(tab_coef))
        tab_coef_neg = -conj(tab_coef)
        return new(tab_coef, tab_coef_neg)
    end
end


nt = 1000
seed = 99
m = ACQR_simula(nt,seed)
y = m["y"]

############################### starting point
wi = collect(range(0.1, length=2, stop=3))
Ai,Ci,Qi,Ri,Si,m1i,P1i = ACQR_em_starting(y,wi,1,true)

#######################
max_iter=10
tol=1e-10
txo=false
print("Testing ACQR_em_s: ")
A,C,Q,R,m1,P1,loglik,aic = ACQR_em_s(y,Ai,Ci,Qi,Ri,m1i,max_iter,tol,txo)
if sum(A.^2) > 1e-6
    println("OK")
end


# Simulation of Linear Model Data

# Starting Simrel Code
function simrel{T}(
    n::T         = 100,
    p::T         = 10,
    q::T         = 5,
    relpos::T    = [1, 2, 3],
    gamma::T     = 0.2,
    R2::T        = 0.9;
    ntest::T     = nothing,
    muY::T       = nothing,
    muX::T       = nothing,
    lambdaMin::T = 10e-4,
    sim::T       = nothing
)
    m = length(relpos) ## Number of relevant components
    if (q < m)
        error("Number of relevant predictors must at least equal to $(m).")
    end
    irrelPosition = setdiff(collect(1:p), relpos)
    extraDim = q - m
    predPosition = vcat(relpos, sample(irrelPosition, extraDim,
                                    replace = false, ordered = true))
    nu = lambdaMin * exp(-gamma)/(1 - lambdaMin)
    if (lambdaMin < 0 || lambdaMin >= 1)
        error("Parameter lambdaMin must be in between 0 and 1")
    end
    lambdas = (exp(-gamma * vec(1:p)) + nu) / (exp(-gamma) + nu)
    sigmaZ = diagm(lambdas)
    sigmaZinv = inv(sigmaZ)
    sigmaZy = zeros(p, 1)
    uniformSample = rand(Uniform(-1, 1), m)
    covarianceSample = sign(uniformSample) .* sqrt(R2 * abs(uniformSample) /
             sum(abs(uniformSample)) .* lambdas[relpos])
    sigmaZy[relpos] = covarianceSample
    sigmaY = 1
    sigma = vcat(hcat(sigmaY, sigmaZy'), hcat(sigmaZy, sigmaZ))

    # Rotation Matrix

    Q = randn(q, q) |> x -> x .- mean(x, 1) # Centering columns
    Rq = qr(Q)[:1]
    R = eye(p)
    R[predPosition, predPosition] = Rq
    if (q < (p - 1))
      Q = randn(p - q, p - q) |> x -> x .- mean(x, 1)
      Rnq = qr(Q)[:1]
      notPredPos = setdiff(1:p, predPosition)
      R[notPredPos, notPredPos] = Rnq
    end

    # Regression coefficients
    betaZ = sigmaZinv * sigmaZy
    betaX = R * betaZ
    beta0 = 0

    if !(muY == nothing)
      beta0 = beta0 + muY
    end
    if !(muX == nothing)
      beta0 = beta0 - transpose(betaX) * muX
    end

    R2 = transpose(sigmaZy) * betaZ
    minerror = sigmaY - R2

    if all(eig(sigma)[:1] .> 0)
      sigmaRot = chol(sigma)
      ucal = randn(n, p + 1)
      ucal_1 = ucal * sigmaRot
      Y = ucal_1[:, 1]
      if !(muY == nothing)
        Y = Y .+ muY
      end
      Z = ucal_1[:, 2:end]
      X = Z * transpose(R)
      if !(muX == nothing)
        X = X .+ muX
      end

      # Test Observations
      if !(ntest == nothing)
        utest = randn(ntest, p + 1)
        utest_1 = utest * sigmaRot
        testY = utest_1[:, 1]
        if !(muY == nothing)
          testY = testY .+ muY
          testZ = utest_1[:, 2:end]
          testX = testZ * transpose(R)
        end
        if !(muX == nothing)
          testX = testX .+ muX
        end
      else
        testX = nothing
        testY = nothing
      end
    else
      return error("Correlation matrix is not positive definite")
    end

    return @NT(
      beta0    = beta0,
      beta     = betaX,
      sigma    = sigma,
      R2       = R2,
      minerror = minerror,
      X        = X,
      Y        = Y,
      testX    = testX,
      testY    = testY
    )
end


function test_user_knn()
    println("-- Testing UserKNN rating prediction")

    m = [NaN 3 NaN 1 2 1 NaN 4
         1 2 NaN NaN 3 2 NaN 3
         NaN 2 3 3 NaN 5 NaN 1]
    data = DataAccessor(m)

    recommender = UserKNN(data, 1)
    build!(recommender)

    @test isapprox(recommender.sim[1, 2], 0.447, atol=1e-3)
    @test isapprox(recommender.sim[2, 3], -0.693, atol=1e-3)

    rec = recommend(recommender, 1, 3, [i for i in 1:8])
    @test rec[1] == (5 => 3.0)
    @test rec[2] == (8 => 3.0)
    @test rec[3] == (2 => 2.0)

    recommender = UserKNN(data, 1, true)
    build!(recommender)

    @test isapprox(recommender.sim[1, 2], 0.447, atol=1e-3)
    @test isapprox(recommender.sim[2, 3], -0.693, atol=1e-3)

    rec = recommend(recommender, 1, 3, [i for i in 1:8])
    @test rec[1] == (5 => 3.0)
    @test rec[2] == (8 => 3.0)
    @test rec[3] == (7 => 2.2)
end

test_user_knn()


@testset "Persistent Image" begin

    intrs = diagram(0.0=>Inf, 0.0=>1.0, 1.0=>1.0, 1.0=>2.0)

    pimg = PersistentImage(intrs...)
    @test length(pimg.intervals) == length(intrs)
    @test pimg(0,0) ≈ 0.0627 atol=1e-3
    @test length(vec(pimg, 5, 6)) == 30

end


"""
    cholUnblocked!(A, Val{:L})

Overwrite the lower triangle of `A` with its lower Cholesky factor.

The name is borrowed from [https://github.com/andreasnoack/LinearAlgebra.jl]
because these are part of the inner calculations in a blocked Cholesky factorization.
"""
function cholUnblocked! end

function cholUnblocked!(D::Diagonal{T}, ::Type{Val{:L}}) where {T<:AbstractFloat}
    map!(sqrt, D.diag, D.diag)
    D
end

function cholUnblocked!(A::StridedMatrix{T}, ::Type{Val{:L}}) where {T<:BlasFloat}
    n = LinearAlgebra.checksquare(A)
    if n == 1
        A[1] < zero(T) && throw(PosDefException(1))
        A[1] = sqrt(A[1])
    elseif n == 2
        A[1] < zero(T) && throw(PosDefException(1))
        A[1] = sqrt(A[1])
        A[2] /= A[1]
        (A[4] -= abs2(A[2])) < zero(T) && throw(PosDefException(2))
        A[4] = sqrt(A[4])
    else
        _, info = LAPACK.potrf!('L', A)
        iszero(info) || throw(PosDefException(info))
    end
    A
end

function cholUnblocked!(D::UniformBlockDiagonal, ::Type{Val{:L}})
    for f in D.facevec
        cholUnblocked!(f, Val{:L})
    end
    D
end


using BenchmarkTools
using Random
using LinearAlgebra

BenchmarkTools.DEFAULT_PARAMETERS.evals = 20

# define arrays and matrices
const rows, cols = 500, 600
const reduceRows, reduceCols = Int(rows / 5), Int(cols / 6)

const float_arrayA = rand(rows * cols)
const float_arrayB = rand(rows * cols)
const small_int_matrixA = rand(1:10, reduceRows, reduceCols)
const small_int_matrixB = rand(1:10, reduceRows, reduceCols)
const float_matrixA = rand(rows, cols)
const float_matrixC = rand(cols, rows)
const small_float_matrixA = rand(reduceRows, reduceCols)
const small_float_matrixB = rand(reduceRows, reduceCols)

test1_name = "Element-wise sum of two $(reduceRows)x$(reduceCols) matrices (int), (1000 loops)"
function test1()
    for _ in 1:1000
        _ = small_int_matrixA + small_int_matrixB
    end
end

test2_name = "Element-wise multiplication of two $(reduceRows)x$(reduceCols) matrices (float64), (1000 loops)"
function test2()
    for _ in 1:1000
        _ = small_float_matrixA .* small_float_matrixB
    end
end

test3_name = "Dot (scalar) product of two $(rows * cols) arrays (float64), (1000 loops)"
function test3()
    for _ in 1:1000
        _ = dot(float_arrayA, float_arrayB)
    end
end

test4_name = "Matrix product of $(rows)x$(cols) and $(cols)x$(rows) matrices (float64)"
function test4()
    _ = float_matrixA * float_matrixC
end

test5_name = "L2 norm of $(rows)x$(cols) matrix (float64), (1000 loops)"
function test5()
    for _ in 1:1000
        _ = norm(float_matrixA)
    end
end

test6_name = "Sort of $(rows)x$(cols) matrix (float64)"
function test6()
    _ = sort(float_matrixA, dims = 2)
end

println(test1_name)
@btime test1()

println(test2_name)
@btime test2()

println(test3_name)
@btime test3()

println(test4_name)
@btime test4()

println(test5_name)
@btime test5()

println(test6_name)
@btime test6()


module HarmonicOrthogonalPolynomials
using FastTransforms, LinearAlgebra, ClassicalOrthogonalPolynomials, ContinuumArrays, DomainSets,
        BlockArrays, BlockBandedMatrices, InfiniteArrays, StaticArrays, QuasiArrays, Base, SpecialFunctions
import Base: OneTo, oneto, axes, getindex, convert, to_indices, _maybetail, tail, eltype, *, ==, ^, copy, -, abs, resize!
import BlockArrays: block, blockindex, unblock, BlockSlice
import DomainSets: indomain, Sphere
import LinearAlgebra: norm, factorize
import QuasiArrays: to_quasi_index, SubQuasiArray, *
import ContinuumArrays: TransformFactorization, @simplify, ProjectionFactorization
import ClassicalOrthogonalPolynomials: checkpoints, _sum, cardinality, increasingtruncations
import BlockBandedMatrices: BlockRange1, _BandedBlockBandedMatrix
import FastTransforms: Plan, interlace
import QuasiArrays: LazyQuasiMatrix, LazyQuasiArrayStyle
import InfiniteArrays: InfStepRange, RangeCumsum

export SphericalHarmonic, UnitSphere, SphericalCoordinate, RadialCoordinate, Block, associatedlegendre, RealSphericalHarmonic, sphericalharmonicy, Laplacian, AbsLaplacianPower, abs, -, ^, AngularMomentum
cardinality(::Sphere) = ℵ₁

include("multivariateops.jl")
include("spheretrav.jl")
include("coordinates.jl")

increasingtruncations(::BlockedUnitRange{<:RangeCumsum{Int,<:AbstractRange}}) = broadcast(n -> Block.(oneto(n)), 4:2:∞)

checkpoints(::UnitSphere{T}) where T = [SphericalCoordinate{T}(0.1,0.2), SphericalCoordinate{T}(0.3,0.4)]

abstract type AbstractSphericalHarmonic{T} <: MultivariateOrthogonalPolynomial{3,T} end
struct RealSphericalHarmonic{T} <: AbstractSphericalHarmonic{T} end
struct SphericalHarmonic{T} <: AbstractSphericalHarmonic{T} end
SphericalHarmonic() = SphericalHarmonic{ComplexF64}()
RealSphericalHarmonic() = RealSphericalHarmonic{Float64}()
copy(a::AbstractSphericalHarmonic) = a

axes(S::AbstractSphericalHarmonic{T}) where T = (Inclusion{SphericalCoordinate{real(T)}}(UnitSphere{real(T)}()), blockedrange(1:2:∞))

associatedlegendre(m) = ((-1)^m*prod(1:2:(2m-1)))*(UltrasphericalWeight((m+1)/2).*Ultraspherical(m+1/2))
lgamma(n) = logabsgamma(n)[1]

function sphericalharmonicy(ℓ, m, θ, φ)
    m̃ = abs(m)
    exp((lgamma(ℓ+m̃+1)+lgamma(ℓ-m̃+1)-2lgamma(ℓ+1))/2)*sqrt((2ℓ+1)/(4π)) * exp(im*m*φ) * sin(θ/2)^m̃ * cos(θ/2)^m̃ * jacobip(ℓ-m̃,m̃,m̃,cos(θ))
end

function getindex(S::SphericalHarmonic{T}, x::SphericalCoordinate, K::BlockIndex{1}) where T
    ℓ = Int(block(K))
    k = blockindex(K)
    m = k-ℓ
    convert(T, sphericalharmonicy(ℓ-1, m, x.θ, x.φ))::T
end

==(::SphericalHarmonic{T},::SphericalHarmonic{T}) where T = true
==(::RealSphericalHarmonic{T},::RealSphericalHarmonic{T}) where T = true

# function getindex(S::RealSphericalHarmonic{T}, x::ZSphericalCoordinate, K::BlockIndex{1}) where T
#     # sorts entries by ...-2,-1,0,1,2... scheme
#     ℓ = Int(block(K))
#     k = blockindex(K)
#     m = k-ℓ
#     m̃ = abs(m)
#     indepm = (-1)^m̃*exp((lgamma(ℓ-m̃)-lgamma(ℓ+m̃))/2)*sqrt((2ℓ-1)/(2π))*associatedlegendre(m̃)[x.z,ℓ-m̃]
#     m>0 && return cos(m*x.φ)*indepm
#     m==0 && return cos(m*x.φ)/sqrt(2)*indepm
#     m<0 && return sin(m̃*x.φ)*indepm
# end

function getindex(S::RealSphericalHarmonic{T}, x::SphericalCoordinate, K::BlockIndex{1}) where T
    # starts with m=0, then alternates between sin and cos terms (beginning with sin).
    ℓ = Int(block(K))
    m = blockindex(K)-1
    z = cos(x.θ)
    if iszero(m)
        return sqrt((2ℓ-1)/(4*π))*associatedlegendre(0)[z,ℓ]
    elseif isodd(m)
        m = (m+1)÷2
        return sin(m*x.φ)*(-1)^m*exp((lgamma(ℓ-m)-lgamma(ℓ+m))/2)*sqrt((2ℓ-1)/(2*π))*associatedlegendre(m)[z,ℓ-m]
    else
        m = m÷2
        return cos(m*x.φ)*(-1)^m*exp((lgamma(ℓ-m)-lgamma(ℓ+m))/2)*sqrt((2ℓ-1)/(2*π))*associatedlegendre(m)[z,ℓ-m]
    end
end

getindex(S::AbstractSphericalHarmonic, x::StaticVector{3}, K::BlockIndex{1}) = S[SphericalCoordinate(x), K]
getindex(S::AbstractSphericalHarmonic, x::StaticVector{3}, K::Block{1}) = S[x, axes(S,2)[K]]
getindex(S::AbstractSphericalHarmonic, x::StaticVector{3}, KR::BlockOneTo) = mortar([S[x, K] for K in KR])
getindex(S::AbstractSphericalHarmonic, x::StaticVector{3}, k::Int) = S[x, findblockindex(axes(S,2), k)]
getindex(S::AbstractSphericalHarmonic, x::StaticVector{3}, kr::AbstractUnitRange{Int}) = [S[x, k] for k in kr]

# @simplify *(Ac::QuasiAdjoint{<:Any,<:SphericalHarmonic}, B::SphericalHarmonic) =

##
# Expansion
##

const FiniteSphericalHarmonic{T} = SubQuasiArray{T,2,SphericalHarmonic{T},<:Tuple{<:Inclusion,<:BlockSlice{BlockRange1{OneTo{Int}}}}}
const FiniteRealSphericalHarmonic{T} = SubQuasiArray{T,2,RealSphericalHarmonic{T},<:Tuple{<:Inclusion,<:BlockSlice{BlockRange1{OneTo{Int}}}}}
copy(a::FiniteRealSphericalHarmonic) = a
copy(a::FiniteSphericalHarmonic) = a

function grid(S::FiniteSphericalHarmonic)
    T = real(eltype(S))
    N = blocksize(S,2)
    # The colatitudinal grid (mod $\pi$):
    θ = ((1:N) .- one(T)/2)/N
    # The longitudinal grid (mod $\pi$):
    M = 2*N-1
    φ = (0:M-1)*2/convert(T, M)
    SphericalCoordinate.(π*θ, π*φ')
end
function grid(S::FiniteRealSphericalHarmonic)
    T = real(eltype(S))
    N = blocksize(S,2)
    # The colatitudinal grid (mod $\pi$):
    θ = ((1:N) .- one(T)/2)/N
    # The longitudinal grid (mod $\pi$):
    M = 2*N-1
    φ = (0:M-1)*2/convert(T, M)
    SphericalCoordinate.(π*θ, π*φ')
end

struct SphericalHarmonicTransform{T} <: Plan{T}
    sph2fourier::FastTransforms.FTPlan{T,2,FastTransforms.SPINSPHERE}
    analysis::FastTransforms.FTPlan{T,2,FastTransforms.SPINSPHEREANALYSIS}
end
struct RealSphericalHarmonicTransform{T} <: Plan{T}
    sph2fourier::FastTransforms.FTPlan{T,2,FastTransforms.SPHERE}
    analysis::FastTransforms.FTPlan{T,2,FastTransforms.SPHEREANALYSIS}
end

SphericalHarmonicTransform{T}(N::Int) where T<:Complex = SphericalHarmonicTransform{T}(plan_spinsph2fourier(T, N, 0), plan_spinsph_analysis(T, N, 2N-1, 0))
RealSphericalHarmonicTransform{T}(N::Int) where T<:Real = RealSphericalHarmonicTransform{T}(plan_sph2fourier(T, N), plan_sph_analysis(T, N, 2N-1))

*(P::SphericalHarmonicTransform{T}, f::Matrix{T}) where T = SphereTrav(P.sph2fourier \ (P.analysis * f))
*(P::RealSphericalHarmonicTransform{T}, f::Matrix{T}) where T = RealSphereTrav(P.sph2fourier \ (P.analysis * f))

factorize(S::FiniteSphericalHarmonic{T}) where T =
    TransformFactorization(grid(S), SphericalHarmonicTransform{T}(blocksize(S,2)))
factorize(S::FiniteRealSphericalHarmonic{T}) where T =
    TransformFactorization(grid(S), RealSphericalHarmonicTransform{T}(blocksize(S,2)))

function _sum(A::AbstractSphericalHarmonic{T}, dims) where T
    @assert dims == 1
    PseudoBlockArray(Hcat(sqrt(4convert(T, π)), Zeros{T}(1,∞)), (Base.OneTo(1),axes(A,2)))
end

include("laplace.jl")
include("angularmomentum.jl")

end # module