ShapeSpaceSaliency.hpp

#pragma once
 
#include "HierarchySaliencyMap.hpp"
#include "../../utils/RegularGridAdjacency2D.hpp"
#include "../../utils/Common.hpp"
#include "../MorphologicalTree.hpp"
#include "../detail/TreeKernelAccess.hpp"
 
#include <algorithm>
#include <bit>
#include <cmath>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <sstream>
#include <span>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
 
namespace mmcfilters {
 
enum class ShapeSpaceExtremaPolarity { Minima, Maxima };
 
template <std::floating_point Real> struct ShapeSpaceExtremum {
    NodeId representative = InvalidNode;
    Real birthLevel{};
    Real deathLevel{};
    Real extinction{};
};
 
template <std::floating_point Real> struct ShapeSpaceExtinctionResult {
    std::vector<ShapeSpaceExtremum<Real>> extrema;
    std::vector<Real> nodeScores;
};
 
template <std::floating_point Real> struct ShapeSpaceSaliencyResult {
    std::vector<ShapeSpaceExtremum<Real>> extrema;
    std::vector<Real> nodeScores;
    EdgeSaliencyMap<Real> edgeMap;
};
 
class ShapeSpaceSaliency {
  private:
    static void validatePolarity(ShapeSpaceExtremaPolarity polarity, const char* context) {
        switch (polarity) {
        case ShapeSpaceExtremaPolarity::Minima:
        case ShapeSpaceExtremaPolarity::Maxima:
            return;
        }
        throw std::invalid_argument(std::string(context) + " received an unknown extrema polarity.");
    }
 
    static void validateRootedTree(const MorphologicalTree& tree, const char* context) { detail::requireCommittedRootedHierarchy(tree, context); }
 
    template <std::floating_point Real>
    static void validateNodeBuffer(const MorphologicalTree& tree, std::span<const Real> values, bool requireNonNegative, const char* valueName,
                                   const char* context) {
        const std::size_t expected = static_cast<std::size_t>(tree.numInternalNodeSlots());
        if (values.size() != expected) {
            std::ostringstream oss;
            oss << context << " requires one " << valueName << " value per dense internal node slot; expected " << expected << " values but got "
                << values.size() << ".";
            throw std::invalid_argument(oss.str());
        }
 
        for (NodeId nodeId : tree.aliveNodeIds()) {
            const Real value = values[static_cast<std::size_t>(nodeId)];
            if (!std::isfinite(value)) {
                std::ostringstream oss;
                oss << context << " requires finite " << valueName << " values; node " << nodeId << " is non-finite.";
                throw std::invalid_argument(oss.str());
            }
            if (requireNonNegative && value < Real{0}) {
                std::ostringstream oss;
                oss << context << " requires non-negative " << valueName << " values; node " << nodeId << " has value " << value << ".";
                throw std::invalid_argument(oss.str());
            }
        }
    }
 
    static void validateTreeAndAdjacency(const MorphologicalTree& tree, const RegularGridAdjacency2D& adjacency, const char* context) {
        validateRootedTree(tree, context);
 
        const int rows = tree.numRows();
        const int columns = tree.numColumns();
        const int numPixels = tree.numPixels();
        if (rows <= 0 || columns <= 0 || numPixels <= 0) {
            throw std::invalid_argument(std::string(context) + " requires a non-empty image/pixel domain.");
        }
 
        if (adjacency.getNumRows() != rows || adjacency.getNumColumns() != columns) {
            std::ostringstream oss;
            oss << context << " adjacency domain must match the tree image domain; got " << adjacency.getNumRows() << "x" << adjacency.getNumColumns()
                << " for a tree domain of " << rows << "x" << columns << ".";
            throw std::invalid_argument(oss.str());
        }
    }
 
    static const RegularGridAdjacency2D& requireStoredAdjacency(const MorphologicalTree& tree, const char* context) {
        const RegularGridAdjacency2D* adjacency = tree.constructionAdjacency();
        if (adjacency == nullptr) {
            throw std::invalid_argument(std::string(context) + " requires an attached adjacency relation; pass an explicit adjacency relation instead.");
        }
        return *adjacency;
    }
 
    template <std::floating_point Real>
    static Real checkedExtinction(Real birthLevel, Real deathLevel, ShapeSpaceExtremaPolarity polarity, NodeId representative, const char* context) {
        const long double birth = static_cast<long double>(birthLevel);
        const long double death = static_cast<long double>(deathLevel);
        const long double difference = polarity == ShapeSpaceExtremaPolarity::Minima ? death - birth : birth - death;
 
        if (difference < 0.0L) {
            std::ostringstream oss;
            oss << context << " produced an invalid extinction interval for representative node " << representative << ".";
            throw std::invalid_argument(oss.str());
        }
 
        const long double maximum = static_cast<long double>(std::numeric_limits<Real>::max());
        if (!std::isfinite(difference) || difference > maximum) {
            std::ostringstream oss;
            oss << context << " extinction for representative node " << representative << " is not representable by the requested floating-point type.";
            throw std::overflow_error(oss.str());
        }
 
        const Real result = static_cast<Real>(difference);
        if (!std::isfinite(result)) {
            std::ostringstream oss;
            oss << context << " extinction for representative node " << representative << " overflowed the requested floating-point type.";
            throw std::overflow_error(oss.str());
        }
        return result;
    }
 
  public:
    template <std::floating_point Real>
    [[nodiscard]] static ShapeSpaceExtinctionResult<Real> computeExtinctionValues(const MorphologicalTree& tree, std::span<const Real> attribute,
                                                                                  ShapeSpaceExtremaPolarity polarity) {
        constexpr const char* context = "ShapeSpaceSaliency::computeExtinctionValues";
        validatePolarity(polarity, context);
        validateRootedTree(tree, context);
        validateNodeBuffer(tree, attribute, false, "attribute", context);
 
        std::vector<NodeId> nodes;
        nodes.reserve(static_cast<std::size_t>(tree.numInternalNodeSlots()));
        for (NodeId nodeId : tree.aliveNodeIds()) {
            nodes.push_back(nodeId);
        }
 
        const int numSlots = tree.numInternalNodeSlots();
        const std::size_t slotCount = static_cast<std::size_t>(numSlots);
 
        std::vector<int> depth(slotCount, -1);
        std::vector<NodeId> stack;
        stack.push_back(tree.root());
        depth[static_cast<std::size_t>(tree.root())] = 0;
        while (!stack.empty()) {
            const NodeId nodeId = stack.back();
            stack.pop_back();
            for (NodeId childId : detail::TreeKernelAccess::children(tree, nodeId)) {
                depth[static_cast<std::size_t>(childId)] = depth[static_cast<std::size_t>(nodeId)] + 1;
                stack.push_back(childId);
            }
        }
 
        const std::span<const PixelId> smallestSupportPixelByNode = detail::TreeKernelAccess::smallestNodeSupportPixels(tree);
        const std::span<const std::int32_t> supportCardinalityByNode = detail::TreeKernelAccess::nodeSupportCardinalities(tree);
        const auto canonicalShapeSpaceNodePrecedes = [&](NodeId lhs, NodeId rhs) {
            if (lhs == rhs) {
                return false;
            }
            const std::size_t lhsIndex = static_cast<std::size_t>(lhs);
            const std::size_t rhsIndex = static_cast<std::size_t>(rhs);
            if (smallestSupportPixelByNode[lhsIndex] != smallestSupportPixelByNode[rhsIndex]) {
                return smallestSupportPixelByNode[lhsIndex] < smallestSupportPixelByNode[rhsIndex];
            }
            if (supportCardinalityByNode[lhsIndex] != supportCardinalityByNode[rhsIndex]) {
                return supportCardinalityByNode[lhsIndex] < supportCardinalityByNode[rhsIndex];
            }
            if (depth[lhsIndex] != depth[rhsIndex]) {
                return depth[lhsIndex] < depth[rhsIndex];
            }
            throw std::logic_error(std::string(context) +
                                   " cannot distinguish two live nodes by spatial support and hierarchy depth.");
        };
 
        std::sort(nodes.begin(), nodes.end(), [&](NodeId lhs, NodeId rhs) {
            const Real lhsLevel = attribute[static_cast<std::size_t>(lhs)];
            const Real rhsLevel = attribute[static_cast<std::size_t>(rhs)];
            if (lhsLevel == rhsLevel) {
                return canonicalShapeSpaceNodePrecedes(lhs, rhs);
            }
            if (polarity == ShapeSpaceExtremaPolarity::Minima) {
                return lhsLevel < rhsLevel;
            }
            return rhsLevel < lhsLevel;
        });
 
        Real globalMinimum = attribute[static_cast<std::size_t>(nodes.front())];
        Real globalMaximum = globalMinimum;
        for (NodeId nodeId : nodes) {
            const Real level = attribute[static_cast<std::size_t>(nodeId)];
            globalMinimum = std::min(globalMinimum, level);
            globalMaximum = std::max(globalMaximum, level);
        }
 
        std::vector<NodeId> componentParent(slotCount, InvalidNode);
        std::vector<int> componentSize(slotCount, 0);
        std::vector<int> survivor(slotCount, -1);
        std::vector<std::uint8_t> active(slotCount, 0);
        std::vector<NodeId> plateauRepresentative(slotCount, InvalidNode);
        std::vector<std::vector<NodeId>> priorComponents(slotCount);
 
        auto findComponent = [&](NodeId nodeId) {
            NodeId root = nodeId;
            while (componentParent[static_cast<std::size_t>(root)] != root) {
                root = componentParent[static_cast<std::size_t>(root)];
            }
            while (componentParent[static_cast<std::size_t>(nodeId)] != nodeId) {
                const NodeId next = componentParent[static_cast<std::size_t>(nodeId)];
                componentParent[static_cast<std::size_t>(nodeId)] = root;
                nodeId = next;
            }
            return root;
        };
 
        auto joinComponents = [&](NodeId lhs, NodeId rhs) {
            lhs = findComponent(lhs);
            rhs = findComponent(rhs);
            if (lhs == rhs) {
                return lhs;
            }
            const int lhsSize = componentSize[static_cast<std::size_t>(lhs)];
            const int rhsSize = componentSize[static_cast<std::size_t>(rhs)];
            if (lhsSize < rhsSize || (lhsSize == rhsSize && canonicalShapeSpaceNodePrecedes(rhs, lhs))) {
                std::swap(lhs, rhs);
            }
            componentParent[static_cast<std::size_t>(rhs)] = lhs;
            componentSize[static_cast<std::size_t>(lhs)] += componentSize[static_cast<std::size_t>(rhs)];
            return lhs;
        };
 
        auto forEachShapeNeighbor = [&](NodeId nodeId, auto&& visitor) {
            if (!tree.isRoot(nodeId)) {
                visitor(detail::TreeKernelAccess::nodeParent(tree, nodeId));
            }
            for (NodeId childId : detail::TreeKernelAccess::children(tree, nodeId)) {
                visitor(childId);
            }
        };
 
        std::vector<ShapeSpaceExtremum<Real>> extrema;
        std::vector<std::uint8_t> finalized;
 
        auto finishExtremum = [&](int extremumIndex, Real deathLevel) {
            if (extremumIndex < 0 || static_cast<std::size_t>(extremumIndex) >= extrema.size() || finalized[static_cast<std::size_t>(extremumIndex)] != 0) {
                throw std::runtime_error(std::string(context) + " encountered inconsistent component-extremum state.");
            }
            ShapeSpaceExtremum<Real>& extremum = extrema[static_cast<std::size_t>(extremumIndex)];
            extremum.deathLevel = deathLevel;
            extremum.extinction = checkedExtinction(extremum.birthLevel, deathLevel, polarity, extremum.representative, context);
            finalized[static_cast<std::size_t>(extremumIndex)] = 1;
        };
 
        auto isStronger = [&](int lhsIndex, int rhsIndex) {
            const auto& lhs = extrema[static_cast<std::size_t>(lhsIndex)];
            const auto& rhs = extrema[static_cast<std::size_t>(rhsIndex)];
            if (lhs.birthLevel != rhs.birthLevel) {
                if (polarity == ShapeSpaceExtremaPolarity::Minima) {
                    return lhs.birthLevel < rhs.birthLevel;
                }
                return rhs.birthLevel < lhs.birthLevel;
            }
            return canonicalShapeSpaceNodePrecedes(lhs.representative, rhs.representative);
        };
 
        std::size_t batchBegin = 0;
        while (batchBegin < nodes.size()) {
            const Real level = attribute[static_cast<std::size_t>(nodes[batchBegin])];
            std::size_t batchEnd = batchBegin + 1;
            while (batchEnd < nodes.size() && attribute[static_cast<std::size_t>(nodes[batchEnd])] == level) {
                ++batchEnd;
            }
 
            for (std::size_t i = batchBegin; i < batchEnd; ++i) {
                const NodeId nodeId = nodes[i];
                const std::size_t index = static_cast<std::size_t>(nodeId);
                active[index] = 1;
                componentParent[index] = nodeId;
                componentSize[index] = 1;
                survivor[index] = -1;
            }
 
            for (std::size_t i = batchBegin; i < batchEnd; ++i) {
                const NodeId nodeId = nodes[i];
                forEachShapeNeighbor(nodeId, [&](NodeId neighborId) {
                    if (active[static_cast<std::size_t>(neighborId)] != 0 && attribute[static_cast<std::size_t>(neighborId)] == level) {
                        static_cast<void>(joinComponents(nodeId, neighborId));
                    }
                });
            }
 
            std::vector<NodeId> plateauRoots;
            plateauRoots.reserve(batchEnd - batchBegin);
            for (std::size_t i = batchBegin; i < batchEnd; ++i) {
                const NodeId nodeId = nodes[i];
                const NodeId plateauRoot = findComponent(nodeId);
                plateauRoots.push_back(plateauRoot);
 
                NodeId& representative = plateauRepresentative[static_cast<std::size_t>(plateauRoot)];
                if (representative == InvalidNode || depth[static_cast<std::size_t>(nodeId)] < depth[static_cast<std::size_t>(representative)] ||
                    (depth[static_cast<std::size_t>(nodeId)] == depth[static_cast<std::size_t>(representative)] &&
                     canonicalShapeSpaceNodePrecedes(nodeId, representative))) {
                    representative = nodeId;
                }
 
                forEachShapeNeighbor(nodeId, [&](NodeId neighborId) {
                    if (active[static_cast<std::size_t>(neighborId)] != 0 && attribute[static_cast<std::size_t>(neighborId)] != level) {
                        priorComponents[static_cast<std::size_t>(plateauRoot)].push_back(findComponent(neighborId));
                    }
                });
            }
 
            std::sort(plateauRoots.begin(), plateauRoots.end(), canonicalShapeSpaceNodePrecedes);
            plateauRoots.erase(std::unique(plateauRoots.begin(), plateauRoots.end()), plateauRoots.end());
 
            for (NodeId plateauRoot : plateauRoots) {
                auto& adjacentComponents = priorComponents[static_cast<std::size_t>(plateauRoot)];
                for (NodeId& component : adjacentComponents) {
                    component = findComponent(component);
                }
                std::sort(adjacentComponents.begin(), adjacentComponents.end(), canonicalShapeSpaceNodePrecedes);
                adjacentComponents.erase(std::unique(adjacentComponents.begin(), adjacentComponents.end()), adjacentComponents.end());
 
                int winningExtremum = -1;
                if (adjacentComponents.empty()) {
                    const NodeId representative = plateauRepresentative[static_cast<std::size_t>(plateauRoot)];
                    winningExtremum = static_cast<int>(extrema.size());
                    extrema.push_back(ShapeSpaceExtremum<Real>{representative, level, level, Real{0}});
                    finalized.push_back(0);
                } else {
                    for (NodeId component : adjacentComponents) {
                        const int candidate = survivor[static_cast<std::size_t>(component)];
                        if (candidate < 0) {
                            throw std::runtime_error(std::string(context) + " found an active level component without a surviving extremum.");
                        }
                        if (winningExtremum < 0 || isStronger(candidate, winningExtremum)) {
                            winningExtremum = candidate;
                        }
                    }
 
                    for (NodeId component : adjacentComponents) {
                        const int candidate = survivor[static_cast<std::size_t>(component)];
                        if (candidate != winningExtremum) {
                            finishExtremum(candidate, level);
                        }
                    }
                }
 
                NodeId combinedRoot = plateauRoot;
                for (NodeId component : adjacentComponents) {
                    combinedRoot = joinComponents(combinedRoot, component);
                }
                survivor[static_cast<std::size_t>(combinedRoot)] = winningExtremum;
 
                plateauRepresentative[static_cast<std::size_t>(plateauRoot)] = InvalidNode;
                adjacentComponents.clear();
            }
 
            batchBegin = batchEnd;
        }
 
        const NodeId finalComponent = findComponent(nodes.front());
        for (NodeId nodeId : nodes) {
            if (findComponent(nodeId) != finalComponent) {
                throw std::runtime_error(std::string(context) + " did not produce one connected final level component.");
            }
        }
 
        const int dominantExtremum = survivor[static_cast<std::size_t>(finalComponent)];
        if (dominantExtremum < 0) {
            throw std::runtime_error(std::string(context) + " did not retain a dominant extremum.");
        }
        finishExtremum(dominantExtremum, polarity == ShapeSpaceExtremaPolarity::Minima ? globalMaximum : globalMinimum);
 
        for (std::uint8_t isFinalized : finalized) {
            if (isFinalized == 0) {
                throw std::runtime_error(std::string(context) + " left an extremum without a death level.");
            }
        }
 
        std::sort(extrema.begin(), extrema.end(),
                  [&](const auto& lhs, const auto& rhs) { return canonicalShapeSpaceNodePrecedes(lhs.representative, rhs.representative); });
 
        ShapeSpaceExtinctionResult<Real> result;
        result.extrema = std::move(extrema);
        result.nodeScores.assign(slotCount, Real{0});
        for (const ShapeSpaceExtremum<Real>& extremum : result.extrema) {
            result.nodeScores[static_cast<std::size_t>(extremum.representative)] = extremum.extinction;
        }
        return result;
    }
 
    template <std::floating_point Real>
    [[nodiscard]] static EdgeSaliencyMap<Real> projectContourScores(const MorphologicalTree& tree, std::span<const Real> nodeScores,
                                                                    const RegularGridAdjacency2D& adjacency) {
        constexpr const char* context = "ShapeSpaceSaliency::projectContourScores";
        validateTreeAndAdjacency(tree, adjacency, context);
        validateNodeBuffer(tree, nodeScores, true, "node-score", context);
 
        return projectContourScoresEstablished(tree, nodeScores, adjacency);
    }
 
    template <std::floating_point Real>
    [[nodiscard]] static EdgeSaliencyMap<Real> projectContourScores(const MorphologicalTree& tree, std::span<const Real> nodeScores) {
        return projectContourScores(tree, nodeScores, requireStoredAdjacency(tree, "ShapeSpaceSaliency::projectContourScores"));
    }
 
    template <std::floating_point Real>
    [[nodiscard]] static ShapeSpaceSaliencyResult<Real> compute(const MorphologicalTree& tree, std::span<const Real> attribute,
                                                                ShapeSpaceExtremaPolarity polarity, const RegularGridAdjacency2D& adjacency) {
        ShapeSpaceSaliencyResult<Real> result;
        ShapeSpaceExtinctionResult<Real> extinction = computeExtinctionValues(tree, attribute, polarity);
        result.extrema = std::move(extinction.extrema);
        result.nodeScores = std::move(extinction.nodeScores);
        validateTreeAndAdjacency(tree, adjacency, "ShapeSpaceSaliency::compute");
        result.edgeMap = projectContourScoresEstablished(tree, std::span<const Real>(result.nodeScores), adjacency);
        return result;
    }
 
    template <std::floating_point Real>
    [[nodiscard]] static ShapeSpaceSaliencyResult<Real> compute(const MorphologicalTree& tree, std::span<const Real> attribute,
                                                                ShapeSpaceExtremaPolarity polarity) {
        return compute(tree, attribute, polarity, requireStoredAdjacency(tree, "ShapeSpaceSaliency::compute"));
    }
 
  private:
    template <std::floating_point Real>
    [[nodiscard]] static EdgeSaliencyMap<Real> projectContourScoresEstablished(const MorphologicalTree& tree,
        std::span<const Real> nodeScores, const RegularGridAdjacency2D& adjacency) {
        EdgeSaliencyMap<Real> edgeMap;
        edgeMap.numRows = tree.numRows();
        edgeMap.numColumns = tree.numColumns();
        edgeMap.adjacencyRadius = adjacency.getRadius();
 
        const std::size_t slotCount = static_cast<std::size_t>(tree.numInternalNodeSlots());
        const std::size_t liftingLevelCount = std::max<std::size_t>(1, static_cast<std::size_t>(std::bit_width(slotCount - 1)));
 
        std::vector<int> depth(slotCount, -1);
        std::vector<std::vector<NodeId>> ancestor(liftingLevelCount, std::vector<NodeId>(slotCount, InvalidNode));
        std::vector<std::vector<Real>> pathMaximum(liftingLevelCount, std::vector<Real>(slotCount, Real{0}));
 
        std::vector<NodeId> stack{tree.root()};
        depth[static_cast<std::size_t>(tree.root())] = 0;
        while (!stack.empty()) {
            const NodeId nodeId = stack.back();
            stack.pop_back();
            const std::size_t nodeIndex = static_cast<std::size_t>(nodeId);
            ancestor[0][nodeIndex] = detail::TreeKernelAccess::nodeParent(tree, nodeId);
            pathMaximum[0][nodeIndex] = nodeScores[nodeIndex];
 
            for (NodeId childId : detail::TreeKernelAccess::children(tree, nodeId)) {
                depth[static_cast<std::size_t>(childId)] = depth[nodeIndex] + 1;
                stack.push_back(childId);
            }
        }
 
        for (std::size_t level = 1; level < liftingLevelCount; ++level) {
            for (NodeId nodeId : tree.aliveNodeIds()) {
                const std::size_t nodeIndex = static_cast<std::size_t>(nodeId);
                const NodeId middle = ancestor[level - 1][nodeIndex];
                const std::size_t middleIndex = static_cast<std::size_t>(middle);
                ancestor[level][nodeIndex] = ancestor[level - 1][middleIndex];
                pathMaximum[level][nodeIndex] = std::max(pathMaximum[level - 1][nodeIndex], pathMaximum[level - 1][middleIndex]);
            }
        }
 
        auto accumulateBranch = [&](NodeId smallestNode, NodeId lca, Real& edgeScore) {
            const int smallestNodeDepth = depth[static_cast<std::size_t>(smallestNode)];
            const int lcaDepth = depth[static_cast<std::size_t>(lca)];
            NodeId current = smallestNode;
            std::size_t remaining = static_cast<std::size_t>(smallestNodeDepth - lcaDepth);
            std::size_t level = 0;
            while (remaining != 0) {
                if ((remaining & std::size_t{1}) != 0) {
                    const std::size_t currentIndex = static_cast<std::size_t>(current);
                    edgeScore = std::max(edgeScore, pathMaximum[level][currentIndex]);
                    current = ancestor[level][currentIndex];
                }
                remaining >>= 1;
                ++level;
            }
        };
 
        const int numPixels = tree.numPixels();
        for (NodeId source = 0; source < numPixels; ++source) {
            const NodeId sourceSmallestNode = detail::TreeKernelAccess::smallestNodeMap(tree, source);
 
            for (int targetValue : detail::CommittedGridAccess::forwardNeighbors(adjacency, source)) {
                const NodeId target = static_cast<NodeId>(targetValue);
 
                const NodeId targetSmallestNode = detail::TreeKernelAccess::smallestNodeMap(tree, target);
 
                const NodeId lca = detail::TreeKernelAccess::lowestCommonAncestor(tree, sourceSmallestNode, targetSmallestNode);
 
                Real edgeScore{0};
                accumulateBranch(sourceSmallestNode, lca, edgeScore);
                accumulateBranch(targetSmallestNode, lca, edgeScore);
 
                edgeMap.sources.push_back(source);
                edgeMap.targets.push_back(target);
                edgeMap.values.push_back(edgeScore);
            }
        }
 
        return edgeMap;
    }
 
};
 
} // namespace mmcfilters