@doc raw"""
    AbstractMetric

Abstract type for the pseudo-Riemannian metric tensor ``g``, a family of smoothly
varying inner products on the tangent space. See [`inner`](@ref).

# Functor

    (metric::Metric)(M::Manifold)

Generate the `MetricManifold` that wraps the manifold `M` with given `metric`.
This works for both a variable containing the metric as well as a subtype `T<:AbstractMetric`,
where a zero parameter constructor `T()` is availabe.
"""
abstract type AbstractMetric end

# piping syntax for decoration
(metric::AbstractMetric)(M::AbstractManifold) = MetricManifold(M, metric)
(::Type{T})(M::AbstractManifold) where {T<:AbstractMetric} = MetricManifold(M, T())

"""
    MetricManifold{𝔽,M<:AbstractManifold{𝔽},G<:AbstractMetric} <: AbstractDecoratorManifold{𝔽}

Equip a [`AbstractManifold`](@ref) explicitly with a [`AbstractMetric`](@ref) `G`.

For a Metric AbstractManifold, by default, assumes, that you implement the linear form
from [`local_metric`](@ref) in order to evaluate the exponential map.

If the corresponding [`AbstractMetric`](@ref) `G` yields closed form formulae for e.g.
the exponential map and this is implemented directly (without solving the ode),
you can of course still implement that directly.

# Constructor

    MetricManifold(M, G)

Generate the [`AbstractManifold`](@ref) `M` as a manifold with the [`AbstractMetric`](@ref) `G`.
"""
struct MetricManifold{𝔽,M<:AbstractManifold{𝔽},G<:AbstractMetric} <:
       AbstractDecoratorManifold{𝔽}
    manifold::M
    metric::G
end

@doc raw"""
    RiemannianMetric <: AbstractMetric

Abstract type for Riemannian metrics, a family of positive definite inner
products. The positive definite property means that for ``X  ∈ T_p \mathcal M``, the
inner product ``g(X, X) > 0`` whenever ``X`` is not the zero vector.
"""
abstract type RiemannianMetric <: AbstractMetric end

@doc raw"""
    christoffel_symbols_first(
        M::MetricManifold,
        p,
        B::AbstractBasis;
        backend::AbstractDiffBackend = diff_backend(),
    )

Compute the Christoffel symbols of the first kind in local coordinates of basis `B`.
The Christoffel symbols are (in Einstein summation convention)

```math
Γ_{ijk} = \frac{1}{2} \Bigl[g_{kj,i} + g_{ik,j} - g_{ij,k}\Bigr],
```

where ``g_{ij,k}=\frac{∂}{∂ p^k} g_{ij}`` is the coordinate
derivative of the local representation of the metric tensor. The dimensions of
the resulting multi-dimensional array are ordered ``(i,j,k)``.
"""
christoffel_symbols_first(::AbstractManifold, ::Any, B::AbstractBasis)
function christoffel_symbols_first(
    M::AbstractManifold,
    p,
    B::AbstractBasis;
    backend::AbstractDiffBackend=diff_backend(),
)
    ∂g = local_metric_jacobian(M, p, B; backend=backend)
    n = size(∂g, 1)
    Γ = allocate(∂g, Size(n, n, n))
    @einsum Γ[i, j, k] = 1 / 2 * (∂g[k, j, i] + ∂g[i, k, j] - ∂g[i, j, k])
    return Γ
end
@decorator_transparent_signature christoffel_symbols_first(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)
function decorator_transparent_dispatch(
    ::typeof(christoffel_symbols_first),
    ::MetricManifold,
    args...,
)
    return Val(:parent)
end

@doc raw"""
    christoffel_symbols_second(
        M::AbstractManifold,
        p,
        B::AbstractBasis;
        backend::AbstractDiffBackend = diff_backend(),
    )

Compute the Christoffel symbols of the second kind in local coordinates of basis `B`.
The Christoffel symbols are (in Einstein summation convention)

````math
Γ^{l}_{ij} = g^{kl} Γ_{ijk},
````

where ``Γ_{ijk}`` are the Christoffel symbols of the first kind, and
``g^{kl}`` is the inverse of the local representation of the metric tensor.
The dimensions of the resulting multi-dimensional array are ordered ``(l,i,j)``.
"""
christoffel_symbols_second(::AbstractManifold, ::AbstractBasis, ::Any)
function christoffel_symbols_second(
    M::AbstractManifold,
    p,
    B::AbstractBasis;
    backend::AbstractDiffBackend=diff_backend(),
)
    Ginv = inverse_local_metric(M, p, B)
    Γ₁ = christoffel_symbols_first(M, p, B; backend=backend)
    Γ₂ = allocate(Γ₁)
    @einsum Γ₂[l, i, j] = Ginv[k, l] * Γ₁[i, j, k]
    return Γ₂
end
@decorator_transparent_signature christoffel_symbols_second(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)
function decorator_transparent_dispatch(
    ::typeof(christoffel_symbols_second),
    ::MetricManifold,
    args...,
)
    return Val(:parent)
end

@doc raw"""
    christoffel_symbols_second_jacobian(
        M::AbstractManifold,
        p,
        B::AbstractBasis;
        backend::AbstractDiffBackend = diff_backend(),
    )

Get partial derivatives of the Christoffel symbols of the second kind
for manifold `M` at `p` with respect to the coordinates of `B`, i.e.

```math
\frac{∂}{∂ p^l} Γ^{k}_{ij} = Γ^{k}_{ij,l}.
```

The dimensions of the resulting multi-dimensional array are ordered ``(i,j,k,l)``.
"""
christoffel_symbols_second_jacobian(::AbstractManifold, ::Any, B::AbstractBasis)
function christoffel_symbols_second_jacobian(
    M::AbstractManifold,
    p,
    B::AbstractBasis;
    backend::AbstractDiffBackend=diff_backend(),
)
    n = size(p, 1)
    ∂Γ = reshape(
        _jacobian(q -> christoffel_symbols_second(M, q, B; backend=backend), p, backend),
        n,
        n,
        n,
        n,
    )
    return ∂Γ
end
@decorator_transparent_signature christoffel_symbols_second_jacobian(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)
function decorator_transparent_dispatch(
    ::typeof(christoffel_symbols_second_jacobian),
    ::MetricManifold,
    args...,
)
    return Val(:parent)
end

Base.copyto!(M::MetricManifold, q, p) = copyto!(M.manifold, q, p)
Base.copyto!(M::MetricManifold, Y, p, X) = copyto!(M.manifold, Y, p, X)

decorator_transparent_dispatch(::typeof(exp), M::MetricManifold, args...) = Val(:parent)
function decorator_transparent_dispatch(::typeof(exp!), M::MetricManifold, args...)
    return Val(:intransparent)
end
decorator_transparent_dispatch(::typeof(exp!), M::MetricManifold, q, p, X, t) = Val(:parent)
decorator_transparent_dispatch(::typeof(flat), M::MetricManifold, args...) = Val(:parent)
function decorator_transparent_dispatch(::typeof(flat!), M::MetricManifold, args...)
    return Val(:intransparent)
end
function decorator_transparent_dispatch(
    ::typeof(get_coordinates),
    M::MetricManifold,
    args...,
)
    return Val(:parent)
end
function decorator_transparent_dispatch(
    ::typeof(get_coordinates!),
    M::MetricManifold,
    args...,
)
    return Val(:intransparent)
end
function decorator_transparent_dispatch(::typeof(get_vector), M::MetricManifold, args...)
    return Val(:parent)
end
function decorator_transparent_dispatch(::typeof(get_vector!), M::MetricManifold, args...)
    return Val(:intransparent)
end
function decorator_transparent_dispatch(::typeof(get_basis), M::MetricManifold, args...)
    return Val(:intransparent)
end
function decorator_transparent_dispatch(::typeof(inner), M::MetricManifold, args...)
    return Val(:intransparent)
end
function decorator_transparent_dispatch(
    ::typeof(inverse_retract),
    M::MetricManifold,
    args...,
)
    return Val(:parent)
end
function decorator_transparent_dispatch(
    ::typeof(inverse_retract!),
    M::MetricManifold,
    args...,
)
    return Val(:intransparent)
end
function decorator_transparent_dispatch(
    ::typeof(inverse_retract!),
    M::MetricManifold,
    X,
    p,
    q,
    m::LogarithmicInverseRetraction,
)
    return Val(:parent)
end
decorator_transparent_dispatch(::typeof(log), M::MetricManifold, args...) = Val(:parent)
function decorator_transparent_dispatch(::typeof(log!), M::MetricManifold, args...)
    return Val(:intransparent)
end
decorator_transparent_dispatch(::typeof(mean), M::MetricManifold, args...) = Val(:parent)
function decorator_transparent_dispatch(::typeof(mean!), M::MetricManifold, args...)
    return Val(:intransparent)
end
decorator_transparent_dispatch(::typeof(median), M::MetricManifold, args...) = Val(:parent)
function decorator_transparent_dispatch(::typeof(median!), M::MetricManifold, args...)
    return Val(:intransparent)
end
function decorator_transparent_dispatch(::typeof(norm), M::MetricManifold, args...)
    return Val(:intransparent)
end
decorator_transparent_dispatch(::typeof(project), M::MetricManifold, args...) = Val(:parent)
function decorator_transparent_dispatch(::typeof(project!), M::MetricManifold, args...)
    return Val(:intransparent)
end
decorator_transparent_dispatch(::typeof(sharp), M::MetricManifold, args...) = Val(:parent)
function decorator_transparent_dispatch(::typeof(sharp!), M::MetricManifold, args...)
    return Val(:intransparent)
end
decorator_transparent_dispatch(::typeof(retract), M::MetricManifold, args...) = Val(:parent)
function decorator_transparent_dispatch(::typeof(retract!), M::MetricManifold, args...)
    return Val(:intransparent)
end
function decorator_transparent_dispatch(
    ::typeof(retract!),
    M::MetricManifold,
    q,
    p,
    X,
    m::ExponentialRetraction,
)
    return Val(:parent)
end
function decorator_transparent_dispatch(
    ::typeof(vector_transport_along),
    M::MetricManifold,
    args...,
)
    return Val(:parent)
end
function decorator_transparent_dispatch(
    ::typeof(vector_transport_along!),
    M::MetricManifold,
    args...,
)
    return Val(:intransparent)
end
function decorator_transparent_dispatch(
    ::typeof(vector_transport_direction),
    M::MetricManifold,
    args...,
)
    return Val(:parent)
end
function decorator_transparent_dispatch(
    ::typeof(vector_transport_direction!),
    M::MetricManifold,
    args...,
)
    return Val(:parent)
end
function decorator_transparent_dispatch(
    ::typeof(vector_transport_to),
    M::MetricManifold,
    args...,
)
    return Val(:parent)
end
function decorator_transparent_dispatch(
    ::typeof(vector_transport_to!),
    M::MetricManifold,
    args...,
)
    return Val(:intransparent)
end

@doc raw"""
    det_local_metric(M::AbstractManifold, p, B::AbstractBasis)

Return the determinant of local matrix representation of the metric tensor ``g``, i.e. of the
matrix ``G(p)`` representing the metric in the tangent space at ``p`` with as a matrix.

See also [`local_metric`](@ref)
"""
det_local_metric(::AbstractManifold, p, ::AbstractBasis)
function det_local_metric(M::AbstractManifold, p, B::AbstractBasis)
    return det(local_metric(M, p, B))
end
@decorator_transparent_signature det_local_metric(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis,
)
function decorator_transparent_dispatch(
    ::typeof(det_local_metric),
    ::MetricManifold,
    args...,
)
    return Val(:parent)
end
"""
    einstein_tensor(M::AbstractManifold, p, B::AbstractBasis; backend::AbstractDiffBackend = diff_backend())

Compute the Einstein tensor of the manifold `M` at the point `p`, see [https://en.wikipedia.org/wiki/Einstein_tensor](https://en.wikipedia.org/wiki/Einstein_tensor)
"""
einstein_tensor(::AbstractManifold, ::Any, ::AbstractBasis)
function einstein_tensor(
    M::AbstractManifold,
    p,
    B::AbstractBasis;
    backend::AbstractDiffBackend=diff_backend(),
)
    Ric = ricci_tensor(M, p, B; backend=backend)
    g = local_metric(M, p, B)
    Ginv = inverse_local_metric(M, p, B)
    S = sum(Ginv .* Ric)
    G = Ric - g .* S / 2
    return G
end
@decorator_transparent_signature einstein_tensor(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)
function decorator_transparent_dispatch(
    ::typeof(einstein_tensor),
    ::MetricManifold,
    args...,
)
    return Val(:parent)
end

@doc raw"""
    exp(N::MetricManifold{M,G}, p, X)

Copute the exponential map on the [`AbstractManifold`](@ref) `M` equipped with the [`AbstractMetric`](@ref) `G`.

If the metric was declared the default metric using [`is_default_metric`](@ref), this method
falls back to `exp(M, p, X)`.

Otherwise it numerically integrates the underlying ODE, see [`solve_exp_ode`](@ref).
Currently, the numerical integration is only accurate when using a single
coordinate chart that covers the entire manifold. This excludes coordinates
in an embedded space.
"""
exp(::MetricManifold, ::Any...)

@decorator_transparent_fallback function exp!(M::MetricManifold, q, p, X)
    tspan = (0.0, 1.0)
    A = get_default_atlas(M)
    i = get_chart_index(M, A, p)
    B = induced_basis(M, A, i, TangentSpace)
    sol = solve_exp_ode(M, p, X, tspan, B; dense=false, saveat=[1.0])
    n = length(p)
    return copyto!(q, sol.u[1][(n + 1):end])
end

@doc raw"""
    flat(N::MetricManifold{M,G}, p, X::FVector{TangentSpaceType})

Compute the musical isomorphism to transform the tangent vector `X` from the
[`AbstractManifold`](@ref) `M` equipped with [`AbstractMetric`](@ref) `G` to a cotangent by
computing

````math
X^♭= G_p X,
````
where ``G_p`` is the local matrix representation of `G`, see [`local_metric`](@ref)
"""
flat(::MetricManifold, ::Any...)

@decorator_transparent_fallback function flat!(
    M::MetricManifold,
    ξ::CoTFVector,
    p,
    X::TFVector,
)
    g = local_metric(M, p, ξ.basis)
    copyto!(ξ.data, g * X.data)
    return ξ
end

"""
    gaussian_curvature(M::AbstractManifold, p, B::AbstractBasis; backend::AbstractDiffBackend = diff_backend())

Compute the Gaussian curvature of the manifold `M` at the point `p` using basis `B`.
"""
gaussian_curvature(::AbstractManifold, ::Any, ::AbstractBasis)
function gaussian_curvature(M::AbstractManifold, p, B::AbstractBasis; kwargs...)
    return ricci_curvature(M, p, B; kwargs...) / 2
end
@decorator_transparent_signature gaussian_curvature(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)
function decorator_transparent_dispatch(
    ::typeof(gaussian_curvature),
    ::MetricManifold,
    args...,
)
    return Val(:parent)
end

function injectivity_radius(M::MetricManifold, p)
    return injectivity_radius(base_manifold(M), p)
end
function injectivity_radius(M::MetricManifold, m::AbstractRetractionMethod)
    return injectivity_radius(base_manifold(M), m)
end
function injectivity_radius(M::MetricManifold, m::ExponentialRetraction)
    return injectivity_radius(base_manifold(M), m)
end
function injectivity_radius(M::MetricManifold, p, m::AbstractRetractionMethod)
    return injectivity_radius(base_manifold(M), p, m)
end
function injectivity_radius(M::MetricManifold, p, m::ExponentialRetraction)
    return injectivity_radius(base_manifold(M), p, m)
end

@doc raw"""
    inverse_local_metric(M::AbstractcManifold, p, B::AbstractBasis)

Return the local matrix representation of the inverse metric (cometric) tensor, usually
written ``g^{ij}``.

See also [`local_metric`](@ref)
"""
inverse_local_metric(::AbstractManifold, ::Any, ::AbstractBasis)
function inverse_local_metric(M::AbstractManifold, p, B::AbstractBasis)
    return inv(local_metric(M, p, B))
end
@decorator_transparent_signature inverse_local_metric(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis,
)
function decorator_transparent_dispatch(
    ::typeof(inverse_local_metric),
    ::MetricManifold,
    args...,
)
    return Val(:parent)
end

default_decorator_dispatch(M::MetricManifold) = default_metric_dispatch(M)

"""
    is_default_metric(M, G)

Indicate whether the [`AbstractMetric`](@ref) `G` is the default metric for
the [`AbstractManifold`](@ref) `M`. This means that any occurence of
[`MetricManifold`](@ref)(M,G) where `typeof(is_default_metric(M,G)) = true`
falls back to just be called with `M` such that the [`AbstractManifold`](@ref) `M`
implicitly has this metric, for example if this was the first one implemented
or is the one most commonly assumed to be used.
"""
function is_default_metric(M::AbstractManifold, G::AbstractMetric)
    return _extract_val(default_metric_dispatch(M, G))
end

default_metric_dispatch(::AbstractManifold, ::AbstractMetric) = Val(false)
function default_metric_dispatch(M::MetricManifold)
    return default_metric_dispatch(base_manifold(M), metric(M))
end

"""
    is_default_metric(MM::MetricManifold)

Indicate whether the [`AbstractMetric`](@ref) `MM.G` is the default metric for
the [`AbstractManifold`](@ref) `MM.manifold,` within the [`MetricManifold`](@ref) `MM`.
This means that any occurence of
[`MetricManifold`](@ref)`(MM.manifold, MM.G)` where `is_default_metric(MM.manifold, MM.G)) = true`
falls back to just be called with `MM.manifold,` such that the [`AbstractManifold`](@ref) `MM.manifold`
implicitly has the metric `MM.G`, for example if this was the first one
implemented or is the one most commonly assumed to be used.
"""
function is_default_metric(M::MetricManifold)
    return _extract_val(default_metric_dispatch(M))
end

function Base.convert(::Type{MetricManifold{𝔽,MT,GT}}, M::MT) where {𝔽,MT,GT}
    return _convert_with_default(M, GT, default_metric_dispatch(M, GT()))
end

function _convert_with_default(
    M::MT,
    T::Type{<:AbstractMetric},
    ::Val{true},
) where {MT<:AbstractManifold}
    return MetricManifold(M, T())
end
function _convert_with_default(
    M::MT,
    T::Type{<:AbstractMetric},
    ::Val{false},
) where {MT<:AbstractManifold}
    return error(
        "Can not convert $(M) to a MetricManifold{$(MT),$(T)}, since $(T) is not the default metric.",
    )
end

@doc raw"""
    inner(N::MetricManifold{M,G}, p, X, Y)

Compute the inner product of `X` and `Y` from the tangent space at `p` on the
[`AbstractManifold`](@ref) `M` using the [`AbstractMetric`](@ref) `G`. If `G` is the default
metric (see [`is_default_metric`](@ref)) this is done using `inner(M, p, X, Y)`,
otherwise the [`local_metric`](@ref)`(M, p)` is employed as

````math
g_p(X, Y) = ⟨X, G_p Y⟩,
````
where ``G_p`` is the loal matrix representation of the [`AbstractMetric`](@ref) `G`.
"""
inner(::MetricManifold, ::Any, ::Any, ::Any)

@decorator_transparent_fallback :intransparent function inner(
    M::MetricManifold,
    p,
    X::TFVector,
    Y::TFVector,
)
    X.basis === Y.basis ||
        error("calculating inner product of vectors from different bases is not supported")
    return dot(X.data, local_metric(M, p, X.basis) * Y.data)
end

@doc raw"""
    local_metric(M::AbstractManifold, p, B::AbstractBasis)

Return the local matrix representation at the point `p` of the metric tensor ``g`` with
respect to the [`AbstractBasis`](@ref) `B` on the [`AbstractManifold`](@ref) `M`, usually written ``g_{ij}``.
The matrix has the property that ``g(X, Y)=X^\mathrm{T} [g_{ij}] Y = g_{ij} X^i Y^j``,
where the latter expression uses Einstein summation convention.
The metric tensor is such that the formula works for the given [`AbstractBasis`](@ref) `B`.
"""
local_metric(::AbstractManifold, ::Any, ::AbstractBasis)
@decorator_transparent_signature local_metric(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)
function decorator_transparent_dispatch(::typeof(local_metric), ::MetricManifold, args...)
    return Val(:parent)
end

@doc raw"""
    local_metric_jacobian(
        M::AbstractManifold,
        p,
        B::AbstractBasis;
        backend::AbstractDiffBackend = diff_backend(),
    )

Get partial derivatives of the local metric of `M` at `p` in basis `B` with respect to the
coordinates of `p`, ``\frac{∂}{∂ p^k} g_{ij} = g_{ij,k}``. The
dimensions of the resulting multi-dimensional array are ordered ``(i,j,k)``.
"""
local_metric_jacobian(::AbstractManifold, ::Any, B::AbstractBasis)
function local_metric_jacobian(
    M::AbstractManifold,
    p,
    B::AbstractBasis;
    backend::AbstractDiffBackend=diff_backend(),
)
    n = size(p, 1)
    ∂g = reshape(_jacobian(q -> local_metric(M, q, B), p, backend), n, n, n)
    return ∂g
end
@decorator_transparent_signature local_metric_jacobian(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)
function decorator_transparent_dispatch(
    ::typeof(local_metric_jacobian),
    ::MetricManifold,
    args...,
)
    return Val(:parent)
end

@doc raw"""
    log(N::MetricManifold{M,G}, p, q)

Copute the logarithmic map on the [`AbstractManifold`](@ref) `M` equipped with the [`AbstractMetric`](@ref) `G`.

If the metric was declared the default metric using [`is_default_metric`](@ref), this method
falls back to `log(M,p,q)`. Otherwise, you have to provide an implementation for the non-default
[`AbstractMetric`](@ref) `G` metric within its [`MetricManifold`](@ref)`{M,G}`.
"""
log(::MetricManifold, ::Any...)

@doc raw"""
    log_local_metric_density(M::AbstractManifold, p, B::AbstractBasis)

Return the natural logarithm of the metric density ``ρ`` of `M` at `p`, which
is given by ``ρ = \log \sqrt{|\det [g_{ij}]|}`` for the metric tensor expressed in basis `B`.
"""
log_local_metric_density(::AbstractManifold, ::Any, ::AbstractBasis)
function log_local_metric_density(M::AbstractManifold, p, B::AbstractBasis)
    return log(abs(det_local_metric(M, p, B))) / 2
end
@decorator_transparent_signature log_local_metric_density(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis,
)
function decorator_transparent_dispatch(
    ::typeof(log_local_metric_density),
    ::MetricManifold,
    args...,
)
    return Val(:parent)
end

@doc raw"""
    metric(M::MetricManifold)

Get the metric ``g`` of the manifold `M`.
"""
metric(::MetricManifold)

function metric(M::MetricManifold)
    return M.metric
end

"""
    ricci_curvature(M::AbstractManifold, p, B::AbstractBasis; backend::AbstractDiffBackend = diff_backend())

Compute the Ricci scalar curvature of the manifold `M` at the point `p` using basis `B`.
"""
ricci_curvature(::AbstractManifold, ::Any, ::AbstractBasis)
function ricci_curvature(
    M::AbstractManifold,
    p,
    B::AbstractBasis;
    backend::AbstractDiffBackend=diff_backend(),
)
    Ginv = inverse_local_metric(M, p, B)
    Ric = ricci_tensor(M, p, B; backend=backend)
    S = sum(Ginv .* Ric)
    return S
end
@decorator_transparent_signature ricci_curvature(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)
function decorator_transparent_dispatch(
    ::typeof(ricci_curvature),
    ::MetricManifold,
    args...,
)
    return Val(:parent)
end
"""
    ricci_tensor(M::AbstractManifold, p, B::AbstractBasis; backend::AbstractDiffBackend = diff_backend())

Compute the Ricci tensor, also known as the Ricci curvature tensor,
of the manifold `M` at the point `p` using basis `B`,
see [https://en.wikipedia.org/wiki/Ricci_curvature#Introduction_and_local_definition](https://en.wikipedia.org/wiki/Ricci_curvature#Introduction_and_local_definition).
"""
ricci_tensor(::AbstractManifold, ::Any, ::AbstractBasis)
function ricci_tensor(M::AbstractManifold, p, B::AbstractBasis; kwargs...)
    R = riemann_tensor(M, p, B; kwargs...)
    n = size(R, 1)
    Ric = allocate(R, Size(n, n))
    @einsum Ric[i, j] = R[l, i, l, j]
    return Ric
end
@decorator_transparent_signature ricci_tensor(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)
function decorator_transparent_dispatch(::typeof(ricci_tensor), ::MetricManifold, args...)
    return Val(:parent)
end
@doc raw"""
    riemann_tensor(M::AbstractManifold, p, B::AbstractBasis; backend::AbstractDiffBackend = diff_backend())

Compute the Riemann tensor ``R^l_{ijk}``, also known as the Riemann curvature
tensor, at the point `p`. The dimensions of the resulting multi-dimensional
array are ordered ``(l,i,j,k)``.
"""
riemann_tensor(::AbstractManifold, ::Any, ::AbstractBasis)
function riemann_tensor(
    M::AbstractManifold,
    p,
    B::AbstractBasis;
    backend::AbstractDiffBackend=diff_backend(),
)
    n = size(p, 1)
    Γ = christoffel_symbols_second(M, p, B; backend=backend)
    ∂Γ = christoffel_symbols_second_jacobian(M, p, B; backend=backend) ./ n
    R = allocate(∂Γ, Size(n, n, n, n))
    @einsum R[l, i, j, k] =
        ∂Γ[l, i, k, j] - ∂Γ[l, i, j, k] + Γ[s, i, k] * Γ[l, s, j] - Γ[s, i, j] * Γ[l, s, k]
    return R
end
@decorator_transparent_signature riemann_tensor(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)
function decorator_transparent_dispatch(::typeof(riemann_tensor), ::MetricManifold, args...)
    return Val(:parent)
end

@doc raw"""
    sharp(N::MetricManifold{M,G}, p, ξ::FVector{CotangentSpaceType})

Compute the musical isomorphism to transform the cotangent vector `ξ` from the
[`AbstractManifold`](@ref) `M` equipped with [`AbstractMetric`](@ref) `G` to a tangent by
computing

````math
ξ^♯ = G_p^{-1} ξ,
````
where ``G_p`` is the local matrix representation of `G`, i.e. one employs
[`inverse_local_metric`](@ref) here to obtain ``G_p^{-1}``.
"""
sharp(::MetricManifold, ::Any, ::CoTFVector)

function sharp!(M::N, X::TFVector, p, ξ::CoTFVector) where {N<:MetricManifold}
    Ginv = inverse_local_metric(M, p, X.basis)
    copyto!(X.data, Ginv * ξ.data)
    return X
end

function Base.show(io::IO, M::MetricManifold)
    return print(io, "MetricManifold($(M.manifold), $(M.metric))")
end

@doc raw"""
    solve_exp_ode(
        M::MetricManifold,
        p,
        X,
        tspan,
        B::AbstractBasis;
        backend::AbstractDiffBackend = diff_backend(),
        solver = AutoVern9(Rodas5()),
        kwargs...,
    )

Approximate the exponential map on the manifold over the provided timespan
assuming the Levi-Civita connection by solving the ordinary differential
equation

```math
\frac{d^2}{dt^2} p^k + Γ^k_{ij} \frac{d}{dt} p_i \frac{d}{dt} p_j = 0,
```

where ``Γ^k_{ij}`` are the Christoffel symbols of the second kind, and
the Einstein summation convention is assumed. The arguments `tspan` and
`solver` follow the `OrdinaryDiffEq` conventions. `kwargs...` specify keyword
arguments that will be passed to `OrdinaryDiffEq.solve`.

Currently, the numerical integration is only accurate when using a single
coordinate chart that covers the entire manifold. This excludes coordinates
in an embedded space.

!!! note
    This function only works when
    [OrdinaryDiffEq.jl](https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl) is loaded with
    ```julia
    using OrdinaryDiffEq
    ```
"""
function solve_exp_ode(M, p, X, tspan, B::AbstractBasis; kwargs...)
    return error(
        "solve_exp_ode not implemented on $(typeof(M)) for point $(typeof(p)), vector $(typeof(X)), and timespan $(typeof(tspan)). For a suitable default, enter `using OrdinaryDiffEq` on Julia 1.1 or greater.",
    )
end


#Tensors in Kruskal (CP) format + functions

export ktensor, randktensor, arrange, arrange!, cp_als, display, extract, fixsigns, fixsigns!, full, innerprod, isequal, krank, kten2TT, minus, mtimes, mttkrp, ncomponents, ndims
export normalize, normalize!, nvecs, permutedims, plus, redistribute, redistribute!, size, tenmat, tocell, ttensor, ttm, ttv, uminus, norm

"""
    ktensor(fmat)
    ktensor(lambda,fmat)

Tensor in Kruskal format defined by its vector of weights lambda and factor matrices. Default lambda: vector of ones.
For ktensor X, X.isorth=true if factor matrices are othonormal.
"""
mutable struct ktensor
	lambda::AbstractVector{<:Number}
	fmat::MatrixCell
	isorth::Bool
	function ktensor(lambda::AbstractVector{<:Number},fmat::MatrixCell,isorth::Bool)
    @assert(length(unique([size(fmat[n],2) for n=1:length(fmat)]))==1,"Factor matrices must have same number of columns.")
    @assert(length(lambda)==size(fmat[1],2),"Dimension mismatch.")
		for A in fmat
			if norm(A'*A-eye(size(A,2)))>(size(A,1)^2)*eps()
				isorth=false
			end
		end
		new(lambda,fmat,isorth)
	end
end
#ktensor(lambda::AbstractArray{T1,1},fmat::MatrixCell,isorth::Bool) where {T1<:Number,T2<:Number}=ktensor{T1,T2}(lambda,fmat,isorth)
ktensor(lambda::AbstractVector{<:Number},fmat::MatrixCell)=ktensor(lambda,fmat,true)
ktensor(lambda::AbstractVector{<:Number},mat::AbstractMatrix{<:Number})=ktensor(lambda,collect(mat),true)
ktensor(fmat::MatrixCell)=ktensor(ones(size(fmat[1],2)),fmat)
# function ktensor(lambda::AbstractVector{<:Number},mat...)
#   fmat=MatrixCell(undef,0)
#   for M in mat
# 			push!(fmat,M)
# 	end
# 	ktensor(lambda,fmat,true)
# end
#If array of matrices isn't defined as MatrixCell, but as M=[M1,M2,...,Mn]:
ktensor(lambda::AbstractVector{<:Number},fmat::Array{Matrix{T},1},isorth::Bool) where {T<:Number}=ktensor(lambda,MatrixCell(fmat),isorth)
ktensor(lambda::AbstractVector{<:Number},fmat::Array{Matrix{T},1}) where {T<:Number}=ktensor(lambda,MatrixCell(fmat),true)
ktensor(fmat::Array{Matrix{T},1}) where {T<:Number}=ktensor(MatrixCell(fmat))
"""
    randktensor(I::Vector,R::Integer)
    randktensor(I::Integer,R::Integer,N::Integer)

Create random Kruskal tensor of size I with R components, or of order N and size I × ⋯ × I.
"""
function randktensor(sz::AbstractVector{<:Integer},R::Integer)
  fmat=[randn(sz[n],R) for n=1:length(sz)] #create random factor matrices
  ktensor(fmat)
end
randktensor(sz::Number,R::Number,N::Integer)=randktensor(repeat([sz],N),R)
#For input defined as tuples or nx1 matrices - ranttensor(([I,I,I],[R,R,R]))

"""
    arrange(X[,n::Integer])
    arrange(X[,P::Vector])

Arrange the rank-1 components of a ktensor: normalize the columns of the factor matrices and then sort the ktensor components by magnitude, greatest to least.

## Arguments:
- `n`: Absorb the weights into the nth factor matrix instead of lambda.
- `P`: Rearrange the components of X according to the permutation P. P should be a permutation of 1 to ncomponents(X).
"""
function arrange(X::ktensor,mode=-1)
    N=ndims(X)
    Xnorm=normalize(X) #Ensure that matrices are normalized
    p=sortperm(Xnorm.lambda,rev=true) #Sort
    lambda=Xnorm.lambda[p]
    fmat=MatrixCell(undef,N)
    [fmat[n]=Xnorm.fmat[n][:,p] for n=1:N]
    #Absorb the weight into one factor, if requested
    if mode>0
        fmat[mode]=fmat[mode].*repeat(lambda',size(X,mode),1) #fmat[mode]*diagm(lambda)
        lambda = fill(1,size(lambda))#ones(lambda)
    end
    ktensor(lambda,fmat)
end
function arrange(X::ktensor,perm::AbstractVector{<:Integer})
    N=ndims(X)
    @assert(collect(1:ncomponents(X))==sort(perm),"Invalid permutation.")
    lambda = X.lambda[perm]
    fmat=MatrixCell(undef,N)
    [fmat[n]=X.fmat[n][:,perm] for n=1:N]
    ktensor(lambda,fmat)
end

"""
    arrange!(X[,n::Integer])
    arrange!(X[,P::Vector])

Arrange the rank-1 components of a ktensor: normalize the columns of the factor matrices and then sort the ktensor components by magnitude, greatest to least.
Rewrite ktensor.

## Arguments:
- `n`: Absorb the weights into the nth factor matrix instead of lambda.
- `P`: Rearrange the components of X according to the permutation P. P should be a permutation of 1 to ncomponents(X).
"""
function arrange!(X::ktensor,mode=-1)
    N=ndims(X)
    normalize!(X) #Ensure that matrices are normalized
    p=sortperm(X.lambda,rev=true) #Sort
    X.lambda=X.lambda[p]
    [X.fmat[n]=X.fmat[n][:,p] for n=1:N]
    #Absorb the weight into one factor, if requested
    if mode>0
        X.fmat[mode]=X.fmat[mode].*repeat(X.lambda',size(X,mode),1) #X.fmat[mode]*diagm(X.lambda)
        X.lambda = fill(1,size(X.lambda))#ones(X.lambda)
    end
    X
end
function arrange!(X::ktensor,perm::AbstractVector{<:Integer})
    N=ndims(X)
    @assert(collect(1:ncomponents(X))==sort(perm),"Invalid permutation.")
    X.lambda = X.lambda[perm]
    [X.fmat[n]=X.fmat[n][:,perm] for n=1:N]
    X
end

#Compute a CP decomposition with R components of a tensor X. **Documentation in tensor.jl.
function cp_als(X::ktensor,R::Integer;init="rand",tol=1e-4,maxit=1000,dimorder=[])
    N=ndims(X)
    nr=norm(X)
    K=ktensor
    if length(dimorder) == 0
        dimorder=collect(1:N)
    end
    fmat=MatrixCell(undef,N)
    if isa(init,Vector) || isa(init,MatrixCell)
        @assert(length(init)==N,"Wrong number of initial matrices.")
        for n in dimorder[2:end]
            @assert(size(init[n])==(size(X,n),R),"$(n)-th initial matrix is of wrong size.")
            fmat[n]=init[n]
        end
    elseif init=="rand"
        [fmat[n]=rand(size(X,n),R) for n in dimorder[2:end]]
    elseif init=="eigs" || init=="nvecs"
        [fmat[n]=nvecs(X,n,R) for n in dimorder[2:end]]
    else
        error("Initialization method wrong.")
    end
    G = zeros(R,R,N); #initalize gramians
    [G[:,:,n]=fmat[n]'*fmat[n] for n in dimorder[2:end]]
    fit=0
    for k=1:maxit
        fitold=fit
        lambda=[]
        for n in dimorder
            fmat[n]=mttkrp(X,fmat,n)
            W=reshape(prod(G[:,:,setdiff(collect(1:N),n)],dims=3),Val(2))
            fmat[n]=fmat[n]/W
            if k == 1
                lambda = sqrt.(sum(fmat[n].^2,dims=1))[:] #2-norm
            else
                lambda = maximum(maximum(abs.(fmat[n]),dims=1),dims=1)[:] #max-norm
            end
            fmat[n] = fmat[n]./lambda'
            G[:,:,n] = fmat[n]'*fmat[n]
        end
        K=ktensor(lambda,fmat)
        if nr==0
            fit=norm(K)^2-2*innerprod(X,K)
        else
            nr_res=sqrt.(nr^2+norm(K)^2-2*innerprod(X,K))
            fir=1-nr_res/nr
        end
        fitchange=abs.(fitold-fit)
        if k>1 && fitchange<tol
            break
        end
    end
    arrange!(K)
    fixsigns!(K)
    K
end

#Display a ktensor. **Documentation in ttensor.jl.
function display(X::ktensor,name="ktensor")
    println("Kruskal tensor of size ",size(X),":\n")
    println("$name.lambda: ")
	flush(stdout)
	display(X.lambda)
    for n=1:ndims(X)
        println("\n\n$name.fmat[$n]:")
		flush(stdout)
		display(X.fmat[n])
    end
end

"""
    extract(X,factors)

Create a new ktensor with only the specified factors.
"""
function extract(X::ktensor,factors::AbstractVector{<:Integer})
    N=ndims(X)
    lambda = X.lambda[factors]
    fmat = MatrixCell(undef,N)
    [fmat[n]=X.fmat[n][:,factors] for n=1:N]
    ktensor(lambda, fmat)
end
extract(X::ktensor,factors::Integer)=extract(X,[factors])

"""
    fixsign(X)

Fix sign ambiguity of a ktensor.
"""
function fixsigns(X::ktensor)
    N=ndims(X)
    sgn=zeros(Int,N)
    fmat=deepcopy(X.fmat)
     for r=1:ncomponents(X)
        for n=1:N
            maxind=argmax(abs.(X.fmat[n][:,r]))
            sgn[n]=sign(X.fmat[n][maxind,r])
        end
        negind=findall(sgn.==-1)
        nflip=2*floor(Int,length(negind)/2)
        for i=1:nflip
            fmat[negind[i]][:,r]=-X.fmat[negind[i]][:,r]
        end
    end
    ktensor(X.lambda,fmat)
end
"""
    fixsign(X)

Fix sign ambiguity of a ktensor. Rewrite ktensor.
"""
function fixsigns!(X::ktensor)
    N=ndims(X)
    sgn=zeros(Int,N)
     for r=1:ncomponents(X)
        for n=1:N
            maxind=argmax(abs.(X.fmat[n][:,r]))
            sgn[n]=sign(X.fmat[n][maxind,r])
        end
        negind=findall(sgn.==-1)
        nflip=2*floor(Int,length(negind)/2)
        for i=1:nflip
            X.fmat[negind[i]][:,r]=-X.fmat[negind[i]][:,r]
        end
    end
end

#Makes full tensor out of a ktensor. **Documentation in ttensor.jl.
function full(X::ktensor)
  reshape(Array(X.lambda'*khatrirao(X.fmat[end:-1:1])'),size(X)...)
end

#Inner product of two ktensors. **Documentation in tensor.jl.
function innerprod(X1::ktensor,X2::ktensor)
  @assert(size(X1)==size(X2),"Dimension mismatch.")
  inpr=X1.lambda*X2.lambda'
  for n=1:ndims(X1)
    inpr=inpr.*(X1.fmat[n]'*X2.fmat[n])
  end
  sum(inpr[:])
end
function innerprod(X1::ktensor,X2)
  @assert(isa(X2,ttensor) || isa(X2,Array), "Inner product not available for type $(typeof(X2)).")
  @assert(size(X1)==size(X2),"Dimension mismatch.")
  inpr=0
  for r=1:ncomponents(X1)
    inpr+=X1.lambda[r]*ttv(X2,[X1.fmat[n][:,r] for n=1:ndims(X1)])[1]
  end
  inpr
end
innerprod(X1,X2::ktensor)=innerprod(X2,X1)

#Two ktensors are equal if they have equal components. **Documentation in ttensor.jl.
function isequal(X1::ktensor,X2::ktensor)
  if (X1.lambda == X2.lambda) && (X1.fmat == X2.fmat)
    true
  else
    false
  end
end
==(X1::ktensor,X2::ktensor)=isequal(X1,X2)

"""
    krank(X::ktensor)

Represenation k-rank. Equal to number of columns of factor matrices if ktensor X.
"""
function krank(X::ktensor)
    size(X.fmat[1],2)
end

"""
    kten2TT(X::ktensor)

Transform ktensor to TTtensor.
"""
function kten2TT(X::ktensor)
    N=ndims(X)
    Isz=size(X)
    r=krank(X)
    G=CoreCell(undef,N)
    G[1]=reshape(X.fmat[1],(1,Isz[1],r))
    G[N]=reshape(X.fmat[N]',(r,Isz[N],1))
    for n=2:N-1
        G[n]=zeros(r,Isz[n],r)
        for i=1:Isz[n]
            G[n][:,i,:]=Diagonal(X.fmat[n][i,:])
        end
    end
    TTtensor(G)
end

#Subtraction of two ktensors. **Documentation in ttensor.jl.
function minus(X1::ktensor,X2::ktensor)
  1*X1+(-1)*X2
end
-(X1::ktensor,X2::ktensor)= minus(X1,X2)

#Scalar times tensor. **Documentation in ttensor.jl.
function mtimes(α::Number,X::ktensor)
	ktensor(α*X.lambda,X.fmat);
end
*(α::Number,X::ktensor)=mtimes(α,X)
*(X::ktensor,α::Number)=*(α,X)

#Matricized ktensor times Khatri-Rao product. **Documentation in tensor.jl.
function mttkrp(X::ktensor,M::MatrixCell,n::Integer)
  N=ndims(X)
  @assert(length(M) == N,"Wrong number of matrices")
  n==1 ? R=size(M[2],2) : R=size(M[1],2)
  modes=setdiff(1:N,n)
  W=repeat(X.lambda,1,R) #matrix of weights
  for m in modes
    W=W.*(X.fmat[m]'*M[m])
  end
  X.fmat[n]*W
end
mttkrp(X::ktensor,M::Array{Matrix{T},1},n::Integer) where {T<:Number}=mttkrp(X,MatrixCell(M),n)

"""
    ncomponents(X)

Number of components of a ktensor.
"""
function ncomponents(X::ktensor)
  length(X.lambda)
end

#Number of modes of a ttensor. **Documentation in Base.
function ndims(X::ktensor)
  length(X.fmat)
end

"""
    normalize(X[,n;normtype,factor])

Normalize columns of factor matrices of a ktensor. Also ensures that lambda is positive.

## Arguments:
- `n`: Absorbe the excess weight into nth factor matrix.
   If n=0, equally divide the weights across the factor matrices and set lambda to vector of ones.
   If n="sort", sort the components by lambda value, from greatest to least.
- `normtype`: Default: 2.
- `factor`: Just normalize specified factor.
"""
function normalize(X::ktensor,mode=-1;normtype=2,factor=-1)
  if mode=="sort"
    mode=-2
  end
  N=ndims(X)
  lambda=deepcopy(float(X.lambda))
  fmat=deepcopy(X.fmat)
  if factor>0
    @assert(0<factor<=N,"Factor exceeds dimension.")
    for r = 1:ncomponents(X)
      nr = norm(fmat[factor][:,r],normtype);
      lambda[r]=lambda[r]*nr
      if nr > 0
         fmat[factor][:,r]=fmat[factor][:,r]/nr
      end
    end
  end
  for r=1:ncomponents(X)
    for n=1:N
      nr=norm(fmat[n][:,r],normtype)
      lambda[r]=lambda[r]*nr
      if nr>0
        fmat[n][:,r]=fmat[n][:,r]/nr
      end
    end
  end
  #Check that all the lambda values are positive
  negind = findall(lambda .< 0)
  for i in negind
    lambda[i] = -1 * lambda[i]
    fmat[1][:,i] = -1 * fmat[1][:,i]
  end
  if mode == 0
    d=(lambda').^(1/N)
   [fmat[n]=fmat[n].*repeat(d,size(X,n),1) for n=1:N] #fmat[n]*diagm(d)
    lambda = fill(1,size(lambda))#ones(lambda)
  elseif mode > 0
    fmat[mode] = fmat[mode].*repeat(lambda',size(X,mode),1) #fmat[mode]*diagm(lambda)
    lambda = fill(1,size(lambda))#ones(lambda)
  elseif mode==-2
    p = sortperm(lambda,rev=true)
    return arrange(ktensor(lambda,fmat),p)
  end
  ktensor(lambda,fmat)
end

"""
    normalize!(X[,n,normtype,factor])

Normalize columns of factor matrices of a ktensor. Rewrite ktensor. Also ensures that lambda is positive.

## Arguments:
- `n`: Absorbe the excess weight into nth factor matrix.
   If n=0, equally divide the weights across the factor matrices and set lambda to vector of ones.
   If n="sort", sort the components by lambda value, from greatest to least.
- `normtype`: Default: 2.
- `factor`: Just normalize specified factor.
"""
function normalize!(X::ktensor,mode=-1;normtype=2,factor=-1)
  if mode=="sort"
    mode=-2
  end
  N=ndims(X)
  X.lambda=float(X.lambda)
  if factor>0
    @assert(0<factor<=N,"Factor exceeds dimension.")
    for r = 1:ncomponents(X)
      nr = norm(X.fmat[factor][:,r],normtype);
      X.lambda[r]=X.lambda[r]*nr
      if nr > 0
         X.fmat[factor][:,r]=X.fmat[factor][:,r]/nr
      end
    end
  end
  for r=1:ncomponents(X)
    for n=1:N
      nr=norm(X.fmat[n][:,r],normtype)
      X.lambda[r]=X.lambda[r]*nr
      if nr>0
        X.fmat[n][:,r]=X.fmat[n][:,r]/nr
      end
    end
  end
  #Check that all the lambda values are positive
  negind = findall(X.lambda .< 0)
  for i in negind
    X.fmat[1][:,i] = -1 * X.fmat[1][:,i]
    X.lambda[i] = -1 * X.lambda[i]
  end
  if mode == 0
     d=(X.lambda').^(1/N)
     [X.fmat[n]=X.fmat[n].*repeat(d,size(X,n),1) for n=1:N] #X.fmat[n]*diagm(d)
     X.lambda = fill(1,size(X.lambda))#ones(X.lambda)
  elseif mode > 0
    X.fmat[mode] = X.fmat[mode].*repeat(X.lambda',size(X,mode),1) #X.fmat[n]*diagm(X.lambda)
    X.lambda = fill(1,size(X.lambda))#ones(X.lambda)
  elseif mode==-2
    p = sortperm(X.lambda,rev=true)
    X = arrange(X,p)
  end
   X
end

#Computes the r leading singular vectors of mode-n matricization of a tensor X.
# **Documentation in tensor.jl.
function nvecs(X::ktensor,n::Integer,r=0;flipsign=false)
    M = X.lambda * X.lambda'
    for m in setdiff(1:ndims(X),n)
        M = M .* (X.fmat[m]' * X.fmat[m])
    end
    #V=eigs(Symmetric(X.fmat[n] * M * X.fmat[n]'),nev=r,which=:LM)[2] #has bugs
    V=eigen(Symmetric(X.fmat[n] * M * X.fmat[n]')).vectors[:,end:-1:end-r+1]
    if flipsign
        maxind = findmax(abs.(V),1)[2]
        for i = 1:r
            ind=ind2sub(size(V),maxind[i])
            if V[ind...] < 0
                V[:,ind[2]] = V[:,ind[2]] * -1
            end
        end
    end
    V
end

#Permute dimensions of a ktensor. **Documentation in Base.
function permutedims(X::ktensor,perm::AbstractVector{<:Integer})
    N=ndims(X)
    @assert(collect(1:N)==sort(perm),"Invalid permutation.")
    ktensor(X.lambda,X.fmat[perm])
end

#Addition of two ktensors. **Documentation in ttensor.jl.
function plus(X1::ktensor,X2::ktensor)
	@assert(size(X1) == size(X2),"Dimension mismatch.")
  	lambda=[X1.lambda; X2.lambda]
	fmat=[[X1.fmat[n] X2.fmat[n]] for n=1:ndims(X1)] #concatenate factor matrices
	ktensor(lambda,fmat)
end
+(X1::ktensor,X2::ktensor)=plus(X1,X2)

"""
    redistribute(X,n)

Distribute lambda values to a mode n.
"""
function redistribute(X::ktensor,mode::Integer)
    lambda=deepcopy(float(X.lambda))
    fmat=deepcopy(X.fmat)
    for r = 1:ncomponents(X)
        fmat[mode][:,r] = fmat[mode][:,r] * lambda[r];
        lambda[r] = 1;
    end
    ktensor(lambda,fmat)
end

"""
    redistribute!(X,n)

Distribute lambda values to a mode n. Rewrite ktensor X.
"""
function redistribute!(X::ktensor,mode::Integer)
    for r = 1:ncomponents(X)
        X.fmat[mode][:,r] = X.fmat[mode][:,r] * X.lambda[r];
        X.lambda[r] = 1;
    end
    X
end

#Size of a ktensor. **Documentation in Base.
function size(X::ktensor)
  tuple([size(X.fmat[n],1) for n=1:length(X.fmat)]...)
end
#Size of n-th mode of a ktensor.
function size(X::ktensor,n::Integer)
  size(X.fmat[n],1)
end

function Base.show(io::IO,X::ktensor)
    display(X)
end

#Mode-n matricization of a ktensor. **Documentation in tensor.jl.
tenmat(X::ktensor,n::Integer)=tenmat(full(X),n)

"""
tocell(X[,n])

Converts ktensor X into MatrixCell. If n specified, absorb the weights into the nth factor matrix.
"""
function tocell(X::ktensor,mode::Integer)
    Y = normalize(X,mode)
    Y.fmat
end
function tocell(X::ktensor)
    if isequal(X.lambda,ones(size(X.lambda)))
        return X.fmat;
    end
    N=ndims(X)
    d = (X.lambda').^(1/N)
    fmat=MatrixCell(undef,N)
    for n = 1:N
        fmat[n] = X.fmat[n].*repeat(d,size(X.fmat[n],1),1)
    end
    fmat
end

"""
---
    ttensor(X::ktensor)

Create ttensor out of a ktensor.
"""
function ttensor(X::ktensor)
    Xred=redistribute(X,1)
    ttensor(neye(ncomponents(X),order=ndims(X)),Xred.fmat)
end

#Mode-n multiplication of ktensor and matrix. **Documentation in tensor.jl.
#t='t' transposes matrices
function ttm(X::ktensor,M::MatrixCell,modes::AbstractVector{<:Integer},t='n')
  if t=='t'
	 [M[n]=M[n]' for n=1:length(M)]
	end
	@assert(length(modes)<=length(M),"Too few matrices")
	@assert(length(M)<=ndims(X),"Too many matrices")
  if length(modes)<length(M)
    M=M[modes]
  end
  fmat=deepcopy(X.fmat)
  for n=1:length(modes)
    fmat[modes[n]]=M[n]*fmat[modes[n]]
  end
    ktensor(X.lambda,fmat)
end
ttm(X::ktensor,M::MatrixCell,modes::AbstractRange{<:Integer},t='n')=ttm(X,M,collect(modes),t)
ttm(X::ktensor,M::Matrix{<:Number},n::Integer,t='n')=ttm(X,MatrixCell([M]),[n],t)
ttm(X::ktensor,M::MatrixCell,t::Char)=ttm(X,M,1:length(M),t)
ttm(X::ktensor,M::MatrixCell)=ttm(X,M,1:length(M))
function ttm(X::ktensor,M::MatrixCell,n::Integer,t='n')
	if n>0
		ttm(X,M[n],n,t)
	else
		modes=setdiff(1:ndims(X),-n)
		ttm(X,M,modes,t)
	end
end
#If array of matrices isn't defined as MatrixCell, but as M=[M1,M2,...,Mn]:
ttm(X::ktensor,M::Array{Matrix{T},1},modes::AbstractVector{<:Integer},t='n') where {T<:Number}=ttm(X,MatrixCell(M),modes,t)
ttm(X::ktensor,M::Array{Matrix{T},1},modes::AbstractRange{<:Integer},t='n') where {T<:Number}=ttm(X,MatrixCell(M),modes,t)
ttm(X::ktensor,M::Array{Matrix{T},1},t::Char) where {T<:Number}=ttm(X,MatrixCell(M),t)
ttm(X::ktensor,M::Array{Matrix{T},1}) where {T<:Number}=ttm(X,MatrixCell(M))
ttm(X::ktensor,M::Array{Matrix{T},1},n::Integer,t='n') where {T<:Number}=ttm(X,MatrixCell(M),n,t)

#Mode-n multiplication ktensor and vector. **Documentation in tensor.jl.
#t='t' transposes matrices
function ttv(X::ktensor,V::VectorCell,modes::AbstractVector{<:Integer})
  N=ndims(X)
  remmodes=setdiff(1:N,modes)
  lambda=deepcopy(float(X.lambda))
  if length(modes) < length(V)
    V=V[modes]
  end
  for n=1:length(modes)
      lambda=lambda.*(X.fmat[modes[n]]'*V[n])
  end
  if isempty(remmodes)
    sum(lambda)
  else
    ktensor(lambda,X.fmat[remmodes])
  end
end
ttv(X::ktensor,v::AbstractVector{<:Number},n::Integer)=ttv(X,VectorCell([v]),[n])
ttv(X::ktensor,V::VectorCell,modes::AbstractRange{<:Integer})=ttv(X,V,collect(modes))
ttv(X::ktensor,V::VectorCell)=ttv(X,V,1:length(V))
function ttv(X::ktensor,V::VectorCell,n::Integer)
	if n>0
		ttm(X,V[n],n)
	else
		modes=setdiff(1:ndims(X),-n)
		ttm(X,A,modes)
	end
end
#If array of vectors isn't defined as VectorCell, but as V=[v1,v2,...,vn]:
ttv(X::ktensor,V::Array{Vector{T},1},modes::AbstractVector{<:Integer}) where {T<:Number}=ttv(X,VectorCell(V),modes)
ttv(X::ktensor,V::Array{Vector{T},1},modes::AbstractRange{<:Integer}) where {T<:Number}=ttv(X,VectorCell(V),modes)
ttv(X::ktensor,V::Array{Vector{T},1}) where {T<:Number}=ttv(X,VectorCell(V))
ttv(X::ktensor,V::Array{Vector{T},1},n::Integer) where {T<:Number}=ttv(X,VectorCell(V),n)

#**Documentation in ttensor.jl.
uminus(X::ktensor)=mtimes(-1,X)
-(X::ktensor)=uminus(X)

#Frobenius norm of a ktensor. **Documentation in Base.
function norm(X::ktensor)
  nrm=X.lambda*X.lambda'
  for n=1:ndims(X)
    nrm=nrm.*(X.fmat[n]'*X.fmat[n])
  end
  sqrt.(abs.(sum(nrm[:])))
end


# using Calculus
# define types first
struct Substance
    name::String
    density::Function
    enthalpy::Function
    heat_capacity::Function
    viscosity::Function
    thermal_conductivity::Function
    alt_heat_capacity::Function
    visc_kin::Function
    exp_coeff::Function
    therm_diff::Function
    stens::Function
    Pr::Function
end

function Substance(name, x::Vector{Function})
    dens, enth, Cp, visc, thermCond, stens = x
    function alt_heat_capacity(p, T)
        return ForwardDiff.derivative(T_ -> enth(p, T_), T)
    end
    function visc_kin(p, T)
        return visc(p, T)/dens(p, T)
    end
    function exp_coeff(p, T)
        dRdT = ForwardDiff.derivative(T_ -> dens(p, T_), T)
        return (-1/dens(p, T))*dRdT
    end
    function thermDiff(p, T)
        if Cp == nothing
            Cp_ = alt_heat_capacity(p, T)
        else
            Cp_ = Cp(p, T)
        end
        return thermCond(p, T) / (dens(p,T)*Cp_)
    end
    function Pr(p, T)
        if Cp == nothing
            Cp_ = alt_heat_capacity(p, T)
        else
            Cp_ = Cp(p, T)
        end
        return Cp_ * visc(p, T) / thermCond(p, T)
    end
    return Substance(name, dens, enth, Cp, visc, thermCond, alt_heat_capacity, visc_kin, exp_coeff, thermDiff, stens, Pr)
end

function na_enthalpy(p, T)
    """This function takes the temperature of liquid sodium in K and returns the enthalpy in kJ/kg.
    The reference state is taken as solid sodium at 298 K. The correlation takes this form:
    H(l, T) - H(s, 298.15) = 365.77 + 1.6582 T - 4.2395 × 10^-4 * T^2 + 1.4847 × 10^-7 * T^3 + 2992.6 * T^-1
    This correlation is recommended by a report prepared by Argonne National Lab that can be found here:
    http://www.ne.anl.gov/eda/ANL-RE-95-2.pdf
    The correlation is valid between 371 and 2000 K."""
    a = -365.77
    b = 1.6582
    c = 4.2395E-4
    d = 1.4847E-7
    e = 2992.6
    return (a + b * T - c * T ^ 2 + d * T ^ 3 + e / T) * 1000
end

function na_Cp(p, T)
    """This function takes the temperature of liquid sodium and returns the heat
    capacity at constant pressure in kJ/kg*K. The Argonne report gives a very
    complicated correlation for Cp but notes that it's almost equal to the
    ForwardDiff.derivative of enthalpy with temperature up to about 1800 K. So I used the ForwardDiff.derivative here."""
    a = -365.77
    b = 1.6582
    c = 4.2395E-4
    d = 1.4847E-7
    e = 2992.6
    return (b - 2*c * T + 3 * d * T ^ 2 - e / (T^2)) * 1000
    #return ForwardDiff.derivative(enthalpy, T, dx = 1E-6)
end

function na_density(p, T)
    """This function takes the temperature of liquid sodium in K and returns the density in kg/m^3.
    The correlation takes this form:
    rho = rho_c + f * (1 - Tr) + g * (1 - Tr)^h
    where rho_c is the critical density of sodium, Tr is the reduced temperature T/T_c,
    and f g and h are constants given in the function.
    This correlation is recommended by a report prepared by Argonne National Lab that can be found here:
    http://www.ne.anl.gov/eda/ANL-RE-95-2.pdf
    The correlation is valid between 371 and ~2000 K."""
    rho_c = 219
    f = 275.32
    g = 511.58
    h = 0.5
    Tc = 2503.7
    Tr = T / Tc
    return rho_c + f * (1 - Tr) + g * (1 - Tr) ^ h
end

function na_thermal_conductivity(p, T)
    """This function takes the temperature of liquid sodium in K and returns the thermal conductivity in W/m*K.
    The reference state is taken as solid sodium at 298 K. The correlation takes this form:
    k = 124.67 - 0.11381 * T + 5.5226 × 10^-5 * T^2 - 1.1842 × 10^-8 * T^3
    This correlation is recommended by a report prepared by Argonne National Lab that can be found here:
    http://www.ne.anl.gov/eda/ANL-RE-95-2.pdf
    The correlation is valid between 371 and 2000 K."""
    a = 124.67
    b = 0.11381
    c = 5.5226E-5
    d = 1.1842E-8
    return a - b * T + c * T ^ 2 - d * T ^ 3
end

function na_dynamic_viscosity(p, T)
    """This function takes the temperature of liquid sodium in K and returns the dynamic viscosity in Pa * s.
    The reference state is taken as solid sodium at 298 K. The correlation takes this form:
    ln eta = - 6.4406 - 0.3958 ln (T) + 556.835 / T
    This correlation is recommended by a report prepared by Argonne National Lab that can be found here:
    http://www.ne.anl.gov/eda/ANL-RE-95-2.pdf
    The correlation is valid between 371 and 2500 K."""
    a = -6.4406
    b = -0.3958
    c = 556.835
    return exp(a + b * log(T) + c / T)
end


@testset "Quadratic and hermitian forms" begin
  include("QuadForm/Basic.jl")
  include("QuadForm/Lattices.jl")
  include("QuadForm/Genus.jl")
  include("QuadForm/Morphism.jl")
  include("QuadForm/Enumeration.jl")
  include("QuadForm/MassQuad.jl")
  include("QuadForm/MassHerm.jl")
  include("QuadForm/Quad.jl")
  include("QuadForm/QuadBin.jl")
  include("QuadForm/Herm.jl")
  include("QuadForm/Torsion.jl")
  include("QuadForm/PadicLift.jl")
end


using CLArrays, GLVisualize, GeometryTypes, GLAbstraction, StaticArrays

TY = Float32
N = 1024
const h    = TY(2*π/N)
const epsn = TY(h * .5)
const C    = TY(2/epsn)
const tau  = TY(epsn * h)
Tfinal = 50.

S(x,y) = exp(-x^2/0.1f0)*exp(-y^2/0.1f0)

ArrayType = CLArray
# real-space and reciprocal-space grids
# the real-space grid is just used for plotting!
X_cpu = convert.(TY, collect(linspace(-pi+h, pi, N)) .* ones(1,N))
X = ArrayType(X_cpu);
k = collect([0:N/2; -N/2+1:-1]);
Â = ArrayType(convert.(TY,kron(k.^2, ones(1,N)) + kron(ones(N), k'.^2)));

# initial condition
uc = ArrayType(TY(2.0)*(rand(TY, N, N)-TY(0.5)))

#################################################################
#################################################################

pow3(u) = complex((u * u * u) - u)
function take_step!(u, Â, t_plot, fftplan!, ifftplan!, u3fft, uc, tmp)
    u3fft .= pow3.(u)
    fftplan! * u3fft
    uc .= complex.(u)
    fftplan! * uc
    @. tmp .= ((1f0+C*tau*Â) .* uc .- tau/epsn * (Â .* u3fft)) ./ (1f0+(epsn*tau)*Â.^2f0+C*tau*Â)
    ifftplan! * tmp
    u .= real.(tmp)
    nothing
end
function normalise_af!(u, out)
    out .= u .- minimum(u)
    out .= out ./ maximum(out)
    nothing
end
#################################################################
#################################################################

n = 1
T_plot = 0.01; t_plot = 0.0
ceil(Tfinal / tau)
up = copy(uc)
ucc = complex.(uc)
fftplan! = plan_fft!(ucc)
ifftplan! = plan_ifft!(ucc)
u3fft = similar(ucc)
tmp = similar(ucc)

w = glscreen(resolution = (N, N))
normalise_af!(uc,up)
up .= up .* 0.1f0
robj = visualize(
    reinterpret(Intensity{Float32}, Array(up)),
    stroke_width = 0f0,
    levels = 20f0,
    color_map = Colors.colormap("RdBu", 100),
    color_norm = Vec2f0(0, 0.1)
).children[]
_view(robj, position = Vec3f0(0, 0.5, 2.5), lookat = Vec3f0(0))
io, buffer = GLVisualize.create_video_stream(homedir()*"/Desktop/pd1.mkv", w)
center!(w, :orthographic_pixel)
idx = 0
while isopen(w)
    take_step!(uc, Â, t_plot, fftplan!, ifftplan!, u3fft, ucc, tmp)
    t_plot += tau
    normalise_af!(uc, up)
    up .= up .* 0.1f0
    GLAbstraction.update!(robj[:intensity], reinterpret(SVector{1, Float32}, Array(up)))
    GLWindow.poll_glfw()
    GLWindow.reactive_run_till_now()
    GLWindow.render_frame(w)
    GLWindow.swapbuffers(w)
    if idx == 4
        GLVisualize.add_frame!(io, w, buffer)
        idx = 0
    end
    idx += 1
end
close(io)
GLWindow.destroy!(w)


using Revise
using ADCME
using ADCMEKit
using NNFEM
using PyPlot
using ProgressMeter
using Printf
using MAT
using Clustering
using JLD2

close("all")
include("load_domain_function.jl")

NT = 100
Δt = 30/NT

mesh = "crack_wider.msh"
c1 = (3.5, 0.0)
c2 = (4.0, 0.5)
slip_scale = 5.0 * sqrt(2)

#################### solve static elastic response  ####################
domain = load_crack_domain(mesh=mesh, c1=c1, c2=c2, slip_scale=slip_scale)
# domain = load_crack_domain()

H = domain.elements[1].mat[1].H
EBC_func = nothing
FBC_func = nothing
body_func = nothing
globaldata = GlobalData(missing, missing, missing, missing, domain.neqs, EBC_func, FBC_func, body_func)
Fext = compute_external_force(globaldata, domain)
d = LinearStaticSolver(globaldata, domain, domain.state, H, Fext)
sess = Session(); init(sess)
domain.state = run(sess, d)

#################### solve viscoelastic relaxation  ####################
strain = getStrain(domain)
stress = getStress(domain)
state = domain.state

domain = load_crack_domain(mesh=mesh, c1=c1, c2=c2, slip_scale=slip_scale, option="mixed")
# domain = load_crack_domain(option="mixed")

ts = GeneralizedAlphaSolverTime(Δt, NT)
ubd, abd = compute_boundary_info(domain, globaldata, ts)
Fext = compute_external_force(domain, globaldata, ts)

d0 = state
v0 = zeros(2domain.nnodes)
a0  = zeros(2domain.nnodes)
# σ0 = zeros(getNGauss(domain),3)
# ϵ0 = zeros(getNGauss(domain),3)
σ0 = stress
ϵ0 = strain

μ = zeros(getNGauss(domain))
λ = zeros(getNGauss(domain))
η = zeros(getNGauss(domain))
k = 0
for i = 1:domain.neles
  e = domain.elements[i]
  for mat in e.mat
    global k += 1
    μ[k] = mat.μ
    λ[k] = mat.λ
    η[k] = mat.η
  end
end

gnodes = getGaussPoints(domain)
η = 1e10 *( constant(3.0) - constant(1.5) * (gnodes[:,2])/maximum(gnodes[:,2]))

d, v, a, σ, ϵ = ViscoelasticitySolver(
  globaldata, domain, d0, v0, a0, σ0, ϵ0, Δt, NT, μ, λ, η, Fext, ubd, abd
)

sess = Session(); init(sess)
d_, σ_, η_ = run(sess, [d, σ, η])
matwrite("data/dippingfault_viscosity_inversion.mat", Dict(
  "sigma"=> σ_,
  "d" => d_,
  "eta" => η_,
  "d0" => d0,
  "v0" => v0,
  "a0" => a0,
  "epsilon0"=> ϵ0,
  "sigma0" => σ0,
  "NT" => NT,
  "dt" => Δt
))
@save "data/dippingfault_viscosity_inversion.jld2" domain


module GaussianProcesses

using Optim, PDMats, ElasticPDMats, ElasticArrays, Distances, FastGaussQuadrature, RecipesBase, Distributions, ProgressMeter
using StaticArrays
using StatsFuns, SpecialFunctions

using LinearAlgebra, Printf, Random, Statistics
import Statistics: mean, cov
import Base: size
import PDMats: dim, Matrix, diag, pdadd!, *, \, inv, logdet, eigmax, eigmin, whiten!, unwhiten!, quad, quad!, invquad, invquad!, X_A_Xt, Xt_A_X, X_invA_Xt, Xt_invA_X

# Functions that should be available to package
# users should be explicitly exported here
export GPBase, GP, GPE, GPA, ElasticGPE, Approx, predict_f, predict_y, Kernel, Likelihood, CompositeKernel, SumKernel, ProdKernel, Masked, FixedKernel, fix, Noise, Const, SE, SEIso, SEArd, Periodic, Poly, RQ, RQIso, RQArd, Lin, LinIso, LinArd, Matern, Mat12Iso, Mat12Ard, Mat32Iso, Mat32Ard, Mat52Iso, Mat52Ard, #kernel functions
    MeanZero, MeanConst, MeanLin, MeanPoly, SumMean, ProdMean, MeanPeriodic, #mean functions
    GaussLik, BernLik, ExpLik, StuTLik, PoisLik, BinLik,       #likelihood functions
    mcmc, ess, lss, optimize!, vi, var_exp, dv_var_exp, elbo, initialise_Q,    #inference functions
    set_priors!,set_params!, update_target!, autodiff, update_Q!, noise_variance
using ForwardDiff: GradientConfig, Dual, partials, copyto!, Chunk
import ForwardDiff: seed!

const φ = normpdf
const Φ = normcdf
const invΦ = norminvcdf

# all package code should be included here
include("means/means.jl")
include("kernels/kernels.jl")
include("likelihoods/likelihoods.jl")
include("common.jl")
include("utils.jl")
include("chol_utils.jl")
include("GP.jl")
include("GPE.jl")
include("GPEelastic.jl")
include("GPA.jl")
include("vi.jl")
include("mcmc.jl")
include("optimize.jl")
include("crossvalidation.jl")
include("plot.jl")
include("sparse/sparseGP.jl")

# ScikitLearnBase, which is a skeleton package.
include("ScikitLearn.jl")

end # module


module GraphHierarchy

using LightGraphs
using SimpleWeightedGraphs
using LinearAlgebra
using SparseArrays
using IterativeSolvers
using StatsBase
using Statistics

export forward_hierarchical_structure
export backward_hierarchical_structure
export hierarchical_structure
export forward_hierarchical_coefficients
export backward_hierarchical_coefficients
export hierarchical_coefficients

function weightedMeanStd(v::Array{T,1} where T <: AbstractFloat, w::Array{T,1} where T <: AbstractFloat)
    difftype = typeof(v[1])
    if difftype == typeof(w[1])
        av0 = mean(v, StatsBase.weights(w))
        return ( 1 - av0 ,std( v, StatsBase.weights(w), mean=av0, corrected=false))
    else
        ww = [ convert( difftype, c) for c in w ]
        av = mean(v, StatsBase.weights(ww))
        return ( 1 - av ,std( v, StatsBase.weights(ww), mean=av, corrected=false))
    end
end

function calculatingStructureAssumingTransposedMatrix(A::SparseMatrixCSC{T1,T2} where T1 <: Number where T2 <: Int)
    d = vec(sum(A, dims=2))
    dout = vec(sum(A, dims=1))
    M = spdiagm(0 => d)-A
    dropzeros!(M)
    HLs = lsqr(M,d)
    pos = findall(!iszero,A)
    HDs = sparse([p[2] for p in pos], [p[1] for p in pos], [HLs[p[1]] - HLs[p[2]] for p in pos], length(d), length(d))
    HDsums = sum(HDs, dims=1)
    ICs = [d[i] > 0 ? 1.0 - sum(HDs[:,i])/d[i] : 1.0 for i in 1:length(d)]
    RCs = [dout[i] > 0 ? sum(HDs[i,:])/dout[i] : 0.0 for i in 1:length(dout)]
    return (HLs, ICs, RCs, HDs )
end

function calculatingStructureAssumingTransposedMatrixBig(A::SparseMatrixCSC{T1,T2} where T1 <: Number where T2 <: Int)
    (x, y, v) = findnz(A)
    vb = -map(BigFloat,v)
    M = sparse(x, y, vb, size(A)[1], size(A)[1])
    d = -vec(sum(M, dims=2))
    dout = -vec(sum(M, dims=1))
    for i in 1:length(d)
        if d[i] > 0
            M[i,i] = d[i]
        end
    end
    HLs = lsqr(M,d)
    HDs = sparse(y, x, [ HLs[x[i]] - HLs[y[i]] for i in 1:length(x) ], length(d), length(d) )
    ICs = [d[i] > 0 ? BigFloat(1) - sum(HDs[:,i])/d[i] : BigFloat(1) for i in 1:length(d)]
    RCs = [dout[i] > 0 ? sum(HDs[i,:])/dout[i] : BigFloat(0)  for i in 1:length(dout)]
    return (HLs, ICs, RCs, HDs )
end

function calculatingStructureAssumingTransposedWeightedMatrix(A::SparseMatrixCSC{T1,T2} where T1 <: Number where T2 <: Int)
    (x, y, v) = findnz(A)
    d = vec(sum(A, dims=2))
    dout = vec(sum(A, dims=1))
    M = spdiagm(0 => d)-A
    dropzeros!(M)
    HLs = lsqr(M,d)
    HDs = sparse(y, x, [ HLs[x[i]] - HLs[y[i]] for i in 1:length(x) ], length(d), length(d) )
    HDsums = sum(HDs, dims=1)
    ICs = [d[i] > 0 ? 1.0 - A[i,:]'*HDs[:,i]/d[i] : 1.0 for i in 1:length(d)]
    RCs = [dout[i] > 0 ? A[:,i]'*HDs[i,:]/dout[i] : 0.0 for i in 1:length(dout)]
    return (HLs, ICs, RCs, HDs )
end

function calculatingStructureAssumingTransposedWeightedMatrixBig(A::SparseMatrixCSC{T1,T2} where T1 <: Number where T2 <: Int)
    (x, y, v) = findnz(A)
    vb = -map(BigFloat,v)
    M = sparse(x, y, vb, size(A)[1], size(A)[1])
    d = -vec(sum(M, dims=2))
    dout = -vec(sum(M, dims=1))
    for i in 1:length(d)
        if d[i] > 0
            M[i,i] = d[i]
        end
    end
    HLs = lsqr(M,d)
    HDs = sparse(y, x, [ HLs[x[i]] - HLs[y[i]] for i in 1:length(x) ], length(d), length(d) )
    HDsums = sum(HDs, dims=1)
    ICs = [d[i] > 0 ? BigFloat(1) - A[i,:]'*HDs[:,i]/d[i] : BigFloat(1) for i in 1:length(d)]
    RCs = [dout[i] > 0 ? A[:,i]'*HDs[i,:]/dout[i] : BigFloat(0) for i in 1:length(dout)]
    return (HLs, ICs, RCs, HDs )
end

function forward_hierarchical_structure(g::SimpleDiGraph{T} where T <: Int; big::Bool=false)
    if big
        return calculatingStructureAssumingTransposedMatrixBig(sparse(transpose(adjacency_matrix(g))))
    else
        return calculatingStructureAssumingTransposedMatrix(sparse(transpose(adjacency_matrix(g))))
    end
end

function backward_hierarchical_structure(g::SimpleDiGraph{T} where T <: Int; big::Bool=false)
    if big
        (xb, yb, wb, zb) = calculatingStructureAssumingTransposedMatrixBig(adjacency_matrix(g))
        return (xb, yb, wb, sparse(transpose(zb)))
    else
        (x, y, w, z) = calculatingStructureAssumingTransposedMatrix(adjacency_matrix(g))
        return (x, y, w, sparse(transpose(z)))
    end
end

function hierarchical_structure(g::SimpleDiGraph{T} where T <: Int; big::Bool=false)
    return (forward_hierarchical_structure(g, big = big),backward_hierarchical_structure(g, big = big))
end

function forward_hierarchical_structure(g::SimpleWeightedDiGraph{T1,T2} where T1 <: Int where T2 <: Number; big::Bool=false, weighted::Bool=true)
    if weighted
        if big
            return calculatingStructureAssumingTransposedWeightedMatrixBig(transpose(LightGraphs.weights(g)))
        else
            return calculatingStructureAssumingTransposedWeightedMatrix(transpose(LightGraphs.weights(g)))
        end
    else
        if big
            return calculatingStructureAssumingTransposedMatrixBig(sparse(transpose((!iszero).(adjacency_matrix(g)))))
        else
            return calculatingStructureAssumingTransposedMatrix(sparse(transpose((!iszero).(adjacency_matrix(g)))))
        end
    end
end

function backward_hierarchical_structure(g::SimpleWeightedDiGraph{T1,T2} where T1 <: Int where T2 <: Number; big::Bool=false, weighted::Bool=true)
    if weighted
        if big
            (xb, yb, wb, zb) = calculatingStructureAssumingTransposedWeightedMatrixBig(sparse(LightGraphs.weights(g)))
            return (xb, yb, wb, sparse(transpose(zb)))
        else
            (x, y, w, z) = calculatingStructureAssumingTransposedWeightedMatrix(sparse(LightGraphs.weights(g)))
            return (x, y, w, sparse(transpose(z)))
        end
    else
        if big
            (xb, yb, wb, zb) = calculatingStructureAssumingTransposedMatrixBig(sparse((!iszero).(adjacency_matrix(g))))
            return (xb, yb, wb, sparse(transpose(zb)))
        else
            (x, y, w, z) = calculatingStructureAssumingTransposedMatrix(sparse((!iszero).(adjacency_matrix(g))))
            return (x, y, w, sparse(transpose(z)))
        end
    end
end

function hierarchical_structure(g::SimpleWeightedDiGraph{T1,T2} where T1 <: Int where T2 <: Number; big::Bool=false, weighted::Bool=true)
    return (forward_hierarchical_structure(g, big = big, weighted = weighted),backward_hierarchical_structure(g, big = big, weighted = weighted))
end

function forward_hierarchical_structure(g::SimpleGraph{T} where T <: Int; big::Bool=false)
    if big
        return calculatingStructureAssumingTransposedMatrixBig(adjacency_matrix(g))
    else
        return calculatingStructureAssumingTransposedMatrix(adjacency_matrix(g))
    end
end

function backward_hierarchical_structure(g::SimpleGraph{T} where T <: Int; big::Bool=false)
    forward_hierarchical_structure(g, big = big)
end

function hierarchical_structure(g::SimpleGraph{T} where T <: Int; big::Bool=false)
    return forward_hierarchical_structure(g, big = big)
end

function forward_hierarchical_structure(g::SimpleWeightedGraph{T1,T2} where T1 <: Int where T2 <: Number; big::Bool=false, weighted::Bool=true)
    if weighted
        if big
            return calculatingStructureAssumingTransposedWeightedMatrixBig(sparse(LightGraphs.weights(g)))
        else
            return calculatingStructureAssumingTransposedWeightedMatrix(sparse(LightGraphs.weights(g)))
        end
    else
        if big
            return calculatingStructureAssumingTransposedMatrixBig(adjacency_matrix(g))
        else
            return calculatingStructureAssumingTransposedMatrix(adjacency_matrix(g))
        end
    end
end

function backward_hierarchical_structure(g::SimpleWeightedGraph{T1,T2} where T1 <: Int where T2 <: Number; big::Bool=false, weighted::Bool=true)
    return forward_hierarchical_structure(g, big = big, weighted = weighted)
end

function hierarchical_structure(g::SimpleWeightedGraph{T1,T2} where T1 <: Int where T2 <: Number; big::Bool=false, weighted::Bool=true)
    return forward_hierarchical_structure(g, big = big, weighted = weighted)
end

function forward_hierarchical_structure(A::SparseMatrixCSC{T1,T2} where T1 <: Number where T2 <: Int; big::Bool=false)
    if big
        return calculatingStructureAssumingTransposedMatrixBig(sparse(transpose(A)))
    else
        return calculatingStructureAssumingTransposedMatrix(sparse(transpose(A)))
    end
end

function backward_hierarchical_structure(A::SparseMatrixCSC{T1,T2} where T1 <: Number where T2 <: Int; big::Bool=false)
    if big
        return calculatingStructureAssumingTransposedMatrixBig(A)
    else
        return calculatingStructureAssumingTransposedMatrix(A)
    end
end

function hierarchical_structure(A::SparseMatrixCSC{T1,T2} where T1 <: Number where T2 <: Int; big::Bool=false)
    return ( forward_hierarchical_structure(A, big=big), backward_hierarchical_structure(A, big=big) )
end

function hierarchical_coefficients(g::Union{SimpleGraph{T},SimpleDiGraph{T}} where T <: Int, diffs::SparseMatrixCSC{T1,T2} where T1 <: AbstractFloat where T2 <: Int)
    diffarray = findnz(diffs)[3]
    av = mean(diffarray)
    return (1 - av, std(diffarray, mean=av, corrected=false))
end

function hierarchical_coefficients(g::Union{SimpleGraph{T},SimpleDiGraph{T}} where T <: Int, hstruc::Tuple{Array{T1,1},Array{T1,1},Array{T1,1},SparseMatrixCSC{T1,T2}} where T1 <: AbstractFloat where T2 <: Int )
    return hierarchical_coefficients(g, hstruc[4])
end

function hierarchical_coefficients(g::SimpleDiGraph{T} where T <: Int , hstruc::Tuple{Tuple{Array{T1,1},Array{T1,1},Array{T1,1},SparseMatrixCSC{T1,T2}},Tuple{Array{T1,1},Array{T1,1},Array{T1,1},SparseMatrixCSC{T1,T2}}} where T1 <: AbstractFloat where T2 <: Int )
    return (hierarchical_coefficients(g, hstruc[1]),hierarchical_coefficients(g, hstruc[2]))
end

function hierarchical_coefficients(g::Union{SimpleWeightedGraph{T1,T2},SimpleWeightedDiGraph{T1,T2}} where T1 <: Int where T2 <: AbstractFloat, diffs::SparseMatrixCSC{T1,T2} where T1 <: AbstractFloat where T2 <: Int; weighted::Bool = true )
    if weighted
        W0 = sparse(LightGraphs.weights(g))
        (xw,yw,vw) = findnz(W0)
        (x,y,v) = findnz(diffs)

        if sum(map(abs,x - xw)) + sum(map(abs,y - yw)) > 0
            error("Hierarchical differences do not match the weighted adjacency matrix:\nDifferences:", diffs, "\nWeighted Adjacency Matrix:", W0)
        end

        return weightedMeanStd(v,vw)
    else
        v = findnz(diffs)[3]
        av = mean(v)
        return (1 - av, std(v, mean=av, corrected=false))
    end
end

function hierarchical_coefficients(g::Union{SimpleWeightedGraph{T1,T2},SimpleWeightedDiGraph{T1,T2}} where T1 <: Int  where T2 <: AbstractFloat, hstruc::Tuple{Array{T1,1},Array{T1,1},Array{T1,1},SparseMatrixCSC{T1,T2}} where T1 <: AbstractFloat where T2 <: Int;  weighted::Bool = true  )
    return hierarchical_coefficients(g, hstruc[4], weighted=weighted)
end

function hierarchical_coefficients(g::SimpleWeightedDiGraph{T1,T2} where T1 <: Int  where T2 <: AbstractFloat , hstruc::Tuple{Tuple{Array{T1,1},Array{T1,1},Array{T1,1},SparseMatrixCSC{T1,T2}},Tuple{Array{T1,1},Array{T1,1},Array{T1,1},SparseMatrixCSC{T1,T2}}} where T1 <: AbstractFloat where T2 <: Int; weighted::Bool = true  )
    return (hierarchical_coefficients(g, hstruc[1], weighted=weighted), hierarchical_coefficients(g, hstruc[2], weighted=weighted))
end

end # module


module Stats
    import Base.quantile

    export

    # scalar_stats
    gmean,
    hmean,
    variation,
    sem,
    mad,
    skewness,
    kurtosis,
    minmax,
    midrange,
    range,
    tiedrank,
    percentile,
    quantile,
    quartile,
    quintile,
    decile,
    iqr,
    table,
    mode,
    modes,
    describe,

    # weighted stats
    wmean,

    # corr
    autocor,
    cor_spearman,
    cor_kendall,

    # others
    ## Types
    StatisticalModel,
    RegressionModel,

    ## Fcuntions
    coef,
    coeftable,
    confint,
    ecdf,
    findat,
    inverse_rle,
    loglikelihood,
    nobs,
    predict,
    residuals,
    model_response,
    rle,
    stderr,
    vcov

    include("scalar_stats.jl")
    include("weighted_stats.jl")
 #   include("corr.jl")
    include("others.jl")

end # module


# Direct product of different sectors
#==============================================================================#
const SectorTuple = Tuple{Vararg{Sector}}

struct ProductSector{T<:SectorTuple} <: Sector
    sectors::T
end
ProductSector{T}(args...) where {T<:SectorTuple} = ProductSector{T}(args)
Base.convert(::Type{ProductSector{T}}, t::Tuple) where {T<:SectorTuple} = ProductSector{T}(convert(T, t))

Base.one(::Type{ProductSector{T}}) where {G<:Sector, T<:Tuple{G}} = ProductSector((one(G),))
Base.one(::Type{ProductSector{T}}) where {G<:Sector, T<:Tuple{G,Vararg{Sector}}} = one(G) ⊗ one(ProductSector{tuple_type_tail(T)})

Base.conj(p::ProductSector) = ProductSector(map(conj, p.sectors))
×(p1::P, p2::P) where {P<:ProductSector} = SectorSet{P}(product(map(×,p1.sectors,p2.sectors)...))
Nsymbol(a::P, b::P, c::P) where {P<:ProductSector} = prod(map(Nsymbol, a.sectors, b.sectors, c.sectors))

fusiontype(::Type{<:ProductSector{T}}) where {T<:SectorTuple} = _fusiontype(T)
_fusiontype(::Type{Tuple{}}) = Abelian
_fusiontype(::Type{T}) where {T<:SectorTuple} = fusiontype(tuple_type_head(T)) & _fusiontype(tuple_type_tail(T))

braidingtype(::Type{<:ProductSector{T}}) where {T<:SectorTuple} = _braidingtype(T)
_braidingtype(::Type{Tuple{}}) = Abelian
_braidingtype(::Type{<:SectorTuple}) = fusiontype(tuple_type_head(T)) & _braidingtype(tuple_type_tail(T))

fermionparity(P::ProductSector) = _fermionparity(P.sectors)
_fermionparity(::Tuple{}) = false
_fermionparity(t::Tuple) = xor(fermionparity(t[1]), _fermionparity(tail(t)))

dim(p::ProductSector) = prod(dim, p.sectors)

# Default construction from tensor product of sectors
#-----------------------------------------------------
×(S1::Sector, S2::Sector) = ProductSector((S1, S2))
×(P1::ProductSector, S2::Sector) = ProductSector(tuple(P1.sectors..., S2))
×(S1::Sector, P2::ProductSector) = ProductSector(tuple(S1, P2.sectors...))
×(P1::ProductSector, P2::ProductSector) = ProductSector(tuple(P1.sectors..., P2.sectors...))

×(G1::Type{ProductSector{Tuple{}}}, G2::Type{ProductSector{T}}) where {T<:SectorTuple} = G2
×(G1::Type{ProductSector{T1}}, G2::Type{ProductSector{T2}}) where {T1<:SectorTuple,T2<:SectorTuple} =
    tuple_type_head(T1) ⊗ (ProductSector{tuple_type_tail(T1)} ⊗ G2)
×(G1::Type{ProductSector{Tuple{}}}, G2::Type{<:Sector}) = ProductSector{Tuple{G2}}
×(G1::Type{ProductSector{T}}, G2::Type{<:Sector}) where {T<:SectorTuple} = tuple_type_head(T) ⊗ (ProductSector{tuple_type_tail(T)} ⊗ G2)
×(G1::Type{<:Sector}, G2::Type{ProductSector{T}}) where {T<:SectorTuple} = ProductSector{tuple_type_cons(G1,T)}
×(G1::Type{<:Sector}, G2::Type{<:Sector}) = ProductSector{Tuple{G1,G2}}

function Base.show(io::IO, P::ProductSector)
    sectors = P.sectors
    sep = get(io, :compact, false) ? ", " : " × "
    print(io,"(")
    for i = 1:length(sectors)
        i == 1 || print(io, sep)
        print(io, sectors[i])
    end
    print(io,")")
end

function Base.show(io::IO, ::Type{ProductSector{T}}) where {T<:SectorTuple}
    sectors = T.parameters
    print(io,"(")
    for i = 1:length(sectors)
        i == 1 || print(io, " × ")
        print(io, sectors[i])
    end
    print(io,")")
end


@symbol_func function K_PB(cur_reactor::AbstractReactor)
  cur_scaling = cur_reactor.mode_scaling

  cur_K = cur_reactor.H

  cur_K *= cur_scaling[:constant]

  cur_K *= K_n(cur_reactor) ^ alpha_n_star(cur_reactor)

  cur_K /= K_kappa(cur_reactor)

  cur_K *= cur_reactor.epsilon ^ cur_scaling[:a]

  cur_K *= kappa_tau(cur_reactor) ^ cur_scaling[:kappa]

  cur_K *= cur_reactor.A ^ cur_scaling[:A]

  cur_K
end


function exp(a::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(a) && return a
    isinf(a) && return(signbit(a) ? zero(DoubleFloat{T}) : a)

    if iszero(HI(a))
    return one(DoubleFloat{T})
    elseif isone(abs(HI(a))) && iszero(LO(a))
    if HI(a) >= zero(T)
        return DoubleFloat{T}(2.718281828459045, 1.4456468917292502e-16)
    else # isone(-HI(a)) && iszero(LO(a))
        return DoubleFloat{T}(0.36787944117144233, -1.2428753672788363e-17)
    end
    elseif abs(HI(a)) >= 709.0
      if (HI(a) <= -709.0)
         return zero(DoubleFloat{T})
      else # HI(a) >=  709.0
         return inf(DoubleFloat{T})
      end
    end

    return calc_exp(a)
end

function exp_taylor(a::DoubleFloat{T}) where {T<:IEEEFloat}
    x = a
    x2 = x*x
    x3 = x*x2
    x4 = x2*x2
    x5 = x2*x3
    x10 = x5*x5
    x15 = x5*x10
    x20 = x10*x10
    x25 = x10*x15

    z = x + inv_fact[2]*x2 + inv_fact[3]*x3 + inv_fact[4]*x4
    z2 = x5 * (inv_fact[5] + x*inv_fact[6] + x2*inv_fact[7] +
         x3*inv_fact[8] + x4*inv_fact[9])
    z3 = x10 * (inv_fact[10] + x*inv_fact[11] + x2*inv_fact[12] +
         x3*inv_fact[13] + x4*inv_fact[14])
    z4 = x15 * (inv_fact[15] + x*inv_fact[16] + x2*inv_fact[17] +
         x3*inv_fact[18] + x4*inv_fact[19])
    z5 = x20 * (inv_fact[20] + x*inv_fact[21] + x2*inv_fact[22] +
         x3*inv_fact[23] + x4*inv_fact[24])
    z6 = x25 * (inv_fact[25] + x*inv_fact[26] + x2*inv_fact[27])

    ((((z6+z5)+z4)+z3)+z2)+z + one(DoubleFloat{T})
end

@inline exp_zero_half(a::DoubleFloat{T}) where {T<:IEEEFloat} = exp_taylor(a)

@inline function exp_half_one(a::DoubleFloat{T}) where {T<:IEEEFloat}
    z = mul_by_half(a)
    z = exp_zero_half(z)
    z = square(z)
    return z
end

function mul_by_half(r::DoubleFloat{T}) where {T<:IEEEFloat}
    frhi, xphi = frexp(HI(r))
    frlo, xplo = frexp(LO(r))
    xphi -= 1
    xplo -= 1
    hi = ldexp(frhi, xphi)
    lo = ldexp(frlo, xplo)
    return DoubleFloat{T}(hi, lo)
end

function mul_by_two(r::DoubleFloat{T}) where {T<:IEEEFloat}
    frhi, xphi = frexp(HI(r))
    frlo, xplo = frexp(LO(r))
    xphi += 1
    xplo += 1
    hi = ldexp(frhi, xphi)
    lo = ldexp(frlo, xplo)
    return DoubleFloat{T}(hi, lo)
end

function mul_pow2(r::DoubleFloat{T}, n::Int) where {T<:IEEEFloat}
    frhi, xphi = frexp(HI(r))
    frlo, xplo = frexp(LO(r))
    xphi += n
    xplo += n
    hi = ldexp(frhi, xphi)
    lo = ldexp(frlo, xplo)
    return DoubleFloat{T}(hi, lo)
end

function mul_pwr2(r::DoubleFloat{T}, n::Real) where {T<:IEEEFloat}
    m = 2.0^n
    return DoubleFloat{T}(HI(r)*m, LO(r)*m)
end

function Base.:(^)(r::DoubleFloat{T}, n::Int) where {T<:IEEEFloat}
    if (n == 0)
        return one(DoubleFloat{T})
    end

    s = one(DoubleFloat{T})
    nabs = abs(n)

    if (nabs > 1)
        while nabs > 0
           if nabs % 2 == 1
              s = s*r
           end
           nabs = div(nabs, 2)
           if nabs > 0
               r = square(r)
           end
        end
    else
        s = r
    end

    if (n < 0)
        s = inv(s)
    end
    return s
end

function Base.:(^)(r::DoubleFloat{T}, n::DoubleFloat{T}) where {T<:IEEEFloat}
   if isinteger(n)
      return r^Int(n)
   else
      return exp(n * log(r))
   end
end

function Base.:(^)(r::Int, n::DoubleFloat{T}) where {T<:IEEEFloat}
   if isinteger(n)
      return r^Int(n)
   else
      return exp(n * log(r))
   end
end

function calc_exp(a::DoubleFloat{T}) where {T<:IEEEFloat}
    is_neg = signbit(HI(a))
    xabs = is_neg ? -a : a
    xintpart = modf(xabs)[2]
    xintpart = xintpart.hi + xintpart.lo
    xint = Int64(xintpart)
    xfrac = xabs - T(xint)

    if 0 < xint <= 64
        zint = exp_int[xint]
    elseif xint === zero(Int64)
        zint = zero(DoubleFloat{T})
    else
        dv, rm = divrem(xint, 64)
        zint = exp_int[64]^dv
        if rm > 0
            zint = zint * exp_int[rm]
        end
    end

    # exp(xfrac)
    if HI(xfrac) < 0.5
        zfrac = exp_zero_half(xfrac)
    elseif HI(xfrac) > 0.5
        zfrac = exp_half_one(xfrac)
    else
        if LO(xfrac) == 0.0
            zfrac = DoubleFloat{T}(1.6487212707001282, -4.731568479435833e-17)
        elseif signbit(LO(xfrac))
            zfrac = exp_zero_half(xfrac)
        else
            zfrac = exp_half_one(xfrac)
        end
    end

    z = HI(zint) == zero(T) ? zfrac : zint * zfrac
    if is_neg
        z = inv(z)
    end

    return z
end

function expm1(a::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(a) && return a
    isinf(a) && return(signbit(a) ? zero(DoubleFloat{T}) : a)
    u = exp(a)
    if (u == one(DoubleFloat{T}))
        x
    elseif (u-1.0 == -one(DoubleFloat{T}))
        -one(DoubleFloat{T})
    else
        a*(u-1.0) / log(u)
    end
end

function exp2(a::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(a) && return a
    isinf(a) && return(signbit(a) ? zero(DoubleFloat{T}) : a)
    return DoubleFloat{T}(2)^a
end

function exp10(a::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(a) && return a
    isinf(a) && return(signbit(a) ? zero(DoubleFloat{T}) : a)
    return DoubleFloat{T}(10)^a
end

#=
# ratio of polys from
# https://github.com/sukop/doubledouble/blob/master/doubledouble.py
function calc_exp_frac(x::DoubleFloat{T}) where {T<:AbstractFloat}
  u = (((((((((((x +
                   156.0)*x + 12012.0)*x +
                   600600.0)*x + 21621600.0)*x +
                   588107520.0)*x + 12350257920.0)*x +
                   201132771840.0)*x + 2514159648000.0)*x +
                   23465490048000.0)*x + 154872234316800.0)*x +
                   647647525324800.0)*x + 1295295050649600.0
   v = (((((((((((x -
                   156.0)*x + 12012.0)*x -
                   600600.0)*x + 21621600.0)*x -
                   588107520.0)*x + 12350257920.0)*x -
                   201132771840.0)*x + 2514159648000.0)*x -
                   23465490048000.0)*x + 154872234316800.0)*x -

function exp2(a::DoubleFloat{T}) where {T<:AbstractFloat}
    isnan(x) && return x
    isinf(x) && return(signbit(x) ? zero(DoubleFloat{T}) : x)
    return 2^a
end
647647525324800.0)*x + 1295295050649600.0
  u = u/v
  return u
end
=#

function log(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && !signbit(x) && return x
    x === zero(DoubleFloat{T}) && return neginf(DoubleFloat{T})
    y = DoubleFloat(log(HI(x)), zero(T))
    z = exp(y)
    adj = (z - x) / (z + x)
    adj = mul_by_two(adj)
    y = y - adj
    return y
end

function log1p(x::DoubleFloat{T})  where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && !signbit(x) && return
    u = 1.0 + x
    if u == one(DoubleFloat{T})
        x
    else
        log(u)*x/(u-1.0)
    end
end

logten(::Type{DoubleFloat{Float64}}) = Double64(2.302585092994046, -2.1707562233822494e-16)
logtwo(::Type{DoubleFloat{Float64}}) = Double64(0.6931471805599453, 2.3190468138462996e-17)
logtwo(::Type{DoubleFloat{Float32}}) = Double32(0.6931472, -1.9046542e-9)
logten(::Type{DoubleFloat{Float32}}) = Double32(2.3025851, -3.1975436e-8)
logtwo(::Type{DoubleFloat{Float16}}) = Double16(0.6934, -0.0002122)
logten(::Type{DoubleFloat{Float16}}) = Double16(2.303, -0.0001493)

function log2(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && !signbit(x) && return x
    log(x) / logtwo(DoubleFloat{T})
end

function log10(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && !signbit(x) && return x
    log(x) / logten(DoubleFloat{T})
end


using QuantumClifford
using Yao
using LinearAlgebra
using Plots

N = 12
M = 1000
sts_clifford = [random_stabilizer(N) for _ = 1:M]
sts_haar = [rand_state(N) for _ = 1:M]
kernel_mat_clifford = ones(M, M)
for i = 1:M, j = i+1:M
    kernel_mat_clifford[i,j] = sts_clifford[i]⋅sts_clifford[j]
    @show i,j
end
for i = 1:M, j = i+1:M
    kernel_mat_clifford[j,i] = kernel_mat_clifford[i,j]
end

kernel_mat_haar = ones(M, M)
for i = 1:M, j = i+1:M
    kernel_mat_haar[i,j] = fidelity(sts_haar[i], sts_haar[j])
    @show i,j
end
for i = 1:M, j = i+1:M
    kernel_mat_haar[j,i] = kernel_mat_haar[i,j]
end

begin
    histogram(eigvals(kernel_mat_haar); label = "haar")
    histogram!(eigvals(kernel_mat_clifford); label = "clifford")
end


using Ripserer
using Test

using Ripserer: is_prime

@testset "is_prime" begin
    @test !is_prime(1)
    @test is_prime(2)
    @test is_prime(3)
    @test !is_prime(4)
    @test is_prime(5)
    @test !is_prime(6)
    @test is_prime(7)
    @test !is_prime(8)
    @test !is_prime(9)
    @test !is_prime(10)
    @test is_prime(11)
    @test !is_prime((1, 2, 3))
    @test !is_prime(:two)
    @test !is_prime(Array{Float64,2})
end

@testset "Mod" begin
    @testset "Arithmetic" begin
        @test Mod{2}(1) + Mod{2}(1) == Mod{2}(0)
        @test Mod{5}(2) - Mod{5}(3) == Mod{5}(4)
        @test Mod{7}(2) / Mod{7}(3) == Mod{7}(3)
        @test Mod{17}(2) / Mod{17}(2) == Mod{17}(1)
        @test Mod{3}(2) * Mod{3}(2) == Mod{3}(1)

        @test Mod{3}(1) + 1 == Mod{3}(2)
        @test Mod{5}(2) * 2 == Mod{5}(4)
        @test Mod{7}(2) - 3 == Mod{7}(6)
        @test Mod{13}(10) / 2 == Mod{13}(5)

        @test sign(Mod{2}(1)) == Mod{2}(1)
        @test sign(Mod{13}(0)) == Mod{13}(0)

        @test zero(Mod{2}) == Mod{2}(0)
        @test zero(Mod{5}(1)) == Mod{5}(0)

        for i in 1:10
            @test inv(Mod{11}(i)) == invmod(i, 11)
        end
    end

    @testset "Errors" begin
        @test_throws DomainError Mod{4}(1)
        @test_throws DomainError Mod{-1}(1)
        @test_throws DomainError Mod{String}(1)
        @test_throws DomainError Mod{-1}(1)

        @test_throws ErrorException Mod{3}(1) < Mod{3}(2)
        @test_throws ErrorException Mod{3}(1) ≥ Mod{3}(2)

        @test_throws StackOverflowError Mod{3}(1) + Mod{2}(1)
    end

    @testset "Printing" begin
        @test sprint(print, Mod{3}(1)) == "1 mod 3"
        @test sprint(print, Mod{1801}(1)) == "1 mod 1801"
    end

    @testset "Type promotion" begin
        for T in (Int8, Int16, Int32, Int64, Int128)
            @test promote_type(T, Mod{2}) ≡ Mod{2}
            @test promote(one(T), Mod{5}(2)) == (Mod{5}(1), Mod{5}(2))

            @test Mod{3}(one(T)) == Mod{3}(1)
            @test Mod{3}(one(unsigned(T))) == Mod{3}(1)
            for op in (+, -, *, /)
                @test op(Mod{3}(1), one(T)) isa Mod{3}
                @test op(Mod{3}(1), one(unsigned(T))) isa Mod{3}
            end
        end
    end
end


using Random
using VegaLite
using StatsFuns
using DataFrames
using Distributions
using TripletEmbeddings
using GreenVideoExperiments

Random.seed!(4)

# n = 100
# Y = 2*rand(1,n) .- 1
Y = GreenVideoExperiments.load("TaskA", "ground_truth", "30Hz").data[begin:30:end]
Y = Matrix(transpose(filter(x -> !isnan(x), Y)))

triplets = Triplets(Y; f = x -> logistic(2x))
loss = STE(σ = 1/sqrt(2))
X = Embedding(size(Y))

@time violations = fit!(loss, triplets, X; verbose=true, max_iterations=1000)
procrustes!(X, Y)

dfX = DataFrame(embedding = "X", time = 1:length(X), value = X[1,:])
dfY = DataFrame(embedding = "Y", time = 1:length(Y), value = Y[1,:])

vcat(dfX, dfY) |> @vlplot(:line, x = :time, y = :value, color=:embedding, width=600, height=400)

# savefig("figures/1D.svg")


@testset "diffusion_matrix.jl" begin
    a = 0.0
    b = 2π
    N = 100

    # Test with ClosedIntervalDomain
    domain = ClosedIntervalDomain(a, b)
    D = diffusion_matrix(domain, N)
    @test D isa SparseMatrixCSC
    @test size(D) == (N + 1, N + 1)
    @test nnz(D) == 3(N + 1) - 2
    @test all(diag(D) .< 0)

    # Test with PeriodicIntervalDomain
    domain = PeriodicIntervalDomain(a, b)
    D = diffusion_matrix(domain, N)
    @test D isa SparseMatrixCSC
    @test size(D) == (N, N)
    @test nnz(D) == 3N
    @test all(diag(D, -1) .> 0)
    @test all(diag(D) .< 0)
    @test all(diag(D, +1) .> 0)
    @test D[1, end] > 0
    @test D[end, 1] > 0

    # Concrete unknown domain type
    domain = UnknownDomain()
    @test_throws Exception diffusion_matrix(domain, N)
end


using ZXCalculus
using ZXCalculus: ancilla_extraction

function gen_phase_gadget()
    zxd = ZXDiagram(2)
    push_gate!(zxd, Val(:Z), 1, 1//2)
    push_gate!(zxd, Val(:CNOT), 1, 2)
    push_gate!(zxd, Val(:Z), 1, 1//4)
    push_gate!(zxd, Val(:CNOT), 1, 2)
    push_gate!(zxd, Val(:Z), 1, 1//2)

    return zxd
end

zxd = gen_phase_gadget()
zxg = full_reduction(zxd)
anc_circ = ancilla_extraction(zxg)

# plot(anc_circ)

zxd_swap = ZXDiagram(2)
pushfirst_gate!(zxd_swap, Val(:SWAP), [1, 2])
# plot(zxd_swap)
convert_to_chain(zxd_swap)

zxg_swap = ZXGraph(zxd_swap)
zxd_anc = ancilla_extraction(zxg_swap)
# plot(zxd_anc)
@test length(ZXCalculus.convert_to_chain(zxd_anc)) == 3


#####
##### An attempt to simulate a scenario similar to that reported by
##### Veron and Melville, JFM (2001)
#####

using Oceananigans
using Oceananigans.Units
using Printf

#####
##### Domain
#####

@info "Building a grid..."

Ny = 128
Ly = 0.3

Nx = Int(Ny/2)
Lx = Ly/2

Nz = 2Ny # number of points in the vertical direction
Lz = Ly # domain depth

refinement = 1.5 # controls spacing near surface (higher means finer spaced)
stretching = 8   # controls rate of stretching at bottom
h(k) = (k - 1) / Nz
ζ₀(k) = 1 + (h(k) - 1) / refinement
Σ(k) = (1 - exp(-stretching * h(k))) / (1 - exp(-stretching))

arch = GPU()

#=
grid = VerticallyStretchedRectilinearGrid(architecture = arch,
                                          size = (Nx, Ny, Nz),
                                          halo = (3, 3, 3),
                                          x = (0, Lx),
                                          y = (0, Ly),
                                          z_faces = k -> Lz * (ζ₀(k) * Σ(k) - 1),
                                          topology = (Periodic, Bounded, Bounded))
=#

grid = RegularRectilinearGrid(size = (Nx, Ny, Nz),
                              halo = (3, 3, 3),
                              x = (0, Lx),  # longer in streamwise direction
                              y = (0, Ly),  # wide enough to avoid finite-width effects?
                              z = (-Lz, 0), # not full depth, but deep enough?
                              topology = (Periodic, Bounded, Bounded))

@show grid

#####
##### Boundary conditions
##### We define two options: one using drag model, one no-slip
#####

@info "Defining boundary conditions..."

#####
##### Stokes drift
#####

# Physical parameters
const g = 9.81 # gravitational acceleration
const T = 7.2e-5 # (kinematic) surface tension, m³ s⁻²

# Surface wave parameters
#const aᵣ = 0.002 # m, reference wave amplitude at t ≈ τᵃ
const aᵣ = 0.000 # m, reference wave amplitude at t ≈ τᵃ
const τᵃ = 20 # s, time-scale over which surface wave amplitude increases
const kᵣ = 2π / 0.02 # m⁻¹, constant wavenumber
const ω′ = sqrt(g * kᵣ) # gravity frequency without surface tension effects
const ωᵣ = sqrt(g * kᵣ + T * kᵣ^3) # gravity wave frequency with surface tension

@show ϵ = aᵣ * kᵣ

@show 1 / kᵣ # surface wave decay scale
@show ωᵣ ω′

#@inline a(t) = aᵣ * tanh(t / τᵃ)
#@inline ∂t_a(t) = aᵣ / τᵃ * sech(t / τᵃ) * sech(t / τᵃ)

@inline a(t) = aᵣ
@inline ∂t_a(t) = 0.0

@inline    uˢ(z, t) = a(t) * a(t) * kᵣ * ωᵣ * exp(2kᵣ * z) # constant k
@inline ∂t_uˢ(z, t) = 2 * ∂t_a(t) * a(t) * kᵣ * ωᵣ * exp(2kᵣ * z)
@inline ∂z_uˢ(z, t) = 2kᵣ * uˢ(z, t)

# ∫ᶻ_∂t_uˢ(t) = 2 * ∂t_a(t) * a(t) * kᵣ * ωᵣ * ∫ᶻ exp(2kᵣ * z) dz = ∂t_a(t) * a(t) * ωᵣ
@inline ∫ᶻ_∂t_uˢ(t) = a(t) * ∂t_a(t) * ωᵣ

# From Veron and Melville 2001: β = α / ρ
#
# α = 10⁻² => β ≈ 10⁻⁵
#
# As a sanity check: estimate τ_dynamic ≈ ρ_air * Cᴰ * u_air²
# Then with ρ_air ≈ 1.2, Cᴰ ≈ 10⁻³, u_air = 10 => τ_dynamic ≈ 0.12

const β = 1e-5 # m² s^(-5/2), from Fabrice Veron (June 30 2021 and May 11 2021)

# u_wind_bc = FluxBoundaryCondition((x, y, t) -> - β * sqrt(t) + ∫ᶻ_∂t_uˢ(t))
u_wind_bc = FluxBoundaryCondition(-1e-4)

u_bcs_free_slip = FieldBoundaryConditions(top = u_wind_bc)
boundary_conditions = (; u = u_bcs_free_slip)

#=
# Here's a few other boundary conditions one might consider:
#   * quadratic drag (for LES)
#   * no-slip (for resolved LES or DNS)
#
# Quadratic drag model for momenum fluxes at solid walls
cᵈ = 2e-3
u_drag(x, y, t, u, v, w, cᵈ) = - cᵈ * u * sqrt(u^2 + v^2 + w^2)
v_drag(x, y, t, u, v, w, cᵈ) = - cᵈ * v * sqrt(u^2 + v^2 + w^2)
w_drag(x, y, t, u, v, w, cᵈ) = - cᵈ * w * sqrt(u^2 + v^2 + w^2)

u_drag_bc = FluxBoundaryCondition(u_drag, field_dependencies=(:u, :v, :w), parameters = cᵈ)
v_drag_bc = FluxBoundaryCondition(v_drag, field_dependencies=(:u, :v, :w), parameters = cᵈ)
w_drag_bc = FluxBoundaryCondition(w_drag, field_dependencies=(:u, :v, :w), parameters = cᵈ)

u_bcs_drag = UVelocityBoundaryConditions(grid, top = u_wind_bc, bottom = u_drag_bc, south = u_drag_bc, north = u_drag_bc)
v_bcs_drag = VVelocityBoundaryConditions(grid, bottom = v_drag_bc)
w_bcs_drag = WVelocityBoundaryConditions(grid, north = w_drag_bc, south = w_drag_bc)
boundary_conditions = (u = u_bcs_drag, v = v_bcs_drag, w = w_bcs_drag)

no_slip_bc = ValueBoundaryCondition(0)
u_bcs_no_slip = UVelocityBoundaryConditions(grid, top = u_wind_bc, bottom = no_slip_bc, south = no_slip_bc, north = no_slip_bc)
v_bcs_no_slip = VVelocityBoundaryConditions(grid, bottom = no_slip_bc)
w_bcs_no_slip = WVelocityBoundaryConditions(grid, north = no_slip_bc, south = no_slip_bc)
boundary_conditions = (u = u_bcs_no_slip, v = v_bcs_no_slip, w = w_bcs_no_slip)
=#

@info "Modeling..."

addzero(args...) = 0

model = NonhydrostaticModel(architecture = arch,
                            advection = WENO5(),
                            timestepper = :RungeKutta3,
                            grid = grid,
                            tracers = :c,
                            boundary_conditions = boundary_conditions,
                            closure = IsotropicDiffusivity(ν=1.05e-6, κ=1.0e-6),
                            #stokes_drift = UniformStokesDrift(∂z_uˢ=∂z_uˢ, ∂t_uˢ=∂t_uˢ),
                            coriolis = nothing,
                            buoyancy = nothing)

#@show Δz = minimum(parent(grid.Δzᵃᵃᶜ)) # for a stretched grid
@show Δz = grid.Δz # for a regular grid
@show uᵢ = 1e-15
#wᵢ(x, y, z) = uᵢ * exp(z / (5 * Δz)) * randn()
wᵢ(x, y, z) = uᵢ * randn()

set!(model, w = wᵢ)

c = model.tracers.c
interior(c)[:, :, grid.Nz] .= 1

@info "Revvving up a simulation..."

simulation = Simulation(model, Δt=0.01, stop_time=1minutes)

wizard = TimeStepWizard(cfl=0.8, max_Δt=1.0)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

function progress(s)

    ν = s.model.closure.ν

    umax = maximum(abs, s.model.velocities.u)
    vmax = maximum(abs, s.model.velocities.v)
    wmax = maximum(abs, s.model.velocities.w)

    t = s.model.clock.time
    h = √(ν * t)
    Re = umax * h / ν

    @info @sprintf("Time: %s, iteration: %d, next Δt: %s, max|U|: (%.2e, %.2e, %.2e)  m s⁻¹, hk: %.2f, Re: %.1f",
                   prettytime(s),
                   iteration(s),
                   prettytime(s.Δt),
                   umax, vmax, wmax,
                   h * kᵣ,
                   Re)

    return nothing
end

simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))

@show simulation

#####
##### Set up output
#####

prefix = @sprintf("veron_and_melville_Nz%d_Ly%.1f_β%.1e_a%.1e", grid.Nz, grid.Ly, β, aᵣ)

outputs = merge(model.velocities, model.tracers)

simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs,
                                                    schedule = TimeInterval(0.1),
                                                    force = true,
                                                    prefix = prefix * "_yz",
                                                    field_slicer = FieldSlicer(i = round(Int, grid.Nx/2)))

simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs,
                                                  schedule = TimeInterval(0.1),
                                                  force = true,
                                                  field_slicer = FieldSlicer(j = round(Int, grid.Ny/2)),
                                                  prefix = prefix * "_xz")

simulation.output_writers[:netcdf_xz] = NetCDFOutputWriter(model, outputs,
                                                           schedule = TimeInterval(0.1),
                                                           mode = "c",
                                                           field_slicer = FieldSlicer(j = round(Int, grid.Ny/2)),
                                                           filepath = prefix * "_xz.nc")

simulation.output_writers[:z] = NetCDFOutputWriter(model, outputs,
                                                   schedule = TimeInterval(0.01),
                                                   mode = "c",
                                                   field_slicer = FieldSlicer(i = round(Int, grid.Nx/2),
                                                                              j = round(Int, grid.Ny/2)),
                                                   filepath = prefix * "_z.nc")

C = AveragedField(model.tracers.c, dims=(1, 2))
U = AveragedField(model.velocities.u, dims=(1, 2))

simulation.output_writers[:averages] = NetCDFOutputWriter(model, (c=C, u=U),
                                                          schedule = TimeInterval(0.1),
                                                          mode = "c",
                                                          filepath = prefix * "_averages.nc")

@info "Running..."

@time run!(simulation)

@info "Simulation complete: $simulation. Output:"

for (name, writer) in simulation.output_writers
    absfilepath = abspath(writer.filepath)
    @info "OutputWriter $name, $absfilepath:\n $writer"
end


#=
  Coins puzzle in Julia + ConstraintSolver.jl.

  Problem from
  Tony Hürlimann: "A coin puzzle - SVOR-contest 2007"
  http://www.svor.ch/competitions/competition2007/AsroContestSolution.pdf
  """
  In a quadratic grid (or a larger chessboard) with 31x31 cells, one
  should place coins in such a way that the following conditions are
  fulfilled:
    1. In each row exactly 14 coins must be placed.
    2. In each column exactly 14 coins must be placed.
    3. The sum of the quadratic horizontal distance from the main
       diagonal of all cells containing a coin must be as small as possible.
    4. In each cell at most one coin can be placed.

   The description says to place 14x31 = 434 coins on the chessboard
   each row containing 14 coins and each column also containing 14 coins.
  """

  Note: This problem is quite hard for CP solvers. A MIP solver solves
  the (14,31) problem in millis.

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

=#

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

function coins_grid(n=31,c=14)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer) # , "logLevel" => 0)
    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> true,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

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

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

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

                                                            "time_limit"=>110,
                                                            # "lp_optimizer" => cbc_optimizer, # 10.7s
                                                            "lp_optimizer" => glpk_optimizer, # 1.1s
                                        ))

    @variable(model, x[1:n,1:n], Bin)
    @variable(model, 0 <= s <= n*n*n, Int)

    @constraint(model,s == sum([(x[i,j] * abs(i-j)*abs(i-j)) for i in 1:n, j in 1:n]))
    for i in 1:n
        @constraint(model, c == sum(x[i,:]))
        @constraint(model, c == sum(x[:,i]))
    end

    @objective(model, Min, s)

    # Solve the problem
    println("solve")
    optimize!(model)

    status = JuMP.termination_status(model)
    println("status:$status")
    if status == MOI.OPTIMAL
        x = convert.(Integer,JuMP.value.(x))
        s = convert.(Integer,JuMP.value.(s))
        for i in 1:n
            println(x[i,:])
        end
        println("s:$s")

        num_sols = MOI.get(model, MOI.ResultCount())
        println("\nnum_sols:$num_sols\n")
    end
end

n = 31
c = 14
coins_grid(n,c)