# Source term

function init_source(param::Param)
   (;GridSize, FullSize, nVar) = param

   source_GV = Array{Float64,4}(undef, GridSize..., nVar)
   div_G = Array{Float64,3}(undef, GridSize...)

   U = Array{Float64,4}(undef, FullSize..., 3)

   return source_GV, U, div_G
end

function calc_source!(param::Param, state_GV, source_GV, U, div_G)
   (;nI, nJ, nK, nG, iMin, iMax, jMin, jMax, kMin, kMax) = param

   # Compute ∇⋅B using central difference
   @views divergence!(param, state_GV[:,:,:,B_], div_G)

   source_GV[:,:,:,Rho_] .= 0.0

   @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI
      source_GV[i,j,k,Mx_] = -state_GV[i+nG,j+nG,k+nG,Bx_]*div_G[i,j,k]
      source_GV[i,j,k,My_] = -state_GV[i+nG,j+nG,k+nG,By_]*div_G[i,j,k]
      source_GV[i,j,k,Mz_] = -state_GV[i+nG,j+nG,k+nG,Bz_]*div_G[i,j,k]
      source_GV[i,j,k,Bx_] =  state_GV[i+nG,j+nG,k+nG,Ux_]*div_G[i,j,k]
      source_GV[i,j,k,By_] =  state_GV[i+nG,j+nG,k+nG,Uy_]*div_G[i,j,k]
      source_GV[i,j,k,Bz_] =  state_GV[i+nG,j+nG,k+nG,Uz_]*div_G[i,j,k]
   end

   # Calculate divergence of U
   for k = 1:size(U,3), j = 1:size(U,2), i = 1:size(U,1)
      U[i,j,k,1] = state_GV[i,j,k,Mx_] / state_GV[i,j,k,Rho_]
      U[i,j,k,2] = state_GV[i,j,k,My_] / state_GV[i,j,k,Rho_]
      U[i,j,k,3] = state_GV[i,j,k,Mz_] / state_GV[i,j,k,Rho_]
   end
   divergence!(param, U, div_G)

   if !param.UseConservative
      @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI
         source_GV[i,j,k,P_] = -(γ-1.0)*state_GV[i+nG,j+nG,k+nG,P_]*div_G[i,j,k]
      end
   else
      @inbounds for k = kMin:kMax, j = jMin:jMax, i = iMin:iMax
         ub = state_GV[i,j,k,Ux_]*state_GV[i,j,k,Bx_] +
            state_GV[i,j,k,Uy_]*state_GV[i,j,k,By_] +
            state_GV[i,j,k,Uz_]*state_GV[i,j,k,Bz_]
         source_GV[i-nG,j-nG,k-nG,E_] = -ub*div_G[i-nG,j-nG,k-nG]
      end
   end

   return source_GV
end


using QuantEcon
using NLsolve
using NLopt

import QuantEcon.simulate

mutable struct Model{TF <: AbstractFloat,
                     TM <: AbstractMatrix{TF},
                     TV <: AbstractVector{TF}}
    β::TF
    Π::TM
    G::TV
    Θ::TV
    transfers::Bool
    U::Function
    Uc::Function
    Ucc::Function
    Un::Function
    Unn::Function
    n_less_than_one::Bool
end

"""
{{ class_word }} returns planner's allocation as a function of the multiplier
on the implementability constraint μ
"""
struct SequentialAllocation{TP <: Model,
                            TI <: Integer,
                            TV <: AbstractVector}
    model::TP
    mc::MarkovChain
    S::TI
    cFB::TV
    nFB::TV
    ΞFB::TV
    zFB::TV
end

"""
Initializes the {{ class_word }} from the calibration model
"""
function SequentialAllocation(model::Model)
    β, Π, G, Θ = model.β, model.Π, model.G, model.Θ
    mc = MarkovChain(Π)
    S = size(Π, 1)   # Number of states
    # Now find the first best allocation
    cFB, nFB, ΞFB, zFB = find_first_best(model, S, 1)

    return SequentialAllocation(model, mc, S, cFB, nFB, ΞFB, zFB)
end

"""
Find the first best allocation
"""
function find_first_best(model::Model, S::Integer, version::Integer)
    if version != 1 && version != 2
        throw(ArgumentError("version must be 1 or 2"))
    end
    β, Θ, Uc, Un, G, Π =
        model.β, model.Θ, model.Uc, model.Un, model.G, model.Π
    function res!(z, out)
        c = z[1:S]
        n = z[S+1:end]
        out[1:S] = Θ .* Uc(c, n) + Un(c, n)
        out[S+1:end] = Θ .* n - c - G
    end
    res = nlsolve(res!, 0.5 * ones(2 * S))

    if converged(res) == false
        error("Could not find first best")
    end

    if version == 1
        cFB = res.zero[1:S]
        nFB = res.zero[S+1:end]
        ΞFB = Uc(cFB, nFB)         # Multiplier on the resource constraint
        zFB = vcat(cFB, nFB, ΞFB)
        return cFB, nFB, ΞFB, zFB
    elseif version == 2
        cFB = res.zero[1:S]
        nFB = res.zero[S+1:end]
        IFB = Uc(cFB, nFB) .* cFB + Un(cFB, nFB) .* nFB
        xFB = \(eye(S) - β * Π, IFB)
        zFB = [vcat(cFB[s], xFB[s], xFB) for s in 1:S]
        return cFB, nFB, IFB, xFB, zFB
    end
end

"""
Computes optimal allocation for time ``t\geq 1`` for a given ``\mu``
"""
function time1_allocation(pas::SequentialAllocation, μ::Real)
    model, S = pas.model, pas.S
    Θ, β, Π, G, Uc, Ucc, Un, Unn =
        model.Θ, model.β, model.Π, model.G,
        model.Uc, model.Ucc, model.Un, model.Unn
    function FOC!(z::Vector, out)
        c = z[1:S]
        n = z[S+1:2S]
        Ξ = z[2S+1:end]
        out[1:S] = Uc(c, n) - μ * (Ucc(c, n) .* c + Uc(c, n)) - Ξ            # FOC c
        out[S+1:2S] = Un(c, n) - μ * (Unn(c, n) .* n + Un(c, n)) + Θ .* Ξ    # FOC n
        out[2S+1:end] = Θ .* n - c - G                                       # Resource constraint
        return out
    end
    # Find the root of the FOC
    res = nlsolve(FOC!, pas.zFB)
    if res.f_converged == false
        error("Could not find LS allocation.")
    end
    z = res.zero
    c, n, Ξ = z[1:S], z[S+1:2S], z[2S+1:end]
    # Now compute x
    I  = Uc(c, n) .* c +  Un(c, n) .* n
    x = \(eye(S) - β * model.Π, I)
    return c, n, x, Ξ
end

"""
Finds the optimal allocation given initial government debt `B_` and state `s_0`
"""
function time0_allocation(pas::SequentialAllocation,
                          B_::AbstractFloat, s_0::Integer)
    model = pas.model
    Π, Θ, G, β =
        model.Π, model.Θ, model.G, model.β
    Uc, Ucc, Un, Unn =
        model.Uc, model.Ucc, model.Un, model.Unn
    # First order conditions of planner's problem
    function FOC!(z, out)
        μ, c, n, Ξ = z[1], z[2], z[3], z[4]
        xprime = time1_allocation(pas, μ)[3]
        out .= vcat(
        Uc(c, n) .* (c - B_) + Un(c, n) .* n + β * dot(Π[s_0, :], xprime),
        Uc(c, n) - μ * (Ucc(c, n) .* (c - B_) + Uc(c, n)) - Ξ,
        Un(c, n) - μ * (Unn(c, n) .* n + Un(c, n)) + Θ[s_0] .* Ξ,
        (Θ .* n - c - G)[s_0]
        )
    end
    # Find root
    res = nlsolve(FOC!, [0.0, pas.cFB[s_0], pas.nFB[s_0], pas.ΞFB[s_0]])
    if res.f_converged == false
        error("Could not find time 0 LS allocation.")
    end
    return (res.zero...)
end

"""
Find the value associated with multiplier `μ`
"""
function time1_value(pas::SequentialAllocation, μ::Real)
    model = pas.model
    c, n, x, Ξ = time1_allocation(pas, μ)
    U_val = model.U.(c, n)
    V = \(eye(pas.S) - model.β*model.Π, U_val)
    return c, n, x, V
end

"""
Computes Τ given `c`, `n`
"""
function Τ(model::Model, c::Union{Real,Vector}, n::Union{Real,Vector})
    Uc, Un = model.Uc.(c, n), model.Un.(c, n)
    return 1+Un./(model.Θ .* Uc)
end

"""
Simulates planners policies for `T` periods
"""
function simulate(pas::SequentialAllocation,
                  B_::AbstractFloat, s_0::Integer, T::Integer,
                  sHist::Union{Vector, Void}=nothing)
    model = pas.model
    Π, β, Uc = model.Π, model.β, model.Uc
    if sHist == nothing
        sHist = QuantEcon.simulate(pas.mc, T, init=s_0)
    end
    cHist = zeros(T)
    nHist = zeros(T)
    Bhist = zeros(T)
    ΤHist = zeros(T)
    μHist = zeros(T)
    RHist = zeros(T-1)
    # time 0
    μ, cHist[1], nHist[1], _  = time0_allocation(pas, B_, s_0)
    ΤHist[1] = Τ(pas.model, cHist[1], nHist[1])[s_0]
    Bhist[1] = B_
    μHist[1] = μ
    # time 1 onward
    for t in 2:T
        c, n, x, Ξ = time1_allocation(pas,μ)
        u_c = Uc(c,n)
        s = sHist[t]
        ΤHist[t] = Τ(pas.model, c, n)[s]
        Eu_c = dot(Π[sHist[t-1],:], u_c)
        cHist[t], nHist[t], Bhist[t] = c[s], n[s], x[s] / u_c[s]
        RHist[t-1] = Uc(cHist[t-1], nHist[t-1]) / (β * Eu_c)
        μHist[t] = μ
    end
    return cHist, nHist, Bhist, ΤHist, sHist, μHist, RHist
end

"""
Bellman equation for the continuation of the Lucas-Stokey Problem
"""
mutable struct BellmanEquation{TP <: Model,
                               TI <: Integer,
                               TV <: AbstractVector,
                               TM <: AbstractMatrix{TV},
                               TVV <: AbstractVector{TV}}
    model::TP
    S::TI
    xbar::TV
    time_0::Bool
    z0::TM
    cFB::TV
    nFB::TV
    xFB::TV
    zFB::TVV
end

"""
Initializes the {{ class_word }} from the calibration `model`
"""
function BellmanEquation(model::Model, xgrid::AbstractVector, policies0::Vector)
    S = size(model.Π, 1) # Number of states
    xbar = [minimum(xgrid), maximum(xgrid)]
    time_0 = false
    cf, nf, xprimef = policies0
    z0 = [vcat(cf[s](x), nf[s](x), [xprimef[s, sprime](x) for sprime in 1:S])
                        for x in xgrid, s in 1:S]
    cFB, nFB, IFB, xFB, zFB = find_first_best(model, S, 2)
    return BellmanEquation(model, S, xbar, time_0, z0, cFB, nFB, xFB, zFB)
end

"""
Finds the optimal policies
"""
function get_policies_time1(T::BellmanEquation,
                        i_x::Integer, x::AbstractFloat,
                        s::Integer, Vf::AbstractArray)
    model, S = T.model, T.S
    β, Θ, G, Π = model.β, model.Θ, model.G, model.Π
    U, Uc, Un = model.U, model.Uc, model.Un

    function objf(z::Vector, grad)
        c, xprime = z[1], z[2:end]
        n=c+G[s]
        Vprime = [Vf[sprime](xprime[sprime]) for sprime in 1:S]
        return -(U(c, n) + β * dot(Π[s, :], Vprime))
    end
    function cons(z::Vector, grad)
        c, xprime = z[1], z[2:end]
        n=c+G[s]
        return x - Uc(c, n) * c - Un(c, n) * n - β * dot(Π[s, :], xprime)
    end
    lb = vcat(0, T.xbar[1] * ones(S))
    ub = vcat(1 - G[s], T.xbar[2] * ones(S))
    opt = Opt(:LN_COBYLA, length(T.z0[i_x, s])-1)
    min_objective!(opt, objf)
    equality_constraint!(opt, cons)
    lower_bounds!(opt, lb)
    upper_bounds!(opt, ub)
    maxeval!(opt, 300)
    maxtime!(opt, 10)
    init = vcat(T.z0[i_x, s][1], T.z0[i_x, s][3:end])
    for (i, val) in enumerate(init)
        if val > ub[i]
            init[i] = ub[i]
        elseif val < lb[i]
            init[i] = lb[i]
        end
    end
    (minf, minx, ret) = optimize(opt, init)
    T.z0[i_x, s] = vcat(minx[1], minx[1] + G[s], minx[2:end])
    return vcat(-minf, T.z0[i_x, s])
end
"""
Finds the optimal policies
"""
function get_policies_time0(T::BellmanEquation,
                        B_::AbstractFloat, s0::Integer, Vf::Array)
    model, S = T.model, T.S
    β, Θ, G, Π = model.β, model.Θ, model.G, model.Π
    U, Uc, Un = model.U, model.Uc, model.Un
    function objf(z, grad)
        c, xprime = z[1], z[2:end]
        n = c+G[s0]
        Vprime = [Vf[sprime](xprime[sprime]) for sprime in 1:S]
        return -(U(c, n) + β * dot(Π[s0, :], Vprime))
    end
    function cons(z::Vector, grad)
        c, xprime = z[1], z[2:end]
        n = c + G[s0]
        return -Uc(c, n) * (c - B_) - Un(c, n) * n - β * dot(Π[s0, :], xprime)
    end
    lb = vcat(0, T.xbar[1] * ones(S))
    ub = vcat(1-G[s0], T.xbar[2] * ones(S))
    opt = Opt(:LN_COBYLA, length(T.zFB[s0])-1)
    min_objective!(opt, objf)
    equality_constraint!(opt, cons)
    lower_bounds!(opt, lb)
    upper_bounds!(opt, ub)
    maxeval!(opt, 300)
    maxtime!(opt, 10)
    init = vcat(T.zFB[s0][1], T.zFB[s0][3:end])
    for (i, val) in enumerate(init)
        if val > ub[i]
            init[i] = ub[i]
        elseif val < lb[i]
            init[i] = lb[i]
        end
    end
    (minf, minx, ret) = optimize(opt, init)
    return vcat(-minf, vcat(minx[1], minx[1]+G[s0], minx[2:end]))
end


using OmegaCore, OmegaCore.Proposals
using OmegaTest
using Distributions
using Test

function test_logpdf_ucond()
  a = 1 ~ Normal(0, 1)
  b = 2 ~ Normal(2, 3)
  z(ω) = (a(ω), b(ω))
  ω, ℓ = propose(Random.GLOBAL_RNG, z)
  @test isinferred(propose, Random.GLOBAL_RNG, z)
  a_, b_ = z(ω)
  @test ℓ == logpdf(Normal(0, 1), a_) + logpdf(Normal(2, 3), b_)
end

# Test doesn't double count
function test_logpdf_ucond_nodouble()
  a = 1 ~ Normal(0, 1)
  b = 2 ~ Normal(2, 3)
  z(ω) = (a(ω), b(ω), a(ω))
  ω, ℓ = propose(Random.GLOBAL_RNG, z)
  @test isinferred(propose, Random.GLOBAL_RNG, z)
  a_, b_ = z(ω)
  @test ℓ == logpdf(Normal(0, 1), a_) + logpdf(Normal(2, 3), b_)
end

function test_logpdf_cond()
  μ = 1 ~ Normal(0, 1)
  x = 2 ~ Normalₚ(μ, 1)
  μ_ = 0.1234
  x_ = -0.54321
  μc = μ |ᶜ x == x_
  ω, ℓ = propose(Random.GLOBAL_RNG, μc)
end

@testset "Logpdf" begin
  test_logpdf_ucond()
  test_logpdf_ucond_nodouble()
end


#
# Date created: 2022-03-10
# Author: aradclif
#
#
############################################################################################

function sumsqbody(N::Int, D)
    body = Expr(:block)
    params = D.parameters
    a = Expr(:ref, :A, ntuple(d -> Symbol(:i_, d), N)...)
    b = Expr(:ref, :B, ntuple(d -> params[d] == Static.One ? 1 : Symbol(:i_, d), N)...)
    δ = Expr(:(=), :Δ, a)
    e = Expr(:(=), b, Expr(:call, :+, b, Expr(:call, :*, :Δ, :Δ)))
    push!(body.args, δ)
    push!(body.args, e)
    body
end

function sumsq_quote(N::Int, D)
    ls = loopgen(N)
    body = sumsqbody(N, D)
    push!(ls.args, body)
    return quote
        @turbo $ls
        return B
    end
end

@generated function sumsq!(B::AbstractArray{T, N}, A::AbstractArray{T, N}, dims::D) where {T, N, D}
    sumsq_quote(N, D)
end

function sqrt!(A::AbstractArray{T, N}) where {T, N}
    @turbo for i ∈ eachindex(A)
        A[i] = sqrt(A[i])
    end
    A
end

function lveuclidean(A::AbstractArray{T, N}, dims::NTuple{M, Int}) where {T, N, M}
    Dᴬ = size(A)
    Dᴮ′ = ntuple(d -> d ∈ dims ? StaticInt(1) : Dᴬ[d], N)
    B = zeros(Base.promote_op(/, T, Int), Dᴮ′)
    sumsq!(B, A, Dᴮ′)
    sqrt!(B)
    B
end


export gumbel_softmax

"""
    gumbel_softmax(x::Var, tau)

* x: logarithm of probability distribution, e.g., output of `logsoftmax`.
* tau: temprature
"""
function gumbel_softmax(x::Var, tau)
    tau > 0.0 || throw("tau: $(tau) must be positive.")
    tau = eltype(x.data)(tau)
    y = Var(nothing, gumbel_softmax, (x,tau))
    g = gumbel_softmax(x.data, tau)
    y.data = vec(argmax(g,1))
    y.grad = zeros(g)
    y.df! = () -> begin
        isvoid(x.grad) && return
        ∇gumbel_softmax!(g, y.grad, x.grad, tau)
    end
    y
end

gumbel{T}(u::T) = -log(-log(u+1e-20)+1e-20)

function gumbel_softmax{T}(x::Matrix{T}, tau::T)
    y = similar(x)
    u = rand(T, size(x))
    tau = T(1) / tau
    @inbounds for i = 1:length(x)
        y[i] = (x[i] + gumbel(u[i])) * tau
    end
    softmax(y)
end

function ∇gumbel_softmax!{T}(y::Vector{Int}, gy::Matrix{T}, work::Matrix{T}, gx::Matrix{T}, tau::T)
    for j = 1:size(y,2)
        sum = T(0)
        for i = 1:size(y,1)
            sum += gy[i,j] * y[i,j]
        end
        y[j]
        for i = 1:size(y,1)
            gx[i,j] += y[i,j] * (gy[i,j] - sum)
        end
    end
end


"""
    pairwise_sqd(x)

Pairwise squared distances between each colum vector in `x`.
"""
function pairwise_sqd(x)
    n = size(x, 2)
    xixj = transpose(x) * x
    xsq = sum(x .^ 2; dims=1)
    return transpose(xsq) .+ xsq - 2xixj
end

"""
    pairwise_sqd(x, y)

Pairwise squared distances between each colum vector in `x` and `y`.
"""
function pairwise_sqd(x, y)
    nx = size(x, 2)
    ny = size(y, 2)
    xiyj = transpose(x) * y
    xsq = sum(x .^ 2; dims=1)
    ysq = sum(y .^ 2; dims=1)
    return transpose(xsq) .+ ysq - 2xiyj
end

gaussian_gram_by_pairwise_sqd(pdot, σ) = exp.(-pdot ./ 2(σ ^ 2))
gaussian_gram(x, σ) = gaussian_gram_by_pairwise_sqd(pairwise_sqd(x), σ)
gaussian_gram(x, y, σ) = gaussian_gram_by_pairwise_sqd(pairwise_sqd(x, y), σ)

abstract type AbstractMMD end

function estimate_ratio(mmd::AbstractMMD, pdot_dede, pdot_denu, σ)
    Kdede = gaussian_gram_by_pairwise_sqd(pdot_dede, σ)
    Kdenu = gaussian_gram_by_pairwise_sqd(pdot_denu, σ)
    return _estimate_ratio(mmd, Kdede, Kdenu)
end

"""
    estimate_ratio(mmd::AbstractMMD, x_de, x_nu; σs=nothing)

Estimate `p_nu(x_de) / p_de(x_de)` using mutiple `σ` in `σs`.
If `σs` is not provided, the median of pairwise distances will be used.
"""
function estimate_ratio(mmd::AbstractMMD, x_de, x_nu; σs=nothing)
    pdot_dede = pairwise_sqd(x_de)
    pdot_denu = pairwise_sqd(x_de, x_nu)

    if isnothing(σs)
        σ = sqrt(median([pdot_dede..., pdot_denu...]))
        @info "Automatically choose σ using the median of pairwise distances: $σ."
        σs = [σ]
    end

    r_de = mapreduce(σ -> estimate_ratio(mmd, pdot_dede, pdot_denu, σ), +, σs)

    return inv(length(σs)) * r_de
end

struct MMDAnalytical{T<:AbstractFloat, S, M} <: AbstractMMD
    ϵ::T
    function MMDAnalytical(ϵ::T=1f-3; method::Symbol=:solve) where {T<:Number}
        @assert method in (:solve, :inv)
        S = iszero(ϵ) ? Val{:false} : Val{:true}
        if T == Int
            ϵ = float(ϵ)
        end
        return new{typeof(ϵ), S, Val{method}}(ϵ)
    end
end

function adddiag(mmd::MMDAnalytical{T, Val{:true}, M}, Kdede) where {T, M}
    return Kdede + diagm(0 => mmd.ϵ * fill(one(T), size(Kdede, 1)))
end
adddiag(mmd::MMDAnalytical{T, Val{:false}, M}, Kdede) where {T, M} = Kdede

function _estimate_ratio(mmd::MMDAnalytical{T, S, Val{:inv}}, Kdede, Kdenu) where {T, S}
    n_de, n_nu = size(Kdenu)
    return convert(T, n_de / n_nu) * inv(adddiag(mmd, Kdede)) * Kdenu * fill(one(T), n_nu)
end

function _estimate_ratio(mmd::MMDAnalytical{T, S, Val{:solve}}, Kdede, Kdenu) where {T, S}
    n_de, n_nu = size(Kdenu)
    return convert(T, n_de / n_nu) * (adddiag(mmd, Kdede) \ sum(Kdenu; dims=2)[:,1])
end


using AsyPlots
using ConformalMaps

struct ConformalMapVisualization
    domain::AsyPlots.Plot2D
    range::AsyPlots.Plot2D
end

show(io::IO,C::ConformalMapVisualization) = print(io,
                                        "ConformalMapVisualization()")

#-----------------------------------------------------------------------------
# FUNCTIONS FOR DISPLAY VISUALIZING CONFORMAL MAPS
#-----------------------------------------------------------------------------

function closepath(γ::Array{<:Complex,1})
    return closepath(map(AsyPlots.Vec2,γ))
end

function closepath(γ::Array{<:AsyPlots.Vec2})
    if γ[1] == γ[end]
        return γ
    else
        return [γ;[γ[1]]]
    end
end

function makegrid(zvals;ringcolor="blue",raycolor="red",args...)
    lines = Path2D[]
    ringpen = Pen(color=ringcolor)
    for i=1:size(zvals)[1]
        push!(lines,Path2D(hcat(real(zvals[i,:])',
                                imag(zvals[i,:])');pen=ringpen,args...))
    end
    raypen = Pen(color=raycolor)
    for j=1:size(zvals)[2]
        push!(lines,Path2D(hcat(real(zvals[:,j]),
                                imag(zvals[:,j]));pen=raypen,args...))
    end
    return lines
end

function hyperbolictiling(f::Union{Function,InverseConformalMap};
                          rings::Integer=9,
                          rays::Integer=16,
                          rotation::Real=0.0,
                          innerradius::Real=1.0/3.0,
                          lwfunction=(i->2^(-i/2)),
                          ringcolor="blue",
                          raycolor="red",
                          spline=true)

   points = [[f((1-(1-innerradius)/2^(k-1))*cos(θ+rotation) +
                  im*(1-(1-innerradius)/2^(k-1))*sin(θ+rotation))
                       for θ=range(0, stop=2π, length=1+rays*2^(k-1))] for k=1:rings]
   return Plot([
            (isa(f,InverseConformalMap) ? domain(f) :
             [Point2D(f(0.0),linewidth=2),
              Polygon2D([f(cis(θ)) for θ=range(0, stop=2π, length=250)];linewidth=0.3)]);
             [Path2D([points[i][k],points[i+1][2*k-1]];
                     color=raycolor,linewidth=lwfunction(i))
                          for i=1:length(points)-1 for k=1:length(points[i])-1]
             [Polygon2D(points[i][1:end-1];
                    color=ringcolor,linewidth=lwfunction(i),spline=spline)
                    for i=1:length(points)]])
end

function domain(f::Union{ConformalMap,InverseConformalMap})
    return [Point2D(f.center,linewidth=2),Path2D(closepath(f.domain),linewidth=0.3)]
end

"""
    intersectQ(p,q,r,s)

Determine whether the segment [p,q] intersects
the segment [r,s].

The four points can be 2-tuples of `Real`s, or
`Complex`es.

    intersectQ(p,q,γ)

Determine whether the segment from p to q intersects
the polygonal path γ
"""
function intersectQ(p::Tuple{<:Real,<:Real},
                    q::Tuple{<:Real,<:Real},
                    r::Tuple{<:Real,<:Real},
                    s::Tuple{<:Real,<:Real})
    a = -((p[2]*(r[1] - s[1]) + r[2]*s[1] - r[1]*s[2] +
          p[1]*(-r[2] + s[2]))/((-p[2] + q[2])*(r[1] - s[1]) + (p[1] -
                 q[1])*(r[2] - s[2])))
    b = (-q[2]*r[1] + p[2]*(-q[1] + r[1]) + p[1]*(q[2] - r[2]) +
         q[1]*r[2])/((p[2] - q[2])*(r[1] - s[1]) - (p[1] - q[1])*(r[2] - s[2]))
    return 0 <= a <= 1 && 0 <= b <= 1
end

function intersectQ(p::Complex,q::Complex,r::Complex,s::Complex)
    # Determines whether the segment from p to q intersects the
    # segment from r to s
    p,q,r,s = map(reim,(p,q,r,s))
    return intersectQ(p,q,r,s)
end

function intersectQ(p::Complex,q::Complex,gamma::Array{<:Complex,1})
    for i=1:length(gamma)-1
        if intersectQ(p,q,gamma[i],gamma[i+1])
            return true
        end
    end
    return false
end

function intersectQ(p::AsyPlots.Vec2,
                    q::AsyPlots.Vec2,
                    γ::Array{<:AsyPlots.Vec2,1})
    intersectQ(map(complex,(p,q))...,map(complex,γ))
end

function makegrid(boundary::Array{<:AsyPlots.Vec2,1},n::Integer)
    grid  = Tuple[]
    xvals = [P.x for P in boundary]
    yvals = [P.y for P in boundary]
    ϵ = max(maximum(xvals)-minimum(xvals),maximum(yvals)-minimum(yvals))/(n-1)
    m = length(minimum(xvals):ϵ:maximum(xvals))
    n = length(minimum(yvals):ϵ:maximum(yvals))
    totalgrid = [(x,y) for
        x = range(minimum(xvals), stop=maximum(xvals), length=m),
        y = range(minimum(yvals), stop=maximum(yvals), length=n)]
    pointsinside = [iswellinside(AsyPlots.Vec2(totalgrid[i,j]),
                                boundary;epsilon=1e-3) ? 1 : 0 for i=1:m,j=1:n]
    lines = Tuple{Tuple{Int64,Int64},Tuple{Int64,Int64}}[]
    for i=1:m
        for j=1:n
            for v in [(i+1,j),(i-1,j),(i,j+1),(i,j-1)]
                if 1 <= v[1] <= m && 1 <= v[2] <= n &&
                    pointsinside[i,j] == 1 && pointsinside[v...] == 1 &&
                       ~intersectQ(AsyPlots.Vec2(totalgrid[i,j]),
                                   AsyPlots.Vec2(totalgrid[v...]),
                                   boundary)
                    push!(lines,((i,j),v))
                end
            end
        end
    end
    return totalgrid, pointsinside, lines
end

makegrid(boundary::Array{T,1},n::Integer) where {T<:Complex} =
    makegrid(hcat(reim(boundary)...),n)

function grid(domain,totalgrid,pointsinside,lines,center)
    return AsyPlots.Plot([[AsyPlots.Point(real(center),imag(center)),
    AsyPlots.Path(closepath(domain))];
    [AsyPlots.Path([totalgrid[line[1]...],
                    totalgrid[line[2]...]],linewidth=0.2) for line in lines]])
end

function gridimage(f,totalgrid,pointsinside,lines,center)
    return Plot([[AsyPlots.Point(0.0,0.0),AsyPlots.Circle((0.0,0.0),1,linewidth=0.3)];
     [AsyPlots.Path([f(totalgrid[line[1]...]),
                     f(totalgrid[line[2]...])],linewidth=0.1) for line in lines]])
end

function visualize(CM::ConformalMap,n::Integer=25)
    totalgrid,pointsinside,lines = makegrid(closepath(CM.domain),n)
    return ConformalMapVisualization(
        grid(closepath(CM.domain),totalgrid,pointsinside,lines,CM.center),
        gridimage(p->CM(p[1] + im*p[2]),totalgrid,pointsinside,lines,CM.center)
    )
end

function visualize(ICM::InverseConformalMap;kwargs...)
    return ConformalMapVisualization(
        hyperbolictiling(z->z;kwargs...),hyperbolictiling(ICM;kwargs...)
    )
end

function combine(V::ConformalMapVisualization;kwargs...)
    domainbb = AsyPlots.boundingbox(V.domain)
    rangebb = AsyPlots.boundingbox(V.range)
    totalwidth = domainbb.xmax - domainbb.xmin + rangebb.xmax - rangebb.xmin
    xshift = domainbb.xmax - rangebb.xmin + 0.2*totalwidth
    domainycenter = mean([domainbb.ymax,domainbb.ymin])
    rangeycenter = mean([rangebb.ymax,rangebb.ymin])
    yshift = domainycenter - rangeycenter
    fullpicture = V.domain + Shift(xshift,yshift)*V.range
    bb = AsyPlots.boundingbox(fullpicture)
    midheight = mean([bb.ymax,bb.ymin])
    arrow = Path([domainbb.xmax+0.04*totalwidth midheight;
                  domainbb.xmax+0.16*totalwidth midheight];arrow=AsyPlots.Arrow(4))
    fullpicture += Plot(arrow)
    Plot(fullpicture.elements;kwargs...)
end

function Base.show(io::IO,mime::MIME"text/plain",V::ConformalMapVisualization)
    Base.show(io,mime,combine(V))
end

function Base.show(io::IO, mime::MIME"image/svg+xml", V::ConformalMapVisualization)
    Base.show(io,mime,combine(V))
end

function Base.show(io::IO, mime::MIME"image/png", V::ConformalMapVisualization)
    Base.show(io,mime,combine(V))
end

#-----------------------------------------------------------------------------


using Test, Random, LinearAlgebra, MPStates

@testset "Tensor factorizations" begin
Random.seed!(0)

@testset "QR decomposition" begin
    (m1, m2, m3) = (4, 5, 6)
    M = randn(m1, m2, m3)
    Q, R = factorize_qr(M)

    # Check dims.
    @test size(Q) == (4, 5, 6)
    @test size(R) == (6, 6)

    # Check Q*R=M.
    Q = reshape(Q, 4*5, 6)
    @test Q*R ≈ reshape(M, 4*5, 6)

    # Check that Q is unitary.
    @test Q'*Q ≈ Matrix(I, 6, 6)

end

@testset "LQ decomposition" begin
    (m1, m2, m3) = (4, 5, 6)
    M = randn(m1, m2, m3)
    L, Q = factorize_lq(M)

    # Check dims.
    @test size(L) == (4, 4)
    @test size(Q) == (4, 5, 6)

    # Check L*Q=M.
    Q = reshape(Q, 4, 5*6)
    @test L*Q ≈ reshape(M, 4, 5*6)

    # Check that Q is unitary.
    @test Q*Q' ≈ Matrix(I, 4, 4)
end

@testset "SVD right decomposition" begin
    (m1, m2, m3, m4) = (3, 4, 5, 3)
    M = randn(m1, m2, m3, m4)

    # Full factorization.
    U, SVt = factorize_svd_right(M, cutoff=0., normalize_S=false)
    @test size(U) == (3, 4, 12)
    @test size(SVt) == (12, 5, 3)
    rU = reshape(U, 12, 12)
    @test rU'*rU ≈ Matrix(I, 12, 12)
    rSVt = reshape(SVt, 12, 15)
    @test rU*rSVt ≈ reshape(M, 12, 15)

    # SVD with cutoff.
    U, SVt = factorize_svd_right(M, cutoff=1., normalize_S=false)
    @test size(U, 1) == 3
    @test size(U, 2) == 4
    @test size(U, 3) < 12
    @test size(SVt, 2) == 5
    @test size(SVt, 3) == 3
    @test size(SVt, 1) < 12
    rU = reshape(U, 12, size(U, 3))
    @test rU'*rU ≈ Matrix(I, size(U, 3), size(U, 3))

    # SVD with dimensional cutoff.
    dimcutoff = 5
    U, SVt = factorize_svd_right(M, cutoff=1., dimcutoff=dimcutoff,
                                 normalize_S=false)
    @test size(U, 1) == 3
    @test size(U, 2) == 4
    @test size(U, 3) == dimcutoff
    @test size(SVt, 2) == 5
    @test size(SVt, 3) == 3
    @test size(SVt, 1) == dimcutoff
    rU = reshape(U, 12, dimcutoff)
    @test rU'*rU ≈ Matrix(I, dimcutoff, dimcutoff)
end

@testset "SVD left decomposition" begin
    (m1, m2, m3, m4) = (3, 4, 5, 3)
    M = randn(m1, m2, m3, m4)

    # Full factorization.
    US, Vt = factorize_svd_left(M, cutoff=0., normalize_S=false)
    @test size(US) == (3, 4, 12)
    @test size(Vt) == (12, 5, 3)
    rUS = reshape(US, 12, 12)
    rVt = reshape(Vt, 12, 15)
    @test rVt*rVt' ≈ Matrix(I, 12, 12)
    @test rUS*rVt ≈ reshape(M, 12, 15)

    # SVD with cutoff.
    US, Vt = factorize_svd_left(M, cutoff=1., normalize_S=false)
    @test size(US, 1) == 3
    @test size(US, 2) == 4
    @test size(US, 3) < 12
    @test size(Vt, 2) == 5
    @test size(Vt, 3) == 3
    @test size(Vt, 1) < 12
    rVt = reshape(Vt, size(Vt, 1), 15)
    @test rVt*rVt' ≈ Matrix(I, size(Vt, 1), size(Vt, 1))

    # SVD with dimensional cutoff.
    dimcutoff = 5
    US, Vt = factorize_svd_left(M, cutoff=1., dimcutoff=dimcutoff,
                                normalize_S=false)
    @test size(US, 1) == 3
    @test size(US, 2) == 4
    @test size(US, 3) == dimcutoff
    @test size(Vt, 2) == 5
    @test size(Vt, 3) == 3
    @test size(Vt, 1) == dimcutoff
    rVt = reshape(Vt, dimcutoff, 15)
    @test rVt*rVt' ≈ Matrix(I, dimcutoff, dimcutoff)
end
end # @testset "Tensor factorizations"


using QSM: GaussSeidel
using QSM: BackwardSweep, ForwardSweep, SymmetricSweep
using QSM: Poisson2, __norm_residual, residual!

@testset "Gauss-Seidel" for T in (Float32, Float64)
    Random.seed!(501)
    for sz in [(16, 17, 18), (19, 20, 21)]
        maxit = 2*prod(sz)
        R = CartesianIndices(map(n -> 2:n-1, sz))

        m = zeros(Bool, sz)
        b = zeros(T, sz)
        x0 = zeros(T, sz)

        m[R] .= true
        b[R] .= randn(size(R))
        x0[R] .= randn(size(R))

        A = Poisson2(T, m, Tuple(rand(T, 3)))
        x1 = mul!(copy(x0), A, x0)

        for S in (SymmetricSweep, ForwardSweep, BackwardSweep)
            for RB in (false, true)
                solver = GaussSeidel(S(); iter=maxit, rb=RB)

                y = solver(copy(x0), A, b)
                @test __norm_residual(b, A, y) / norm(b) ≤ sqrt(eps(T))
                @test norm(residual!(copy(y), b, A, y)) / norm(b) ≤ sqrt(eps(T))

                # Fixed-point test.
                # Relaxation should not alter the exact solution.
                # Use exact solution as initial guess.
                solver = GaussSeidel{S, 1, RB}()
                @test solver(copy(x0), A, x1) ≈ x0

                solver = GaussSeidel{S, 25, RB}()
                @test solver(copy(x0), A, x1) ≈ x0
            end
        end
    end
end


#=

  Mr Greenguest puzzle (fancy dress) in ConstraintSolver.jl

  Problem (and LPL) code in

  http://diuflx71.unifr.ch/lpl/GetModel?name=/demo/demo2

  """
  (** Mr. Greenfan wants to give a dress party where the male guests
   * must wear green dresses. The following rules are given:
   * 1 If someone wears a green tie he has to wear a green shirt.
   * 2 A guest may only wear green socks and a green shirt
   *   if he wears a green tie or a green hat.
   * 3 A guest wearing a green shirt or a green hat or who does
   *   not wear green socks must wear a green tie.
   * 4 A guest who is not dressed according to rules 1-3 must
   *   pay a $11 entrance fee.
   * Mr Greenguest wants to participate but owns only a green shirt
   * (otherwise he would have to pay one for $9). He could buy
   * a green tie for $10, a green hat (used) for $2 and green socks
   * for $12.
   * What is the cheapest solution for Mr Greenguest to participate?
   *)
  """

  The solution:
    t = 1
    h = 0
    r = 1
    s = 0
    n = 0
    cost = 10

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function fancy(print_solutions=true,all_solutions=true,timeout=6)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> all_solutions,
                                                            "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>timeout,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    @variable(model, x[1:5], Bin)
    T,H,R,S,N = x
    @variable(model, 0 <= cost <= 10000, Int)

    #  This is a straight translation from the LPL
    #
    # ( ((T ==1) => (R ==1)) #\/ N ==1)  % Picat constraint
    @constraint(model, T := { R ==1 || N ==1 } )

    # ( ((S ==1 \/ R==1) => (T==1 \/ H==1)) \/ N==1)
    b1 = @variable(model, binary=true)
    @constraint(model, b1 := {S ==1 || R==1})
    b2 = @variable(model, binary=true)
    @constraint(model, b2 := {N==1})
    @constraint(model, b1 => {T ==1 || H ==1 } )
    @constraint(model, b1 + b2 >= 1)

    # ( ((R ==1 \/ H==1 \/ (S!=1)) => T==1) \/ N==1)
    b3 = @variable(model, binary=true)
    @constraint(model, b3 := {R ==1 || H==1 || S!=1})
    @constraint(model, b1 + b3 >= 1)

    @constraint(model, cost == 10*T + 2*H + 12*S + 11*N)

    @objective(model,Min,cost)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    num_sols = 0
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                cost_val = convert.(Integer,JuMP.value.(cost; result=sol))
                println("x:$x_val  cost:$cost_val")
                println("t: $(x_val[1])")
                println("h: $(x_val[2])")
                println("r: $(x_val[3])")
                println("s: $(x_val[4])")
                println("n: $(x_val[5])")
                println("cost: $cost_val")
                println()

            end
        end
    else
        println("status:$status")
    end

    return status, num_sols
end

@time fancy()


include("../src/twoscales_pure_ab.jl")
include("../src/henon_heiles.jl")
using DifferentialEquations
using LinearAlgebra
using Plots
using Random

function getmindif(tab::Vector{Vector{BigFloat}})
    nmmin = Inf
    ret = (0, 0)
    for i = 1:size(tab, 1)
        for j = (i+1):size(tab, 1)
            nm = norm(tab[i] - tab[j], Inf)
            if nm < nmmin
                nmmin = nm
                ret = i, j
            end
        end
    end
    return ret, nmmin
end

function fctmain(n_tau, prec)

    setprecision(prec)
    Random.seed!(5612)
    u0 = rand(BigFloat, 4)

    t_max = big"1.0"
    epsilon = big"1e-3"
    println("epsilon=$epsilon")
    nbmaxtest = 10
    ordmax = 12
    debord = 3
    pasord = 1
    y = ones(Float64, nbmaxtest, div(ordmax - debord, pasord) + 1)
    fct = u -> [u[2]^2, 1 / (1 + u[3]^2), 0.5 - u[4], 0.8u[1]]
    res_gen = Array{Array{BigFloat,1},2}(undef, nbmaxtest, div(ordmax - debord, pasord) + 1)
    x = zeros(Float64, nbmaxtest)
    A = [0 0 1 0; 0 0 0 0; -1 0 0 0; 0 0 0 0]
    parphi = PreparePhi(epsilon, n_tau, A, fct)
    nm = NaN
    ordmax += 1
    ordprep = undef
    nb = 20 * 2^(nbmaxtest)
    solref = undef
    while isnan(nm)
        ordmax -= 1
        ordprep = ordmax + 2
        @time par_u0 = PrepareU0(parphi, ordprep, u0)
        @time pargen = PrepareTwoScalePureAB(nb, t_max, ordmax, par_u0)
        @time solref = twoscales_pure_ab(pargen, only_end = true, diff_fft = false)
        nm = norm(solref, Inf)
    end

    tabsol = Array{Array{BigFloat,1},1}(undef, 1)

    tabsol[1] = solref

    println("ordmax=$ordmax")

    println("solref=$solref")
    solref_gen = solref
    indref = 1

    ind = 1
    for order = debord:pasord:ordmax
        ordprep = order + 2
        println("eps=$epsilon solRef=$solref order=$order")
        nb = 10
        indc = 1
        labels = Array{String,2}(undef, 1, order - debord + 1)
        resnorm = 0
        resnormprec = 1
        sol = undef
        println("preparation ordre $order + 2")
        @time par_u0 = PrepareU0(parphi, ordprep, u0)
        while indc <= nbmaxtest
            @time pargen = PrepareTwoScalePureAB(nb, t_max, order, par_u0)
            @time sol = twoscales_pure_ab(pargen, only_end = true, diff_fft = false)
            push!(tabsol, sol)
            res_gen[indc, ind] = sol
            (a, b), nm = getmindif(tabsol)
            if a != indref
                println("New solref !!!! a=$a, b=$b nm=$nm")
                indref = a
                solref = tabsol[a]
                for i = 1:ind
                    borne = (i < ind) ? size(res_gen, 1) : indc
                    for j = 1:borne
                        nm2 = min(norm(res_gen[j, i] - solref, Inf), 1.1)
                        y[j, i] = nm2 == 0 ? nm : nm2
                    end
                end
            else
                diff = solref - sol
                y[indc, ind] = min(norm(diff, Inf), 1.1)
            end
            x[indc] = 1.0 / nb
            println("result=$y")
            println("log2(y)=$(log2.(y))")
            nb *= 2
            indc += 1
        end
        for i = debord:pasord:order
            labels[1, (i-debord)÷pasord+1] = " eps,order=$(convert(Float32,epsilon)),$i "
        end
        p = Plots.plot(
            x,
            view(y, :, 1:ind),
            xlabel = "delta t",
            xaxis = :log,
            ylabel = "error",
            yaxis = :log,
            legend = :topleft,
            label = labels,
            marker = 2,
        )
        prec_v = precision(BigFloat)
        eps_v = convert(Float32, epsilon)
        Plots.savefig(p, "out/r2_$(prec_v)_$(eps_v)_$(order)_$(n_tau)_fct3.pdf")
        if resnorm > resnormprec
            break
        end
        resnormprec = resnorm
        ind += 1
    end
end

# testODESolver()

# for i=3:9
#     fctMain(2^i)
# end
fctmain(256, 512)


cluster = [1 1 1 1; 2 2 3 3; 4 4 5 5]
X = Array(reshape(1:24, 2, 3, 4))

@testset "pool" begin
    @testset "GlobalPool" begin
        glb_cltr = [1 1 1 1; 1 1 1 1; 1 1 1 1]
        p = GlobalPool(+, 3, 4)
        @test p(X) == NNlib.scatter(+, X, glb_cltr)
    end

    @testset "LocalPool" begin
        p = LocalPool(+, cluster)
        @test p(X) == NNlib.scatter(+, X, cluster)
    end

    @testset "TopKPool" begin
        N = 10
        k, in_channel = 4, 7
        X = rand(in_channel, N)
        for T = [Bool, Float64]
            adj = rand(T, N, N)
            p = TopKPool(adj, k, in_channel)
            @test eltype(p.p) === Float32
            @test size(p.p) == (in_channel,)
            @test eltype(p.Ã) === T
            @test size(p.Ã) == (k, k)
            y = p(X)
            @test size(y) == (in_channel, k)
        end
    end

    @testset "topk_index" begin
        X = [8,7,6,5,4,3,2,1]
        @test topk_index(X, 4) == [1,2,3,4]
        @test topk_index(X', 4) == [1,2,3,4]
    end
end


import MitosisStochasticDiffEq as MSDE
using Test, Random
using StaticArrays
using Statistics
using LinearAlgebra

@testset "static array tests" begin
seed = 1234
Random.seed!(seed)

# define SDE function
f(u,p,t) = p[1]*u .+ p[2]
g(u,p,t) = p[3]

function fstat(u,p,t)
  dx = p[1]*u[1] + p[2]
  @SVector [dx]
end

function gstat(u,p,t)
  @SVector [p[3]]
end

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

# initial condition
u0 = [1.1]
u0stat = @SVector [1.1]

# set true model parameters
par = [-0.2, 0.1, 0.9]

# set of linear parameters Eq.~(2.2)
plin = copy(par)

sdekernel1 = MSDE.SDEKernel(f,g,trange,par)
sdekernel2 = MSDE.SDEKernel(fstat,gstat,trange,par)

Random.seed!(seed)
samples1 = MSDE.sample(sdekernel1, u0, save_noise=false)
Random.seed!(seed)
samples2 = MSDE.sample(sdekernel2, u0stat, save_noise=false)

@test isapprox(samples1[1], samples2[1], rtol=1e-10)
@test isapprox(samples1[2][3], samples2[2][3], rtol=1e-10)
@test typeof(samples2[2][2][end]) <: SArray

# initial values for ODE
mynames = (:logscale, :μ, :Σ);
myvalues1 =  [0.0, 1.5, 0.1];
NT1 = NamedTuple{mynames}(myvalues1)

#myvalues2 = @SVector [0.0, 1.5, 0.1] # will be just converted to a standard array
myvalues2 = [0.0, (@SVector [1.5]), (@SMatrix[0.1])]
NT2 = NamedTuple{mynames}(myvalues2)

sdetildekernel1 = MSDE.SDEKernel(f,g,trange,plin)
sdetildekernel2 = MSDE.SDEKernel(fstat,gstat,trange,plin)
message1, backward1 = MSDE.backwardfilter(sdetildekernel1, NT1)
message2, backward2 = MSDE.backwardfilter(sdetildekernel2, NT2)

@test isapprox(backward1, backward2, rtol=1e-10)
@test typeof(backward2.x[1]) <: SArray
@test typeof(backward2.x[2]) <: SArray
@test typeof(backward2.x[3]) <: SArray

x0 = [1.34]
x0stat = @SVector [1.34]
ll0 = randn()

Random.seed!(seed)
samples1 = MSDE.forwardguiding(sdekernel1, message1,
  (x0, ll0), inplace=true)
Random.seed!(seed)
samples2 = MSDE.forwardguiding(sdekernel2, message2,
  (x0stat, ll0), inplace=false)

@test isapprox(samples1[2][1], samples2[2][1], rtol=1e-10)
@test isapprox(samples1[2][2], samples2[2][2], rtol=1e-10)
@test_broken typeof(samples2[2][end]) <: SArray
@test_broken typeof(samples2[3]) <: SArray
end

@testset "static tilde parameter tests" begin
  d = 2

  𝕏 = SVector{d,Float64}
  𝕏_ = SVector{d+1,Float64}
  bθ = @SVector [.5, 0.9]
  σ0 = @SVector [1.25, 1.35]
  θ0 = (bθ, σ0)

  M = @SMatrix [-1.0 1.0; 1.0 -1.0]
  f(u,θ,t) = 𝕏(tanh.(Diagonal(θ[1]) * M * u))  # f(u,θ,t) = Diagonal(θ[1]) * M * u
  g(u,θ,t) = θ[2]

  dt0 = 0.001
  trange = 0.0:dt0:1.0

  u0 = zero(𝕏)  # value at root node

  B(θ) = (Diagonal(θ[1]) * M)
  Σ(θ) = SMatrix{2,2}(Diagonal(θ[2]))
  b̃ = @SVector [.1, 0.1]  #θ̃ = θ0
  θ̃ = (b̃, σ0)

  θlinstat = ((B(θ̃), @SVector(zeros(d)), Σ(θ̃)))
  θlin =  Array.(θlinstat)
  @show typeof(θlin), typeof(θlinstat)

  # initial values for ODE
  Random.seed!(123)
  logscale = randn()
  ν = @SVector randn(d)
  P = @SMatrix randn(d,d)
  myvalues = [logscale, ν, P]
  mynames = (:logscale, :μ, :Σ)
  NTstat = NamedTuple{mynames}(myvalues)
  NT = NamedTuple{mynames}([logscale, Array(ν), Array(P)])

  tildekernelstat = MSDE.SDEKernel(f,g,trange,θlinstat)
  messagestat, backwardstat = MSDE.backwardfilter(tildekernelstat, NTstat)

  @test typeof(backwardstat.x[1]) <: SArray
  @test typeof(backwardstat.x[2]) <: SArray
  @test typeof(backwardstat.x[3]) <: SArray

  tildekernel = MSDE.SDEKernel(f,g,trange,θlin)
  message, backward = MSDE.backwardfilter(tildekernel, NT)

  @test !(typeof(backward.x[1]) <: SArray)
  @test !(typeof(backward.x[2]) <: SArray)
  @test !(typeof(backward.x[3]) <: SArray)

  @test isapprox(backwardstat.x[1], backward.x[1], rtol=1e-10)
  @test isapprox(backwardstat.x[2], backward.x[2], rtol=1e-10)
  @test isapprox(backwardstat.x[3], backward.x[3], rtol=1e-10)

  # mixed cases
  message1, backward1 = MSDE.backwardfilter(tildekernel, NTstat)
  @test typeof(backward1.x[1]) <: SArray
  @test typeof(backward1.x[2]) <: SArray
  @test typeof(backward1.x[3]) <: SArray

  @test isapprox(backward1.x[1], backward.x[1], rtol=1e-10)
  @test isapprox(backward1.x[2], backward.x[2], rtol=1e-10)
  @test isapprox(backward1.x[3], backward.x[3], rtol=1e-10)

  message2, backward2 = MSDE.backwardfilter(tildekernelstat, NT)
  @test !(typeof(backward2.x[1]) <: SArray)
  @test !(typeof(backward2.x[2]) <: SArray)
  @test !(typeof(backward2.x[3]) <: SArray)

  @test isapprox(backward2.x[1], backward.x[1], rtol=1e-10)
  @test isapprox(backward2.x[2], backward.x[2], rtol=1e-10)
  @test isapprox(backward2.x[3], backward.x[3], rtol=1e-10)

  # backward
  (c, μ, Pmat) = NT
  p = tildekernel.p
  u0 = MSDE.mypack(μ, Pmat, c)
  MSDE.CovarianceFilterODE(u0, p, 0.0)
  b, β, σtil = p
  MSDE.compute_dν(b,μ,β)
  MSDE.compute_dP(b,Pmat,σtil)
  @test typeof(MSDE.compute_dν(b,μ,β)) == typeof(μ)
  @test typeof(MSDE.compute_dP(b,Pmat,σtil)) == typeof(Pmat)

  # backwardstat
  (c, μ, Pmat) = NTstat
  p = tildekernelstat.p
  u0 = MSDE.mypack(μ, Pmat, c)
  MSDE.CovarianceFilterODE(u0, p, 0.0)
  b, β, σtil = p
  MSDE.compute_dν(b,μ,β)
  MSDE.compute_dP(b,Pmat,σtil)
  @test typeof(MSDE.compute_dν(b,μ,β)) == typeof(μ)
  @test typeof(MSDE.compute_dP(b,Pmat,σtil)) == typeof(Pmat)

  # backward1
  (c, μ, Pmat) = NTstat
  p = tildekernel.p
  u0 = MSDE.mypack(μ, Pmat, c)
  MSDE.CovarianceFilterODE(u0, p, 0.0)
  b, β, σtil = p
  MSDE.compute_dν(b,μ,β)
  MSDE.compute_dP(b,Pmat,σtil)
  @test typeof(MSDE.compute_dν(b,μ,β)) == typeof(μ)
  @test typeof(MSDE.compute_dP(b,Pmat,σtil)) == typeof(Pmat)

  # backward2
  (c, μ, Pmat) = NT
  p = tildekernelstat.p
  u0 = MSDE.mypack(μ, Pmat, c)
  MSDE.CovarianceFilterODE(u0, p, 0.0)
  b, β, σtil = p
  @test typeof(MSDE.compute_dν(b,μ,β)) == typeof(μ)
  @test typeof(MSDE.compute_dP(b,Pmat,σtil)) == typeof(Pmat)
end


function sweepline(V,EV)
	# event generation and ordering
	segments = presortedlines(V,EV)
	evpairs = [[(v1,v2,"start",k), (v2,v1,"end",k)] for (k,(v1,v2)) in enumerate(segments)]
	isless0 = (x,y) -> [x[1][1],x[1][2],x[2][2]] < [y[1][1],y[1][2],y[2][2]]
	events = sort(cat(evpairs),lt=isless0)
	# Initialize event queue ξ = all segment endpoints; Sort ξ by increasing x and y
	pqkeys = [(e[1],e[3],e[4]) for e in events]
	pqvalues = events
	ξ = PriorityQueue(zip(pqkeys,pqvalues))
@show ξ,0 println()
	# Proper ordering object for the `SortedMultiDic` used by sweep line SL
	##o = Lt((x,y) -> (x,y))
	# Initialize sweep line SL to be empty
	SL = SortedMultiDict{Any,Any}()##o)
@show SL, 0 println()
	# Initialized output intersection list Λ to be empty
	Λ = Array{Array{Float64,1},1}[]
@show Λ,0 println()

	while length(ξ) ≠ 0 # (ξ is nonempty)
		E = peek(ξ)[2]   # the next value of event (k => v) from ξ
@show E, 1 println()
		(v1, v2, nodetype, edgeid) = E
		# segE = E's segment
		e, segE = v1, (v1, v2, nodetype, edgeid) # key, value
		if E[3] == "start" # (E is a left endpoint)
			# Add segE to SL
			i = insert!(SL, e, segE)  # SortedMultiDict, key, value -> semitoken
@show SL, 1 println()
			global xe,ye = e
			if first(SL) !== last(SL) # more than one segment in SL
				# compute segA and/or segB
				segA, segB = selectsegmentneighbor(SL,i)
				if segA ≠ []
				# segA = the segment above segE in SL
					a = Int(segA[3][1]); segA = segA[1:2] # edgeid of segA
					I = intersection(segE,segA)
					# (if Intersect( segE with segA) exists)
					if typeof(I) ≠ Nothing
						# Insert I into ξ
						#(keyE,newsegE),(keyA,newsegA) = swapsegments(SL,I,segE,segA)
						key = (I,"int",edgeid); val = (I,I,"int",a)
						enqueue!(ξ, key,val)
						#enqueue!(ξ, keyA,newsegA)
@show ξ,6 println()
					end
				end
				if segB ≠ []
				# segB = the segment below segE in SL
					b = Int(segB[3][1]); segB = segB[1:2] # edgeid of segB
					I = intersection(segE,segB)
					# (if Intersect( segE with segB) exists)
					if typeof(I) ≠ Nothing
						# Insert I into ξ
						#(keyE,newsegE),(keyB,newsegB) = swapsegments(SL,I,segE,segB)
						key = (I,"int",edgeid); val = (I,I,"int",b)
						enqueue!(ξ, key,val)
						#enqueue!(ξ, keyB,newsegB)
@show ξ,5 println()
					end
				end
			end
		elseif E[3] == "end" # (E is a right endpoint)
			# segE = E's segment
			# e = key(segE); st1 = corresponding semitoken in SL
			st1,st2 = searchequalrange(SL::SortedMultiDict, e)
			# compute segA and/or segB
			segA, segB = selectsegmentneighbor(SL,st1)
			# segA = the segment above segE in SL
			if segA ≠ []
				a = Int(segA[3][1]) # edgeid of segA
				segA = segA[1:2]
			end
			# segB = the segment below segE in SL
			if segB ≠ []
				b = Int(segB[3][1])
				segB = segB[1:2] # edgeid of segB
			end
			# Remove segE from SL
@show st1,st2 println()
			##delete!((SL,st1))
			# (I = Intersect( segA with segB) exists)
			if segA ≠ [] && segB ≠ []
				I = intersection(segA, segB)
				if typeof(I) ≠ Nothing
					# (I is not in ξ already)
					# no problem anyway (in case I is overwritten)
					# Insert I into ξ
					#(keyA,newsegA),(keyB,newsegB) = swapsegments(SL,I,segA,segB)
					key = (I,"int",a); val = (I,I,"int",b)
					enqueue!(ξ, key,val)
					#enqueue!(ξ, keyB,newsegB)
@show ξ,4 println()
				end
			end
		else # E is an intersection event
		@assert E[3] == "int"
			# Add E to the output list Λ
			push!(Λ, [E[1],E[2]])
@show Λ,1 println()
			# segE1 above segE2 be E's intersecting segments in SL
			i = insert!(SL, e, segE) ### ????
			segE1, segE2 = selectsegmentneighbor(SL,i)
			if segE1 ≠ []
				h = Int(segE1[3][1]) # index of segment segE1
				v1 = V[:,EV[h][1]] # key of segE1 in SL
				hst1,hst2 = searchequalrange(SL::SortedMultiDict, v1)
				# extreme semitokens of segE1 in SL
				# segB = the segment below segE1 in SL
				segE1, segB = selectsegmentneighbor(SL,hst1)
				if segE1 ≠ [] && segB ≠ []
					I = intersection(segE1,segB)
					# (I = Intersect(segE1 with segB) exists)
					if typeof(I) ≠ Nothing
						# (if I is not in ξ already) no problem in case
						# Insert I into ξ
						(keyE1,newsegE1),(keyB,newsegB) = swapsegments(SL,I,segE1,segB)
						key = (I,"int",edgeid); val = (I,I,"int",b) # edgeid ??
						enqueue!(ξ, key,val)
						#enqueue!(ξ, keyB,newsegB)
@show ξ,5 println()
					end
				end
			end
			if segE2 ≠ []
				k = Int(segE2[3][1])  # index of segment segE2
				v1 = V[:,EV[k][1]]  # key of segE2 in SL
				kst1,kst2 = searchequalrange(SL::SortedMultiDict, v1)
				# extreme semitokens of segE2 in SL
				# Swap their positions so that segE2 is now above segE1
				# ??? HOW TO DO IT ???
				# segA = the segment above segE2 in SL
@show SL,kst1 println()
				segA, segE2 = selectsegmentneighbor(SL,kst1)
				if segA ≠ [] && segE2 ≠ []
					I = intersection(segA, segE2)
					# (I = Intersect(segE2 with segA) exists)
					if typeof(I) ≠ Nothing
						# (if I is not in ξ already) no problem in case
						# Insert I into ξ
						(keyA,newsegA),(keyE2,newsegE2) = swapsegments(SL,I,segA,segE2)
						##key = (I,"int",edgeid); val = (I,I,"int",a) # edgeid ??
						enqueue!(ξ, keyA,newsegA)
						enqueue!(ξ, keyE2,newsegE2)
@show ξ,4 println()
					end
				end
			end
			delete!((SL, i)) ### ????
		end
		# remove E from ξ
@show ξ,1 println()
		dequeue!(ξ)
@show ξ,2 println()
	end
	return Λ
end


function ThermalGround(; name)
    @named hp = HeatPort()
    eqs = [hp.T ~ 0]
    ODESystem(eqs, t, systems=[hp], name=name)
end

function HeatCapacitor(; name, C=1.0)
    c_th = C

    @named hp = HeatPort()
    @parameters C
    @variables T(t) dt(t)

    D = Differential(t)
    eqs = [
        T ~ hp.T
        dt ~ D(T)
        D(T) ~ hp.Q_flow / C
        ]
    ODESystem(eqs, t, [T, dt], [C], systems=[hp], defaults=Dict(C => c_th), name=name)
end

function ThermalConductor(; name, G=1.0)
    g_th = G

    @named hp1 = HeatPort()
    @named hp2 = HeatPort()
    @parameters G
    @variables Q_flow(t) T(t)

    eqs = [
        T ~ hp1.T - hp2.T
        Q_flow ~ G*T
        Q_flow ~ hp1.Q_flow
        -Q_flow ~ hp2.Q_flow
    ]
    ODESystem(eqs, t, [Q_flow, T], [G], systems=[hp1, hp2], defaults=Dict(G => g_th), name=name)
end

function ThermalResistor(; name, R=1.0)
    r_th = R
    @named hp1 = HeatPort()
    @named hp2 = HeatPort()
    @parameters R
    @variables Q_flow(t) T(t)

    eqs = [
        T ~ R*Q_flow
        T ~ hp1.T - hp2.T
        hp1.Q_flow ~ Q_flow
        hp2.Q_flow ~ -Q_flow
    ]
    ODESystem(eqs, t, [Q_flow, T], [R], systems=[hp1, hp2], defaults=Dict(R => r_th), name=name)
end

function ConvectiveConductor(; name, G=1.0)
    g_c = G

    @named solidport = HeatPort()
    @named fluidport = HeatPort()
    @parameters G # Convective thermal conductance
    @variables Q_flow(t) dT(t)

    eqs = [
        dT ~ solidport.T - fluidport.T
        solidport.Q_flow ~ Q_flow
        fluidport.Q_flow ~ -Q_flow
        dT ~ G*Q_flow
    ]
    ODESystem(eqs, t, [Q_flow, dT], [G], systems=[solidport, fluidport], defaults=Dict(G => g_c), name=name)
end

function ConvectiveResistor(; name, R=1.0)
    r_c = R

    @named solidport = HeatPort()
    @named fluidport = HeatPort()
    @parameters R # Convective thermal resistance
    @variables Q_flow(t) dT(t)

    eqs = [
        dT ~ solidport.T - fluidport.T
        solidport.Q_flow ~ Q_flow
        fluidport.Q_flow ~ -Q_flow
        dT ~ R*Q_flow
    ]
    ODESystem(eqs, t, [Q_flow, dT], [R], systems=[solidport, fluidport], defaults=Dict(R => r_c), name=name)
end

function BodyRadiation(; name, G=1.0)
    g_r = G
    σ   = 5.6703744191844294e-8 # Stefan-Boltzmann constant

    @named hp1 = HeatPort()
    @named hp2 = HeatPort()
    @parameters G # Net radiation conductance between two surfaces
    @variables Q_flow(t)

    eqs = [
        Q_flow ~ G*σ*(hp1.T^4 - hp2.T^4)
    ]
    ODESystem(eqs, t, [Q_flow], [G], systems=[hp1, hp2], defaults=Dict(G => g_r), name=name)
end

function ThermalCollector(; name, N=1)
    hp = []
    for i in 1:N
        _hp = HeatPort(name=Symbol(:hp, i))
        push!(hp, _hp)
    end
    @named collector_port = HeatPort()

    eqs = [
        collector_port.Q_flow + sum(k -> k.Q_flow, hp) ~ 0
        collector_port.T ~ hp[1].T
    ]
    for i in 1:N-1
        push!(eqs, hp[i].T ~ hp[i+1].T)
    end
    ODESystem(eqs, t, [], [], systems=[hp..., collector_port], name=name)
end


using GeoAcceleratedArrays
using Test

@testset "GeoAcceleratedArrays.jl" begin
    # Accelerate 1000 points
    points = rand(1, 1000)
    acc_points = accelerate(points, KDTreeIndex)

    center = [0.0]
    radius = 0.1
    sphere = HyperSphere(center, radius)

    indices = findall(in(sphere), acc_points)
    in_sphere = acc_points[:, indices]
    @test all(in_sphere .<= radius)
end


using DataFrames, LowRankModels

println("cross validation example")
srand(5)

do_cv = true
do_cv_by_iter = true
do_reg_path = true
do_plot = false

m,n,k = 50,50,3
A = randn(m,k)*randn(k,n) + k*sprandn(m,n,.05)
losses = fill(HuberLoss(),n)
r = QuadReg(.1)
glrm = GLRM(A,losses,r,r,k+2)

if do_cv
    println("Computing cross validation error for each fold")
    params = Params(1.0, max_iter=100, abs_tol=0.0, min_stepsize=.001)
    train_error, test_error, train_glrms, test_glrms = cross_validate(glrm, nfolds=5, params=params)
    df = DataFrame(train_error = train_error, test_error = test_error)
end

if do_cv_by_iter
    println("Computing training and testing error as a function of iteration number")
    train_error, test_error = cv_by_iter(glrm)
    df = DataFrame(train_error = train_error, test_error = test_error)
end

if do_reg_path
    println("Computing regularization path")
    params = Params(1.0, max_iter=50, abs_tol=.00001, min_stepsize=.01)
    train_error, test_error, train_time, reg_params =
    regularization_path(glrm, params=params, reg_params=logspace(2,-2,15))
    df = DataFrame(train_error = train_error, test_error = test_error,
                   train_time = train_time, reg_param = reg_params)
	if do_plot
		p = plot(df, :reg_param, [:train_error, :test_error]; scale = :log, filename = None)
	end
end

println(df)


# Holly Jackson, 2020

using BenchmarkTools
using LazySets
using LinearAlgebra
using Plots
using SparseArrays

## SET UP EXAMPLE SYSTEMS

example_num = 2

if (example_num == 1)
    vertices_x = [0.5, 1.5, 2.5, 0, 1, 2]
    all_vertices_x = [0.5, 1.5, 2.5, 0, 1, 2, 0.5, 1.5, 2.5]
    vertices_y = [2, 2, 2, 1, 1, 1]
    all_vertices_y = [2, 2, 2, 1, 1, 1, 0, 0, 0]
    edges = [[1,4],[1,5],[1,2],[2,5],[2,6],[2,3],[3,6],
        [4,7],[4,5],[5,7],[5,8],[5,6],[6,8],[6,9]]
elseif (example_num == 2)
    all_vertices_x = []
    all_vertices_y = []
    edges = []
    w = 50
    for i = 1:100
        for j = 1:w
            if (i % 2 == 0)
                push!(all_vertices_x, j-1)
                if (i > 1 && j > 1)
                    push!(edges,[(i-2)*w+j-1,(i-1)*w+j])
                end
            else
                # on odd rows offset by 0.5
                push!(all_vertices_x, j-1+0.5)
                if (i > 1 && j < w)
                    push!(edges,[(i-2)*w+j+1,(i-1)*w+j])
                end
            end
            push!(all_vertices_y, abs(i-100))
            if (j > 1)
                push!(edges,[(i-1)*w+j-1,(i-1)*w+j])
            end
            if (i > 1)
                push!(edges,[(i-2)*w+j,(i-1)*w+j])
            end

        end
    end
    vertices_x = all_vertices_x[1:length(all_vertices_x)-w]
    vertices_y = all_vertices_y[1:length(all_vertices_y)-w]
else
    # 3rd example, smaller lattice with several supernodes
    all_vertices_x = []
    all_vertices_y = []
    edges = []
    w = 20
    for i = 1:w
        for j = 1:w
            if (i % 2 == 0)
                push!(all_vertices_x, j-1)
                if (i > 1 && j > 1)
                    push!(edges,[(i-2)*w+j-1,(i-1)*w+j])
                end
            else
                # on odd rows offset by 0.5
                push!(all_vertices_x, j-1+0.5)
                if (i > 1 && j < w)
                    push!(edges,[(i-2)*w+j+1,(i-1)*w+j])
                end
            end
            push!(all_vertices_y, abs(i-100))
            if (j > 1)
                push!(edges,[(i-1)*w+j-1,(i-1)*w+j])
            end
            if (i > 1)
                push!(edges,[(i-2)*w+j,(i-1)*w+j])
            end

        end
    end
    vertices_x = all_vertices_x[1:length(all_vertices_x)-w]
    vertices_y = all_vertices_y[1:length(all_vertices_y)-w]
    # connect a few interior nodes to many other nodes
    supernodes = [[w/2,w/2],[w/2,w/4],[w/2,3*w/4]]
    # edges already connected to supernodes
    for i = 1:3
        s = Int((supernodes[i][1]-1)*w+supernodes[i][2])
        # connect to a bunch of other vertices
        for j = 1:length(vertices_x)
            if (j != s && rand() < 0.8)
                new_edge = [s,j]
                if (j < s)
                    new_edge = [j,s]
                end
                push!(edges, new_edge)
            end
        end
    end
end

N = length(vertices_x) # number of unfixed joints
S = length(edges)      # number of struts

v = zeros(2*N)         # original vertex positions
for i = 1:N
    v[2*(i-1)+1] = vertices_x[i]
    v[2*(i-1)+2] = vertices_y[i]
end

if (example_num == 1)
    f_load = 0.5
elseif (example_num == 2)
    f_load = 2.5
else
    f_load = 2
end
η = 5   # material stiffness

# -----------------------------
# Incidence matrix construction
# -----------------------------

A = spzeros(2*N, 2*S)
b_fix = spzeros(2*S)
for i = 1:S
    n1 = edges[i][1]#Int(edges[i,1])
    if (n1 <= N) # not fixed
        A[2*(n1-1)+1,2*(i-1)+1] = 1
        A[2*(n1-1)+2,2*(i-1)+2] = 1
    else
        b_fix[2*(i-1)+1] = all_vertices_x[n1]
        b_fix[2*(i-1)+2] = all_vertices_y[n1]
    end
    n2 = edges[i][2]#Int(edges[i,2])
    if (n2 <= N) # not fixed
        A[2*(n2-1)+1,2*(i-1)+1] = -1
        A[2*(n2-1)+2,2*(i-1)+2] = -1
    else
        b_fix[2*(i-1)+1] = -all_vertices_x[n2]
        b_fix[2*(i-1)+2] = -all_vertices_y[n2]
    end
end

# print(A)
A_T = transpose(A)
# spy(A,legend=false)#, markersize = 6, c = :reds)

f_L = spzeros(2*N)
if (example_num == 1)
    f_L[2*3] = -f_load
    f_L[2*3+1] = -f_load
else # (example_num == 2 or 3)
    f_L[2*w] = -f_load
    f_L[2*w+1] = -f_load
end

# --------------------------
# Helper matrix construction
# --------------------------

α = spzeros(2*S, 2*S)
for i = 1:S
    α[2*(i-1)+1,2*(i-1)+1] = 1
    α[2*(i-1)+1,2*(i-1)+2] = α[2*(i-1)+2,2*(i-1)+1] = 1
    α[2*(i-1)+2,2*(i-1)+2] = 1
end
α_T = transpose(α)
# print(α)
# spy(α, markersize = 6, c = :reds)

# CALCULATE NOMINAL LENGTHS
L0 = sqrt.(α*(A_T*v + b_fix).^2) # nominal lengths

# ------------------------
# Final system formulation
# ------------------------

function f(u)
    L = sqrt.(α * (A_T * (u + v) + b_fix).^2)
    A * (η * (A_T * (u + v) + b_fix) ./ L .* (L0 - L)) - f_L
end

A_dense = Array(A)
A_T_dense = Array(A_T)
v_dense = Array(v)
b_fix_dense = Array(b_fix)
α_dense = Array(α)
L0_dense = Array(L0)
f_L_dense = Array(f_L)
function f_dense(u)
    L = sqrt.(α_dense * (A_T_dense * (u + v_dense) + b_fix_dense).^2)
    A * (η * (A_T_dense * (u + v_dense) + b_fix_dense) ./ L .* (L0_dense - L)) - f_L_dense
end

println("A forward pass through f(u) takes...")
@btime f(rand(2*N))

## FORWARD MOE AD

function forward_mode_ad(D0, Σ)
    # w = A_T * (u .+ v) .+ b_fix
    w_D0 = A_T * (D0 .+ v) .+ b_fix
    w_Σ = A_T * Σ
    # L = sqrt.(α * (w .* w))
    D0 = α * (w_D0 .* w_D0)
    L_Σ = (1/2) .* (D0) .^ (-1/2) .* (α * (2 .* (w_D0) .* w_Σ))
    L_D0 = sqrt.(D0)
    # s = η .* w ./ L .* (L0 .- L)
    # f = A * s .- f_L
    D0 = A * (η .* w_D0 ./ L_D0 .* (L0 .- L_D0)) .- f_L
    Σ =  A * (η ./ L_D0 .* (L0 .- L_D0) .* w_Σ + η .* w_D0 .* L0 .* (-1) .* (L_D0) .^(-2) .* L_Σ)
    return [D0, Σ]
end

using FiniteDifferences
using Distributed
function forward_mode_ad!(J, D0, Σ)
    @distributed for i = 1:size(Σ,2)
        J[i,:] = central_fdm(5, 1)(ε -> f(D0 .+ ε .* Σ[:,i]), 0)
    end
end

u0 = rand(2*N)
using ForwardDiff
println("Julia ForwardDiff.jl takes...")
@btime ForwardDiff.jacobian(u -> f(u), u0)
println("My forward mode AD takes...")
Σ = spzeros(2*N,2*N) + I
@btime forward_mode_ad(u0, Σ)

## REVERSE MODE AD

# Pullback function on vector λ
function pullback(λ, u)
    w_bar = η .* A_T * λ
    z = A_T * (u .+ v) + b_fix
    x = z.^2
    y = α * x
    y_bar = w_bar .* (z .* (-1/2) .* y.^(-3/2) .* (L0 - sqrt.(y))
        + z ./ sqrt.(y) .* (-(1/2) * y.^(-1/2)))
    x_bar = α_T * y_bar
    z_bar = x_bar .* 2 .* z + w_bar .* (1 ./ sqrt.(y) .* (L0 - sqrt.(y)))
    return A * z_bar
end

function reverse_mode_ad(f, u)
    return pullback(spzeros(2*N,2*N) + I,u)
end

using ReverseDiff
println("Julia ReversedDiff.jl takes...")
@btime ReverseDiff.jacobian(u -> f_dense(u), u0)
println("My reverse mode AD takes...")
@btime reverse_mode_ad(f, u0)

## ASSIGN MATRIX COLORS

function assign_matrix_colors(S, rows_to_ignore, forward)
    if (!forward)
        S = transpose(S)
    end
    colors = Array{Int64,1}(zeros(2*N))
    color_num = 1
    colors[1] = color_num # column to color_num
    @inbounds for i = 2:2*N
        # check collisions in every row
        colliding_colors = Set()
        @inbounds for j = 1:2*N
            if (j in rows_to_ignore)
                continue
            end
            if (S[j,i] != 0)
                # col = find_nz(S[j,1:(i-1)])
                # col = S[j,1:(i-1)]
                for k = 1:(i-1)
                    if (S[j,k] != 0)
                        push!(colliding_colors,colors[k])
                    end
                end
            end
        end
        # are there any colors that didn't collide
        need_new_col = true
        @inbounds for j = 1:color_num
            if !(j in colliding_colors)
                valid_color = j
                need_new_col = false
                break
            end
        end
        if (need_new_col)
            valid_color = color_num + 1
            color_num = color_num + 1
        end
        colors[i] = valid_color
    end
    return colors, color_num
end

function decompress!(J, J_compress, colors, S, forward)
    @inbounds for i = 1:2*N
        col = S[:,i]
        J[:,i] = J_compress[:,colors[i]] .* col
    end
end

function matrix_colors(SP, rows_to_ignore, forward)
    # figure out most dense rows, let's say rows that are more than 80% full
    c, max_color = assign_matrix_colors(SP, rows_to_ignore, forward)
    if (forward)
        Σ_color = spzeros(2*N,max_color)
    else
        Σ_color = spzeros(max_color,2*N)
    end
    @inbounds for i = 1:2*N
        if (forward)
            Σ_color[i,c[i]] = 1.0
        else
            Σ_color[c[i],i] = 1.0
        end
    end
    return [Σ_color, c]
end

## FORWARD MODE AD WITH MATRIX COLORS

# Get sparsity pattern
Σ =  spzeros(2*N,2*N) + I
SP = Array(forward_mode_ad(rand(2*N), Σ)[2])#ForwardDiff.jacobian(u -> f(u), rand(2*N))ReverseDiff.jacobian(f, input
SP[SP .> 0] .= 1.0; SP[SP .< 0] .= 1.0;
spy(SP,legend=false)
# savefig("bone_sparsity2.png")

MC = matrix_colors(SP,Set(),true)
Σ_color = MC[1]
c = MC[2]

MC_rev = matrix_colors(SP,Set(),false)
Σ_color_rev = MC[1]
c_rev = MC[2]

function forward_mode_AD_mat_color!(J, u, c, Σ_color)
    D = forward_mode_ad(u, Σ_color)
    decompress!(J, D[2], c, SP, true)
end

function reverse_mode_AD_mat_color!(J, u, c, Σ_color)
    D = pullback(Σ_color, u)
    decompress!(J, D, c, SP, false)
end

println("FORWARD SPARSE")
J = spzeros(2*N,2*N)
@btime forward_mode_AD_mat_color!(J, u0, c, Σ_color)
println("REVERSE SPARSE")
J = spzeros(2*N,2*N)
@btime reverse_mode_AD_mat_color!(J, u0, c_rev, Σ_color_rev)

## COMBINED MODE AD

function combined_mode_AD(u, c, Σ_color, λ, dense_rows)
    D = forward_mode_ad(u, Σ_color)
    J = spzeros(2*N,2*N)
    decompress!(J, D[2], c, SP, true)
    # fill in dense rows
    R = pullback(λ,u)
    count = 1
    for r in dense_rows
        J[r,:] = R[:,count]
        count = count + 1
    end
    return J
end

# find rows to ignore that are dense
rows_to_ignore = Set()
for i = 1:2*N
    if (sum(SP[i,:])/(2*N) > 0.5)
        push!(rows_to_ignore,i)
    end
end
println("DENSE ROWS: ",rows_to_ignore)

MC = matrix_colors(SP,rows_to_ignore,true)
Σ_color = MC[1]
c = MC[2]

# Stack vectors so all dense rows can be calculated at once
λ = zeros(length(rows_to_ignore),2*N)
count = 1
for r in rows_to_ignore
    λ[count,:] = LazySets.Arrays.SingleEntryVector(r,2*N,1)
    global count = count + 1
end
λ = transpose(λ)

println("COMBINED MODE AD, FORWARD")
@btime combined_mode_AD(u0, c, Σ_color, λ, rows_to_ignore)

function combined_mode_AD_reverse(u, c, Σ_color, λ, dense_cols)
    J = spzeros(2*N,2*N)
    reverse_mode_AD_mat_color!(J, u, c, Σ_color)
    # fill in dense rows
    C = forward_mode_ad(u,λ)
    C = C[2]
    count = 1
    for c in dense_cols
        J[:,c] = C[:,count]
        count = count + 1
    end
    return J
end

# find rows to ignore that are dense
cols_to_ignore = Set()
for i = 1:2*N
    if (sum(SP[:,i])/(2*N) > 0.5)
        push!(cols_to_ignore,i)
    end
end
println("DENSE COLS: ",cols_to_ignore)

MC_rev = matrix_colors(SP,cols_to_ignore,false)
Σ_color_rev = MC[1]
c_rev = MC[2]

# Stack vectors so all dense rows can be calculated at once
λ = zeros(2*N,length(cols_to_ignore))
count = 1
for c in cols_to_ignore
    λ[:,count] = LazySets.Arrays.SingleEntryVector(c,2*N,1)
    global count = count + 1
end

println("COMBINED MODE AD, REVERSE")
@btime combined_mode_AD_reverse(u0, c_rev, Σ_color_rev, λ, cols_to_ignore)

## SOLVE WITH NETWON'S METHOD

# ----------------
# Newton iteration
# ----------------

function newton_step(f, u0)
    # -----------------------------------------------
    # -----------------------------------------------
    # REMEMBER TO SET THIS ACCORDING TO METHOD YOU WANT TO USE
    J = reverse_mode_ad(f, u0)
    # -----------------------------------------------
    # -----------------------------------------------
    # J = combined_mode_ad(u0, c, Σ_color, rows_to_ignore)
    # create the (sparse) LU (etc.) factorization
    F = factorize(J)
    # solve J d = f(u0) with LU decompsition
    d = zeros(2*N)
    ldiv!(d, F, f(u0))
    return u0 - d
end

function newton(f, u0, max_iter)
    u_last = u0

    for i in 1:max_iter
        u = newton_step(f, u_last)
        @show u
        if LinearAlgebra.norm(f(u),2) < 1e-10 && LinearAlgebra.norm(u-u_last,2) < 1e-10
            print("converged in ",i," iterations")
            return u
        else
            println(LinearAlgebra.norm(f(u),2))
            println(LinearAlgebra.norm(u-u_last,2))
        end
        u_last = u
    end

    return u_last
end

## EVALUATE AND PLOT

# initial guess
u0 = 0.0001*ones(2*N)
u = newton(f, u0, 100)

# Plot results
new_vertices_x = zeros(length(all_vertices_x))
new_vertices_y = zeros(length(all_vertices_y))
for i = 1:N
    new_vertices_x[i] = all_vertices_x[i] - u[2*(i-1)+1]
    new_vertices_y[i] = all_vertices_y[i] - u[2*(i-1)+2]
end
for i = N+1:length(all_vertices_x)
    new_vertices_x[i] = all_vertices_x[i]
    new_vertices_y[i] = all_vertices_y[i]
end

gr()
p = plot(title = string("Deformation of Lattice Under ",f_load, " Units Force"), legend = false, aspect_ratio = 1)
for i = 1:S
    strut1 = edges[i][1]
    strut2 = edges[i][2]
    plot!(p, [all_vertices_x[strut1], all_vertices_x[strut2]], [all_vertices_y[strut1], all_vertices_y[strut2]], color = :steelblue)
    plot!(p, [new_vertices_x[strut1] .+ 100, new_vertices_x[strut2] .+ 100], [new_vertices_y[strut1], new_vertices_y[strut2]], color = :green)
end
# scatter!(p, all_vertices_x[N+1:end], all_vertices_y[N+1:end], markersize = 1, color = :red)
# scatter!(p, vertices_x, vertices_y, markersize = 1, color = :steelblue)
# scatter!(p, new_vertices_x[1:N] .+ 100, new_vertices_y[1:N], markersize = 1, color = :green)
p

# plot
savefig(p, "example.png")


# -----------------------------------------------------------
#
# Using boundary terms and integration (summation) by parts
# to construct more accurate derivative matrices for time
# evolution using Lax-Friedrichs-type fluxes
#
# All derivatives are computed here in pos basis first
# Then conjugated into hier basis
#
# -----------------------------------------------------------

# Efficiency criticality: MED
# Most computations time is taken up by functions called from
# other script files

# Accuracy criticality: LOW
# The accuracy is limited only by functions in other scripts
# and ODE.jl

# Construct the time evolution matrix for the wave equation
# given a Kaplacian and the initial position and velocity f0, v0
function wave_data(laplac::A, f0coeffs::AbstractArray{T, 1},
        v0coeffs::AbstractArray{T, 1}) where {A <: AbstractArray, T <: Real}

    len = length(f0coeffs)
    rows = rowvals(laplac)
    vals = nonzeros(laplac)

    I = Int[]
    J = Int[]
    V = T[]
    for col = 1:len
        for i in nzrange(laplac, col)
            row = rows[i]
            val = vals[i]
            push!(I,row+len)
            push!(J,col)
            push!(V,val)
        end
    end
    for i = 1:len
        push!(I, i)
        push!(J, i+len)
        push!(V, 1.0)
    end
    # println("len = ", len) # to make debugging quicker
    # println("max(rows) = ", maximum(rows))
    RHS = sparse(I, J, V, 2*len, 2*len, +)
    y0 = Array{T}([i<=len ? f0coeffs[i] : v0coeffs[i-len] for i in 1:2*len])
    return RHS, y0
end

# Evolves the wave equation from time t0 to t1
# given coefficients for an initial position profile of f0
# and an initial velocity profile of v0
function wave_evolve(D::Int, k::Int, n::Int,
                              f0coeffs::Array{T, 1}, v0coeffs::Array{T, 1},
                              time0::Real, time1::Real;
                              order="45", scheme="sparse", kwargs...) where T <: Real

    laplac   = laplacian_matrix(D, k, n; scheme=scheme)

    RHS, y0 = wave_data(laplac, f0coeffs, v0coeffs)
    if order == "45"
        soln = ode45((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    elseif order == "78"
        soln = ode78((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    else
        throw(ArgumentError(:order))
    end
    return soln
end

# Same function but for anisotropic grids
function wave_evolve(D::Int, k::Int, n::Vector{Int},
                    f0coeffs::Array{T, 1}, v0coeffs::Array{T, 1},
                    time0::Real, time1::Real;
                    order="45", scheme="full", kwargs...) where T <: Real

    laplac   = laplacian_matrix(D, k, n; scheme=scheme)

    RHS, y0 = wave_data(laplac, f0coeffs, v0coeffs)
    if order == "45"
        soln = ode45((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    elseif order == "78"
        soln = ode78((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    else
        throw(ArgumentError(:order))
    end
    return soln
end

# Evolves the wave equation from time t0 to t1
# given initial position profile f0 and velocity profile v0
function wave_evolve(D::Int, k::Int, n::Int,
                              f0::Function, v0::Function,
                              time0::Real, time1::Real;
                              order="45", scheme="sparse", kwargs...)

    f0coeffs = vcoeffs_DG(D, k, n, f0; scheme=scheme)
    v0coeffs = vcoeffs_DG(D, k, n, v0; scheme=scheme)
    return wave_evolve(D, k, n, f0coeffs, v0coeffs, time0, time1;
                        order=order, scheme=scheme, kwargs...)
end

# wave_evolve for anisotropic grids
function wave_evolve(D::Int, k::Int, n::Vector{Int},
                                f0::Function, v0::Function,
                                time0::Real, time1::Real;
                                order="45", scheme="full", kwargs...)

    f0coeffs = vcoeffs_DG(D, k, n, f0; scheme=scheme)
    v0coeffs = vcoeffs_DG(D, k, n, v0; scheme=scheme)
    return wave_evolve(D, k, n, f0coeffs, v0coeffs, time0, time1;
                         order=order, scheme=scheme, kwargs...)
end

# The reason we have a wave evolution in 1D is to test
# that the standard "position" and multiresolution "heirarchical"
# bases agree for wave PDE evolution
function wave_evolve_1D(k::Int, max_level::Int,
                            f0::Function, v0::Function,
                            time0::Real, time1::Real; basis="hier", order = "45",
                            kwargs...)
    if basis == "pos"
        f0coeffs = pos_vcoeffs_DG(k, max_level, f0)
        v0coeffs = pos_vcoeffs_DG(k, max_level, v0)
    elseif basis == "hier"
        f0coeffs = vcoeffs_DG(1, k, max_level, f0)
        v0coeffs = vcoeffs_DG(1, k, max_level, v0)
    elseif basis == "nodal"
        throw(MethodError("Nodal basis wave evolution not yet implemented"))
    elseif basis == "point"
        m2p = modal2points_1D(k, max_level)
        f0coeffs = m2p*vcoeffs_DG(1, k, max_level, f0)
        v0coeffs = m2p*vcoeffs_DG(1, k, max_level, v0)
    else
        throw(ArgumentError(:basis))
    end

    D_op = periodic_DLF_matrix(k, max_level; basis=basis)
    laplac= *(D_op,D_op)

    RHS, y0 = wave_data(laplac, f0coeffs, v0coeffs)
    if order == "45"
        soln = ode45((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    elseif order == "78"
        soln = ode78((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    else
        throw(ArgumentError(:order))
    end
    return soln
end

# Evolve the Vlassov-Poisson (AKA collisionless Boltzmann) equation
# for a mass distribution in D dimensions
# -> in a 2*D-dimensional phase space
# D, k, n (good)
# f0_coeffs_modal (start with anything- like a gaussian) (good)
# F_point (Depends on the potential I pick. Use Isothermal Truncated)
# v_point (its a constant and simple vector for any potential) (good)
# time0, time1, order, scheme (good)
function vlasov_evolve(D::Int, k::Int, n::Int,
            m2n::SparseMatrixCSC{T, Int}, n2p::SparseMatrixCSC{T, Int},
            p2n::SparseMatrixCSC{T, Int}, n2m::SparseMatrixCSC{T, Int},
            f0_modal::Array{T,1}, F_point::Array{Array{T,1}, 1},
            time0::T, time1::T, nout::Int=2;
            order="45", scheme="sparse", kwargs...) where T <: Real

    # The grad matrix- using the same derivative operator as we did
    # for the wave equation.
    # We may need to apply filtering intermittently in the evolution
    Ds = grad_matrix(2*D, k, n; scheme=scheme)

    println("Writing operators...")
    flush(stdout)
    jldopen("vlasov.h5", "w") do f
        f["dimensions"] = D
        f["order"] = k
        f["levels"] = n
        function w(n::AbstractString, A::SparseMatrixCSC)
            f["$n.m"] = A.m
            f["$n.n"] = A.n
            f["$n.colptr"] = A.colptr
            f["$n.rowval"] = A.rowval
            f["$n.nzval"] = A.nzval
        end
        w("m2n", m2n)
        w("n2p", n2p)
        w("p2n", p2n)
        w("n2m", n2m)
        for (d, ds) in enumerate(Ds)
            w("Ds[$d]", ds)
        end
    end

    # Coeffs for the constant 1 in D-dim space
    one_1D = get_one_modal(1, k, n)
    # Coeffs for velocity in D-dim space
    v_modal_1D = get_xi_modal(1, 1, k, n)
    # Coeffs for velocity in 2*D-dim phase space - the tensor product of the above
    v_modal = [tensor_construct(2*D, k, n, [j-D == i ? v_modal_1D : one_1D for j in 1:2*D])
              for i in 1:D]
    v_point = [n2p * (m2n * v) for v in v_modal]

    function steprule(t::T, f_modal::Array{T, 1})
        println("    evaluating RHS at t=$t...")
        flush(stdout)
        # We want this in a modal basis for differentiation to work fast
        # Matrix multiplication is the bottle neck here.
        dfdxs_modal = [Ds[d] * f_modal for d in 1:D]
        dfdps_modal = [Ds[d] * f_modal for d in (D+1):(2*D)]

        dfdxs_point = [n2p * (m2n * x) for x in dfdxs_modal]
        dfdps_point = [n2p * (m2n * x) for x in dfdps_modal]

        # dH/dp * df/dx
        contrib1 = sum(v_point[d] .* dfdxs_point[d] for d in 1:D)
        # dH/dx * df/dp
        contrib2 = sum(F_point[d] .* dfdps_point[d] for d in 1:D)

        # df/dt = - (dH/dp * df/dx - dH/dx * df/dp)
        return n2m * (p2n * (- contrib1 + contrib2))
    end

    println("State vector size: $(length(f0_modal))")
    println("Operator sizes:")
    for d in 1:D
        println("    Ds[$d]: $(nnz(Ds[d]))")
    end
    println("    m2n:   $(nnz(m2n))")
    println("    n2p:   $(nnz(n2p))")
    println("    p2n:   $(nnz(p2n))")
    println("    n2m:   $(nnz(n2m))")
    println("Beginning PDE evolution...")
    flush(stdout)

    if order == "45"
        solver = ode45
    elseif order == "78"
        solver = ode78
    else
        throw(ArgumentError(:order))
    end
    soln = solver(steprule, f0_modal, range(time0, time1, length=nout); kwargs...)

    println("Writing solution...")
    flush(stdout)
    jldopen("vlasov.h5", "r+") do f
        f["times"] = soln[1]
        for (i,sol) in enumerate(soln[2])
            ifmt = @sprintf "%06d" i
            f["f_modal.$ifmt"] = sol
        end
    end

    println("Done.")
    flush(stdout)
end

function norm_squared(coeffs::AbstractArray{T},; basis = "hier") where T <: Real

    if basis=="hier" || basis=="pos"
        return sum(i^2 for i in coeffs)
    else
        throw(MethodError("Nodal and point basis norms not yet implemented"))
    end
end

function energy_func_1D(k, level, soln::Tuple{Array{T, 1}, Array{Array{T, 1}, 1}};
    basis = "hier") where T <: Real

    len            = length(soln[1])
    num_coeffs    = Int(round(length(soln[2][1])/2))
    times        = copy(soln[1])
    energies    = Array{T}(undef, len)
    D_op        = periodic_DLF_matrix(k, level; basis=basis)

    for i in 1:len
        ux   = D_op *  (@view soln[2][i][1:num_coeffs])
        udot = @view soln[2][i][num_coeffs+1:end]
        energies[i] = norm_squared(ux) + norm_squared(udot)

    end
    return (times, energies)
end

function energy_func(D::Int, k::Int, n::Int, soln::Tuple{Array{T, 1}, Array{Array{T, 1}, 1}};
        scheme = "sparse") where T <: Real

    len         = length(soln[1])
    num_coeffs    = div(length(soln[2][1]), 2)
    times         = copy(soln[1])
    energies    = Array{T}(undef, len)
    D_ops        = grad_matrix(D, k, n; scheme=scheme)

    for i in 1:len
        ux   = [D_ops[j] * (@view soln[2][i][1:num_coeffs]) for j in 1:D]
        udot = @view soln[2][i][num_coeffs+1:end]
        energies[i] = sum(norm_squared(ux[d]) for d in 1:D) + norm_squared(udot)
    end
    return (times, energies)
end


#!/usr/bin/env julia

using SphericalFourierBesselDecompositions
SFB = SphericalFourierBesselDecompositions

using Test
using LinearAlgebra

@testset "Window Chains" begin
    # Note: a lot of the differences here come down to the choice of nside and
    # details of how healpy is used.

    @testset "Wmix" begin
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(3, 5, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=1)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :radial, :ang_quarter, :rotate, :separable)

        cache2 = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
        cache3 = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
        cache4 = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
        cache5 = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)

        ## test cache2 consistency
        #@show real.(cache2.wmix') real.(cache2.wmix)
        #@show imag.(cache2.wmix') imag.(cache2.wmix)
        @test cache2.wmix' ≈ cache2.wmix
        #@show real.(cache2.wmix_negm') real.(cache2.wmix_negm)
        #@show imag.(cache2.wmix_negm') imag.(cache2.wmix_negm)
        nlmsize = SFB.getnlmsize(amodes)
        for i=1:nlmsize, i′=1:nlmsize
            n1, l1, m1 = SFB.getnlm(amodes, i)
            n2, l2, m2 = SFB.getnlm(amodes, i′)
            @test cache2.wmix_negm[i′,i] ≈ (-1)^(m1+m2) * cache2.wmix_negm[i,i′]
        end

        get_win(nlm1, nlm2; verbose=false) = begin
            n1, l1, m1 = nlm1
            n2, l2, m2 = nlm2

            nl1 = SFB.getidx(cache2.amodes, n1, l1, 0)
            nl2 = SFB.getidx(cache2.amodes, n2, l2, 0)
            w2 = SFB.WindowChains.get_wmix(cache2.wmix, cache2.wmix_negm,
                                           nl1, m1, nl2, m2)
            verbose && @show cache2.wmix[nl1+abs(m1),nl2+abs(m2)]
            verbose && @show cache2.wmix'[nl1+abs(m1),nl2+abs(m2)]
            verbose && @show conj(cache2.wmix[nl2+abs(m2),nl1+abs(m1)])
            verbose && @show cache2.wmix_negm[nl1+abs(m1),nl2+abs(m2)]
            verbose && @show cache2.wmix_negm'[nl1+abs(m1),nl2+abs(m2)]
            verbose && @show conj(cache2.wmix_negm[nl2+abs(m2),nl1+abs(m1)])

            wlmlm4 = SFB.WindowChains.get_wlmlm(cache4, l1, m1, l2, m2)
            inlnl4 = cache4.Ilnln[l1+1,n1,l2+1,n2]
            w4 = inlnl4 * wlmlm4

            wlmlm5 = SFB.WindowChains.window_wmix(l1, m1, l2, m2,
                                                  cache5.Wlm, cache5.LMcache)
            inlnl5 = cache5.Ilnln[l1+1,n1,l2+1,n2]
            w5 = inlnl5 * wlmlm5
            return w2, w4, w5
        end

        test_mode(nlm1, nlm2; verbose=true) = begin
            w2a, w4a, w5a = get_win(nlm1, nlm2; verbose=verbose)
            verbose && @show nlm1,nlm2
            verbose && @show w2a w4a w5a
            @test w2a ≈ w4a ≈ w5a

            w2b, w4b, w5b = get_win(nlm2, nlm1; verbose=verbose)
            verbose && @show nlm2,nlm1
            verbose && @show w2b w4b w5b
            @test w2b ≈ w4b ≈ w5b

            nlm1[3] = -nlm1[3]
            nlm2[3] = -nlm2[3]
            w2c, w4c, w5c = get_win(nlm1, nlm2; verbose=verbose)
            verbose && @show nlm1,nlm2,nlm1[3]+nlm2[3]
            verbose && @show w2c w4c w5c
            @test w2c ≈ w4c ≈ w5c

            @test w2b ≈ conj(w2a)
            @test w2c ≈ (-1)^(nlm1[3] + nlm2[3]) * conj(w2a)
            @test w4b ≈ conj(w4a)
            @test w4c ≈ (-1)^(nlm1[3] + nlm2[3]) * conj(w4a)
            @test w5b ≈ conj(w5a)
            @test w5c ≈ (-1)^(nlm1[3] + nlm2[3]) * conj(w5a)
        end

        test_mode([1, 1, -1], [1, 0, 0])
        test_mode([1, 1, 1], [1, 1, 1])
        test_mode([1, 1, -1], [1, 1, 1])

        nlmsize = SFB.getnlmsize(amodes)
        long_tests && for j=1:nlmsize, i=1:nlmsize
            n1, l1, m1 = SFB.getnlm(amodes, i)
            n2, l2, m2 = SFB.getnlm(amodes, j)
            test_mode([n1, l1, m1], [n2, l2, m2]; verbose=false)
            test_mode([n1, l1, -m1], [n2, l2, m2]; verbose=false)
            test_mode([n1, l1, m1], [n2, l2, -m2]; verbose=false)
            test_mode([n1, l1, -m1], [n2, l2, -m2]; verbose=false)
        end
    end

    @testset "k = 1" begin
        k = 1
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(3, 5, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=1)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :radial, :ang_quarter, :rotate, :separable)

        wlnn = SFB.win_lnn(win, wmodes, cmodes)
        @show wlnn

        cache1 = SFB.WindowChains.WindowChainsCacheWignerChain(win, wmodes, amodes)
        cache2 = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
        cache3 = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
        cache4 = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
        cache5 = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)
        cache = SFB.WindowChainsCache(win, wmodes, amodes)

        # single modes
        ell = [1]
        n1 = [1]
        n2 = [1]
        wk0 = (2*ell[1] + 1) * wlnn[SFB.getidx(cmodes, ell[1], n1[1], n2[1])]
        wk1 = SFB.window_chain(ell, n1, n2, cache1)
        wk2 = SFB.window_chain(ell, n1, n2, cache2)
        wk3 = SFB.window_chain(ell, n1, n2, cache3)
        wk4 = SFB.window_chain(ell, n1, n2, cache4)
        wk5 = SFB.window_chain(ell, n1, n2, cache5)
        wk = SFB.window_chain(ell, n1, n2, cache)
        @show wk wk0 wk1 wk2 wk3 wk4 wk5
        @test wk ≈ wk0
        @test_skip wk ≈ wk1
        @test wk ≈ wk2
        @test wk ≈ wk3
        @test wk ≈ wk4
        @test wk ≈ wk5

        # all modes
        wk_lnni = SFB.window_chain(k, win, wmodes, cmodes)
        @show wk_lnni
        @test all(isreal.(wk_lnni))
        wk_lnni = real.(wk_lnni)

        @show size(wk_lnni) size(wlnn)
        long_tests && for i=1:SFB.getlnnsize(cmodes)
            l, n1, n2 = SFB.getlnn(cmodes, i)
            wk = wk_lnni[i] / (2*l+1)
            @show l,n1,n2,wk_lnni[i],wk,wlnn[i]
            @test wk ≈ wlnn[i]
        end
    end

    @testset "k = 2" begin
        k = 2
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(2, 2, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=1)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :radial, :ang_quarter, :rotate, :separable)

        # comparison
        @time M = SFB.power_win_mix(win, wmodes, cmodes)
        @time M′ = SFB.power_win_mix(win, wmodes, cmodes, interchange_NN′=true)

        cache1 = SFB.WindowChains.WindowChainsCacheWignerChain(win, wmodes, amodes)
        cache2 = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
        cache3 = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
        cache4 = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
        cache5 = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)
        cache = SFB.WindowChainsCache(win, wmodes, amodes)

        # single modes
        test_single(ell, n1, n2) = begin
            @show ell,n1,n2
            lnnA = SFB.getidx(cmodes,ell[2],n2[2],n1[2])
            lnnB = SFB.getidx(cmodes,ell[1],n1[1],n2[1])
            wk0 = (2*ell[2] + 1) * M′[lnnA, lnnB]
            wk1 = SFB.window_chain(ell, n1, n2, cache1)
            wk2 = SFB.window_chain(ell, n1, n2, cache2)
            wk3 = SFB.window_chain(ell, n1, n2, cache3)
            wk4 = SFB.window_chain(ell, n1, n2, cache4)
            wk5 = SFB.window_chain(ell, n1, n2, cache5)
            wk = SFB.window_chain(ell, n1, n2, cache)
            @show wk wk0 wk1 wk2 wk3 wk4 wk5
            @test wk ≈ wk0
            @test_skip wk ≈ wk1
            @test wk ≈ wk2
            @test wk ≈ wk3
            @test wk ≈ wk4
            @test wk ≈ wk5
        end
        test_single([1,1], [1,1], [1,1])
        test_single([2,0], [1,1], [2,1])
        test_single([0,0], [1,1], [1,2])
        test_single([0,0], [1,1], [2,2])
        test_single([0,1], [1,1], [1,1])

        # all modes
        @time wk_lnni = SFB.window_chain(k, win, wmodes, cmodes)
        @time wk_lnni = SFB.window_chain(k, win, wmodes, cmodes)
        @test all(isreal.(wk_lnni))

        @show size(wk_lnni) size(M)
        @test size(wk_lnni) == size(M)
        long_tests && for i2=1:SFB.getlnnsize(cmodes), i1=1:SFB.getlnnsize(cmodes)
            l1, n1, n1′ = SFB.getlnn(cmodes, i1)
            l2, n2, n2′ = SFB.getlnn(cmodes, i2)
            i2′ = SFB.getidx(cmodes, l2, n2′, n2)
            wk = wk_lnni[i1,i2′] / (2*l2+1)
            @show l1,n1,n1′,l2,n2,n2′,wk,M′[i2,i1]
            @test wk ≈ M′[i2,i1]
        end
    end

    @testset "Wk method agreement" begin
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(8, 4, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=0)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :radial, :ang_quarter, :rotate, :separable)

        cache1 = SFB.WindowChains.WindowChainsCacheWignerChain(win, wmodes, amodes)
        cache2 = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
        cache3 = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
        cache4 = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
        cache5 = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)
        cache = SFB.WindowChainsCache(win, wmodes, amodes)

        test_individual(ell, n1, n2) = begin
            @time wk1 = SFB.window_chain(ell, n1, n2, cache1)
            @time wk2 = SFB.window_chain(ell, n1, n2, cache2)
            @time wk3 = SFB.window_chain(ell, n1, n2, cache3)
            @time wk4 = SFB.window_chain(ell, n1, n2, cache4)
            @time wk5 = SFB.window_chain(ell, n1, n2, cache5)
            @time wk = SFB.window_chain(ell, n1, n2, cache)
            @show ell n1 n2
            @show wk1 wk2 wk3 wk4 wk5 wk
            @test_skip wk1 ≈ wk
            @test wk2 ≈ wk
            @test wk3 ≈ wk
            @test wk4 ≈ wk
            @test wk5 ≈ wk
            return wk
        end

        println("Testing individual modes...")
        long_tests && for l1=0:amodes.lmax
            test_individual([l1], [1], [1])
            for l2=0:amodes.lmax
                test_individual([l1, l2], [1, 1], [1, 1])
                for l3=0:amodes.lmax
                    test_individual([l1, l2, l3], [1, 1, 1], [1, 1, 1])
                end
            end
        end
        #test_individual([0, 0, 0, 1], [1, 1, 1, 1], [1, 1, 1, 1])
        #test_individual([1, 3, 4, 5], [1, 1, 1, 1], [1, 1, 1, 1])
        #test_individual([2, 3, 4, 5], [1, 1, 1, 1], [1, 1, 1, 1])
        #test_individual([3, 3, 3, 3], [1, 1, 2, 2], [2, 2, 1, 1])

        # check n ↔ n′ symmetry doesn't exist
        ell = [2, 1, 0]
        n1 = [4, 2, 2]
        n2 = [2, 3, 3]
        test_individual(ell, n1, n2)
        test_individual(ell, n2, n1)

        # check antisymmetry doesn't exist
        ell = [2, 1, 0, 1]
        n1 = [4, 2, 2, 2]
        n2 = [2, 3, 3, 2]
        test_individual(ell, n1, n2)
        test_individual(ell, n2, n1)
        ell[1:2] .= ell[2:-1:1]
        n1[1:2] .= n1[2:-1:1]
        n2[1:2] .= n2[2:-1:1]
        test_individual(ell, n1, n2)

        # check anticyclic symmetry
        #ell = [2, 1, 0, 1]
        #n1 = [2, 3, 3, 2]
        #n2 = [4, 2, 2, 2]
        for i=1:10
            k = 4
            p1 = k:-1:1
            ell = rand(0:amodes.lmax, k)
            n1 = rand(1:amodes.nmax, k)
            n2 = rand(1:amodes.nmax, k)
            wk = test_individual(ell, n1, n2)
            wkp1 = test_individual(ell[p1], n2[p1], n1[p1])
            @show wk wkp1
            @test wk ≈ wkp1
        end

        return

        # These checks can fail:
        # check k=3 n,N ↔ n′,N′ symmetry, Nah, not a symmetry
        for i=1:10
            k = 3
            ell = rand(0:amodes.lmax, k)
            n1 = rand(1:amodes.nmax, k)
            n2 = rand(1:amodes.nmax, k)
            wk = test_individual(ell, n1, n2)
            n1[2:3] .= n1[3:-1:2]
            n2[2:3] .= n2[3:-1:2]
            wkp1 = test_individual(ell, n2, n1)
            @show wk wkp1
            @test wk ≈ wkp1
        end
    end

    @testset "Symmetricizing f=$features" for features in [(), (:separable,)]
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(8, 4, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=0)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :radial, :ang_quarter, :rotate, features...)
        cache = SFB.WindowChainsCache(win, wmodes, amodes)

        ell = [1, 2, 3, 4]
        n1 = [1, 2, 3, 4]
        n2 = [5, 6, 7, 8]
        symmetries = [1=>0, 2=>1, 3=>2]
        wkfull = SFB.window_chain(ell, n1, n2, cache, symmetries)
        @show wkfull
        @test wkfull ≈ -1.2707026010569427e-8  # only useful for detecting changes
    end

    @testset "Wk single pixel mask" begin
        rmin = 500.0
        rmax = 1000.0
        nside = 16
        amodes = SFB.AnlmModes(8, 4, rmin, rmax; nside=nside)
        cmodes = SFB.ClnnModes(amodes, Δnmax=0)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        mask = fill(0.0, SFB.hp.nside2npix(amodes.nside))
        mask[805+1] = 1.0
        phi = fill(1.0, wmodes.nr)
        win = SFB.SeparableArray(phi, mask, name1=:phi, name2=:mask)

        cache1 = SFB.WindowChains.WindowChainsCacheWignerChain(win, wmodes, amodes)
        cache2 = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
        cache3 = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
        cache4 = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
        cache5 = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)
        cache = SFB.WindowChainsCache(win, wmodes, amodes)

        test_individual(ell, n1, n2) = begin
            @time wk1 = SFB.window_chain(ell, n1, n2, cache1)
            @time wk2 = SFB.window_chain(ell, n1, n2, cache2)
            @time wk3 = SFB.window_chain(ell, n1, n2, cache3)
            @time wk4 = SFB.window_chain(ell, n1, n2, cache4)
            @time wk5 = SFB.window_chain(ell, n1, n2, cache5)
            @time wk = SFB.window_chain(ell, n1, n2, cache)
            @show ell n1 n2
            @show wk1 wk2 wk3 wk4 wk5 wk
            @test_skip wk1 ≈ wk
            @test wk2 ≈ wk
            @test wk3 ≈ wk
            @test wk4 ≈ wk
            @test wk5 ≈ wk
            return wk
        end

        println("Testing individual modes...")
        for l1=0:amodes.lmax
            test_individual([l1], [1], [1])
            long_tests && for l2=0:amodes.lmax
                test_individual([l1, l2], [1, 1], [1, 1])
                for l3=0:amodes.lmax
                    test_individual([l1, l2, l3], [1, 1, 1], [1, 1, 1])
                end
            end
        end
    end

    @testset "Other tests" begin
        rmin = 0.0
        rmax = 2000.0
        amodes = SFB.AnlmModes(0.03, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=0)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        @show amodes.lmax_n

        phi = fill(1.0, wmodes.nr)
        mask = fill(1.0, SFB.hp.nside2npix(amodes.nside))
        win = SFB.SeparableArray(phi, mask, name1=:phi, name2=:mask)

        wlm, LMcache = SFB.WindowChains.calc_Wlm(mask, amodes.lmax, amodes.nside)
        lrange = 48:48
        long_tests && for l1=lrange, l2=lrange, m1=-l1:l1, m2=-l2:l2
            w = SFB.WindowChains.window_wmix(l1, m1, l2, m2, wlm, LMcache)
            if l1 == l2 && m1 == m2
                @test w ≈ 1
                if !(w ≈ 1)
                    @error "Wlmlm unexpected result" l1,l2 m1,m2 w
                end
            else
                @test w ≈ 0  atol=1e-10
                if !(abs(w) <= 1e-10)
                    @error "Wlmlm unexpected result" l1,l2 m1,m2 w
                end
            end
        end

        #return

        # Note: the following tests can only be done reasonably fast with
        # WignerFamilies instead of WignerSymbols in the window chain.

        cache = SFB.WindowChainsCache(win, wmodes, amodes)
        @show typeof(cache)
        for l=0:amodes.lmax
            nmax = amodes.nmax_l[l+1]
            @test cache.Ilnln[l+1,1:nmax,l+1,1:nmax] ≈ I  atol=1e-4
            Ik = SFB.WindowChains.calc_kprod(cache.Ilnln, [l,l] .+ 1, [1,1], [1,1])
            @test Ik ≈ 1  atol=1e-4
        end

        ell = [48, 48]
        n1 = [1, 1]
        n2 = [1, 1]
        wk = SFB.window_chain(ell, n1, n2, cache)
        @show wk wk/(2*ell[1]+1)
        @test wk/(2*ell[1]+1) ≈ 1

        long_tests && for m1=-ell[1]:ell[1], m2=-ell[2]:ell[2]
            w = SFB.WindowChains.get_wlmlm(cache, ell[end], m2, ell[1], m1)
            w2 = SFB.WindowChains.window_wmix(ell[end], m2, ell[1], m1, wlm, LMcache)
            if m1 != m2 || ell[1] != ell[2]
                @test w ≈ 0  atol=1e-10
                @test w2 ≈ 0  atol=1e-10
            else
                @test w ≈ 1
                @test w2 ≈ 1
                if !(w ≈ 1)
                    @error "Wlmlm unexpected result" ell m1,m2 w w2
                end
            end
        end

        ell = [48, 0]
        n1 = [1, 1]
        n2 = [2, 2]
        wk = SFB.window_chain(ell, n1, n2, cache)
        @show wk
    end
end

# vim: set sw=4 et sts=4 :


# The multi distribution represents `n` independent draws of the
# distribution given as a parameter
struct Multi{T<:Distribution}
    n::Int
    distribution::T
end

function Distributions.logpdf(multi::Multi, samples)
    sum([logpdf(multi.distribution, sample) for sample in samples])
end

function Base.rand(multi::Multi)
    rand(multi.distribution, multi.n)
end