ExtremumResidualTreeByCTsBuilder.hpp

#pragma once
 
#include "ComponentTreeBuilder.hpp"
#include "NativeHierarchyOwned.hpp"
#include "ValuedMorphologicalTree.hpp"
#include "detail/DualMinMaxTreeIncrementalFilterLeaf.hpp"
#include "detail/TreeKernelAccess.hpp"
#include "../utils/Common.hpp"
 
#include <cassert>
#include <cstddef>
#include <iterator>
#include <list>
#include <optional>
#include <queue>
#include <span>
#include <utility>
#include <vector>
 
namespace mmcfilters {
namespace detail {
template <AltitudeValue T> struct ExtremumResidualTreeTestAccess;
} // namespace detail
 
template <AltitudeValue T> class ExtremumResidualTreeByCTsBuilder {
  public:
    explicit ExtremumResidualTreeByCTsBuilder(RegularGridAdjacency2D adjacency) : adjacency_(std::move(adjacency)) {}
 
    [[nodiscard]] NativeHierarchyOwned<T> build(const ImagePtr<T>& image) {
        auto minTree = detail::materializeNativeHierarchy<T>(ComponentTreeBuilder(adjacency_, false).build(image), std::nullopt,
                                                             detail::MaterializedVersionPolicy::PreserveLinkedConstruction);
        auto maxTree = detail::materializeNativeHierarchy<T>(ComponentTreeBuilder(adjacency_, true).build(image), std::nullopt,
                                                             detail::MaterializedVersionPolicy::PreserveLinkedConstruction);
        const auto maxTreeNodeSlots = maxTree.topology().numInternalNodeSlots();
        const auto minTreeNodeSlots = minTree.topology().numInternalNodeSlots();
        detail::DualMinMaxTreeIncrementalFilterLeaf<T> updater(&minTree, &maxTree, adjacency_);
        ExtremumAgenda agenda;
        SelectedSupport selectedSupport;
        ResidualTreeConstruction extTree(image->getSize(), maxTreeNodeSlots, minTreeNodeSlots);
 
        agenda.initialize(maxTree, minTree, updater);                 // Line 2.
        while (const auto selected = agenda.next(maxTree, minTree)) { // Lines 3-5.
            const auto& selectedExtremum = *selected;
            const bool isMaximum = selectedExtremum.polarity == Polarity::Maximum;
            ValuedMorphologicalTree<T>& primal = isMaximum ? maxTree : minTree;
            ValuedMorphologicalTree<T>& dual = isMaximum ? minTree : maxTree;
            selectedSupport.collectFromLeaf(selectedExtremum.nodeId, primal, dual);
            const NodeId primalParentNode = detail::TreeKernelAccess::nodeParent(primal.topology(), selectedExtremum.nodeId);
            const T nodeAltitude = primal.nodeAltitudes()[static_cast<std::size_t>(selectedExtremum.nodeId)];
 
            // Lines 6-11: add X_s and attach the pending supports contained in X_s as children.
            const NodeId residualNode = extTree.addSelectedSupport(selectedSupport, selectedExtremum, nodeAltitude);
 
            // Line 12: apply the elementary leveling.
            applyElementaryLeveling(updater, isMaximum, selectedExtremum.nodeId, selectedSupport);
 
            // Line 13: pending supports follow the proper-part transfers.
            extTree.updatePendingSupports(residualNode, selectedExtremum.polarity, primalParentNode, updater.getLastSelectedSupportNode(),
                                          updater.getLastProperPartTransfers());
 
            const Polarity dualPolarity = oppositePolarity(selectedExtremum.polarity);
            for (NodeId node : updater.getLastCandidateNodes())
                agenda.enqueue(dual, dualPolarity, node, updater);
            agenda.enqueue(primal, selectedExtremum.polarity, primalParentNode, updater);
        }
 
        // Lines 16-18: attach the remaining pending supports to the terminal root.
        const T terminalAltitude = minTree.nodeAltitudes()[static_cast<std::size_t>(minTree.topology().root())];
        assert(minTree.topology().numNodes() == 1 && maxTree.topology().numNodes() == 1);
        assert(maxTree.nodeAltitudes()[static_cast<std::size_t>(maxTree.topology().root())] == terminalAltitude);
        return std::move(extTree).completeTree(
            image, makeMorphologicalTreeSemantics(MorphologicalTreeKind::ExtremumResidualTree, SharedAdjacencyContext{adjacency_}), terminalAltitude);
    }
 
    [[nodiscard]] const RegularGridAdjacency2D& adjacency() const noexcept { return adjacency_; }
 
  private:
    friend struct detail::ExtremumResidualTreeTestAccess<T>;
    enum class Polarity {
        Maximum, 
        Minimum  
    };
    [[nodiscard]] static constexpr Polarity oppositePolarity(Polarity polarity) {
        return polarity == Polarity::Maximum ? Polarity::Minimum : Polarity::Maximum;
    }
 
    struct ExtremumCandidate {
        std::size_t supportCardinality = 0;  
        PixelId minimumPixel = InvalidPixel; 
        Polarity polarity;                   
        NodeId nodeId;                       
    };
    struct ExtremumCandidatePriority {
        bool operator()(const ExtremumCandidate& lhs, const ExtremumCandidate& rhs) const {
            if (lhs.supportCardinality != rhs.supportCardinality)
                return lhs.supportCardinality > rhs.supportCardinality;
            return lhs.minimumPixel > rhs.minimumPixel;
        }
    };
    struct ExtremumAgenda {
        std::priority_queue<ExtremumCandidate, std::vector<ExtremumCandidate>, ExtremumCandidatePriority> candidates; 
 
        void enqueue(const ValuedMorphologicalTree<T>& tree, Polarity polarity, NodeId node, const detail::DualMinMaxTreeIncrementalFilterLeaf<T>& updater) {
            if (!isNonterminalLeaf(tree, node))
                return;
            candidates.push({static_cast<std::size_t>(detail::TreeKernelAccess::properPartCardinality(tree.topology(), node)), updater.minPixel(tree, node),
                             polarity, node});
        }
 
        void initialize(const ValuedMorphologicalTree<T>& maxTree, const ValuedMorphologicalTree<T>& minTree,
                        const detail::DualMinMaxTreeIncrementalFilterLeaf<T>& updater) {
            for (NodeId node : maxTree.topology().aliveNodeIds())
                enqueue(maxTree, Polarity::Maximum, node, updater);
            for (NodeId node : minTree.topology().aliveNodeIds())
                enqueue(minTree, Polarity::Minimum, node, updater);
        }
 
        [[nodiscard]] std::optional<ExtremumCandidate> next(const ValuedMorphologicalTree<T>& maxTree, const ValuedMorphologicalTree<T>& minTree) {
            while (!candidates.empty()) {
                const ExtremumCandidate selected = candidates.top();
                candidates.pop();
                const ValuedMorphologicalTree<T>& primal = selected.polarity == Polarity::Maximum ? maxTree : minTree;
                if (!isNonterminalLeaf(primal, selected.nodeId))
                    continue;
                // Flat zones only merge. Equal cardinality and retention of the
                // recorded minimum certify that the selection key is unchanged.
                if (static_cast<std::size_t>(detail::TreeKernelAccess::properPartCardinality(primal.topology(), selected.nodeId)) !=
                    selected.supportCardinality)
                    continue;
                if (detail::TreeKernelAccess::smallestNodeMap(primal.topology(), selected.minimumPixel) != selected.nodeId)
                    continue;
                return selected;
            }
            return std::nullopt;
        }
    };
 
    struct SelectedSupport {
        std::vector<PixelId> pixels;                  
        NodeId dualProperPartNode = InvalidNode;      
        NodeId wholeDualProperPartNode = InvalidNode; 
 
        void collectFromLeaf(NodeId leaf, const ValuedMorphologicalTree<T>& primal, const ValuedMorphologicalTree<T>& dual) {
            const auto& dualTopology = dual.topology();
            pixels.clear();
            wholeDualProperPartNode = InvalidNode;
            for (PixelId pixel : detail::TreeKernelAccess::properParts(primal.topology(), leaf)) {
                pixels.push_back(pixel);
            }
            // An extremal flat zone is nonempty and lies in one dual proper part.
            dualProperPartNode = detail::TreeKernelAccess::smallestNodeMap(dualTopology, pixels.front());
#ifndef NDEBUG
            for (PixelId pixel : pixels) {
                assert(detail::TreeKernelAccess::smallestNodeMap(dualTopology, pixel) == dualProperPartNode);
            }
#endif
            if (detail::TreeKernelAccess::properPartCardinality(dualTopology, dualProperPartNode) == static_cast<int>(pixels.size())) {
                wholeDualProperPartNode = dualProperPartNode;
            }
        }
    };
    [[gnu::always_inline]] static void applyElementaryLeveling(detail::DualMinMaxTreeIncrementalFilterLeaf<T>& updater, bool isMaximum, NodeId selectedLeafNode,
                                                               const SelectedSupport& selectedSupport) {
        if (isMaximum)
            updater.pruneMaxLeafAndUpdateMinTree(selectedLeafNode, selectedSupport.pixels, selectedSupport.dualProperPartNode,
                                                 selectedSupport.wholeDualProperPartNode);
        else
            updater.pruneMinLeafAndUpdateMaxTree(selectedLeafNode, selectedSupport.pixels, selectedSupport.dualProperPartNode,
                                                 selectedSupport.wholeDualProperPartNode);
    }
 
    class PendingSupportsByProperPart {
        std::vector<std::list<NodeId>> residualNodesByProperPart_;        
        std::vector<std::list<NodeId>::iterator> positionByResidualNode_; 
 
      public:
        explicit PendingSupportsByProperPart(std::size_t componentNodeSlots) : residualNodesByProperPart_(componentNodeSlots) {}
 
        std::list<NodeId>& pendingSupportsIn(NodeId properPartNode) { return residualNodesByProperPart_[static_cast<std::size_t>(properPartNode)]; }
        const std::list<NodeId>& pendingSupportsIn(NodeId properPartNode) const { return residualNodesByProperPart_[static_cast<std::size_t>(properPartNode)]; }
 
        void addPendingSupport(NodeId residualNode, NodeId properPartNode) {
            auto& pendingNodes = pendingSupportsIn(properPartNode);
            if (static_cast<std::size_t>(residualNode) >= positionByResidualNode_.size())
                positionByResidualNode_.resize(static_cast<std::size_t>(residualNode) + 1);
            pendingNodes.push_back(residualNode);
            positionByResidualNode_[static_cast<std::size_t>(residualNode)] = std::prev(pendingNodes.end());
        }
 
        void removePendingSupport(NodeId residualNode, NodeId properPartNode) {
            pendingSupportsIn(properPartNode).erase(positionByResidualNode_[static_cast<std::size_t>(residualNode)]);
        }
 
        void transferPendingSupports(NodeId sourceNode, NodeId targetNode) {
            auto& targetPendingSupports = pendingSupportsIn(targetNode);
            targetPendingSupports.splice(targetPendingSupports.end(), pendingSupportsIn(sourceNode));
        }
    };
 
    struct ResidualTreeConstruction {
        PendingSupportsByProperPart maxTreePendingSupports; 
        PendingSupportsByProperPart minTreePendingSupports; 
        std::vector<NodeId> nodeParents{InvalidNode};       
        std::vector<NodeId> smallestNodeMap;                
        std::vector<T> nodeAltitudes{T{}};                  
        ResidualTreeConstruction(std::size_t numPixels, std::size_t maxTreeNodeSlots, std::size_t minTreeNodeSlots)
            : maxTreePendingSupports(maxTreeNodeSlots), minTreePendingSupports(minTreeNodeSlots), smallestNodeMap(numPixels, InvalidNode) {
            nodeParents.reserve(numPixels);
            nodeAltitudes.reserve(numPixels);
        }
 
        NodeId addSelectedSupport(const SelectedSupport& selectedSupport, const ExtremumCandidate& selectedExtremum, T supportAltitude) {
            // Property 6 gives each recorded support a new nonempty proper part;
            // together with the terminal root this bounds the node count by pixels.
            PendingSupportsByProperPart& primalPendingSupports = pendingSupportsFor(selectedExtremum.polarity);
            PendingSupportsByProperPart& dualPendingSupports = pendingSupportsFor(oppositePolarity(selectedExtremum.polarity));
            auto& pendingChildren = primalPendingSupports.pendingSupportsIn(selectedExtremum.nodeId);
            const NodeId residualNode = static_cast<NodeId>(nodeParents.size());
            nodeParents.push_back(InvalidNode);       // Add X_s to implicit Xi, line 7.
            nodeAltitudes.push_back(supportAltitude); // eta(X_s), Eqs. 25, 33.
            for (PixelId pixel : selectedSupport.pixels) {
                if (smallestNodeMap[static_cast<std::size_t>(pixel)] == InvalidNode) {
                    smallestNodeMap[static_cast<std::size_t>(pixel)] = residualNode;
                }
            }
            while (!pendingChildren.empty()) {
                const NodeId childResidualNode = pendingChildren.front();
                nodeParents[static_cast<std::size_t>(childResidualNode)] = residualNode; // Eq. (32); remove child from implicit Xi, lines 9-10.
                dualPendingSupports.removePendingSupport(childResidualNode, selectedSupport.dualProperPartNode);
                primalPendingSupports.removePendingSupport(childResidualNode, selectedExtremum.nodeId);
            }
            return residualNode;
        }
 
        template <class ProperPartTransfer>
        void updatePendingSupports(NodeId residualNode, Polarity selectedPolarity, NodeId primalDestinationNode, NodeId dualDestinationNode,
                                   std::span<const ProperPartTransfer> dualProperPartTransfers) {
            PendingSupportsByProperPart& dualPendingSupports = pendingSupportsFor(oppositePolarity(selectedPolarity));
            // Pending supports follow whole proper-part transfers. The partial transfer
            // of X_s contains only its consumed children, already removed from Xi_s.
            for (const auto& transfer : dualProperPartTransfers)
                dualPendingSupports.transferPendingSupports(transfer.source, transfer.target);
            // Primal pruning moves the complete leaf to its parent. The dual updater
            // tracks the complete selected support through pixel moves and whole-part merges.
            pendingSupportsFor(selectedPolarity).addPendingSupport(residualNode, primalDestinationNode);
            dualPendingSupports.addPendingSupport(residualNode, dualDestinationNode);
        }
 
        [[nodiscard]] NativeHierarchyOwned<T> completeTree(const ImagePtr<T>& image, MorphologicalTreeSemantics treeSemantics, T rootAltitude) && {
            nodeParents[0] = 0;
            nodeAltitudes[0] = rootAltitude;
            // The remaining implicit Xi nodes become children of the terminal root.
            for (NodeId& parentNode : nodeParents)
                if (parentNode == InvalidNode)
                    parentNode = 0;
            for (NodeId& smallestNode : smallestNodeMap) {
                if (smallestNode == InvalidNode) {
                    smallestNode = 0;
                }
            }
            return makeNativeHierarchyOwned<T>(std::move(nodeParents), std::move(smallestNodeMap), std::move(nodeAltitudes), NodeId{0},
                                               GridDomain2D{image->getNumRows(), image->getNumColumns()}, std::move(treeSemantics));
        }
 
      private:
        PendingSupportsByProperPart& pendingSupportsFor(Polarity polarity) {
            return polarity == Polarity::Maximum ? maxTreePendingSupports : minTreePendingSupports;
        }
    };
 
    RegularGridAdjacency2D adjacency_; 
 
    [[nodiscard]] static bool isNonterminalLeaf(const ValuedMorphologicalTree<T>& tree, NodeId node) {
        const auto& topology = tree.topology();
        return detail::TreeKernelAccess::isAlive(topology, node) && node != topology.root() && detail::TreeKernelAccess::numChildren(topology, node) == 0;
    }
};
 
} // namespace mmcfilters