function barycenter_matrices(d, models, normalize=true; allow_shortcut=true)
    r = roots.(SpectralDistances.Continuous(), models)
    w = [Float64.(d.weight(r)) for r in r]
    realpoles = any(any(iszero ∘ imag, r) for r in r)

    realpoles = realpoles || !allow_shortcut # if we don't allow the shortcut it's the  same as if there are real poles.

    if !realpoles
        r = [r[1:end÷2] for r in r]
        w = [2w[1:end÷2] for w in w]
    end
    X = [[Float64.(real(r)'); Float64.(imag(r)')] for r in r]
    if !all(sum.(w) .≈ 1)
        if normalize
            w = s1.(w)
        else
            @warn "sum.(w) ≠ 1" sum.(w)
        end
    end

    X, w, realpoles
end

function embedding(model::AR,args...)
    embedding(roots(SpectralDistances.Continuous(), model), args...)
end

"""
    embedding([::Type{Vector}], m, [full=true])

Returns a `Vector/Matrix` containing the roots of `m`.
`full` indicates whether or not to use all poles or only one half-plane.
"""
embedding(::Type{Vector}, model, args...) = embedding(model, args...)
embedding(::Type{Matrix}, model, args...) = reshape(embedding(model, args...), :, 2)

function embedding(r::AbstractRoots, full::Bool=true)
    if full
        [real.(r); imag.(r)]
    else
        realpoles = any(any(iszero ∘ imag, r) for r in r)
        realpoles && @error("Model contained real poles, this produces weird results when used in an embedding")
        @views [real.(r[1:end÷2]); imag.(r[1:end÷2])]
    end
end

"""
    barycenter(d::OptimalTransportRootDistance, models; normalize = true, kwargs...)
    $(SIGNATURES)

**Approximately** calculate the barycenter supported on the same number of atoms as the number of poles in the models.

The solver can be selected by providing a keword argument, example: `solver=IPOT`.

Uses the algorithms from ["Fast Computation of Wasserstein Barycenters"](https://arxiv.org/pdf/1310.4375.pdf)

# Example:
```julia
models = examplemodels(10)

d = OptimalTransportRootDistance(domain=SpectralDistances.Continuous(),p=2, weight=residueweight, β=0.01)
Xe = barycenter(d, models, solver=sinkhorn_log!, uniform=true)

plot()
pzmap!.(models)
pzmap!(tf(Xe), m=:c, title="Barycenter OptimalTransportRootDistance", lab="BC")
```

# Arguments:
- `models`: vector of AR models
- `normalize`: make sure weights sum to 1
- `kwargs`: are sent to the solver
"""
function barycenter(d::OptimalTransportRootDistance,models, λ=s1(ones(length(models))); normalize=true, uniform=true, solver=sinkhorn_log!, kwargs...)
    d.p == 2 || throw(ArgumentError("p must be 2"))
    X, w, realpoles = barycenter_matrices(d, models, normalize)
    if uniform
        bc = barycenter(X, λ; uniform=true, solver=solver, β=d.β, kwargs...)
    else
        bc,w = barycenter(X,w, λ; uniform=false, solver=solver, β=d.β, kwargs...)
    end
    bc2model(bc, realpoles)

end

function bc2model(bc, realpoles)
    r1 = complex.(bc[1,:],bc[2,:])
    if realpoles
        bcr = ContinuousRoots(r1)
    else
        @assert !any(iszero ∘ imag, r1) "A real root was found in barycenter even though inputs had no real roots"
        bcr = ContinuousRoots([r1; conj.(r1)])
    end
    AR(bcr)
end

"""
    barycenter(d::EuclideanRootDistance, models::AbstractVector, [λ])

# Example:
```julia
models = examplemodels(10)

Xe = barycenter(EuclideanRootDistance(domain=SpectralDistances.Continuous(),p=2), models)

G = tf.(models)
plot()
pzmap!.(G)
pzmap!(tf(Xe), m=:c, title="Barycenter EuclideanRootDistance")
current()
```
"""
function barycenter(d::EuclideanRootDistance,models::AbstractVector;kwargs...)
    r = roots.(SpectralDistances.Continuous(), models)
    w = map(d.weight, r)
    bc = map(1:length(r[1])) do pi
        sum(w[pi]*r[pi] for (w,r) in zip(w,r))/sum(w[pi] for w in w)
    end
    AR(ContinuousRoots(bc))
end

function barycenter(d::EuclideanRootDistance,models::AbstractVector, λ;kwargs...)
    N = length(models)
    r = roots.(SpectralDistances.Continuous(), models)
    w = map(d.weight, r)
    bc = map(1:length(r[1])) do pi
        sum(λ[i]*w[pi]*r[pi] for (i,w,r) in zip(1:N,w,r))/sum(λ[i]*w[pi] for (i,w) in enumerate(w))
    end
    AR(ContinuousRoots(bc))
end

function distmat_euclidean(X::AbstractMatrix,Y::AbstractMatrix)
    C = [mean(abs2, c1-c2) for c1 in eachcol(X), c2 in eachcol(Y)]
end

function distmat_euclidean!(C,X::AbstractMatrix,Y::AbstractMatrix)
    for (j,c2) in enumerate(eachcol(Y))
        for (i,c1) in enumerate(eachcol(X))
            C[i,j] = mean(((c1,c2),) -> abs2(c1-c2), zip(c1,c2))
        end
    end
    C
end

"""
    barycenter(X::Vector{<:AbstractArray}, λ)

Calculate the weighted barycenter for point clouds in `X`.
Each `X[i]` has the shame `n_dims × n_atoms`
`λ` is the weight vector that should sum to 1.
"""
function barycenter(X::Vector{<:AbstractArray}, λ; uniform=true, solver=sinkhorn_log, kwargs...)
    N = length(X)
    n = size(X[1],2)
    w = s1(ones(n))

    # X0 = mean(X)
    ind = rand(1:length(X))
    X0 = X[ind] .- mean(X[ind],dims=2) .+ mean(mean.(X, dims=2))
    perturb!(X0,X)
    bc = alg2(X0,X,w,fill(w,N); solver=solver, weights=λ, uniform=uniform, kwargs...)[1]
end

function perturb!(X0,X)
    m = 0.1minimum(std, X)
    X0 .+= m .* randn.()
end

function barycenter(X::Vector{<:AbstractArray}, p, λ; uniform=true, solver=sinkhorn_log!, kwargs...)
    N = length(X)
    n = size(X[1],2)
    w = s1(ones(n))
    m = 0.1minimum(std, X)
    # X0 = mean(X)
    ind = rand(1:length(X))
    X0 = X[ind] .- mean(X[ind],dims=2) .+ mean(mean.(X, dims=2))
    X0 .+= m .* randn.()
    alg2(X0,X,w,p; solver=solver, weights=λ, uniform=uniform, kwargs...)
end

barycentric_weighting(X,λ,sw) = sum(λ[i].*X[i][:,sw[i]] for i in eachindex(sw))

function softmax(x)
    e = exp.(x)
    e ./= sum(e)
end

using Random

"""
    λ = barycentric_coordinates(pl, ql, p, q; options, kwargs...)

Compute the barycentric coordinates `λ` such that
sum(λᵢ W(pᵢ,q) for i in eachindex(p)) is minimized.

This function works best with the `sinkhorn_log!` solver, a large β (around 1) and small tolerance. These are set using `kwargs...`.

# Arguments:
- `pl`: Atoms in measures `p`, vector, length `n_measures`, of matrices of size `n_dims × n_atoms`
- `ql`: Atoms in measure `q`
- `p`: Measures `p`, a matrix of weight vectors, size `n_atoms × n_measures` that sums to 1
- `q`: the veight vector for measure `q`, length is `n_atoms`
- `options`: For the Optim solver. Defaults are `options = Optim.Options(store_trace=false, show_trace=false, show_every=0, iterations=20, allow_f_increases=true, time_limit=100, x_tol=1e-5, f_tol=1e-6, g_tol=1e-6, f_calls_limit=0, g_calls_limit=0)`
- `solver`: = [`sinkhorn_log!`](@ref) solver
- `tol`:    = 1e-7 tolerance
- `β`:      = 0.1 entropy regularization. This function works best with rather large regularization, hence the large default value.
- `kwargs`: these are sent to the solver algorithm.

# Example:
```julia
using SpectralDistances, ControlSystems, Optim
models = examplemodels(10)

d = OptimalTransportRootDistance(
    domain = SpectralDistances.Continuous(),
    p      = 2,
    weight = residueweight,
    β      = 0.01,
)
Xe = barycenter(d, models, solver=sinkhorn_log!)

G = tf.(models)
plot()
pzmap!.(G)
pzmap!(tf(Xe), m=:c, title="Barycenter OptimalTransportRootDistance", lab="BC")

options = Optim.Options(store_trace       = true,
                        show_trace        = false,
                        show_every        = 1,
                        iterations        = 50,
                        allow_f_increases = true,
                        time_limit        = 100,
                        x_tol             = 1e-7,
                        f_tol             = 1e-7,
                        g_tol             = 1e-7,
                        f_calls_limit     = 0,
                        g_calls_limit     = 0)

method = LBFGS()
λ = barycentric_coordinates(d, models, Xe, method,
    options = options,
    solver  = sinkhorn_log!,
    robust  = true,
    uniform = true,
    tol     = 1e-6,
)
bar(λ, title="Barycentric coorinates")

G = tf.(models)
plot()
pzmap!.(G, lab="")
pzmap!(tf(Xe), m = :c, title = "Barycenter OptimalTransportRootDistance", lab = "BC")
# It's okay if the last system dot does not match the barycenter exactly, there are limited models to choose from.
pzmap!(G[argmax(λ)], m = :c, lab = "Largest bc coord", legend = true)
```
"""
function barycentric_coordinates(pl, ql, p, q::AbstractVector{T}, method=LBFGS();
    options = Optim.Options(store_trace       = false,
                            show_trace        = false,
                            show_every        = 1,
                            iterations        = 20,
                            allow_f_increases = true,
                            time_limit        = 100,
                            x_tol             = 1e-5,
                            f_tol             = 1e-6,
                            g_tol             = 1e-6,
                            f_calls_limit     = 0,
                            g_calls_limit     = 0),
    robust = true,
    solver = sinkhorn_log!,
    tol    = 1e-7,
    β      = 0.1,
    plot   = nothing,
    kwargs...) where T

    # C = [[mean(abs2, x1-x2) for x1 in eachcol(Xi), x2 in eachcol(ql)] for Xi in pl]

    S = length(pl)
    k = length(p[1])
    # λl = 1e-8randn(S)
    # dists = map(i->sum(C[i].*IPOT(C[i],p[i],q)[1]), eachindex(p))
    # λl = -v1(dists) # Initial guess based on distances between anchors and query point
    # λl .*= 0sqrt(length(λl)) # scale so that softmax(λ) is reasonably sparse.
    # randn!(λl)
    # return softmax(λl)
    # function fg!(F,G,λl)
    #     λ = softmax(λl) # optimization done in log domain
    #     local cost = 0.0
    #     if G !== nothing
    #         cost,P,∇ℰ = sinkhorn_diff(pl,ql,p,q,C,λ; L=L, kwargs...)
    #         G .= ∇ℰ
    #     end
    #     if F !== nothing
    #         if G === nothing
    #             cost = sinkhorn_cost(pl,ql,p,q,C,λ; L=L, kwargs...)
    #         end
    #         return cost
    #     end
    # end
    # if robust
    #     res = Optim.optimize(Optim.only_fg!(fg!), λl, NelderMead(), Optim.Options(iterations=60, store_trace=false))
    #     λl = res.minimizer
    # end
    # res = Optim.optimize(Optim.only_fg!(fg!), λl, method, options)

    λl = zeros(T,S)
    C = zeros(T,k,k)
    costfun = λ -> sinkhorn_cost(pl, ql, p, q, softmax(λ);
        solver = solver,
        tol    = tol,
        β      = β,
        kwargs...)

    if plot !== nothing
        @assert S == 2 "Can only plot for two anchor measures"
        cf = x -> costfun(log.([x,1-x]))
        plot(LinRange(1e-3, 1-1e-3, 100), cf) |> display
    end
    if robust
        res = Optim.optimize(costfun, λl, ParticleSwarm(), Optim.Options(iterations=100, store_trace=false))
        λl = res.minimizer
    end
    local λh
    try
        res = Optim.optimize(costfun, λl, method, options, autodiff=:forward)
        λh = softmax(res.minimizer)
    catch err
        @error("Barycentric coordinates: optimization failed, retrying with robust options: ", err)
        res = Optim.optimize(costfun, λl, NelderMead(), options)
        λh = softmax(res.minimizer)
    end

    # cost,P,∇ℰ = sinkhorn_diff(pl,ql,p,q,C,λh; kwargs...)
    λh
end

function barycentric_coordinates(d::OptimalTransportRootDistance,models, qmodel, method=BFGS(); kwargs...)

    d.p == 2 || throw(ArgumentError("p must be 2"))
    pl, p, realpolesp = barycenter_matrices(d, models)
    ql, q, realpolesq = barycenter_matrices(d, [qmodel])
    if realpolesp != realpolesq
        pl, p, realpolesp = barycenter_matrices(d, models, allow_shortcut=false)
        ql, q, realpolesq = barycenter_matrices(d, [qmodel], allow_shortcut=false)
    end

    for pi in p
        pi ./= sum(pi)
    end
    q = q[1] ./= sum(q[1])
    @assert sum(q) ≈ 1
    @assert all(sum(p) ≈ 1 for p in p)

    if d.divergence !== nothing
        if :solver ∈ keys(kwargs)
            throw(ArgumentError("Solver choice not available for unbalanced distances, solver `sinkhorn_unbalanced` will be used."))
        end
        solver, kwargs = unbalanced_solver_closure(d, kwargs...)
    else
        solver = get(kwargs, :solver, IPOT)
    end
    λ = barycentric_coordinates(pl, ql[1], p, vec(q), method; solver=solver, β=d.β, kwargs...)
    return λ
end

function barycentric_coordinates(d::EuclideanRootDistance, models, qmodel, args...; kwargs...)
    d.p == 2 || throw(ArgumentError("p must be 2"))
    N = length(models)
    pl = embedding.(models)
    ql = embedding(qmodel)
    if d.weight != unitweight
        return _wrd_barycentric_coordinates(d,models,pl,ql)
    end
    ⅅ  = reduce(hcat, pl) # Dictionary matrix
    TotalLeastSquares.sls(Float64.(ⅅ), Float64.(ql); kwargs...)
end

function _wrd_barycentric_coordinates(d,models,pl,ql)
    N = length(models)

    R = roots.(domain(d), models)
    w = map(d.weight, R)
    w2 = [Float64.(getindex.(w, i)) for i in 1:length(w[1])]
    w2 = [w2;w2]

    P = [Float64.(getindex.(pl, i)) for i in 1:length(pl[1])]
    n = length(ql)

    function scostfun(λ)
        λ = softmax(λ)
        sum(abs2(dot(λ,Diagonal(P[j]),w2[j])/dot(λ,w2[j]) - ql[j]) for j in 1:n)
    end

    res = Optim.optimize(scostfun, zeros(N), LBFGS(m=20), Optim.Options(store_trace=false, show_trace=false, show_every=1, iterations=100, allow_f_increases=false, time_limit=100, x_tol=0, f_tol=0, g_tol=1e-8, f_calls_limit=0, g_calls_limit=0), autodiff=:forward)

    softmax(res.minimizer)
end

##
"Sum over j≠i. Internal function."
function ∑jni(X,i,S,k)
    s = zero(X[1][:,1])
    d = size(X[1],1)
    @inbounds for j in eachindex(X)
        j == i && continue
        for l in 1:d
            s[l] += X[j][l,S[j][k]]
        end
    end
    s
end

"""
    ISA(X, w = nothing; iters = 100, printerval = typemax(Int))

Iterative swapping algorithm from "On the Computation of Wasserstein barycenters", Giovanni Puccetti et al.

# Arguments:
- `X`: vector of d×k matrices where d is dimension and k number of atoms
- `w`: weights. See the files `test_barycenter.jl` for different uses.
- `iters`: maximum number of iterations
- `printerval`: print this often
"""
function ISA(X, w=nothing; iters=100, printerval = typemax(Int))
    n = length(X)
    d,k = size(X[1])

    if w !== nothing
        # X = deepcopy(X)
        X = map(eachindex(X)) do i # This does both copy and handles weird input types
            X[i] .* w[i] # This should work for both w[i] scalar and vector
        end
    end

    fudgefactor = 1.0

    σ = [collect(1:k) for _ in 1:n] # let σᵢ = Id, 1 ≤ i ≤ n.
    σ′ = deepcopy(σ)
    @inbounds for iter = 1:iters
        swaps = 0
        for i = 1:n
            σᵢ = σ[i]
            σᵢ′ = σ′[i]
            for k₁ = 1:k-1, k₂ = k₁+1:k
                Xik1 = @view X[i][:,σᵢ[k₁]]
                Xik2 = @view X[i][:,σᵢ[k₂]]
                if dot(Xik1, ∑jni(X,i,σ,k₁)) + dot(Xik2, ∑jni(X,i,σ,k₂)) < (dot(Xik2, ∑jni(X,i,σ,k₁)) + dot(Xik1, ∑jni(X,i,σ,k₂)))*fudgefactor
                    σᵢ′[k₁],σᵢ′[k₂] = σᵢ[k₂],σᵢ[k₁]
                    swaps += 1
                end
            end
        end
        iter % printerval == 0 && @show iter, swaps
        swaps == 0 && (return σ)
        σ = deepcopy(σ′) # Update assignment
        fudgefactor *= 1-1/iters^3
    end
    σ
end

## ======================
"""
    alg1(X, Y, â, b; β = 1, printerval = typemax(Int), tol = 1.0e-5, iters = 10000, solver = IPOT)

Algorithm 1 from ["Fast Computation of Wasserstein Barycenters"](https://arxiv.org/pdf/1310.4375.pdf) Notation is the same as in the paper.

# Arguments:
- `X`: Initial guess for barycenter support points
- `Y`: Support points for measures to calc barycenter of
- `a`: initial guess of barycenter weights
- `b`: Weigts of measures in Y
- `β`: Reg param, higher is more reg (inverse of λ in paper)
- `printerval`: DESCRIPTION
- `tol`: DESCRIPTION
- `iters`: DESCRIPTION
- `solver`: any of [`IPOT`](@ref) (default), [`sinkhorn`](@ref), [`sinkhorn_log`](@ref)
"""
function alg1(X,Y,â::AbstractVector{T},b;β=1, printerval=typemax(Int), tol=1e-5, iters=10000, solver=IPOT, weights=nothing)::Vector{T} where T
    N = length(Y)
    â = copy(â)
    a = copy(â)
    ã = copy(â)
    t0 = 1
    t = 0
    𝛂 = similar(a, length(a), N)
    Mth = [distmat_euclidean(X,Y[1]) for i in 1:Threads.nthreads()]
    for outer t = 1:iters
        B = (t0+t)/2
        a .= (1-inv(B)).*â .+ inv(B).*ã
        @sync for i in 1:N
            Threads.@spawn begin
                M = distmat_euclidean!(Mth[Threads.threadid()], X,Y[i])
                ai = solver(M,a,b[i]; iters=50000, β=β, tol=tol)[2]
                if !all(isfinite, a)
                    @warn "Got nan in inner sinkhorn alg 1, increasing precision"
                    ai = solver(M,big.(â),big.(b[i]); iters=50000, β=β, tol=tol)[2]
                    ai = eltype(â).(ai)
                end
                scale!(ai, i, weights)
                𝛂[:,i] .= ai
            end
        end

        ã .= ã .* exp.((-t0*B).*vec(mean(𝛂, dims=2)))
        ã ./= sum(ã)
        aerr = sum(abs2,â-ã)
        t % printerval == 0 && @info "Sinkhorn alg1:  iter: $t, aerr: $aerr"
        â .= (1-inv(B)).*â .+ inv(B).*ã
        â ./ sum(â)
        if aerr < tol
            t > printerval && @info "Sinkhorn alg1 done at iter $t"
            return â
        end
    end
    t > printerval && @info "Sinkhorn alg1 maximum number of iterations reached: $iters"
    â
end

"""
    alg2(X, Y, a, b;
            β          = 1/10,
            θ          = 0.5,
            printerval = typemax(Int),
            tol        = 1.0e-6,
            innertol   = 1.0e-5,
            iters      = 500,
            inneriters = 1000,
            atol       = 1.0e-32,
            solver     = IPOT,
            γ          = 1,
        )

Algorithm 2 from ["Fast Computation of Wasserstein Barycenters"](https://arxiv.org/pdf/1310.4375.pdf) Notation is the same as in the paper.

# Arguments:
- `X`: Initial guess for barycenter support points
- `Y`: Support points for measures to calc barycenter of
- `a`: initial guess of barycenter weights
- `b`: Weigts of measures in Y
- `β`: Reg param, higher is more reg
- `θ`: step size ∈ [0,1]
- `printerval`: print this often
- `tol`: outer tolerance
- `innertol`: inner tolerance
- `solver`: any of [`IPOT`](@ref) (default), [`sinkhorn`](@ref), [`sinkhorn_log`](@ref)
- `γ`: Sparsity parameter, if <1, encourage a uniform weight vector, if >1, do the opposite. Kind of like the inverse of α in the Dirichlet distribution.
"""
function alg2(X,Y,a,b; β = 1/10, θ = 0.5, printerval=typemax(Int), tol=1e-6, innertol=1e-4, iters=500, inneriters=10000, atol=1e-32, solver=IPOT, γ=0.0, weights=nothing, uniform=false)
    uniform || @warn("This function is known to be buggy when not enforcing uniform weights", maxlog=10)
    N  = length(Y)
    a  = copy(a)
    ao = copy(a)
    X  = copy(X)
    if weights !== nothing && eltype(weights) != eltype(X)
        X = convert(Matrix{eltype(weights)}, X)
    end
    Xo = copy(X)
    weights === nothing && (weights = fill(1/N, N))
    fill!(ao, 1/length(ao))
    for iter = 1:iters
        uniform || (a = alg1(X,Y,ao,b,β=β, printerval=printerval, tol=innertol, iters=inneriters, solver=solver, weights=weights))
        if γ > 0 && γ != 1
            a .= softmax(γ.*log.(a))
        else
            a = s1(a)
        end
        YT = sum(1:N) do i
            M = distmat_euclidean(X,Y[i])
            T,_ = solver(M,a,b[i]; iters=inneriters, β=β, tol=innertol)
            @assert !any(isnan, T) "Got nan in sinkhorn alg 2"
            scale!(Y[i]*T', i, weights)
        end

        X .= (1-θ).*X .+ θ.*(YT ./ (a' .+ eps()))
        aerr = mean(abs2, a-ao)
        xerr = mean(abs2, X-Xo)
        iter % printerval == 0 && @info "Sinkhorn alg2:  iter: $iter, aerr: $aerr, xerr: $xerr"
        if xerr < tol
            iter > printerval && @info "Sinkhorn alg2 done at iter $iter"
            return X,(a./=sum(a))
        end
        copyto!(ao,a)
        copyto!(Xo,X)
        ao ./= sum(ao)
        # θ *= 0.999
    end
    iters > printerval && @info "Sinkhorn alg2 maximum number of iterations reached: $iters"
    X,a
end

scale!(x,_,::Nothing) = x
scale!(x::AbstractArray{T},i,w::AbstractArray{T}) where T = (x .*= w[i])
scale!(x,i,w::AbstractArray) = (x * w[i]) # for dual numbers etc.

struct BCWorkspace{T,TK}
    K::TK
    KV::Array{T,3}
    U::Array{T,3}
    b::Array{T,2}
    bold::Array{T,2}
    S::Array{T,2}
end

function BCWorkspace(A, β)
    N    = length(A)
    m, n = size(A[1])
    T    = eltype(eltype(A))
    β    = T(β)

    b    = zeros(T, m, n)
    bold = zeros(T, m, n)
    U    = ones(T, m, n, N)
    KV   = ones(T, m, n, N)
    S1   = Matrix{T}(undef, m, n)
    S2   = Matrix{T}(undef, m, n)

    xi1 = Matrix{T}(undef, m, m)
    xi2 = Matrix{T}(undef, n, n)
    _initialize_conv_op!(xi1, xi2, β)

    function K(u,x)
        # xi1 * x * xi2
        mul!(S2,x,xi2)
        mul!(u,xi1,S2)
    end
    BCWorkspace(K, KV, U, b, bold, S1)
end

"""
    barycenter_convolutional(A, [λ]; β = 0.01, iters = 1000, tol = 1e-9, ϵ = 1e-90, verbose = false)

Convolutional barycenters.

´β` is the regularization and `λ` (optional) are the weights (barycentric coordinates). To reuse allocated space between successive calls, use the "workspace" method from the example below. `ϵ` is a truncation parameter for numerical stability.
```julia
a1      = zeros(10, 10)
a1[2,2] = 1
a2      = zeros(10, 10)
a2[6,6] = 1
A       = [a1,a2]
β       = 0.01
λ       = [0.5, 0.5] # Barycentric coordinates, must sum to 1
w       = BCWorkspace(A, β)
b       = barycenter_convolutional(w,A,λ)
plot(heatmap(a1), heatmap(a2), heatmap(b))
```

Ref: J. Solomon, F. de Goes, G. Peyré, M. Cuturi, A. Butscher, A. Nguyen, T. Du, L. Guibas. Convolutional Wasserstein Distances: Efficient Optimal Transportation on Geometric Domains.  2015 https://people.csail.mit.edu/jsolomon/assets/convolutional_w2.compressed.pdf
"""
function barycenter_convolutional(
    w::BCWorkspace{T,TK},
    A::AbstractVector{<:AbstractMatrix},
    λ::Union{Vector{<:Union{Float64,Float32}}, <: Fill} = Fill(1 / length(models), length(models));
    iters   = 1000,
    tol     = 1e-6,
    ϵ       = 1e-90,
    verbose = false,
) where {T,TK}

    length(A) < 2 && return A[]
    @fastmath sum(A[1]) ≈ sum(A[2]) || @warn "Input matrices do not appear to have the same mass (sum)"
    0.99 < sum(λ) < 1.01 || throw(ArgumentError("sum of barycentric coordinates λ was $(sum(λ)) but should be 1"))
    N = length(A)
    @unpack K, KV, U, b, bold, S = w
    iter = 0
    err = one(T)

    @views while err > tol && iter < iters
        copyto!(bold, b)
        iter = iter + 1
        b .= 0
        for r = 1:N # TODO: if S storage is expanded, this loop can be split into several and @avx or @tullio can be put on some. Maybe even K can be manually inlined here to make everythin avx-able. Can also be run on GPU easily.
            K(S, U[:, :, r])
            @avx S .= A[r] ./ max.(ϵ, S)
            K(KV[:, :, r], S)
            @avx @. b += λ[r] * log(max(ϵ, U[:, :, r] * KV[:, :, r]))
        end
        @avx b .= exp.(b)
        for r = 1:N
            @avx @. U[:, :, r] = b / max(ϵ, KV[:, :, r])
        end
        @avx b ./= sum(b)

        if iter % 10 == 1
            @fastmath err = sum(abs(bold - b) for (bold, b) in zip(bold, b))
            verbose && @info "Sinkhorn conv barycenters: iter = $iter, error = $err"
        end
    end
    b#, cost
end

# function barycenter_convolutional_diff(
#     A::AbstractVector{<:AbstractMatrix},
#     λ::AbstractVector{T};
#     β       = 0.01,
#     iters   = 1000,
#     tol     = 1e-6,
#     ϵ       = 1e-90,
#     verbose = false,
# ) where {T}
#
#     length(A) < 2 && return A[]
#     @fastmath sum(A[1]) ≈ sum(A[2]) ||
#               @warn "Input matrices do not appear to have the same mass (sum)"
#     0.99 < sum(λ) < 1.01 ||
#         throw(ArgumentError("sum of barycentric coordinates λ was $(sum(λ)) but should be 1"))
#
#     N    = length(A)
#     err  = one(T)
#     m, n = size(A[1])
#     b    = similar(λ, m, n) .= 0
#     # S  = similar(b)
#     U    = similar(λ, m, n, N) .= 1
#     KV   = similar(λ, m, n, N) .= 1
#
#     xi1  = Matrix{T}(undef, m, m)
#     xi2  = Matrix{T}(undef, n, n)
#     _initialize_conv_op!(xi1, xi2, β)
#
#
#
#     for i = 1:iters
#         b .= 0
#         for r = 1:N
#             S = xi1 * U[:, :, r] * xi2
#             # mul!(S,U,xi2)
#             # mul!(V,xi1,S) # TODO: this is by far the most expensive operation
#             @. S = A[r] / max(ϵ, S)
#             KV[:, :, r] = xi1 * S * xi2
#             # mul!(S,V,xi2)
#             # mul!(U,xi1,S)
#             @. b += λ[r] * log(max(ϵ, U[:, :, r] * KV[:, :, r]))
#         end
#         @. b = exp(b)
#         for r = 1:N
#             @. U[:, :, r] = b / max(ϵ, KV[:, :, r])
#         end
#
#     end
#     b
# end

function barycenter_convolutional(
    A::AbstractVector{<:AbstractMatrix},
    λ = Fill(1 / length(A), length(A));
    β = 0.001,
    kwargs...,
)
    w = BCWorkspace(A, β)
    barycenter_convolutional(w, A, λ; kwargs...)
end

"""
    barycenter_convolutional(models::Vector{<:DSP.Periodograms.TFR}, λ = Fill(1 / length(models), length(models)); dynamic_floor = default_dynamic_floor(models), kwargs...)

Covenience function for the calculation of spectrograms. This function transforms the spectrograms to log-power and adjusts the floor to `dynamic_floor`, followed by a normalization to sum to 1.

This function will be called if [`barycenter`](@ref) is called with [`ConvOptimalTransportDistance`](@ref) as first argument.

# Arguments:
- `dynamic_floor`: Sets the floor of the spectrogram in log-domain, i.e., all values below this will be truncated. The default value is based on a quantile of the spectrogram powers. If your spectrograms are mostly low entropy, you can try to increase this number to get sharper results.
- `kwargs`: Same as for the base method

# Example:
```julia
using SpectralDistances, DSP, Plots
N     = 24_000
t     = 1:N
f     = range(0.8, stop=1.2, length=N)
y1    = sin.(t .* f) .+ 0.1 .* randn.()
y2    = sin.(t .* reverse(f .+ 0.5)) .+ 0.1 .* randn.()
S1,S2 = spectrogram.((y1,y2), 1024)

A = [S1,S2]
β = 0.0001     # Regularization parameter (higher implies more smoothing and a faster, more stable solution)
λ = [0.5, 0.5] # Barycentric coordinates (must sum to 1)
B = barycenter_convolutional(A, β=β, tol=1e-6, iters=200, ϵ=1e-100, dynamic_floor=-2)
plot(plot(S1, title="S1"), plot(B, title="Barycenter"), plot(S2, title="S2"), layout=(1,3), colorbar=false)
```
"""
function barycenter_convolutional(
    models::Vector{<:DSP.Periodograms.TFR},
    λ = Fill(1 / length(models), length(models));
    dynamic_floor = default_dynamic_floor(models),
    kwargs...,
)

    A  = normalize_spectrogram.(models, dynamic_floor)
    ms = mean(sum, A)
    sumnorm = x -> x .* (ms/sum(x))
    A .= sumnorm.(A)
    b  = barycenter_convolutional(A, λ; kwargs...)
    B  = deepcopy(models[1])
    Bp = power(B)
    Bp .= exp.(b .+ dynamic_floor)
    B
end

function barycenter(d::ConvOptimalTransportDistance, args...; kwargs...)
    @assert d.invariant_axis == 0 "Invariant axis not supported"
    barycenter_convolutional( args...; β = d.β, kwargs...)
end

function barycenter(d::ConvOptimalTransportDistance, A::Vector{<:DSP.Periodograms.TFR}, args...; kwargs...)
    @assert d.invariant_axis == 0 "Invariant axis not supported"
    barycenter_convolutional(
        A, args...;
        dynamic_floor = isnan(d.dynamic_floor) ? default_dynamic_floor(A) : d.dynamic_floor,
        β = d.β,
        kwargs...,
    )
end

# function barycenter(d::OptimalTransportRootDistance,models; normalize=true, kwargs...)
#     X, w, realpoles = barycenter_matrices(d, models, normalize)
#     S = ISA(X; kwargs...)
#     bc = barycentric_weighting(X,w,S)
#
#     r1 = complex.(bc[1,:],bc[2,:])
#     if realpoles
#         bcr = ContinuousRoots(r1)
#     else
#         @assert !any(iszero ∘ imag, r1) "A real root was found in barycenter even though inputs had no real roots"
#         bcr = ContinuousRoots([r1; conj.(r1)])
#     end
#     AR(bcr)
# end
# function barycenter(d::OptimalTransportRootDistance,models,w; kwargs...)
#     r = roots.(SpectralDistances.Continuous(), models)
#     realpoles = any(any(iszero ∘ imag, r) for r in r)
#
#     if !realpoles
#         r = [r[1:end÷2] for r in r]
#     end
#     X = [[real(r)'; imag(r)'] for r in r]
#
#     S = ISA(X, w; kwargs...)
#     X̂ = [X[i][:,S[i]] for i in eachindex(X)]
#     bc = sum(w.*X̂)
#     r1 = complex.(bc[1,:],bc[2,:])
#     if realpoles
#         bcr = ContinuousRoots(r1)
#     else
#         @assert !any(iszero ∘ imag, r1) "A real root was found in barycenter even though inputs had no real roots"
#         bcr = ContinuousRoots([r1; conj.(r1)])
#     end
#
#     AR(bcr)
# end
# the method below does not work very well due to likely bug in ISA
# function barycenter(X::Vector{<:AbstractArray}, λ; iters=100, kwargs...)
#     sw = ISA(X, λ; iters=iters, kwargs...)
#     barycentric_weighting(X,λ,sw)
# end
# function barycentric_weighting2(X,λ,sw)
#     sum(λ[sw[i]]'.*X[i][:,sw[i]] for i in eachindex(sw))
# end # This one was definitely bad
## ==========================================================
# function simplex_ls(ⅅ, ql; iters=2000, verbose=false, kwargs...)
#
#     N = length(ql)
#     α = 1.0
#     ε = 0.1/N
#     @show λ = max.(ⅅ\ql .+ 0.01 .* randn.(), 0)
#     λ ./= sum(λ)
#     λo = copy(λ)
#     g = similar(λ)
#     local ng, err
#     err = 0.0
#     for iter = 1:iters # Do a few projected gradient iterations with soft thresholding
#         g .= ⅅ'*(ql-ⅅ*λ) # This is the negative gradient so + below
#         ng = norm(g)
#         ng < 1e-10 && break
#         λ .+= α.*g
#         λ .-= minimum(λ)
#         λ .= soft_th.(λ, ε)
#         λ ./= sum(λ)
#         α *= 0.999
#         err = norm(λ-λo)
#         verbose && @info "Iter $iter norm(g): $ng norm(λ-λo): $err"
#         err < 1e-5 && break
#         λo .= λ
#     end
#     verbose &&  @info "Converged norm(g): $ng norm(λ-λo): $err"
#     λ
# end

struct BCCWorkspace{T}
    w::Vector{T}
    b::Vector{Array{T,3}}
    r::Array{T,3}
    φ::Vector{Array{T,3}}
    C::Matrix{T}
    C2::Matrix{T}
    C3::Matrix{T}
    C4::Matrix{T}
    C5::Matrix{T}
    S2::Matrix{T}
    xi1::Matrix{T}
    xi2::Matrix{T}
    scw::SCWorkspace{T}
end

"""
    BCCWorkspace(X::Vector{<:AbstractMatrix{T}}, L, β) where T

Create a workspace cache for [`barycentric_coordinates`](@ref) with the convolutional distance.

# Arguments:
- `X`: Input matrices
- `L`: Number of iterations
- `β`: Regularization factor
"""
function BCCWorkspace(p::Vector{<:AbstractMatrix{T}}, L, β) where T
    N   = length(p)
    S   = length(p)
    m,n = size(p[1])
    w   = zeros(T,S)
    b   = [fill(1/N, m, n, S) for _ in 0:L+1]
    r   = zeros(T,m,n,S)
    φ   = [Array{T}(undef,m,n,S) for _ in 1:L]
    C   = zeros(T,m,n)
    C2  = zeros(T,m,n)
    C3  = zeros(T,m,n)
    C4  = zeros(T,m,n)
    C5  = zeros(T,m,n)
    S2  = zeros(T,m,n)
    scw = SCWorkspace(p[1], p[1], β)

    xi1 = Matrix{T}(undef, m, m)
    xi2 = Matrix{T}(undef, n, n)
    _initialize_conv_op!(xi1, xi2, β)
    BCCWorkspace{T}(w,b,r,φ,C,C2,C3,C4,C5,S2,xi1,xi2,scw)
end

"""
    cost, B, grad = sinkhorn_convolutional_diff(workspace::BCCWorkspace{T}, X, q::AbstractMatrix{T}, λ::AbstractVector; β = 0.01) where T

Returns cost, barycenter and gradient of `λ`. Called from within [`barycentric_coordinates`](@ref)`(
    d::ConvOptimalTransportDistance,
    X::Vector{<:AbstractMatrix},
    q::AbstractMatrix)`

# Arguments:
- `workspace`: See [`BCCWorkspace`](@ref)
- `X`: Input matrices (anchors)
- `q`: Query matrix
- `λ`: barycentric coordinates
- `β`: reg parameter
"""
function sinkhorn_convolutional_diff(workspace::BCCWorkspace{T}, p, q::AbstractMatrix{T}, λ::AbstractVector; β=0.01) where T

    @unpack w,b,r,φ,C,C2,C3,C4,C5,S2,xi1,xi2,scw = workspace
    w .= 0
    r .= 0

    L = length(φ)
    N = length(p)
    S = length(p)
    m,n = size(p[1])

    function K(u,x)
        # xi1 * x * xi2
        mul!(S2,x,xi2)
        mul!(u,xi1,S2)
    end

    local P
    @views for l = 1:L
        for s in 1:S
            K(C,b[l][:,:,s])
            @avx @. C = p[s] / C
            K(φ[l][:,:,s], C)
        end
        P = dropdims(prod(φ[l].^reshape(λ,1,1,:), dims=3), dims=3)
        @avx b[l+1] .= P ./ φ[l]
    end

    cost,a,_ = sinkhorn_convolutional(scw, P, q; β=β)

    @avx ∇W = @. (a = β * a)
    # ∇W = bb
    @avx g = (∇W .= ∇W .* P)

    for l = L:-1:1
        @views for s in 1:S
            @avx S2 .= log.(φ[l][:,:,s])
            w[s] += dot(S2, g)
            K(C,b[l][:,:,s])
            @avx @. C = abs2(C)
            @avx @. C4 = (λ[s]* g - r[:,:,s]) / φ[l][:,:,s]
            K(C2,C4)
            @avx @. C2 = C2 * p[s] / C
            K(C3,C2)
            @avx @. r[:,:,s] = -C3 * b[l][:,:,s]
        end
        g = dropdims(sum(r, dims=3), dims=3)
    end
    cost, P,w
end

"""
    sinkhorn_convolutional_diff(p::Vector, q::AbstractMatrix, λ::AbstractVector; β = 0.01, L = 32, kwargs...)
"""
function sinkhorn_convolutional_diff(p::Vector, q::AbstractMatrix, λ::AbstractVector; β=0.01, L = 32, kwargs...)
    w = BCCWorkspace(p, L, β)
    sinkhorn_convolutional_diff(w, p, q, λ; β=β, kwargs...)
end

"""
    barycentric_coordinates(d::ConvOptimalTransportDistance, X::Vector{<:AbstractMatrix}, q::AbstractMatrix; method = LBFGS(), kwargs...)

Calculate the barycentric coordinates of a vector of matrices `X` using the convolutional method.

# Arguments:
- `q`: Query matrix
- `method`: The optimizer from Optim
- `kwargs`: Are sent to [`sinkhorn_convolutional`](@ref)

## Optim options
The default options are
```
options = Optim.Options(
        store_trace       = true,
        show_trace        = false,
        show_every        = 1,
        iterations        = 10,
        allow_f_increases = false,
        time_limit        = 150,
        x_tol             = 1e-3,
        f_tol             = 1e-3,
        g_tol             = 1e-4,
    )
```
"""
function barycentric_coordinates(
    d::ConvOptimalTransportDistance,
    X::Vector{<:AbstractMatrix},
    q::AbstractMatrix;
    L = 40,
    method  = LBFGS(),
    options = Optim.Options(
        store_trace       = true,
        show_trace        = false,
        show_every        = 1,
        iterations        = 10,
        allow_f_increases = false,
        time_limit        = 150,
        x_tol             = 1e-3,
        f_tol             = 1e-3,
        g_tol             = 1e-4,
    ),
    kwargs...,
)
    workspace = BCCWorkspace(X, L, d.β)
    @assert d.invariant_axis == 0 "Invariant axis not supported"

    function fg!(F, G, λ)
        λ = softmax(λ)
        cost, B, g = sinkhorn_convolutional_diff(workspace, X, q, λ; β=d.β)

        if G != nothing
            G .= λ .* (g .- dot(g, λ)) # Chain rule for softmax
            # G .= g
        end
        return cost
    end

    λ = zeros(length(X))

    res = Optim.optimize(Optim.only_fg!(fg!), λ, method, options)
    softmax(res.minimizer), res
end


include("utils_dg.jl")

K = 7  # Number of elements

n = 4
nn = 1
r = jacobiGL(0, 0, n)
V = vandermonde(r, 0, 0, n)
fr = collect(range(-1,1,length=100))
fV = vandermonde(fr, 0, 0, n)
refine = fV * inv(V)
# cell average
avg = Diagonal(zeros(n+1))
avg[1:nn] .= 1
p0 = V * avg * inv(V)
# polyone
nn = 2
tmp = zeros(n+1)
tmp[1:nn] .= 1
avg = Diagonal(tmp)
p1 = V * avg * inv(V)
# poly two
nn = 2+1
tmp = zeros(n+1)
tmp[1:nn] .= 1
avg = Diagonal(tmp)
p2 = V * avg * inv(V)

# poly 6
Ω = (; a=0, b=2π, periodic = true)
a = Ω.a
b = Ω.b
mesh = Mesh(K, n, Ω.a, Ω.b, periodic = Ω.periodic)
x = mesh.x
weights = sum(mesh.M, dims = 1)
Δx = x[2] - x[1]

# inexact integration
DM = Diagonal(sum(mesh.M, dims = 1)[:])
mesh.M .= DM
mesh.Mi .= inv(DM)
mesh.lift[:,1] .= mesh.Mi[1,:]
mesh.lift[:,end] .= mesh.Mi[end,:]

b = Ω.b
a = Ω.a

u_first  = @. sin(2π/(b-a) * x) *( 1+ 0.9 * sin(8 * 2π/(b-a) * x) )
using Random
Random.seed!(123)
uo = refine * u_first
x = refine * x

Duu  = refine * compute_volume_terms(u_first .* u_first, mesh)
uDu  = 2.0 .* refine * ( u_first .* compute_volume_terms(u_first, mesh) )
ESDu = 2/3 .* Duu + 1/3 .* uDu

## Plotting
fig = Figure(resolution = (1700+600, 1000+400))

lims = (extrema(x)..., extrema(uo)...)
axo = Axis(fig[1,1], xlabel = L"x", ylabel = L"y", ylabelsize = 22,
    xlabelsize= 22, xgridstyle=:dash, ygridstyle=:dash, xtickalign = 1,
    xticksize=10, ytickalign=1, yticksize=10,  xlabelpadding = -10)

axo.limits = (extrema(x)..., (extrema(uo) .* 1.1)...)
axo.title = "Function"
axo.titlesize = 32
for i in 1:K
    lines!(axo, x[:,i], uo[:,i], linewidth = 5)
end

ax0 = Axis(fig[1,2], xlabel = L"x", ylabel = L"y", ylabelsize = 22,
    xlabelsize= 22, xgridstyle=:dash, ygridstyle=:dash, xtickalign = 1,
    xticksize=10, ytickalign=1, yticksize=10,  xlabelpadding = -10, ylims = (-1.1, 1.1))
lims2 = (extrema(x)..., (extrema(Duu) .* 1.1)...)
ax0.limits = lims2
ax0.title = "Conservative "
ax0.titlesize = 32
for i in 1:K
    lines!(ax0, x[:,i], Duu[:,i], linewidth = 5)
end

ax1 = Axis(fig[2,1], xlabel = L"x", ylabel = L"y", ylabelsize = 22,
    xlabelsize= 22, xgridstyle=:dash, ygridstyle=:dash, xtickalign = 1,
    xticksize=10, ytickalign=1, yticksize=10,  xlabelpadding = -10, ylims = (-1.1, 1.1))
ax1.limits = lims2
ax1.title = "Advective"
ax1.titlesize = 32
for i in 1:K
    lines!(ax1, x[:,i], uDu[:,i], linewidth = 5)
end

ax3 = Axis(fig[2,2], xlabel = L"x", ylabel = L"y", ylabelsize = 22,
    xlabelsize= 22, xgridstyle=:dash, ygridstyle=:dash, xtickalign = 1,
    xticksize=10, ytickalign=1, yticksize=10,  xlabelpadding = -10, ylims = (-1.1, 1.1))
ax3.limits = lims2
ax3.title = "Entropy Stable"
ax3.titlesize = 32
for i in 1:K
    lines!(ax3, x[:,i], ESDu[:,i], linewidth = 5)
end


using RandomizedLinAlg

include("factorization.jl")
include("rlinalg.jl")
include("rsvd.jl")
include("rsvd_fnkz.jl")


#========================================================================================#
"""
	NBodies

A system of N interacting bodies moving in dim dimensions over time T in res timesteps.

Author: Niall Palfreyman, 24/3/2022.
"""
module NBodies

# Externally callable methods of NBodies
export NBody

#-----------------------------------------------------------------------------------------
# Module types:

"""
	NBody

An NBody system capable of containing multiple (N) bodies that gravitationally interact
with each other.
"""
struct NBody
	N							# Number of bodies
	T							# Duration of simulation
	res							# Timestep resolution
	xi							# Initial positions of bodies
	pi							# Initial momenta of bodies
	m							# Masses of bodies

	"Construct a new NBody"
	function NBody( T=40, resolution=20000)
		# Initialise all fields of the decoding apparatus:
		new(
			0,					# Initially no bodies in the system
			T,					# Duration
			resolution,			# Timestep resolution
			zeros(1,N),			# Initial positions
			zeros(1,N),			# Initial momenta
			zeros(1,N),			# Masses of bodies
		)
	end
end

#-----------------------------------------------------------------------------------------
# Module methods:

"""
	addbody( nbody::NBody, xi, pi, m=1)

Add to the system a new body with initial position and momentum xi, pi, and with mass m.
"""
function addbody( nbody, xi, pi, m=1)
	nbody.xi = [nbody.xi xi]
	nbody.pi = [nbody.pi pi]
	nbody.m  = [nbody.m m]
	nbody.N  += 1
end

"""
	simulate( nbody::NBody)

Run a simulation of the given NBody system over duration divided into res timesteps.
"""
function simulate( nbody)
	dt = nbody.T / nbody.res
	t = 0:dt:nbody.T
	x = zeros(1,nbody.N,nbody.res+1)
	p = zeros(1,nbody.N,nbody.res+1)
	G = 1
	Gmm = repmat(G*nbody.m,[1 1 nbody.N])
	Gmm = Gmm .* permute(Gmm,[1 3 2])

	# Relative positions of bodies at positions r:
	relPos = @(r) repmat(permute(r,[1 3 2]),[1 nbody.N 1]) - repmat(r,[1 1 nbody.N])

	# Initialisation:
	x[1,:,1] = nbody.xi
	p[1,:,1] = nbody.pi

	# Simulation:
	for n = 1:nbody.res
		# Simulate timestep n:
		x[1,:,n+1] = x[1,:,n] + dt * p[1,:,n]./nbody.m

		diff = relPos(x(1,:,n))
		absDiff = sqrt(sum(diff.*diff,1))
		invSq = absDiff.^3 + permute(eye(nbody.N),[3 2 1])
		forces = Gmm.*diff ./ invSq
		forces = sum(forces,3)
		p[1,:,n+1] = p[1,:,n] + forces * dt
	end

	(t,x,p)
end

#-----------------------------------------------------------------------------------------
"""
	demo()

Demonstrate simulation of a simple 2-body problem.
"""
function demo()
	nb = NBody( 2.25, 10000)
	nb.addbody( -1, 0)
	nb.addbody( 1, 0)

	[t,x] = nb.simulate()
	plot( ...
		t,squeeze(x(1,1,:)),'b', ...
		t,squeeze(x(1,2,:)),'r' ...
	);
end

end		# of NBodies


using IPIL8

# --------------------------------------------------------------------------------------------------
@testset "Тесты на регрессию прямой задачи                  " begin

    # Получим стандартные аргументы
    eval(IPIL8.default_parameters());

    # Загрузим зафиксированные данные регрессионных тестов
    include("direct_regression_tests_data.jl");

    # Тестируем внутреннюю функцию, поэтому входные и выходные векторы размерности N-2;
    y = u₀[2:end-1];
    RP = zero(y);

# --------------------------------------------------------------------------------------------------
@testset "Вектор правой части                               " begin

    IPIL8.directRP!(RP, y, 1, h, N, ε, ulₘ, urₘ);

    @test(isapprox(known_RP, RP));

end # testset Вектор правой части
# --------------------------------------------------------------------------------------------------

# --------------------------------------------------------------------------------------------------
@testset "Три диагонали якобиана правой части               " begin

    dl = zeros(N-2);
    dd = zeros(N-1);
    du = zeros(N-2);

    IPIL8.DRP!(dl, dd, du, y, 1, h, N, ε, ulₘ, urₘ)

    @test(isapprox(known_dl, dl));
    @test(isapprox(known_dd, dd));
    @test(isapprox(known_du, du));

end
# --------------------------------------------------------------------------------------------------

# --------------------------------------------------------------------------------------------------
@testset "Решение прямой задачи                             " begin

    solution = zeros(N+1, M+1);

    IPIL8.solve!(solution, u₀, N, h, Tₘ, M, τ, ε, ulₘ, urₘ);

    @test(isapprox(known_u, solution));

end
# --------------------------------------------------------------------------------------------------

end # testset Регрессия прямой задачи
# --------------------------------------------------------------------------------------------------


using ChebyshevQuantum
using Test

tol_val = 1e-11

deriv=:forwarddiff

@testset "find edge function " begin

    #step function
    function f(x)
        if x < 0.2
            return 0.0
        else
            return 1.0
        end
    end

    c, stddev = ChebyshevQuantum.EdgeDetect.find_edges(f, deriv=deriv)

    @test isapprox(c[1,1], 0.2,  atol=1e-4)
    @test isapprox(c[1,2], 0.2,  atol=1e-4)

end

@testset "find edge f, f' " begin

    #two edges
    function f(x)
        if x < 0.3
            return 1.0 * abs(x-0.1) + x^3
        else
            return 1.0 * abs(x-0.1) + x^3 - 1
        end
    end
    #f(0.1)
    #f.([0.1,0.2])
    c, contin = ChebyshevQuantum.EdgeDetect.find(f, deriv=deriv)

    @test isapprox(c[1,1], 0.1,  atol=1e-4)
    @test isapprox(c[1,2], 0.1,  atol=1e-4)

    @test isapprox(c[2,1], 0.3,  atol=1e-4)
    @test isapprox(c[2,2], 0.3,  atol=1e-4)

    @test contin[1] == true
    @test contin[2] == false

end

@testset "find edge singularity " begin

    #singularity
    function f(x)
        return 2.0 * sqrt(abs(x-0.1))
    end

    c, contin = ChebyshevQuantum.EdgeDetect.find(f)

    @test isapprox(c[1,1], 0.1,  atol=1e-4)
    @test isapprox(c[1,2], 0.1,  atol=1e-4)

    @test contin[1] == true

    #singularity
    function f2(x)
        return -max(1.0 / (x-0.1), -1e10)
    end

    c2, contin2 = ChebyshevQuantum.EdgeDetect.find(f2)

    @test isapprox(c2[1,1], 0.1,  atol=1e-4)
    @test isapprox(c2[1,2], 0.1,  atol=1e-4)

    @test contin2[1] == false

end

@testset "don't find  " begin

    #no edges
    function f(x)
        return sin(30*x)
    end

    c, contin = ChebyshevQuantum.EdgeDetect.find(f)

    @test size(c)[1] == 0

    #no edges
    function f2(x)
        return exp(-1/x^2)
    end

    c2, contin2 = ChebyshevQuantum.EdgeDetect.find(f2)

    @test size(c2)[1] == 0

    #edge at end, don't find
    function f3(x)
        return max(-1/x, -1e10)
    end

    c3, contin3 = ChebyshevQuantum.EdgeDetect.find(f3, a=0, b =1)

    @test size(c3)[1] == 0

end


# Module Definition
# #+HTML: Source
# #+HTML:

# [[file:~/.julia/dev/MatrixProductStates/README.org::*Module%20Definition][Module Definition:1]]
module MatrixProductStates

using LinearAlgebra, TensorOperations, TensorCast, LowRankApprox, Arpack, Strided, SparseArrays
#using ProgressMeter

export *, /, ==, ≈, isequal, adjoint, getindex, randn
export MPS, MPO, left, right, compress, imag_time_evolution, real_time_evolution, real_timedep_time_evolution, rightcanonical, leftcanonical
export ground_state, two_point_correlator, realize

include("utils.jl")
include("MPS.jl")
include("MPO.jl")
include("compression.jl")
include("contraction.jl")
include("groundstate.jl")
include("correlation.jl")
include("timeevolution.jl")

end
# Module Definition:1 ends here


#= make_scenarios_rejection_sample.jl
make timeseries of DR performance with a mean performace near 100%
Give a set of outcomes that each have an associated probability
And information on how often to resample from those outcomes

inputs:
n of total timesteps
prob distribution of performance and outcomes, vdr, p (original)
number of timesteps in which DR performance is constant
should scenarios be randomly selected
how many scenarios are selected?
threshold for rejection

outputs:
vdr[t,o]
p[o]
=#
using StatsBase
function make_scenarios(n_timesteps,v_og,p_og,int_length; randsel = true, nrand = 5, thresh = 0.005)
    # determine total number of possible scenarios
    Tp = convert(Int64,floor(n_timesteps/int_length)) # number of time blocks
    n_omega = length(p_og) # number of possilbe outcomes
    n_new_scenarios = convert(Int64,n_omega^Tp)

    # make vdr placeholder and all_s
    vdr = fill(0.0,(n_timesteps,nrand))
    p = fill(0.0,(nrand))
    all_s = collect(1:n_new_scenarios)

    # rejection sample
    i=1;
    n_sampled = 0
    while i<n_new_scenarios && n_sampled < nrand
        # draw a single sample
        s = sample(all_s,1)
        println("sampled ",s)

        # make vdr data for that sample
        s_vdr = fill(0.0,(n_timesteps))
        s_p = fill(0.0,(n_timesteps))
        for j in 1:Tp #row chunk
            # identify which rows belong to this time period
            if j == Tp
                last_period = n_timesteps
            else
                last_period = int_length*j #t_firsts[i+1]-1
            end
            rows = (int_length*(j-1) + 1):last_period

            # sample correct performance for that time period
            # imagine all possible combos are listed out in an array that is Tp x n_new_scenarios
            ind = mod(fld(s[1]-1,n_omega^(Tp-j)),n_omega)+1
                # fld: largest integer <= x/y
                # n_omega^(Tp-j) is the number of row-wise 'blocks' of the same response in our imagined array
                #           on the last block, there is just one block (j=Tp, this = 1)
                #           on the first block,   (j=1, this = n_omega^(Tp-1), e.g. n of blocks if every n_omega blocks were grouped)
                # divide sample number by above;
                # find remainer of that with number of possible responses
                # add 1 since above is 0-indexed
            s_vdr[rows] = v_og[ind]
            s_p[rows] = p_og[ind]
        end

        # assess if s_vdr is ok
        if abs(1-mean(s_vdr,weights(s_p))) < thresh #mean is weighted by probability
            # add to vdr and p
            vdr[:,n_sampled+1] = s_vdr
            p[n_sampled+1] = prod(s_p)
            # update n_sampled
            n_sampled = n_sampled+1
            println("added sample ",s)
        end

        # remove from possible scenarios to sample
        deleteat!(all_s,find(all_s .== s))

        i=i+1;
    end
    return(vdr,p)
end


module MyModel

using Mimi

include("region_parameters.jl")
include("gross_economy.jl")
include("emissions.jl")

export construct_MyModel

function construct_MyModel()

	m = Model()

	set_dimension!(m, :time, collect(2015:5:2110))
	set_dimension!(m, :regions, [:Region1, :Region2, :Region3])	 # Note that the regions of your model must be specified here

	add_comp!(m, grosseconomy)
	add_comp!(m, emissions)

	# set parameters for grosseconomy component
	update_param!(m, :grosseconomy, :l, l)
	update_param!(m, :grosseconomy, :tfp, tfp)
	update_param!(m, :grosseconomy, :s, s)
	update_param!(m, :grosseconomy, :depk,depk)
	update_param!(m, :grosseconomy, :k0, k0)
	update_param!(m, :grosseconomy, :share, 0.3)

	# set and connect parameters for emissions component
	update_param!(m, :emissions, :sigma, sigma)
	connect_param!(m, :emissions, :YGROSS, :grosseconomy, :YGROSS)

    return m

end

end #module


"""
    A = JopGradient(dom, δ)

Gradient for a 2D or 3D arrays. The range will have one more
dimension than the domain and contain the components of the gradient
in each dimension.

* `dom::JetSpace{T,N}` is the domain of the operator.
* `δ::NTuple{T,N}` is the grid spacing in each dimension.

# Examples:

## 2D
```julia
nz,nx = 11,12
dom = JetSpace(Float32, nz, nx)
A = JopGradient(dom, (1.0,1.0))
size(domain(A)) # (11,12)
size(range(A))  # (11,12,2)
```

## 3D
```julia
nz,ny,nx = 11,12,13
dom = JetSpace(Float32, nz, ny, nx)
A = JopGradient(dom,(1.0,1.0,1.0))
size(domain(A)) # (11,12,13)
size(range(A))  # (11,12,13,3)
```
"""
function JopGradient(dom::JetSpace{T,N}, δ::NTuple{N,T}) where {T,N}
    rng = JetSpace(T, (size(dom)...,length(size(dom))))
    if length(size(dom)) != length(δ)
        error("dimensionality of length(size(dom)) ($(length(size(dom)))) and length(δ) ($(length(δ))) do not agree!")
    end
    JopLn(dom = dom, rng = rng, df! = JopGradient_df!, df′! = JopGradient_df′!, s = (δ=δ,))
end

export JopGradient

# ....................................................................................
# 2D
# ....................................................................................
function JopGradient_df!(d::AbstractArray{T,3}, m::AbstractArray{T,2}; δ, kwargs...) where {T}
    d .= 0
    n1,n2 = size(m)
    for k1=1:n1-1, k2=1:n2
        d[k1,k2,1] += m[k1+1,k2] / δ[1]
        d[k1,k2,1] -= m[k1+0,k2] / δ[1]
    end
    for k1=1:n1, k2=1:n2-1
        d[k1,k2,2] += m[k1,k2+1] / δ[2]
        d[k1,k2,2] -= m[k1,k2+0] / δ[2]
    end
    d
end

function JopGradient_df′!(m::AbstractArray{T,2}, d::AbstractArray{T,3}; δ, kwargs...) where {T}
    m .= 0
    n1,n2 = size(m)
    for k1=1:n1-1, k2=1:n2
        m[k1+1,k2] += d[k1,k2,1] / δ[1]
        m[k1+0,k2] -= d[k1,k2,1] / δ[1]
    end
    for k1=1:n1, k2=1:n2-1
        m[k1,k2+1] += d[k1,k2,2] / δ[2]
        m[k1,k2+0] -= d[k1,k2,2] / δ[2]
    end
    m
end

# ....................................................................................
# 3D
# ....................................................................................
function JopGradient_df!(d::AbstractArray{T,4}, m::AbstractArray{T,3}; δ, kwargs...) where {T}
    d .= 0
    n1,n2,n3 = size(d)
    for k1=1:n1-1, k2=1:n2, k3=1:n3
        d[k1,k2,k3,1] += m[k1+1,k2,k3] / δ[1]
        d[k1,k2,k3,1] -= m[k1+0,k2,k3] / δ[1]
    end
    for k1=1:n1, k2=1:n2-1, k3=1:n3
        d[k1,k2,k3,2] += m[k1,k2+1,k3] / δ[2]
        d[k1,k2,k3,2] -= m[k1,k2+0,k3] / δ[2]
    end
    for k1=1:n1, k2=1:n2, k3=1:n3-1
        d[k1,k2,k3,3] += m[k1,k2,k3+1] / δ[3]
        d[k1,k2,k3,3] -= m[k1,k2,k3+0] / δ[3]
    end
    d
end

function JopGradient_df′!(m::AbstractArray{T,3}, d::AbstractArray{T,4}; δ, kwargs...) where {T}
    m .= 0
    n1,n2,n3 = size(m)
    for k1=1:n1-1, k2=1:n2, k3=1:n3
        m[k1+1,k2,k3] += d[k1,k2,k3,1] / δ[1]
        m[k1+0,k2,k3] -= d[k1,k2,k3,1] / δ[1]
    end
    for k1=1:n1, k2=1:n2-1, k3=1:n3
        m[k1,k2+1,k3] += d[k1,k2,k3,2] / δ[2]
        m[k1,k2+0,k3] -= d[k1,k2,k3,2] / δ[2]
    end
    for k1=1:n1, k2=1:n2, k3=1:n3-1
        m[k1,k2,k3+1] += d[k1,k2,k3,3] / δ[3]
        m[k1,k2,k3+0] -= d[k1,k2,k3,3] / δ[3]
    end
    m
end


using Test

include("loaddata.jl")

@testset "encode by farthest points" begin
    using KernelMethods
    using SimilaritySearch
    using StatsBase: mean

    X, ylabels = loadiris()
    dist = lp_distance(3.3)
    L = Float64[]
    for label in ["Iris-setosa", "Iris-versicolor", "Iris-virginica"]
        y = ylabels .== label
        A = X[y]
        B = X[.~y]
        occ = fit(OneClassClassifier, dist, A, 21)
        ypred = [predict(occ, dist, x).similarity > 0 for x in X]
        push!(L, mean(ypred .== y))
        println(stderr, "==> $label: $(L[end])")
    end

    macrorecall = mean(L)
    println(stderr, "===> macro-recall: $macrorecall")
    @test macrorecall > 0.9
end


"""
Updates the accumulated cost.
"""
function update_cost(acc_cost::Float64, new_cost::Union{Int64,Float64},
                     reduction_type::Symbol)
    if reduction_type == :sum
        return acc_cost + new_cost
    elseif reduction_type == :max
        return max(acc_cost, new_cost)
    end
    throw("Unsupported cost reduction strategy: $(reduction_type)")
end

"""
Numerically stable softmax.
"""
function softmax(x, θ::AbstractFloat=1.0, dim::Integer=1)
    x_exps = exp.((x .- maximum(x)) * θ)
    return x_exps ./ sum(x_exps, dims=dim)
end

"""
Various strategies for computing action selection probabilities.

TODO(kykim): Consider putting these out into a separate jl.
"""
function ucb_probs(ucb_scores)
    return ucb_scores /= sum(ucb_scores)
end

function ucb_softmax_probs(ucb_scores)
    return softmax(ucb_scores, θ=1.0)
end

function expected_cost_probs(est_α_probs, est_α_costs)
    exp_α_costs = est_α_probs .* est_α_costs
    return exp_α_costs /= sum(exp_α_costs)
end

function adaptive_probs(est_α_probs, est_α_costs, nominal_probs, β, γ)
    max_α_cost = maximum(est_α_costs)
    max_α_prob = maximum(est_α_probs)
    max_nominal_prob = maximum(nominal_probs)

    cvar_strategy = est_α_costs .* nominal_probs .+ (max_α_cost / 20) .+ eps()
    cdf_strategy = est_α_probs .* nominal_probs .+ (max_α_prob * max_nominal_prob / 20) .+ eps()
    cvar_strategy /= sum(cvar_strategy)
    cdf_strategy /= sum(cdf_strategy)

    # Mixture weighting.
    a_probs = β * nominal_probs .+ γ * cdf_strategy .+ (1 - β - γ) * cvar_strategy
    return a_probs
end


using Revise
using JuMP, EAGO

m = Model(optimizer_with_attributes(EAGO.Optimizer, "verbosity" => 1,
                                                    "output_iterations" => 1000,
                                                    "iteration_limit" => 100000,

                                                    "cp_depth" => -1,
                                                    "cp_repetitions" => -1,
                                                    "cp_forward_reverse_limit" => 2,

                                                    "cut_min_iterations" => 2,
                                                    "cut_max_iterations" => 2,
                                                    "objective_cut_on" => true,
                                                    "subgrad_tighten" => true,

                                                    "obbt_depth" => 4,
                                                    "obbt_repetitions" => 4,
                                                    "obbt_aggressive_on" => true,

                                                    "upper_bounding_depth" => 4,

                                                    "fbbt_lp_depth" => 100000,
                                                    "fbbt_lp_repetitions" => 3))

# ----- Variables ----- #
x_Idx = Any[2, 3, 4, 5]
@variable(m, x[x_Idx])
JuMP.set_lower_bound(x[2], 1.0e-7)
JuMP.set_upper_bound(x[2], 0.5)
JuMP.set_lower_bound(x[3], 1.0e-7)
JuMP.set_upper_bound(x[3], 0.5)
JuMP.set_lower_bound(x[4], 1.0e-7)
JuMP.set_upper_bound(x[4], 0.5)
JuMP.set_lower_bound(x[5], 1.0e-7)
JuMP.set_upper_bound(x[5], 0.5)

# ----- Constraints ----- #
@constraint(m, e2, x[2]+x[3] == 0.5)
@constraint(m, e3, x[4]+x[5] == 0.5)

# ----- Objective ----- #
@NLobjective(m, Min, (x[2]*(log(x[2]/(x[2]+x[4]))+log(x[2]/(x[2]+0.095173*x[4])))+x[4]*(log(x[4]/(x[2]+x[4]))+log(x[4]/(0.30384*x[2]+x[4])))+(x[2]+2.6738*x[4])*log(x[2]+2.6738*x[4])+(0.374*x[2]+x[4])*log(0.374*x[2]+x[4])+2.6738*x[4]*log(x[4]/(x[2]+2.6738*x[4]))+0.374*x[2]*log(x[2]/(0.374*x[2]+x[4]))+x[3]*(log(x[3]/(x[3]+x[5]))+log(x[3]/(x[3]+0.095173*x[5])))+x[5]*(log(x[5]/(x[3]+x[5]))+log(x[5]/(0.30384*x[3]+x[5])))+(x[3]+2.6738*x[5])*log(x[3]+2.6738*x[5])+(0.374*x[3]+x[5])*log(0.374*x[3]+x[5])+2.6738*x[5]*log(x[5]/(x[3]+2.6738*x[5]))+0.374*x[3]*log(x[3]/(0.374*x[3]+x[5]))-3.6838*x[2]*log(x[2])-1.59549*x[4]*log(x[4])-3.6838*x[3]*log(x[3])-1.59549*x[5]*log(x[5])))

JuMP.optimize!(m)

backend_opt = backend(m).optimizer.model.optimizer
last_relaxed_opt = backend_opt.relaxed_optimizer
run_time = backend_opt._run_time
println("run time: $run_time")


"""
    ComputeLimberArray!(CosmologicalGrid::CosmologicalGrid,
    BackgroundQuantities::BackgroundQuantities)

This function compute the Limber grid. In the Limber approximation, ``k`` and
``z`` are related by the following relation:
```math
k_\\ell(z)=\\frac{\\ell+1/2}{r(z)}.
```
"""
function ComputeLimberArray!(CosmologicalGrid::CosmologicalGrid,
    BackgroundQuantities::BackgroundQuantities)
    for (idx_l, myl) in enumerate(CosmologicalGrid.ℓBinCenters)
        CosmologicalGrid.KLimberArray[idx_l, :] =  (myl + 0.5) ./
        BackgroundQuantities.rZArray
    end
end

"""
    InterpolatePowerSpectrumLimberGrid!(CosmologicalGrid::CosmologicalGrid,
    BackgroundQuantities::BackgroundQuantities, PowerSpectrum::PowerSpectrum,
    ::InterpolationMethod)

This function interpolates the Power Spectrum on the ``k-z`` grid and evaluates
it on the Limber grid.
"""
function InterpolatePowerSpectrumLimberGrid!(CosmologicalGrid::CosmologicalGrid,
    BackgroundQuantities::BackgroundQuantities, PowerSpectrum::PowerSpectrum,
    ::RectBivSplineDierckx)
    InterpPmm = Dierckx.Spline2D(log10.(CosmologicalGrid.KArray),
    CosmologicalGrid.ZArray, log10.(PowerSpectrum.PowerSpectrumNonlinArray);
    kx=5, ky=5, s=0.0)
    for (idx_l, myl) in enumerate(CosmologicalGrid.ℓBinCenters)
        PowerSpectrum.InterpolatedPowerSpectrum[idx_l, :] =
        10 .^(InterpPmm.(log10.(CosmologicalGrid.KLimberArray[idx_l, :]),
        CosmologicalGrid.ZArray))
    end
end

function InterpolatePowerSpectrumLimberGrid!(CosmologicalGrid::CosmologicalGrid,
    BackgroundQuantities::BackgroundQuantities, PowerSpectrum::PowerSpectrum,
    ::BSplineCubic)
    x = LinRange(log10(first(CosmologicalGrid.KArray)),
    log10(last(CosmologicalGrid.KArray)), length(CosmologicalGrid.KArray))
    y = LinRange(first(CosmologicalGrid.ZArray), last(CosmologicalGrid.ZArray),
    length(CosmologicalGrid.ZArray))
    InterpPmm = Interpolations.interpolate(
    log10.(PowerSpectrum.PowerSpectrumNonlinArray),
    BSpline(Cubic(Line(OnGrid()))))
    InterpPmm = scale(InterpPmm, x, y)
    InterpPmm = Interpolations.extrapolate(InterpPmm, Line())
    for (idx_l, myl) in enumerate(CosmologicalGrid.ℓBinCenters)
        PowerSpectrum.InterpolatedPowerSpectrum[idx_l, :] =
        10 .^(InterpPmm.(log10.(CosmologicalGrid.KLimberArray[idx_l, :]),
        CosmologicalGrid.ZArray))
    end
end

function InterpolatePowerSpectrumLimberGrid!(CosmologicalGrid::CosmologicalGrid,
    BackgroundQuantities::BackgroundQuantities, PowerSpectrum::PowerSpectrum,
    ConvolvedDensity::ConvolvedDensity)
    x = LinRange(log10(first(CosmologicalGrid.KArray)),
    log10(last(CosmologicalGrid.KArray)), length(CosmologicalGrid.KArray))
    y = LinRange(first(CosmologicalGrid.ZArray), last(CosmologicalGrid.ZArray),
    length(CosmologicalGrid.ZArray))
    InterpPmm = Interpolations.interpolate(
    log10.(PowerSpectrum.PowerSpectrumLinArray),
    BSpline(Cubic(Line(OnGrid()))))
    InterpPmm = scale(InterpPmm, x, y)
    InterpPmm = Interpolations.extrapolate(InterpPmm, Line())
    for iidx in 1:length(ConvolvedDensity.ZBinArray)-1
        for lidx in 1:length(CosmologicalGrid.ℓBinCenters)
            PowerSpectrum.InterpolatedPowerSpectrumBeyondLimber[lidx, :] =
            10 .^(InterpPmm.(log10.(
            CosmologicalGrid.KBeyondLimberArray[lidx, :]),
            CosmologicalGrid.ZArray[1]))
        end
    end
end

function ExtractGrowthFactor!(PowerSpectrum::PowerSpectrum)
    PowerSpectrum.GrowthFactor = zeros(
        length(PowerSpectrum.PowerSpectrumLinArray[1,:]))
    for zidx in 1:length(PowerSpectrum.GrowthFactor)
        PowerSpectrum.GrowthFactor[zidx] =
        sqrt(PowerSpectrum.PowerSpectrumLinArray[1,zidx]./
        PowerSpectrum.PowerSpectrumLinArray[1,1])
    end
end


function t_test(V::VAR)
  K = size(V.Σ,1)
  Kp = size(V.X,1)
  H = kron(inv(V.X*V.X'),V.Σ)
  SE = sqrt(diag(H))
  T = reshape(vec(V.β)./SE,K,Kp)
  return T
end


getsize(width, height, depth) =
    2*(width*height+width*depth+height*depth),
    width*height*depth


function cheb2ultra(c_cheb::AbstractVector{T},λ::T,plan::ChebyshevUltrasphericalPlan{BACKWARD,T}) where T<:AbstractFloat
    M,N,nM₀,αN,K = getconstants(plan)
    i₁,i₂,j₁,j₂ = getindices(plan)
    p₁,p₂,rp,c₁,c₂,um,vm,θ,tempsin,tempsin2,tempsinλ,tempsinλm,tempmindices,cnλ,cnmλ,w,anλ,c_cheb2,pr = getplan(plan)

    # Initialize c_cheb2
    @inbounds for i=1:N+1 c_cheb2[i] = c_cheb[i] end
    @inbounds for i=N+2:2N+1 c_cheb2[i] = 0 end

    # Initialize values at Clenshaw-Curtis points
    v_ultra = zero(c_cheb2)

    # Perform transposed DCT-I
    applyTN!(c_cheb2,p₁)

    # Scale values by Clenshaw-Curtis weights
    @inbounds for i=1:2N+1 c_cheb2[i] *= w[i] end

    k=0
    for κ=1:K
        k = κ
        copyto!(cnmλ,cnλ)
        copyto!(tempsinλm,tempsinλ)

        j₁[k] < nM₀ && (k-=1; break)

        @inbounds for i=i₁[k]:i₁[k+1]-1
            forward_recurrence!(pr,j₂[k],θ[i],tempsin2[2N+2-i],tempsin2[i],rp)
            cc2i = c_cheb2[i]
            for j=1:j₂[k] v_ultra[j] += cc2i*pr[j] end
        end

        @inbounds for i=i₂[k+1]+1:i₂[k]
            forward_recurrence!(pr,j₂[k],θ[i],tempsin2[2N+2-i],tempsin2[i],rp)
            cc2i = c_cheb2[i]
            for j=1:j₂[k] v_ultra[j] += cc2i*pr[j] end
        end

        for m=0:M-1
            # Compute u_m(θ) and v_m(θ)
            compute_umvm!(um,vm,λ,tempsinλm,m,θ,i₁[k+1]:i₂[k+1])

            # Multiply point-wise by u_m(θ) and v_m(θ) for valid indices
            init_c₁c₂!(c₁,c₂.parent,um,vm,c_cheb2,i₁[k+1],i₂[k+1])

            # Apply planned DCT-I and DST-I in-place
            applyTN!(c₁,p₁);applyUN!(c₂,p₂)

            # Compute diagonal 2N-scaling multiplied by local coefficients and zero out excess
            @inbounds for j=j₁[k]:j₂[k] v_ultra[j] += cnmλ[j]*(c₁[j]+c₂.parent[j]) end

            # Update C_{n,m}^λ by recurrence in m
            cst = (λ+m)/2/(m+1)*(m+1-λ)
            @inbounds for j=1:2N+1 cnmλ[j] *= cst/(j+λ+m) end

            # Update sin^{λ+m}(θ)
            @inbounds for j=1:2N+1 tempsinλm[j] *= tempsin[j] end
        end
    end

    # Finish off recurrence
    @inbounds for i=i₁[k+1]:i₂[k+1]
        forward_recurrence!(pr,j₂[k+1],θ[i],tempsin2[2N+2-i],tempsin2[i],rp)
        cc2i = c_cheb2[i]
        for j=1:j₂[k+1] v_ultra[j] += cc2i*pr[j] end
    end

    # Scale by orthonormality constants
    ret = zero(c_cheb)
    @inbounds for i=1:N+1 ret[i] = v_ultra[i]/anλ[i] end

    ret
end


# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Function to obtain the price index
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

"""
    create_price_index_hat(w_hat::Vector, τ_hat_Z::Matrix, τ_hat_F::Matrix, VA_coeff::Vector, γ::Matrix, π_Z::Matrix, π_F::Matrix, θ::Vector)

The function computes the optimal price indices (change) conditional on the economys' wage structure by minimizing the distance between successive iterations
    in the intermediate input price index.

# Arguments
- `w_hat::Vector`: N×1, country wage (change) vector w_hat.
- `τ_hat_Z::Matrix`: NS×NS, origin country-industry destination country-industry intermediate goods demand trade cost (change) matrix τ_hat_Z.
- `τ_hat_F::Matrix`: NS×N, origin country-industry destination country final goods trade cost (change) matrix τ_hat_F.
- `VA_coeff::Vector`: NS×1, country-industry value added coefficients vector VA_coeff.
- `γ::Matrix`: S×NS, country-industry intermediate input expenditure share matrix γ.
- `π_Z::Matrix`: NS×NS, origin country-industry destination country-industry intermediate goods trade share matrix π_Z.
- `π_F::Matrix`: NS×N, origin country-industry destination country final goods trade share matrix π_F.
- `θ::Vector`: NS×1, country-industry trade elasticities vector θ.

# Output
- `P_hat_Z::Matrix{Float64}`: S×NS, country-industry composite intermediate goods price index (change) per industry matrix P_hat_Z.
- `P_hat_F::Matrix{Float64}`: S×N, country-industry composite final goods price index (change) matrix P_hat_F.
- `π_hat_Z::Matrix{Float64}`: NS×NS, origin country-industry destination country-industry intermediate goods trade share (change) matrix π_hat_Z.
- `π_hat_F::Matrix{Float64}`: NS×N, origin country-industry destination country final goods trade share (change) matrix π_hat_F.
- `cost_hat::Matrix{Float64}`: NS×1, country-industry production cost (change) vector cost_hat.

# Examples
```julia-repl
julia> P_hat_Z, P_hat_F, π_hat_Z, π_hat_F, cost_hat = create_price_index_hat(w_hat, τ_hat_Z, τ_hat_F, VA_coeff, γ, π_Z, π_F, θ)
```
"""

function create_price_index_hat(w_hat::Vector, τ_hat_Z::Matrix, τ_hat_F::Matrix, VA_coeff::Vector, γ::Matrix, π_Z::Matrix, π_F::Matrix, θ::Vector)

    # initialize
    w0_hat = copy(w_hat) # N×1, update wage vector
    P0_hat_Z = ones(S, N*S) # S×NS, intermediate goods price indices

    # iteration parameters
    tolerance_p = 1e-3
    max_iteration_p = 100
    iteration_p = 0
    max_error_p = 1.0

    while max_error_p > tolerance_p && iteration_p <= max_iteration_p

        # # cost structure of the economy, equation (47)
        # cost_hat_w = [w0_hat[ceil(Int,i/S)]^VA_coeff[i] for i in 1:N*S] # NS×1, wages
        # cost_hat_p = P0_hat_Z.^γ
        # cost_hat_p = [prod(cost_hat_p[:,i]) for i in 1:N*S] # NS×1, price indices (prod is the same as sum just for multiplication)
        # cost_hat = cost_hat_w .* cost_hat_p # NS×1
        # global cost_hat = reshape(cost_hat, S, N) # S×N

        global cost_hat = ones(S,N)
        VAcoeff = reshape(VA_coeff, S, N)
        for s in 1:S
            for j in 1:N
                cost_hat[s,j] = w0_hat[j]^VAcoeff[s,j]
                for r in 1:S
                    cost_hat[s,j] = cost_hat[s,j]*P0_hat_Z[r,(j-1)*S+s]^γ[r,(j-1)*S+s]
                end
            end
        end

        # ----------
        # price index for intermediate goods from equation (48)
        global cost_hat_Z = zeros(N*S, N*S)
        θ = reshape(θ, S, N) # S×N

        for i in 1:N
            for r in 1:S
                for j in 1:N
                    for s in 1:S
                        iirr = (i-1)*S+r
                        jjss = (j-1)*S+s
                        cost_hat_Z[iirr,jjss] = (cost_hat[r,i]*τ_hat_Z[iirr,jjss])^(-θ[r,i])
                    end
                end
            end
        end

        # price index calculation
        global P_hat_Z = zeros(S, N*S)
        for r in 1:S
            for j in 1:N
                for s in 1:S
                    jjss = (j-1)*S+s
                    P_hat_Z[r,jjss] = (sum(π_Z[r:S:(N-1)*S+r,jjss] .* cost_hat_Z[r:S:(N-1)*S+r,jjss]))^(-1/θ[r,j])
                end
            end
        end

        P_hat_Z = ifelse.(isinf.(P_hat_Z), 0.0, P_hat_Z)

        # update iteration parameters
        error = abs.(P_hat_Z .- P0_hat_Z)
        max_error_p = maximum(error) # update error
        P0_hat_Z[:] = copy(P_hat_Z) # update to new price index
        iteration_p += 1 # update iteration count

        println(" - Inner loop: Iteration $iteration_p completed with error $max_error_p") # print update on progress

    end

    # ----------
    # price index for final goods from equation (49) --- not used in optimization, i.e. residual optimum (since cost is same in F as in Z!)
    cost_hat_F = zeros(N*S,N)
    for i in 1:N
        for r in 1:S
            for j in 1:N
                iirr = (i-1)*S+r
                cost_hat_F[iirr,j] = (cost_hat[r,i]*τ_hat_F[iirr,j])^(-θ[r,i])
            end
        end
    end

    global P_hat_F = zeros(S, N)
    for r in 1:S
        for j in 1:N
            P_hat_F[r,j] = (sum(π_F[r:S:(N-1)*S+r,j] .* cost_hat_F[r:S:(N-1)*S+r,j]))^(-1/θ[r,j])
        end
    end

    P_hat_F = ifelse.(isinf.(P_hat_F), 0.0, P_hat_F)

    # ----------
    # trade shares from equation (45) and (46)
    θ = vec(θ)
    θ = repeat(θ, 1, N*S)
    global π_hat_Z = cost_hat_Z ./ repeat(P_hat_Z, N) .^ (.-θ) # NS×NS, - cost_hat_Z is already to power of -θ => only P_hat_Z^-θ

    θ = θ[:, 1:N] # reduce to one NS×N again
    global π_hat_F = cost_hat_F ./ repeat(P_hat_F, N) .^ (.-θ) # NS×N,  - cost_hat_Z is already to power of -θ => only P_hat_Z^-θ

    return P_hat_Z, P_hat_F, π_hat_Z, π_hat_F, cost_hat
end


# This contains several objectives

##
# TODO make objectives have ! names
##
# Objective: L∞ norm of the disturbance
function max_disturbance!(model::Model, var)
    o = symbolic_infty_norm(var)
    @objective(model, Min, o)
    return o
end

function min_sum!(model::Model, var)
    o = sum(sum.(var))
    @objective(model, Min, o)
    return o
end

function max_sum!(model::Model, var)
    o = sum(sum.(var))
    @objective(model, Max, o)
    return o
end

function zero_objective!(model::Model)
    @objective(model, Max, 0.0)
end

function linear_objective!(mode::Model, map::HPolytope, var)
    c, d = tosimplehrep(map)
    o = c * var - d
    @objective(model, Min, o)
    return o
end


# this example is transported from Andrew Delong and Anton Osokin's Matlab wrapper. All credit goes to them.
using GCoptimization

h = gco_create(4,3)     # % Create new object with NumSites=4, NumLabels=3

gco_setdatacost(h, [
    0 9 2 0;            # % Sites 1,4 prefer  label 1
    3 0 3 3;            # % Site  2   prefers label 2 (strongly)
    5 9 0 5])           # % Site  3   prefers label 3

gco_setsmoothcost(h, [
    0 1 2;              # %
    1 0 1;              # % Linear (Total Variation) pairwise cost
    2 1 0])             # %

gco_setneighbors(h, [
    0 1 0 0;            # % Sites 1 and 2 connected with weight 1
    0 0 1 0;            # % Sites 2 and 3 connected with weight 1
    0 0 0 2;            # % Sites 3 and 4 connected with weight 2
    0 0 0 0])

gco_expansion(h)        # % Compute optimal labeling via alpha-expansion

gco_getlabeling(h)      # % Optimal labeling is (1,2,1,1)

gco_energy(h)           # % Energy = Data Energy + Smooth Energy + Label Energy


# Compare Jain-Neal sampler to Reversible Jump on the galaxy dataset.
module GalaxyCompareJNtoRJ

using BayesianMixtures
B = BayesianMixtures
can_plot = (Pkg.installed("PyPlot")!=nothing)
can_plot ? using PyPlot : warn("Skipping plots since PyPlot is not installed.")

# ==================== Run Jain-Neal sampler ====================

# Read the data from file
x = readdlm("datasets/galaxy.dat",' ',Float64)[:]

# Specify model, data, and MCMC options
n_total = 1000000  # total number of MCMC sweeps to run
log_pk = "k -> log(k in (1:30)? 1/30 : 0)"   # log prior on k is Uniform{1,...,30}
options = B.options("Normal","MFM",x,n_total; n_keep=1000,log_pk=log_pk,t_max=30)

# Run MCMC sampler
result = B.run_sampler(options)

# ==================== Run RJMCMC sampler using Nmix ====================

# Remove old output files
if (!isdir("galx")); mkdir("galx"); end
outfiles = ["1.bdlog","1.bk","1.ent","1.k","1.log","1.out","1.pe"]
for outfile in outfiles; try rm(joinpath("galx",outfile)); end; end

# Run Peter Green's Nmix program
if is_apple() # if running Mac OSX
    run(`chmod +x Nmix`)
    tic()
    run(`./Nmix -n$n_total -nb0 -ns$n_total galx`)
    elapsed_time_RJ = toq()
elseif is_windows() # if running Windows
    tic()
    run(`Nmix.exe -n$n_total -nb0 -ns$n_total galx`)
    elapsed_time_RJ = toq()
else
    error("Only Mac and Windows are currently supported for running Green's Nmix program.")
end

# Read and delete the output
output = readdlm(joinpath("galx","1.out"))
for outfile in outfiles; rm(joinpath("galx",outfile)); end

# Parse the output
k_RJ = zeros(Int64,n_total)
t_RJ = zeros(Int64,n_total)
i = 1
for j = 1:n_total
    k,t = output[i,find(output[i,:].!="")]
    k_RJ[j] = k
    t_RJ[j] = t
    i += k+1
end

# ==================== Compare posterior on k ====================

# Compute the posterior on k from Jain-Neal sampler
k_posterior_JN = B.k_posterior(result)

# Compute posterior on k from RJMCMC sampler
use = options.n_burn+1:n_total
k_posterior_RJ = B.histogram(k_RJ[use], 0:options.t_max)[1]/length(use)

# Print posteriors on k
println("Posteriors on k:")
println("  Jain-Neal  RJMCMC")
for k=1:15
    @printf "k=%d: %.4f %.4f\n" k k_posterior_JN[k] k_posterior_RJ[k]
end

if can_plot

    # ==================== Compare traceplots of t ====================
    # Traceplot of t for Jain-Neal
    B.open_figure(1)
    B.traceplot(result.t[1:min(5000,n_total)])
    B.labels("Traceplot for t (Jain-Neal)","sweep number","t")

    # Traceplot of t for RJMCMC
    B.open_figure(2)
    B.traceplot(t_RJ[1:min(5000,n_total)])
    B.labels("Traceplot for t (RJMCMC)","sweep number","t")

    # ==================== Compare autocorrelation of t ====================
    # Plot autocorrelation for Jain-Neal
    B.open_figure(3)
    B.plot_autocorrelation(result,0.01)

    # Plot autocorrelation for RJMCMC
    B.plot_autocorrelation(t_RJ[use],0.01,elapsed_time_RJ/n_total; color="r")
    legend(["Jain-Neal","RJMCMC"],loc="upper right",fontsize=12)

    # ==================== Plot density estimate ====================
    # Plot density estimate using Jain-Neal output
    B.open_figure(4)
    B.plot_histogram(x; color="w",edgecolor="g",linewidth=0.5)
    B.rug_plot(x)
    B.plot_density_estimate(result; resolution=500)
end

end # module GalaxyCompareJNtoRJ


#= src/util.jl
=#

crandn(::Type{T}, dims::Integer...) where {T<:Real} = convert(Array{T}, randn(dims...))
for (elty, relty) in ((:ComplexF32, :Float32), (:ComplexF64, :Float64))
  @eval begin
    crandn(::Type{$elty}, dims::Integer...) =
      reshape(reinterpret($elty, vec(crandn($relty, 2*dims[1], dims[2:end]...))),
              (dims...,))
  end
end

function findmaxabs(x::AbstractVecOrMat{T}) where T
  m  = zero(real(T))
  mi = 0
  @inbounds for i = 1:length(x)
    t = abs(x[i])
    t < m && continue
    m  = t
    mi = i
  end
  m, mi
end

function getcols(trans::Symbol, A::AbstractMatrix, cols)
  if     trans == :n  return A[:,cols]
  elseif trans == :c  return collect(A[cols,:]')
  end
end
function getcols(trans::Symbol, A::AbstractLinearOperator, cols)
  if     trans == :n  return A[:,cols]
  elseif trans == :c  return collect(A'[:, cols])
  end
end

function hermitianize!(A::AbstractMatrix, uplo::Symbol=:U)
  checksquare(A)
  n = size(A, 2)
  if uplo == :U
    for j = 1:n, i = 1:j
      @inbounds A[i,j] = 0.5*(A[i,j] + conj(A[j,i]))
    end
  elseif uplo == :L
    for i = 1:n, j = 1:i
      @inbounds A[i,j] = 0.5*(A[i,j] + conj(A[j,i]))
    end
  end
  Hermitian(A, uplo)
end

function iscale!(A::AbstractMatrix, b::AbstractVector)
  m, n = size(A)
  length(b) == n || throw(DimensionMismatch)
  @inbounds for j = 1:n
    bj = b[j]
    @simd for i = 1:m
      A[i,j] /= bj
    end
  end
  A
end
function iscale!(b::AbstractVector, A::AbstractMatrix)
  m, n = size(A)
  length(b) == m || throw(DimensionMismatch)
  @inbounds for j = 1:n
    @simd for i = 1:m
      A[i,j] /= b[i]
    end
  end
  A
end

function orthcols!(
    A::AbstractMatrix{T}, tau::Vector{T}, work::Vector{T}; thin::Bool=true) where T<:BlasFloat
  m, n = size(A)
  k = min(m, n)
  A, tau, work = _LAPACK.geqrf!(A, tau, work)
  A, tau, work = _LAPACK.orgqr!(A, tau, k, work)
  if thin  A = k < n ? A[:,1:k] : A
  else     A[:,k+1:end] .= 0
  end
  A, tau, work
end
orthcols!(A::AbstractMatrix{T}; thin::Bool=true) where {T} =
  orthcols!(A, Array{T}(undef, 1), Array{T}(undef, 1), thin=thin)[1]

function orthrows!(
    A::AbstractMatrix{T}, tau::Vector{T}, work::Vector{T}; thin::Bool=true) where T<:BlasFloat
  m, n = size(A)
  k = min(m, n)
  A, tau, work = _LAPACK.gelqf!(A, tau, work)
  A, tau, work = _LAPACK.orglq!(A, tau, k, work)
  if thin  A = k < m ? A[1:k,:] : A
  else     A[k+1:end,:] .= 0
  end
  A, tau, work
end
orthrows!(A::AbstractMatrix{T}; thin::Bool=true) where {T} =
  orthrows!(A, Array{T}(undef, 1), Array{T}(undef, 1), thin=thin)[1]

function scalevec!(s::AbstractVector, x::AbstractVector)
  n = length(x)
  length(s) == n || throw(DimensionMismatch)
  @inbounds @simd for i = 1:n
    x[i] *= s[i]
  end
  x
end
function iscalevec!(s::AbstractVector, x::AbstractVector)
  n = length(x)
  length(s) == n || throw(DimensionMismatch)
  @inbounds @simd for i = 1:n
    x[i] /= s[i]
  end
  x
end

symrelerr(x, y) = 2*abs((x - y)/(x + y))