/**
 * @author Tomas Polasek, David Hrusa
 * @date 1.14.2020
 * @version 1.0
 * @brief gui rendering and helper classes.
 */

#include "TreeStats.h"

#include <queue>

#include <boost/random.hpp>
#include <boost/math/distributions.hpp>

#include "TreeScene.h"
#include "TreeRenderSystemRT.h"

namespace treestat
{

namespace impl
{

/// @brief Internal implementation of the DistributionEngine.
struct DistributionEngineImpl
{
    /// @brief Randomness engine used.
    using EngineT = boost::mt19937;

    /// @brief Type used for variate generators.
    template <typename DistT>
    using VariateGenerator = boost::variate_generator<EngineT, DistT>;

    /// Instance of the engine.
    EngineT engine{ };
}; // struct DistributionEngineImpl

/// @brief Normal Gaussian distribution.
struct NormalDistribution : public DistributionProperties
{
    /// Unique identifier of this distribution.
    static constexpr auto IDENTIFIER{ "Normal" };
    /// @brief Type used for sampling the distribution.
    using DistT = boost::math::normal_distribution<RealType>;

    /// @brief Initialize normal distribution with mean and standard deviation.
    NormalDistribution(DistributionEngineImpl &engine,
        RealType mean = 0, RealType sigma = 1);

    /// @brief Initialize normal distribution from serialized representation.
    NormalDistribution(DistributionEngineImpl &engine,
        const std::string &serialized);

    // Implement interface:
    virtual std::string serialize() const override final;
    virtual void deserialize(const std::string &serialized) override final;
    virtual void setEngine(DistributionEngineImpl &engine) override final;
    virtual RealType cdf(const RealType &x) const override final;
    virtual RealType sample() override final;

    /// @brief Intitialize the internal structure using provided parameters.
    void initialize(RealType mean, RealType sigma);

    /// Distribution implementation.
    DistT distribution{ };
    /// Sample generator for the distribution.
    DistributionEngineImpl::EngineT generator{ };
}; // struct NormalDistribution

NormalDistribution::NormalDistribution(DistributionEngineImpl &engine,
    RealType mean, RealType sigma) :
    generator{ engine.engine }
{ initialize(mean, sigma); }

NormalDistribution::NormalDistribution(DistributionEngineImpl &engine,
    const std::string &serialized) :
    generator{ engine.engine }
{ deserialize(serialized); }

std::string NormalDistribution::serialize() const
{
    return treeio::json{
        { "distribution", IDENTIFIER },
        { "mean", distribution.mean() },
        { "sigma", distribution.standard_deviation() },
    }.dump();
}

void NormalDistribution::deserialize(const std::string &serialized)
{
    const auto data{ treeio::json::parse(serialized) };

    if (data.at("distribution").get<std::string>() != IDENTIFIER)
    { throw StochasticException("Deserializing with unknown distribution identifier!"); }

    initialize(
        data.at("mean").get<RealType>(),
        data.at("sigma").get<RealType>()
    );
}

void NormalDistribution::setEngine(DistributionEngineImpl &engine)
{ generator = engine.engine; }

NormalDistribution::RealType NormalDistribution::cdf(const RealType &x) const
{ return boost::math::cdf<RealType>(distribution, x); }

NormalDistribution::RealType NormalDistribution::sample()
{ return boost::math::quantile<RealType>(distribution, generator()); }

void NormalDistribution::initialize(RealType mean, RealType sigma)
{ distribution = DistT{ mean, sigma }; }

std::shared_ptr<DistributionProperties>
    DistributionProperties::normalDistribution(DistributionEngineImpl &engine,
        RealType mean, RealType sigma)
{ return std::make_shared<NormalDistribution>(engine, mean, sigma); }

std::shared_ptr<DistributionProperties>
    DistributionProperties::deserializeDistribution(
        DistributionEngineImpl &engine, const std::string &serialized)
{
    const auto data{ treeio::json::parse(serialized) };
    const auto name{ data.at("distribution").get<std::string>() };

    if (name == NormalDistribution::IDENTIFIER)
    { return std::make_shared<NormalDistribution>(engine, serialized); }
    else
    { throw StochasticException("Unknown distribution is being deserialized!"); }
}

} // namespace impl

StochasticEngine::StochasticEngine(uint32_t seed)
{ reset(seed); }
StochasticEngine::~StochasticEngine()
{ /* Automatic */ }

void StochasticEngine::reset(uint32_t seed)
{
    if (!mImpl)
    { mImpl = std::make_shared<impl::DistributionEngineImpl>(); }

    if (seed)
    { mImpl->engine.seed(seed); }
    else
    { mImpl->engine.seed(); }
}

std::string ImageData::valueTypeToStr(const ValueType &type)
{
    switch (type)
    {
        default:
        case ValueType::UInt:
        { return "UInt"; }
        case ValueType::Float:
        { return "Float"; }
    }
}

void ImageData::saveTo(treeio::json &out) const
{
    out["image"] = {
        { "width", width },
        { "height", height },
        { "channels", channels },
        { "valueType", valueTypeToStr(valueType) },
        { "data", data.empty() ? std::string{ } : treeutil::encodeBinaryJSON(data) },
    };
}

std::size_t branchingToOrdinal(const Branching &branching)
{
    switch (branching)
    {
        default:
        case Branching::Monopodial:
        { return 0u; }
        case Branching::SympodialMonochasial:
        { return 1u; }
        case Branching::SympodialDichasial:
        { return 2u; }
    }
}

Branching ordinalToBranching(std::size_t ordinal)
{
    switch (ordinal)
    {
        default:
        case 0u:
        { return Branching::Monopodial; }
        case 1u:
        { return Branching::SympodialMonochasial; }
        case 2u:
        { return Branching::SympodialDichasial; }
    }
}

std::size_t ramificationToOrdinal(const Ramification &ramification)
{
    switch (ramification)
    {
        default:
        case Ramification::Continuous:
        { return 0u; }
        case Ramification::Rhythmic:
        { return 1u; }
        case Ramification::Diffuse:
        { return 2u; }
    }
}

Ramification ordinalToRamification(std::size_t ordinal)
{
    switch (ordinal)
    {
        default:
        case 0u:
        { return Ramification::Continuous; }
        case 1u:
        { return Ramification::Rhythmic; }
        case 2u:
        { return Ramification::Diffuse; }
    }
}

void TreeStatValues::saveTo(treeio::json &out) const
{
    forEach([&out] (const auto &var)
    {
        auto &varOut{ out[var.name] };
        varOut["name"] = var.name;
        varOut["description"] = var.description;

        var.var.saveTo(varOut);
        var.hist.saveTo(varOut);
    });

    auto &commonOut{ out["common"] };
    commonOut = {
        { "treeAge", treeAge },
        { "trunkLength", trunkLength },
        { "interNodeLength", interNodeLengthEstimate },
        { "interNodeVariance", interNodeLengthVariance },
        { "leafCount", leafCount },
        { "branching", branchingToOrdinal(branching) },
        { "ramification", ramificationToOrdinal(ramification) },
    };

    auto &visualOut{ out["visual"] };
    shadowImprint.saveTo(visualOut["shadowImprint"]);
}

TreeStats::TreeStats()
{ /* Automatic */ }

TreeStats::TreeStats(const treeio::ArrayTree &tree)
{ calculateStatistics(tree); }

TreeStats::~TreeStats()
{ /* Automatic */ }

void printChains(std::size_t currentChainIdx, const treeutil::TreeChains &chains, std::size_t maxDepth,
    const std::string &indent = "", std::size_t currentDepth = 0u)
{
    const auto &currentChain{ chains.chains()[currentChainIdx] };

    const auto newline{ std::string{ "\n" } + indent + "\t" };
    Info << indent << "[ Chain " << currentChainIdx
        << newline << "depth: " << currentDepth
        << "; nodeCount: " << currentChain.nodes.size()
        << "; gOrder: " << currentChain.graveliusOrder
        << "; gDepth: " << currentChain.graveliusDepth
        << newline << "children (" << currentChain.childChains.size() << "): [ \n";

    if (currentDepth + 1u < maxDepth)
    {
        if (currentChain.childChains.empty())
        { Info << indent << "\tNo Child Chains"; }
        else
        {
            for (const auto &childChainIdx: currentChain.childChains)
            { printChains(childChainIdx, chains, maxDepth, indent + "\t", currentDepth + 1u); Info << "\n"; }
        }
    }
    else
    { Info << indent << "\tMaximum Depth Reached"; }

    Info << newline << " ]"
        << "\n" << indent << " ]";

    if (currentDepth == 0u)
    { Info << std::endl; }
}

void TreeStats::calculateStatistics(const treeio::ArrayTree &tree)
{
    treeutil::Timer profilingTimer{ };

    /* Phase 0 Initialize structures */
    mStats = { };
    if (tree.empty()) { return; }

    /* Phase 1 Prepare helper tree data */
    const auto [chains, chainDataStorage]{ prepareTree(mStats, mSettings, tree) };
    if (chainDataStorage.empty()) { return; }

    const auto timePreparation{ profilingTimer.reset() };

    /* Phase 2 Calculate basic statistics */
    calculateBasicStats(mStats, mSettings, chains, chainDataStorage);

    const auto timeBasicStats{ profilingTimer.reset() };

    /* Phase 3 Calculate visual statistics */
    calculateVisualStats(mStats, mSettings, tree);

    const auto timeVisualStats{ profilingTimer.reset() };

    /* Phase 4 Calculate derived statistics */
    calculateDerivedStats(mStats, mSettings, chains, chainDataStorage);

    const auto timeDerivedStats{ profilingTimer.reset() };

    /* Phase 5 Calculate growth statistics */
    calculateGrowthStats(mStats, mSettings, chains, chainDataStorage);

    const auto timeGrowthStats{ profilingTimer.reset() };

    /* Phase 6 Finalize and clean up */
    finalizeStats(mStats, mSettings);

    const auto timeFinalizeStats{ profilingTimer.reset() };

    const auto timeTotalNoVisual{
        timePreparation + timeBasicStats +
        timeDerivedStats + timeGrowthStats +
        timeFinalizeStats
    };
    const auto timeTotal{ timeTotalNoVisual + timeVisualStats };

    Info << "[Prof] timePreparation: " << timePreparation << std::endl;
    Info << "[Prof] timeBasicStats: " << timeBasicStats << std::endl;
    Info << "[Prof] timeVisualStats: " << timeVisualStats << std::endl;
    Info << "[Prof] timeDerivedStats: " << timeDerivedStats << std::endl;
    Info << "[Prof] timeGrowthStats: " << timeGrowthStats << std::endl;
    Info << "[Prof] timeFinalizeStats: " << timeFinalizeStats << std::endl;
    Info << "[Prof] timeTotalNoVisual: " << timeTotalNoVisual << std::endl;
    Info << "[Prof] timeTotal: " << timeTotal << std::endl;
}

void TreeStats::saveStatisticsToDynamic(treeio::ArrayTree &tree)
{
    auto &dynamic{ tree.metaData().dynamicData() };

    auto &statsOut{ dynamic["stats"] };
    mStats.saveTo(statsOut);
}

std::pair<treeutil::TreeChains, std::vector<TreeStats::ChainData>>
    TreeStats::prepareTree(TreeStatValues &stats, const StatsSettings &settings,
        const treeio::ArrayTree &tree)
{
    treeutil::TreeChains chains{ tree };

    //printChains(0u, chains, 3u);

    if (chains.chains().empty())
    { return { }; }

    std::vector<ChainData> chainDataStorage{ };
    chainDataStorage.resize(chains.chains().size());

    // Calculate forward pass information - root -> leaves.
    prepareTreeForwardPass(chains, chainDataStorage, stats, settings);

    // Calculate backward pass information - leaves -> root.
    prepareTreeBackwardPass(chains, chainDataStorage, stats, settings);

    return { chains, chainDataStorage };
}

void TreeStats::prepareTreeForwardPass(treeutil::TreeChains &chains, std::vector<ChainData> &chainDataStorage,
    TreeStatValues &stats, const StatsSettings &settings)
{
    const auto &chainsTree{ chains.internalTree() };

    for (std::size_t chainIdx = 0u; chainIdx < chains.chains().size(); ++chainIdx)
    { // Go through chains and pre-calculate data.
        const auto &chain{ chains.chains()[chainIdx] };
        auto &chainData{ chainDataStorage[chainIdx] };

        float length{ 0.0f };
        float volume{ 0.0f };
        float angleSum{ 0.0f };
        float angleAbsSum{ 0.0f };
        float lengthAverageThickness{ 0.0f };
        Vector3D startToEnd{ };
        if (chain.nodes.size() > 1u)
        {
            const auto &firstNodeData{ chainsTree.getNode(chain.nodes[0u]).data() };
            const auto &secondNodeData{ chainsTree.getNode(chain.nodes[1u]).data() };
            const auto &lastNodeData{ chainsTree.getNode(chain.nodes.back()).data() };

            chainData.segmentData.reserve(chain.nodes.size() - 1u);
            auto lastNodePos{ firstNodeData.pos };
            auto lastNodeThickness{ firstNodeData.thickness };
            startToEnd = (lastNodeData.pos - lastNodePos);
            auto lastToCurrent{ (secondNodeData.pos - lastNodePos).normalized() };

            chainData.minThickness = std::min<float>(
                chainData.minThickness, firstNodeData.thickness);
            chainData.maxThickness = std::max<float>(
                chainData.maxThickness, firstNodeData.thickness);
            chainData.firstDirection = lastToCurrent;

            for (auto it = chain.nodes.begin() + 1u; it != chain.nodes.end(); ++it)
            {
                const auto currentNodePos{ chainsTree.getNode(*it).data().pos };
                const auto currentDirection{ (currentNodePos - lastNodePos).normalized() };
                const auto lastToCurrentAngle{
                    treeutil::angleBetweenNormVectorsRad<float>(lastToCurrent, currentDirection)
                };

                const auto segmentLength{ lastNodePos.distanceTo(currentNodePos) };
                length += segmentLength;
                angleSum += lastToCurrentAngle;
                angleAbsSum += std::abs(lastToCurrentAngle);

                auto thickness{ chainsTree.getNode(*it).data().thickness };
                if (std::isnan(thickness) || std::isinf(thickness))
                { thickness = 0.0f; }

                volume += treeutil::circularConeFrustumVolume(
                    segmentLength, lastNodeThickness, thickness);

                lengthAverageThickness += segmentLength * thickness;
                stats.segmentThickness.observeSample(thickness);
                chainData.minThickness = std::min<float>(
                    chainData.minThickness, thickness);
                chainData.maxThickness = std::max<float>(
                    chainData.maxThickness, thickness);

                chainData.segmentData.emplace_back(SegmentData{
                    segmentLength, lastNodeThickness, thickness
                });

                lastNodePos = currentNodePos;
                lastNodeThickness = thickness;
                lastToCurrent = currentDirection;

                // Try to find a better first direction, when first few nodes are at the exact same position.
                if (!treeutil::aboveEpsilon(chainData.firstDirection.length()))
                { chainData.firstDirection = lastToCurrent; }
            }

            chainData.lastDirection = lastToCurrent;
        }

        chainData.length = length;
        chainData.angleSum = angleSum;
        chainData.angleAbsSum = angleAbsSum;
        chainData.lengthAverageThickness = lengthAverageThickness;
        chainData.volume = volume;
        chainData.startToEnd = startToEnd;
    }
}

void TreeStats::prepareTreeBackwardPass(treeutil::TreeChains &chains, std::vector<ChainData> &chainDataStorage,
    TreeStatValues &stats, const StatsSettings &settings)
{
    using ChainIdxT = treeutil::TreeChains::NodeChain::ChainIdxT;
    static constexpr auto INVALID_CHAIN_IDX{ treeutil::TreeChains::NodeChain::INVALID_CHAIN_IDX };

    std::vector<std::size_t> finalizedChildChains{ };
    finalizedChildChains.resize(chains.chains().size(), 0u);
    std::queue<ChainIdxT> toProcess{ };
    for (const auto &leafChainIdx : chains.leafChains())
    { toProcess.push(leafChainIdx); }

    while (!toProcess.empty())
    { // Process all chains from leaves to root.
        const auto currentChainIdx{ toProcess.front() }; toProcess.pop();
        auto &currentChain{ chains.chains()[currentChainIdx] };
        auto &currentChainData{ chainDataStorage[currentChainIdx] };

        for (const auto &ccIdx : currentChain.childChains)
        { // Aggregate data from child chains.
            const auto &childChain{ chains.chains()[ccIdx] };
            const auto &childChainData{ chainDataStorage[ccIdx] };

            currentChainData.subtreeLength += childChainData.subtreeLength;
            currentChainData.subtreeVolume += childChainData.subtreeVolume;
            currentChainData.subtreeLeafCount += childChainData.subtreeLeafCount;
        }

        // Set fixed value for leaf chains.
        if (currentChain.childChains.empty())
        {
            currentChainData.subtreeLeafCount = 1u;
            currentChainData.subtreeLength = currentChainData.length;
            currentChainData.subtreeVolume = currentChainData.volume;
        }

        // Check parent node finalization..
        const auto parentChainIdx{ currentChain.parentChain };
        if (parentChainIdx != INVALID_CHAIN_IDX)
        { // We have parent -> Not currently at the root node.
            const auto &parentChain{ chains.chains()[parentChainIdx] };
            finalizedChildChains[parentChainIdx]++;

            // Check if all of the parents children have necessary calculations finished.
            if (finalizedChildChains[parentChainIdx] >= parentChain.childChains.size())
            { // All children are finished -> Add parent for processing.
                toProcess.push(parentChainIdx);
            }
        }
    }
}

void TreeStats::calculateBasicStats(TreeStatValues &stats, const StatsSettings &settings,
    const treeutil::TreeChains &chains, const std::vector<ChainData> &chainDataStorage)
{
    for (std::size_t chainIdx = 0u; chainIdx < chains.chains().size(); ++chainIdx)
    { // Go through the tree using non-branching sequences of segments -> chains.
        const auto &chain{ chains.chains()[chainIdx] };
        const auto &chainData{ chainDataStorage[chainIdx] };

        const auto segmentCount{ static_cast<uint32_t>(chain.nodes.size() - 1u) };
        const auto chainDepth{ static_cast<uint32_t>(chain.chainDepth) };

        stats.segmentsPerChain.observeSample(segmentCount);
        stats.chainsPerDepth.observeSample(chainDepth);
        if (treeutil::aboveEpsilon(chainData.length))
        { stats.chainLength.observeSample(chainData.length); }
        if (treeutil::aboveEpsilon(chainData.startToEnd.length()))
        { stats.chainTotalLength.observeSample(chainData.startToEnd.length()); }
        stats.chainDeformation.observeSample(chainData.angleAbsSum);

        for (const auto &segment : chainData.segmentData)
        { // Approximate volume of each segment with circular cone frustum volume.
            const auto volume{
                treeutil::circularConeFrustumVolume(
                    segment.length, segment.startThickness, segment.endThickness)
            };
            stats.segmentVolume.observeSample({ segment.startThickness, volume });
            stats.segmentVolume.observeSample({ segment.endThickness, volume });
        }

        if (chain.nodes.size() > 2u && treeutil::aboveEpsilon(chainData.startToEnd.length()) &&
            treeutil::aboveEpsilon(chainData.length))
        { stats.chainStraightness.observeSample(chainData.startToEnd.length() / chainData.length); }
        if (treeutil::aboveEpsilon(chainData.startToEnd.length()))
        {
            stats.chainSlope.observeSample(
                treeutil::angleBetweenVectorsRad<float>(
                    chainData.startToEnd, Vector3D{ 0.0f, 1.0f, 0.0f }
                )
            );
        }

        if (chainData.minThickness < std::numeric_limits<float>::max())
        { stats.chainMinThickness.observeSample(chainData.minThickness); }
        if (chainData.maxThickness > std::numeric_limits<float>::min())
        { stats.chainMaxThickness.observeSample(chainData.maxThickness); }
        if (chainData.minThickness < std::numeric_limits<float>::max() &&
            chainData.maxThickness > std::numeric_limits<float>::min())
        { stats.chainMinMaxThicknessRatio.observeSample(chainData.minThickness / chainData.maxThickness); }

        if (chain.parentChain != chain.INVALID_CHAIN_IDX)
        { // Calculate child-parent statistics.
            const auto &parentChain{ chains.chains()[chain.parentChain] };
            const auto &parentChainData{ chainDataStorage[chain.parentChain] };

            const auto parentChildAngle{ treeutil::angleBetweenVectorsRad<float>(
                parentChainData.startToEnd, chainData.startToEnd
            ) };

            if (treeutil::aboveEpsilon(chainData.length) && treeutil::aboveEpsilon(parentChainData.length))
            { stats.chainLengthRatio.observeSample(chainData.length / parentChainData.length); }
            if (treeutil::aboveEpsilon(chainData.angleSum) && treeutil::aboveEpsilon(parentChainData.angleSum))
            { stats.chainAngleSumDelta.observeSample(chainData.angleSum - parentChainData.angleSum); }
            stats.chainParentChildAngle.observeSample(parentChildAngle);

            const auto tropismBranch{ Vector3D::dotProduct(
                (chainData.startToEnd - chainData.firstDirection).normalized(),
                StatsSettings::DEFAULT_HORIZON_UP_DIRECTION.normalized()
            ) };
            stats.chainTropismBranch.observeSample(tropismBranch);

            const auto tropismHorizon{ Vector3D::dotProduct(
                chainData.firstDirection.normalized(),
                StatsSettings::DEFAULT_HORIZON_UP_DIRECTION.normalized()
            ) };
            stats.chainTropismHorizon.observeSample(tropismHorizon);

            const auto tropismParent{ Vector3D::dotProduct(
                    parentChainData.lastDirection.normalized(),
                    chainData.startToEnd.normalized()
            ) };
            stats.chainTropismParent.observeSample(tropismParent);
        }
        if (chain.childChains.size() > 1u)
        { // Calculate child statistics.
            auto minChildAngle{ std::numeric_limits<float>::max() };
            auto maxChildAngle{ std::numeric_limits<float>::lowest() };
            auto minChildAngleTotal{ std::numeric_limits<float>::max() };
            auto maxChildAngleTotal{ std::numeric_limits<float>::lowest() };

            auto minParentAngle{ std::numeric_limits<float>::max() };
            auto maxParentAngle{ std::numeric_limits<float>::lowest() };
            auto minParentAngleTotal{ std::numeric_limits<float>::max() };
            auto maxParentAngleTotal{ std::numeric_limits<float>::lowest() };

            for (std::size_t idx1 = 0u; idx1 < chain.childChains.size(); ++idx1)
            {
                const auto &cChainIdx1{ chain.childChains[idx1] };
                const auto &cChainData1{ chainDataStorage[cChainIdx1] };

                auto minSiblingAngle{ std::numeric_limits<float>::max() };
                auto maxSiblingAngle{ std::numeric_limits<float>::lowest() };
                auto minSiblingAngleTotal{ std::numeric_limits<float>::max() };
                auto maxSiblingAngleTotal{ std::numeric_limits<float>::lowest() };

                const auto parentChainAngle{ treeutil::angleBetweenVectorsRad<float>(
                    chainData.lastDirection, cChainData1.firstDirection)};
                const auto parentChainAngleTotal{ treeutil::angleBetweenVectorsRad<float>(
                    chainData.startToEnd, cChainData1.startToEnd)};

                if (treeutil::aboveEpsilon(cChainData1.firstDirection.length()) &&
                    treeutil::aboveEpsilon(chainData.lastDirection.length()))
                {
                    stats.chainParentAngle.observeSample(parentChainAngle);
                    minParentAngle = std::min<float>(minParentAngle, parentChainAngle);
                    maxParentAngle = std::max<float>(maxChildAngle, parentChainAngle);
                }

                if (treeutil::aboveEpsilon(cChainData1.startToEnd.length()) &&
                    treeutil::aboveEpsilon(chainData.startToEnd.length()))
                {
                    stats.chainParentAngleTotal.observeSample(parentChainAngleTotal);
                    minParentAngleTotal = std::min<float>(minParentAngleTotal, parentChainAngleTotal);
                    maxParentAngleTotal = std::max<float>(maxParentAngleTotal, parentChainAngleTotal);
                }

                for (std::size_t idx2 = 0u; idx2 < chain.childChains.size(); ++idx2)
                {
                    const auto &cChainIdx2{ chain.childChains[idx2] };
                    const auto &cChainData2{ chainDataStorage[cChainIdx2] };

                    const auto childChainAngle{ treeutil::angleBetweenVectorsRad<float>(
                        cChainData1.firstDirection, cChainData2.firstDirection)};
                    const auto childChainAngleTotal{ treeutil::angleBetweenVectorsRad<float>(
                        cChainData1.startToEnd, cChainData2.startToEnd)};

                    if (treeutil::aboveEpsilon(cChainData1.firstDirection.length()) &&
                        treeutil::aboveEpsilon(cChainData2.firstDirection.length()))
                    {
                        stats.chainSiblingAngle.observeSample(childChainAngle);
                        minSiblingAngle = std::min<float>(minSiblingAngle, childChainAngle);
                        maxSiblingAngle = std::max<float>(maxSiblingAngle, childChainAngle);
                    }

                    if (treeutil::aboveEpsilon(cChainData1.startToEnd.length()) &&
                        treeutil::aboveEpsilon(cChainData2.startToEnd.length()))
                    {
                        stats.chainSiblingAngleTotal.observeSample(childChainAngleTotal);
                        minSiblingAngleTotal = std::min<float>(minSiblingAngleTotal, childChainAngleTotal);
                        maxSiblingAngleTotal = std::max<float>(maxSiblingAngleTotal, childChainAngleTotal);
                    }

                    if (treeutil::aboveEpsilon(chainData.subtreeLeafCount))
                    {
                        const auto leafAsymmetry{
                            (cChainData1.subtreeLeafCount - cChainData2.subtreeLeafCount) /
                            static_cast<float>(chainData.subtreeLeafCount)
                        };
                        stats.chainAsymmetryLeaf.observeSample(leafAsymmetry);
                    }

                    if (treeutil::aboveEpsilon(chainData.subtreeLength))
                    {
                        const auto lengthAsymmetry{
                            (cChainData1.subtreeLength - cChainData2.subtreeLength) /
                            static_cast<float>(chainData.subtreeLength)
                        };
                        stats.chainAsymmetryLength.observeSample(lengthAsymmetry);
                    }

                    if (treeutil::aboveEpsilon(chainData.subtreeVolume))
                    {
                        const auto volumeAsymmetry{
                            (cChainData1.subtreeVolume - cChainData2.subtreeVolume) /
                            static_cast<float>(chainData.subtreeVolume)
                        };
                        stats.chainAsymmetryVolume.observeSample(volumeAsymmetry);
                    }
                }
                if (minSiblingAngle < std::numeric_limits<float>::max())
                {
                    stats.chainMinSiblingAngle.observeSample(minSiblingAngle);
                    minChildAngle = std::min<float>(minChildAngle, minSiblingAngle);
                }
                if (maxSiblingAngle > std::numeric_limits<float>::lowest())
                {
                    stats.chainMaxSiblingAngle.observeSample(maxSiblingAngle);
                    maxChildAngle = std::max<float>(maxChildAngle, maxSiblingAngle);
                }
                if (minSiblingAngleTotal < std::numeric_limits<float>::max())
                {
                    stats.chainMinSiblingAngleTotal.observeSample(minSiblingAngleTotal);
                    minChildAngleTotal = std::min<float>(minChildAngleTotal, minSiblingAngleTotal);
                }
                if (maxSiblingAngleTotal > std::numeric_limits<float>::lowest())
                {
                    stats.chainMaxSiblingAngleTotal.observeSample(maxSiblingAngleTotal);
                    maxChildAngleTotal = std::max<float>(maxChildAngleTotal, maxSiblingAngleTotal);
                }
            }

            if (minChildAngle < std::numeric_limits<float>::max())
            { stats.chainMinChildAngle.observeSample(minChildAngle); }
            if (maxChildAngle > std::numeric_limits<float>::lowest())
            { stats.chainMaxChildAngle.observeSample(maxChildAngle); }
            if (minChildAngleTotal < std::numeric_limits<float>::max())
            { stats.chainMinChildAngleTotal.observeSample(minChildAngleTotal); }
            if (maxChildAngleTotal > std::numeric_limits<float>::lowest())
            { stats.chainMaxChildAngleTotal.observeSample(maxChildAngleTotal); }

            if (minParentAngle < std::numeric_limits<float>::max())
            { stats.chainMinParentAngle.observeSample(minParentAngle); }
            if (maxParentAngle > std::numeric_limits<float>::lowest())
            { stats.chainMaxParentAngle.observeSample(maxParentAngle); }
            if (minParentAngleTotal < std::numeric_limits<float>::max())
            { stats.chainMinParentAngleTotal.observeSample(minParentAngleTotal); }
            if (maxParentAngleTotal > std::numeric_limits<float>::lowest())
            { stats.chainMaxParentAngleTotal.observeSample(maxParentAngleTotal); }
        }
    }
}

void TreeStats::calculateGrowthStats(TreeStatValues &stats, const StatsSettings &settings,
    const treeutil::TreeChains &chains, const std::vector<ChainData> &chainDataStorage)
{
    const auto branchingPointDelta{ stats.interNodeLengthEstimate * settings.growthBranchingPointDelta };
    const auto compactChains{ chains.generateCompactChains(branchingPointDelta) };
    const auto &chainsTree{ chains.internalTree() };
    const auto &fullChains{ chains.chains() };

    for (const auto &compactChain : compactChains)
    { // Pre-calculate statistics on the compact chains, skipping leaves.
        if (compactChain.childChains.empty())
        { continue; }

        const auto maxCompactedChainLength{ calculateMaxCompactChainLength(
            fullChains, chainsTree, compactChain) };
        const auto internodesInChain{ maxCompactedChainLength / stats.interNodeLengthEstimate };
        stats.chainInternodeCount.observeSample(internodesInChain);
    }

    // Calculate required statistics from the first run.
    const auto averageInternodeStats{ stats.chainInternodeCount.var.properties() };

    for (const auto &compactChain : compactChains)
    { // Go through all compacted chains, skipping leaves.
        if (compactChain.childChains.empty())
        { continue; }
        if (compactChain.childChains.size() == 1u)
        { Warning << "Found compact chain with one child chain!" << std::endl; }

        std::vector<float> branchContinuationWeights{ };
        std::vector<float> branchSizeWeights{ };

        for (const auto &childChainRec : compactChain.childChains)
        { // Calculate weights for child chains.
            const auto parentChainInfo{ compactChain.compactedChains[childChainRec.originIdx] };
            const auto parentChainIdx{ parentChainInfo.chainIdx };
            const auto &parentChain{ fullChains[parentChainIdx] };
            const auto &parentChainData{ chainDataStorage[parentChainIdx] };

            const auto childChainIdx{ childChainRec.chainIdx };
            const auto &childChain{ fullChains[childChainIdx] };
            const auto &childChainData{ chainDataStorage[childChainIdx] };

            const auto childContinuationWeight{ Vector3D::dotProduct(
                parentChainData.lastDirection, childChainData.firstDirection) };
            branchContinuationWeights.push_back(childContinuationWeight);

            const auto childSizeWeight{ childChainData.subtreeLeafCount };
            branchSizeWeights.push_back(childSizeWeight);
        }

        // Determine branching categories.
        const auto [ mpFitness, mpSameAxis, mpHigherAxis ]{
            calculateMonopodialFitness(branchContinuationWeights, branchSizeWeights) };
        const auto [ smFitness, smSameAxis, smHigherAxis ]{
            calculateSympodialMonochasialFitness(branchContinuationWeights, branchSizeWeights) };
        const auto [ sdFitness, sdSameAxis, sdHigherAxis ]{
            calculateSympodialDichasialFitness(branchContinuationWeights, branchSizeWeights) };

        // Observe branching results.
        stats.chainMonopodialBranchingFitness.observeSample(mpFitness);
        stats.aggregateBranchingFitness.observeSample(
            { branchingToOrdinal(Branching::Monopodial), mpFitness });
        stats.chainSympodialMonochasialBranchingFitness.observeSample(smFitness);
        stats.aggregateBranchingFitness.observeSample(
            { branchingToOrdinal(Branching::SympodialMonochasial), smFitness });
        stats.chainSympodialDichasialBranchingFitness.observeSample(sdFitness);
        stats.aggregateBranchingFitness.observeSample(
            { branchingToOrdinal(Branching::SympodialDichasial), sdFitness });

        // Calculate how many internodes does this chain contain before branching.
        const auto maxCompactedChainLength{ calculateMaxCompactChainLength(
            fullChains, chainsTree, compactChain) };
        const auto internodesInChain{ maxCompactedChainLength / stats.interNodeLengthEstimate };

        // Determine ramification categories.
        const auto continuousFitness{ calculateContinuousFitness(internodesInChain,
            averageInternodeStats.mean, averageInternodeStats.variance,
            averageInternodeStats.min, averageInternodeStats.max) };
        const auto rhythmicFitness{ calculateRhythmicFitness(internodesInChain,
            averageInternodeStats.mean, averageInternodeStats.variance,
            averageInternodeStats.min, averageInternodeStats.max) };
        const auto diffuseFitness{ calculateDiffuseFitness(internodesInChain,
            averageInternodeStats.mean, averageInternodeStats.variance,
            averageInternodeStats.min, averageInternodeStats.max) };

        // Observe ramification results.
        stats.chainContinuousRamificationFitness.observeSample(continuousFitness);
        stats.aggregateRamificationFitness.observeSample(
            { ramificationToOrdinal(Ramification::Continuous), continuousFitness });
        stats.chainRhythmicRamificationFitness.observeSample(rhythmicFitness);
        stats.aggregateRamificationFitness.observeSample(
            { ramificationToOrdinal(Ramification::Rhythmic), rhythmicFitness });
        stats.chainDiffuseRamificationFitness.observeSample(diffuseFitness);
        stats.aggregateRamificationFitness.observeSample(
            { ramificationToOrdinal(Ramification::Diffuse), diffuseFitness });
    }
}

/*(
 * prumer kmene
 * krivost retezu
 * oproti horizontu
 * Minimalni pocet featur ktere koreluji
 * vetvici uhel pro ruzne urovne by mel byt konstantni
 *   vypocitat rozdil oproti prumerne hodnote -> variance
 *   vynechat prvni uroven
 * Vypocitat featury pouze pro uroven 1 a 2
 * Vaha - tluste vetve pod kmenem jsou dulezitejsi.
 *      - vahovat pomoci objemu.
 * Apicalni dominance explicitni feature
 *
 * Pomery bb v shadow obrazu a ve 3D
 *
 * Overleaf:
 *  Tabulka
 *  Delky/tloustky, uhly, krivosti, objemy.
 *  Globalni, lokalni
 */

float TreeStats::calculateMaxCompactChainLength(
    const treeutil::TreeChains::ChainStorage &fullChains,
    const treeutil::TreeChains::InternalArrayTree &chainsTree,
    const treeutil::TreeChains::CompactNodeChain &compactChain)
{
    auto maxCompactedChainLength{ std::numeric_limits<float>::lowest() };

    /// @brief Helper for calculating lengths of compacted chains.
    struct ChainLengthHelper
    {
        /// Next index to be processed.
        std::size_t nextIdx{ 0u };
        /// Currently accumulated length.
        float accumulatedLength{ 0.0f };
    }; // struct ChainLengthHelper

    std::queue<ChainLengthHelper> compactedChainQueue{ };
    compactedChainQueue.push(ChainLengthHelper{ 0u, 0.0f });
    while (!compactedChainQueue.empty())
    {
        const auto chainLengthHelper{ compactedChainQueue.front() }; compactedChainQueue.pop();
        const auto subChainRec{ compactChain.compactedChains[chainLengthHelper.nextIdx] };
        const auto &subChain{ fullChains[subChainRec.chainIdx] };

        const auto chainLength{
            chainLengthHelper.accumulatedLength +
            subChain.calculateChainLength(chainsTree)
        };

        auto foundChild{ false };
        for (std::size_t iii = 0u; iii < compactChain.compactedChains.size(); ++iii)
        {
            const auto &subChainChildRec{ compactChain.compactedChains[iii] };
            if (subChainChildRec.originIdx == chainLengthHelper.nextIdx)
            { foundChild = true; compactedChainQueue.push(ChainLengthHelper{ iii, chainLength }); }
        }

        if (!foundChild)
        { maxCompactedChainLength = std::max(maxCompactedChainLength, chainLength); }
    }

    return maxCompactedChainLength;
}

std::tuple<float, TreeStats::BranchIdxList, TreeStats::BranchIdxList>
TreeStats::calculateMonopodialFitness(
        const std::vector<float> &branchContinuationWeights,
        const std::vector<float> &branchSizeWeights)
{
    /*
     * Monopodial:
     *
     *  \ | /
     *   \|/
     *    |
     *
     *  * One primary branch with high size/continuation - same axis.
     *  * Two or more child branches with the same size weights - higher axis.
     */

    BranchIdxList sameAxis{ };
    BranchIdxList higherAxis{ };

    const auto [minPrimary, argMinPrimary, maxPrimary, argMaxPrimary]{
        treeutil::argMinMax(branchSizeWeights) };
    const auto [minContinuation, argMinContinuation, maxContinuation, argMaxContinuation]{
        treeutil::argMinMax(branchContinuationWeights) };
    const auto allSame{ !treeutil::aboveEpsilon(maxPrimary - minPrimary) };
    const auto primaryBranchIdx{ allSame ? argMaxContinuation : argMaxPrimary };
    const auto maxDifference{
        allSame ?
        1.0f :
        (maxPrimary - minPrimary)
    };

    const auto secondaryBranchCount{ branchSizeWeights.size() - 1u };
    auto meanSecondarySizeWeight{ 0.0f };
    auto secondMaxPrimary{ minPrimary };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx) { continue; }
        meanSecondarySizeWeight += branchSizeWeights[iii];
        secondMaxPrimary = std::max(secondMaxPrimary, branchSizeWeights[iii]);
    }
    meanSecondarySizeWeight /= static_cast<float>(secondaryBranchCount);

    const auto primaryBranchFitness{
        (allSame ? 1.0f : ((branchSizeWeights[primaryBranchIdx] - secondMaxPrimary) / maxDifference)) *
        ((branchContinuationWeights[primaryBranchIdx] + 1.0f) / 2.0f)
    };
    sameAxis.push_back(primaryBranchIdx);

    auto secondaryBranchFitness{ 0.0f };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx) { continue; }
        const auto branchFitness{ branchSizeWeights[iii] - meanSecondarySizeWeight };
        secondaryBranchFitness += 1.0f - (branchFitness * branchFitness) / (maxDifference * maxDifference);
        higherAxis.push_back(iii);
    }
    secondaryBranchFitness /= static_cast<float>(std::max<std::size_t>(2u, secondaryBranchCount));

    const auto monopodialFitness{ primaryBranchFitness * secondaryBranchFitness };

    return { monopodialFitness, sameAxis, higherAxis };
}

std::tuple<float, TreeStats::BranchIdxList, TreeStats::BranchIdxList>
TreeStats::calculateSympodialMonochasialFitness(
    const std::vector<float> &branchContinuationWeights,
    const std::vector<float> &branchSizeWeights)
{
    /*
     * Sympodial Monochasial:
     *
     *  \ |/
     *   \|
     *    |
     *
     *  * One primary branch - same axis.
     *  * One secondary branch - higher axis.
     */

    BranchIdxList sameAxis{ };
    BranchIdxList higherAxis{ };

    const auto [minPrimary, argMinPrimary, maxPrimary, argMaxPrimary]{
        treeutil::argMinMax(branchSizeWeights) };
    const auto [minContinuation, argMinContinuation, maxContinuation, argMaxContinuation]{
        treeutil::argMinMax(branchContinuationWeights) };
    const auto allSame{ !treeutil::aboveEpsilon(maxPrimary - minPrimary) };
    const auto primaryBranchIdx{ allSame ? argMaxContinuation : argMaxPrimary };
    const auto maxDifference{
        allSame ?
        1.0f :
        (maxPrimary - minPrimary)
    };

    const auto secondaryBranchCount{ branchSizeWeights.size() - 1u };
    auto secondMaxPrimary{ minPrimary };
    std::size_t argSecondMaxPrimary{ 1u };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx) { continue; }
        const auto sizeValue{ branchSizeWeights[iii] };
        if (secondMaxPrimary >= sizeValue)
        { secondMaxPrimary = sizeValue; argSecondMaxPrimary = iii; }
    }
    const auto secondaryBranchIdx{ argSecondMaxPrimary };

    const auto primaryBranchFitness{
        (allSame ? 1.0f : ((branchSizeWeights[primaryBranchIdx] - secondMaxPrimary) / maxDifference)) *
        ((branchContinuationWeights[primaryBranchIdx] + 1.0f) / 2.0f)
    };
    sameAxis.push_back(primaryBranchIdx);
    higherAxis.push_back(secondaryBranchIdx);

    auto secondaryBranchFitnessPenalization{ 0.0f };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx || iii == secondaryBranchIdx) { continue; }
        const auto branchFitness{ branchSizeWeights[iii] };
        secondaryBranchFitnessPenalization += branchFitness / maxDifference;
        higherAxis.push_back(iii);
    }
    const auto secondaryBranchFitness{
        secondaryBranchCount >= 2u ?
        1.0f - secondaryBranchFitnessPenalization / static_cast<float>(secondaryBranchCount - 1u) :
        1.0f
    };

    const auto sympodialMonochasialFitness{ primaryBranchFitness * secondaryBranchFitness };

    return { sympodialMonochasialFitness, sameAxis, higherAxis };
}

std::tuple<float, TreeStats::BranchIdxList, TreeStats::BranchIdxList>
TreeStats::calculateSympodialDichasialFitness(
    const std::vector<float> &branchContinuationWeights,
    const std::vector<float> &branchSizeWeights)
{
    /*
     * Sympodial Dichasial:
     *
     *  \   /
     *   \|/
     *    |
     *
     *  * Two secondary branches with same size weight, low continuation - higher axis.
     *  * Optional one lower size weight branch - same axis.
     */

    BranchIdxList sameAxis{ };
    BranchIdxList higherAxis{ };

    const auto [minPrimary, argMinPrimary, maxPrimary, argMaxPrimary]{
        treeutil::argMinMax(branchSizeWeights) };
    const auto allSame{ !treeutil::aboveEpsilon(maxPrimary - minPrimary) };
    const auto primaryBranchIdx{ argMaxPrimary };
    const auto maxDifference{
        allSame ?
        1.0f :
        (maxPrimary - minPrimary)
    };

    const auto secondaryBranchCount{ branchSizeWeights.size() - 1u };
    auto secondMaxPrimary{ minPrimary };
    std::size_t argSecondMaxPrimary{ 1u };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx) { continue; }
        const auto sizeValue{ branchSizeWeights[iii] };
        if (secondMaxPrimary >= sizeValue)
        { secondMaxPrimary = sizeValue; argSecondMaxPrimary = iii; }
    }
    const auto secondaryBranchIdx{ argSecondMaxPrimary };

    const auto primaryBranchFitness{
        (secondMaxPrimary / branchSizeWeights[primaryBranchIdx]) *
        (1.0f - (std::abs(
            branchContinuationWeights[primaryBranchIdx] -
            branchContinuationWeights[secondaryBranchIdx]
        ) / 2.0f))
    };
    higherAxis.push_back(primaryBranchIdx);
    higherAxis.push_back(secondaryBranchIdx);

    auto secondaryBranchFitnessPenalization{ 0.0f };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx || iii == secondaryBranchIdx) { continue; }
        const auto branchFitness{ branchSizeWeights[iii] };
        secondaryBranchFitnessPenalization += (branchFitness / maxDifference) *
                                              ((branchContinuationWeights[iii] + 1.0f) / 2.0f);
        sameAxis.push_back(iii);
    }
    const auto secondaryBranchFitness{
        secondaryBranchCount >= 2u ?
        1.0f - secondaryBranchFitnessPenalization / static_cast<float>(secondaryBranchCount - 1u) :
        1.0f
    };

    const auto sympodialDichasialFitness{ primaryBranchFitness * secondaryBranchFitness };

    return { sympodialDichasialFitness, sameAxis, higherAxis };
}

ImageData copyModality(const treert::RayTracer::OutputModality<float> &modality)
{
    ImageData result{ };

    result.width = modality.width;
    result.height = modality.height;
    result.channels = 1u;
    result.valueType = ImageData::ValueType::Float;

    const auto dataSize{ modality.dataBuffer.size() * sizeof(float) };
    result.data.resize(dataSize);
    std::memcpy(result.data.data(), modality.dataBuffer.data(), dataSize);

    return result;
}

ImageData copyModality(const treert::RayTracer::OutputModality<Vector3D> &modality)
{
    ImageData result{ };

    result.width = modality.width;
    result.height = modality.height;
    result.channels = 3u;
    result.valueType = ImageData::ValueType::Float;

    const auto dataSize{ modality.dataBuffer.size() * sizeof(Vector3D) };
    result.data.resize(dataSize);
    std::memcpy(result.data.data(), modality.dataBuffer.data(), dataSize);

    return result;
}

ImageData copyModality(const treert::RayTracer::OutputModality<uint32_t> &modality)
{
    ImageData result{ };

    result.width = modality.width;
    result.height = modality.height;
    result.channels = 1u;
    result.valueType = ImageData::ValueType::UInt;

    const auto dataSize{ modality.dataBuffer.size() * sizeof(uint32_t) };
    result.data.resize(dataSize);
    std::memcpy(result.data.data(), modality.dataBuffer.data(), dataSize);

    return result;
}

template <typename ModalityT>
ImageData createEmpty(std::size_t w, std::size_t h)
{ ModalityT modality(w, h, 1); return copyModality(modality); }

void TreeStats::calculateVisualStats(TreeStatValues &stats,
    const StatsSettings &settings, const treeio::ArrayTree &inputTree)
{
    if (!settings.visualEnabled)
    {
        stats.shadowImprint = createEmpty<treert::RayTracer::VolumeModality>(1u, 1u);
        return;
    }

    // Prepare ray-tracer with tree reconstruction placed in the scene.
    const auto rayTracerPtr{
        treerndr::RenderSystemRT::prepareTreeRayTracer(
            inputTree, settings.visualTreeScale)
    };
    auto &rayTracer{ *rayTracerPtr };

    // Setup ray-tracer.
    rayTracer.setVerbose(settings.visualVerbose);
    rayTracer.setSampling(settings.visualSampleCount);

    // Prepare modalities.
    rayTracer.traceModalities({ treert::TracingModality::Volume });
    treert::RayTracer::VolumeModality tracedAccumulator{ };

    // Prepare progress printing for long jobs.
    treeutil::ProgressBar progressBar{ "Views " };
    static constexpr auto PROGRESS_PRINT_STEP{ 10u };
    treeutil::ProgressPrinter progressPrinter{ progressBar,
        std::max<std::size_t>(PROGRESS_PRINT_STEP, settings.visualViewCount),
        PROGRESS_PRINT_STEP
    };

    // Ray-trace requested views.
    const auto basePos{ settings.visualTreeScale / 2.0f };
    for (std::size_t viewIdx = 0u; viewIdx < settings.visualViewCount; ++viewIdx)
    { // Ray-trace each view independently.
        const auto t{ viewIdx / static_cast<float>(settings.visualViewCount) };

        // Position the camera.
        const auto cameraPosition{
            Vector3D{
                std::sin(t * 2.0f * treeutil::PI<float>) * basePos,
                basePos,
                std::cos(t * 2.0f * treeutil::PI<float>) * basePos
            }
        };
        const auto cameraFocus{ Vector3D{ 0.0f, basePos, 0.0f } };

        // Prepare ray-tracing context for current view.
        const auto ctx{
            rayTracer.generateOrthoContext(
                settings.visualViewWidth, settings.visualViewHeight,
                cameraPosition, cameraFocus
            )
        };

        // Perform ray tracing.
        rayTracer.traceRays(ctx);

        // Accumulate results.
        const auto &tracedData{ rayTracer.volumeModality() };
        if (tracedAccumulator.dataBuffer.empty())
        { tracedAccumulator = tracedData; }
        else
        {
            for (std::size_t jjj = 0u; jjj < tracedAccumulator.dataBuffer.size(); ++jjj)
            { tracedAccumulator.dataBuffer[jjj] += tracedData.dataBuffer[jjj]; }
        }

        // Report on progress.
        if (settings.visualViewCount >= PROGRESS_PRINT_STEP)
        { progressPrinter.printProgress(Info, viewIdx + 1u); }
    }

    // Finalize processing of modalities.
    for (auto &volume : tracedAccumulator.dataBuffer)
    { volume /= static_cast<float>(settings.visualViewCount); }

    // Save results into buffers.
    stats.shadowImprint = copyModality(tracedAccumulator);

    if (settings.visualExportResults)
    { // Export the results.
        tracedAccumulator.exportToFile(settings.visualExportPath + "/shadowImprint.png");
    }
}

void TreeStats::calculateDerivedStats(TreeStatValues &stats, const StatsSettings &settings,
    const treeutil::TreeChains &chains, const std::vector<ChainData> &chainDataStorage)
{
    const auto &chainLengths{ stats.chainLength.var.observations() };
    // Get trunk chain length, which is always the first.
    stats.trunkLength = chainLengths.empty() ? 0.0f : chainLengths.front();

    // Estimate inter-node length, without including the trunk.
    const auto [lengthMean, lengthVariance]{
        chainLengths.empty() ?
            std::make_pair(0.0f, 0.0f) :
            utils::sampleMeanVariance<float>(
                chainLengths.begin() + 1u,
                chainLengths.end()
            )
    };
    stats.interNodeLengthEstimate = lengthMean;
    stats.interNodeLengthVariance = lengthVariance;

    // Find the deepest leaf chain.
    std::size_t deepestLeafChain{ 0u };
    for (const auto &leafChainIdx : chains.leafChains())
    { deepestLeafChain = std::max(deepestLeafChain, chains.chains()[leafChainIdx].chainDepth); }

    // Estimate crown age by using the deepest leaf chain, without the trunk.
    const auto crownAge{ deepestLeafChain - 1u };
    // Estimate trunk age by dividing it into corresponding inter-node sizes.
    const auto trunkAge{ static_cast<std::size_t>(std::ceil(stats.trunkLength / stats.interNodeLengthEstimate)) };

    // Total age is the age of the crown + age of the trunk...
    stats.treeAge = crownAge + trunkAge;

    // Determine number of leaves.
    stats.leafCount = chains.leafChains().size();
}

float TreeStats::calculateContinuousFitness(float internodeCount, float mean, float var, float min, float max)
{
    /*
     * Continuous ramification is defined by regular branching, specifically
     * 1 internode per branch.
     */

    // Calculate fitness as distance from optimal value - 1 internode per branching.
    const auto diff{ std::abs(1.0f - internodeCount) };
    const auto fitness{ treeutil::smoothstep<float>(1.0f - diff) };

    return fitness;
}

float TreeStats::calculateRhythmicFitness(float internodeCount, float mean, float var, float min, float max)
{
    /*
     * Rhythmic ramification is defined by regular branching, which is other
     * than 1 internode per branch.
     */

    // Calculate regularity as distance from optimal value - average internodes per branching.
    const auto regularityDiff{ std::abs(mean - internodeCount) };
    const auto regularityFitness{ treeutil::smoothstep<float>(1.0f - regularityDiff) };

    // Calculate rhythmic nature as distance from 1 internode per branch.
    const auto rhythmicDiff{ std::abs(1.0f - internodeCount) };
    const auto rhythmicFitness{ treeutil::smoothstep<float>(rhythmicDiff) };

    // Resulting fitness is combination of the two attributes.
    const auto fitness{ regularityFitness * rhythmicFitness };
    return fitness;
}

float TreeStats::calculateDiffuseFitness(float internodeCount, float mean, float var, float min, float max)
{
    /*
     * Diffuse ramification is defined by irregular branching.
     */

    // Calculate irregularity as distance from optimal value - average internodes per branching.
    const auto diff{ std::abs(mean - internodeCount) };
    const auto fitness{ treeutil::smoothstep<float>(diff) };

    return fitness;
}

void TreeStats::finalizeStats(TreeStatValues &stats, const StatsSettings &settings)
{
    // Calculate the rest of the histograms.
    stats.forEach([] (auto &var)
    { var.var.properties(); if (!var.histPrepared) { var.calculateHistogram(MIN_BUCKETS); } });

    // Calculate properties and clean up unnecessary data:
    stats.forEach([] (auto &var)
    // TODO - Clear observations after saving?
    { /* var.var.clearObservations(true); */ });

    // Determine aggregate stats.
    stats.branching = ordinalToBranching(stats.aggregateBranchingFitness.hist.maxBucket().first);
    stats.ramification = ordinalToRamification(stats.aggregateRamificationFitness.hist.maxBucket().first);
}

TreeStats::StatsSettings &TreeStats::settings()
{ return mSettings; }
const TreeStats::StatsSettings &TreeStats::settings() const
{ return mSettings; }

} // namespace treestat


#ifndef ZSVM_SPHERICAL_ECG_JACOBI_VARIATIONAL_OPTIMIZER_HPP_INCLUDED
#define ZSVM_SPHERICAL_ECG_JACOBI_VARIATIONAL_OPTIMIZER_HPP_INCLUDED

// C++ standard library headers
#include <chrono>
#include <cstddef> // for std::size_t
#include <random> // for std::random_device
#include <string>
#include <vector>

// OpenMP multithreading headers
#include <omp.h>

// Boost library headers
#include <boost/multiprecision/cpp_int.hpp>
#include <boost/random.hpp> // for boost::random::independent_bits_engine etc.

// Project-specific headers
#include "Restrict.hpp"
#include "AmoebaOptimizer.hpp"
#include "RealVariationalSolver.hpp"
#include "SphericalECGJacobiContext.hpp"

namespace zsvm {

    template <typename T>
    class SphericalECGJacobiVariationalOptimizer {

        typedef boost::random::mt19937 random_generator_t;

        typedef std::array<
                random_generator_t::result_type,
                random_generator_t::state_size> random_state_t;

        typedef boost::random::independent_bits_engine<
                random_generator_t,
                std::numeric_limits<T>::digits,
                boost::multiprecision::cpp_int> random_engine_t;

    private: // ================================================================

        const std::size_t num_particles;
        const std::size_t num_parameters;
        SphericalECGJacobiContext<T> context;
        RealVariationalSolver<T> solver;
        std::vector<std::vector<T>> basis;
        std::vector<std::vector<T>> basis_matrices;
        random_engine_t random_engine;
        boost::random::normal_distribution<T> correlation_distribution;

        static random_engine_t properly_seeded_random_engine() {
            random_state_t seed_data;
            std::random_device nondet_random_source;
            std::generate(seed_data.begin(), seed_data.end(),
                          std::ref(nondet_random_source));
            seed_data[0] = static_cast<random_generator_t::result_type>(
                    std::chrono::duration_cast<std::chrono::nanoseconds>(
                            std::chrono::system_clock::now().time_since_epoch()
                    ).count() % std::numeric_limits<
                            random_generator_t::result_type>::max());
            boost::random::seed_seq seed_sequence(seed_data.begin(),
                                                  seed_data.end());
            random_engine_t random_engine(seed_sequence);
            return random_engine;
        }

    public: // =================================================================

        explicit SphericalECGJacobiVariationalOptimizer(
                long long int space_dimension,
                const std::vector<zsvm::Particle<T>> &particles,
                const std::string &mass_carrier,
                const std::string &charge_carrier)
                : num_particles(particles.size()),
                  num_parameters(num_particles * (num_particles - 1) / 2),
                  context(SphericalECGJacobiContext<T>::create(
                          particles, mass_carrier, charge_carrier,
                          space_dimension)),
                  random_engine(properly_seeded_random_engine()),
                  correlation_distribution(0, 3) {}

        T get_ground_state_energy() {
            return solver.get_eigenvalue(0);
        }

    public: // =================================================================

        const std::vector<std::vector<T>> &get_basis() {
            return basis;
        }

        void set_basis(const std::vector<std::vector<T>> &input_basis) {
            basis = input_basis;
            recompute_basis_matrices();
            recompute_solver_matrices();
        }

        T augmented_ground_state_energy(
                const T *RESTRICT new_basis_matrix) {
            if (solver.empty()) {
                T overlap_element, hamiltonian_element;
                context.evaluate_matrix_elements(
                        overlap_element, hamiltonian_element,
                        new_basis_matrix, new_basis_matrix);
                return hamiltonian_element / overlap_element;
            }
            const std::size_t basis_size = basis_matrices.size();
            std::vector<T> new_overlap_column(basis_size + 1);
            std::vector<T> new_hamiltonian_column(basis_size + 1);
            for (std::size_t i = 0; i < basis_size; ++i) {
                context.evaluate_matrix_elements(
                        new_overlap_column[i], new_hamiltonian_column[i],
                        basis_matrices[i].data(), new_basis_matrix);
            }
            context.evaluate_matrix_elements(
                    new_overlap_column[basis_size],
                    new_hamiltonian_column[basis_size],
                    new_basis_matrix, new_basis_matrix);
            return solver.minimum_augmented_eigenvalue(
                    new_overlap_column.data(),
                    new_hamiltonian_column.data());
        }

        void random_basis_element(T *RESTRICT basis_element) {
            for (std::size_t i = 0; i < num_parameters; ++i) {
                basis_element[i] = correlation_distribution(random_engine);
            }
        }

        void construct_basis_matrix(
                T *RESTRICT basis_matrix,
                const T *RESTRICT basis_element) {
            context.gaussian_parameter_matrix(basis_element, basis_matrix);
        }

        void add_basis_element(const std::vector<T> &basis_element) {
            basis.push_back(basis_element);
            std::vector<T> basis_matrix(num_parameters);
            construct_basis_matrix(basis_matrix.data(), basis_element.data());
            basis_matrices.push_back(basis_matrix);
            const std::size_t basis_size = basis.size();
            solver.set_basis_size_conservative(basis_size);
            for (std::size_t i = 0; i < basis_size - 1; ++i) {
                context.evaluate_matrix_elements(
                        solver.overlap_matrix_element(i, basis_size - 1),
                        solver.hamiltonian_matrix_element(i, basis_size - 1),
                        basis_matrices[i].data(),
                        basis_matrices[basis_size - 1].data());
                solver.overlap_matrix_element(basis_size - 1, i) =
                        solver.overlap_matrix_element(i, basis_size - 1);
                solver.hamiltonian_matrix_element(basis_size - 1, i) =
                        solver.hamiltonian_matrix_element(i, basis_size - 1);
            }
            context.evaluate_matrix_elements(
                    solver.overlap_matrix_element(
                            basis_size - 1, basis_size - 1),
                    solver.hamiltonian_matrix_element(
                            basis_size - 1, basis_size - 1),
                    basis_matrices[basis_size - 1].data(),
                    basis_matrices[basis_size - 1].data());
        }

        void replace_basis_element(const std::vector<T> &basis_element) {
            const std::size_t basis_size = basis.size();
            basis[basis_size - 1] = basis_element;
            construct_basis_matrix(basis_matrices[basis_size - 1].data(),
                                   basis_element.data());
            for (std::size_t i = 0; i < basis_size - 1; ++i) {
                context.evaluate_matrix_elements(
                        solver.overlap_matrix_element(i, basis_size - 1),
                        solver.hamiltonian_matrix_element(i, basis_size - 1),
                        basis_matrices[i].data(),
                        basis_matrices[basis_size - 1].data());
                solver.overlap_matrix_element(basis_size - 1, i) =
                        solver.overlap_matrix_element(i, basis_size - 1);
                solver.hamiltonian_matrix_element(basis_size - 1, i) =
                        solver.hamiltonian_matrix_element(i, basis_size - 1);
            }
            context.evaluate_matrix_elements(
                    solver.overlap_matrix_element(
                            basis_size - 1, basis_size - 1),
                    solver.hamiltonian_matrix_element(
                            basis_size - 1, basis_size - 1),
                    basis_matrices[basis_size - 1].data(),
                    basis_matrices[basis_size - 1].data());
        }

        std::vector<T> last_basis_element() {
            return basis.back();
        }

        bool expand_amoeba(std::size_t num_trials,
                           const T &initial_step_size,
                           std::size_t max_steps) {
            std::vector<T> new_basis_element(num_parameters);
            std::vector<T> best_basis_element;
            T best_energy = solver.empty()
                            ? std::numeric_limits<T>::max()
                            : solver.get_eigenvalue(0);
            for (std::size_t trial = 0; trial < num_trials; ++trial) {
                random_basis_element(new_basis_element.data());
                const T new_energy = refine_amoeba(
                        new_basis_element.data(), initial_step_size, max_steps);
                if (new_energy < best_energy) {
                    best_basis_element = new_basis_element;
                    best_energy = new_energy;
                }
            }
            if (!best_basis_element.empty()) {
                add_basis_element(best_basis_element);
                return true;
            } else {
                return false;
            }
        }

    public: // =================================================================

        T refine_amoeba(T *RESTRICT basis_element,
                        const T &initial_step_size,
                        std::size_t max_steps) {
            std::vector<T> basis_matrix(num_parameters);
            dznl::AmoebaOptimizer<T> amoeba(
                    basis_element, num_parameters, initial_step_size,
                    [&](const T *b) {
                        construct_basis_matrix(basis_matrix.data(), b);
                        return augmented_ground_state_energy(
                                basis_matrix.data());
                    });
            for (std::size_t step = 0; step < max_steps; ++step) {
                amoeba.step();
            }
            return amoeba.current_minimum(basis_element);
        }

        void recompute_basis_matrices() {
            basis_matrices.clear();
            for (const auto &basis_element : basis) {
                std::vector<T> basis_matrix(num_parameters);
                construct_basis_matrix(basis_matrix.data(),
                                       basis_element.data());
                basis_matrices.push_back(basis_matrix);
            }
        }

        void recompute_solver_matrices() {
            const std::size_t num_basis_matrices = basis_matrices.size();
            solver.set_basis_size_destructive(num_basis_matrices);
            for (std::size_t i = 0; i < num_basis_matrices; ++i) {
                for (std::size_t j = 0; j < num_basis_matrices; ++j) {
                    context.evaluate_matrix_elements(
                            solver.overlap_matrix_element(i, j),
                            solver.hamiltonian_matrix_element(i, j),
                            basis_matrices[i].data(),
                            basis_matrices[j].data());
                }
            }
        }

    }; // class SphericalECGJacobiVariationalOptimizer

} // namespace zsvm

#endif // ZSVM_SPHERICAL_ECG_JACOBI_VARIATIONAL_OPTIMIZER_HPP_INCLUDED


/**
 * @file
 * @brief Interface describing a three-layer feedforward neural network.
 * @author Arno Bastenhof
 */

#ifndef NEURAL_HPP_
#define NEURAL_HPP_

#include <cmath>
#include <stdexcept>
#include <cstdint>

#include <armadillo>

namespace mnist {

class NeuralNet {
public:
  explicit          NeuralNet();
                    NeuralNet(const NeuralNet &) = delete;
                    NeuralNet(NeuralNet &&) = delete;
  NeuralNet&        operator=(const NeuralNet &) = delete;
  NeuralNet&        operator=(NeuralNet&&) = delete;
  void              LearnWeights(const arma::Mat<uint8_t>& data,
                        const double rate, const double reg, int epochs);
  double            Evaluate(const arma::Mat<uint8_t>& data) const;
private:
  enum {
    kBatchSz = 50,
    kInputLayerSz = 784,
    kHiddenLayerSz = 30,
    kOutputLayerSz = 10,
    kWeightsHeadSz = kHiddenLayerSz * (kInputLayerSz + 1),  // rows * cols
    kWeightsTailSz = kOutputLayerSz * (kHiddenLayerSz + 1), // rows * cols
    kWeightsSz = kWeightsHeadSz + kWeightsTailSz
  };
  arma::vec         weights_;
  mutable arma::mat activ_l1_;
  mutable arma::mat activ_l2_;
  mutable arma::mat activ_l3_;
  static void       ValidateSize(const arma::Mat<uint8_t>&);
  void              InitWeights();          // randomly initializes weights
  void              ForwardProp(const arma::Mat<uint8_t>&) const;
  arma::vec         BackProp(const arma::Col<uint8_t>&, const double) const;
};

/**
 * @brief Default constructor.
 */
inline NeuralNet::NeuralNet()
  : weights_(kWeightsSz)
  , activ_l1_(kBatchSz, kInputLayerSz + 1)
  , activ_l2_(kBatchSz, kHiddenLayerSz + 1)
  , activ_l3_(kBatchSz, kOutputLayerSz)
{
  // set bias activations
  activ_l1_.col(0).ones();
  activ_l2_.col(0).ones();
}

inline void NeuralNet::ValidateSize(const arma::Mat<uint8_t>& data)
{
  // Training- and test set sizes known from the MNIST database
  if (data.n_rows % kBatchSz != 0) {
    throw std::runtime_error{"Unexpected dimensions of input data"};
  }
}

} // namespace mnist

#endif // NEURAL_HPP_


#include <iostream>
#include <boost/numeric/mtl/mtl.hpp>

using namespace std;
int main(int, char**)
{
    mtl::dense2D<double>    A(3, 4), A_t(3, 4), S(3, 3), V(3, 4), D(4,4);
    A= 0;

    A[0][0]=1;  A[0][1]=1;  A[0][2]=1;  A[0][3]=4;
    A[1][0]=1;  A[1][1]=2;  A[1][2]=2;  A[1][3]=3;
    A[2][0]=9;  A[2][1]=3;  A[2][2]=2;  A[2][3]=4;
    std::cout<<"A=\n"<< A <<"\n";

    boost::tie(S, V, D)= svd(A, 1.e-10)= svd(A);  //second argument is optional (missmatch of upper R (A= Q*R))
    std::cout<<"Matrix  S=\n"<< S <<"\n";
    std::cout<<"Matrix  V=\n"<< V <<"\n";
    std::cout<<"Matrix  D=\n"<< D <<"\n";
    A_t= S*V*trans(D);
    std::cout<<"Matrix  A=S*V*D'=\n"<< A_t <<"\n";

    return 0;
}


#ifndef H_CONVEX_MPC
#define H_CONVEX_MPC

// Standard
#include <math.h>
#include <iostream>
#include <stdio.h>

// Timer
#include <chrono>

// QP Solver
#include <ExternalSource/myOptimizer/Goldfarb/QuadProg++.hh>

// Eigen
#include <Eigen/Dense>
#include <eigen3/unsupported/Eigen/MatrixFunctions>

// Gait Cycle
#include <PnC/GaitCycle/GaitCycle.hpp>
#include <memory>

// Reaction Force Schedule
#include <PnC/MPC/ReactionForceSchedule.hpp>

// #define MPC_PRINT_ALL
// #define MPC_TIME_ALL

// We are following the MPC formulation from:
// Di Carlo, Jared, et al. "Dynamic locomotion in the mit cheetah 3 through convex model-predictive control."
// 2018 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). IEEE, 2018.

// We modify the formulation to smoothen the solutions at each iteration.
//  There is also an optional setting to smoothen the current solution from the previous solution

class CMPC{
public:
	CMPC();
	~CMPC();

	double robot_mass; // kg mass of the robot
	Eigen::MatrixXd I_robot; // Robot inertia

	int horizon;// mpc horizon (number of steps);
	double mpc_dt; // mpc time interval per horizon steo
	double mu; // coefficient of friction
	double fz_max; // maximum z reaction force for one force vector.
	double control_alpha_; // Regularization term on the controls
	double delta_smooth_; // Parameter to smoothen the reaction force results

	bool smooth_from_prev_result_; // Whether to smooth the solution of the MPC iteration using the previous result

	bool custom_smoothing_; // Enable user-speicified custom smoothing matrix
	Eigen::MatrixXd Dc1; // Custom smoothing for the entire horizon

	bool rotate_inertia; // whether or not the inertia needs to be rotated to the world frame.
						 // If the inertia is expressed in body frame this needs to be true.
						 // Otherwise if the inertia is already updated to be in the world frame, this can be set to false.

	Eigen::VectorXd latest_f_vec_out; // Container holding the the latest computed force output for control (not to be confused with the force at the end of the horizon)
	Eigen::VectorXd f_prev; // Previous reaction force vector for the first horizon.

	Eigen::VectorXd X_pred; // dimension: (13*horizon) state prediction over the horizon. [x_{k+1}, x_{k+2}, ..., x_{k+horizon}]
	Eigen::VectorXd F_out;  // dimension: (n_Fr*horizon) forces to exert over the horizon. [f_{k}, f_{k+1}, ..., f_{k - 1 + horizon}]

	std::shared_ptr<GaitCycle> gait_cycle_ptr; // pointer to the gait cycle object
	std::shared_ptr<ReactionForceSchedule> reaction_force_schedule_ptr; // pointer to the reaction force schedule

  	// Vector cost for the MPC: <<  th1,  th2,  th3,  px,  py,  pz,   w1,  w2,   w3,   dpx,  dpy,  dpz,  g
	// last term is gravity and should always be 0,0
	// cost_vec << 0.25, 0.25, 10.0, 2.0, 2.0, 50.0, 0.0, 0.0, 0.30, 0.20, 0.2, 0.10, 0.0;
	Eigen::VectorXd cost_vec;

	bool use_terminal_cost; // if true, will use the custom terminal cost
	Eigen::VectorXd terminal_cost_vec;

	// Returns the state vector size of the MPC per horizon.
	int getStateVecDim(){return 13;}

	// Helper function which returns a 13-vector given the state of the robot.
	// x = [Theta, p, omega, pdot, g] \in \mathbf{R}^13
	Eigen::VectorXd getx0(const double roll, const double pitch, const double yaw,
						  const double com_x, const double com_y, const double com_z,
						  const double roll_rate, const double pitch_rate, const double yaw_rate,
						  const double com_x_rate, const double com_y_rate, const double com_z_rate);

	// Gets the latest computed ground forces for control
	Eigen::VectorXd getComputedGroundForces(){return latest_f_vec_out;}
	// Get computed ground forces in matrix form where the columns are the ground reaction forces
	// in the same order as r_feet.
	Eigen::MatrixXd getMatComputedGroundForces();

	// Gets the computed forces over the horizon
	Eigen::VectorXd getForcesOverHorizon(){return F_out;}
	// Gets the predicted state evolution over the horizon
	Eigen::VectorXd getXpredOverHorizon(){return X_pred;}

	void setRobotMass(const double robot_mass_in){ robot_mass = robot_mass_in; }
	// Whether or not the inertia needs to be rotated to the world frame.
	void rotateBodyInertia(bool rotate_inertia_in){ rotate_inertia = rotate_inertia_in; }

	void setRobotInertia(const Eigen::MatrixXd & inertia_in){ I_robot = inertia_in; }
	void setDt(const double mpc_dt_in){mpc_dt = mpc_dt_in;}	// MPC dt interval per horizon
	void setHorizon(const int & horizon_in);	// MPC horizon (number of steps)

	void setMu(const double mu_in){ mu = mu_in;} // Set the coefficient of friction
	void setMaxFz(const double fz_max_in); // Set the maximum z reaction force for one force vector
	void setSmoothFromPrevResult(const bool smooth_prev_in){smooth_from_prev_result_ = smooth_prev_in;} // Sets whether to smoothen the current solution using the previous solution

	// Set preview start time for the gate cycle
	void setPreviewStartTime(const double t_preview_start_in);
	void setCustomGaitCycle(std::shared_ptr<GaitCycle> gait_cycle_ptr_in);
	void setCustomReactionForceSchedule(std::shared_ptr<ReactionForceSchedule> reaction_force_schedule_ptr_in);

  	// Vector cost for the MPC: <<  th1,  th2,  th3,  px,  py,  pz,   w1,  w2,   w3,   dpx,  dpy,  dpz,  g
	void setCostVec(const Eigen::VectorXd & cost_vec_in); // Sets the cost vector.

  	// Set Terminal vector cost for the MPC: <<  th1,  th2,  th3,  px,  py,  pz,   w1,  w2,   w3,   dpx,  dpy,  dpz,  g
	// use_terminal_cost_in : if true, will use the terminal cost specified.
	void setTerminalCostVec(const bool use_terminal_cost_in, const Eigen::VectorXd & terminal_cost_vec_in); // Sets the cost vector.

	void setControlAlpha(const double control_alpha_in){control_alpha_ = control_alpha_in;}
	void setDeltaSmooth(const double delta_smooth_in){delta_smooth_ = delta_smooth_in;}

	void enableCustomSmoothing(const bool enable_custom_smoothing){custom_smoothing_ = enable_custom_smoothing;}
	void setCustomSmoothing(const Eigen::MatrixXd & Dc1_in){ Dc1 = Dc1_in; }

	// Human readable prints out of f_vec_out.
	void print_f_vec(int & n_Fr, const Eigen::VectorXd & f_vec_out);

	// Helper function which transforms a constant x_des to a constant trajectory reference
	// Input: x_des \in \mathbf{R}^{13}
	// Ouput: X_ref \in \mathbf{R}^{13*horizon}
	void get_constant_desired_x(const Eigen::VectorXd & x_des, Eigen::VectorXd & X_ref);

	// The main solve MPC routine
	// All the values are in world frame
	// Inputs:
	// 		x0 = [Theta, p, omega, pdot, g] \in \mathbf{R}^{13}    (Starting state of the system)
	// 		X_des \in \mathbf{R}^{13*horizon} 				     (Reference or desired state evolution)
	// 		r_feet \in \mathbf{R}^{3 x number_of_point_contacts} (Matrix of point contact locations where each column is the xyz location of the contact)
	// Outputs:
	//		x_pred \in \mathbf{R}^{13} - The predicted state after a time interval of mpc_dt has dimension
	//		f_vec_out \mathbf{R}^{3 x number_of_point_contacts} The xyz forces needed to be exerted at the contact points to track X_des.

	void solve_mpc(const Eigen::VectorXd & x0, const Eigen::VectorXd & X_des, const Eigen::MatrixXd & r_feet,
	               Eigen::VectorXd & x_pred, Eigen::VectorXd & f_vec_out);

	// For testing
	void simulate_toy_mpc();

	void assemble_vec_to_matrix(const int & n, const int & m, const Eigen::VectorXd & vec, Eigen::MatrixXd & mat_out);
	void integrate_robot_dynamics(const double & dt, const Eigen::VectorXd & x_current, const Eigen::MatrixXd & f_Mat, const Eigen::MatrixXd & r_feet,
	                              Eigen::VectorXd & x_next);

private:
	Eigen::MatrixXd r_feet_; // Store the value of r_feet locally.
	double t_preview_start; // store start time of the preview.

	double gravity_acceleration;
  	Eigen::MatrixXd R_roll(const double & phi);
	Eigen::MatrixXd R_pitch(const double & theta);
	Eigen::MatrixXd R_yaw(const double & psi);
	Eigen::MatrixXd skew_sym_mat(const Eigen::VectorXd & v);

	void cont_time_state_space(const Eigen::VectorXd & x_current,
	                           const Eigen::MatrixXd & r_feet,
	                           Eigen::MatrixXd & A, Eigen::MatrixXd & B);
	void discrete_time_state_space(const Eigen::MatrixXd & A, const Eigen::MatrixXd & B, Eigen::MatrixXd & Adt, Eigen::MatrixXd & Bdt);

	void qp_matrices(const Eigen::MatrixXd & Adt, const Eigen::MatrixXd & Bdt, Eigen::MatrixXd & Aqp, Eigen::MatrixXd & Bqp);
	void get_force_constraints(const int & n_Fr, Eigen::MatrixXd & CMat, Eigen::VectorXd & cvec);
	void get_qp_constraints(const Eigen::MatrixXd & CMat, const Eigen::VectorXd & cvec, Eigen::MatrixXd & Cqp, Eigen::VectorXd & cvec_qp);
	void get_qp_costs(const int & n, const int & m, const Eigen::VectorXd & vecS_cost, const double & control_alpha, Eigen::MatrixXd & Sqp, Eigen::MatrixXd & Kqp, Eigen::MatrixXd& D0);

	// Converts location of the feet expressed in world frame to the CoM frame. We assume the CoM orientation frame is always aligned with world.
	void convert_r_feet_to_com_frame(const Eigen::VectorXd & p_com, const Eigen::MatrixXd & r_feet, Eigen::MatrixXd & r_feet_com);

	void solve_mpc_qp(const Eigen::MatrixXd & Aqp,  const Eigen::MatrixXd & Bqp, const Eigen::VectorXd & X_ref,
					  const Eigen::VectorXd & x0,   const Eigen::MatrixXd & Sqp, const Eigen::MatrixXd & Kqp,
                      const Eigen::MatrixXd& D0, const Eigen::VectorXd& f_prev_in,
					  const Eigen::MatrixXd & Cqp, const Eigen::VectorXd & cvec_qp,
					  Eigen::VectorXd & f_vec_out);

};

#endif


#include <SBGATPolyhedronGravityModel.hpp>
#include <SBGATPolyhedronGravityModelUQ.hpp>
#include <SBGATObjWriter.hpp>

#include <vtkCleanPolyData.h>
#include <vtkOBJReader.h>

#include <json.hpp>
#include <boost/progress.hpp>

#include <armadillo>

int main(){

	std::ifstream i("input_file.json");
	nlohmann::json input_data;
	i >> input_data;

	std::string PATH_SHAPE = input_data["PATH_SHAPE"];
	double CORRELATION_DISTANCE =  input_data["CORRELATION_DISTANCE"];

	double ERROR_STANDARD_DEV  = input_data["ERROR_STANDARD_DEV"];
	double DENSITY  = input_data["DENSITY"];

	bool UNIT_IN_METERS  = input_data["UNIT_IN_METERS"];
	bool HOLD_MASS_CONSTANT  = input_data["HOLD_MASS_CONSTANT"];

	int N_MONTE_CARLO = input_data["N_MONTE_CARLO"];

	std::string OUTPUT_DIR = input_data["OUTPUT_DIR"];

	std::cout << "- Path to shape: " << PATH_SHAPE << std::endl;
	std::cout << "- Standard deviation on point coordinates (m) : " << ERROR_STANDARD_DEV << std::endl;
	std::cout << "- Correlation distance (m) : " << CORRELATION_DISTANCE << std::endl;
	std::cout << "- Density (kg/m^3) : " << DENSITY << std::endl;
	std::cout << "- Monte Carlo Draws : " << N_MONTE_CARLO << std::endl;

	// Reading
	vtkSmartPointer<vtkOBJReader> reader = vtkSmartPointer<vtkOBJReader>::New();
	reader -> SetFileName(PATH_SHAPE.c_str());
	reader -> Update();

	// An instance of SBGATPolyhedronGravityModel is created to evaluate the PGM of
	// the considered polytdata
	vtkSmartPointer<SBGATPolyhedronGravityModel> pgm_filter = vtkSmartPointer<SBGATPolyhedronGravityModel>::New();
	pgm_filter -> SetInputConnection(reader -> GetOutputPort());
	pgm_filter -> SetDensity(DENSITY);

	if(UNIT_IN_METERS){
		pgm_filter -> SetScaleMeters();
	} else{
		pgm_filter -> SetScaleKiloMeters();
	}

	std::cout << "Building pgm ...\n";
	pgm_filter -> Update();

	// An instance of SBGATPolyhedronGravityModelUQ is created to perform
	// uncertainty quantification from the PGM associated to the shape
	SBGATPolyhedronGravityModelUQ pgm_uq;
	pgm_uq.SetModel(pgm_filter);
	pgm_uq.PrecomputeMassPropertiesPartials();
	std::cout << "Populating shape covariance ...\n";

	// Populate the shape vertices covariance
	pgm_uq.ComputeVerticesCovarianceGlobal(ERROR_STANDARD_DEV,CORRELATION_DISTANCE);
	// Regularizing the covariance
	int regularized_eigen_values = pgm_uq.RegularizeCovariance();

	std::cout << regularized_eigen_values << " eigenvalues were regularized\n";

	std::cout << "Saving non-zero partition of shape covariance ...\n";

	// Save the covariance
	pgm_uq.SaveNonZeroVerticesCovariance(OUTPUT_DIR + "shape_covariance.json");

	// Saving baseline slices
	pgm_uq.TakeAndSaveSlice(0,OUTPUT_DIR + "baseline_slice_x.txt",0);
	pgm_uq.TakeAndSaveSlice(1,OUTPUT_DIR + "baseline_slice_y.txt",0);
	pgm_uq.TakeAndSaveSlice(2,OUTPUT_DIR + "baseline_slice_z.txt",0);

	std::vector<arma::vec::fixed<3> > all_positions = {
		arma::vec::fixed<3>({300,0,0}),
		arma::vec::fixed<3>({400,0,0}),
		arma::vec::fixed<3>({500,0,0}),
		arma::vec::fixed<3>({-300,0,0}),
		arma::vec::fixed<3>({-400,0,0}),
		arma::vec::fixed<3>({-500,0,0}),
		arma::vec::fixed<3>({0,300,0}),
		arma::vec::fixed<3>({0,400,0}),
		arma::vec::fixed<3>({0,500,0}),
		arma::vec::fixed<3>({0,-300,0}),
		arma::vec::fixed<3>({0,-400,0}),
		arma::vec::fixed<3>({0,-500,0}),
		arma::vec::fixed<3>({0,0,300}),
		arma::vec::fixed<3>({0,0,400}),
		arma::vec::fixed<3>({0,0,500}),
		arma::vec::fixed<3>({0,0,-300}),
		arma::vec::fixed<3>({0,0,-400}),
		arma::vec::fixed<3>({0,0,-500}),
	};

	// Analytical UQ

	std::vector<arma::mat::fixed<3,3> > analytical_covariances_acc(all_positions.size());
	std::vector<double> analytical_variances_pot(all_positions.size());

	std::cout << "Computing analytical uncertainties ... ";
	auto start = std::chrono::system_clock::now();
	#pragma omp parallel for
	for (int e = 0; e < all_positions.size(); ++e){
		analytical_variances_pot[e] = pgm_uq.GetVariancePotential(all_positions[e],HOLD_MASS_CONSTANT);
		analytical_covariances_acc[e] = pgm_uq.GetCovarianceAcceleration(all_positions[e],HOLD_MASS_CONSTANT);
	}
	auto end = std::chrono::system_clock::now();

	std::chrono::duration<double> elapsed_seconds = end-start;
	std::cout << "Done computing analytical uncertainties in " << elapsed_seconds.count() << " s\n";

	// Running a Monte Carlo to compare againnst
	std::vector<arma::vec> deviations;
	std::vector<double> densities;
	std::vector<std::vector<arma::vec::fixed<3> > >  all_accelerations;
	std::vector < std::vector<double> > all_potentials;
	std::vector<vtkSmartPointer<vtkPolyData > > saved_shapes(10);

	std::cout << "Running MC ... ";

	start = std::chrono::system_clock::now();
	SBGATPolyhedronGravityModelUQ::RunMCUQPotentialAccelerationInertial(PATH_SHAPE,DENSITY,
		UNIT_IN_METERS,
		HOLD_MASS_CONSTANT,
		pgm_uq.GetCovarianceSquareRoot(),
		N_MONTE_CARLO,
		all_positions,
		OUTPUT_DIR,
		std::min(30,N_MONTE_CARLO),
		deviations,
		densities,
		all_accelerations,
		all_potentials);

	end = std::chrono::system_clock::now();

	elapsed_seconds = end-start;

	std::cout << "Done running MC in " << elapsed_seconds.count() << " s\n";

	std::vector<double> mc_variances_pot(all_positions.size());
	std::vector<arma::mat > mc_covariances_acc(all_positions.size());

	std::cout << "Computing MC dispersions...\n";

	#pragma omp parallel for
	for (int e = 0; e < mc_variances_pot.size(); ++e){

		arma::vec potentials_mc(N_MONTE_CARLO);
		arma::mat accelerations_mc(3,N_MONTE_CARLO);

		for (int sample = 0; sample < N_MONTE_CARLO; ++sample){
			potentials_mc(sample) = all_potentials[sample][e];
			accelerations_mc.col(sample) = all_accelerations[sample][e];
		}

		mc_variances_pot[e] = arma::var(potentials_mc);
		mc_covariances_acc[e] = arma::cov(accelerations_mc.t());

	}

	std::cout << "\t After " << N_MONTE_CARLO << " MC outcomes:\n";

	for (int e = 0; e < all_positions.size(); ++e){
		all_positions[e].t().print("\t At: ");
		std::cout << "\t\tMC variance in potential: " << mc_variances_pot[e] << std::endl;
		std::cout << "\t\tAnalytical variance in potential: " << analytical_variances_pot[e] << std::endl;

		std::cout << "\t\tError (%): " << (mc_variances_pot[e] - analytical_variances_pot[e])/analytical_variances_pot[e] * 100 << std::endl;

		std::cout << "\t\tMC Covariance in acceleration: \n" << mc_covariances_acc[e] << std::endl;
		std::cout << "\t\tAnalytical covariance in acceleration: \n" << analytical_covariances_acc[e] << std::endl;

		std::cout << "\t\tError (%): " << arma::norm(mc_covariances_acc[e] - analytical_covariances_acc[e])/arma::trace(mc_covariances_acc[e]) * 100 << std::endl;
	}

	return 0;
}


/* boost random/lognormal_distribution.hpp header file
 *
 * Copyright Jens Maurer 2000-2001
 * Permission to use, copy, modify, sell, and distribute this software
 * is hereby granted without fee provided that the above copyright notice
 * appears in all copies and that both that copyright notice and this
 * permission notice appear in supporting documentation,
 *
 * Jens Maurer makes no representations about the suitability of this
 * software for any purpose. It is provided "as is" without express or
 * implied warranty.
 *
 * See http://www.boost.org for most recent version including documentation.
 *
 * $Id: lognormal_distribution.hpp 11696 2001-11-14 21:53:38Z jmaurer $
 *
 * Revision history
 *  2001-02-18  moved to individual header files
 */

#ifndef BOOST_RANDOM_LOGNORMAL_DISTRIBUTION_HPP
#define BOOST_RANDOM_LOGNORMAL_DISTRIBUTION_HPP

#include <cmath>      // std::exp, std::sqrt
#include <cassert>
#include <boost/random/normal_distribution.hpp>

#ifdef BOOST_NO_STDC_NAMESPACE
namespace std {
  using ::log;
  using ::sqrt;
}
#endif

namespace boost {

#if defined(__GNUC__) && (__GNUC__ < 3)
// Special gcc workaround: gcc 2.95.x ignores using-declarations
// in template classes (confirmed by gcc author Martin v. Loewis)
  using std::sqrt;
  using std::exp;
#endif

template<class UniformRandomNumberGenerator, class RealType = double>
class lognormal_distribution
{
public:
  typedef UniformRandomNumberGenerator base_type;
  typedef RealType result_type;
  lognormal_distribution(base_type & rng, result_type mean,
                         result_type sigma)
    : _rng(rng, std::log(mean*mean/std::sqrt(sigma*sigma + mean*mean)),
           std::sqrt(std::log(sigma*sigma/mean/mean+1)))
  {
    assert(mean > 0);
  }
  // compiler-generated copy constructor is fine
  // normal_distribution cannot be assigned, neither can this class
  result_type operator()()
  {
#ifndef BOOST_NO_STDC_NAMESPACE
    // allow for Koenig lookup
    using std::exp;
#endif
    return exp(_rng());
  }

#ifndef BOOST_NO_OPERATORS_IN_NAMESPACE
  friend bool operator==(const lognormal_distribution& x,
                         const lognormal_distribution& y)
  { return x._rng == y._rng; }
#else
  // Use a member function
  bool operator==(const lognormal_distribution& rhs) const
  { return _rng == rhs._rng;  }
#endif
private:
  normal_distribution<base_type, result_type> _rng;
};

} // namespace boost

#endif // BOOST_RANDOM_LOGNORMAL_DISTRIBUTION_HPP


#include "frc971/control_loops/drivetrain/improved_down_estimator.h"

#include <Eigen/Geometry>
#include <random>

#include "frc971/control_loops/quaternion_utils.h"
#include "aos/testing/random_seed.h"
#include "frc971/control_loops/drivetrain/drivetrain_test_lib.h"
#include "frc971/control_loops/runge_kutta.h"
#include "glog/logging.h"
#include "gtest/gtest.h"

namespace frc971 {
namespace control_loops {
namespace testing {

namespace {
// Check if two quaternions are logically equal, to within some reasonable
// tolerance. This is needed because a single rotation can be represented by two
// quaternions.
bool QuaternionEqual(const Eigen::Quaterniond &a, const Eigen::Quaterniond &b,
                     double tolerance) {
  // If a == b, then a.inverse() * b will be the identity. The identity
  // quaternion is the only time where the vector portion of the quaternion is
  // zero.
  return (a.inverse() * b).vec().norm() <= tolerance;
}
}  // namespace

// Do a known transformation to see if quaternion integration is working
// correctly.
TEST(DownEstimatorTest, QuaternionIntegral) {
  Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
  Eigen::Vector3d uy = Eigen::Vector3d::UnitY();
  Eigen::Vector3d uz = Eigen::Vector3d::UnitZ();

  Eigen::Quaternion<double> q(
      Eigen::AngleAxis<double>(0.5 * M_PI, Eigen::Vector3d::UnitY()));

  Eigen::Quaternion<double> q0(
      Eigen::AngleAxis<double>(0, Eigen::Vector3d::UnitY()));

  auto qux = q * ux;

  VLOG(1) << "Q is w: " << q.w() << " vec: " << q.vec();
  VLOG(1) << "ux is " << ux;
  VLOG(1) << "qux is " << qux;

  // Start by rotating around the X body vector for pi/2
  Eigen::Quaternion<double> integral1(
      RungeKutta(std::bind(&drivetrain::DrivetrainUkf::QuaternionDerivative, ux,
                           std::placeholders::_1),
                 q0.coeffs(), 0.5 * M_PI));

  VLOG(1) << "integral1 * uz => " << integral1 * uz;

  // Then rotate around the Y body vector for pi/2
  Eigen::Quaternion<double> integral2(
      RungeKutta(std::bind(&drivetrain::DrivetrainUkf::QuaternionDerivative, uy,
                           std::placeholders::_1),
                 integral1.normalized().coeffs(), 0.5 * M_PI));

  VLOG(1) << "integral2 * uz => " << integral2 * uz;

  // Then rotate around the X body vector for -pi/2
  Eigen::Quaternion<double> integral3(
      RungeKutta(std::bind(&drivetrain::DrivetrainUkf::QuaternionDerivative,
                           -ux, std::placeholders::_1),
                 integral2.normalized().coeffs(), 0.5 * M_PI));

  integral1.normalize();
  integral2.normalize();
  integral3.normalize();

  VLOG(1) << "Integral is w: " << integral1.w() << " vec: " << integral1.vec()
          << " norm " << integral1.norm();

  VLOG(1) << "Integral is w: " << integral3.w() << " vec: " << integral3.vec()
          << " norm " << integral3.norm();

  VLOG(1) << "ux => " << integral3 * ux;
  EXPECT_NEAR(0.0, (ux - integral1 * ux).norm(), 5e-2);
  EXPECT_NEAR(0.0, (uz - integral1 * uy).norm(), 5e-2);
  EXPECT_NEAR(0.0, (-uy - integral1 * uz).norm(), 5e-2);

  EXPECT_NEAR(0.0, (uy - integral2 * ux).norm(), 5e-2);
  EXPECT_NEAR(0.0, (uz - integral2 * uy).norm(), 5e-2);
  EXPECT_NEAR(0.0, (ux - integral2 * uz).norm(), 5e-2);

  EXPECT_NEAR(0.0, (uy - integral3 * ux).norm(), 5e-2);
  EXPECT_NEAR(0.0, (-ux - integral3 * uy).norm(), 5e-2);
  EXPECT_NEAR(0.0, (uz - integral3 * uz).norm(), 5e-2);
}

TEST(DownEstimatorTest, UkfConstantRotation) {
  drivetrain::DrivetrainUkf dtukf(
      drivetrain::testing::GetTestDrivetrainConfig());
  const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
  EXPECT_EQ(0.0,
            (Eigen::Vector3d(0.0, 0.0, 1.0) - dtukf.H(dtukf.X_hat().coeffs()))
                .norm());
  Eigen::Matrix<double, 3, 1> measurement;
  measurement.setZero();
  for (int ii = 0; ii < 200; ++ii) {
    dtukf.Predict(ux * M_PI_2, measurement, std::chrono::milliseconds(5));
  }
  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(M_PI_2, ux));
  EXPECT_TRUE(QuaternionEqual(expected, dtukf.X_hat(), 0.01))
      << "Expected: " << expected.coeffs()
      << " Got: " << dtukf.X_hat().coeffs();
  EXPECT_NEAR(
      0.0,
      (Eigen::Vector3d(0.0, 1.0, 0.0) - dtukf.H(dtukf.X_hat().coeffs())).norm(),
      1e-10);
}

// Tests that the euler angles in the status message are correct.
TEST(DownEstimatorTest, UkfEulerStatus) {
  drivetrain::DrivetrainUkf dtukf(
      drivetrain::testing::GetTestDrivetrainConfig());
  const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
  const Eigen::Vector3d uy = Eigen::Vector3d::UnitY();
  const Eigen::Vector3d uz = Eigen::Vector3d::UnitZ();
  // First, rotate 3 radians in the yaw axis, then 0.5 radians in the pitch
  // axis, and then 0.1 radians about the roll axis.
  // The down estimator should ignore any of the pitch movement.
  constexpr double kYaw = 3.0;
  constexpr double kPitch = 0.5;
  constexpr double kRoll = 0.1;
  Eigen::Matrix<double, 3, 1> measurement;
  measurement.setZero();
  aos::monotonic_clock::time_point now = aos::monotonic_clock::epoch();
  const std::chrono::milliseconds dt(5);
  // Run a bunch of one-second rotations at the appropriate rate to cause the
  // total pitch/roll/yaw to be kPitch/kRoll/kYaw.
  for (int ii = 0; ii < 200; ++ii) {
    dtukf.UpdateIntegratedPositions(now);
    now += dt;
    dtukf.Predict(uz * kYaw, measurement, dt);
  }
  for (int ii = 0; ii < 200; ++ii) {
    dtukf.UpdateIntegratedPositions(now);
    now += dt;
    dtukf.Predict(uy * kPitch, measurement, dt);
  }
  EXPECT_FLOAT_EQ(kYaw, dtukf.yaw());
  for (int ii = 0; ii < 200; ++ii) {
    dtukf.UpdateIntegratedPositions(now);
    now += dt;
    dtukf.Predict(ux * kRoll, measurement, dt);
  }
  EXPECT_FLOAT_EQ(kYaw, dtukf.yaw());
  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(kPitch, uy) *
                                    Eigen::AngleAxis<double>(kRoll, ux));
  flatbuffers::FlatBufferBuilder fbb;
  fbb.ForceDefaults(true);
  fbb.Finish(dtukf.PopulateStatus(&fbb, now));

  aos::FlatbufferDetachedBuffer<drivetrain::DownEstimatorState> state(
      fbb.Release());
  EXPECT_EQ(kPitch, state.message().longitudinal_pitch());
  // The longitudinal pitch is not actually the same number as the roll, so we
  // don't check it here.

  EXPECT_TRUE(QuaternionEqual(expected, dtukf.X_hat(), 0.0001))
      << "Expected: " << expected.coeffs()
      << " Got: " << dtukf.X_hat().coeffs();
}

// Tests that if the gyro indicates no movement but that the accelerometer shows
// that we are slightly rotated, that we eventually adjust our estimate to be
// correct.
TEST(DownEstimatorTest, UkfAccelCorrectsBias) {
  drivetrain::DrivetrainUkf dtukf(
      drivetrain::testing::GetTestDrivetrainConfig());
  const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
  Eigen::Matrix<double, 3, 1> measurement;
  // Supply the accelerometer with a slightly off reading to ensure that we
  // don't require exactly 1g to work.
  measurement << 0.01, 0.99, 0.0;
  EXPECT_TRUE(
      QuaternionEqual(Eigen::Quaterniond::Identity(), dtukf.X_hat(), 0.0))
      << "X_hat: " << dtukf.X_hat().coeffs();
  EXPECT_EQ(0.0,
            (Eigen::Vector3d(0.0, 0.0, 1.0) - dtukf.H(dtukf.X_hat().coeffs()))
                .norm());
  for (int ii = 0; ii < 200; ++ii) {
    dtukf.Predict({0.0, 0.0, 0.0}, measurement, std::chrono::milliseconds(5));
  }
  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(M_PI_2, ux));
  EXPECT_TRUE(QuaternionEqual(expected, dtukf.X_hat(), 0.01))
      << "Expected: " << expected.coeffs()
      << " Got: " << dtukf.X_hat().coeffs();
}

// Tests that if the accelerometer is reading values with a magnitude that isn't
// ~1g, that we are slightly rotated, that we eventually adjust our estimate to
// be correct.
TEST(DownEstimatorTest, UkfIgnoreBadAccel) {
  drivetrain::DrivetrainUkf dtukf(
      drivetrain::testing::GetTestDrivetrainConfig());
  const Eigen::Vector3d uy = Eigen::Vector3d::UnitY();
  Eigen::Matrix<double, 3, 1> measurement;
  // Set up a scenario where, if we naively took the accelerometer readings, we
  // would think that we were rotated. But the gyro readings indicate that we
  // are only rotating about the Y (pitch) axis.
  measurement << 0.3, 1.0, 0.0;
  for (int ii = 0; ii < 200; ++ii) {
    dtukf.Predict({0.0, M_PI_2, 0.0}, measurement,
                  std::chrono::milliseconds(5));
  }
  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(M_PI_2, uy));
  EXPECT_TRUE(QuaternionEqual(expected, dtukf.X_hat(), 1e-1))
      << "Expected: " << expected.coeffs()
      << " Got: " << dtukf.X_hat().coeffs();
  EXPECT_NEAR(
      0.0,
      (Eigen::Vector3d(-1.0, 0.0, 0.0) - dtukf.H(dtukf.X_hat().coeffs()))
          .norm(),
      1e-10)
      << dtukf.H(dtukf.X_hat().coeffs());
}

// Tests that computing sigma points, and then computing the mean and covariance
// returns the original answer.
TEST(DownEstimatorTest, SigmaPoints) {
  const Eigen::Quaternion<double> mean(
      Eigen::AngleAxis<double>(M_PI / 2.0, Eigen::Vector3d::UnitX()));

  Eigen::Matrix<double, 3, 3> covariance;
  covariance << 0.4, -0.1, 0.2, -0.1, 0.6, 0.0, 0.2, 0.0, 0.5;
  covariance *= 0.1;

  const Eigen::Matrix<double, 4, 3 * 2 + 1> vectors =
      drivetrain::GenerateSigmaPoints(mean, covariance);

  const Eigen::Matrix<double, 4, 1> calculated_mean =
      frc971::controls::QuaternionMean(vectors);

  VLOG(1) << "actual mean: " << mean.coeffs();
  VLOG(1) << "calculated mean: " << calculated_mean;

  Eigen::Matrix<double, 3, 3 * 2 + 1> Wprime;
  Eigen::Matrix<double, 3, 3> calculated_covariance =
      drivetrain::ComputeQuaternionCovariance(
          Eigen::Quaternion<double>(calculated_mean), vectors, &Wprime);

  EXPECT_NEAR(1.0,
              (mean.conjugate().coeffs() * calculated_mean.transpose()).norm(),
              1e-4);

  EXPECT_NEAR(0.0, (calculated_covariance - covariance).norm(), 1e-8);
}

// Tests that computing sigma points with a large covariance that will precisely
// wrap, that we do clip the perturbations.
TEST(DownEstimatorTest, ClippedSigmaPoints) {
  const Eigen::Quaternion<double> mean(
      Eigen::AngleAxis<double>(M_PI / 2.0, Eigen::Vector3d::UnitX()));

  Eigen::Matrix<double, 3, 3> covariance;
  covariance << 0.4, -0.1, 0.2, -0.1, 0.6, 0.0, 0.2, 0.0, 0.5;
  covariance *= 100.0;

  const Eigen::Matrix<double, 4, 3 * 2 + 1> vectors =
      drivetrain::GenerateSigmaPoints(mean, covariance);

  const Eigen::Matrix<double, 4, 1> calculated_mean =
      frc971::controls::QuaternionMean(vectors);

  Eigen::Matrix<double, 3, 3 * 2 + 1> Wprime;
  Eigen::Matrix<double, 3, 3> calculated_covariance =
      drivetrain::ComputeQuaternionCovariance(
          Eigen::Quaternion<double>(calculated_mean), vectors, &Wprime);

  EXPECT_NEAR(1.0,
              (mean.conjugate().coeffs() * calculated_mean.transpose()).norm(),
              1e-4);

  const double calculated_covariance_norm = calculated_covariance.norm();
  const double covariance_norm = covariance.norm();
  EXPECT_LT(calculated_covariance_norm, covariance_norm / 2.0)
      << "Calculated covariance should be much smaller than the original "
         "covariance.";
}

}  // namespace testing
}  // namespace control_loops
}  // namespace frc971