using Plots, Distributions
using IPMeasures
using IPMeasures: crit_mmd2_var, crit_mxy_over_mltpl, crit_mxy_over_sum
plotlyjs()

function sample_blobs(n, ratio, rows=5, cols=5, sep=10)
    correlation = (ratio - 1) / (ratio + 1)
    # generate within-blob variation
    X = rand(Distributions.Normal(0,1), 2, n)

    corr_sigma = [1  correlation;  correlation 1]
    Y = rand(Distributions.MultivariateNormal([0,0], corr_sigma), n)

    # assign to blobs
    X[1, :] .+= rand(1:rows,n) .* sep
    X[2, :] .+= rand(1:cols,n) * sep
    Y[1, :] .+= rand(1:rows,n) * sep
    Y[2, :] .+= rand(1:cols,n) * sep

    (X, Y)
end

@testset "Criteria comparison" begin
    γs = -3:0.1:2
    x = randn(2,1000);
    y = rand(2,1000) .- [0.5,0.5];
    y = randn(2,1000) .* [0.01,0.01];
    x, y = sample_blobs(2000, 6.0, 5, 5, 10)
    plot(γs, [crit_mmd2_var(IPMeasures.IMQKernel(10.0^γ), x, y, 1000) for γ in γs], label = "MMD2 / √VAR")
    plot!(γs, [crit_mxy_over_mltpl(IPMeasures.IMQKernel(10.0^γ), x, y) for γ in γs], label = "M(X,Y) / SQRT(M(X,X) * M(Y,Y)")
    p = plot!(γs, [crit_mxy_over_sum(IPMeasures.IMQKernel(10.0^γ), x, y) for γ in γs], label = "M(X,Y) / (M(X,X) + M(Y,Y)")
    @test p != nothing
end


function _adapt_ϵ(logϵ, Hbar, logϵbar, da_stat, m, M_adapt, δ, μ;
    γ=0.05, t0=10, κ=0.75)

if m <= M_adapt

Hbar = (1.0 - 1.0 / (m + t0)) * Hbar + (1 / (m + t0)) * (δ - da_stat)
logϵ = μ - sqrt(m) / γ * Hbar
logϵbar = m^(-κ) * logϵ + (1 - m^(-κ)) * logϵbar

else

logϵ = logϵbar

end

return logϵ, Hbar, logϵbar

end


module RegularizedGLM

using Reexport
@reexport using PenaltyFunctions

# linreg: μ = identity,                 μ′ = x -> 1
# logreg: μ = x -> 1 / (1 + exp(-x)),   μ′ = x -> exp(-x) / (1 + exp(-x))^2
# poisreg: μ = exp,                     μ′ = exp

function glmfit(x, y, μ::Function, μ′::Function)
    n, p = size(x)
    β = zeros(p)
    xβ = zeros(n)
    r = zeros(n)
    xtr = zeros(p)
    x2 = vec(sum(abs2, x, 1))
    for iter in 1:100
        for i in 1:n
            r[i] = y[i] - μ(xβ[i])
        end
        At_mul_B!(xtr, x, r)
        for j in 1:p
            β[j] = β[j] + p * xtr[j] / x2[j]
        end
        A_mul_B!(xβ, x, β)
    end
    β
end

function logreg(x, y)
    n, p = size(x)
    β = zeros(p)
    buffer_n = zeros(n)
    buffer_p = zeros(p)
    buffer_p2 = zeros(p)
    xtx_inv = 4inv(Symmetric(x'x))
    for iter in 1:50
        for i in 1:n
            buffer_n[i] = y[i] - 1 / (1 + exp(-buffer_n[i]))
        end
        At_mul_B!(buffer_p, x, buffer_n)
        A_mul_B!(buffer_p2, xtx_inv, buffer_p)
        β[:] = β + buffer_p2
        A_mul_B!(buffer_n, x, β)
    end
    β
end

# (exp(x) / (1 + exp(x))) / (1 + exp(x))

end # module


using Plots
using ComplexPortraits: generate_spike

xx = collect(LinRange(0.0, 2.0, 801))
sp1 = generate_spike(a=-9.0, b=0.5)
sp2 = generate_spike(a=-4.0, b=0.5)
sp3 = generate_spike(a=-15.0, b=1.5)

plot(xx, sp1.(xx), label="a= -9, b=0.5", linewidth=3, legend=:bottom, framestyle=:box, size=(800,400),
  tickfontsize=11, legendfontsize=11)
plot!(xx, sp2.(xx), label="a= -4, b=0.5", linewidth=3)
pl = plot!(xx, sp3.(xx), label="a=-15, b=1.5", linewidth=3)
savefig(pl, "./images/spike.png")


# --------------------------------------------------------------------------
# ACE1.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

"""
`mutable struct ToyModel` : a tight-binding type toy model admitting
multiple species, for testing the completeness of basis sets.
"""
mutable struct ToyModel{NZ} <: AbstractCalculator
   β::Float64
   eF::Float64
   zlist::SZList{NZ}
   α::Matrix{Float64}
end

function ToyModel(species; α = :rand, β = 1.0)
   zlist = ZList(species, static=true)
   if α == :rand
      α = 1.0 .+ 3 * rand(length(zlist), length(zlist))
      α = 0.5 * (α + α')
   end
   V = ToyModel(β, 0.0, ZList(species, static=true), α)
   V.eF = find_eF(V)
   return V
end

function find_eF(V)
   H = hamiltonian(rand_config(V), V)
   emin, emax = extrema(eigvals(H))
   for n = 1:30
      emin, emax = extrema([eigvals(H); [emin, emax]])
   end
   return 0.5 * (emin+emax)
end

fd(e, β) = 1/(1 + exp(β * e))
dfd(e, β) = (- β * exp(β * e)) / (1 + exp(β * e))^2

hop(r, α) = exp(-α * (r/3-1))
dhop(r, α) = - α/3 * exp(-α * (r/3-1))

function hamiltonian(at, V)
   X = positions(at)
   Z = at.Z
   H = zeros(length(at), length(at))
   for i = 1:length(at)-1, j = (i+1):length(at)
      α = V.α[ z2i(V.zlist, Z[i]), z2i(V.zlist, Z[j]) ]
      H[j, i] = H[i, j] = hop(norm(X[i]-X[j]), α)
   end
   return H
end

function hamiltonian_d(at::Atoms{T}, V) where {T}
   X = positions(at)
   Z = at.Z
   dH = zeros(JVec{T}, length(at), length(at))
   for i = 1:length(at), j = (i+1):length(at)
      α = V.α[ z2i(V.zlist, Z[i]), z2i(V.zlist, Z[j]) ]
      Rij = X[i]-X[j]
      rij = norm(Rij)
      dH[i, j] = dhop(rij, α) * (Rij / rij)
      dH[j, i] = - dH[i, j]
   end
   return dH
end

function energy(V::ToyModel, at::Atoms)
   H = hamiltonian(at, V)
   e = eigvals(H)
   # @show extrema(e)
   return sum(fd.(e .- V.eF, V.β))
end

function forces(V::ToyModel, at::Atoms)
   H = hamiltonian(at, V)
   e, Ψ = eigen(H)
   dH = hamiltonian_d(at, V)
   F = 2 * [ sum( dfd(e[n] - V.eF, V.β) * sum(Ψ[:, n] .* dH[:, j]) * Ψ[j, n]
                for n = 1:length(e) )
           for j = 1:length(at) ]
   # df_H = sum( dfd(e[n] - V.eF, V.β) * Ψ[:,n] * Ψ[:,n]'  for n = 1:length(e) )
   return F #  - sum(df_H .* dH; dims = 1)[:]
end

Base.rand(V::ToyModel; kwargs...) = rand_config(V; kwargs...)


# following Asymptotic Gaussian law for
# noninteracting indistinguishable
# particles in random networks
# Valery S. Shchesnovich

"""
    partition_expectation_values(partition_size_vector, partition_counts)
    partition_expectation_values(part_occ::PartitionOccupancy)

Return the Haar averaged probability of photon number count in binned outputs
for [`Distinguishable`](@ref) and [`Bosonic`](@ref) particles.

!!! note "Reference"
    [https://www.nature.com/articles/s41598-017-00044-8.pdf](https://www.nature.com/articles/s41598-017-00044-8.pdf)
"""
function partition_expectation_values(partition_size_vector, partition_counts)

    m = sum(partition_size_vector)
    number_partitions = length(partition_size_vector)
    partition_size_ratio = partition_size_vector ./ m # q
    n = sum(partition_counts)
    @test all(partition_counts .>= 0)

    # factorials can quickly get pretty big so this is required
    if any(partition_counts .> 20) || n > 20
        n = big(n)
        partition_counts = big.(partition_counts)
    end

    proba_dist = factorial(n) / vector_factorial(partition_counts) * prod(partition_size_ratio .^ partition_counts)

    proba_bosonic = proba_dist * prod([partition_counts[i] > 1 ? prod([1+l/partition_size_vector[i] for l in 1:partition_counts[i]-1]) : 1 for i in 1:length(partition_size_vector)]) /prod([1+l/m for l in 1:n-1])

    proba_dist, proba_bosonic

end

partition_expectation_values(part_occ::PartitionOccupancy) = partition_expectation_values(partition_occupancy_to_partition_size_vector_and_counts(part_occ)...)

"""
    subset_expectation_value(subset_size, k, n, m)

Return the Haar averaged probability to find `k` from `n` photons inside a subset of
binned output modes of length `size_subset` among `m` modes for [`Distinguishable`](@ref)
and [`Bosonic`](@ref) cases.
"""
function subset_expectation_value(subset_size, k,n,m)

    partition_size_vector = [subset_size, m-subset_size]
    partition_counts = [k,n-k]

    partition_expectation_values(partition_size_vector, partition_counts)

end

"""
    subset_relative_distance_of_averages(subset_size, n, m)
"""
function subset_relative_distance_of_averages(subset_size,n,m)

    @warn "check tvd conventions"
    0.5*abs(sum([abs(subset_expectation_value(subset_size, k,n,m)[1] - subset_expectation_value(subset_size, k,n,m)[2]) for k in 0:n]))

end

"""
    choose_best_average_subset(;m, n, distance=tvd)

Return the ideal subset size on average and its total variance distance.
"""
function choose_best_average_subset(;m,n, distance = tvd)

    """returns the ideal subset size on average and its TVD,
    to compare with the full bunching test of shsnovitch"""

    function distance_this_subset_size(subset_size)

        proba_dist = [subset_expectation_value(subset_size,k,n,m)[1] for k in 0:n]
        proba_bos = [subset_expectation_value(subset_size,k,n,m)[2] for k in 0:n]

        distance(proba_dist, proba_bos)

    end

    max_distance = 0
    subset_size_max_ratio = nothing

    for subset_size in 1:m-1

        if distance_this_subset_size(subset_size) > max_distance
            max_distance = distance_this_subset_size(subset_size)
            subset_size_max_ratio = subset_size
        end
    end

    subset_size_max_ratio, distance_this_subset_size(subset_size_max_ratio)

end

"""
    best_partition_size(;m, n, n_subsets, distance=tvd)

Return the ideal `partition_size_vector` for a given number of subsets `n_subsets`
and the Haar averaged TVD in second parameter.

!!! note
    For a single subset, `n_subsets`=2 as we need a complete partition, occupying all modes.
"""
function best_partition_size(;m,n, n_subsets, distance = tvd)

    """return the ideal partition_size_vector for a given number
    of subsets n_subsets as well as the haar averaged tvd in second parameter

    for a single subset, n_subsets = 2 as we need a complete partition, occupying all modes"""

    @argcheck n_subsets >= 2 "we need a complete partition, occupying all modes"
    @argcheck n_subsets <= m "more partition bins than output modes"

    part_list = all_mode_configurations(m,n_subsets, only_photon_number_conserving = true)

    remove_trivial_partitions!(part_list)

    max_distance = 0
    part_max_ratio = nothing

    for part in ranked_partition_list(part_list)

        events = all_mode_configurations(n,n_subsets, only_photon_number_conserving = true)

        pdf = [partition_expectation_values(part, event) for event in events]

        pdf_dist = hcat(collect.(pdf)...)[1,:]
        pdf_bos = hcat(collect.(pdf)...)[2,:]

        this_distance = distance(pdf_bos,pdf_dist)

        println(part, " : ", this_distance)

        if this_distance > max_distance
            max_distance = this_distance
            part_max_ratio = part
        end
    end

    part_max_ratio, max_distance

end


"""
OED.jl - Optimal Experimental Design for Inverse Problems with State Constraints

Codes used to generate the results presented in

@article{ruthotto2017optimal,
  title={Optimal Experimental Design for Inverse Problems with State Constraints},
  author={Ruthotto, Lars and Chung, Julianne and Chung, Matthias},
  journal={SIAM Journal on Scientific Computing},
  year={2018}
}

Highlights:
- We solve the OED problem involving constraints on the states
                         (bound and linear equality constraints in inversion)
- This code parallelizes ove the training data using jInv's parallelization.
"""
module OED

	using jInv.ForwardShare
	using jInv.Utils
	using jInv.Mesh
	using jInv.InverseSolve
	using LinearOperators
	using MAT

	# descriptions of OED problems
	include("AbstractOEDParam.jl")
	include("oedTypeA.jl")
	include("oedTypeB.jl")
	include("unconstrainedOED.jl")
	include("eqconstrainedOED.jl")
	include("ineqconstrainedOED.jl")
	include("sensMat.jl")

	# optimization routines
	include("ipqp.jl")
	include("regFun.jl")
	include("models.jl")

	# codes to set up OED problem B
	include("rotation2D.jl")
	include("linearInter.jl")
	include("getLinearInterMatrix.jl")
	include("getLinearInterMatrixTranspose.jl")

	# main drivers
	include("runOEDTypeA.jl")
	include("runOEDTypeB.jl")
end


using Test, Random, AdvancedHMC
using Statistics: mean
include("common.jl")

@testset "`gdemo`" begin
    rng = MersenneTwister(0)

    n_samples = 5_000
    n_adapts = 1_000

    θ_init = randn(rng, 2)

    metric = DiagEuclideanMetric(2)
    h = Hamiltonian(metric, ℓπ_gdemo, ∂ℓπ∂θ_gdemo)
    init_eps = Leapfrog(0.1)
    prop = NUTS(init_eps)
    adaptor = StanHMCAdaptor(n_adapts, Preconditioner(metric), NesterovDualAveraging(0.8, prop.integrator))

    samples, _ = sample(rng, h, prop, θ_init, n_samples, adaptor, n_adapts)

    m_est = mean(map(_s -> [exp(_s[1]), _s[2]], samples[1000:end]))

    @test m_est ≈ [49 / 24, 7 / 6] atol=RNDATOL
end


using OmegaCore, Distributions

# A model where x --causes-> y
function m1()
  x = 1 ~ Normal(0, 1)
  y(ω) = (2 ~ Normal(x(ω), 1))(ω)
  (x = x, y = y)
end

# A model where y --causes-> x
function m2()
  x = 3 ~ Normal(0, 1)
  y(ω) = (4 ~ Normal(x(ω), 1))(ω)
  (x = y, y = x)
end

is_model_1 = 3 ~ Bernoulli(0.5)
model = ifelseₚ(is_model_1, m1(), m2())
mx(ω) = model(ω).x
mxω(ω) = mx(ω)(ω)
my(ω) = model(ω).y

x_ = 1.3
y_ = 2.0

# This won't do what we want!
condition1(ω) = (mxω |ᵈ (my => 100.0))(ω) >= 101.0

condition(ω) = (mxω |ᵈ (my =>ˡ y_))(ω) >= x_

# condition = (mxω |ᵈ (my =>ˡ y_)) >=ₚ x_

prob = mean(randsample(is_model_1 |ᶜ condition, 1000))
println("""The model1 is correct (x causes y), given after forcing y to 2 x,
we observed that x was 1.3 is""", prob)


@doc raw"""
    LogisticSoftMaxLikelihood(num_class::Int) -> MultiClassLikelihood

## Arguments
- `num_class::Int` : Total number of classes

    LogisticSoftMaxLikelihood(labels::AbstractVector) -> MultiClassLikelihood

## Arguments
- `labels::AbstractVector` : List of classes labels

---

The multiclass likelihood with a logistic-softmax mapping: :
```math
p(y=i|\{f_k\}_{1}^{K}) = \frac{\sigma(f_i)}{\sum_{k=1}^k \sigma(f_k)}
```
where ``\sigma`` is the logistic function.
This likelihood has the same properties as [softmax](https://en.wikipedia.org/wiki/Softmax_function).
---

For the analytical version, the likelihood is augmented multiple times.
More details can be found in the paper [Multi-Class Gaussian Process Classification Made Conjugate: Efficient Inference via Data Augmentation](https://arxiv.org/abs/1905.09670).
"""
LogisticSoftMaxLikelihood(x) = MultiClassLikelihood(LogisticSoftMaxLink(), x)

struct LogisticSoftMaxLink <: AbstractLink end

function (::LogisticSoftMaxLink)(f::AbstractVector{<:Real})
    return normalize(logistic.(f), 1)
end

function implemented(
    ::MultiClassLikelihood{<:LogisticSoftMaxLink},
    ::Union{<:AnalyticVI,<:MCIntegrationVI,<:GibbsSampling},
)
    return true
end

Base.show(io::IO, ::LogisticSoftMaxLink) = print(io, "Logistic-SoftMax Link")

## Local Updates ##
function init_local_vars(
    l::MultiClassLikelihood{<:LogisticSoftMaxLink}, batchsize::Int, T::DataType=Float64
)
    num_class = n_class(l)
    c = [ones(T, batchsize) for _ in 1:num_class] # Second moment of fₖ
    α = num_class * ones(T, batchsize) # First variational parameter of Gamma distribution
    β = num_class * ones(T, batchsize) # Second variational parameter of Gamma distribution
    θ = [rand(T, batchsize) * 2 for _ in 1:num_class] # Variational parameter of Polya-Gamma distribution
    γ = [rand(T, batchsize) for _ in 1:num_class] # Variational parameter of Poisson distribution
    return (; c, α, β, θ, γ)
end

function local_updates!(
    local_vars,
    ::MultiClassLikelihood{<:LogisticSoftMaxLink},
    y,
    μ::NTuple{N,<:AbstractVector},
    Σ::NTuple{N,<:AbstractVector},
) where {N}
    @. local_vars.c = broadcast((Σ, μ) -> sqrt.(Σ + abs2.(μ)), Σ, μ)
    for _ in 1:2
        broadcast!(
            (β, c, μ, ψα) -> 0.5 / β * exp.(ψα) .* safe_expcosh.(-0.5 * μ, 0.5 * c),
            local_vars.γ,
            Ref(local_vars.β),
            local_vars.c,
            μ,
            Ref(digamma.(local_vars.α)),
        ) # Update γ
        local_vars.α .= 1 .+ (local_vars.γ...)
    end
    broadcast!(
        (y, γ, c) -> 0.5 * (y + γ) ./ c .* tanh.(0.5 .* c),
        local_vars.θ, # target
        eachcol(y), # argument 1
        local_vars.γ, # argument 2
        local_vars.c, # argument 3
    ) # update θ
    return local_vars
end

function sample_local!(local_vars, ::MultiClassLikelihood{<:LogisticSoftMaxLink}, y, f)
    broadcast!(
        f -> rand.(Poisson.(0.5 * local_vars.α .* safe_expcosh.(-0.5 * f, 0.5 * f))),
        local_vars.γ,
        f,
    )
    local_vars.α .= rand.(Gamma.(1 .+ (local_vars.γ...), inv.(local_vars.β)))
    local_vars.θ .= broadcast(
        (y, γ, f) -> rand.(PolyaGamma.(y .+ Int.(γ), abs.(f))), eachcol(y), local_vars.γ, f
    )
    return local_vars
end

## Global Gradient Section ##

@inline function ∇E_μ(::MultiClassLikelihood{<:LogisticSoftMaxLink}, ::AOptimizer, y, state)
    return 0.5 .* (eachcol(y) .- state.γ)
end
@inline function ∇E_Σ(::MultiClassLikelihood{<:LogisticSoftMaxLink}, ::AOptimizer, y, state)
    return 0.5 .* state.θ
end

## ELBO Section ##
function expec_loglikelihood(
    ::MultiClassLikelihood{<:LogisticSoftMaxLink}, ::AnalyticVI, y, μ, Σ, state
)
    tot = -length(y) * logtwo
    tot += -sum(sum(state.γ .+ eachcol(y))) * logtwo
    tot += 0.5 * sum(zip(state.θ, state.γ, eachcol(y), μ, Σ)) do (θ, γ, y, μ, Σ)
        dot(μ, (y - γ)) - dot(θ, abs2.(μ)) - dot(θ, Σ)
    end
    return tot
end

function AugmentedKL(l::MultiClassLikelihood{<:LogisticSoftMaxLink}, state, y)
    return PolyaGammaKL(l, state, y) + PoissonKL(l, state) + GammaEntropy(l, state)
end

function PolyaGammaKL(::MultiClassLikelihood, state, y)
    return sum(broadcast(PolyaGammaKL, eachcol(y) .+ state.γ, state.c, state.θ))
end

function PoissonKL(::MultiClassLikelihood, state)
    return sum(
        broadcast(
            PoissonKL,
            state.γ,
            Ref(state.α ./ state.β),
            Ref(digamma.(state.α) .- log.(state.β)),
        ),
    )
end

##  Compute the equivalent of KL divergence between an improper prior p(λ) (``1_{[0,\\infty]}``) and a variational Gamma distribution ##
function GammaEntropy(::MultiClassLikelihood, state)
    return -sum(state.α) + sum(log, first(state.β)) - sum(first ∘ logabsgamma, state.α) -
           dot(1 .- state.α, digamma.(state.α))
end

## Numerical Gradient Section ##

function grad_samples!(
    model::AbstractGPModel{T,<:MultiClassLikelihood{<:LogisticSoftMaxLink},<:NumericalVI},
    samples::AbstractMatrix{T},
    opt_state,
    y,
    index,
) where {T}
    l = likelihood(model)
    grad_μ = zeros(T, n_latent(model))
    grad_Σ = zeros(T, n_latent(model))
    g_μ = similar(grad_μ)
    num_sample = size(samples, 1)
    @views @inbounds for i in 1:num_sample
        σ = logistic.(samples[i, :])
        samples[i, :] .= l(samples[i, :])
        s = samples[i, y][1]
        g_μ .= grad_logisticsoftmax(samples[i, :], σ, y) / s
        grad_μ += g_μ
        grad_Σ += diaghessian_logisticsoftmax(samples[i, :], σ, y) / s - abs2.(g_μ)
    end
    for k in 1:n_latent(model)
        opt_state[k].ν[index] = -grad_μ[k] / num_sample
        opt_state[k].λ[index] = grad_Σ[k] / num_sample
    end
end

function log_like_samples(
    ::AbstractGPModel{T,<:MultiClassLikelihood{<:LogisticSoftMaxLink}},
    samples::AbstractMatrix,
    y,
) where {T}
    num_sample = size(samples, 1)
    loglike = zero(T)
    for i in 1:num_sample
        σ = logistic.(samples[i, :])
        loglike += log(σ[y][1]) - log(sum(σ))
    end
    return loglike / num_sample
end

function grad_logisticsoftmax(s::AbstractVector{T}, σ::AbstractVector{T}, y) where {T<:Real}
    return s[y][1] * (y .- s) .* (one(T) .- σ)
end

function diaghessian_logisticsoftmax(
    s::AbstractVector{T}, σ::AbstractVector{T}, y
) where {T<:Real}
    return s[y][1] * (one(T) .- σ) .*
           (abs2.(y - s) .* (1 .- σ) - s .* (1 .- s) .* (1 .- σ) - σ .* (y - s))
end

function hessian_logisticsoftmax(
    s::AbstractVector{T}, σ::AbstractVector{T}, y
) where {T<:Real}
    m = length(s)
    i = findfirst(y)
    hessian = zeros(T, m, m)
    @inbounds for j in 1:m
        for k in 1:m
            hessian[j, k] =
                (one(T) - σ[j]) *
                s[i] *
                (
                    (δ(T, i, k) - s[k]) * (one(T) - σ[k]) * (δ(T, i, j) - s[j]) -
                    s[j] * (δ(T, j, k) - s[k]) * (one(T) - σ[k]) -
                    δ(T, k, j) * σ[j] * (δ(T, i, j) - s[j])
                )
        end
    end
    return hessian
end


using MathOptInterfaceSCS
using Base.Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MOI.Test
const MOIU = MOI.Utilities

using MathOptInterfaceBridges
const MOIB = MathOptInterfaceBridges

MOIU.@model SCSModelData () (EqualTo, GreaterThan, LessThan) (Zeros, Nonnegatives, Nonpositives, SecondOrderCone, ExponentialCone, PositiveSemidefiniteConeTriangle) () (SingleVariable,) (ScalarAffineFunction,) (VectorOfVariables,) (VectorAffineFunction,)

MOIB.@bridge SplitInterval MOIB.SplitIntervalBridge () (Interval,) () () () (ScalarAffineFunction,) () ()
MOIB.@bridge LogDet MOIB.LogDetBridge () () (LogDetConeTriangle,) () () () (VectorOfVariables,) (VectorAffineFunction,)
MOIB.@bridge RootDet MOIB.RootDetBridge () () (RootDetConeTriangle,) () () () (VectorOfVariables,) (VectorAffineFunction,)

# linear9test needs 1e-3 with SCS < 2.0 and 5e-1 with SCS 2.0
# linear2test needs 1e-4
const config = MOIT.TestConfig(atol=1e-4, rtol=1e-4)

@testset "Continuous linear problems" begin
    # AlmostSuccess for linear9 with SCS 2
    MOIT.contlineartest(SplitInterval{Float64}(MOIU.CachingOptimizer(SCSModelData{Float64}(), SCSOptimizer())), config, ["linear9"])
end

@testset "Continuous conic problems" begin
    MOIT.contconictest(RootDet{Float64}(LogDet{Float64}(MOIU.CachingOptimizer(SCSModelData{Float64}(), SCSOptimizer()))), config, ["rsoc", "geomean", "psds", "rootdet", "logdets"])
end


####bandwidth selector for kernel density

#rule of thumb
function bwnormal(xdata::RealVector)
    0.9 * min((quantile(xdata, .75) - quantile(xdata, .25)) / 1.34, std(xdata)) * length(xdata) ^ (-0.2)
end
function midrange(x::RealVector)
    lq, uq = quantile(x, [.25, .75])
    uq - lq
end

# J(h)=∑ᵢⱼK'((xᵢ - xⱼ)/h) /(n²h) + 2K(0)/(nh) =
# where K'(u) = invsqrt2π(exp(-0.25u*u)/sqrt(2) - 2exp(-0.5u*u))
#J(h) = invsqrt2π/(n²h) ∑ᵢⱼ (exp(-0.25u*u)/sqrt(2) - 2exp(-0.5u*u)) + 2 * invsqrt2π/nh
#J(h) = 2*invsqrt2π/(n²h) ∑{i<j} (exp(-0.25u*u)/sqrt(2) - 2exp(-0.5u*u)) + invsqrt2π/sqrt(2)nh
#For normal kernel
# function Jh(xdata::RealVector, h::Real, n::Int)
#     tmp = 0.0
#     @inbounds for i in 1:(n-1)
#         for j in (i+1):n
#             u = (xdata[i] - xdata[j])/h
#             # u = exp(-0.25*u*u)
#             u = -0.25*u*u
#             u = exp(u)
#             tmp += u*invsqrt2 - 2*u*u
#         end
#     end
#     2*tmp / (n*n*h) + 1/(sqrt2*n*h)
# end
#For gaussiankernel, equivalent to the above one. The above version requires
#less computing but is inefficient in practice due to the exp function.
function Jh(xdata::RealVector, h::Real, w::Vector, n::Int)
    ll = 0.0
    @inbounds for ind in 1:n
        gaussiankernel(xdata[ind], xdata, sqrt2*h, w, n)
        ll += mean(w)
        # gaussiankernel(xdata[ind], xdata, h, w, n)
        # w[ind] = 0.0
        # ll -= 2*mean(w)
    end
    ll / n - leaveoneout(xdata, gaussiankernel, h, w, n)
end

#for general kernel
function Jh(xdata::RealVector, kernel::Function, h::Real, w::Vector, n::Int, xlb::Real, xub::Real)
    hquadrature(x->begin kernel(x, xdata,h,w,n); mean(w)^2; end, xlb, xub, maxevals=200)[1] - leaveoneout(xdata, kernel, h, w, n)
end
function leaveoneout(xdata::RealVector, kernel::Function, h::Real, w::Vector, n::Int)

    ll = 0.0
    @inbounds for ind in 1:n
        kernel(xdata[ind], xdata, h, w, n)
        w[ind] = 0.0
        ll += mean(w)
    end
    ll * 2 / (n-1)
end
#For betakernel
function Jh(xdata::RealVector, logxdata::RealVector,log1_xdata::RealVector, kernel::Function, h::Real, w::Vector, n::Int, xlb::Real, xub::Real)
    hquadrature(x->begin kernel(x, logxdata, log1_xdata, h,w,n); mean(w)^2; end, xlb, xub, maxevals=200)[1] - leaveoneout(xdata, logxdata, log1_xdata, kernel, h, w, n)
end
function leaveoneout(xdata::RealVector, logxdata::RealVector, log1_xdata::RealVector, kernel::Function, h::Real, w::Vector, n::Int)

    ll = 0.0
    @inbounds for ind in 1:n
        kernel(xdata[ind], logxdata, log1_xdata, h, w, n)
        w[ind] = 0.0
        ll += mean(w)
    end
    ll * 2 / (n-1)
end
#For gammakernel
function Jh(xdata::RealVector, logxdata::RealVector, kernel::Function, h::Real, w::Vector, n::Int, xlb::Real, xub::Real)
    hquadrature(x->begin kernel(x, xdata, logxdata, h,w,n); mean(w)^2; end, xlb, xub, maxevals=200)[1] - leaveoneout(xdata, logxdata, kernel, h, w, n)
end
function leaveoneout(xdata::RealVector, logxdata::RealVector, kernel::Function, h::Real, w::Vector, n::Int)

    ll = 0.0
    @inbounds for ind in 1:n
        kernel(xdata[ind], xdata, logxdata, h, w, n)
        w[ind] = 0.0
        ll += mean(w)
    end
    ll * 2 / (n-1)
end
#Least Squares cross validation for Gaussian Kernel
#May fail to work if there are multiple equial x_i
# Silverman suggest search interval be (0.25, 1.5)n^(-0.2)σ
function bwlscv(xdata::RealVector, kernel::Function)
    n=length(xdata)
    w  = zeros(n)
    h0=bwnormal(xdata)
    #when n is large, leaveoneout is more expensive than numeric integration
    if kernel == gaussiankernel && n<200
        return Optim.minimizer(Optim.optimize(h -> Jh(xdata, h, w, n), 0.01*h0, 10*h0, iterations=200, abs_tol=h0/n))
    end

    xlb, xub = extrema(xdata)
    h0=midrange(xdata)
    hlb = h0/n
    hub = h0
    if kernel == betakernel
        xlb = 0.0
        xub = 1.0
        hub = 0.25
        logxdata = log.(xdata)
        log1_xdata = log.(1.0 .- xdata)
        return Optim.minimizer(Optim.optimize(h -> Jh(xdata, logxdata, log1_xdata, kernel, h, w, n, xlb,xub), hlb, hub, iterations=200,abs_tol=h0/n^2))
    elseif kernel == gammakernel
        xlb = 0.0
        logxdata = log.(xdata)
        return Optim.minimizer(Optim.optimize(h -> Jh(xdata, logxdata, kernel, h, w, n, xlb,xub), hlb, hub, iterations=200,abs_tol=h0/n^2))
    end
    return Optim.minimizer(Optim.optimize(h -> Jh(xdata, kernel, h, w, n, xlb,xub), hlb, hub, iterations=200,abs_tol=h0/n^2))
end

# likelihood cross validation for beta and gamma kernel
#there seems no other easy way; least square cross validation can be formidable because their convolution have no close form
#may also work for other kernels, but likelihood cv has some known disadvantages.
#Not recommended, but easy to implement for arbitrary kernel
function lcv(xdata::RealVector, kernel::Function, h::Real, w::Vector, n::Int)
#     -mean(kerneldensity(xdata,xdata,kernel,h)) + mean(map(kernel, xdata, xdata, h))
    ind = 1
    ind_end = 1+n
    ll = 0.0
    @inbounds while ind < ind_end
        kernel(xdata[ind], xdata, h, w, n)
        w[ind] = 0.0
        ll += log(mean(w))
        ind += 1
    end
    -ll
end
function bwlcv(xdata::RealVector, kernel::Function)
    n = length(xdata)
    w = zeros(n)
    h0=midrange(xdata)
    hlb = h0/n^2
    hub = h0
    if kernel==betakernel
        hub = 0.25
    end
    return Optim.minimizer(Optim.optimize(h->lcv(xdata,kernel,h,w,n), hlb, hub, iterations=200,abs_tol=h0/n^2))
end

function lcv(xdata::RealMatrix, kernel::Vector, h::RealVector, w::Vector, n::Int)
#     -mean(kerneldensity(xdata,xdata,kernel,h)) + mean(map(kernel, xdata, xdata, h))
    if any(h .<= 0.0)
        return Inf
    end
    ind = 1
    ind_end = 1+n
    ll = 0.0
    p = size(xdata)[2]
    @inbounds while ind < ind_end
        fill!(w, 1.0)
        wtmp = ones(n)
        for j in 1:p
            kernel[j](xdata[ind, j], xdata[:, j], h[j], wtmp, n)
            for k in 1:n
                w[k] = w[k] * wtmp[k]
            end
        end
        divide!(w, sum(w))

        w[ind] = 0.0
        ll += log(mean(w))
        ind += 1
    end
    -ll
end
function bwlcv(xdata::RealMatrix, kernel::Vector)
    n, p = size(xdata)
    w = ones(n)
    h0 = zeros(p)
    hlb = zeros(p)
    hub = zeros(p)
    for j in 1:p
        if kernel[j] == gaussiankernel
            h0[j] = bwnormal(xdata[:, j])
            hlb[j] = 0.1*h0[j]
            hub[j] = 10*h0[j]
        elseif kernel[j] == betakernel
            h0[j] = midrange(xdata[:, j])
            hlb[j] = h0[j]/n
            hub[j] = 0.25
        elseif kernel[j] == gammakernel
            h0[j] = midrange(xdata[:, j])
            hlb[j] = h0[j]/n
            hub[j] = h0[j]
        end
    end
    h = Optim.minimizer(Optim.optimize(h->lcv(xdata, kernel, h, w, n), h0))
    if all(hlb .<= h .<= hub)
        return h
    else
        return h0
    end
end

#leave-one-out LSCV. 1/n \sum((yi - yihat)/(1-wi))
function lscvlocalconstant(xdata::RealVector, ydata::RealVector, kernel::Function, h::Real, w::Vector, n::Int)
    tmp = 0.0
    ind = 1
    ind_end = 1 + n
    @inbounds while ind < ind_end
#         map!(kernel, w, xdata, xdata[i], h)
        kernel(xdata[ind], xdata, h, w, n)
        divide!(w, sum(w))
        tmp += abs2((wsum(w, ydata) - ydata[ind])/(1-w[ind]))
        ind += 1
    end
    tmp/n
end

function bwlocalconstant(xdata::RealVector, ydata::RealVector, kernel::Function=gaussiankernel)
    n=length(xdata)
    length(ydata)==n || error("length(ydata) != length(xdata)")
    w = ones(n)
    if kernel == gaussiankernel
        h0= bwnormal(xdata)
        hlb = 0.1*h0
        hub = 10*h0
    elseif kernel == betakernel
        h0 = midrange(xdata)
        hlb = h0/n
        hub = 0.25
    elseif kernel == gammakernel
        h0 = midrange(xdata)
        hlb = h0/n
        hub = h0
    end
    Optim.minimizer(Optim.optimize(h->lscvlocalconstant(xdata, ydata, kernel, h, w, n), hlb, hub))
end

#see reference:Smoothing Parameter Selection in Nonparametric Regression Using an Improved Akaike Information Criterion
# Clifford M. Hurvich, Jeffrey S. Simonoff and Chih-Ling Tsai
# Journal of the Royal Statistical Society. Series B (Statistical Methodology), Vol. 60, No. 2 (1998), pp. 271-293
#http://www.jstor.org/stable/2985940
function AIClocallinear(xdata::RealVector, ydata::RealVector, kernel::Function, h::Real, w::Vector, n::Int)
    tmp = 0.0
    traceH = 0.0
    ind = 1
    ind_end = 1+n
    @inbounds while ind < ind_end
#           map!(kernel, w, xdata, xdata[i], h)
          kernel(xdata[ind], xdata, h, w, n)
          s0 = sum(w)
          s1 = s0*xdata[ind] - wsum(w, xdata)
          s2 = wsumsqdiff(w, xdata, xdata[ind], n)
          sy0 = wsum(w, ydata)
          sy1 = wsumyxdiff(w, xdata, xdata[ind], ydata, n)
          tmp+=abs2((s2 * sy0 - s1 * sy1) /(s2 * s0 - s1 * s1) - ydata[ind])
          traceH += s0*w[ind]/(s2 * s0 - s1 * s1)
          ind += 1
    end
    tmp/n  + 2*(traceH+1)/(n-traceH-2)
end

function bwlocallinear(xdata::RealVector, ydata::RealVector, kernel::Function=gaussiankernel)
    n=length(xdata)
    length(ydata)==n || error("length(ydata) != length(xdata)")
    w = ones(n)
    if kernel == gaussiankernel
        h0= bwnormal(xdata)
        hlb = 0.1*h0
        hub = 10*h0
    elseif kernel == betakernel
        h0 = midrange(xdata)
        hlb = h0/n
        hub = 0.25
    elseif kernel == gammakernel
        h0 = midrange(xdata)
        hlb = h0/n
        hub = h0
    end
    Optim.minimizer(Optim.optimize(h->AIClocallinear(xdata, ydata, kernel,h, w, n), hlb, hub))
end

function bwreg(xdata::RealVector, ydata::RealVector, reg::Function, kernel::Function=gaussiankernel)

    if (reg == locallinear) || (reg == lp1)
        return bwlocallinear(xdata, ydata, kernel)
    elseif (reg == localconstant) || (reg == lp0)
        return bwlocalconstant(xdata, ydata, kernel)
    else
        error("I don't know what is $reg")
    end
end

# #leave-one-out LSCV for multivariate NW
function lscvlocalconstant(xdata::RealMatrix, ydata::RealVector, kernel::Array{Function, 1}, h::RealVector, w::Vector, n::Int)
    if any(h .<= 0.0)
        return Inf
    end
    tmp = 0.0
    ind = 1
    ind_end = 1+n
    p=size(xdata)[2]
    @inbounds while ind < ind_end
        fill!(w, 1.0)
        wtmp = ones(n)
        for j in 1:p
            kernel[j](xdata[ind, j], xdata[:, j], h[j], wtmp, n)
            for k in 1:n
                w[k] = w[k] * wtmp[k]
            end
        end
        divide!(w, sum(w))
        tmp += abs2((wsum(w, ydata) - ydata[ind])/(1-w[ind]))
        ind += 1
    end
    tmp/n
end

function bwlocalconstant(xdata::RealMatrix, ydata::RealVector, kernel::Array{Function, 1} = Function[gaussiankernel for i in 1:size(xdata)[2]])
    n, p = size(xdata)
    w = ones(n)
    h0 = zeros(p)
    hlb = zeros(p)
    hub = zeros(p)
    for j in 1:p
        if kernel[j] == gaussiankernel
            h0[j] = bwnormal(xdata[:, j])
            hlb[j] = 0.1*h0[j]
            hub[j] = 10*h0[j]
        elseif kernel[j] == betakernel
            h0[j] = midrange(xdata[:, j])
            hlb[j] = h0[j]/n
            hub[j] = 0.25
        elseif kernel[j] == gammakernel
            h0[j] = midrange(xdata[:, j])
            hlb[j] = h0[j]/n
            hub[j] = h0[j]
        end
    end
    h_output=Optim.minimizer(Optim.optimize(h->lscvlocalconstant(xdata, ydata, kernel, h, w, n), h0))
    if any(h_output .<= 0.0)
        for j in 1:p
            if h_output[j] .<= 0.0
                h_output[j] = 2 .* h0[j]
            end
        end
        h_output = Optim.minimizer(Optim.optimize(h->lscvlocalconstant(xdata, ydata, kernel, h, w, n), h_output))
    end
    h_output
end


using luteos
using SymPy

# let # limit scope

Ps = [P1(), P2(), P3()] # Range of polynomial order
Ns = [9, 17, 33] # Range of grid size

dim = 2

# Material
mat = Material( ν=0.33, E=1.0, dim = dim )

## Set up problem
#   Get functions for exact solution
function ExactSol( Cstiff )

  dim = 2

  @syms x1 x2

  u1 = (x2 - x2.^2) .* cos.(pi*x1) .* (1/2 - x1) .* (1 - x2/2)   #10 * (x2 - x2.^2) .* sin.(pi*x1) .* (1 - x1) .* (1 - x2/2)
  u2 = 2 * (x1 - x1.^2) .* sin.(pi*x2) .* (1 - x2) .* (1 - x1/2) # 2 * (x1 - x1.^2) .* sin.(pi*x2) .* (1 - x2) .* (1 - x1/2)

  ∂u1∂x1 = diff( u1, x1 )
  ∂u1∂x2 = diff( u1, x2 )
  ∂u2∂x1 = diff( u2, x1 )
  ∂u2∂x2 = diff( u2, x2 )

  ∇u = [ ∂u1∂x1 ∂u2∂x1; ∂u1∂x2 ∂u2∂x2 ]

  ϵ = 1/2 * ( ∇u + ∇u.' )

  σ = fill( 0*x1, dim, dim )

  for ii in 1:dim, jj in 1:dim, kk in 1:dim, ll in 1:dim
    σ[ii,jj] += (Cstiff[ii,jj,kk,ll] * ϵ[kk,ll])[1,1,1,1]
  end

  F  = fill( 0*x1, dim )

  ∇σ = fill( 0*x1, dim, dim, dim )

  for ii in 1:dim, jj in 1:dim
    ∇σ[ii,jj,1] = diff( σ[ii,jj], x1 )
    ∇σ[ii,jj,2] = diff( σ[ii,jj], x2 )
  end

  for ii in 1:dim, jj in 1:dim
    F[ii] -= ∇σ[ii,jj,jj]
  end

  u1func_org = lambdify( u1 )
  u2func_org = lambdify( u2 )

  F1func_org = lambdify( F[1] )
  F2func_org = lambdify( F[2] )
  σ1func_org = lambdify( σ[1] )
  σ2func_org = lambdify( σ[2] )
  σ4func_org = lambdify( σ[4] )
  ϵ1func_org = lambdify( ϵ[1] )
  ϵ2func_org = lambdify( ϵ[2] )
  ϵ4func_org = lambdify( ϵ[4] )

  # make sure you can call it with matrices
  u1func = x -> u1func_org.(x[:,1], x[:,2])
  u2func = x -> u2func_org.(x[:,1], x[:,2])

  F1func = x -> F1func_org.(x[:,1], x[:,2])
  F2func = x -> F2func_org.(x[:,1], x[:,2])
  σ1func = x -> σ1func_org.(x[:,1], x[:,2])
  σ2func = x -> σ2func_org.(x[:,1], x[:,2])
  σ4func = x -> σ4func_org.(x[:,1], x[:,2])
  ϵ1func = x -> ϵ1func_org.(x[:,1], x[:,2])
  ϵ2func = x -> ϵ2func_org.(x[:,1], x[:,2])
  ϵ4func = x -> ϵ4func_org.(x[:,1], x[:,2])

  return( u1func, u2func, F1func, F2func, σ1func, σ2func, σ4func, ϵ1func, ϵ2func, ϵ4func )

end

( u1func, u2func, F1func, F2func, σ1func, σ2func, σ4func, ϵ1func, ϵ2func, ϵ4func ) = ExactSol( mat.Cstiff )

#   Setup boundary conditions
function funcB( p::Array{Float64} )
  return fill(0.0, size(p,1), 2)
end
funcNE = (p) -> [ σ1func(p)  σ2func(p)]
funcNW = (p) -> [-σ1func(p) -σ2func(p)]
bctype = [1,2,1,2] # All Dirichlet
#   Setup source function
source = (p) -> [F1func(p) F2func(p)]

## Compute error for different polynomial orders and grid sizes
#   Initialize arrays
Err_uh1 = fill( 0.0, length(Ps), length(Ns) )
Err_uh2 = fill( 0.0, length(Ps), length(Ns) )
Err_σh1 = fill( 0.0, length(Ps), length(Ns) )
Err_σh2 = fill( 0.0, length(Ps), length(Ns) )
Err_σh4 = fill( 0.0, length(Ps), length(Ns) )
# Err_ϵh1 = fill( 0.0, length(Ps), length(Ns) )
# Err_ϵh2 = fill( 0.0, length(Ps), length(Ns) )
# Err_ϵh4 = fill( 0.0, length(Ps), length(Ns) )

#   Loop over polynomial order and grid size
for ii in 1:length(Ps), jj in 1:length(Ns)

  P = Ps[ii]; N = Ns[jj]
  println("P ", P.p)

  mesh   = Mesh2D( "square", P, N = N)
  master = Master2D( P )

  prob = Elas( @sprintf("Reg %i %i", P.p, N), mat, source, bctype, false, [funcB, funcNE, funcB, funcNW] )

  (uhathTri, uh, σh) = hdgSolve( master, mesh, prob )

  #   Initialize arrays
  err_uh = fill( 0.0, dim )
  err_σh = fill( 0.0, dim^2 )
  # err_ϵh = fill( 0.0, dim^2 )

  # preallocate
  jcwd = fill( 0.0, size(master.∇ϕ,2), size(master.∇ϕ,2) )
  ∂ξ∂x = fill( 0.0, size(master.∇ϕ,2), dim^2 )
  ∂x∂ξ = fill( 0.0, size(master.∇ϕ,2), dim^2 )

  for kk in 1:size(mesh.t,1)

    # Compute Jacobians
    luteos.compJacob!( master, mesh.nodes[:,:,kk], ∂ξ∂x, jcwd, ∂x∂ξ )

    # u
    Δuh1       = master.ϕ' * ( uh[:,1,kk] - u1func( mesh.nodes[:,:,kk] ) )
    Δuh2       = master.ϕ' * ( uh[:,2,kk] - u2func( mesh.nodes[:,:,kk] ) )
    err_uh[1] += Δuh1' * jcwd * Δuh1
    err_uh[2] += Δuh2' * jcwd * Δuh2

    # σ
    Δσh1       = master.ϕ' * ( σh[:,1,kk] - σ1func( mesh.nodes[:,:,kk] ) )
    Δσh2       = master.ϕ' * ( σh[:,2,kk] - σ2func( mesh.nodes[:,:,kk] ) )
    Δσh4       = master.ϕ' * ( σh[:,4,kk] - σ4func( mesh.nodes[:,:,kk] ) )
    err_σh[1] += Δσh1' * jcwd * Δσh1
    err_σh[2] += Δσh2' * jcwd * Δσh2
    err_σh[4] += Δσh4' * jcwd * Δσh4

    # # ϵ
    # Δϵh1       = master.ϕ' * ( ϵh[:,1,kk] - ϵ1func( mesh.nodes[:,:,kk] ) )
    # Δϵh2       = master.ϕ' * ( ϵh[:,2,kk] - ϵ2func( mesh.nodes[:,:,kk] ) )
    # Δϵh4       = master.ϕ' * ( ϵh[:,4,kk] - ϵ4func( mesh.nodes[:,:,kk] ) )
    # err_ϵh[1] += Δϵh1' * jcwd * Δϵh1
    # err_ϵh[2] += Δϵh2' * jcwd * Δϵh2
    # err_ϵh[4] += Δϵh4' * jcwd * Δϵh4

  end

  Err_uh1[ii,jj] = sqrt(err_uh[1])
  Err_uh2[ii,jj] = sqrt(err_uh[2])
  Err_σh1[ii,jj] = sqrt(err_σh[1])
  Err_σh2[ii,jj] = sqrt(err_σh[2])
  Err_σh4[ii,jj] = sqrt(err_σh[4])
  # Err_ϵh1[ii,jj] = sqrt(err_ϵh[1])
  # Err_ϵh2[ii,jj] = sqrt(err_ϵh[2])
  # Err_ϵh4[ii,jj] = sqrt(err_ϵh[4])

end

# Compute convergence rates
h = 1 ./ ( Ns - 1 )
conv_uh1 = (log.( Err_uh1[:,end-2]) - log.( Err_uh1[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));
conv_uh2 = (log.( Err_uh2[:,end-2]) - log.( Err_uh2[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));

conv_σh1 = (log.( Err_σh1[:,end-2]) - log.( Err_σh1[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));
conv_σh2 = (log.( Err_σh2[:,end-2]) - log.( Err_σh2[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));
conv_σh4 = (log.( Err_σh4[:,end-2]) - log.( Err_σh4[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));

# conv_ϵh1 = (log.( Err_ϵh1[:,end-2]) - log.( Err_ϵh1[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));
# conv_ϵh2 = (log.( Err_ϵh2[:,end-2]) - log.( Err_ϵh2[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));
# conv_ϵh4 = (log.( Err_ϵh4[:,end-2]) - log.( Err_ϵh4[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));

# Output to terminal
@printf("\n")
@printf("   Convergence rates for 2D Dirichlet/Neumann problem\n\n")
@printf("   --------------------------------------------------\n\n")
@printf( "P   ")
for jj in 1:size(Ps,1)
  @printf( " %6i", Ps[jj].p )
end
@printf( "\n" )

#   u
@printf( "u₁  ")
for jj in 1:size(Ps,1)
  @printf( " %6.4f", conv_uh1[jj] )
end
@printf( "\n" )
@printf( "u₂  ")
for jj in 1:size(Ps,1)
  @printf( " %6.4f", conv_uh2[jj] )
end
@printf( "\n" )

#   σ
@printf( "σ₁  ")
for jj in 1:length(Ps)
  @printf( " %6.4f", conv_σh1[jj] )
end
@printf( "\n" )
@printf( "σ₂  ")
for jj in 1:length(Ps)
  @printf( " %6.4f", conv_σh2[jj] )
end
@printf( "\n" )
@printf( "σ₄  ")
for jj in 1:length(Ps)
  @printf( " %6.4f", conv_σh4[jj] )
end
@printf( "\n" )
#
# #   ϵ
# @printf( "ϵ₁  ")
# for jj in 1:length(Ps)
#   @printf( " %6.4f", conv_ϵh1[jj] )
# end
# @printf( "\n" )
# @printf( "ϵ₂  ")
# for jj in 1:length(Ps)
#   @printf( " %6.4f", conv_ϵh2[jj] )
# end
# @printf( "\n" )
# @printf( "ϵ₄  ")
# for jj in 1:length(Ps)
#   @printf( " %6.4f", conv_ϵh4[jj] )
# end
# @printf( "\n" )

open("errors_Elas_Neumann2D.dat", "w") do f
  @printf(f, "P \t N \t E_uh1 \t E_uh2 \t E_σh1 \t E_σh2 \t E_σh4 \n")#\t E_ϵh1 \t E_ϵh2 \t E_ϵh4 \n")#\t E_J\n")
  for ii in 1:length(Ps), jj in 1:length(Ns)
    @printf(f, "%i \t %i \t %16.15e \t %16.15e \t %16.15e \t %16.15e \t %16.15e \n",#\t %16.15e \t %16.15e \t %16.15e\n",
      Ps[ii].p, Ns[jj], Err_uh1[ii,jj], Err_uh2[ii,jj],
      Err_σh1[ii,jj], Err_σh2[ii,jj], Err_σh4[ii,jj] )#,
      #Err_ϵh1[ii,jj], Err_ϵh2[ii,jj], Err_ϵh4[ii,jj] )
  end
end

# end # limit scope


# Computes r_hat from Section 3.1.1 of the paper
function rhat(θ, Ψy)
    r = transpose(Ψy) * θ
    r .= exp.(r)
    r ./ mean(r)
end

# Computes the Hessian of the KLIEP loss
KLIEP_Hessian(θ, Ψy) = StatsBase.cov(Ψy, weights(rhat(θ, Ψy)), 2; corrected=false)

# Computes the sufficient statistics for an Ising mode
#
#    Ψ(x, y) = -1. if x == y
#    Ψ(x, y) =  1. if x != y
#
# Observations are columns of X
#
# Output
#    matrix with columns corresponding to observations
#    rows are sufficients statistics Ψ(xa, xb)
#    the order corresponds to the packed order where
#    lower triangular part is mapped row by row
function Ψising(X)
    m, n = size(X)
    p = div(m * (m - 1), 2)
    out = zeros(Float64, p, n)
    for i = 1:n
        ind = 0
        for row = 2:m
            for col = 1:row-1
                ind += 1
                out[ind, i] = xor(X[col, i], X[row, i]) ? -1. : 1.
            end
        end
    end

    # test that all the sufficient statistics
    # have variance different from zero
    for j=1:p
        if iszero(var(view(out, j, :)))
            @warn "a row of Ψ has zero variance"
        end
    end
    out
end

function _find_supp(k::Int, x)
    supp = findall(!iszero, x)
    pos_k = findfirst(isequal(k), supp)
    if pos_k == nothing
        push!(supp, k)
    else
        supp[pos_k], supp[end] = supp[end], supp[pos_k]
    end
    supp
end

function _find_supp(k::Int, x, y)
    supp1 = findall(!iszero, x)
    supp2 = findall(!iszero, y)
    supp3 = union(supp1, supp2)
    pos_k = findfirst(isequal(k), supp3)
    if pos_k == nothing
        push!(supp3, k)
    else
        supp3[pos_k], supp3[end] = supp3[end], supp3[pos_k]
    end
    supp3
end

# maps (i, j) coordinate of a symmetric matrix to a packed format
# diagonal of the matrix is not included
# upper triangular part is mapped col by col
function trimap(i::Integer, j::Integer)
    if i > j
        trimap(j, i)
    else
        i + div((j - 1) * (j - 2), 2)
    end
end

# mapping from packed storage to (i, j)
function itrimap(k::Integer)
    j = convert(Int, ceil((1. + sqrt(1. + 8. * k)) / 2.))
    i = k - div((j - 1) * (j - 2), 2)
    CartesianIndex(i, j)
end

# vector -> symmetric matrix
function unpack(θ)
    m = length(θ)
    p = convert(Int, ceil((1. + sqrt(1. + 8. * m)) / 2.))
    out = zeros(Float64, p, p)
    for i=1:m
        out[itrimap(i)] = θ[i]
    end
    out + out'
end

# symmetric matrix -> vector
function pack(Θ)
    m = size(Θ, 1)
    p = div(m * (m - 1), 2)
    out = zeros(Float64, p)
    for j = 1:(m-1)
        for i = (j+1):m
            out[trimap(i,j)] = Θ[i,j]
        end
    end
    out
end

####################################
#
# loss En[-θ'Ψx(i)] + log( En[exp(θ'Ψy(i))] )
#
####################################
struct CDKLIEPLoss <: CoordinateDifferentiableFunction
    μx::Vector{Float64}   # mean(Ψx, dims = 2)
    Ψy::Matrix{Float64}
    r::Vector{Float64}    # ny dim vector --- stores θ'Ψy(i)
    p::Int64
end

function CDKLIEPLoss(Ψx::Matrix{Float64}, Ψy::Matrix{Float64})
    (p = size(Ψx, 1)) == size(Ψy, 1) || throw(DimensionMismatch())

    CDKLIEPLoss(vec(mean(Ψx, dims=2)), Ψy, zeros(size(Ψy, 2)), p)
end

CoordinateDescent.numCoordinates(f::CDKLIEPLoss) = f.p

function CoordinateDescent.initialize!(f::CDKLIEPLoss, x::SparseIterate)

    mul!(f.r, transpose(f.Ψy), x)

    nothing
end

function CoordinateDescent.gradient(
    f::CDKLIEPLoss,
    x::SparseIterate,
    j::Int64)

    μx, Ψy, r = f.μx, f.Ψy, f.r

    -μx[j] + mean( exp.(r) .* Ψy[j, :] ) / mean(exp, r)
end

function CoordinateDescent.descendCoordinate!(
    f::CDKLIEPLoss,
    g::Union{ProxL1, ProxZero},
    x::SparseIterate,
    j::Int64)

    μx, Ψy, r = f.μx, f.Ψy, f.r

    mean_exp_r = mean(exp, r)
    mean_exp_r_Ψy = mean( exp.(r) .* Ψy[j, :] )
    mean_exp_r_Ψy2 = mean( exp.(r) .* Ψy[j, :].^2. )

    @inbounds grad = -μx[j] + mean_exp_r_Ψy / mean_exp_r
    H = mean_exp_r_Ψy2 / mean_exp_r - mean_exp_r_Ψy^2 / mean_exp_r^2.

    @inbounds oldVal = x[j]
    @inbounds x[j] -= grad / H
    newVal = cdprox!(g, x, j, 1. / H)
    h = newVal - oldVal

    # update internals
    for i=1:length(r)
        @inbounds r[i] += h*Ψy[j, i]
    end
    h
end


"""
    power_spectrum_3d(::AtmosLESConfigType, u, v, w, L, dim, nor)

Calculates the Powerspectrum from the 3D velocity fields `u`, `v`, `w`. Inputs
need to be equi-spaced and the domain is assumed to be the same size and
have the same number of points in all directions.

# Arguments
- L: size domain
- dim: number of points
- nor: normalization factor
"""
function power_spectrum_3d(::AtmosLESConfigType, u, v, w, L, dim, nor)
    u_fft = fft(u)
    v_fft = fft(v)
    w_fft = fft(w)

    mu = Array((abs.(u_fft) / size(u, 1)^3))
    mv = Array((abs.(v_fft) / size(v, 1)^3))
    mw = Array((abs.(w_fft) / size(w, 1)^3))
    if mod(dim, 2) == 0
        rx = range(0, stop = dim - 1, step = 1) .- dim / 2 .+ 1
        ry = range(0, stop = dim - 1, step = 1) .- dim / 2 .+ 1
        rz = range(0, stop = dim - 1, step = 1) .- dim / 2 .+ 1
        R_x = circshift(rx', (1, dim / 2 + 1))
        R_y = circshift(ry', (1, dim / 2 + 1))
        R_z = circshift(rz', (1, dim / 2 + 1))
        k_nyq = dim / 2
    else
        rx = range(0, stop = dim - 1, step = 1) .- (dim - 1) / 2
        ry = range(0, stop = dim - 1, step = 1) .- (dim - 1) / 2
        rz = range(0, stop = dim - 1, step = 1) .- (dim - 1) / 2
        R_x = circshift(rx', (1, (dim + 1) / 2))
        R_y = circshift(ry', (1, (dim + 1) / 2))
        R_z = circshift(rz', (1, (dim + 1) / 2))
        k_nyq = (dim - 1) / 2
    end
    r = zeros(size(rx, 1), size(ry, 1), size(rz, 1))
    for i in 1:size(rx, 1), j in 1:size(ry, 1), l in 1:size(rz, 1)
        r[i, j, l] = sqrt(R_x[i]^2 + R_y[j]^2 + R_z[l]^2)
    end
    dx = 2 * pi / L
    k = range(1, stop = k_nyq, step = 1) .* dx
    endk = size(k, 1)
    contribution = zeros(size(k, 1))
    spectrum = zeros(size(k, 1))
    for N in 2:Int64(k_nyq - 1)
        for i in 1:size(rx, 1), j in 1:size(ry, 1), l in 1:size(rz, 1)
            if (r[i, j, l] * dx <= (k'[N + 1] + k'[N]) / 2) &&
               (r[i, j, l] * dx > (k'[N] + k'[N - 1]) / 2)
                spectrum[N] =
                    spectrum[N] + mu[i, j, l]^2 + mv[i, j, l]^2 + mw[i, j, l]^2
                contribution[N] = contribution[N] + 1
            end
        end
    end
    for i in 1:size(rx, 1), j in 1:size(ry, 1), l in 1:size(rz, 1)
        if (r[i, j, l] * dx <= (k'[2] + k'[1]) / 2)
            spectrum[1] =
                spectrum[1] + mu[i, j, l]^2 + mv[i, j, l]^2 + mw[i, j, l]^2
            contribution[1] = contribution[1] + 1
        end
    end
    for i in 1:size(rx, 1), j in 1:size(ry, 1), l in 1:size(rz, 1)
        if (r[i, j, l] * dx <= k'[endk]) &&
           (r[i, j, l] * dx > (k'[endk] + k'[endk - 1]) / 2)
            spectrum[endk] =
                spectrum[endk] + mu[i, j, l]^2 + mv[i, j, l]^2 + mw[i, j, l]^2
            contribution[endk] = contribution[endk] + 1
        end
    end
    spectrum = spectrum .* 2 .* pi .* k .^ 2 ./ (contribution .* dx .^ 3) ./ nor

    return spectrum, k
end


function solve_bvp(problem::FiniteDifferenceBVPProblem{TF, TBCL, TBCR}) where {TF, TBCL<:DirichletBC, TBCR<:DirichletBC}
    u = problem.A\problem.rhs;
    return u
end

function solve_bvp(problem::FiniteDifferenceBVPProblem{TF, TBCL, TBCR}) where {TF, TBCL<:NeumannBC, TBCR<:NeumannBC}
    u = qr(problem.A)\problem.rhs;
    return u
end

function solve_bvp(problem::FiniteDifferenceBVPProblem{TF, TBCL, TBCR}) where {TF, TBCL<:PeriodicBC, TBCR<:PeriodicBC}
    u = qr(problem.A)\problem.rhs;
    return u
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    intersecttype(f, segment, triangle)

Calculate the intersection type of `segment and `triangle`
and apply function `f` to it.

## References

* Jiménez, J., Segura, R. and Feito, F. 2009. [A robust segment/triangle
  intersection algorithm for interference tests. Efficiency study]
  (https://www.sciencedirect.com/science/article/pii/S0925772109001448?via%3Dihub)
"""
function intersecttype(f::Function, s::Segment{3,T}, t::Triangle{3,T}) where {T}
  vₛ = vertices(s)
  vₜ = vertices(t)

  A = vₛ[1] - vₜ[3]
  B = vₜ[1] - vₜ[3]
  C = vₜ[2] - vₜ[3]

  W₁ = B × C
  wᵥ = A ⋅ W₁

  D = vₛ[2] - vₜ[3]
  sᵥ = D ⋅ W₁

  if wᵥ > atol(T)
    # rejection 2
    if sᵥ > atol(T)
      return NoIntersection() |> f
    end

    W₂ = A × D
    tᵥ = W₂ ⋅ C

    # rejection 3
    if tᵥ < -atol(T)
      return NoIntersection() |> f
    end

    uᵥ = -(W₂ ⋅ B)

    # rejection 4
    if uᵥ < -atol(T)
      return NoIntersection() |> f
    end

    # rejection 5
    if wᵥ < (sᵥ + tᵥ + uᵥ)
      return NoIntersection() |> f
    end
  elseif wᵥ < -atol(T)
    # rejection 2
    if sᵥ < -atol(T)
      return NoIntersection() |> f
    end

    W₂ = A × D
    tᵥ = W₂ ⋅ C

    # rejection 3
    if tᵥ > atol(T)
      return NoIntersection() |> f
    end

    uᵥ = -(W₂ ⋅ B)

    # rejection 4
    if uᵥ > atol(T)
      return NoIntersection() |> f
    end

    # rejection 5
    if wᵥ > (sᵥ + tᵥ + uᵥ)
      return NoIntersection() |> f
    end
  else
    if sᵥ > atol(T)
      W₂ = D × A
      tᵥ = W₂ ⋅ C

      # rejection 3
      if tᵥ < -atol(T)
        return NoIntersection() |> f
      end

      uᵥ = -(W₂ ⋅ B)

      # rejection 4
      if uᵥ < -atol(T)
        return NoIntersection() |> f
      end
      # rejection 5
      if -sᵥ < (tᵥ + uᵥ)
        return NoIntersection() |> f
      end
    elseif sᵥ < -atol(T)
      W₂ = D × A
      tᵥ = W₂ ⋅ C

      # rejection 3
      if tᵥ > atol(T)
        return NoIntersection() |> f
      end

      uᵥ = -(W₂ ⋅ B)

      # rejection 4
      if uᵥ > atol(T)
        return NoIntersection() |> f
      end

      # rejection 5
      if -sᵥ > (tᵥ + uᵥ)
        return NoIntersection() |> f
      end
    else
      # rejection 1, coplanar segment
      return NoIntersection() |> f
    end
  end

  λ = clamp(wᵥ / (wᵥ - sᵥ), zero(T), one(T))

  return IntersectingSegmentTriangle(s(λ)) |> f
end

"""
    intersecttype(s, t)

Calculate the intersection of a ray and triangle in 3D
and apply function `f` to it.

## References

* Möller, T. and Trumbore, B., 1997. [Fast, minimum storage ray-triangle
  intersection. Journal of graphics tools]
  (https://www.tandfonline.com/doi/abs/10.1080/10867651.1997.10487468)
"""
function intersecttype(f::Function, r::Ray{3,T}, t::Triangle{3,T}) where {T}
  vs = vertices(t)
  o = origin(r)
  d = direction(r)

  e₁ = vs[3] - vs[1]
  e₂ = vs[2] - vs[1]
  p = d × e₂
  det = e₁ ⋅ p

  # keep det > 0, modify T accordingly
  if det > atol(T)
    τ = o - vs[1]
  else
    τ = vs[1] - o
    det = -det
  end

  if det < atol(T)
    # This ray is parallel to the plane of the triangle.
    return NoIntersection() |> f
  end

  # calculate u parameter and test bounds
  u = τ ⋅ p
  if u < -atol(T) || u > det
    return NoIntersection() |> f
  end

  q = τ × e₁

  # calculate v parameter and test bounds
  v = d ⋅ q
  if v < -atol(T) || u + v > det
    return NoIntersection() |> f
  end

  λ = (e₂ ⋅ q) * (one(T) / det)

  if λ < -atol(T)
    return NoIntersection() |> f
  end

  λ = clamp(λ, zero(T), typemax(T))

  return IntersectingRayTriangle(r(λ)) |> f
end

"""
    intersecttype(s, p)

Calculate the intersection of a ray and polygon in 3D
and apply function `f` to it.
"""
function intersecttype(f::Function, r::Ray{3,T}, p::Ngon{N,3,T}) where {N,T}
  for t in triangulate(p)
    res = intersecttype(identity, r, t)
    if res != NoIntersection()
      return res |> f
    end
  end
  return NoIntersection() |> f
end


function nonsym_xxz_ham(;spin = 1,delta = 1,zfield = 0.0)
    (sx,sy,sz,id) = nonsym_spintensors(spin)

    mpo=Array{typeof(sx),3}(undef,1,5,5)
    mpo[1,:,:]=[id sx sy delta*sz zfield*sz;
                0*id 0*id 0*id 0*id sx;
                0*id 0*id 0*id 0*id sy;
                0*id 0*id 0*id 0*id sz;
                0*id 0*id 0*id 0*id id]

    return MPOHamiltonian(mpo)
end

function su2_xxx_ham(;spin = 1//2)
    #only checked for spin = 1 and spin = 2...
    ph = ℂ[SU₂](spin=>1)

    Sl1 = TensorMap(ones, Defaults.eltype, ℂ[SU₂](0=>1)*ph , ℂ[SU₂](1=>1)*ph)*sqrt(spin^2+spin)
    Sr1 = TensorMap(ones, Defaults.eltype, ℂ[SU₂](1=>1)*ph , ℂ[SU₂](0=>1)*ph)*sqrt(spin^2+spin)

    return MPOHamiltonian([Sl1,Sr1]);
end

function u1_xxz_ham(;spin = 1,delta = 1,zfield = 0.0)
    (sxd,syd,szd,idd) = spinmatrices(spin);
    @tensor ham[-1 -2;-3 -4]:=sxd[-1,-3]*sxd[-2,-4]+syd[-1,-3]*syd[-2,-4]+(delta*szd)[-1,-3]*szd[-2,-4]+zfield*0.5*szd[-1,-3]*idd[-2,-4]+zfield*0.5*idd[-1,-3]*szd[-2,-4]

    indu1map = [U₁(v) for v in -spin:1:spin];
    pspace = U1Space((v=>1 for v in indu1map));

    symham = TensorMap(zeros,eltype(ham),pspace*pspace,pspace*pspace)

    for (i,j,k,l) in Iterators.product(1:size(ham,1),1:size(ham,1),1:size(ham,1),1:size(ham,1))
        if ham[i,j,k,l]!=0
            copyto!(symham[(indu1map[i],indu1map[j],indu1map[k],indu1map[l])],ham[i,j,k,l])
        end
    end

    return MPOHamiltonian(decompose_localmpo(add_util_leg(symham)))
end


println("Testing...")
using Test
using SeqAdaptiveIS
using StatsFuns
using Distributions: MvNormal, logpdf
using Statistics: mean, cov

# These are fairly rudimentary 2D tests with easy targets: I guess their primary
# purpose is to ensure that nothing has *broken* the procedures, but is not
# sufficiently strong enough yet to tell if it breaks in difficult or edge
# cases.

# ======================== AMIS inference tests ================================
p_x = MvNormal(zeros(2), Float64[1 0; 0 1])
test_lpdf64_v0 = (x-> logpdf(MvNormal(zeros(2), Float64[1 0; 0 1]), x'), p_x, zeros(2), Float64[1 0; 0 1])
test_lpdf64_v1 = (x-> logpdf(MvNormal(ones(2), Float64[1 0; 0 1]), x'), p_x, ones(2), Float64[1 0; 0 1])

function test_amis(test_case, k::Int)
    log_target, prior, μ, Σ = test_case
    S, logW, pis, mus, covs = amis(log_target, prior, k)
    W = softmax(logW)
    μ̂ = vec(sum(S .* W, dims=1))
    Σ̂ = (S .- μ̂')' * ((S .- μ̂') .* W)
    @test maximum(abs.(μ - μ̂)) < 0.1
    @test maximum(vec(Σ - Σ̂)) < 0.3
end

test_amis(test_lpdf64_v0, 1)
test_amis(test_lpdf64_v0, 2)
test_amis(test_lpdf64_v0, 5)
test_amis(test_lpdf64_v1, 1)
test_amis(test_lpdf64_v1, 2)
test_amis(test_lpdf64_v1, 5)

# ============================= Seq AMIS test ==================================

n_targets = 8
true_pars = [(zeros(2) .+ i, Float64[1 0; 0 1] ./ i) for i in 1:n_targets];
seq_targets = [x->logpdf(MvNormal(μ, Σ), x') for (μ, Σ) in true_pars]

seq_smps, seq_gmms = seq_amis(seq_targets, 2, 3)
print("Sequential targets complete: ")
for i in 1:n_targets
    μ, Σ = true_pars[i]
    S = resample(seq_smps[i], 1000)
    μ̂ = vec(mean(S, dims=1))
    Σ̂ = cov(S, dims=1)
    print("|")
    @test maximum(abs.(μ - μ̂)) < 0.1
    @test maximum(vec(Σ - Σ̂)) < 0.3
end
println()


@testset "553.optimal-division.jl" begin
    @test optimal_division([1000, 100, 10, 2]) == "1000/(100/10/2)"
    @test optimal_division([1000, 100, 10, 2]) == "1000/(100/10/2)"
end


# Adjoint test for F and J
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January 2017
#

using PyCall
using PyPlot
using opesciSLIM.TimeModeling

## Set up model structure
n = (160, 170)	# (x,y,z) or (x,z)
d = (15.,15.)
o = (0.,0.)

# Velocity [km/s]
v = ones(Float32,n) + 0.5f0
v[:,Int(round(end/3)):end] = 4f0
v0 = ones(Float32,n) + 0.5f0

# Slowness squared [s^2/km^2]
m = (1f0./v).^2
m0 = (1f0./v0).^2
dm = vec(m - m0)
rho =  v - .5f0
rho0 = v0 - .5f0

# Setup info and model structure
nsrc = 1
model = Model(n,d,o,m, rho)
model0 = Model(n,d,o,m0, rho0)

## Set up receiver geometry
nxrec = 160
xrec = linspace(100f0,1500f0,nxrec)
yrec = 0f0
zrec = linspace(30f0,30f0,nxrec)

# receiver sampling and recording time
timeR = 1500f0	# receiver recording time [ms]
dtR = 4.2f0	# receiver sampling interval

# Set up receiver structure
recGeometry = Geometry(xrec,yrec,zrec;dt=dtR,t=timeR,nsrc=nsrc)

## Set up source geometry (cell array with source locations for each shot)
xsrc = 1200f0
ysrc = 0f0
zsrc = 30f0

# source sampling and number of time steps
timeS = 1500f0
dtS = 4.2f0

# Set up source structure
srcGeometry = Geometry(xsrc,ysrc,zsrc;dt=dtS,t=timeS)

# Info structure
ntComp = get_computational_nt(srcGeometry,recGeometry,model)
info = Info(prod(n), nsrc, ntComp)

# setup wavelet
f0 = 0.01f0
wavelet = source(timeS,dtS,f0)
wave_rand = wavelet.*rand(Float32,size(wavelet))
###################################################################################################

# Modeling operators
F = joModeling(info,model,srcGeometry,recGeometry)
q = joData(srcGeometry,wavelet)
#
# Nonlinear modeling
d_hat = F*q

# Generate random noise data vector with size of d_hat
# in the range of F
qr = joData(srcGeometry,wave_rand)
d1 = F*qr

# Adjoint computation
q_hat = F'*d1

# Result F
println(abs(dot(d1,d_hat)))
println(abs(dot(q,q_hat)))
println("Residual: ", abs(dot(d1,d_hat) - dot(q,q_hat)))
println("Ratio: ", abs(dot(d1,d_hat)/dot(q,q_hat)))

# Linearized modeling
F0 = joModeling(info,model0,srcGeometry,recGeometry)
J = joJacobian(F0,q)

dD_hat = J*dm
dm_hat = J'*dD_hat

# Result J
println(dot(dD_hat,dD_hat))
println(dot(dm,dm_hat))
println("Residual: ", abs(dot(dD_hat,dD_hat) - dot(dm,dm_hat)))
println("Ratio: ", abs(dot(dD_hat,dD_hat)/dot(dm,dm_hat)))


# ---
# title: 982. Triples with Bitwise AND Equal To Zero
# id: problem982
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/triples-with-bitwise-and-equal-to-zero/description/>
# hidden: true
# ---
#
# Given an array of integers `A`, find the number of triples of indices (i, j,
# k) such that:
#
#   * `0 <= i < A.length`
#   * `0 <= j < A.length`
#   * `0 <= k < A.length`
#   * `A[i] & A[j] & A[k] == 0`, where `&` represents the bitwise-AND operator.
#
#
#
# **Example 1:**
#
#
#
#     Input: [2,1,3]
#     Output: 12
#     Explanation: We could choose the following i, j, k triples:
#     (i=0, j=0, k=1) : 2 & 2 & 1
#     (i=0, j=1, k=0) : 2 & 1 & 2
#     (i=0, j=1, k=1) : 2 & 1 & 1
#     (i=0, j=1, k=2) : 2 & 1 & 3
#     (i=0, j=2, k=1) : 2 & 3 & 1
#     (i=1, j=0, k=0) : 1 & 2 & 2
#     (i=1, j=0, k=1) : 1 & 2 & 1
#     (i=1, j=0, k=2) : 1 & 2 & 3
#     (i=1, j=1, k=0) : 1 & 1 & 2
#     (i=1, j=2, k=0) : 1 & 3 & 2
#     (i=2, j=0, k=1) : 3 & 2 & 1
#     (i=2, j=1, k=0) : 3 & 1 & 2
#
#
#
#
# **Note:**
#
#   1. `1 <= A.length <= 1000`
#   2. `0 <= A[i] < 2^16`
#
#
## @lc code=start
using LeetCode

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


# caeser cipher
const n = length('A':'Z')
# function encrypt(ch::Char, shift::Integer)
#     if ch in 'A':'Z'
#         'A' + mod((ch - 'A') + shift, n)
#     elseif ch in 'a':'z'
#         'a' + mod((ch - 'a') + shift, n)
#     else
#         ch
#     end
# end

# function decrypt(ch::Char, shift::Integer)
#     if ch in 'A':'Z'
#         'A' + mod((ch - 'A') - shift, n)
#     elseif ch in 'a':'z'
#         'a' + mod((ch - 'a') - shift, n)
#     else
#         ch
#     end
# end

# more elegant solution based on discussion here:
# https://discourse.julialang.org/t/a-switch-type-function/13491/15
function encryptc(ch::Char, shift::Integer)
    ch in 'A':'Z' ? 'A' + mod((ch - 'A') + shift, n) :
    ch in 'a':'z' ? 'a' + mod((ch - 'a') + shift, n) :
    #= otherwise =# ch
end

function decryptc(ch::Char, shift::Integer)
    ch in 'A':'Z' ? 'A' + mod((ch - 'A') - shift, n) :
    ch in 'a':'z' ? 'a' + mod((ch - 'a') - shift, n) :
    #= otherwise =# ch
end

# Example
# julia> message = "This is the end of the world!"
# "This is the end of the world"
#
# julia> shift = 3
# 3
#
# julia> emessage = map(message) do char
#     encryptc(char, shift)
# end
# "Wklv lv wkh hqg ri wkh zruog!"

# julia> map(emessage) do char
#     decryptc(char, shift)
# end
# "This is the end of the world!"

# substition cipher
using Random
# Engheim's definition is wrong - he uses 'char' for 'ch'
function encrypts(ch::Char, updict::Dict{Char, Char}, lowdict::Dict{Char, Char})
    ch in 'A':'Z' ? get(updict, ch, ch) :
    ch in 'a':'z' ? get(lowdict, ch, ch) :
    #= otherwise =# ch
end

# Also don't think decrypt is confusing as Engheim claims
function decrypts(ch::Char, rupdict::Dict{Char, Char}, rlowdict::Dict{Char, Char})
    ch in 'A':'Z' ? get(rupdict, ch, ch) :
    ch in 'a':'z' ? get(rlowdict, ch, ch) :
    #= otherwise =# ch
end

# example
# upshuf=shuffle('A':'Z')
# lowshuf=shuffle('a':'z')
# julia> updict = Dict(zip('A':'Z', upshuf))
# Dict{Char, Char} with 26 entries:
#   'E' => 'D'
#   'Z' => 'C'
#   'X' => 'G'
#   'B' => 'U'
#   'C' => 'T'
#   'D' => 'Y'
#   'A' => 'S'
#   'R' => 'F'
#   'G' => 'A'
#   'F' => 'P'
#   'N' => 'Q'
#   'M' => 'M'
#   'K' => 'V'
#   ⋮   => ⋮
#
#   julia> lowdict = Dict(zip('a':'z', lowshuf))
#   Dict{Char, Char} with 26 entries:
#     'n' => 'x'
#     'f' => 'b'
#     'w' => 'y'
#     'd' => 'a'
#     'e' => 'f'
#     'o' => 'l'
#     'h' => 'c'
#     'j' => 'm'
#     'i' => 't'
#     'k' => 'z'
#     'r' => 'v'
#     's' => 'q'
#     't' => 'w'
#     ⋮   => ⋮
#
#   julia> encrypts('a', updict, lowdict)
#   'i': ASCII/Unicode U+0069 (category Ll: Letter, lowercase)
#
#   julia> encrypts('?', updict, lowdict)
#   '?': ASCII/Unicode U+003F (category Po: Punctuation, other)
#
#   julia> encrypts('Z', updict, lowdict)
#   'C': ASCII/Unicode U+0043 (category Lu: Letter, uppercase)
#
# The idea for how to reverse the dictionary comes from here:
# https://discourse.julialang.org/t/is-something-like-reversed-dict-findall-x-x-house-cc-is-to-slow/16443/2
# julia> rupdict = Dict(values(updict) .=> keys(updict))
# Dict{Char, Char} with 26 entries:
#   'E' => 'P'
#   'Z' => 'O'
#   'X' => 'Q'
#   'C' => 'Z'
#   'B' => 'Y'
#   'D' => 'E'
#   'A' => 'G'
#   'R' => 'T'
#   'G' => 'X'
#   'F' => 'R'
#   'Q' => 'N'
#   'P' => 'F'
#   'M' => 'M'
#   ⋮   => ⋮
#
# julia> rlowdict = Dict(values(lowdict) .=> keys(lowdict))
# Dict{Char, Char} with 26 entries:
#   'n' => 'g'
#   'f' => 'e'
#   'w' => 't'
#   'd' => 'q'
#   'e' => 'v'
#   'o' => 'm'
#   'h' => 'z'
#   'y' => 'w'
#   'i' => 'a'
#   't' => 'i'
#   's' => 'l'
#   'r' => 'u'
#   'j' => 'x'
#   ⋮   => ⋮
# julia> message = "This is the end of the world!"
# "This is the end of the world!"

# julia> emessage = map(message) do char
#            encrypts(char,updict,lowdict)
#         end
# "Rctq tq wcf fxa lb wcf ylvsa!"

# julia> map(emessage) do char
#            decrypts(char,rupdict,rlowdict)
#         end
# "This is the end of the world!"


# Utilities for Robust Control Toolbox
using LinearAlgebra

function eye(n)::Matrix{Float64}
	return Matrix{Float64}(I, n, n)
end

function linspace(x1::Real,x2::Real,N::Integer)::Vector{Float64}
    if N>1
        dx = (x2-x1)/(N-1);
        return(x1:dx:x2);
    else
        return([]);
    end
end

function logspace(x1::Real,x2::Real,N::Integer)::Vector{Float64}
    if min(x1,x2) <= 0
        error("logspace is undefined for negative values");
        return([]);
    else
        x = linspace(log(x1),log(x2),N);
        return(exp.(x));
    end
end


using Test
using Quiqbox
using Quiqbox: isFull, BasisFuncMix, unpackBasisFuncs, inSymbols, varVal, ElementNames,
               sortBasisFuncs, ParamList, sumOf, expressionOfCore, mergeGaussFuncs,
               ijkOrbitalList, gaussProd, normOfGTOin
using Symbolics
using LinearAlgebra

@testset "Basis.jl" begin

# struct GaussFunc
exp1, con1 = (2, 1.0)
gf1 = GaussFunc(exp1, con1)
gf2 = GaussFunc(exp1, con1)
@test typeof(gf1) == typeof(gf2)
@test (gf1, gf2) isa Tuple{GaussFunc, GaussFunc}
@test (gf1, gf2) isa NTuple{2, GaussFunc}
gf1_ref = gf1
@test gf1.xpn == ParamBox(exp1, :xpn)
@test gf1.con == ParamBox(con1, :con)
@test gf1.param == (gf1.xpn, gf1.con)
@test hasEqual(gf1, gf2)
@test hasIdentical(gf1, gf1)
@test !hasIdentical(gf1, gf2)
@test gf1 !== gf2
@test (gf1 === gf1_ref)
@test hasIdentical(gf1, gf1_ref)
gf1_ref.xpn[] = 3.0
@test gf1.xpn[] == gf1.param[1][] == 3.0
gf1.param[2][] = 4
@test gf1_ref.con[] == gf1_ref.param[2][] == 4
@test gf1_ref === gf1
@test gf2 != gf1
x = [2.0, 1]
for i=1:2 gf1.param[i][] = x[i] end
@test gf1 === gf1_ref !== gf2
@test hasEqual(gf1, gf2)

# function genExponent genContraction
m1 = x->x^2
v1 = round(rand(), digits=6)
e1 = genExponent(v1, m1)
@test typeof(e1).parameters[2] == ParamList[:xpn]
@test e1[] == v1
@test e1() == m1(v1)

c1 = genContraction(v1, m1)
@test typeof(c1).parameters[2] == ParamList[:con]
@test c1[] == v1
@test e1() == m1(v1)

e2 = genExponent(e1)
e3 = genExponent(c1)
@test hasIdentical(e1, e2)
@test !hasEqual(e1, e3)
@test e1[] == e3[]
@test e1() == e3()
@test e1.map !== e3.map
v2 = rand()
@test e1.map(v2) == e3.map(v2)

c2 = genExponent(c1)
c3 = genExponent(e1)
@test hasIdentical(c1, c2)
@test !hasEqual(c1, c3)
@test c1[] == c3[]
@test c1() == c3()
@test c1.map !== c3.map
@test c1.map(v2) == c3.map(v2)

# function genSpatialPoint
k = fill(1.0)
v1 = genSpatialPoint([k,2,3  ])
v2 = genSpatialPoint([k,2,3.0])
@test v1 == v2
@test hasEqual(v1, v2)
@test !hasIdentical(v1, v2)
@test hasIdentical(v1[1], v2[1])
@test !hasIdentical(v1[2:3], v2[2:3])

# struct BasisFunc and function genBasisFunc
cen = [1,2,3]
cenParNames = [ParamList[:X], ParamList[:Y], ParamList[:Z]]
cenPar = ParamBox.(cen, cenParNames) |> Tuple
bf1 = genBasisFunc(cen, [gf1])
bf1_1 = genBasisFunc(cen, [gf1])
bf1_2 = genBasisFunc(bf1)
bf1_3 = genBasisFunc(cenPar, gf1)
bf1_4 = genBasisFunc(cenPar, gf1)
bf2 = genBasisFunc(cen, [gf2])
@test bf1 !== bf2
@test hasEqual(bf1, bf1_1, bf1_2, bf1_3, bf1_4, bf2)
@test hasIdentical(bf1, bf1_2)
@test !hasIdentical(bf1, bf1_1)
@test !hasIdentical(bf1, bf2)
@test !hasIdentical(bf1, bf1_3)
@test hasIdentical(bf1_3, bf1_4)
@test bf1.subshell == "S"
@test bf1 isa BasisFunc

bf11 = genBasisFunc(cen, [gf1, gf1])
bf1_3 = genBasisFunc(cen, (exp1, con1))
@test hasEqual(bf1, bf1_3)
@test !hasIdentical(bf1, bf1_3)
bf11_2 = genBasisFunc(cen, ([exp1, exp1], [con1, con1]))
@test hasEqual(bf11, bf11_2)
@test !hasIdentical(bf11, bf11_2)
bf2_P_norm2 = genBasisFunc(cen, [gf2], "P")
@test bf2_P_norm2.subshell == "P"
@test bf2_P_norm2 isa BasisFuncs

bf3_1 = genBasisFunc([0,0,0], (3,1))
bf3_2 = genBasisFunc([1,0,0], (2,1))
bf3_3 = genBasisFunc([0,0,0], (2,1), "P")
bf3_4 = genBasisFunc([0,0,0], (1,1), "D")
@test genBasisFunc([bf3_1, bf3_4, bf3_3, bf3_2]) == [bf3_1, bf3_3, bf3_4, bf3_2]

bf4_1 = genBasisFunc([0,0,0], "STO-3G")[]
bf4_2 = genBasisFunc([0,0,0], ("STO-3G", "He"))[]
bf4_3 = genBasisFunc([0,0,0], "STO-3G", nucleus = "He")[]
bf4s1 = genBasisFunc([0,0,0], ["STO-3G", "STO-3G"])
bf4s2 = genBasisFunc([0,0,0], ["STO-3G", "STO-3G"], nucleus = "He")
bf4s3 = genBasisFunc([0,0,0], [("STO-3G", "H"), ("STO-3G", "He")])
@test hasEqual.(Ref(bf4_1), bf4s1, Ref(bf4s3[1])) |> prod
@test hasEqual.(Ref(bf4_2), Ref(bf4_3), bf4s2, Ref(bf4s3[2])) |> prod
@test isapprox(1, overlap(bf4_1, bf4_1)[], atol=1e-8)
@test isapprox(1, overlap(bf4_2, bf4_2)[], atol=1e-8)

errorThreshold1 = 1e-11
bf3_2_2 = genBasisFunc([1,0,0], (2,1), normalizeGTO=true)
@test isapprox(1, overlap(bf3_2_2, bf3_2_2)[], atol=errorThreshold1)
@test isapprox(0.0553891828418, overlap(bf3_2, bf3_2)[], atol=errorThreshold1)
bf3_3_2 = genBasisFunc([0,0,0], (2,1), "P", normalizeGTO=true)
@test isapprox(LinearAlgebra.I, overlap(bf3_3_2, bf3_3_2), atol=errorThreshold1)
@test isapprox(0.0207709435653*LinearAlgebra.I, overlap(bf3_3, bf3_3), atol=errorThreshold1)
bf3_4_2 = genBasisFunc([0,0,0], (1,1), "D", normalizeGTO=true)
@test isapprox([1.0 0.0 0.0 1/3 0.0 1/3;
                0.0 1/3 0.0 0.0 0.0 0.0;
                0.0 0.0 1/3 0.0 0.0 0.0;
                1/3 0.0 0.0 1.0 0.0 1/3;
                0.0 0.0 0.0 0.0 1/3 0.0;
                1/3 0.0 0.0 1/3 0.0 1.0],
                overlap(bf3_4_2, bf3_4_2), atol=errorThreshold1)
@test isapprox([0.3691314831028692 0.0 0.0 0.1230438277009564 0.0 0.1230438277009564;
                0.0 0.1230438277009564 0.0 0.0 0.0 0.0;
                0.0 0.0 0.1230438277009564 0.0 0.0 0.0;
                0.1230438277009564 0.0 0.0 0.3691314831028692 0.0 0.1230438277009564;
                0.0 0.0 0.0 0.0 0.1230438277009564 0.0;
                0.1230438277009564 0.0 0.0 0.1230438277009564 0.0 0.3691314831028692],
                overlap(bf3_4, bf3_4), atol=errorThreshold1)

# function sortBasisFuncs
bfs1 = [genBasisFunc([1,1,1], (2,1), [1,0,0]), genBasisFunc([1,1,1], (3,1), [2,0,0]),
        genBasisFunc([1,1,2], (3,1), [0,0,0]), genBasisFunc([1,1,1], (3,1), "P")]
bfs2 = [genBasisFunc([1,1,1], (2,1), [1,0,0]), genBasisFunc([1,1,1], (3,1), "P"),
        genBasisFunc([1,1,1], (3,1), [2,0,0]), genBasisFunc([1,1,2], (3,1), [0,0,0])]
bfs3 = sortBasisFuncs(bfs1, groupCenters=true)
@test length.(bfs3) == [3,1]
bfs3 = bfs3 |> flatten
@test !hasEqual(bfs1, bfs2)
@test  hasEqual(bfs3, bfs2)

# function isFull
@test isFull(bf3_1) == true
@test isFull(bf3_3) == true
bf5 = genBasisFunc([0,0,0], (2,1), [1,0,0])
@test isFull(bf5) == false
@test isFull(1) == false

# function centerOf centerCoordOf
@test centerOf(bf5) == bf5.center
@test centerCoordOf(bf5) == [0,0,0]

# struct BasisFuncMix
bfm1 = BasisFuncMix(bf1)
@test bfm1 == BasisFuncMix([bf1])
@test BasisFuncMix(BasisFuncMix(bf3_1)) == BasisFuncMix(bf3_1)
bf5_2 = genBasisFunc([0,0,0], (2,1), [[1,0,0]])
bfm2 = BasisFuncMix(bf5)
@test hasEqual(bfm2, BasisFuncMix(bf5_2)[])

errorThreshold2 = 1e-15
bs1 = genBasisFunc.(gridCoords(GridBox(1,1.5)), Ref(GaussFunc(1.0, 0.5)))
nuc = ["H", "H"]
nucCoords = [rand(3), rand(3)]
bfm = BasisFuncMix(bs1)
S = overlaps([bfm])[]
@test S == overlap(bfm, bfm)[]
@test isapprox(S, overlaps(bs1) |> sum, atol=errorThreshold2)
T = elecKinetics([bfm])[]
@test T == elecKinetic(bfm, bfm)[]
@test isapprox(T, elecKinetics(bs1) |> sum, atol=errorThreshold2)
V = nucAttractions([bfm], nuc, nucCoords)[]
@test V == nucAttraction(bfm, bfm, nuc, nucCoords)[]
@test isapprox(V, nucAttractions(bs1, nuc, nucCoords) |> sum, atol=errorThreshold2)
eeI = eeInteractions([bfm])[]
@test eeI == eeInteraction(bfm, bfm, bfm, bfm)[]
@test isapprox(eeI, eeInteractions(bs1) |> sum, atol=errorThreshold2)

# function sumOf
bs2 = [genBasisFunc([1,1,1], (2,1), [1,0,0], normalizeGTO=true),
       genBasisFunc([1,1,1], (3,1), [2,0,0], normalizeGTO=true),
       genBasisFunc([1,1,2], (3,1), [0,0,0], normalizeGTO=true),
       genBasisFunc([1,1,1], (3,1), [0,1,0], normalizeGTO=true)]
bs2_2 = [genBasisFunc([1,1,1], (2,1), [1,0,0]),
         genBasisFunc([1,1,1], (3,1), [2,0,0]),
         genBasisFunc([1,1,2], (3,1), [0,0,0]),
         genBasisFunc([1,1,1], (3,1), [0,1,0])]
bs2_3 = [genBasisFunc([1,1,1], (2,1), [1,0,0]),
         genBasisFunc([1,1,1], (3,1), [2,0,0], normalizeGTO=true),
         genBasisFunc([1,1,2], (3,1), [0,0,0]),
         genBasisFunc([1,1,1], (3,1), [0,1,0], normalizeGTO=true)]
bs3 = [genBasisFunc([1,1,1], (2,1), [1,0,0], normalizeGTO=true),
       genBasisFunc([1,1,1], (3,1), [0,1,0], normalizeGTO=true),
       genBasisFunc([1,1,1], (3,1), [2,0,0], normalizeGTO=true),
       genBasisFunc([1,1,2], (3,1), [0,0,0], normalizeGTO=true)]
bs3_2 = [genBasisFunc([1,1,1], (2,1), [1,0,0]),
       genBasisFunc([1,1,1], (3,1), [0,1,0]),
       genBasisFunc([1,1,1], (3,1), [2,0,0]),
       genBasisFunc([1,1,2], (3,1), [0,0,0])]

bfm_1 = +(bs2...,)
bfm_2 = sumOf(bs2)
bfm_3 = BasisFuncMix(bs3)
bfm_4 = +(bs2_2...,)
bfm_5 = sumOf(bs2_2)
bfm_6 = BasisFuncMix(bs3_2)
bfm_7 = sumOf([bfm_6])
@test hasEqual(bfm_1, bfm_3)
@test hasEqual(bfm_2, bfm_3)
@test hasEqual(bfm_4, bfm_6)
@test hasEqual(bfm_5, bfm_6)
@test hasEqual(bfm_6, bfm_7)

# function mergeGaussFuncs
gf_merge1 = GaussFunc(2,1)
gf_merge2 = GaussFunc(2,1)
gf_merge3 = GaussFunc(2,1)

@test mergeGaussFuncs(gf_merge1) === gf_merge1

mgf1 = mergeGaussFuncs(gf_merge1, gf_merge1)[]
mgf2 = mergeGaussFuncs(gf_merge1, gf_merge2)[]
@test mgf1.xpn() == 2 == mgf1.con()
@test !hasIdentical(mgf1, mgf2)
gf_merge1_2 = GaussFunc(gf_merge1.xpn, gf_merge2.con)
gf_merge1_3 = GaussFunc(gf_merge2.xpn, gf_merge1.con)
mgf1_2 = mergeGaussFuncs(gf_merge1, gf_merge1_2)[]
mgf1_3 = mergeGaussFuncs(gf_merge1, gf_merge1_3)[]
@test mgf1_2.xpn() == 2 == mgf1_2.con()
@test !hasIdentical(mgf1, mgf1_2)
@test !hasIdentical(mgf1_2, mgf1_3)
@test hasEqual(mgf1, mgf1_2, mgf1_3)

gf_merge3 = GaussFunc(1.5,1)
@test hasIdentical(mergeGaussFuncs(gf_merge1, gf_merge3), [gf_merge1, gf_merge3])

# function add, mul, gaussProd
@test add(bs2[1]) === bs2[1]
bf1s = BasisFuncs(bf1.center, bf1.gauss, [ijkOrbitalList[bf1.ijk[1]]], bf1.normalizeGTO)
@test hasIdentical(add(bf1s), bf1)
bfm_add1 = BasisFuncMix(bs2)
@test add(bfm_add1) == sumOf(bs2)
bf_add1 = genBasisFunc([1,2,1], (2,1.1))
bf_add2 = genBasisFunc([1,1,1], (1,1.2))
bf_add1_2 = BasisFuncMix(bf_add1)
bf_add2_2 = BasisFuncMix(bf_add2)
@test hasEqual(add(bf_add1, bf_add2), add(bf_add2, bf_add1))
@test hasEqual(add(bf_add1_2, bf_add2_2), add(bf_add2_2, bf_add1_2))
@test hasEqual(add(bf_add1_2, bf_add2), add(bf_add1, bf_add2))
@test hasEqual(add(bf_add1, bf_add2_2), add(bf_add1, bf_add2))
@test hasEqual(add(bf_add1_2, bf_add2_2), add(bf_add1, bf_add2))

for bs in (bs2, bs2_2, bs2_3)
    X = overlaps(bs)^(-0.5)
    bsNew = [mul.(bs, @view X[:,i]) for i in 1:size(X, 2)] .|> sum
    SNew = overlaps(bsNew)
    @test isapprox(SNew, LinearAlgebra.I, atol=1e-14)
end
@test hasEqual(mul(1.0, bf1), bf1)
@test hasEqual(mul(1.1, bf1), mul(bf1, 1.1))

α₁, α₂ = rand(1:0.01:10, 2)
d₁, d₂ = rand(2)
R₁ = rand(-2:0.01:2, 3)
R₂ = rand(-2:0.01:2, 3)
xr = -10:0.1:10
yr = -10:0.1:10
zr = -10:0.1:10
gf_mul = mul(GaussFunc(α₁, d₁), GaussFunc(α₂, d₂))
@test isapprox(gf_mul.xpn[], α₁ + α₂, atol=1e-14)
@test isapprox(gf_mul.con[], d₁ * d₂, atol=1e-14)
bl = true
for x in xr, y in yr, z in zr
    xv = [x, y, z]
    α₃, d₃, R₃ = gaussProd((α₁, d₁, R₁), (α₂, d₂, R₂))
    bl *= isapprox(d₁*exp(-α₁*sum(abs2, xv-R₁)) * d₂*exp(-α₂*sum(abs2, xv-R₂)),
                   d₃*exp(-α₃*sum(abs2, xv-R₃)), atol=1e-16) # limit: atol=1e-20
end
@test bl

@test hasEqual(mul(bf1s, bf1s), [mul(bf1s[1], bf1s[1])])
@test hasEqual(mul(bf1s, bf1s, normalizeGTO=true)[],
               mul(bf1s[1], bf1s[1], normalizeGTO=true))
bf_mul1 = genBasisFunc([1,0,0], (2,3))
bf_mul2 = genBasisFunc([1,0,0], (1.5,1))
bf_mul2_2 = genBasisFunc([1,0,0], (1.5,1), normalizeGTO=true)
@test hasEqual(mul(bf_mul1, bf_mul2),  mul(BasisFuncMix([bf_mul1]), BasisFuncMix([bf_mul2])))
@test hasEqual(mul(bf_mul1, bf_mul2, normalizeGTO=true),
               mul(BasisFuncMix([bf_mul1]), BasisFuncMix([bf_mul2]), normalizeGTO=true))
bf_mul3 = genBasisFunc([1,0,0], ([1.0, 1.5], [0.3, 0.4]))
bf_mul4 = mul(bf_mul1, bf_mul2)
bf_mul4_0 = genBasisFunc([1,0,0], (3.5,3))
@test hasEqual(bf_mul4, bf_mul4_0)
bf_mul5 = mul(bf_mul1, bf_mul2, normalizeGTO=true)
bf_mul5_0 = genBasisFunc([1,0,0], (3.5,3), normalizeGTO=true)
@test hasEqual(bf_mul5, bf_mul5_0)
bf_mul6 = mul(bf_mul1, bf_mul2_2)
bf_mul6_0 = genBasisFunc([1,0,0], (3.5,3*normOfGTOin(bf_mul2)[]))
@test hasEqual(bf_mul6, bf_mul6_0)
bf_mul7 = mul(bf_mul1, bf_mul2_2, normalizeGTO=true)
bf_mul7_0 = genBasisFunc([1,0,0], (3.5,3*normOfGTOin(bf_mul2)[]), normalizeGTO=true)
@test hasEqual(bf_mul7, bf_mul7_0)
bf_mul8 = mul(bf_mul1, bf_mul3)
bf_mul8_0 = genBasisFunc([1,0,0], ([3, 3.5], [0.9, 1.2]))
@test hasEqual(bf_mul8, bf_mul8_0)
bf_mul9 = mul(bf_mul1, bf_mul3, normalizeGTO=true)
bf_mul9_0 = genBasisFunc([1,0,0], ([3, 3.5], [0.9, 1.2]), normalizeGTO=true)
@test hasEqual(bf_mul9, bf_mul9_0)
bf_mul10 = genBasisFunc([1,0,0], (1.2,0.3), [1,0,0])
bf_mul11 = genBasisFunc([0,1,0], (1.5,1))
bf_mul12 = genBasisFunc([0,1,0], (1.5,1), [1,0,0])
bf_mul13 = genBasisFunc([1,1,1], ([1.2, 1.0], [0.2, 0.5]))
bf_mul14 = genBasisFunc([1,0.2,1], ([1.2, 1.0], [0.2, 0.5]), [1,1,1])
bf_mul15 = mul(bf_mul10, BasisFuncMix([bf_mul1, bf_mul2]))
bf_mul15_0 = mul(bf_mul10, BasisFuncMix([bf_mul1, bf_mul2]), normalizeGTO=true)
bf_mul16 = mul(bf_mul15, bf_mul15_0)
bf_mul16_0 = mul(bf_mul15, bf_mul15_0, normalizeGTO=true)

testNorm = function (bf1, bf2)
    bf3 = mul(bf1, bf2)
    n1 = Quiqbox.getOverlap(bf3, bf3)
    n2 = Quiqbox.getOverlap(mul(bf1, bf1), mul(bf2, bf2))
    @test isapprox(n1, n2, atol=1e-12)
end

testNorm(bf_mul1, bf_mul10)
testNorm(bf_mul1, bf_mul11)
testNorm(bf_mul1, bf_mul12)
testNorm(bf_mul1, bf_mul13)
testNorm(bf_mul1, bf_mul14)
testNorm(bf_mul1, bf_mul15)
testNorm(bf_mul1, bf_mul15_0)
testNorm(bf_mul1, bf_mul16)
testNorm(bf_mul1, bf_mul16_0)
testNorm(bf_mul10, bf_mul11)
testNorm(bf_mul11, bf_mul13)
testNorm(bf_mul12, bf_mul13)
testNorm(bf_mul12, bf_mul14)
testNorm(bf_mul13, bf_mul14)
testNorm(bf_mul12, bf_mul15)
testNorm(bf_mul13, bf_mul15)
testNorm(bf_mul14, bf_mul15)
testNorm(bf_mul15, bf_mul15_0)
testNorm(bf_mul12, bf_mul16)
testNorm(bf_mul13, bf_mul16)
testNorm(bf_mul14, bf_mul16)
testNorm(bf_mul16, bf_mul16_0)

# (dev) function getOverlap
@test isapprox(Quiqbox.getOverlap(bf_mul2_2, bf_mul2_2), 1.0, atol=1e-10)
@test isapprox(Quiqbox.getOverlap(bf_mul5_0, bf_mul5_0), 9.0, atol=1e-10)
bf_mul18 = genBasisFunc(rand(3), "STO-3G")[]
@test isapprox(Quiqbox.getOverlap(bf_mul18, bf_mul18), 1.0, atol=1e-10)

# function shift
ijk = [1,0,0]
didjdk = [0,1,1]
bf_os1 = genBasisFunc([0,0,0], (2,1), ijk)
bf_os2 = genBasisFunc([0,0,0], (2,1), ijk, normalizeGTO=true)
bf_os1S = genBasisFunc([0,0,0], (2,1), ijk+didjdk)
bf_os2S = genBasisFunc([0,0,0], (2,1), ijk+didjdk, normalizeGTO=true)
@test hasEqual(shift(bf_os1, didjdk), bf_os1S)
@test hasEqual(shift(bf_os2, didjdk), bf_os2S)

# function unpackBasisFuncs
@test unpackBasisFuncs(bfm1)[1] == bf1
@test unpackBasisFuncs(bf1) == [bf1]
@test unpackBasisFuncs(0) == []

# function decompose
bf_d_1 = genBasisFunc([1,0,0], (1, 0.5))
bf_d_2 = genBasisFunc([1,0,0], ([2,1], [0.1, 0.5]))
bf_d_3 = genBasisFunc([1,1,0], ([2,1], [0.1, 0.2]), "P")
bm_d_1 = BasisFuncMix([bf_d_1, bf_d_2])
dm1 = reshape([bf_d_1], 1, 1)
@test hasEqual(decompose(bf_d_1), dm1)
@test hasEqual(decompose(bf_d_1, splitGaussFunc=true), dm1)
@test hasEqual(decompose(bf_d_2), reshape([bf_d_2], 1, 1))
dm2 = reshape([genBasisFunc([1,0,0], (2, 0.1)), bf_d_1], 2, 1)
@test hasEqual(decompose(bf_d_2, splitGaussFunc=true), dm2)
@test hasEqual(decompose(bf_d_3), reshape(bf_d_3[:], 1, 3))
@test hasEqual(decompose(bf_d_3, splitGaussFunc=true),
               hcat(decompose.(bf_d_3[:], splitGaussFunc=true)...))
@test decompose(bm_d_1) === bm_d_1
@test hasEqual(decompose(bm_d_1, splitGaussFunc=true),
              reshape(vcat(decompose.(bm_d_1.BasisFunc, splitGaussFunc=true)...), 3, 1))

# function basisSize
@test basisSize("P") == (3,)
@test basisSize(["S", "P", "D"]) == (1, 3, 6)
@test basisSize(bf1) == (1,)
@test basisSize(bfm1) == (1,) == basisSize(bfm2)
@test basisSize([bf1, bf2, bf3_3, bf5]) == (1,1,3,1)

# function genGaussFuncText
bfCoeff = [[6.163845031, 1.097161308], [0.4301284983, 0.6789135305],
           [0.245916322, 0.06237087296], [0.0494717692, 0.9637824081],
           [0.245916322, 0.06237087296], [0.5115407076, 0.6128198961]]
bfCoeff2 = vcat([[bfCoeff[2i+1]'; bfCoeff[2i+2]']' for i=0:2]...)
content = """
S    2   1.0
         6.163845031                   0.4301284983
         1.097161308                   0.6789135305
S    2   1.0
         0.245916322                   0.0494717692
         0.06237087296                 0.9637824081
P    2   1.0
         0.245916322                   0.5115407076
         0.06237087296                 0.6128198961
"""
lines = (content |> IOBuffer |> readlines)
@test map(i->Quiqbox.genGaussFuncText(bfCoeff2[i,:]...), 1:size(bfCoeff2)[1] |> collect) ==
      [lines[2], lines[3], lines[5], lines[6], lines[8], lines[9]].*"\n"

# function genBasisFuncText & genBFuncsFromText
randElement = ElementNames[rand(1:length(ElementNames))]
bs1 = genBasisFunc(missing, ("6-31G", "H"))
cens = [rand(3) for _=1:length(bs1)]
txt1 = genBasisFuncText(bs1, printCenter=false, groupCenters=false) |> join
txt2 = genBasisFuncText(bs1, printCenter=false) |> join
bs2_1 = genBFuncsFromText(txt1)
bs2_2 = genBFuncsFromText(txt2)
assignCenter!.(cens, bs1)
assignCenter!.(cens, bs2_1)
assignCenter!.(cens, bs2_2)
txt3 = genBasisFuncText(bs1) |> join
bs2_3 = genBFuncsFromText(txt3)
@test hasEqual(bs1, bs2_1, ignoreContainer=true)
@test hasEqual(bs1, bs2_2, ignoreContainer=true)
@test hasEqual.(sortBasisFuncs(bs1), bs2_3, ignoreFunction=true) |> prod
@test hasEqual.(bs1, bs2_1, ignoreFunction=true) |> prod
@test hasEqual.(bs1, bs2_2, ignoreFunction=true) |> prod
@test hasEqual.(sortBasisFuncs(bs1), bs2_3, ignoreFunction=true) |> prod

# function assignCenter!
bf6 = genBasisFunc(missing, "STO-3G")[]
coord = [1,0,0]
bf6_1 = genBasisFunc(coord, "STO-3G")[]
@test !hasEqual(bf6, bf6_1)
assignCenter!(coord, bf6)
@test hasEqual(bf6, bf6_1)

# function getParams
pb1 = ParamBox(2, :p, canDiff=false)
@test getParams(pb1) == pb1
@test getParams(pb1, :p) == pb1
@test getParams(pb1, :P) === nothing
@test getParams(pb1, onlyDifferentiable=true) === nothing

pb2 = ParamBox(2, :q)
@test getParams([pb1, pb2]) == [pb1, pb2]
@test getParams([pb1, pb2], :p) == [pb1]

gf_pbTest1 = GaussFunc(2,1)
@test getParams(gf_pbTest1) == gf_pbTest1.param |> collect

gf_pbTest2 = GaussFunc(1.5,0.5)
bf_pbTest1 = genBasisFunc([1,0,0], [gf_pbTest1, gf_pbTest2])
@test getParams(bf_pbTest1) == [bf_pbTest1.center...,
                                gf_pbTest1.param..., gf_pbTest2.param...]

alpha = ParamList[:xpn]
@test getParams(bf_pbTest1, alpha) ==
      vcat(getParams(gf_pbTest1, alpha), getParams(gf_pbTest2, alpha)) ==
      [gf_pbTest1.param[1], gf_pbTest2.param[1]]

cs = [pb1, bf_pbTest1]
ss = [:X, :Y, :Z]
@test [getParams(i, j) for i in cs, j in ss] |> flatten ==
      [nothing, bf_pbTest1.center[1], nothing, bf_pbTest1.center[2],
       nothing, bf_pbTest1.center[3]]
@test getParams(cs) == vcat(getParams(pb1), getParams(bf_pbTest1))
@test (getParams.(Ref(cs), ss) |> flatten) == (bf_pbTest1.center |> collect)

# function copyBasis
e = genExponent(3.0, x -> x^2 + 1)
c = genContraction(2.0)
gf_cb1 = GaussFunc(e, c)
gf_cb2 = copyBasis(gf_cb1)
@test !hasEqual(gf_cb1, gf_cb2)
@test hasEqual(gf_cb1.con, gf_cb2.con)
@test gf_cb1.xpn() == gf_cb2.xpn() == gf_cb2.xpn[] != gf_cb1.xpn[]
gf_cb3 = copyBasis(gf_cb1, false)
@test !hasEqual(gf_cb1, gf_cb3)
@test gf_cb1.xpn() != gf_cb3.xpn() == gf_cb3.xpn[] == gf_cb1.xpn[]
cen_cb1 = rand(3)

bf_cb1 = genBasisFunc(cen_cb1, [gf_cb1, gf_cb3])
bf_cb2 = copyBasis(bf_cb1)
bf_cb3 = copyBasis(bf_cb1, false)
testbf_cb = function (bf1, bf2)
    @test bf1.center == bf2.center
    @test bf1.ijk == bf2.ijk
    @test bf1.normalizeGTO == bf2.normalizeGTO
end
testbf_cb(bf_cb1, bf_cb2)
testbf_cb(bf_cb1, bf_cb3)
@test hasEqual(bf_cb2.gauss, (gf_cb2, gf_cb3))
@test hasEqual(bf_cb3.gauss, (gf_cb3, gf_cb3))

bfm_cb1 = BasisFuncMix([bf_cb1, bf_cb3])
bfm_cb2 = copyBasis(bfm_cb1)
bfm_cb3 = copyBasis(bfm_cb1, false)
testbf_cb.(bfm_cb1.BasisFunc, bfm_cb2.BasisFunc)
testbf_cb.(bfm_cb1.BasisFunc, bfm_cb3.BasisFunc)
@test hasEqual(bfm_cb2.BasisFunc |> collect, [bf_cb2, bf_cb3])
@test hasEqual(bfm_cb3.BasisFunc |> collect, [bf_cb3, bf_cb3])

# function uniqueParams! getVar getVarDict
e_gv1 = genExponent(2.0)
c_gv1 = genContraction(1.0)
gf_gv1 = GaussFunc(e_gv1, c_gv1)
e_gv2 = genExponent(2.5)
c_gv2 = genContraction(0.5)
gf_gv2 = GaussFunc(e_gv2, c_gv2)
e_gv3 = genExponent(1.05, x->x^2)
c_gv3 = genContraction(1.5)
gf_gv3 = GaussFunc(e_gv3, c_gv3)
gf_gv12 = GaussFunc(e_gv1, c_gv2)
gf_gv23 = GaussFunc(e_gv2, c_gv3)
gf_gv31 = GaussFunc(e_gv3, c_gv1)
x_gv1 = fill(1.0)
y_gv1 = fill(5.0)
cen_gv1 = genSpatialPoint([x_gv1, 2, 3])
cen_gv2 = genSpatialPoint([4, y_gv1, 6])
cen_gv3 = genSpatialPoint([x_gv1, y_gv1, 6])

bf_gv1 = genBasisFunc(cen_gv1, gf_gv1)
bf_gv2 = genBasisFunc(cen_gv2, gf_gv2)
bf_gv3 = genBasisFunc(cen_gv2, [gf_gv1, gf_gv2])
bf_gv4 = genBasisFunc(cen_gv1, gf_gv3)
bf_gv5 = genBasisFunc(cen_gv3, gf_gv1)
bf_gv6 = genBasisFunc(cen_gv3, [gf_gv12, gf_gv23, gf_gv31])
bfm_gv = BasisFuncMix([bf_gv2, bf_gv4, bf_gv6])

@test uniqueParams!(bf_gv4) == [cen_gv1..., e_gv3, c_gv3]

pbs_gv1 = uniqueParams!([bf_gv1, bf_gv2])
pbs_gv1_0 = [cen_gv1[1], cen_gv2[1],
             cen_gv1[2], cen_gv2[2],
             cen_gv1[3], cen_gv2[3],
                  e_gv1,      e_gv2,
                  c_gv1,      c_gv2]
@test pbs_gv1 == pbs_gv1_0
@test hasIdentical(pbs_gv1, pbs_gv1_0)

pbs_gv2 = uniqueParams!([bf_gv2, bf_gv1, bf_gv3])

pbs_gv2_0 = [cen_gv2[1], cen_gv1[1], cen_gv2[1],
             cen_gv2[2], cen_gv1[2], cen_gv2[2],
             cen_gv2[3], cen_gv1[3], cen_gv2[3],
             e_gv2, e_gv1, e_gv1, e_gv2,
             c_gv2, c_gv1, c_gv1, c_gv2]
@test pbs_gv2 == pbs_gv2_0
@test hasIdentical(pbs_gv2, pbs_gv2_0)

bfs_gv = sortBasisFuncs([bf_gv1, bf_gv2, bf_gv3, bf_gv4, bf_gv5, bf_gv6])
pbs_gv3 = uniqueParams!(bfs_gv, filterMapping=true)
pbs_gv4 = uniqueParams!(vcat(bfs_gv, bfm_gv), filterMapping=true)
@test pbs_gv3 == pbs_gv4
@test hasIdentical(pbs_gv3, pbs_gv4)

pbs_gv0 = [cen_gv1[1], cen_gv2[1],
           cen_gv1[2], cen_gv2[2],
           cen_gv1[3], cen_gv2[3], cen_gv3[3],
           e_gv1, e_gv2, e_gv3,
           c_gv1, c_gv2, c_gv3]
pbs_gv0_2 = sort(pbs_gv0, by=x->(typeof(x).parameters[2], x.index[]))
pbs_gv3_2 = sort(pbs_gv3, by=x->(typeof(x).parameters[2], x.index[]))
@test pbs_gv3_2 == pbs_gv0_2
@test hasEqual(pbs_gv3_2, pbs_gv0_2)
@test hasIdentical(pbs_gv3_2, pbs_gv0_2)

@test getVar(e_gv1).val.name == :α₁
@test getfield.(getfield.(getVar(bf_gv6), :val), :name) ==
      [:X₁, :Y₂, :Z₂, :α₁, :d₃, :α₃, :d₂, :x_α₂, :d₁]
@test getfield.(getfield.(getVar(bfm_gv), :val), :name) ==
      [:X₂, :Y₂, :Z₃, :α₃, :d₃, :X₁, :Y₁, :Z₁, :x_α₂,
       :d₂, :X₁, :Y₂, :Z₂, :α₁, :d₃, :α₃, :d₂, :x_α₂, :d₁]

@test getVarDict(e_gv1) == Dict(Symbolics.variable(:α, 1)=>2.0)

pairs_gv1 = sort(getVarDict(e_gv3, includeMapping=true) |> collect, by=x->string(x[1]))
strs_gv1 = getindex.(pairs_gv1, 1) .|> string
strs_gv2 = getindex.(pairs_gv1, 2) .|> string

@test strs_gv1[1][1:3] == "f_α" && strs_gv1[1][end-8:end] == "(x_α₂)"
@test strs_gv1[2:end] == ["x_α₂", "x_α₂^2"]
@test strs_gv2 == ["x_α₂^2", "1.05", "1.1025"]

@test getVarDict(gf_gv1) == Dict( Pair.(Symbolics.variable.([:α, :d], fill(1, 2)),
                                        [2.0, 1.0]) )

@test getVarDict(bf_gv1) == Dict( Pair.(Symbolics.variable.([:X, :Y, :Z, :α, :d],
                                                            fill(1, 5)),
                                        [1.0, 2.0, 3.0, 2.0, 1.0]) )

@test getVarDict(bf_gv1) == getVarDict([bf_gv1, gf_gv1])
@test getVarDict([bf_gv1, bf_gv2]) ==
      merge(getVarDict(bf_gv1), getVarDict(bf_gv2)) ==
      getVarDict([cen_gv1..., cen_gv2..., gf_gv1, gf_gv2]) ==
      getVarDict(vcat(bf_gv1.param|>collect, bf_gv2.param|>collect))
@test getVarDict(bfm_gv) ==
      getVarDict([bf_gv2, bf_gv4, bf_gv6]) ==
      getVarDict(pbs_gv0)

# function expressionOf, expressionOfCore
expr1 = expressionOf(gf1) |> string
@test expr1 == "exp(-2.0(r₁^2) - 2.0(r₂^2) - 2.0(r₃^2))" ||
      expr1 == "exp(-2.0(r₁^2) - (2.0(r₂^2)) - (2.0(r₃^2)))"

expr2 = expressionOf(bf1)[]|>string
@test expr2 == "exp(-2.0((r₁ - 1.0)^2) - 2.0((r₂ - 2.0)^2) - 2.0((r₃ - 3.0)^2))" ||
      expr2 == "exp(-2.0((r₁ - 1.0)^2) - (2.0((r₂ - 2.0)^2)) - (2.0((r₃ - 3.0)^2)))"
@test expressionOfCore(gf1) |> string == "d*exp(-α*(r₁^2 + r₂^2 + r₃^2))"
@test expressionOfCore(bf1)[]|>string == "d*exp(-α*((r₁ - X)^2 + (r₂ - Y)^2 + (r₃ - Z)^2))"

expStr = expressionOfCore(bf6)[] |> string
idx = findfirst('d', expStr)
@test isapprox(parse(Float64, expStr[1:idx-1]), 7.579425332952777, atol=1e-12)
@test expStr[idx:end] == "d*exp(-α*((r₁ - X)^2 + (r₂ - Y)^2 + (r₃ - Z)^2))*(α^0.75)" ||
      expStr[idx:end] == "d*(α^0.75)*exp(-α*((r₁ - X)^2 + (r₂ - Y)^2 + (r₃ - Z)^2))"

@test expressionOfCore(bfm1)[]|>string == "d*exp(-α*((r₁ - X)^2 + (r₂ - Y)^2 + (r₃ - Z)^2))"
expr_bfm2 =  expressionOfCore(bfm2)[] |> string
@test expr_bfm2 == "d*exp(-α*((r₁ - X)^2 + (r₂ - Y)^2 + (r₃ - Z)^2))*(r₁ - X)" ||
      expr_bfm2 == "d*(r₁ - X)*exp(-α*((r₁ - X)^2 + (r₂ - Y)^2 + (r₃ - Z)^2))"

# function inSymbols
sym1 = :a1
sym2 = :d
syms = [:a, :b, :c]
@test inSymbols(sym1, syms) == :a
@test !inSymbols(sym2, syms)

sym3 = Symbolics.variable(:a)
@test inSymbols(sym3, syms) == :a
@test inSymbols(sym3.val, syms) == :a
@test inSymbols(log(sym3).val, syms) == false

@test inSymbols(abs, syms) == false

# function varVal
vars = @variables X, Y, Z, F(X)
F = vars[end]
vars = vars[1:end-1]
f1(X,Y,Z) = log(X+0.5Y*Z)
expr = f1(X,Y,Z)
vals = [1, 3.334, -0.2]
d1 = Dict(vars .=> vals)
errT = 1e-10
@test varVal.(vars, Ref(d1)) == vals
@test isapprox(varVal(sum(vars), d1), sum(vals), atol=errT)
@test isapprox(varVal(prod(vars), d1), prod(vals), atol=errT)
@test isapprox(varVal(vars[1]^vars[2], d1), vals[1]^vals[2], atol=errT)
@test isapprox(varVal(expr, d1), f1(vals...), atol=errT)

f1 = Symbolics.variable(abs, T=Symbolics.FnType{Tuple{Any}, Real})(Z)
f2 = Symbolics.variable(:abs, T=Symbolics.FnType{Tuple{Any}, Real})(Z)
f1sym = f1.val
f2sym = f2.val
@test varVal(f1sym, d1) == abs(d1[Z])
@test varVal(f2sym, d1) == abs(d1[Z])

sp = 1.5
gb1 = GridBox(1,sp)
d2 = getVarDict(gb1.box |> flatten)
l = Symbolics.variable(:L,0)
vars2 = [i.val for i in keys(d2)]
diffs2 = Differential(l).(vars2)
exprs2 = map(x -> (x isa SymbolicUtils.Term) ? d2[x] : x, vars2)
diffExprs2 = Symbolics.derivative.(exprs2, Ref(l))
excs2 = Symbolics.build_function.(exprs2, l) .|> eval
vals2 = [f(sp) for f in excs2]
excdiffs2 = Symbolics.build_function.(diffExprs2, l) .|> eval
diffvals2 = [f(sp) for f in excdiffs2]

@test varVal.(vars2, Ref(d2)) == vals2
@test varVal.(diffs2, Ref(d2)) == diffvals2

end


####################################
#
# options
#
####################################

struct CDOptions
  maxIter::Int64
  optTol::Float64
  randomize::Bool           # coordinate are visitied in a randomized order or not
  warmStart::Bool           # when running CD, should we do path following or not
  numSteps::Int64           # when pathFollowing, how many points are there on the path
end

CDOptions(;
  maxIter::Int64=2000,
  optTol::Float64=1e-7,
  randomize::Bool=true,
  warmStart::Bool=true,
  numSteps::Int=50) = CDOptions(maxIter, optTol, randomize, warmStart, numSteps)

# this is used for ScaledLasso and FeasibleLasso
struct IterLassoOptions
  maxIter::Int64
  optTol::Float64
  initProcedure::Symbol    # :Screening, :InitStd, :WarmStart
  sinit::Int64             # how many columns of X to use to estimate initial variance or obtain initial residuals
  σinit::Float64
  optionsCD::CDOptions
end

IterLassoOptions(;
  maxIter::Int64=20,
  optTol::Float64=1e-2,
  initProcedure::Symbol=:Screening,
  sinit::Int64=5,
  σinit::Float64=1.,
  optionsCD::CDOptions=CDOptions()) = IterLassoOptions(maxIter, optTol, initProcedure, sinit, σinit, optionsCD)

####################################
#
# helper functions
#
####################################

"""
Helper function that finds an initial estimate for σ that is needed
for Lasso and ScaledLasso procedures.

The procedure works as follows:

* s input variables that are most correlated with the response y are found
* y is regressed on those s features
* σ is estimated based on the residuals, which gives an upper bound on the true sigma
"""
_findInitSigma!(
  X::AbstractMatrix{T},
  y::AbstractVector{T},
  s::Int,
  storage::Vector{T}) where {T <: AbstractFloat} =  std(_findInitResiduals!(X, y, s, storage))

function _findInitResiduals!(
  X::AbstractMatrix{T},
  y::AbstractVector{T},
  s::Int,
  storage::Vector{T}) where {T <: AbstractFloat}

  S = _findLargestCorrelations(X, y, s)
  Xs = view(X, :, S)
  mul!(storage, Xs, Xs \ y)
  @. storage = y - storage
  return storage
end

function _findInitResiduals!(
  w::AbstractVector{T},
  X::AbstractMatrix{T},
  y::AbstractVector{T},
  s::Int,
  storage::Vector{T}) where {T <: AbstractFloat}

  S = _findLargestCorrelations(w, X, y, s)

  Xs = view(X, :, S)
  mul!(storage, Xs, (Xs' * Diagonal(w) * Xs) \ (Xs' * Diagonal(w) * y))
  @. storage = y - storage
  return storage
end

# return a bit array containing indices of columns
function _findLargestCorrelations(
  X::AbstractMatrix{T},
  y::AbstractVector{T},
  s::Int) where {T <: AbstractFloat}

  p = size(X, 2)
  storage = Array{T}(undef, p)
  mul!(storage, transpose(X), y)
  @. storage = abs(storage)
  S = storage .>= nlargest(s, storage)[end]
end

function _findLargestCorrelations(
  w::AbstractVector{T},
  X::AbstractMatrix{T},
  y::AbstractVector{T},
  s::Int) where {T <: AbstractFloat}

  n, p = size(X)
  storage = Array{T}(undef, p)
  @inbounds for j=1:p
    val = zero(T)
    @simd for i=1:n
      val += X[i,j] * w[i] * y[i]
    end
    storage[j] = abs(val)
  end
  S = storage .>= nlargest(s, storage)[end]
end

function _stdX!(out::Vector{T}, X::AbstractMatrix{T}) where {T <: AbstractFloat}
  n, p = size(X)

  @inbounds for j=1:p
    v = zero(T)
    @simd for i=1:n
      v += X[i, j]^2.
    end
    out[j] = sqrt(v / n)
  end
  out
end

function _stdX!(out::Vector{T}, w::AbstractVector{T}, X::AbstractMatrix{T}) where {T <: AbstractFloat}
  n, p = size(X)

  @inbounds for j=1:p
    v = zero(T)
    @simd for i=1:n
      v += w[i] * X[i, j]^2.
    end
    out[j] = sqrt(v / n)
  end
  out
end

function _getLoadings!(out::Vector{T}, X::AbstractMatrix{T}, e::AbstractVector{T}) where {T <: AbstractFloat}
  n, p = size(X)

  @inbounds for j=1:p
    v = zero(T)
    @simd for i=1:n
      v += (X[i, j]*e[i])^2.
    end
    out[j] = sqrt(v / n)
  end
  out
end

function _getSigma(w::AbstractVector{T}, r::AbstractVector{T}) where {T <: AbstractFloat}
    n = length(w)
    σ = zero(T)
    for ii = 1:n
        σ += r[ii]^2 * w[ii]
    end
    σ /= sum(w)
    sqrt(σ)
end