@doc doc"""
    ProductManifold{TM<:Tuple, TRanges<:Tuple, TSizes<:Tuple} <: Manifold

Product manifold $M_1 \times M_2 \times \dots \times M_n$ with product geometry.
`TRanges` and `TSizes` statically define the relationship between representation
of the product manifold and representations of point, tangent vectors
and cotangent vectors of respective manifolds.

# Constructor

    ProductManifold(M_1, M_2, ..., M_n)

generates the product manifold $M_1 \times M_2 \times \dots \times M_n$.
Alternatively, the same manifold can be contructed using the `×` operator:
`M_1 × M_2 × M_3`.
"""
struct ProductManifold{TM<:Tuple} <: Manifold
    manifolds::TM
end

ProductManifold(manifolds::Manifold...) = ProductManifold{typeof(manifolds)}(manifolds)

struct ProductMetric <: Metric end

"""
    ProductFVectorDistribution([type::VectorBundleFibers], [x], distrs...)

Generates a random vector at point `x` from vector space (a fiber of a tangent
bundle) of type `type` using the product distribution of given distributions.

Vector space type and `x` can be automatically inferred from distributions `distrs`.
"""
struct ProductFVectorDistribution{
    TSpace<:VectorBundleFibers{<:VectorSpaceType,<:ProductManifold},
    TD<:(NTuple{N,Distribution} where {N}),
    TX,
} <: FVectorDistribution{TSpace,TX}
    type::TSpace
    x::TX
    distributions::TD
end

"""
    ProductPointDistribution(M::ProductManifold, distributions)

Product distribution on manifold `M`, combined from `distributions`.
"""
struct ProductPointDistribution{
    TM<:ProductManifold,
    TD<:(NTuple{N,Distribution} where {N}),
} <: MPointDistribution{TM}
    manifold::TM
    distributions::TD
end

"""
    ProductRetraction(retractions::AbstractRetractionMethod...)

Product retraction of `retractions`. Works on [`ProductManifold`](@ref).
"""
struct ProductRetraction{TR<:Tuple} <: AbstractRetractionMethod
    retractions::TR
end

function ProductRetraction(retractions::AbstractRetractionMethod...)
    return ProductRetraction{typeof(retractions)}(retractions)
end

"""
    InverseProductRetraction(retractions::AbstractInverseRetractionMethod...)

Product inverse retraction of `inverse retractions`. Works on [`ProductManifold`](@ref).
"""
struct InverseProductRetraction{TR<:Tuple} <: AbstractInverseRetractionMethod
    inverse_retractions::TR
end

function InverseProductRetraction(inverse_retractions::AbstractInverseRetractionMethod...)
    return InverseProductRetraction{typeof(inverse_retractions)}(inverse_retractions)
end

"""
    PrecomputedProductOrthonormalBasis(parts::NTuple{N,AbstractPrecomputedOrthonormalBasis} where N, F::AbstractNumbers = ℝ)

A precomputed orthonormal basis of a tangent space of a product manifold.
The tuple `parts` stores bases corresponding to multiplied manifolds.

The type parameter `F` denotes the [`AbstractNumbers`](@ref) that will be used as scalars.
"""
struct PrecomputedProductOrthonormalBasis{
    T<:NTuple{N,AbstractPrecomputedOrthonormalBasis} where {N},
    F,
} <: AbstractPrecomputedOrthonormalBasis{F}
    parts::T
end

function PrecomputedProductOrthonormalBasis(
    parts::NTuple{N,AbstractPrecomputedOrthonormalBasis},
    F::AbstractNumbers = ℝ,
) where {N}
    return PrecomputedProductOrthonormalBasis{typeof(parts),F}(parts)
end

"""
    check_manifold_point(M::ProductManifold, x; kwargs...)

Check whether `x` is a valid point on the [`ProductManifold`](@ref) `M`.

The tolerance for the last test can be set using the `kwargs...`.
"""
function check_manifold_point(M::ProductManifold, x::ProductRepr; kwargs...)
    for t ∈ ziptuples(M.manifolds, submanifold_components(M, x))
        err = check_manifold_point(t...; kwargs...)
        err === nothing || return err
    end
    return nothing
end

function check_manifold_point(M::ProductManifold, x::ProductArray; kwargs...)
    for t ∈ ziptuples(M.manifolds, submanifold_components(M, x))
        err = check_manifold_point(t...; kwargs...)
        err === nothing || return err
    end
    return nothing
end

"""
    check_tangent_vector(M::ProductManifold, x, v; kwargs... )

Check whether `v` is a tangent vector to `x` on the [`ProductManifold`](@ref)
`M`, i.e. atfer [`check_manifold_point`](@ref)`(M, x)`, and all projections to
base manifolds must be respective tangent vectors.

The tolerance for the last test can be set using the `kwargs...`.
"""
function check_tangent_vector(M::ProductManifold, x::ProductRepr, v::ProductRepr; kwargs...)
    perr = check_manifold_point(M, x)
    perr === nothing || return perr
    ts = ziptuples(M.manifolds, submanifold_components(M, x), submanifold_components(M, v))
    for t ∈ ts
        err = check_tangent_vector(t...; kwargs...)
        err === nothing || return err
    end
    return nothing
end
function check_tangent_vector(
    M::ProductManifold,
    x::ProductArray,
    v::ProductArray;
    kwargs...,
)
    perr = check_manifold_point(M, x)
    perr === nothing || return perr
    ts = ziptuples(M.manifolds, submanifold_components(M, x), submanifold_components(M, v))
    for t ∈ ts
        err = check_tangent_vector(t...; kwargs...)
        err === nothing || return err
    end
    return nothing
end

@doc doc"""
    cross(M,N)
    cross(M1, M2, M3,...)

Return the [`ProductManifold`](@ref) For two [`Manifold`](@ref)s `M` and `N`,
where for the case that one of them is a [`ProductManifold`](@ref) itself,
the other is either prepended (if `N` is a product) or appenden (if `M`) is.
If both are product manifold, they are combined into one product manifold,
keeping the order.

For the case that more than one is a product manifold of these is build with the
same approach as above
"""
cross(::Manifold...)
cross(M1::Manifold, M2::Manifold) = ProductManifold(M1, M2)
cross(M1::ProductManifold, M2::Manifold) = ProductManifold(M1.manifolds..., M2)
cross(M1::Manifold, M2::ProductManifold) = ProductManifold(M1, M2.manifolds...)
function cross(M1::ProductManifold, M2::ProductManifold)
    return ProductManifold(M1.manifolds..., M2.manifolds...)
end

function det_local_metric(M::MetricManifold{ProductManifold,ProductMetric}, x::ProductArray)
    dets = map(det_local_metric, M.manifolds, submanifold_components(M, x))
    return prod(dets)
end

@doc doc"""
    distance(M::ProductManifold, x, y)

compute the distance two points `x` and `y` on the [`ProductManifold`](@ref) `M`, which is
the 2-norm of the elementwise distances on the internal manifolds that build `M`.
"""
function distance(M::ProductManifold, x, y)
    return sqrt(sum(
        map(
            distance,
            M.manifolds,
            submanifold_components(M, x),
            submanifold_components(M, y),
        ) .^ 2,
    ))
end

@doc doc"""
    exp(M::ProductManifold, x, v)

compute the exponential map from `x` towards `v` on the [`ProductManifold`](@ref) `M`,
which is the elementwise exponential map on the internal manifolds that build `M`.
"""
exp(::ProductManifold, ::Any...)
function exp(M::ProductManifold, x::ProductRepr, v::ProductRepr)
    return ProductRepr(map(
        exp,
        M.manifolds,
        submanifold_components(M, x),
        submanifold_components(M, v),
    )...)
end

function exp!(M::ProductManifold, y, x, v)
    map(
        exp!,
        M.manifolds,
        submanifold_components(M, y),
        submanifold_components(M, x),
        submanifold_components(M, v),
    )
    return y
end

@doc doc"""
    flat(M::ProductManifold, x, w::FVector{TangentSpaceType})

use the musical isomorphism to transform the tangent vector `w` from the tangent space at
`x` on the [`ProductManifold`](@ref) `M` to a cotangent vector.
This can be done elementwise, so for every entry of `w` (and `x`) sparately
"""
flat(::ProductManifold, ::Any...)

function flat!(M::ProductManifold, v::CoTFVector, x, w::TFVector)
    vfs = map(u -> FVector(CotangentSpace, u), submanifold_components(v))
    wfs = map(u -> FVector(TangentSpace, u), submanifold_components(w))
    map(flat!, M.manifolds, vfs, submanifold_components(M, x), wfs)
    return v
end

function get_basis(M::ProductManifold, x, B::AbstractBasis)
    parts = map(t -> get_basis(t..., B), ziptuples(M.manifolds, submanifold_components(x)))
    return PrecomputedProductOrthonormalBasis(parts)
end
function get_basis(M::ProductManifold, x, B::DiagonalizingOrthonormalBasis)
    vs = map(ziptuples(
        M.manifolds,
        submanifold_components(x),
        submanifold_components(B.v),
    )) do t
        return get_basis(t[1], t[2], DiagonalizingOrthonormalBasis(t[3]))
    end
    return PrecomputedProductOrthonormalBasis(vs)
end
function get_basis(M::ProductManifold, x, B::ArbitraryOrthonormalBasis)
    parts = map(t -> get_basis(t..., B), ziptuples(M.manifolds, submanifold_components(x)))
    return PrecomputedProductOrthonormalBasis(parts)
end

function get_coordinates(M::ProductManifold, x, v, B::PrecomputedProductOrthonormalBasis)
    reps = map(
        get_coordinates,
        M.manifolds,
        submanifold_components(x),
        submanifold_components(v),
        B.parts,
    )
    return vcat(reps...)
end
function get_coordinates(M::ProductManifold, x, v, B::ArbitraryOrthonormalBasis)
    reps = map(
        t -> get_coordinates(t..., B),
        ziptuples(M.manifolds, submanifold_components(x), submanifold_components(v)),
    )
    return vcat(reps...)
end

function get_vector(
    M::ProductManifold{<:NTuple{N,Any}},
    x::ProductRepr,
    v,
    B::PrecomputedProductOrthonormalBasis,
) where {N}
    dims = map(manifold_dimension, M.manifolds)
    dims_acc = accumulate(+, [1, dims...])
    parts = ntuple(N) do i
        get_vector(
            M.manifolds[i],
            submanifold_component(x, i),
            v[dims_acc[i]:dims_acc[i]+dims[i]-1],
            B.parts[i],
        )
    end
    return ProductRepr(parts)
end
function get_vector(
    M::ProductManifold{<:NTuple{N,Any}},
    x::ProductRepr,
    v,
    B::ArbitraryOrthonormalBasis,
) where {N}
    dims = map(manifold_dimension, M.manifolds)
    dims_acc = accumulate(+, [1, dims...])
    parts = ntuple(N) do i
        get_vector(
            M.manifolds[i],
            submanifold_component(x, i),
            v[dims_acc[i]:dims_acc[i]+dims[i]-1],
            B,
        )
    end
    return ProductRepr(parts)
end

function get_vectors(
    M::ProductManifold{<:NTuple{N,Manifold}},
    x::ProductRepr,
    B::PrecomputedProductOrthonormalBasis,
) where {N}
    xparts = submanifold_components(x)
    BVs = map(t -> get_vectors(t...), ziptuples(M.manifolds, xparts, B.parts))
    zero_tvs = map(t -> zero_tangent_vector(t...), ziptuples(M.manifolds, xparts))
    vs = typeof(ProductRepr(zero_tvs...))[]
    for i = 1:N, k = 1:length(BVs[i])
        push!(vs, ProductRepr(zero_tvs[1:i-1]..., BVs[i][k], zero_tvs[i+1:end]...))
    end
    return vs
end

function hat!(M::ProductManifold, v, x, vⁱ)
    dim = manifold_dimension(M)
    @assert length(vⁱ) == dim
    i = one(dim)
    ts = ziptuples(M.manifolds, submanifold_components(M, v), submanifold_components(M, x))
    for t ∈ ts
        dim = manifold_dimension(first(t))
        tvⁱ = @inbounds view(vⁱ, i:(i+dim-1))
        hat!(t..., tvⁱ)
        i += dim
    end
    return v
end

@doc doc"""
    injectivity_radius(M::ProductManifold[, x])

Compute the injectivity radius on the [`ProductManifold`](@ref), which is the
minimum of the factor manifolds.
"""
injectivity_radius(::ProductManifold, ::Any...)
function injectivity_radius(M::ProductManifold, x)
    return min(map(injectivity_radius, M.manifolds, submanifold_components(M, x))...)
end
injectivity_radius(M::ProductManifold) = min(map(injectivity_radius, M.manifolds)...)

@doc doc"""
    inner(M::ProductManifold, x, v, w)

compute the inner product of two tangent vectors `v`, `w` from the tangent space
at `x` on the [`ProductManifold`](@ref) `M`, which is just the sum of the
internal manifolds that build `M`.
"""
function inner(M::ProductManifold, x, v, w)
    subproducts = map(
        inner,
        M.manifolds,
        submanifold_components(M, x),
        submanifold_components(M, v),
        submanifold_components(M, w),
    )
    return sum(subproducts)
end

@doc doc"""
    inverse_retract(M::ProductManifold, x, y, m::InverseProductRetraction)

Compute the inverse retraction from `x` with respect to `y` on the [`ProductManifold`](@ref)
`M` using an [`InverseProductRetraction`](@ref), which by default encapsulates a inverse
retraction for each manifold of the product. Then this method is performed elementwise,
so the encapsulated inverse retraction methods have to be available per factor.
"""
inverse_retract(::ProductManifold, ::Any, ::Any, ::Any, ::InverseProductRetraction)

function inverse_retract!(M::ProductManifold, v, x, y, method::InverseProductRetraction)
    map(
        inverse_retract!,
        M.manifolds,
        submanifold_components(M, v),
        submanifold_components(M, x),
        submanifold_components(M, y),
        method.inverse_retractions,
    )
    return v
end

is_default_metric(::ProductManifold, ::ProductMetric) = Val(true)

function isapprox(M::ProductManifold, x, y; kwargs...)
    return all(
        t -> isapprox(t...; kwargs...),
        ziptuples(M.manifolds, submanifold_components(M, x), submanifold_components(M, y)),
    )
end
function isapprox(M::ProductManifold, x, v, w; kwargs...)
    return all(
        t -> isapprox(t...; kwargs...),
        ziptuples(
            M.manifolds,
            submanifold_components(M, x),
            submanifold_components(M, v),
            submanifold_components(M, w),
        ),
    )
end

@doc doc"""
    log(M::ProductManifold, x, y)

Compute the logarithmic map from `x` to `y` on the [`ProductManifold`](@ref) `M`,
which can be computed using the logarithmic maps of the manifolds elementwise.
"""
log(::ProductManifold, ::Any...)
function log(M::ProductManifold, x::ProductRepr, y::ProductRepr)
    return ProductRepr(map(
        log,
        M.manifolds,
        submanifold_components(M, x),
        submanifold_components(M, y),
    )...)
end

function log!(M::ProductManifold, v, x, y)
    map(
        log!,
        M.manifolds,
        submanifold_components(M, v),
        submanifold_components(M, x),
        submanifold_components(M, y),
    )
    return v
end

@doc doc"""
    manifold_dimension(M::ProductManifold)

Return the manifold dimension of the [`ProductManifold`](@ref), which is the sum of the
manifold dimensions the product is made of.
"""
manifold_dimension(M::ProductManifold) = mapreduce(manifold_dimension, +, M.manifolds)

@doc doc"""
    norm(M::PowerManifold, x, v)

Compute the norm of `v` from the tangent space of `x` on the [`ProductManifold`](@ref),
i.e. from the element wise norms the 2-norm is computed.
"""
function norm(M::ProductManifold, x, v)
    norms_squared = (
        map(
            norm,
            M.manifolds,
            submanifold_components(M, x),
            submanifold_components(M, v),
        ) .^ 2
    )
    return sqrt(sum(norms_squared))
end

function ProductFVectorDistribution(
    type::VectorBundleFibers{<:VectorSpaceType,<:ProductManifold},
    x::Union{AbstractArray,MPoint,ProductRepr},
    distributions::FVectorDistribution...,
)
    return ProductFVectorDistribution{typeof(type),typeof(distributions),typeof(x)}(
        type,
        x,
        distributions,
    )
end
function ProductFVectorDistribution(
    type::VectorBundleFibers{<:VectorSpaceType,<:ProductManifold},
    distributions::FVectorDistribution...,
)
    x = ProductRepr(map(d -> support(d).x, distributions))
    return ProductFVectorDistribution(type, x, distributions...)
end
function ProductFVectorDistribution(distributions::FVectorDistribution...)
    M = ProductManifold(map(d -> support(d).space.M, distributions)...)
    VS = support(distributions[1]).space.VS
    if !all(d -> support(d).space.VS == VS, distributions)
        error("Not all distributions have support in vector spaces of the same type, which is currently not supported")
    end
    # Probably worth considering sum spaces in the future?
    x = ProductRepr(map(d -> support(d).x, distributions)...)
    return ProductFVectorDistribution(VectorBundleFibers(VS, M), x, distributions...)
end

function ProductPointDistribution(M::ProductManifold, distributions::MPointDistribution...)
    return ProductPointDistribution{typeof(M),typeof(distributions)}(M, distributions)
end
function ProductPointDistribution(distributions::MPointDistribution...)
    M = ProductManifold(map(d -> support(d).manifold, distributions)...)
    return ProductPointDistribution(M, distributions...)
end

function rand(rng::AbstractRNG, d::ProductPointDistribution)
    return ProductRepr(map(d -> rand(rng, d), d.distributions)...)
end
function rand(rng::AbstractRNG, d::ProductFVectorDistribution)
    return ProductRepr(map(d -> rand(rng, d), d.distributions)...)
end

function _rand!(rng::AbstractRNG, d::ProductPointDistribution, x::AbstractArray{<:Number})
    return copyto!(x, rand(rng, d))
end
function _rand!(rng::AbstractRNG, d::ProductPointDistribution, x::ProductRepr)
    map(
        t -> _rand!(rng, t[1], t[2]),
        d.distributions,
        submanifold_components(d.manifold, x),
    )
    return x
end
function _rand!(rng::AbstractRNG, d::ProductFVectorDistribution, v::AbstractArray{<:Number})
    return copyto!(v, rand(rng, d))
end
function _rand!(rng::AbstractRNG, d::ProductFVectorDistribution, v::ProductRepr)
    map(t -> _rand!(rng, t[1], t[2]), d.distributions, submanifold_components(d.space.M, v))
    return v
end

@doc doc"""
    retract(M::ProductManifold, x, v, m::ProductRetraction)

Compute the retraction from `x` with tangent vector `v` on the [`ProductManifold`](@ref) `M`
using an [`ProductRetraction`](@ref), which by default encapsulates retractions of the
base manifolds. Then this method is performed elementwise, so the encapsulated retractions
method has to be one that is available on the manifolds.
"""
retract(::ProductManifold, ::Any...)

function retract!(M::ProductManifold, y, x, v, method::ProductRetraction)
    map(
        retract!,
        M.manifolds,
        submanifold_components(M, y),
        submanifold_components(M, x),
        submanifold_components(M, v),
        method.retractions,
    )
    return y
end

function representation_size(M::ProductManifold)
    return (mapreduce(m -> prod(representation_size(m)), +, M.manifolds),)
end

@doc doc"""
    sharp(M::ProductManifold, x, w::FVector{CotangentSpaceType})

Use the musical isomorphism to transform the cotangent vector `w` from the tangent space at
`x` on the [`ProductManifold`](@ref) `M` to a tangent vector.
This can be done elementwise, so vor every entry of `w` (and `x`) sparately
"""
sharp(::ProductManifold, ::Any...)

function sharp!(M::ProductManifold, v::TFVector, x, w::CoTFVector)
    vfs = map(u -> FVector(TangentSpace, u), submanifold_components(v))
    wfs = map(u -> FVector(CotangentSpace, u), submanifold_components(w))
    map(sharp!, M.manifolds, vfs, submanifold_components(M, x), wfs)
    return v
end

"""
    submanifold(M::ProductManifold, i::Integer)

Extract the `i`th factor of the product manifold `M`.
"""
submanifold(M::ProductManifold, i::Integer) = M.manifolds[i]

"""
    submanifold(M::ProductManifold, i::Val)
    submanifold(M::ProductManifold, i::AbstractVector)

Extract the factor of the product manifold `M` indicated by indices in `i`.
For example, for `i` equal to `Val((1, 3))` the product manifold constructed
from the first and the third factor is returned.

The version with `AbstractVector` is not type-stable, for better preformance use `Val`.
"""
submanifold(M::ProductManifold, i::Val) = ProductManifold(select_from_tuple(M.manifolds, i))
submanifold(M::ProductManifold, i::AbstractVector) = submanifold(M, Val(tuple(i...)))

support(d::ProductPointDistribution) = MPointSupport(d.manifold)
function support(tvd::ProductFVectorDistribution)
    return FVectorSupport(
        tvd.type,
        ProductRepr(map(d -> support(d).x, tvd.distributions)...),
    )
end

function vee!(M::ProductManifold, vⁱ, x, v)
    dim = manifold_dimension(M)
    @assert length(vⁱ) == dim
    i = one(dim)
    ts = ziptuples(M.manifolds, submanifold_components(M, x), submanifold_components(M, v))
    for t ∈ ts
        SM = first(t)
        dim = manifold_dimension(SM)
        tvⁱ = @inbounds view(vⁱ, i:(i+dim-1))
        vee!(SM, tvⁱ, Base.tail(t)...)
        i += dim
    end
    return vⁱ
end

function zero_tangent_vector!(M::ProductManifold, v, x)
    map(
        zero_tangent_vector!,
        M.manifolds,
        submanifold_components(M, v),
        submanifold_components(M, x),
    )
    return v
end


# the channel example using an activity-based approach
#
using DiscreteEvents, Printf, Random

mutable struct Server
  id::Int64
  name::AbstractString
  input::Channel
  output::Channel
  op     # operation
  token  # current token

  Server(id, name, input, output, op) = new(id, name, input, output, op, nothing)
end

arrive(S::Server) = event!(SF(serve, S), SF(isready, S.input))

function serve(S::Server)
    S.token = take!(S.input)
    @printf("%5.2f: %s %d took token %d\n", tau(), S.name, S.id, S.token)
    event!((SF(put!, S.output, S.op(S.id, S.token)), SF(arrive, S)), after, rand())
end

reset!(𝐶)
Random.seed!(123)

ch1 = Channel(32)  # create two channels
ch2 = Channel(32)

s = shuffle(1:8)
for i in 1:2:8
    arrive(Server(s[i], "foo", ch1, ch2, +))
    arrive(Server(s[i+1], "bar", ch2, ch1, *))
end

put!(ch1, 1) # put first token into channel 1

run!(𝐶, 10)


using GalerkinSparseGrids
using LinearAlgebra
using Plots

function get_coeffs(D,n)

    n_vec = fill(n,D)
    ls = ntuple(i->(n+1), D)
    coeffs = Dict()

    for q in 0:(D-1)
        factor = (-1)^q * binomial(D-1,q)

        for level in CartesianIndices(ls)
            cutoff(Val(:full),level,n_vec) && continue # welcher cutoff? (entlang diagonale vs. rechteck links oben)
            # hier cutoff entlang Diagonale -> n::Int64

            #CartesianIndex als Vektor darstellen
            level_n = Vector{Int64}() #[]
            for j in 1:D
                append!(level_n, level[j])
            end

            if sum(level_n) == n - q + D
                coeffs[level] = factor
            end
        end
    end
    println("Coefficients for CT: ", coeffs)
    return coeffs
end

function get_test_coeffs(D,n)
    level = CartesianIndex(ntuple(i->(n),D))
    coeffs = Dict()
    coeffs[level] = 1
    return coeffs
end

# Ein Kombinationsschritt, wenn t0 und t1 ganzes Zeitintervall abdecken, dann Kombitechnik, wo erst nach allen Zeitschritten kombiniert wird
# generate_coeffs ist boolean, ob f0coeffs & v0coeffs neu initialisiert werden (über traveling_wave) (--> true) oder ob bereits Koeffizienten in Dictionary Form (f_result, v_result)
# übergeben wurden (--> false)
function combine_end(D,k,n,m,f_result,v_result,t0,t1,generate_coeffs)
    # coeffs = get_coeffs(D,n)
    coeffs = get_test_coeffs(D,n)

    f_exact = x->sin(2*pi*x[1])*sin(2*pi*x[2])

    # # Convert n::Int to n::Vector
    # n_vec = fill(n,D)

    # Leere Vorlagen für soln erstellen
    len = get_size(Val(D), k, n+D, Val(:full)) # get_size mit n::Int
    # vec = Vector{Float64}(0.,len)
    vec = zeros(len)
    sparsegrid_f = V2D(D,k,n+D,vec,scheme="sparse") # welches scheme? spzeros(len,len) oder vec von welcher Länge?
    # println("sparsegrid_f", sparsegrid_f)
    sparsegrid_v = V2D(D,k,n+D,vec,scheme="sparse")
    err_all = []
    err_all_dict = []

    for (level,factor) in coeffs # geht nur über Level mit Koeffizienten != 0

        # CartesianIndex als Vektor darstellen
        level_n = Vector{Int64}() #[]
        for j in 1:D
            append!(level_n, level[j])
        end

        # Unterscheidung in 1. Zeitschritt und restliche Schritte
        if generate_coeffs # == true : 1. Zeitschritt --> initial position f und velocity v werden über traveling_wave generiert

            f0coeffs, v0coeffs = traveling_wave(k, level_n, m, scheme="full") #geht wsl nur, wenn f0 = x->sin(2*pi*x[1])*sin(2*pi*x[2]) verwendet wird

        else # 2+ Zeitschritt --> benötigen zuerst noch Dekombinierung (Zwischenergebnis auf Gitter aufteilen)
            # println("1")
            # Leere Vorlagen für subgrid erstellen
            len = get_size(Val(D), k, level_n, Val(:full)) # get_size mit n::Vector
            # vec = Vector{Float64}(undef,len)
            vec = zeros(len)
            subgrid_f = V2D(D,k,level_n,vec,scheme="full")
            subgrid_v = V2D(D,k,level_n,vec,scheme="full")
            # println("2")
            # über subgrid_f iterieren
            for (key,value) in subgrid_f
                for c in eachindex(value), m in eachindex(value[c])
                    subgrid_f[key][c][m] = f_result[key][c][m]
                end
            end
            # über subgrid_v iterieren #alternativ: f und v in eine for schleife und vorsichtshalber assert subgrid_f = subgrid_v
            for (key,value) in subgrid_v
                for c in eachindex(value), m in eachindex(value[c])
                    subgrid_v[key][c][m] = v_result[key][c][m]
                end
            end
            # println("3")
            f0coeffs = D2V(D,k,level_n,subgrid_f,scheme="full")
            v0coeffs = D2V(D,k,level_n,subgrid_v,scheme="full")
        end
        println("4: level_n = ", level_n)
        # Evolve function
        soln = wave_evolve(D, k, level_n, f0coeffs, v0coeffs, t0, t1; order="78", scheme="full")
        # println("5")
        dict_f = V2D(D,k,level_n,soln[2][1],scheme="full")
        # println("dict_f", dict_f)
        dict_v = V2D(D,k,level_n,soln[2][2],scheme="full")
        # println("6")
        # über dict_f iterieren
        for (key,value) in dict_f
            for c in eachindex(value), m in eachindex(value[c])
                sparsegrid_f[key][c][m] += factor * dict_f[key][c][m]
            end
        end
        # über dict_v iterieren
        for (key,value) in dict_v
            for c in eachindex(value), m in eachindex(value[c])
                sparsegrid_v[key][c][m] += factor * dict_v[key][c][m]
            end
        end
        # println("7")
        #reconstruct_DG
        f_rep_dict = x -> reconstruct_DG(dict_f, [x...])
        f_rep = x -> reconstruct_DG(sparsegrid_f, [x...])
        #mcerr
        f_exact = x -> cos(2*pi*(dot(m,x) - sqrt(dot(m,m))*t1))
        err_dict = mcerr(f_exact, f_rep_dict, D)
        err = mcerr(f_exact, f_rep, D)
        println("k = ", k, ", n = ", level_n, ", mcerr (dict) = ", err_dict)
        println("k = ", k, ", n = ", level_n, ", mcerr (kombiniert) = ", err)
        append!(err_all_dict, err_dict)
        append!(err_all, err)
        # println("8")
    end
    # println("9")
    f_rep = x -> reconstruct_DG(sparsegrid_f, [x...])
    err = mcerr(f_exact, f_rep, D)
    println("k = ", k, ", n = ", n, ", mcerr = ", err)
    # x = 1:n;
    # plot(x,err_all[1:n],title="Monte Carlo Error",xlabel="n",ylabel="mcerr",label="k=$k")
    return sparsegrid_f, sparsegrid_v
end

# Gitter nach jedem Zeitschritt kombinieren
function combine_between(D,k,n,m,f_result,v_result,t0,t1)
    tstep = t1-t0
    sparsegrid_f, sparsegrid_v = combine_end(D,k,n,m,f_result,v_result,t0,t0+tstep,true)
    for t_start in t0+tstep:tstep:(t1-tstep)
        t_end = t_start + tstep
        println("\nAktuelles t_start = ", t_start)
        sparsegrid_f, sparsegrid_v = combine_end(D,k,n,m,sparsegrid_f,sparsegrid_v,t_start,t_end,false)
    end
end

# Test:
println("\n\n-------------------------------------------------------------------")
D = 2;
k = 3;
n = 3; # Gitter ist immer isotrop
m = [1,2];
f0 = x->sin(2*pi*x[1])*sin(2*pi*x[2])
v0 = x->0
t0 = 0
t1 = 1

# get_coeffs(D,n)
combine_end(D,k,n,m,f0,v0,t0,t1,true)
# combine_between(D,k,n,m,0,0,t0,t1)
# get_test_coeffs(D,n)


export graph_bbcs_cheb_exp

"""
    (graph,cref)=graph_bbcs_cheb_exp(k;T=Complex{BigFloat})

Computes a polynomial evaluation approximating the exponential for
skew-Hermitian matrices using `k` matrix multiplications following the procedure
in the reference.

For `k > 5` it resorts to scaling-and-squaring of the `k = 5` graph.
The graph is in the degopt format, see [`graph_degopt`](@ref).

**Reference**

1. P. Bader, S. Blanes, F. Casas, M. Seydaoglu. "An efficient algorithm to
   compute the exponential of skew-Hermitian matrices for the time integration
   of the Schrödinger equation".
   [arXiv:2103.10132 [math.NA]](https://arxiv.org/abs/2103.10132), 2021.
"""
function graph_bbcs_cheb_exp(k; T = Complex{BigFloat})
    CBF = Complex{BigFloat}
    if (k <= 5) # Basic implementation from paper
        if (k == 1)
            x1a = vcat(zero(CBF), one(CBF))
            x1b = vcat(zero(CBF), one(CBF))
            xv = [(x1a, x1b)]
            y = [
                BigFloat("0.9999999999999999999998"),
                BigFloat("-0.9999999999761950000001")im,
                BigFloat("-0.4999999999920650000000"),
            ]
        elseif (k == 2)
            x1a = vcat(zero(CBF), one(CBF))
            x1b = vcat(zero(CBF), one(CBF))
            x2a = vcat(zeros(CBF, 2), one(CBF))
            x2b = vcat(
                zero(CBF),
                BigFloat("0.16666657785001893215")im,
                BigFloat("0.04166664890333648869"),
            )
            xv = [(x1a, x1b), (x2a, x2b)]
            y = [
                BigFloat("0.99999999999999999997"),
                BigFloat("-0.99999999999981067844")im,
                BigFloat("-0.49999999999994320353"),
                one(CBF),
            ]
        elseif (k == 3) # Equation (22)
            x1a = vcat(zero(CBF), one(CBF))
            x1b = vcat(zero(CBF), one(CBF))
            x2a = vcat(zeros(CBF, 2), one(CBF))
            x2b = vcat(
                zero(CBF),
                BigFloat("431") / BigFloat("4000"),
                BigFloat("-0.02693906873598870733")im,
            )
            x3a = vcat(
                zeros(CBF, 2),
                BigFloat("0.66321004441662438593")im,
                one(CBF),
            )
            x3b = vcat(
                BigFloat("0.54960853911436015786")im,
                BigFloat("0.16200952846773660904"),
                BigFloat("-0.01417981805211804396")im,
                BigFloat("-0.03415953916892111403"),
            )
            xv = [(x1a, x1b), (x2a, x2b), (x3a, x3b)]
            y = [
                BigFloat("0.99999999999999999928"),
                BigFloat("-0.99999999999999233987")im,
                BigFloat("-0.13549409636220703066"),
                zero(CBF),
                one(CBF),
            ]
        elseif (k == 4) # Equation (23)
            x1a = vcat(zero(CBF), one(CBF))
            x1b = vcat(zero(CBF), one(CBF))
            x2a = vcat(zeros(CBF, 2), one(CBF))
            x2b = vcat(zero(CBF), one(CBF), zero(CBF))
            # B_4^2
            x3a = vcat(
                zero(CBF),
                BigFloat("0.13340427306445612526")im,
                BigFloat("0.02022602029818310774"),
                BigFloat("-0.00674638241111650999")im,
            )
            x3b = x3a
            # (B_2 + A_6)*A_6 = (B_2 + B_3 + B_4^2)*(B_3 + B_4^2)
            b2 = vcat(
                zero(CBF),
                BigFloat("1.41183797496250375498")im,
                zero(CBF),
                BigFloat("-0.00866935318616372016")im,
            )
            b3 = vcat(
                BigFloat("2.69584306915332564689"),
                BigFloat("-1.35910926168869260391")im,
                BigFloat("-0.09896214548845831754"),
                BigFloat("0.01596479463299466666")im,
            )
            x4a = vcat((b2 + b3), one(CBF))
            x4b = vcat(b3, one(CBF))
            xv = [(x1a, x1b), (x2a, x2b), (x3a, x3b), (x4a, x4b)]
            # B_1 + (B_2 + A_6)*A_6 = B_1 + previous
            b1 = vcat(
                BigFloat("-6.26756985350202252845"),
                BigFloat("2.52179694712098096140")im,
                BigFloat("0.05786296656487001838"),
                BigFloat("-0.07766686408071870344")im,
            )
            y = vcat(b1, zero(CBF), one(CBF))
        elseif (k == 5) # Equation (25)
            x1a = vcat(zero(CBF), one(CBF))
            x1b = vcat(zero(CBF), one(CBF))
            x2a = vcat(zeros(CBF, 2), one(CBF))
            x2b = vcat(zero(CBF), one(CBF), zero(CBF))
            x3a = vcat(zeros(CBF, 3), one(CBF))
            x3b = vcat(zeros(CBF, 3), one(CBF))
            # B_1 * B_5
            x4a = vcat(
                zero(CBF),
                BigFloat("3") / BigFloat("25"),
                BigFloat("-0.00877476096879703859")im,
                BigFloat("-0.00097848453523780954"),
                zero(CBF),
            )
            x4b = vcat(
                zeros(CBF, 2),
                BigFloat("-0.123953695858283131480")im,
                BigFloat("-0.011202694841085592373"),
                BigFloat("-0.000012367240538259896")im,
            )

            # (B_3 + A_9)*A_9 = (B_3 + B_4 + previous)*(B_4 + previos)
            b3 = vcat(
                BigFloat("-2.58175430371188142440"),
                BigFloat("-1.73033278310812419209")im,
                BigFloat("-0.07673476833423340755"),
                BigFloat("-0.00261502969893897079")im,
                BigFloat("-0.00003400011993049304"),
            )
            b4 = vcat(
                BigFloat("2.92377758396553673559"),
                BigFloat("1.44513300347488268510")im,
                BigFloat("0.12408183566550450221"),
                BigFloat("-0.01957157093642723948")im,
                BigFloat("0.00002425253007433925"),
            )
            x5a = vcat((b3 + b4), one(CBF))
            x5b = vcat(b4, one(CBF))
            xv = [(x1a, x1b), (x2a, x2b), (x3a, x3b), (x4a, x4b), (x5a, x5b)]
            # B_2 + (B_3 + A_9)*A_9 = B_2 + previous
            b2 = vcat(
                zero(CBF),
                BigFloat("-0.66040840760771318751")im,
                BigFloat("-1.09302278471564897987"),
                BigFloat("0.25377155817710873323")im,
                BigFloat("0.00054374267434731225"),
            )
            y = vcat(b2, zero(CBF), one(CBF))
        end
        # Force convert to type
        xv = map(
            i -> (convert.(T, xv[i][1]), convert.(T, xv[i][2])),
            1:size(xv, 1),
        )
        y = convert.(T, y)
        return graph_degopt(xv, y)
    else  # Scalign-and-squaring phase
        s = k - 5
        degopt = Degopt(graph_bbcs_cheb_exp(5; T = T)[1])
        scale!(degopt, BigFloat("2.0")^(-big(s)))
        for i = 1:s
            square!(degopt)
        end
        return graph_degopt(degopt)
    end
end


using CausticsEngineering
using Test

"""
$(SIGNATURES)

"""
function testSquareMesh!()
    mesh = create_mesh(100, 50)

    println(mesh.nodeArray[1, 1])
    println(mesh.nodes[1])

    mesh.nodeArray[1, 1].x = 8
    println(mesh.nodeArray[1, 1])
    println(mesh.nodes[1])

    mesh.nodes[1].y += 12
    println(mesh.nodeArray[1, 1])
    println(mesh.nodes[1])
end

"""
$(SIGNATURES)

TO REFACTOR.
"""
function testSolidify!()
    println("Testing solidification")
    width = 100
    height = 100
    origMesh = create_mesh(width, height)

    for y = 1:height, x = 1:width
        x2 = (x - width / 2) / width
        y2 = (y - height / 2) / height
        value = x2 * x2 + y2 * y2
        origMesh.nodeArray[x, y].z = 15 - value * 25
    end

    save_stl!(origMesh, "./examples/testSolidify.obj")
    solidMesh = solidify(origMesh, 0)
    save_stl!(solidMesh, "./examples/testSolidify2.obj")
end

function testAreaTriangle(x1, y1, x2, y2, x3, y3)
    a = sqrt((x2 - x1)^2 + (y2 - y1)^2)
    b = sqrt((x3 - x2)^2 + (y3 - y2)^2)
    c = sqrt((x1 - x3)^2 + (y1 - y3)^2)
    s = (a + b + c) / 2.0

    surface_sq = (a + b + c) * (-a + b + c) * (a - b + c) * (a + b - c) / 16.0
    return surface_sq <= 1e-100 ? 0.0 : surface_sq
end

function doTestAreaTriangle(N = 10_000)
    X1 = 10_000 .* rand(N)
    Y1 = 10_000 .* rand(N)
    X2 = 10_000 .* rand(N)
    Y2 = 10_000 .* rand(N)
    X3 = 10_000 .* rand(N)
    Y3 = 10_000 .* rand(N)

    for x1 ∈ X1, y1 ∈ Y1, x2 ∈ X2, y2 ∈ Y2, x3 ∈ X3, y3 ∈ Y3

        _ = testAreaTriangle(x1, y1, x2, y2, x3, y3) > 0.0
    end
end

@testset "triangle area" begin
    for x1 ∈ 10_000 .* rand(10_000),
        y1 ∈ 10_000 .* rand(10_000),
        x2 ∈ 10_000 .* rand(10_000),
        y2 ∈ 10_000 .* rand(10_000),
        x3 ∈ 10_000 .* rand(10_000),
        y3 ∈ 10_000 .* rand(10_000)

        @test testAreaTriangle(x1, y1, x2, y2, x3, y3) >= 0.0
    end
end

@testset "More triangle area" begin
    @test area(0, 1, 1, 0, 0, 0)
end

v1 = CausticsEngineering.Vertex3D(0.0, 1.0, 0.0, 0.0, 0.0)
v2 = CausticsEngineering.Vertex3D(1.0, 0.0, 0.0, 0.0, 0.0)
v3 = CausticsEngineering.Vertex3D(0.0, 0.0, 0.0, 0.0, 0.0)
CausticsEngineering.area(v1, v2, v3)

v1 = CausticsEngineering.Vertex3D(0.0, 2.0, 0.0, 0.0, 0.0)
v2 = CausticsEngineering.Vertex3D(2.0, 0.0, 0.0, 0.0, 0.0)
v3 = CausticsEngineering.Vertex3D(0.0, 0.0, 0.0, 0.0, 0.0)
CausticsEngineering.area(v1, v2, v3)

fm = CausticsEngineering.FaceMesh(512, 512);
fm.toptriangles[1, 1]
CausticsEngineering.triangle3D(fm, 1, 1, :top)
CausticsEngineering.area(CausticsEngineering.triangle3D(fm, 1, 1, :top)...)
CausticsEngineering.area(CausticsEngineering.triangle3D(fm, 1, 1, :bottom)...)

CausticsEngineering.area(CausticsEngineering.triangle3D(fm, 126, 13, :top)...)
CausticsEngineering.area(CausticsEngineering.triangle3D(fm, 126, 13, :bottom)...)

v1 = [7250.768146661309, 4849.994680605491]
v2 = [1725.2665442814496, 8169.783790398901]
v3 = [6560.1725812433, 5264.912964907599]

a = sqrt(sum((v2 - v1) .^ 2))
b = sqrt(sum((v3 - v2) .^ 2))
c = sqrt(sum((v1 - v3) .^ 2))
s = (a + b + c) / 2.0

s * (s - a) * (s - b) * (s - c)
(a + b + c) * (-a + b + c) * (a - b + c) * (a + b - c) / 16.0

b + c
a

10^8 * (b + c - a)


"""
    Zarate

A Sugiyaqma style layout for DAGs and Sankey diagrams.
Based on

Zarate, D. C., Le Bodic, P., Dwyer, T., Gange, G., & Stuckey, P. (2018, April).
Optimal sankey diagrams via integer programming.
In 2018 IEEE Pacific Visualization Symposium (PacificVis) (pp. 135-139). IEEE.

# Fields
 - `time_limit::Dates.Period`: how long to spend trying alternative orderings.
   There are often many possible orderings that have the same amount of crossings.
   There is a chance that one of these will allow a better arrangement than others.
   Mostly the first solution tends to be very good.
   By setting a `time_limit > Second(0)`, multiple will be tried until the time limit is exceeded.
   (Note: that this is not a maximum time, but rather a limit that once exceeded no more
   attempts will be made.).
   If you have a `time_limit` greater than `Second(0)` set then the result is no longer deterministic.
   Note also that this is heavily affected by first call compilation time.
"""
Base.@kwdef struct Zarate <: AbstractLayout
    time_limit::Dates.Period = Dates.Second(0)
    ordering_solver::Any = ()->Cbc.Optimizer(; randomSeed=1, randomCbcSeed=1)
    arranging_solver::Any = ECOS.Optimizer
end

"""
    solve_positions(::Zarate, graph)

Returns:
 - `xs`: the xs coordinates of vertices in the layout
 - `ys`: the ys coordinates of vertices in the layout
 - `paths`: a Dict which, for each edge in `graph`, contains a Tuple of coordinate vectors (xs, ys).

The layout is calculated on a graph where dummy nodes can be added to the different layers.
As a result, plotting edges as straight lines between two nodes can result in
more crossings than optimal : edges should instead be routed through these different dummy nodes.
`paths` contains for each edge, a Tuple of vectors, representing that route through the
different nodes as x and y coordinates.

# Example:
```julia
using LightGraphs, Plots

g = random_orientation_dag(complete_graph(5))
xs, ys, paths = solve_positions(Zarate(), g)
scatter(xs, ys)
for e in edges(g)
    e_xs, e_ys = paths[e]
    plot!(e_xs, e_ys)
end
```
"""
function solve_positions(layout::Zarate, original_graph)
    graph = copy(original_graph)

    # 1. Layer Assigment
    layer2nodes = layer_by_longest_path_to_source(graph)
    is_dummy_mask, edge_to_path = add_dummy_nodes!(graph, layer2nodes)

    # 2. Layer Ordering
    start_time = Dates.now()
    min_total_distance = Inf
    min_num_crossing = Inf
    local best_pos
    ordering_model, is_before = ordering_problem(layout, graph, layer2nodes)
    for round in 1:typemax(Int)
        round > 1 && forbid_solution!(ordering_model, is_before)

        optimize!(ordering_model)
        # No need to keep banning solutions if not finding optimal ones anymore
        round > 1 && termination_status(ordering_model) != MOI.OPTIMAL && break
        num_crossings = objective_value(ordering_model)
        # we are not interested in arrangements that have more crossings, only in
        # alternatives with same number of crossings.
        num_crossings > min_num_crossing && break
        min_num_crossing = num_crossings
        order_layers!(layer2nodes, is_before)

        # 3. Node Arrangement
        xs, ys, total_distance = assign_coordinates(layout, graph, layer2nodes)
        if total_distance < min_total_distance
            min_total_distance = total_distance
            best_pos = (xs, ys)
        end
        Dates.now() - start_time > layout.time_limit && break
    end
    xs, ys = best_pos
    return xs[.!is_dummy_mask], ys[.!is_dummy_mask], Dict(edge => (xs[path], ys[path]) for (edge, path) in edge_to_path)
end

"""
    ordering_problem(::Zarate, graph, layer2nodes)

Formulates the problem of working out optimal ordering as a MILP.

Returns:
 - `model::Model`: the JuMP model that when optized will find the optimal ordering
 - `is_before::AbstractVector{AbstractVector{Variable}}`: the variables of the model,
   which once solved will have `value(is_before[n1][n2]) == true`
   if `n1` is best arrange before `n2`.
"""
function ordering_problem(layout::Zarate, graph, layer2nodes)
    m = Model(layout.ordering_solver)
    set_silent(m)

    node_is_before = Vector{Any}(undef, nv(graph))
    for (layer, nodes) in enumerate(layer2nodes)
        before = @variable(m, [nodes, nodes], Bin, base_name="befores_$layer")
        for n1 in nodes
            node_is_before[n1] = before[n1, :]
            for n2 in nodes
                n1 === n2 && continue
                # can't have n1<n2 and n2<n1
                @constraint(m, before[n1, n2] + before[n2, n1] == 1)
                for n3 in nodes
                    (n1 === n3 || n2 === n3) && continue
                    # at most two of these 3 hold
                    @constraint(m , before[n1, n2] + before[n2, n3] + before[n3, n1] <= 2)
                end
            end
        end
    end

    weights_mat = weights(graph)

    function crossings(src_layer)
        total = AffExpr(0)
        TDict{V} = Dict{Tuple{Int,Int}, V}
        crossing_vars = TDict{TDict{VariableRef}}()
        for src1 in src_layer, src2 in src_layer
            for dst1 in outneighbors(graph, src1), dst2 in outneighbors(graph, src2)
                # Can't cross if share end-point
                (src1 === src2 || dst1 === dst2) && continue
                # Zarate et al section "Further improvements with branching priorities"
                # we don't make this binary even though it is, because that makes the branchs happen
                # at wrong places. It only needs to branch the `before_*` variables.
                # the crossing will be binary as a result
                crossing = @variable(m, binary=false, integer=false, base_name="cross $src1-$dst2 x $src1-$dst2")
                get!(TDict{VariableRef}, crossing_vars, (src1, dst1))[src2, dst2] = crossing
                # two edges cross if the src1 is before scr2; but dst1 is after dest2
                @constraint(m, node_is_before[src1][src2] + node_is_before[dst2][dst1] - 1 <= crossing)

                # for Sankey diagrams we minimise not just crossing but area of crossing
                # treating the weights of the graph as the widths of the line and ignoring skew
                w1 = weights_mat[src1, dst1]
                w2 = weights_mat[src2, dst2]
                area = w1*w2
                add_to_expression!(total, area, crossing)
            end
        end

        # perfomance optimizations based on Zarate eq 8 and 9
        for ((u1,v1), other_edges) in crossing_vars
            for ((u2, v2), c1) in other_edges
                c2 = crossing_vars[(u2,v2)][(u1,v1)]
                @constraint(m, c1==c2)  # Zarate eq 8

                cx = try_get(try_get(crossing_vars, (u1,v2)), (u2,v1))
                cx === nothing && continue
                @constraint(m, c1 + cx == 1)# Zarate eq 9
            end
        end
        return total
    end

    @objective(m, Min, sum(crossings, layer2nodes))
    #@show m
    return m, node_is_before
end

function order_layers!(layer2nodes, is_before)
    # Cunning trick: we need to go from the `before` matrix to an actual order list
    # we can do this by sorting when having `lessthan` read from the `before` matrix
    is_before_func(n1, n2) = Bool(value(is_before[n1][n2]))
    for layer in layer2nodes
        sort!(layer; lt=is_before_func)
    end
    return layer2nodes
end

"""
    forbid_solution!(m, is_before)

Modifies the JuMP model `m` to forbid the solution present in `is_before`.
`is_before` must contain solved variables.
"""
function forbid_solution!(m, is_before)
    cur_trues = AffExpr(0)
    total_trues = 0
    for var_list in is_before
        for var in var_list
            if value(var) > 0.5
                add_to_expression!(cur_trues, var)
                total_trues += 1
            end
        end
    end
    # for it to be a different order some of the ones that are currently true must swap to being false
    @constraint(m, sum(cur_trues) <= total_trues - 1)

    return m
end

"""
    assign_coordinates(graph, layer2nodes)

Works out the `x` and `y` coordinates for each node in the `graph`.
This is via formulating the problem as a QP, and minimizing total distance
of links.
It maintains the order given in `layer2nodes`s.
returns: `xs, ys, total_distance`
"""
function assign_coordinates(layout, graph, layer2nodes)
    m = Model(layout.arranging_solver)
    set_silent(m)
    set_optimizer_attribute(m, "print_level", 0)  # TODO this can be deleted once the version of IPOpt that actually supports `set_silent` is released

    node2y = Dict{Int, VariableRef}()
    for (layer, nodes) in enumerate(layer2nodes)
        first_node, other_nodes = Iterators.peel(nodes)
        prev_y = node2y[first_node] = @variable(m, base_name="y_$first_node")
        for node in other_nodes
            # each must be greater than the last, with a gap of 1 unit
            y = @variable(m, base_name="y_$node")
            @constraint(m, prev_y + 1.0 <= y)
            prev_y = node2y[node] = y
        end
    end

    all_distances = AffExpr[]
    # minimize link distances
    for cur in vertices(graph)
        for out in outneighbors(graph, cur)
            distance = (node2y[cur] - node2y[out])
            push!(all_distances, distance)
        end
        # to prevent going way off-scale, also keep things near x-axis
    end
    # for first layer minimize distance to origin
    for cur in first(layer2nodes)
        distance = node2y[cur]
        push!(all_distances, distance)
    end
    @variable(m, total_distance)
    distance_cone = [1; total_distance; all_distances]
    @constraint(m, distance_cone in RotatedSecondOrderCone())
    @objective(m, Min, total_distance);
    optimize!(m)
    #termination_status(m)
    score = objective_value(m)

    ###
    xs = Vector{Float64}(undef, nv(graph))
    ys = Vector{Float64}(undef, nv(graph))
    for (xi, nodes) in enumerate(layer2nodes)
        for node in nodes
            xs[node] = xi
            ys[node] = value(node2y[node])
        end
    end
    return xs, ys, score
end


# kernels for normal vectors from:
# Snape, Zafeiriou, "Kernel-PCA Analysis of Surface Normals for Shape-from-Shading", 2014

using Debug

# IPKernel = Inner Product kernel - this is linear kernel w/o centering
function IPKernel(X1,X2)
	K = X1*X2'
	#K = zeros(size(X1,1), size(X2,1))
	#d = size(X1,2)
	#for i=1:size(X1,1)
	#	for j=1:size(X2,1)
	#		# sum of inner products of the normals for each "pin list"
	#		for k=1:3:d
	#			K[i,j] += (X1[i,k:k+2]' * X2[j,k:k+2])[1]
	#		end
	#	end
	#end
	return K
end

# Given a list of points in feature space (projected onto PCA subspace and back)
# this function computes what the corresponding pre-image points would be
function InverseIPKernel(Y)
	n = size(Y,1)
	d = size(Y,2)
	Xp = zeros(size(Y))
	for i=1:n
		for j=1:3:d
			Xp[i,j:j+2] = Y[i,j:j+2] / norm(Y[i,j:j+2])
		end
	end
	return Xp
end

# Compute the Azimuthal Equidistant Projection (AEP) feature vectors
function AEPFeature(X,mus)
	feature = zeros(size(mus))
	d = length(X)
	muk = 1
	for k=1:3:d
		mphi = mus[muk]
		mtheta = mus[muk+1]

		x = X[k]
		y = X[k+1]
		z = X[k+2]
		phi = atan( y/x )
		theta = pi/2 - asin( z )
		cosc = sin(mtheta)*sin(theta) + cos(mtheta)*cos(theta)*cos(phi-mphi)
		c = acos(cosc)
		sinc = sin(c)
		k = c / sinc

		tx = k * cos(theta) * sin(phi - mphi)
		ty = k * (cos(mtheta) * sin(phi) - sin(mphi)*cos(theta)*cos(phi-mphi))
		feature[muk] = tx
		feature[muk+1] = ty

		muk += 2
	end
	feature
end
# interface that computes mu as needed.
function AEPFeature(X)
	mus = LinearMean(X)
	xaep1 = AEPFeature(X[1,:],mus)
	F = zeros(size(X,1), length(xaep1))
	for i=1:size(X,1)
		F[i,:] = AEPFeature(X[i,:],mus)
	end
	F
end

# Azimuthal Equidistant Projection (AEP) kernel
# X1,X2 - data sets
# mus - mean for each point as elevation and azimuthal angles
function AEPKernel(X1,X2,mus)
	K = zeros(size(X1,1), size(X2,1))
	d = size(X1,2)
	for i=1:size(X1,1)
		feature1 = AEPFeature(X1[i,:],mus)
		for j=1:size(X2,1)
			feature2 = AEPFeature(X2[i,:],mus)
			K[i,j] = (feature1 * feature2')[1]
		end
	end
	return K
end

# interface that computes mu as needed.
function AEPKernel(X1,X2)
	if X1 == X2
		mus = LinearMean(X1)
	else
		mus = LinearMean([X1; X2])
	end
	K = AEPKernel(X1,X2,mus)
	K,mus
end

# Azimuthal Equidistant Projection (AEP) covariance - perform computation direction in feature space
# NOTE: use this when dimension of feature < n
#       (for AEP dimension of feature <=> dimension of input)
# X1,X2 - data sets
# mus - mean for each point as elevation and azimuthal angles
function AEPCovariance(X,mus)
	ftmp = AEPFeature(X[1,:],mus)
	fl = length(ftmp)
	cmat = zeros(fl,fl) # it appears that this will be smaller thann
	d = size(X,2)
	n = size(X,1)
	for i=1:size(X1,1)
		feature = AEPFeature(X[i,:],mus)
		cmat += feature * feature' # outer product
	end
	cmat / n
end

# interface that computes mu as needed.
function AEPCovariance(X)
	mus = LinearMean(X)
	cmat = AEPCovariance(X1,X2,mus)
	cmat,mus
end

# Inverse Azimuthal Equidistant Projection (AEP) kernel
# Y - data in feature space
# mus - mean for each point as elevation and azimuthal angles
function InverseAEPKernel(Y,mus)
	n = size(Y,1)
	d = size(Y,2)
	X = zeros(n,(d/2)*3)
	for i=1:n
		Xk = 1
		for k=1:2:d
			mphi = mus[k]
			mtheta = mus[k+1]

			tx = Y[i,k]
			ty = Y[i,k+1]

			c = sqrt( tx^2 + ty^2 )
			if abs(mtheta-pi/2) < eps(typeof(mtheta)) # mtheta == pi/2
				psi = -tx/ty
			elseif abs(mtheta+pi/2) < eps(typeof(mtheta)) #mtheta == -pi/2
				psi = tx/ty
			else
				psi = (tx * sin(c)) / (c * cos(mtheta) * cos(c) - ty * sin(mtheta) * sin(c))
			end

			theta = asin(cos(c) * sin(mtheta) - (1/c) * ty * sin(c) * cos(mtheta))
			phi = mphi + atan(psi)

			nz = sin(pi/2 - theta)
			# ny/nx = tan(phi)
			h = sqrt(1 - nz^2) # hypotenuse
			nx = h * cos(phi)
			ny = h * sin(phi)

			# make sure it's unit length:
			len = sqrt(nx^2 + ny^2 + nz^2)
			nx /= len
			ny /= len
			nz /= len

			X[i,Xk] = nx
			X[i,Xk+1] = ny
			X[i,Xk+2] = nz

			Xk += 3
		end
	end
	X
end

# compute the linear mean and then reproject onto the sphere
function LinearMean(X)
	mu = mean(X,1)
	d = size(X,2)
	for i=1:3:d
		mu[i:i+2] /= norm(mu[i:i+2])
	end
	mu
end

# project onto tangent plane of mus
# see Srivastava,Jermyn,Joshi, "Riemannian Analysis of Probability Density Functions with Applications in Vision", equation (9)
# X - matrix: (num_of_data, 3*num_of_positions) - so each row has a normal for each positions, and each column is the list of normals for that position
# mus - vector: (3*num_of_positions) - tangent point for each position
# undirected - true => ignore direction (sign of vector)
function LogMap(X,mus, undirected=false)
	dim = 3
	n = size(X,1)
	d = size(X,2)
	logmap = zeros(size(X))
	for position = 1:dim:d
		Xp = X[:,position:position+2]
		mu = reshape(mus[position:position+2],1,dim)
		if undirected
			# flip for consistent (undirected) orientation
			flip = find((Xp * mu) .< 0)
			Xp[flip,:] *= -1
		end
		inner = max( -1, min(1, Xp * mu ) )
		u = Xp - ((inner*ones(1,dim)) .* (ones(n,1)*mu))
		unorm = sqrt(sum(u.^2,2))
		coef = acos(inner) ./ max(eps(typeof(unorm[1])), unorm)
		logmap[:,position:position+2] = coef .* u
	end
	logmap
end

# project onto sphere from tangent plane of mus
# see Srivastava,Jermyn,Joshi, "Riemannian Analysis of Probability Density Functions with Applications in Vision", equation (8)
# X - matrix: (num_of_data, 3*num_of_positions) - so each row has a normal for each positions, and each column is the list of normals for that position
# mus - vector: (3*num_of_positions) - tangent point for each position
# undirected - true => ignore direction (sign of vector)
@debug function ExpMap(X,mus,undirected=false)
	dim = 3
	n = size(X,1)
	d = size(X,2)
	expmap = zeros(size(X))
	for position = 1:dim:d
		Xp = X[:,position:position+2]
		mu = reshape(mus[position:position+2],1,dim)

		dist = sqrt(sum( Xp.^2, 2))

		if undirected
			# flip to shortest orientation (for consistent undirected vectors)
			shorten = find(dist .> pi/2)
			Xp[shorten,:] ./= (dist[shorten]*ones(1,dim))
			dist[shorten] .-= pi/2
			Xp[shorten,:] = -Xp[shorten,:] .* (dist[shorten]*ones(1,dim))
		end

		@bp
		expmap[:,position:position+2] = (cos(dist)*mu) + (sin(dist) * (Xp ./ (max(eps(typeof(dist[1])), dist)*ones(1,dim))))
	end
	expmap
end

# Frechet Mean using gradient descent
# X - the data
# mu0 - initial guess at mean
# stepsize - step size in gradient descent
# maxiters - maximum number of steps
# stopping - convergence of mu below this value
function FrechetMean(X, mu0, stepsize, maxiters, stopping)
	n = size(X,1)
	d = size(X,2)
	dim = 3
	mu = copy(mu0)

	for iter = 1:maxiters
		gradient = -mean(LogMap(X, mu),1)

		lastmu = copy(mu)
		mu = ExpMap(-stepsize * gradient, lastmu )

		if norm(mu-lastmu) < stopping
			break
		end
	end

	mu
end

# feature vector for PGA Kernel
function PGAFeature(X,mus)
	LogMap(X,mus)
end

# Kernel Based on PGA given the intrinsic means
function PGAKernel(X1,X2,mus)
	d = length(mus)
	K = zeros(size(X1,1),size(X2,1))
	for i=1:size(X1,1)
		feature1 = PGAFeature(X1[i,:],mus)
		for j=1:size(X2,1)
			feature2 = PGAFeature(X2[i,:],mus)
			K[i,j] = feature1' * feature2
		end
	end
	K
end

# Kernel Based on PGA - interface that computes mus as needed
function PGAKernel(X1,X2)

	stepsize = 1e-1
	maxiters = 1000
	stopping = 1e-7

	if X1 == X2
		mus = LinearMean(X1)
		mus = FrechetMean(X1,mus,stepsize,maxiters,stopping)
	else
		mus = LinearMean([X1;X2])
		mus = FrechetMean([X1;X2],mus,stepsize,maxiters,stopping)
	end

	K = PGAKernel(X1,X2,mus)
	K, mus
end

# Covariance Based on PGA given the intrinsic means - compute in feature space directly
# NOTE: use this when PGA feature dimension < size of data
#       for PGA feature dimension <=> input dimension
function PGACovariance(X,mus)
	d = length(mus)
	n = size(X,1)
	ftmp = PGAFeature(X[1,:],mus)
	fl = length(ftmp)
	cmat = zeros(fl,fl)
	for i=1:n
		feature = PGAFeature(X[i,:],mus)
		cmat += feature * feature'
	end
	cmat / n
end

# Covariance Based on PGA - interface that computes mus as needed
function PGACovariance(X)

	stepsize = 1e-1
	maxiters = 1000
	stopping = 1e-7

	if X1 == X2
		mus = LinearMean(X1)
		mus = FrechetMean(X1,mus,stepsize,maxiters,stopping)
	else
		mus = LinearMean([X1;X2])
		mus = FrechetMean([X1;X2],mus,stepsize,maxiters,stopping)
	end

	cmat = PGACovariance(X,mus)
	cmat, mus
end

# Inverse PGA kernel
function InversePGAKernel(Y,mus)
	X = ExpMap(Y,mus)
end

# Spherical Feature
function SphericalFeature(X)
	n = size(X,1)
	d = size(X,2)
	dim = 3
	feature = zeros(n,int(4*(d/dim)))
	for i=1:n
		fk = 1
		for k=1:dim:d
			nrm = norm(X[i,k:k+2])
			if nrm > 0
				feature[i,fk:fk+2] = X[i,k:k+2] / nrm
				feature[i,fk+3] = sqrt( 1 - feature[i,fk+2]^2 )
			end
			fk += 4
		end
	end
	feature
end
function SphericalFeatureInverse(Phi)
	dim = 3
	X = zeros(size(Phi,1),int(dim*(size(Phi,2)/4)))
	n = size(X,1)
	d = size(X,2)
	for i=1:n
		fk = 1
		for k=1:dim:d
			gx = Phi[i,fk]
			gy = Phi[i,fk+1]
			gz = Phi[i,fk+2]
			sgz = Phi[i,fk+3]
			#rho = atan2( (gy/sqrt((gx^2)+(gy^2))) , (gx/sqrt((gx^2)+(gy^2))) )
			rho = atan2( gy , gx )
			#psi = atan2( (sgz/sqrt((gz^2)+(sgz^2))) , (gz/sqrt((gz^2)+(sgz^2))) )
			psi = atan2( sgz, gz )
			#v = [(cos(psi)*sin(rho)), (sin(psi)*sin(rho)), cos(psi)]
			#gzh = sqrt(gz^2 + sgz^2)
			#v = [(cos(psi)*sin(rho)), (sin(psi)*sin(rho)), cos(psi)*gzh]
			v = [gx,gy,gz]
			X[i,k:k+2] = v / norm(v)
			fk += 4
		end
	end
	X
end

# Spherical Kernel (aka Cosine kernel)
function SphericalKernel(X1,X2)
	Phi1 = SphericalFeature(X1)
	Phi2 = SphericalFeature(X2)
	K = Phi1 * Phi2'
end
function SphericalKernelLowMem(X1,X2)
	K = zeros(size(X1,1),size(X2,1))
	for i=1:size(X1,1)
		feature1 = SphericalFeature(X1[i,:])
		feature1[find(isnan(feature1))] = 0
		for j=1:size(X2,1)
			feature2 = SphericalFeature(X2[i,:])
			feature2[find(isnan(feature2))] = 0
			K[i,j] = (feature1 * feature2')[1]
		end
	end
	return K
end

# Equivalent to Spherical Kernel above, but uses cos,acos,atan instead of constructing feature vectors..
function SphericalCosKernel(X1,X2)
	error("does not match SphericalKernel() result")
	d = size(X1,2)
	dim = 3
	dpr = int(d / dim)
	K = zeros(size(X1,1),size(X2,1))
	for i=1:size(X1,1)
		phi1 = zeros(dpr)
		theta1 = zeros(dpr)
		dpri = 1
		for k = 1:dim:d
			v = X1[i,k:k+2] / norm(X1[i,k:k+2])
			phi1[dpri] = atan( v[2]/v[1] )
			theta1[dpri] = acos( v[3] )
			dpri += 1
		end
		for j=1:size(X2,1)
			phi2 = zeros(dpr)
			theta2 = zeros(dpr)
			dpri = 1
			for k = 1:dim:d
				v = X2[i,k:k+2] / norm(X2[i,k:k+2])
				phi2[dpri] = atan( v[2]/v[1] )
				theta2[dpri] = acos( v[3] )
				dpri += 1
			end
			K[i,j] = sum( cos(phi1-phi2) ) + sum( cos(theta1-theta2) )
		end
	end
	K
end
# Covariance Based on Spherical kernel
# NOTE: use this when [spherical feature dimension] < [size of data]
#       for [spherical feature dimension] <=> (4/3)*[input dimension]
function SphericalCovariance(X)
	n = size(X,1)
	Phi = SphericalFeature(X)
	cmat = (Phi'*Phi) / n
end
function SphericalCovarianceLowMem(X)
	d = size(X,2)
	n = size(X,1)
	ftmp = SphericalFeature(X[1,:])
	fl = length(ftmp[:])
	cmat = zeros(fl,fl)
	for i=1:n
		feature = SphericalFeature(X[i,:])
		cmat += feature * feature'
	end
	cmat / n
end


using MultivariateMeasures
using Test

@testset "MultivariateMeasures.jl" begin

    @testset "LKJCholesky" begin
        D = LKJCholesky{(:k,:η)}
        par = transform(asparams(D, (k=4,)), randn(1))
        d = D(merge((k=4,),par))
        # @test params(d) == par

        η  = par.η
        logη = log(η)

        y = rand(d)
        η = par.η
        ℓ = logdensity(LKJCholesky(4,η), y)
        @test ℓ ≈ logdensity(LKJCholesky(k=4,logη=logη), y)
    end
end


#--------------------------------------------------------------------
# Spacetime Discretization methods in Julia
# Soham 01-2018
#--------------------------------------------------------------------

function distribute{T<:Integer}(fbndr::Function, fbndc::Function, frhs::Function, Nx::T, Ny::T, M::T)::Dict{Array{Int,1}, Patch}
    dop = derivOP(Nx, Ny)/M^4
    bop = boundaryOP(Nx, Ny)
    dbase  = Dict{Array{Int, 1}, Patch}()
    for i in 2:2M, k in i-min(i-1,M):min(i-1,M)
        loc  = [k, i-k]
        rhs  = RHS(frhs, Nx, Ny, M, loc)
        bndx = (loc[2]==1) ? (getPatchIC(fbndr, 0, Nx, M, loc[1])) : (getPatchBnd(dbase[loc-[0,1]], 0))
        bndy = (loc[1]==1) ? (getPatchIC(fbndc, 1, Ny, M, loc[2])) : (getPatchBnd(dbase[loc-[1,0]], 1))
        dbase[loc] = calcPatch(bndx, bndy, rhs, dop, bop, loc)
    end
    return dbase
end


struct Diffusion{T1<:Real,T2<:Real,T3<:Real} <: ContinuousUnivariateDistribution
    D::T1
    δ::T2
    ϵ::T3
    Diffusion{T1,T2,T3}(D::T1, δ::T2, ϵ::T3) where {T1<:Real,T2<:Real,T3<:Real} =
        new{T1,T2,T3}(D, δ, ϵ)
end

function Diffusion(
    D::T1,
    δ::T2,
    ϵ::T3;
    check_args = true,
) where {T1<:Real,T2<:Real,T3<:Real}
    check_args && @check_args(Diffusion, D > zero(D) && δ > zero(δ) && ϵ >= zero(ϵ))
    return Diffusion{T1,T2,T3}(D, δ, ϵ)
end

@distr_support Diffusion 0 +Inf

# minimum(d::Diffusion) = 0.0
# maximum(d::Diffusion) = +Inf
cdf(d::Diffusion, x::Real) = 1 - exp(-x^2 / 4 * (d.D * d.δ + d.ϵ^2))
quantile(d::Diffusion, p) = sqrt(-4 * (d.D * d.δ + d.ϵ^2) * log(1 - p))
pdf(d::Diffusion, x::Float64) =
    0 ≤ x ? x / (2 * (d.D * d.δ + d.ϵ^2)) * exp(-x^2 / (4 * (d.D * d.δ + d.ϵ^2))) : zero(x)
logpdf(d::Diffusion, x::AbstractVector{<:Real}) =
    0 ≤ x ? log(x) - log(2 * (d.D * d.δ + d.ϵ^2)) - (x^2 / (4 * (d.D * d.δ + d.ϵ^2))) :
    zero(x)
# Distributions.pdf(d::Diffusion, x::AbstractArray{<:Real}) =
#     0 ≤ x ? x / 2(d.D * d.δ + d.ϵ^2) * exp(-x^2 / 4(d.D * d.δ + d.ϵ^2)) : zero(x)
quantile(d::Diffusion, x::AbstractVector{<:Real}) =
    sqrt(-4 * (d.D * d.δ + d.ϵ^2) * log(1 - x))
rand(d::Diffusion, rng::AbstractVector{<:Real}) = quantile(d::Diffusion, rng)

### MLE fitting
struct DiffusionStats <: SufficientStats
    r::Real # (weighted) sum of sqared r
    w::Float64 # total sample weight
    δ::Float64
    ϵ::Float64

    DiffusionStats(r::Real, w::Real, δ::Float64, ϵ::Float64) = new(r, w, δ, ϵ)
end

function Distributions.suffstats(
    ::Type{<:Diffusion},
    x::AbstractArray
)
    n = length(x)
    r = 0.0
    for i = 1:n
        @inbounds r += x[i]^2
    end
    DiffusionStats(r, Float64(n), 0.001, 0.0)
end

function Distributions.suffstats(d::Diffusion, x::AbstractArray)
    n = length(x)
    r = 0.0
    for i = 1:n
        @inbounds r += x[i]^2
    end
    DiffusionStats(r, Float64(n), d.δ, d.ϵ)
end

function Distributions.suffstats(::Type{<:Diffusion}, x::AbstractArray{T}, w::AbstractArray{Float64}) where T <: Real
    n = length(x)
    if length(w) != n
        throw(DimensionMismatch("Inconsistent argument dimensions."))
    end
    tw = w[1]
    r = 0.0
    for i = 1:n
        @inbounds wi = w[i]
        @inbounds xi = x[i]
        r += wi * xi^2
        tw += wi
    end
    DiffusionStats(r, tw, 0.001, 0.0)
end

function Distributions.fit_mle(::Type{<:Diffusion}, ss::DiffusionStats)
        # D = (ss.r - 4 * ss.w * ss.δ * ss.ϵ^2) / (4 * ss.w * ss.δ)
    D = ss.r / (4 * ss.w * ss.δ) - ss.ϵ^2 / ss.δ
    Diffusion(D, ss.δ, ss.ϵ)
end


#==============================================================================
    Code for solving the Hamiltonian Jacobi Bellman for
	   a Ramsey Model with a diffusion process for capital

 	Based on  Matlab code from Ben Moll:
        http://www.princeton.edu/~moll/HACTproject.htm

        Updated to julia 1.0.0
==============================================================================#

using Distributions, Plots, LinearAlgebra, SparseArrays
using Random

Random.seed!(12364)

γ= 2.0 #gamma parameter for CRRA utility
ρ = 0.05 #the discount rate
α = 1/3 # the curvature of the production function (cobb-douglas)
δ = 0.05 # the depreciation rate
σ = 0.5 #variance term
n=0.51
T = 1000 # set the number of external time periods for misspecification loop

# Add in the misspecification
σ_g =0.02
Σ_g=[σ_g] #create cell with the agent's guess for σ

#=============================================================================
 	k our capital follows a process that depends on
            dk = (s⋅f(k) - (δ-σ^2)k)dt - σ⋅k⋅dW_[t]
        Where W is a Wiener process
=============================================================================#
#K_starting point, for mean of z process

k_st = ((α)/(ρ+δ+n-σ^2))^(1/(1-α))

# create the grid for k
H = 100 #number of points on grid
k_min = 0.1*k_st # min value
k_max = 10*k_st # max value
k = LinRange(k_min, k_max, H)
dk = (k_max-k_min)/(H-1)
k_min_1=k_min
k_max_1=k_max

#create matrices for k and z
kk = k*1

# use Ito's lemma to find the drift and variance of our optimization equation

# parameters for our loops
max_it = 100
ε = 0.1^(6)
Δ = 1000

# create a random distribution to pull from for the updating rule
Dist = Bernoulli(.45)

# set up all of these empty matrices for the misspecified version later on
Vaf_2, Vab_2, c_2 = [zeros(H,1) for i in 1:3]
Val_2 =[]
savings=[]

# set up all of these empty matrices for the steady state algorithm
Vaf_1, Vab_1, c_1 = [zeros(H,1) for i in 1:3]

y = pdf.(Normal(.99, σ), k) #check
plot(k,y, grid=false,
		xlabel="k", ylabel="Probability",
		legend=false, color="purple", title="PDF of k")
#png("PDF_k")

#============================================================================#
# Now it's time to solve the model, first put in a guess for the value function
v0 = (kk.^α).^(1-γ)/(1-γ)/ρ
v_1=v0

maxit= 30 #set number of iterations (only need 6 to converge)
dist = [] # set up empty array for the convergence criteria

for j = 1:maxit
    V=v_1

    #Now set up the forward difference

    Vaf_1[1:H-1,:] = (V[2:H, :] - V[1:H-1,:])/dk
    Vaf_1[H,:] .= (k_max.^α - (δ+n-σ^2).*k_max).^(-γ) # imposes a constraint

    #backward difference
    Vab_1[2:H,:] = (V[2:H, :] - V[1:H-1,:])/dk
    Vab_1[1,:] .= (k_min.^α - (δ+n-σ^2).*k_min).^(-γ)

    #I_concave = Vab .> Vaf # indicator for whether the value function is concave

    # Consumption and savings functions
    cf = Vaf_1.^(-1/γ)
    drf =(1 .-(cf/kk.^α))*kk.^α-(δ+n-σ^2).*kk

    # consumption and saving backwards difference
    cb = Vab_1.^(-1.0/γ)
    drb = (1 .-(cb/k.^α))*kk.^α-(δ+n-σ^2).*kk
    #println(sb)
    #consumption and derivative of the value function at the steady state

    c0 = kk.^α - (δ+n-σ^2).*kk
    Va0 = c0.^(-γ)

    # df chooses between the forward or backward difference

    If = drf.>0 # positive drift will ⇒ forward difference
    Ib = drb.<0 # negative drift ⇒ backward difference
    I0=(1.0.-If-Ib) # at steady state

    Va_upwind = Vaf_1.*If + Vab_1.*Ib + Va0.*I0 # need to include SS term

    global c_1 = Va_upwind.^(-1/γ)
    u = (c_1.^(1-γ))/(1-γ)

    # Now to constuct the A matrix
    X = -min.(drb, 0)/dk + (σ.*kk/(2*dk))
    Y = -max.(drf, 0)/dk + min.(drb, 0)/dk - (σ.*kk/(dk))
    Z = max.(drf, 0)/dk + (σ.*kk/(2*dk))

    updiag = 0
    updiag =[updiag; Z[1:H-1,1]; 0]
    updiag =(updiag[:])

    centerdiag=reshape(Y, H, 1)
    centerdiag = (centerdiag[:]) # for tuples

   lowdiag = X[2:H, 1]
   lowdiag=(lowdiag)

   # spdiags in Matlab allows for automatic trimming/adding of zeros
       # spdiagm does not do this
  A = sparse(Diagonal(centerdiag))+ [zeros(1, H); sparse(Diagonal(lowdiag)) zeros(H-1,1)] + sparse(Diagonal(updiag))[2:end, 1:H] # trim first element

  B = (1/Δ + ρ)*sparse(I,H,H) - A

  u_stacked= reshape(u, H, 1)
  V_stacked = reshape(V,H, 1)

  b = u_stacked + (V_stacked./Δ)

  V_stacked = B\b

  V = reshape(V_stacked, H, 1)

  V_change = V-v_1

  global v_1= V
  # need push function to add to an already existing array
  push!(dist, findmax(abs.(V_change))[1])
  if dist[j].< ε
      println("Value Function Converged Iteration=")
      println(j)
      break
  end

end

# calculate the savings for kk
ss_1 = kk.^α - (δ+n-σ^2).*kk

#==================== Now resolve the model with σ misspecified =======#
for t in 1:T
    # Now it's time to solve the model, first put in a guess for the value function
    global v0 = (kk.^α).^(1-γ)/(1-γ)/ρ
    v_2=v0
    global dist = [] # needs to be in the time loop to work properly
    indicator = rand(Dist)
    for j = 1:maxit
        V=v_2

        #Now set up the forward difference

        Vaf_2[1:H-1,:] = (V[2:H, :] - V[1:H-1,:])/dk
        Vaf_2[H,:] .= (k_max.^α - (δ+n-σ_g^2).*k_max).^(-γ) # imposes a constraint

        #backward difference
        Vab_2[2:H,:] = (V[2:H, :] - V[1:H-1,:])/dk
        Vab_2[1,:] .= (k_min.^α - (δ+n-σ_g^2).*k_min).^(-γ)

        #I_concave = Vab .> Vaf # indicator for whether the value function is concave

        # Consumption and savings functions
        cf = Vaf_2.^(-1/γ)
        drf =(1 .-(cf/kk.^α))*kk.^α-(δ+n-σ_g^2).*kk

        # consumption and saving backwards difference
        cb = Vab_2.^(-1.0/γ)
        drb = (1 .-(cb/k.^α))*kk.^α-(δ+n-σ_g^2).*kk
        #println(sb)
        #consumption and derivative of the value function at the steady state

        c0 = kk.^α - (δ+n-σ_g^2).*kk
        Va0 = c0.^(-γ)

        # df chooses between the forward or backward difference

        If = drf.>0 # positive drift will ⇒ forward difference
        Ib = drb.<0 # negative drift ⇒ backward difference
        I0=(1.0.-If-Ib) # at steady state

        Va_upwind = Vaf_2.*If + Vab_2.*Ib + Va0.*I0 # need to include SS term

        global c_2 = Va_upwind.^(-1/γ)
        u = (c_2.^(1-γ))/(1-γ)

        # Now to constuct the A matrix
        X = -min.(drb, 0)/dk + (σ_g.*kk/(2*dk))
        Y = -max.(drf, 0)/dk + min.(drb, 0)/dk - (σ_g.*kk/(dk))
        Z = max.(drf, 0)/dk + (σ_g.*kk/(2*dk))

        updiag = 0
        updiag =[updiag; Z[1:H-1,1]; 0]
        updiag =(updiag[:])

        centerdiag=reshape(Y, H, 1)
        centerdiag = (centerdiag[:]) # for tuples

       lowdiag = X[2:H, 1]
       lowdiag=(lowdiag)

       # spdiags in Matlab allows for automatic trimming/adding of zeros
           # spdiagm does not do this
      A = sparse(Diagonal(centerdiag))+ [zeros(1, H); sparse(Diagonal(lowdiag)) zeros(H-1,1)] + sparse(Diagonal(updiag))[2:end, 1:H] # trim first element

      B = (1/Δ + ρ)*sparse(I,H,H) - A

      u_stacked= reshape(u, H, 1)
      V_stacked = reshape(V,H, 1)

      b = u_stacked + (V_stacked./Δ)

      V_stacked = B\b

      V = reshape(V_stacked, H, 1)

      V_change = V-v_2

      v_2= V
      # need push function to add to an already existing array
      push!(dist, findmax(abs.(V_change))[1])
      if dist[j].< ε
        push!(Val_2, v_2)
        break
      end
    end
    push!(Σ_g, σ_g)
    # calculate the savings for kk
    ss_2 = kk.^α - (δ+n-σ_g^2).*kk
    push!(savings, ss_2)
    # add in updating
    if indicator ==1
		σ_noise = σ + rand(Normal(0,.1),1)[1] # generate a noisy observation
        global σ_g = σ_g + .01(σ_noise-σ_g) # use RLS to update using the noisy obs.
    end
    println(t)
end

# Plot the savings vs. k
plot(k[1:H-1], c_2[1:H-1], grid=false,
		xlabel="k", ylabel="c(k)",
        xlims=(k[1],k[end]),legend=:bottomright,
		label="Guess", title="Optimal Consumption Policies")
plot!(k[1:H-1],c_1[1:H-1], label="True Value", line=:dash)
#png("OptimalCons")

plot(k, savings[1], grid=false,
		xlabel="k", ylabel="s(k)",
        xlims=(k[1],k[end]),label="Period 1",
        title="Optimal Savings Policies",
		color=:hotpink, legend=:topleft, line=:dashdot)
plot!(k, zeros(H,1), color=:black, label="")
plot!(k,savings[100], label="Period 100", color=:blue,line=:dot)
plot!(k,savings[500], label="Period 500", color=:green, line=:dashdotdot)
plot!(k,savings[1000], label="Period 1,000", color=:orange)
plot!(k,ss_1, label="True Value", line=:dash, color=:black, legend=:bottomleft)
png("OptimalSavings_2")

plot(k, savings, grid=false,
		xlabel="k", ylabel="s(k)",
        xlims=(k[1],k[end]),label="", title="Optimal Savings Policies")
plot!(k, zeros(H,1), line=:dash, color=:black, label="", legend=:topleft)
plot!(k,ss_1, color=:black, label="True Value", legend=:bottomleft)
#png("OptimalSavings_All")

plot(k,Val_2[1], label="Period 1 ", color=:hotpink, line=:dashdot)
plot!(k,Val_2[100], label="Period 100", color=:blue, line=:dot)
plot!(k,Val_2[500], label="Period 500", color=:green, line=:dashdotdot)
plot!(k,Val_2[end], label="Period 1,000", color=:orange )
plot!(k,v_1, grid=false,
		xlabel="k", ylabel="V(k)",
		xlims=(k[1],k[end]), title="Value Functions",
        legend=:bottomright, color=:black, line=:dash,
        label="True Value")
png("Value_function_vs_k_2")

plot(k, Val_2, grid=false,
		xlabel="k", ylabel="V(k)",
		xlims=(k[1],k[end]), title="Value Functions",
		label="", legend=:bottomright)
plot!(k,v_1, label="True Value", color=:black, line=:dash)
#png("Value_function_vs_k_all")

plot([1:T], Σ_g[2:end], grid=false,
		xlabel="time", ylabel="\$\\sigma\$",
		title="\$ \\textrm{Estimate of } \\sigma \\textrm{ over time}\$",
		label="Misspecfication", legend=:bottomright)
plot!([1:T],ones(T).*σ, label="True Value", color=:black, line=:dash)
png("sigma")


@testset "Meataxe" begin

  @testset "cleanvect" begin

    F, a = Nemo.FiniteField(3, 1, "a")
    M=MatrixSpace(F,2,3)([1,1,0,0,1,0])
    v=MatrixSpace(F,1,3)([2,2,0])
    @test iszero(Hecke.cleanvect(M,v))
    v[1,3]=1
    @test !iszero(Hecke.cleanvect(M,v))

  end

  @testset "closure and spinning" begin

   F, a = Nemo.FiniteField(3, 1, "a")
   G=[MatrixSpace(F,4,4)([1,2,0,0,1,1,0,0,0,0,1,2,0,0,1,1])]
   M=FqGModule(G)
   v=MatrixSpace(F,1,4)([1,0,0,0])
   X=Hecke.closure(v,M.G)
   Y=Hecke.spinning(v,M.G)
   @test rows(X)==2
   @test rows(Y)==2
   rref!(Y)
   @test X==Y
   w=MatrixSpace(F,1,4)([0,0,0,1])
   v=vcat(v,w)
   @test rows(Hecke.closure(v,M.G))==4

  end

  @testset "meataxe" begin

    F, a = Nemo.FiniteField(3, 1, "a")
    G=[MatrixSpace(F,4,4)([1,2,0,0,1,1,0,0,0,0,1,2,0,0,1,1])]
    M=FqGModule(G)
    bool,B=meataxe(M)
    @test !bool
    @test rows(B)==2
    @test rows(Hecke.closure(B, M.G))==2

    N=Hecke.actsub(B,G)
    bool,B=meataxe(N)
    @test bool

    G=[MatrixSpace(F,3,3)([1,0,0,0,0,1,0,1,0]), MatrixSpace(F,3,3)([0,0,1,1,0,0,0,1,0])]
    M=FqGModule(G)
    bool,B=meataxe(M)
    @test !bool
    @test rows(Hecke.closure(B, M.G))==rows(B)

  end

  @testset "composition factors and series" begin

    F, a = Nemo.FiniteField(3, 1, "a")
    G=[MatrixSpace(F,4,4)([1,2,0,0,1,1,0,0,0,0,1,2,0,0,1,1])]
    M=FqGModule(G)
    lf=composition_factors(M)
    cs=composition_series(M)
    @test length(lf)==1
    @test length(cs)==2
    x=Hecke.actsub(cs[1],M.G)
    @test Hecke.isisomorphic(lf[1][1],x)
    x,_=Hecke.actquo(cs[1],M.G)
    @test Hecke.isisomorphic(lf[1][1],x)

    M=FqGModule([matrix(F,2,2,[0,1,2,0])])
    N=FqGModule([matrix(F,2,2,[0,2,1,0])])
    M.isirreducible=true
    @test Hecke.isisomorphic(M,N)

    M1=matrix(F,2,2,[1,0,1,1])
    M2=matrix(F,2,2,[1,1,0,1])
    M=FqGModule([M1,M2])
    M.isirreducible=true

    N1=matrix(F,2,2,[2,2,1,0])
    N2=matrix(F,2,2,[1,1,0,1])
    N=FqGModule([N1,N2])

    @test Hecke.isisomorphic(M,N)

  end

  @testset "Submodules" begin

    F, a = Nemo.FiniteField(3, 1, "a")
    A=MatrixSpace(F,3,3)(1)
    M=FqGModule([A])
    ls=minimal_submodules(M)
    @test length(ls)==13

    F, a = Nemo.FiniteField(2, 1, "a")
    A=MatrixSpace(F,6,6)(1)
    A[5,6]=1
    M=FqGModule([A])
    ls=minimal_submodules(M)
    @test length(ls)==31
    ls=submodules(M,4)
    @test length(ls)==171

  end

end


include("C:/Users/stravaglino3/Downloads/julia_code/julia_code/model_fit_global.jl")
include("C:/Users/stravaglino3/Downloads/julia_code/models.jl")
include("C:/Users/stravaglino3/Downloads/julia_code/helper_functions.jl")

f = (du,u,p,t) -> bimolecular_diss!(du,u,p,t,dens,cons) # enclose constants

# u0 = [0.0;0.0]
u₀ = [1.0]          # Initial condition
tspan = (0.0,10.0)  # Simulation interval
p = 0.5      # equation parameter. p = [k₁, k₋₁]
dens = []
cons = []

# Setup the ODE problem, then solve
prob = ODEProblem(f,u₀,tspan,p)
sol = solve(prob,Tsit5())
# sol = solve(prob,Tsit5())
plot(sol)
dataset1 = generate_data(sol)
# plot(sol.t,log.(Array(sol))')
#
# fs = collect(0:20)
# ks = bell_diss.(0.5,10,fs)
# plot(fs,ks)
model_opt2 = (1.0e-6,10.0)
pinit = 0.0
t = collect(range(0,stop=10,length=101))
rates,hes,l = optimization(bimolecular_diss!,[vec(dataset1)],
                            [[]],[[]],model_opt2,[t])

Juno.@enter optimization(bimolecular_diss!,[vec(dataset1)],
            [[]],[[]],model_opt2,[t])


## Convplutional layers

expand(N, i::Tuple) = i
expand(N, i::Integer) = ntuple(_ -> i, N)

"""
Convolutional layer that serves as the base for Conv2D and Conv3D
"""
mutable struct Conv <: LazyLayer
    channels::Int
    kernel_size::Union{Int,Tuple}
    activation
    padding::Union{Int,Tuple}
    strides::Union{Int,Tuple}
    dilation
    built::Bool
    use_bias::Bool
	init::NamedTuple
    params::NamedTuple
end

function Conv(
    channels::Int,
    kernel_size::Union{Int,Tuple},
    activation = identity;
    padding = 0,
    strides = 1,
    dilation = 1,
    use_bias = true,
	initw = Knet.xavier,
	initb = zeros
)
    @assert channels > 0 "Conv layer should have more then 0 channels"
    Conv(
        channels,
        kernel_size,
        activation,
        padding,
        strides,
        dilation,
        false,
        use_bias,
		(w=initw, b=initb),
        (w=nothing, b=nothing)
    )
end

function build(layer::Conv, shape::Shape)
    input_channels = shape[end]
    rank = length(shape) - 1
    kernel_size = expand(rank, layer.kernel_size)
    w = getparam(kernel_size..., input_channels, layer.channels, init=layer.init.w)
    b = nothing
    if layer.use_bias
        b = getparam(repeat([1],rank)..., layer.channels, 1, init = layer.init.b)
    end
    layer.params = (w=w,b=b)
end

function call(c::Conv, x::Tensor)
    w,b = c.params

    x = Knet.conv4(
        w,
        x,
        padding = c.padding,
        stride = c.strides,
        dilation = c.dilation,
    )
    if c.use_bias
        c.activation.(x .+ b)
    else
        c.activation.(x)
    end
end

"""
Utility function to determine the output size of a convolutional
layer given a certain input size and configuration of a convolutional layer.
This function works even if the weights of the layer are not initialized.

For example:
	c = Conv2D(16, (3,3), strides=3, padding=1)
	output_size(c, (224,224)) # output is (75, 75)

"""
function output_size(c::Conv, input_size)
    result = []
    l = length(input_size)
    padding = expand(l, c.padding)
    strides = expand(l, c.strides)
    dilation = expand(l, c.dilation)
    kernel_size = expand(l, c.kernel_size)

    for (i, elem) in enumerate(input_size)
        dimension = floor(
            Int,
            (elem + 2 * padding[i] - dilation[i] * (kernel_size[i] - 1) - 1) /
            strides[i],
        ) + 1
        push!(result, dimension)
    end
    tuple(result...)
end

"""
2D convolution layer (e.g. spatial convolution over images).

This layer creates a convolution kernel that is convolved with the layer
input to produce a tensor of outputs. If use_bias is true, a bias vector
is created and added to the outputs. Finally, if activation is not nothing,
it is applied to the outputs as well.
"""
Conv2D = Conv

"""
3D convolution layer (e.g. spatial convolution over volumes).

This layer creates a convolution kernel that is convolved with the layer input
to produce a tensor of outputs. If use_bias is true, a bias vector is created
and added to the outputs. Finally, if activation is not nothing, it is
applied to the outputs as well.
"""
Conv3D = Conv

"""
ConvTranspose layer
"""
mutable struct ConvTranspose <: LazyLayer
    channels::Int
    kernel_size::Union{Int,Tuple}
    activation
    padding::Union{Int,Tuple}
    strides::Union{Int,Tuple}
    dilation
    built::Bool
    use_bias::Bool
	init::NamedTuple
    params::NamedTuple
end

function ConvTranspose(
    channels::Int,
    kernel_size::Union{Int,Tuple},
    activation = identity;
    padding = 0,
    strides = 1,
    dilation = 1,
    use_bias = true,
	initw = Knet.xavier,
	initb = zeros
)
    @assert channels > 0
    ConvTranspose(
        channels,
        kernel_size,
        activation,
        padding,
        strides,
        dilation,
        false,
        use_bias,
		(w=initw, b=initb),
        (w=nothing, b=nothing)
    )
end

function build(layer::ConvTranspose, shape::Tuple)
    input_channels = shape[end]
    rank = length(shape) - 1
    kernel_size = expand(rank, layer.kernel_size)
    w = getparam(kernel_size..., layer.channels, input_channels,init=layer.init.w)
    b = nothing
    if layer.use_bias
        b = getparam(repeat([1],rank)..., layer.channels, 1, init=layer.init.b)
    end
    layer.params = (w=w,b=b)
end

function call(c::ConvTranspose, x::Tensor)
    w,b = c.params

    x = Knet.deconv4(
        w,
        x,
        padding = c.padding,
        stride = c.strides,
        dilation = c.dilation,
    )
    if c.use_bias
        c.activation.(x .+ b)
    else
        c.activation.(x)
    end
end

"""
Transposed 2D convolution layer (sometimes called Deconvolution).

The need for transposed convolutions generally arises from the desire to use a
transformation going in the opposite direction of a normal convolution,
i.e., from something that has the shape of the output of some convolution
to something that has the shape of its input while maintaining a connectivity
pattern that is compatible with said convolution.
"""
Conv2DTranspose = ConvTranspose

"""
Transposed 3D convolution layer (sometimes called Deconvolution).

The need for transposed convolutions generally arises from the desire to use a
transformation going in the opposite direction of a normal convolution,
i.e., from something that has the shape of the output of some convolution
to something that has the shape of its input while maintaining a connectivity
pattern that is compatible with said convolution.
"""
Conv3DTranspose = ConvTranspose

## AvgPool layer
"""
Pooling layers
"""
abstract type PoolingLayer <: Layer end

struct AvgPool <: PoolingLayer
    pool_size::Union{Int,Tuple}
    padding::Union{Int,Tuple}
    strides::Union{Int,Tuple}
end

function AvgPool(pool_size=2; padding = 0, strides = pool_size)
    AvgPool(pool_size, padding, strides)
end

function call(p::AvgPool, X)
    Knet.pool(
        X,
        window = p.pool_size,
        padding = p.padding,
        stride = p.strides,
        mode = 0,
    )
end

AvgPool1D = AvgPool

"""Average pooling operation for two dimensional (spatial) data."""
AvgPool2D = AvgPool

"""Average pooling operation for three dimensional (spatial) data."""
AvgPool3D = AvgPool

"""
Adaptive Average Pool has a fixed size output and enables creating a convolutional
network that can be used for multiple image formats.
"""
struct AdaptiveAvgPool <: PoolingLayer
    output_size::Tuple
end

function call(m::AdaptiveAvgPool, x::Tensor)
    s = []
    for idx = 1:length(m.output_size)
        push!(s, size(x, idx) - m.output_size[idx] + 1)
    end
    s = tuple(s...)
    return Knet.pool(x, window = s, padding = 0, stride = 1, mode = 0)
end

## MaxPool layer

struct MaxPool <: PoolingLayer
    pool_size::Union{Int,Tuple}
    padding::Union{Int,Tuple}
    strides::Union{Int,Tuple}
    nanOpt::Int
end

function MaxPool(pool_size=2; padding = 0, strides = pool_size, nanOpt = 0)
    MaxPool(pool_size, padding, strides, nanOpt)
end

function call(p::MaxPool, X::Tensor)
    Knet.pool(
        X,
        window = p.pool_size,
        padding = p.padding,
        maxpoolingNanOpt = p.nanOpt,
        stride = p.strides,
        mode = 1,
    )
end

MaxPool1D = MaxPool

"""Max pooling operation for two dimensional (spatial) data."""
MaxPool2D = MaxPool

"""Max pooling operation for three dimensional (spatial) data."""
MaxPool3D = MaxPool

"""
Adaptive MaxPool has a fixed size output and enables creating a convolutional
network that can be used for different image sizes.
"""
struct AdaptiveMaxPool <: PoolingLayer
    output_size::Tuple
end

function call(m::AdaptiveMaxPool, x::Tensor)
    s = []
    for idx = 1:length(m.output_size)
        push!(s, size(x, idx) - m.output_size[idx] + 1)
    end
    s = tuple(s...)
    return Knet.pool(x, window = s, padding = 0, stride = 1, mode = 1)
end

@debug "Loaded Convolutional modules"


"""
    beamforming(csm,v)

Calculate frequency-domain beamforming using cross-spectral matrix `csm` and steering vector `v`.
`csm` must be a square (Hermitian) matrix optionally with a third (frequency) dimension.
First dimension of `v` and `csm` must be equal.
"""
function beamforming(csm::A,v::B) where {T1<:AbstractFloat,T2<:AbstractFloat, N,  A<:AbstractArray{Complex{T1},N},B<:AbstractArray{Complex{T2},N}}
    #@assert size(csm,1) == size(csm,2) "csm must be square with dimensions M x M (x Nf)!"
    #@assert size(v,1) == size(csm,1) "First dimension of v and csm must be equal!"
    b = Array{T2, 2}(undef, size(v,2), size(csm,3))
    for j in axes(csm,3)
        csmd = @view csm[:,:,j]
        for i in axes(v,2)
            vd = @view v[:,i,j]
            b[i,j] = real(vd'*csmd*vd)
        end
    end
    return b
end

"""
    beamforming(Environment)

Calculate frequency-domain beamforming using cross-spectral matrix `csm` and steering vector `v`
stored in Environment struct
"""
function beamforming(E::Environment)
    @unpack CSM_s,w,steeringvec,M,N,Nf,fn = E
    b = Array{Float64}(undef, N, Nf)
    @views @inbounds for j in 1:Nf
        st = E.steeringvec.arr[:,:,j]
        @. st = complex(w[:,j].*real(st), imag(st))
        bf_col!(b[:,j],st,selectdim(CSM_s.arr, 3, j))
    end
    return FreqArray(b,fn)
end

function bf_col!(b::A,st::B,csm::C) where {T1<:AbstractFloat,T2<:AbstractFloat, N, A<:AbstractArray{T1,1}, B<:AbstractArray{Complex{T1},N},C<:AbstractArray{Complex{T2},N}}
    @views @inbounds for i in 1:length(b)
        b[i] = real(mydot(st[:,i]',csm,st[:,i]))
    end
end

function mydot(xt, A, x)
    s = zero(promote_type(eltype(x), eltype(A)))
    @inbounds @simd for m in 1:size(A,1)
        for n in 1:size(A,2)
            s += xt[m] * A[m,n] * x[n]
        end
    end
    s
end


using QuasinormalModes
using Test

function TestIfValueOfRadialWorks()
    a = -0.402546147858408 - 0.06278653515072422*im;
    Ψ = qnmfunction(s=-2,l=2,n=2,m=-2,a=0.4)
    (abs(Ψ(2.0)-a) < 10^(-9))
end

@testset "QuasinormalModes.jl" begin
    # Write your tests here.
    @test TestIfValueOfRadialWorks()
end


## test_separable.jl : test separable kernels for GRF generation

@testset "separable GRFs           " begin

    ## 1d Exponential ##
    cov = SeparableCovarianceFunction(Exponential(0.1))
    pts = range(0,stop = 1,length = 1001)
    grf = GaussianRandomField(cov,KarhunenLoeve(1000),pts)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,SeparableCovarianceFunction)
    @test isa(grf.cov.cov,Vector)
    @test isa(grf.cov.cov[1],Exponential)
    @test ndims(grf.cov) == 1
    @test isa(grf,GaussianRandomField{KarhunenLoeve{1000}})
    @test length(grf.pts[1]) == length(pts)
    @test length(sample(grf)) == length(pts)

    ## 2d Exponential ##
    e1 = Exponential(0.1)
    e2 = Exponential(0.01)
    cov = SeparableCovarianceFunction(e1,e2)
    pts1 = 0:0.05:1
    pts2 = 0:0.05:1
    grf = GaussianRandomField(cov,KarhunenLoeve(500),pts1,pts2)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,SeparableCovarianceFunction)
    @test isa(grf.cov.cov,Vector)
    @test isa(grf.cov.cov[1],Exponential)
    @test isa(grf.cov.cov[2],Exponential)
    @test ndims(grf.cov) == 2
    @test isa(grf,GaussianRandomField{KarhunenLoeve{500}})
    @test length(grf.pts) == 2
    @test length(grf.pts[1]) == length(pts1)
    @test length(grf.pts[2]) == length(pts2)
    @test size(sample(grf),1) == length(pts1)
    @test size(sample(grf),2) == length(pts2)

    ## 2d Exponential ##
    e1 = Exponential(0.1)
    e2 = Exponential(0.01)
    cov = SeparableCovarianceFunction(e1,e2)
    pts1 = 0:0.05:1
    pts2 = 0:0.05:1
    grf = GaussianRandomField(cov,KarhunenLoeve(500),pts1,pts2)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,SeparableCovarianceFunction)
    @test isa(grf.cov.cov,Vector)
    @test isa(grf.cov.cov[1],Exponential)
    @test isa(grf.cov.cov[2],Exponential)
    @test ndims(grf.cov) == 2
    @test isa(grf,GaussianRandomField{KarhunenLoeve{500}})
    @test length(grf.pts) == 2
    @test length(grf.pts[1]) == length(pts1)
    @test length(grf.pts[2]) == length(pts2)
    @test size(sample(grf),1) == length(pts1)
    @test size(sample(grf),2) == length(pts2)

    # matern/exponential 2d
    m1 = Matern(0.2,1.0)
    e1 = Exponential(0.01)
    cov = SeparableCovarianceFunction([m1,e1])
    pts1 = 0:0.05:1
    pts2 = 0:0.05:1
    grf = GaussianRandomField(cov,KarhunenLoeve(250),pts1,pts2,quad=EOLE())
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,SeparableCovarianceFunction)
    @test isa(grf.cov.cov,Vector)
    @test isa(grf.cov.cov[1],Matern)
    @test isa(grf.cov.cov[2],Exponential)
    @test ndims(grf.cov) == 2
    @test isa(grf,GaussianRandomField{KarhunenLoeve{250}})
    @test length(grf.pts) == 2
    @test length(grf.pts[1]) == length(pts1)
    @test length(grf.pts[2]) == length(pts2)
    @test size(sample(grf),1) == length(pts1)
    @test size(sample(grf),2) == length(pts2)

    ## 3d separable Matern ##
    m1 = Matern(0.2,1.0)
    m2 = Matern(0.4,3.0)
    m3 = Matern(0.6,4.0)
    cov = SeparableCovarianceFunction(m1,m1,m2)
    pts1 = 0:0.1:1
    pts2 = 0:0.1:1
    pts3 = 0:0.1:1
    grf = GaussianRandomField(cov,KarhunenLoeve(100),pts1,pts2,pts3)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,SeparableCovarianceFunction)
    @test isa(grf.cov.cov,Vector)
    @test isa(grf.cov.cov[1],Matern)
    @test isa(grf.cov.cov[2],Matern)
    @test isa(grf.cov.cov[3],Matern)
    @test ndims(grf.cov) == 3
    @test isa(grf,GaussianRandomField{KarhunenLoeve{100}})
    @test length(grf.pts) == 3
    @test length(grf.pts[1]) == length(pts1)
    @test length(grf.pts[2]) == length(pts2)
    @test length(grf.pts[3]) == length(pts3)
    @test size(sample(grf),1) == length(pts1)
    @test size(sample(grf),2) == length(pts2)
    @test size(sample(grf),3) == length(pts3)

end


using Test
using LinearAlgebra
import TensorTrains:one_body_to_matrix, one_body_mpo, two_body_to_matrix, two_body_mpo
using Random

@testset "MPO constructors" begin
    L=6
    p=rand(1:L)
    q=rand(1:L)
    #one-body MPO constructors
    H = one_body_to_matrix(p,q,L)
    Hmpo = one_body_mpo(p,q,L)
    H2 = tto_to_tensor(Hmpo)
    @test isapprox(norm(H-reshape(H2,2^L,2^L)),0.0,atol=1e-10)
    #two-body MPO constructors
#    k,l,m,n=1,2,3,4
#    H = two_body_to_matrix(k,l,m,n,L) #fails but probably bug in two_body_to_matrix
#    Hmpo = two_body_mpo(k,l,m,n,L)
#    H2 = tto_to_tensor(Hmpo)
#    @test isapprox(norm(H-reshape(H2,2^L,2^L)),0.0,atol=1e-10)
end

@testset "PPP_C_NH_N" begin
    N = 6
    H_tto = PPP_C_NH_N(N,order=randperm(12))
    ψ_0 = tt_up_rks(half_filling(N),32;ϵ_wn=1e-2)
    ψ_0 = orthogonalize(ψ_0)
    E_mals, ψ_tt = mals_eigsolv(H_tto,ψ_0)
    @test isapprox(E_mals[end],-0.46752535023652453,atol=1e-6) #value in Bendazzoli, G. L., & Evangelisti, S. (1991). Full-CI calculations of alternant cyclic polyenes (CH) N, N= 2, 4, 6, ƒ 18, in the PPP approximation. Chemical physics letters, 185(1-2), 125-130.
end


"""
Parameters for NFW Profile, includes useful quantities defined by parameters
"""
type NFWParameters{G <: AbstractFloat}
    a::G
    d::G
    e::G
    Ec::G
    rlim::G
    b::G
    q::G
    Jb::G
    vmax::G
    rmax::G

    #secondary parameters
    rs::G
    Φs::G
    xlim::G
    Φlim::G

    #adjusted parameter values
    adjEc::G
    adjJb::G
end

function NFWParameters(a::G, d::G, e::G, Ec::G, rlim::G, b::G, q::G, Jb::G, vmax::G, rmax::G) where G <: AbstractFloat

    rs = rmax / 2.16
    Φs = (vmax / 0.465)^2
    xlim = rlim / rs
    adjEc = Φs * Ec
    adjJb = Jb * rs * vmax / 0.465

    p = NFWParameters(a, d, e, Ec, rlim, b, q, Jb, vmax, rmax, rs, Φs, xlim, 0.0, adjEc, adjJb)
    p.Φlim = gravitational_potential(xlim, p)

    return p
end


using Parameters
include("common.jl")

# Creatine kinase and creatine shuttle parameters
@with_kw struct CKParams
    KCK_IN= 1.4E-4
    KCK_MT = 1.33E-6
    KTR_CR = 2E-3
    KEQ_CR = 0.0095
    # Basal consumption of cytosolic ATP (mM/ms)
    V_ATPASE_CYTO = 1E-5
    # Pool of creatine and creatine phosphate (mM)
    CR_T = 25.0
end

# CK shuttle rates
function ck_system(atp_i, atp_ic, crp_i, crp_ic, V_ATPASE_CYTO, KTR_CR, KCK_IN,
                  KCK_MT, KEQ, CR_T, AXP_I_T)
    cr_ic = CR_T - crp_ic
    cr_i = CR_T - crp_i
    adp_ic = AXP_I_T - atp_ic
    adp_i = AXP_I_T - atp_i

    vck_cyto = KCK_IN * (atp_ic * cr_ic - adp_ic * crp_ic / KEQ)
    vck_mito = KCK_MT * (atp_i * cr_i - adp_i * crp_i / KEQ)
    vtr_crp = KTR_CR * (crp_i - crp_ic)

    d_atp_in_ck = -vck_mito
    d_atp_ic = -vck_cyto - V_ATPASE_CYTO
    d_crp_i = vck_mito - vtr_crp
    d_crp_ic = vtr_crp + vck_cyto
    return (d_atp_in_ck, d_atp_ic, d_crp_i, d_crp_ic)
end

function ck_system(atp_i, atp_ic, crp_i, crp_ic, pCK::CKParams, AXP_I_T)
    @unpack V_ATPASE_CYTO, KTR_CR, KCK_IN, KCK_MT, KEQ_CR, CR_T = pCK
    ck_system(atp_i, atp_ic, crp_i, crp_ic, V_ATPASE_CYTO, KTR_CR, KCK_IN, KCK_MT, KEQ_CR, CR_T, AXP_I_T)
end


"""
    _variable_dbbt!

Tightens the bounds of the `_current_node` using the current global upper bound
and the duality information obtained from the relaxation.
"""
function _variable_dbbt!(m::GlobalOptimizer{N,T,S}) where {N,T<:AbstractFloat,S}

    LBD = m._lower_objective_value
    UBD = m._global_upper_bound
    delta = UBD - LBD
    if LBD <= UBD
        @inbounds for i = 1:m._branch_variable_num
            ml = m._lower_lvd[i]
            if ml > zero(T)
                cut = _lower_bound(BranchVar,m,i) + delta/ml
                if cut < _upper_bound(BranchVar,m,i)
                    _set_upper_bound!(BranchVar,m,i,cut)
                end
            else
                mh = m._lower_uvd[i]
                if mh > zero(T)
                    cut = _upper_bound(BranchVar,m,i) - delta/mh
                    if cut > _lower_bound(BranchVar,m,i)
                        _set_lower_bound!(BranchVar,m,i,cut)
                    end
                end
            end
        end
    end
    return nothing
end

"""
$(SIGNATURES)

Default postprocess perfoms duality-based bound tightening only.
"""
function postprocess!(t::ExtensionType, m::GlobalOptimizer)
    (m.dbbt_depth > m._iteration_count) && _variable_dbbt!(m)
    return nothing
end


using BenchmarkTools, AutoARIMA

const SUITE = BenchmarkGroup()

z = cos.(1:10000) .* log.(1:10000) .* randn()

SUITE["fit"] = BenchmarkGroup()

SUITE["fit"]["ar"] = BenchmarkGroup()
for f in (yule_walker, levinson_durbin, least_squares)
    for p in (1, 10)
        SUITE["fit"]["ar"][string(f),p] = @benchmarkable $f($z, $p)
    end
end

SUITE["fit"]["ma"] = BenchmarkGroup()
for f in (innovations,)
    for q in (1, 10)
        SUITE["fit"]["ma"][string(f),q] = @benchmarkable $f($z, $q, n=20)
    end
end

SUITE["fit"]["arma"] = BenchmarkGroup()
for f in (hannan_rissanen,)
    for p in (1, 10), q in (1, 10)
        SUITE["fit"]["arma"][string(f),p,q] = @benchmarkable $f($z, $p, $q)
    end
end

SUITE["forecast"] = BenchmarkGroup()

SUITE["forecast"]["ar"] = BenchmarkGroup()
for p in (1, 10, 100)
    ar = fit(ARParams(p), z)
    SUITE["forecast"]["ar"][p] = @benchmarkable forecast($ar, $z)
end

SUITE["forecast"]["ma"] = BenchmarkGroup()
for q in (1, 10, 100)
    ma = fit(MAParams(q), z, n=200)
    SUITE["forecast"]["ma"][q] = @benchmarkable forecast($ma, $z)
end

SUITE["forecast"]["arma"] = BenchmarkGroup()
for p in (1, 10, 100), q in (1, 10, 100)
    arma = fit(ARMAParams(p, q), z)
    SUITE["forecast"]["arma"][p,q] = @benchmarkable forecast($arma, $z)
end


####
# запустить: ]dev TestFilter

using TestFilter
using DSP
using DelimitedFiles
using Plots

x = readdlm("examples/ecg.csv")[:, 1]

f = DSPFilter()
fcut = 30
fs = 257
init(f, Lowpass(2 * fcut / fs), Butterworth(2))

y = similar(x)
run_filter(f, y, x)

plot(x)
plot!(y)


using Revise
using Pkg
Pkg.activate(".")
using MCPhylo
using Random
using LinearAlgebra

Random.seed!(42)

s = rand(InverseWishart(10, diagm(ones(3))))
#s = rand(Exponential())
mu = [0.1,0.4,0.3,0.15,0.05]
μ = randn(3)
mu1 = mu + rand(5)
mu1 /= sum(mu1)
mu1 = rand(Dirichlet(5,1))
my_data = Dict{Symbol, Any}(
  :df => collect(transpose(rand(Multinomial(1, mu1), 500)))
  #collect(transpose(rand(MultivariateNormal(μ, s), 500)))
);

# model setup
model =  Model(
    df = Stochastic(2, (mu1) ->  MultivariateDistribution[Multinomial(1, mu1) for i in 1:500], false),
    #MultivariateDistribution[MultivariateNormal(mu1, cm) for i in 1:500], false),
    mu1 = Stochastic(1, () -> Dirichlet(5,1.0), true),
    cm = Stochastic(2, () -> InverseWishart(10, diagm(ones(3))), false)
     )
# intial model values
inits = [ Dict{Symbol, Union{Any, Real}}(
    :mu1 => rand(Dirichlet(5, 1.0)),
    :cm => rand(InverseWishart(3, diagm(ones(3)))),
    :df => my_data[:df],

    ) for i in 1:2
    ]

scheme = [#Slice([:mu1, :cm], 1.0),
          #RWM(:mu1, 1.0)
          SliceSimplex(:mu1, scale = 1)
          #NUTS([:mu1, :cm], tree_depth=10, target=0.7)
          #HMC([:mu1, :cm], 0.001, 10)
]

setsamplers!(model, scheme);

# do the mcmc simmulation. if trees=true the trees are stored and can later be
# flushed ot a file output.

sim = mcmc(model, my_data, inits, 25000, burnin=12500,thin=10, chains=2, trees=true)