Divergence.h

Go to the documentation of this file.

#pragma once

 
#include <concepts>
#include <expected>
#include <limits>
#include <random>
#include <type_traits>
#include <utility>
 
#include <Eigen/Cholesky>
#include <Eigen/SVD>
 
#include <EmrysFilter/Core/Execution.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Multiply.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Core/Triangularization.h>
#include <EmrysFilter/Core/Weights.h>
#include <EmrysFilter/Distributions/Categorical.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/GaussianForm.h>
#include <EmrysFilter/Distributions/Hybrid.h>
#include <EmrysFilter/Distributions/Mixture.h>
#include <EmrysFilter/Distributions/Moments.h>
#include <EmrysFilter/Distributions/Product.h>
#include <EmrysFilter/Distributions/Sampling.h>
#include <EmrysFilter/Distributions/SqrtGaussian.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Spaces/TagPath.h>
 
namespace Emrys
{


struct MonteCarlo
{
    Eigen::Index count;  
};

template <typename Scalar = double>
struct MonteCarloEstimate
{
    Scalar value;          
    Scalar standardError;  
    Eigen::Index count;    
};
 
namespace internal
{

template <typename Indices>
[[nodiscard]] Eigen::Index
indexCount(Indices const & indices)
{
    if constexpr (std::is_integral_v<Indices>)
    {
        return 1;
    }
    else
    {
        return Eigen::Index(indices.size());
    }
}

template <typename Indices>
[[nodiscard]] Eigen::Index
indexAt(Indices const & indices, Eigen::Index k)
{
    if constexpr (std::is_integral_v<Indices>)
    {
        return Eigen::Index(indices);
    }
    else
    {
        return Eigen::Index(indices[k]);
    }
}
 
}  // namespace internal

template <Space Space_, RealScalar Scalar_ = double>
class GroundMetric
{
public:
 
    using SpaceType = Space_;                                          
    using Scalar = Scalar_;                                            
    using Index = Eigen::Index;                                        
    using MatrixType = typename Moments<Space_, Scalar_>::MatrixType;  

    explicit GroundMetric(Space_ space)
        : m_matrix{MatrixType::Identity(Emrys::differenceDimension(space), Emrys::differenceDimension(space))}
        , m_space{std::move(space)}
    {}

    template <typename Path>
    void
    scale(Scalar factor)
    {
        auto const rows = diffIdx<Path>(m_space);
        for (Index k = 0; k < internal::indexCount(rows); ++k)
        {
            m_matrix.row(internal::indexAt(rows, k)) *= factor;
        }
    }

    template <typename Into, typename From>
    void
    combine(Scalar factor)
    {
        auto const into = diffIdx<Into>(m_space);
        auto const from = diffIdx<From>(m_space);
        EMRYS_FILTER_ASSERT(
            internal::indexCount(into) == internal::indexCount(from)
            && "combined blocks must have the same difference dimension"
        );
        MatrixType const before = m_matrix;
        for (Index k = 0; k < internal::indexCount(into); ++k)
        {
            m_matrix.row(internal::indexAt(into, k)) += factor * before.row(internal::indexAt(from, k));
        }
    }

    [[nodiscard]] MatrixType const &
    matrix() const
    {
        return m_matrix;
    }

    [[nodiscard]] Space_ const &
    space() const
    {
        return m_space;
    }
 
private:
 
    MatrixType m_matrix;                            
    EMRYS_FILTER_NO_UNIQUE_ADDRESS Space_ m_space;  
};
 
namespace internal
{

template <typename>
inline constexpr bool alwaysFalse = false;

struct NoClosedForm
{
    template <typename T>
    operator T() const;
};

inline void
assertSameDimensions(auto const & p, auto const & q)
{
    EMRYS_FILTER_ASSERT(
        p.valueDimension() == q.valueDimension() && p.differenceDimension() == q.differenceDimension()
        && "the distributions must have the same dimensions"
    );
}

template <typename Scalar>
class RunningMean
{
public:

    void
    add(Scalar x)
    {
        if (x == Eigen::NumTraits<Scalar>::infinity())
        {
            m_infinite = true;
        }
        if (m_infinite)
        {
            return;
        }
        ++m_count;
        Scalar const delta = x - m_mean;
        m_mean += delta / Scalar(m_count);
        m_sumOfSquares += delta * (x - m_mean);
    }

    [[nodiscard]] bool
    infinite() const
    {
        return m_infinite;
    }

    [[nodiscard]] Scalar
    mean() const
    {
        return m_infinite ? Eigen::NumTraits<Scalar>::infinity() : m_mean;
    }

    [[nodiscard]] Scalar
    standardError() const
    {
        return m_infinite ? Eigen::NumTraits<Scalar>::infinity()
                          : Eigen::numext::sqrt(m_sumOfSquares / Scalar(m_count - 1) / Scalar(m_count));
    }
 
private:
 
    Eigen::Index m_count = 0;   
    Scalar m_mean = 0;          
    Scalar m_sumOfSquares = 0;  
    bool m_infinite = false;    
};

template <typename Scalar>
[[nodiscard]] Scalar
infinity()
{
    return Eigen::NumTraits<Scalar>::infinity();
}

template <typename Scalar>
[[nodiscard]] Scalar
logEvenMixture(Scalar a, Scalar b)
{
    auto sum = LogSumExp<Scalar>{};
    sum.add(a);
    sum.add(b);
    return sum.value() - Eigen::numext::log(Scalar(2));
}

template <typename Scalar>
[[nodiscard]] Scalar
hellingerFromBhattacharyya(Scalar distance)
{
    return Eigen::numext::sqrt(Eigen::numext::maxi(-Eigen::numext::expm1(-distance), Scalar(0)));
}

template <typename Scalar, typename DifferenceDerived, typename FirstDerived, typename SecondDerived>
[[nodiscard]] Scalar
buresWasserstein(
    Eigen::MatrixBase<DifferenceDerived> const & difference,
    Eigen::MatrixBase<FirstDerived> const & first,
    Eigen::MatrixBase<SecondDerived> const & second
)
{
    using Cross = MatrixOf<
        Scalar,
        FirstDerived::ColsAtCompileTime,
        SecondDerived::ColsAtCompileTime,
        FirstDerived::MaxColsAtCompileTime,
        SecondDerived::MaxColsAtCompileTime>;
    Cross const cross = multiply(first.transpose(), second);
    Scalar const nuclearNorm = Eigen::JacobiSVD<Cross>{cross}.singularValues().sum();
    Scalar const squared =
        difference.squaredNorm() + first.squaredNorm() + second.squaredNorm() - Scalar(2) * nuclearNorm;
    return Eigen::numext::sqrt(Eigen::numext::maxi(squared, Scalar(0)));
}

template <typename D>
inline constexpr bool isGaussianDensity = false;

template <typename Space, typename Scalar>
inline constexpr bool isGaussianDensity<Gaussian<Space, Scalar>> = true;

template <typename Params>
inline constexpr bool isGaussianDensity<FormGaussian<Params>> = true;

template <typename P, typename Q>
concept GaussianDensitiesOfOneSpace =
    isGaussianDensity<P> && isGaussianDensity<Q> && std::same_as<typename P::SpaceType, typename Q::SpaceType>
    && std::same_as<typename P::Scalar, typename Q::Scalar>;

template <typename D>
inline constexpr bool isCovarianceDensity = std::same_as<D, Gaussian<typename D::SpaceType, typename D::Scalar>>;

template <typename D>
inline constexpr int covarianceRootTriangle = D::ParamsType::IsInfoForm ? Eigen::Upper : Eigen::Lower;

template <typename D>
[[nodiscard]] typename D::MatrixType
covarianceRootOf(D const & d)
{
    using MatrixType = typename D::MatrixType;
    if constexpr (isCovarianceDensity<D>)
    {
        return MatrixType{d.llt().matrixL()};
    }
    else if constexpr (D::ParamsType::IsInfoForm)
    {
        Eigen::Index const n = d.differenceDimension();
        return d.colorNormals(MatrixType::Identity(n, n));
    }
    else
    {
        return d.root();
    }
}

template <typename D, typename Derived>
void
whitenInPlace(D const & d, Eigen::MatrixBase<Derived> & matrix)
{
    if constexpr (isCovarianceDensity<D>)
    {
        solveInPlace(d.llt().matrixL(), matrix);
    }
    else if constexpr (D::ParamsType::IsInfoForm)
    {
        auto const root = covarianceRootOf(d);
        solveInPlace(root.template triangularView<Eigen::Upper>(), matrix);
    }
    else
    {
        solveInPlace(d.root().template triangularView<Eigen::Lower>(), matrix);
    }
}

template <typename P, typename Q>
[[nodiscard]] typename P::MatrixType
averageFactor(P const & p, Q const & q)
{
    using Scalar = typename P::Scalar;
    using MatrixType = typename P::MatrixType;
    constexpr int Rows = MatrixType::RowsAtCompileTime;
    constexpr int MaxRows = MatrixType::MaxRowsAtCompileTime;
    using Stack = MatrixOf<
        Scalar,
        Rows,
        Rows == Eigen::Dynamic ? Eigen::Dynamic : 2 * Rows,
        MaxRows,
        MaxRows == Eigen::Dynamic ? Eigen::Dynamic : 2 * MaxRows>;
    Eigen::Index const n = p.differenceDimension();
    auto stack = Stack{n, 2 * n};
    Scalar const half = Eigen::numext::sqrt(Scalar(0.5));
    stack.leftCols(n) = half * covarianceRootOf(p);
    stack.rightCols(n) = half * covarianceRootOf(q);
    return lowerTriangularRoot(stack);
}

template <typename P, typename G>
concept SampleWithLogDensityFrom = requires(P const & p, G & rng) { drawSampleWithLogDensity(p, rng); };

template <typename P, typename Q, typename G>
concept MonteCarloComparable =
    SampleWithLogDensityFrom<P, G> && HasLogDensity<Q> && std::same_as<typename P::SpaceType, typename Q::SpaceType>
    && std::same_as<typename P::Scalar, typename Q::Scalar>;

template <typename Scalar, typename Policy, typename P, typename G, typename TermOf>
[[nodiscard]] RunningMean<Scalar>
meanOverDraws(Policy && policy, P const & p, Eigen::Index count, G & rng, TermOf const & termOf)
{
    EMRYS_FILTER_ASSERT(count >= 2 && "a Monte Carlo estimate needs at least two draws");
    auto terms = RunningMean<Scalar>{};
    if constexpr (isSequential<Policy>)
    {
        for (Eigen::Index j = 0; j < count && !terms.infinite(); ++j)
        {
            auto const draw = drawSampleWithLogDensity(p, rng);
            terms.add(termOf(draw.sample, draw.logDensity));
        }
    }
    else
    {
        auto samples = SampleSet<typename P::SpaceType, Scalar>{p.valueDimension(), count};
        auto logP = Eigen::Array<Scalar, Eigen::Dynamic, 1>{count};
        for (Eigen::Index j = 0; j < count; ++j)
        {
            auto const draw = drawSampleWithLogDensity(p, rng);
            samples.col(j) = draw.sample;
            logP(j) = draw.logDensity;
        }
        auto values = Eigen::Array<Scalar, Eigen::Dynamic, 1>{count};
        forEachIndex(std::forward<Policy>(policy), count, [&](Eigen::Index j) {
            values(j) = termOf(samples.col(j), logP(j));
        });
        for (Eigen::Index j = 0; j < count && !terms.infinite(); ++j)
        {
            terms.add(values(j));
        }
    }
    return terms;
}

template <typename Scalar, typename Policy, typename P, typename Q, typename G>
[[nodiscard]] RunningMean<Scalar>
bhattacharyyaCoefficientTerms(Policy && policy, MonteCarlo const & method, P const & p, Q const & q, G & rng)
{
    return meanOverDraws<Scalar>(
        std::forward<Policy>(policy),
        p,
        method.count,
        rng,
        [&](auto const & sample, Scalar logP) {
            return Eigen::numext::exp(Scalar(0.5) * (q.logDensity(sample) - logP));
        }
    );
}

template <typename Scalar, typename Policy, typename P, typename Q, typename G>
[[nodiscard]] RunningMean<Scalar>
mixtureDivergenceTerms(Policy && policy, MonteCarlo const & method, P const & p, Q const & q, G & rng)
{
    return meanOverDraws<Scalar>(
        std::forward<Policy>(policy),
        p,
        method.count,
        rng,
        [&](auto const & sample, Scalar logP) {
            return Scalar(logP - logEvenMixture(logP, Scalar(q.logDensity(sample))));
        }
    );
}

template <typename Space, typename MomentsType, typename TruthDerived>
[[nodiscard]] typename MomentsType::Scalar
expectedSquaredErrorOf(Space const & space, MomentsType const & moments, Eigen::MatrixBase<TruthDerived> const & truth)
{
    EMRYS_FILTER_ASSERT(truth.size() == moments.valueDimension() && "the truth must be a value of the space");
    typename MomentsType::DifferenceVectorType const error = Emrys::minus(space, moments.mean(), truth);
    return error.squaredNorm() + moments.covarianceLower().trace();
}

template <typename Space, typename MomentsType, typename TruthDerived, typename WeightsDerived>
[[nodiscard]] typename MomentsType::Scalar
expectedSquaredErrorOf(
    Space const & space,
    MomentsType const & moments,
    Eigen::MatrixBase<TruthDerived> const & truth,
    Eigen::MatrixBase<WeightsDerived> const & weights
)
{
    EMRYS_FILTER_ASSERT(truth.size() == moments.valueDimension() && "the truth must be a value of the space");
    EMRYS_FILTER_ASSERT(
        weights.cols() == moments.differenceDimension() && "the weights must have a column per difference coordinate"
    );
    typename MomentsType::DifferenceVectorType const error = Emrys::minus(space, moments.mean(), truth);
    // tr(W Σ W^T) is the sum of the entries of (W Σ) ∘ W.
    return multiply(weights, error).squaredNorm()
         + multiply(weights, moments.covarianceLower()).cwiseProduct(weights).sum();
}

template <typename Distribution, typename Compute>
[[nodiscard]] auto
fromMoments(Distribution const & distribution, Compute compute)
{
    using Result = std::remove_cvref_t<decltype(distribution.moments())>;
    if constexpr (isExpected<Result>)
    {
        using Value = decltype(compute(std::declval<typename Result::value_type const &>()));
        using Expected = std::expected<Value, typename Result::error_type>;
        Result const moments = distribution.moments();
        if (!moments)
        {
            return Expected{std::unexpect, moments.error()};
        }
        return Expected{compute(*moments)};
    }
    else
    {
        return compute(distribution.moments());
    }
}
 
}  // namespace internal

template <typename P, typename Q>
    requires internal::GaussianDensitiesOfOneSpace<P, Q>
[[nodiscard]] typename P::Scalar
kullbackLeibler(P const & p, Q const & q)
{
    using Scalar = typename P::Scalar;
    internal::assertSameDimensions(p, q);
    typename P::MatrixType whitened = internal::covarianceRootOf(p);
    internal::whitenInPlace(q, whitened);
    Scalar const divergence = Scalar(0.5)
                            * (whitened.squaredNorm() + q.squaredMahalanobisDistance(p.mean())
                               - Scalar(p.differenceDimension()) + q.logDeterminant() - p.logDeterminant());
    return Eigen::numext::maxi(divergence, Scalar(0));
}

template <DiscreteSpace Space, typename Scalar>
[[nodiscard]] Scalar
kullbackLeibler(Categorical<Space, Scalar> const & p, Categorical<Space, Scalar> const & q)
{
    EMRYS_FILTER_ASSERT(p.cardinality() == q.cardinality() && "the distributions must have the same values");
    auto const logP = p.logProbabilities();
    auto const logQ = q.logProbabilities();
    Scalar divergence = 0;
    for (Eigen::Index i = 0; i < p.cardinality(); ++i)
    {
        if (logP(i) == -internal::infinity<Scalar>())
        {
            continue;
        }
        if (logQ(i) == -internal::infinity<Scalar>())
        {
            return internal::infinity<Scalar>();
        }
        divergence += Eigen::numext::exp(logP(i)) * (logP(i) - logQ(i));
    }
    return Eigen::numext::maxi(divergence, Scalar(0));
}

template <typename P, typename Q, std::uniform_random_bit_generator G>
    requires internal::MonteCarloComparable<P, Q, G>
[[nodiscard]] MonteCarloEstimate<typename P::Scalar>
kullbackLeibler(MonteCarlo const & method, P const & p, Q const & q, G & rng)
{
    return kullbackLeibler(execution::sequential, method, p, q, rng);
}

template <ExecutionPolicy Policy, typename P, typename Q, std::uniform_random_bit_generator G>
    requires internal::MonteCarloComparable<P, Q, G>
[[nodiscard]] MonteCarloEstimate<typename P::Scalar>
kullbackLeibler(Policy && policy, MonteCarlo const & method, P const & p, Q const & q, G & rng)
{
    using Scalar = typename P::Scalar;
    auto const terms = internal::meanOverDraws<Scalar>(
        std::forward<Policy>(policy),
        p,
        method.count,
        rng,
        [&](auto const & sample, Scalar logP) {
            Scalar const logQ = q.logDensity(sample);
            return logQ == -internal::infinity<Scalar>() ? internal::infinity<Scalar>() : Scalar(logP - logQ);
        }
    );
    return {terms.mean(), terms.standardError(), method.count};
}

template <typename P, typename Q>
[[nodiscard]] internal::NoClosedForm
kullbackLeibler(P const &, Q const &)
{
    static_assert(
        internal::alwaysFalse<P>,
        "kullbackLeibler(p, q) has a closed form only between two Gaussian densities of one space, in any of the four "
        "forms (call toDensity() on the parameters, and copy() a GaussianMarginal first), or two Categorical "
        "views of one space; "
        "estimate it otherwise with kullbackLeibler(MonteCarlo{count}, p, q, rng)"
    );
    return {};
}

template <typename Scalar = double>
struct CollapseBounds
{
    Scalar lower;  
    Scalar upper;  
};
 
namespace internal
{

template <typename Scalar, typename CovarianceDerived, typename ForEach>
[[nodiscard]] CollapseBounds<Scalar>
collapseBounds(Eigen::MatrixBase<CovarianceDerived> const & covarianceLower, ForEach forEach)
{
    using MatrixType = typename CovarianceDerived::PlainObject;
    auto const llt = Eigen::LLT<MatrixType>{covarianceLower};
    EMRYS_FILTER_ASSERT(llt.info() == Eigen::Success && "the mixture's covariance must be positive definite");
    Scalar componentLogDeterminants = 0;
    Scalar weightEntropy = 0;
    forEach([&](Scalar const logWeight, auto const & gaussian) {
        if (logWeight > -Eigen::NumTraits<Scalar>::infinity())
        {
            Scalar const weight = Eigen::numext::exp(logWeight);
            componentLogDeterminants += weight * FactorizedGaussian::logDeterminant(gaussian);
            weightEntropy -= weight * logWeight;
        }
    });
    Scalar const logDeterminant = Scalar(2) * llt.matrixLLT().diagonal().array().log().sum();
    Scalar const upper = Eigen::numext::maxi(Scalar(0.5) * (logDeterminant - componentLogDeterminants), Scalar(0));
    return CollapseBounds<Scalar>{Eigen::numext::maxi(upper - weightEntropy, Scalar(0)), upper};
}
 
}  // namespace internal

template <typename Space, typename Scalar, typename Allocator>
[[nodiscard]] CollapseBounds<Scalar>
kullbackLeiblerCollapseBounds(GaussianMixture<Space, Scalar, Allocator> const & p)
{
    auto const moments = p.moments();
    return internal::collapseBounds<Scalar>(moments.covarianceLower(), [&](auto const visit) {
        for (auto const & entry : p)
        {
            visit(entry.logWeight - p.logTotalWeight(), entry.component);
        }
    });
}

template <typename Space, typename ModeSpace, typename Scalar, typename Allocator>
[[nodiscard]] CollapseBounds<Scalar>
kullbackLeiblerCollapseBounds(HybridGaussian<Space, ModeSpace, Scalar, Allocator> const & p)
{
    auto const moments = p.moments();
    return internal::collapseBounds<Scalar>(moments.covarianceLower(), [&](auto const visit) {
        auto const logProbabilities = internal::logProbabilitiesOf(p.modes());
        for (Eigen::Index m = 0; m < p.size(); ++m)
        {
            visit(Scalar(logProbabilities(m)), p.conditional(m));
        }
    });
}

template <typename Space, typename ModeSpace, typename Scalar, typename ComponentAllocator, typename Allocator>
[[nodiscard]] CollapseBounds<Scalar>
kullbackLeiblerCollapseBounds(
    Mixture<HybridGaussian<Space, ModeSpace, Scalar, ComponentAllocator>, Allocator> const & p
)
{
    auto const moments = p.moments();
    return internal::collapseBounds<Scalar>(moments.covarianceLower(), [&](auto const visit) {
        for (auto const & entry : p)
        {
            Scalar const logWeight = entry.logWeight - p.logTotalWeight();
            auto const logProbabilities = internal::logProbabilitiesOf(entry.component.modes());
            for (Eigen::Index m = 0; m < entry.component.size(); ++m)
            {
                visit(Scalar(logWeight + logProbabilities(m)), entry.component.conditional(m));
            }
        }
    });
}

template <typename P, typename Q>
    requires internal::GaussianDensitiesOfOneSpace<P, Q>
[[nodiscard]] typename P::Scalar
bhattacharyya(P const & p, Q const & q)
{
    using Scalar = typename P::Scalar;
    using MatrixType = typename P::MatrixType;
    internal::assertSameDimensions(p, q);
    typename P::DifferenceVectorType difference = Emrys::minus(p.space(), p.mean(), q.mean());
    Scalar logDeterminant{};
    if constexpr (internal::isCovarianceDensity<P> && internal::isCovarianceDensity<Q>)
    {
        MatrixType const average = Scalar(0.5) * (p.covarianceLower() + q.covarianceLower());
        auto const llt = Eigen::LLT<MatrixType>{average};
        EMRYS_FILTER_ASSERT(llt.info() == Eigen::Success && "the average covariance must be positive definite");
        llt.matrixL().solveInPlace(difference);
        logDeterminant = Scalar(2) * llt.matrixLLT().diagonal().array().log().sum();
    }
    else
    {
        MatrixType const average = internal::averageFactor(p, q);
        EMRYS_FILTER_ASSERT(
            (average.diagonal().array() != Scalar(0)).all() && "the average covariance must be positive definite"
        );
        average.template triangularView<Eigen::Lower>().solveInPlace(difference);
        logDeterminant = Scalar(2) * average.diagonal().array().abs().log().sum();
    }
    Scalar const distance = Scalar(0.125) * difference.squaredNorm()
                          + Scalar(0.5) * (logDeterminant - Scalar(0.5) * (p.logDeterminant() + q.logDeterminant()));
    return Eigen::numext::maxi(distance, Scalar(0));
}

template <DiscreteSpace Space, typename Scalar>
[[nodiscard]] Scalar
bhattacharyya(Categorical<Space, Scalar> const & p, Categorical<Space, Scalar> const & q)
{
    EMRYS_FILTER_ASSERT(p.cardinality() == q.cardinality() && "the distributions must have the same values");
    auto logCoefficient = LogSumExp<Scalar>{};
    for (Eigen::Index i = 0; i < p.cardinality(); ++i)
    {
        logCoefficient.add(Scalar(0.5) * (p.logProbability(i) + q.logProbability(i)));
    }
    return Eigen::numext::maxi(-logCoefficient.value(), Scalar(0));
}

template <typename P, typename Q, std::uniform_random_bit_generator G>
    requires internal::MonteCarloComparable<P, Q, G>
[[nodiscard]] MonteCarloEstimate<typename P::Scalar>
bhattacharyya(MonteCarlo const & method, P const & p, Q const & q, G & rng)
{
    return bhattacharyya(execution::sequential, method, p, q, rng);
}

template <ExecutionPolicy Policy, typename P, typename Q, std::uniform_random_bit_generator G>
    requires internal::MonteCarloComparable<P, Q, G>
[[nodiscard]] MonteCarloEstimate<typename P::Scalar>
bhattacharyya(Policy && policy, MonteCarlo const & method, P const & p, Q const & q, G & rng)
{
    using Scalar = typename P::Scalar;
    auto const terms = internal::bhattacharyyaCoefficientTerms<Scalar>(std::forward<Policy>(policy), method, p, q, rng);
    return {-Eigen::numext::log(terms.mean()), terms.standardError() / terms.mean(), method.count};
}

template <typename P, typename Q>
[[nodiscard]] internal::NoClosedForm
bhattacharyya(P const &, Q const &)
{
    static_assert(
        internal::alwaysFalse<P>,
        "bhattacharyya(p, q) has a closed form only between two Gaussian densities of one space, in any of the four "
        "forms (call toDensity() on the parameters, and copy() a GaussianMarginal first), or two Categorical "
        "views of one space; "
        "estimate it otherwise with bhattacharyya(MonteCarlo{count}, p, q, rng)"
    );
    return {};
}

template <typename P, typename Q>
    requires internal::GaussianDensitiesOfOneSpace<P, Q>
[[nodiscard]] typename P::Scalar
hellinger(P const & p, Q const & q)
{
    return internal::hellingerFromBhattacharyya(bhattacharyya(p, q));
}

template <DiscreteSpace Space, typename Scalar>
[[nodiscard]] Scalar
hellinger(Categorical<Space, Scalar> const & p, Categorical<Space, Scalar> const & q)
{
    return internal::hellingerFromBhattacharyya(bhattacharyya(p, q));
}

template <typename P, typename Q, std::uniform_random_bit_generator G>
    requires internal::MonteCarloComparable<P, Q, G>
[[nodiscard]] MonteCarloEstimate<typename P::Scalar>
hellinger(MonteCarlo const & method, P const & p, Q const & q, G & rng)
{
    return hellinger(execution::sequential, method, p, q, rng);
}

template <ExecutionPolicy Policy, typename P, typename Q, std::uniform_random_bit_generator G>
    requires internal::MonteCarloComparable<P, Q, G>
[[nodiscard]] MonteCarloEstimate<typename P::Scalar>
hellinger(Policy && policy, MonteCarlo const & method, P const & p, Q const & q, G & rng)
{
    using Scalar = typename P::Scalar;
    auto const terms = internal::bhattacharyyaCoefficientTerms<Scalar>(std::forward<Policy>(policy), method, p, q, rng);
    Scalar const distance = Eigen::numext::sqrt(Eigen::numext::maxi(Scalar(1) - terms.mean(), Scalar(0)));
    // dH/dc = -1 / (2H); at H = 0 the first-order error vanishes, and H is of the order of the
    // square root of the coefficient's error.
    Scalar const standardError = distance > Scalar(0) ? terms.standardError() / (Scalar(2) * distance)
                                                      : Eigen::numext::sqrt(terms.standardError());
    return {distance, standardError, method.count};
}

template <typename P, typename Q>
[[nodiscard]] internal::NoClosedForm
hellinger(P const &, Q const &)
{
    static_assert(
        internal::alwaysFalse<P>,
        "hellinger(p, q) has a closed form only between two Gaussian densities of one space, in any of the four "
        "forms (call toDensity() on the parameters, and copy() a GaussianMarginal first), or two Categorical "
        "views of one space; "
        "estimate it otherwise with hellinger(MonteCarlo{count}, p, q, rng)"
    );
    return {};
}

template <DiscreteSpace Space, typename Scalar>
[[nodiscard]] Scalar
jensenShannon(Categorical<Space, Scalar> const & p, Categorical<Space, Scalar> const & q)
{
    EMRYS_FILTER_ASSERT(p.cardinality() == q.cardinality() && "the distributions must have the same values");
    Scalar divergence = 0;
    for (Eigen::Index i = 0; i < p.cardinality(); ++i)
    {
        Scalar const logP = p.logProbability(i);
        Scalar const logQ = q.logProbability(i);
        Scalar const logMixture = internal::logEvenMixture(logP, logQ);
        if (logP != -internal::infinity<Scalar>())
        {
            divergence += Eigen::numext::exp(logP) * (logP - logMixture);
        }
        if (logQ != -internal::infinity<Scalar>())
        {
            divergence += Eigen::numext::exp(logQ) * (logQ - logMixture);
        }
    }
    Scalar const bound = Eigen::numext::log(Scalar(2));
    return Eigen::numext::mini(Eigen::numext::maxi(Scalar(0.5) * divergence, Scalar(0)), bound);
}

template <typename P, typename Q, std::uniform_random_bit_generator G>
    requires internal::MonteCarloComparable<P, Q, G> && internal::MonteCarloComparable<Q, P, G>
[[nodiscard]] MonteCarloEstimate<typename P::Scalar>
jensenShannon(MonteCarlo const & method, P const & p, Q const & q, G & rng)
{
    return jensenShannon(execution::sequential, method, p, q, rng);
}

template <ExecutionPolicy Policy, typename P, typename Q, std::uniform_random_bit_generator G>
    requires internal::MonteCarloComparable<P, Q, G> && internal::MonteCarloComparable<Q, P, G>
[[nodiscard]] MonteCarloEstimate<typename P::Scalar>
jensenShannon(Policy && policy, MonteCarlo const & method, P const & p, Q const & q, G & rng)
{
    using Scalar = typename P::Scalar;
    auto const fromP = internal::mixtureDivergenceTerms<Scalar>(policy, method, p, q, rng);
    auto const fromQ = internal::mixtureDivergenceTerms<Scalar>(policy, method, q, p, rng);
    Scalar const standardError =
        Scalar(0.5)
        * Eigen::numext::sqrt(
            fromP.standardError() * fromP.standardError() + fromQ.standardError() * fromQ.standardError()
        );
    return {Scalar(0.5) * (fromP.mean() + fromQ.mean()), standardError, method.count};
}

template <typename P, typename Q>
[[nodiscard]] internal::NoClosedForm
jensenShannon(P const &, Q const &)
{
    static_assert(
        internal::alwaysFalse<P>,
        "jensenShannon(p, q) has a closed form only between two Categorical views of one space; between Gaussians "
        "or anything else, estimate it with jensenShannon(MonteCarlo{count}, p, q, rng)"
    );
    return {};
}

template <typename P, typename Q>
    requires internal::GaussianDensitiesOfOneSpace<P, Q>
[[nodiscard]] typename P::Scalar
wasserstein(P const & p, Q const & q)
{
    using MatrixType = typename P::MatrixType;
    internal::assertSameDimensions(p, q);
    typename P::DifferenceVectorType const difference = Emrys::minus(p.space(), p.mean(), q.mean());
    MatrixType const first = internal::covarianceRootOf(p);
    MatrixType const second = internal::covarianceRootOf(q);
    return internal::buresWasserstein<typename P::Scalar>(difference, first, second);
}

template <typename P, typename Q, typename WeightsDerived>
    requires internal::GaussianDensitiesOfOneSpace<P, Q>
[[nodiscard]] typename P::Scalar
wasserstein(P const & p, Q const & q, Eigen::MatrixBase<WeightsDerived> const & weights)
{
    using Scalar = typename P::Scalar;
    using DensityType = P;
    using Measured =
        internal::MatrixOf<Scalar, WeightsDerived::RowsAtCompileTime, 1, WeightsDerived::MaxRowsAtCompileTime, 1>;
    using Factor = internal::MatrixOf<
        Scalar,
        WeightsDerived::RowsAtCompileTime,
        DensityType::DifferenceDimensionAtCompileTime,
        WeightsDerived::MaxRowsAtCompileTime,
        DensityType::MaxDifferenceDimensionAtCompileTime>;
    internal::assertSameDimensions(p, q);
    EMRYS_FILTER_ASSERT(
        weights.cols() == p.differenceDimension() && "the weights must have a column per difference coordinate"
    );
    typename DensityType::DifferenceVectorType const difference = Emrys::minus(p.space(), p.mean(), q.mean());
    Measured const measured = internal::multiply(weights, difference);
    Factor const first = internal::multiply(
        weights,
        internal::asTriangular<internal::covarianceRootTriangle<P>>(internal::covarianceRootOf(p))
    );
    Factor const second = internal::multiply(
        weights,
        internal::asTriangular<internal::covarianceRootTriangle<Q>>(internal::covarianceRootOf(q))
    );
    return internal::buresWasserstein<Scalar>(measured, first, second);
}

template <typename P, typename Q>
    requires internal::GaussianDensitiesOfOneSpace<P, Q>
[[nodiscard]] typename P::Scalar
wasserstein(P const & p, Q const & q, GroundMetric<typename P::SpaceType, typename P::Scalar> const & metric)
{
    return wasserstein(p, q, metric.matrix());
}

template <typename P, typename Q>
[[nodiscard]] internal::NoClosedForm
wasserstein(P const &, Q const &)
{
    static_assert(
        internal::alwaysFalse<P>,
        "wasserstein(p, q) has a closed form only between two Gaussian densities of one space, in any of the "
        "four forms (call toDensity() on the parameters, and copy() a GaussianMarginal first); between "
        "mixtures or particle sets it needs an optimal-transport solver, "
        "which the library does not have. expectedSquaredError(p, truth) gives the distance to a point"
    );
    return {};
}

template <internal::HasMoments Distribution, typename TruthDerived>
[[nodiscard]] auto
expectedSquaredError(Distribution const & distribution, Eigen::MatrixBase<TruthDerived> const & truth)
{
    return internal::fromMoments(distribution, [&](auto const & moments) {
        return internal::expectedSquaredErrorOf(distribution.space(), moments, truth);
    });
}

template <internal::HasMoments Distribution, typename TruthDerived, typename WeightsDerived>
[[nodiscard]] auto
expectedSquaredError(
    Distribution const & distribution,
    Eigen::MatrixBase<TruthDerived> const & truth,
    Eigen::MatrixBase<WeightsDerived> const & weights
)
{
    return internal::fromMoments(distribution, [&](auto const & moments) {
        return internal::expectedSquaredErrorOf(distribution.space(), moments, truth, weights);
    });
}

template <internal::HasMoments Distribution, typename TruthDerived>
[[nodiscard]] auto
expectedSquaredError(
    Distribution const & distribution,
    Eigen::MatrixBase<TruthDerived> const & truth,
    GroundMetric<typename Distribution::SpaceType, typename Distribution::Scalar> const & metric
)
{
    return expectedSquaredError(distribution, truth, metric.matrix());
}
 
}  // namespace Emrys