AttributeFilters.hpp

#pragma once
 
#include "../utils/Image.hpp"
#include "../utils/Common.hpp"
#include "../trees/TreeAltitudeAlgorithms.hpp"
#include "../trees/ValuedMorphologicalTree.hpp"
#include "../trees/ValuedMorphologicalTreeView.hpp"
#include "../trees/detail/CommittedTreeAccess.hpp"
#include "../trees/detail/TreeTraversalDetail.hpp"
#include "../utils/CommittedImageAccess.hpp"
#include "../utils/Contract.hpp"
#include "AttributeReconstructionFilters.hpp"
#include "detail/ViterbiDecision.hpp"
 
#include <cmath>
#include <concepts>
#include <memory>
#include <stack>
#include <stdexcept>
#include <string>
#include <vector>
 
namespace mmcfilters {
 
template <AltitudeValue T> class AttributeFilters {
  protected:
 
    using AltitudeView = ValuedMorphologicalTreeView<T>;
 
    AltitudeView view_;
    const ValuedMorphologicalTree<T>* valuedTree_ = nullptr;
    const MorphologicalTree& tree;
    std::size_t treeMutationVersion_ = 0;
 
    AltitudeView view() const noexcept { return view_; }
 
    void requireStableTree(const char* context) const { MMCFILTERS_CONTRACT_CHECKED_ONLY(tree.requireMutationVersion(treeMutationVersion_, context)); }
 
    template <std::floating_point Real> static void requireAttributePointer(const Real* attribute, const char* context) {
        MMCFILTERS_CONTRACT_REQUIRE(attribute != nullptr,
                                    throw std::invalid_argument(std::string(context) + " requires a non-null attribute buffer."));
    }
 
    static void requireNodePreservationMaskSize(const MorphologicalTree& tree, const NodePreservationMask& nodePreservationMask,
                                                const char* context) {
        MMCFILTERS_CONTRACT_REQUIRE(nodePreservationMask.size() == static_cast<std::size_t>(tree.numInternalNodeSlots()),
                                    throw std::invalid_argument(std::string(context) +
                                                                " node-preservation-mask size must match the internal node slot count."));
    }
 
    template <typename TImagePtr> static void requireOutputImage(const MorphologicalTree& tree, const TImagePtr& image, const char* context) {
        MMCFILTERS_CONTRACT_REQUIRE(image != nullptr, throw std::invalid_argument(std::string(context) + " requires a non-null output image."));
        MMCFILTERS_CONTRACT_REQUIRE(image->getNumRows() == tree.numRows() && image->getNumColumns() == tree.numColumns(),
                                    throw std::invalid_argument(std::string(context) + " output image shape must match the tree image domain."));
    }
 
    static T altitudeOf(const AltitudeView& view, NodeId nodeId) noexcept { return view.nodeAltitudes()[static_cast<std::size_t>(nodeId)]; }
 
    template <typename TValue> static void writeProperParts(const MorphologicalTree& tree, NodeId nodeId, TValue* output, TValue value) {
        for (PixelId pixel : detail::CommittedTreeAccess::properParts(tree, nodeId)) {
            output[pixel] = value;
        }
    }
 
    template <typename TValue> static void writeSubtreeProperParts(const MorphologicalTree& tree, NodeId nodeId, TValue* output, TValue value) {
        for (NodeId subtreeNodeId : detail::CommittedTreeAccess::subtree(tree, nodeId)) {
            writeProperParts(tree, subtreeNodeId, output, value);
        }
    }
 
    static void filteringByPruningMinMaskImpl(AltitudeView view, const NodePreservationMask& nodePreservationMask, ImagePtr<T> imgOutputPtr) {
        const char* context = "AttributeFilters::filteringByPruningMin";
        view.requireTopologyUnchanged(context);
        const MorphologicalTree& tree = view.topology();
        requireNodePreservationMaskSize(tree, nodePreservationMask, context);
        requireOutputImage(tree, imgOutputPtr, context);
        std::stack<NodeId> stack;
        stack.push(tree.root());
        auto imgOutput = imgOutputPtr->rawData();
 
        while (!stack.empty()) {
            const NodeId nodeId = stack.top();
            stack.pop();
            writeProperParts(tree, nodeId, imgOutput, altitudeOf(view, nodeId));
            for (NodeId childNodeId : detail::CommittedTreeAccess::children(tree, nodeId)) {
                if (nodePreservationMask[static_cast<std::size_t>(childNodeId)]) {
                    stack.push(childNodeId);
                } else {
                    writeSubtreeProperParts(tree, childNodeId, imgOutput, altitudeOf(view, nodeId));
                }
            }
        }
    }
 
    static void filteringByPruningMaxMaskImpl(AltitudeView view, const NodePreservationMask& nodePreservationMask, ImagePtr<T> imgOutputPtr) {
        const char* context = "AttributeFilters::filteringByPruningMax";
        view.requireTopologyUnchanged(context);
        const MorphologicalTree& tree = view.topology();
        requireNodePreservationMaskSize(tree, nodePreservationMask, context);
        requireOutputImage(tree, imgOutputPtr, context);
        // This internal marker means "this whole subtree can be collapsed",
        // which is the opposite of the caller's preservation decision at the leaves
        // before descendant information is merged.
        std::vector<uint8_t> collapsibleRejectedSubtree(tree.numInternalNodeSlots(), false);
        detail::traversePostOrder(
            tree, tree.root(),
            [&collapsibleRejectedSubtree, &nodePreservationMask](NodeId nodeId) -> void {
                collapsibleRejectedSubtree[nodeId] = !nodePreservationMask[static_cast<std::size_t>(nodeId)];
            },
            [&collapsibleRejectedSubtree](NodeId parentNodeId, NodeId childNodeId) -> void {
                collapsibleRejectedSubtree[parentNodeId] =
                    (collapsibleRejectedSubtree[parentNodeId] & collapsibleRejectedSubtree[childNodeId]);
            },
            [](NodeId) -> void {});
 
        auto imgOutput = imgOutputPtr->rawData();
        std::stack<NodeId> stack;
        stack.push(tree.root());
        while (!stack.empty()) {
            const NodeId nodeId = stack.top();
            stack.pop();
            writeProperParts(tree, nodeId, imgOutput, altitudeOf(view, nodeId));
            for (NodeId childNodeId : detail::CommittedTreeAccess::children(tree, nodeId)) {
                if (!collapsibleRejectedSubtree[childNodeId]) {
                    stack.push(childNodeId);
                } else {
                    writeSubtreeProperParts(tree, childNodeId, imgOutput, altitudeOf(view, childNodeId));
                }
            }
        }
    }
 
    template <std::floating_point Real>
    static void filteringByPruningMinAttributeImpl(AltitudeView view, const Real* attribute, Real threshold, ImagePtr<T> imgOutputPtr) {
        const char* context = "AttributeFilters::filteringByPruningMin";
        view.requireTopologyUnchanged(context);
        const MorphologicalTree& tree = view.topology();
        requireAttributePointer(attribute, context);
        requireOutputImage(tree, imgOutputPtr, context);
        auto imgOutput = imgOutputPtr->rawData();
        std::stack<NodeId> stack;
        stack.push(tree.root());
        while (!stack.empty()) {
            const NodeId nodeId = stack.top();
            stack.pop();
            writeProperParts(tree, nodeId, imgOutput, altitudeOf(view, nodeId));
            for (NodeId childNodeId : detail::CommittedTreeAccess::children(tree, nodeId)) {
                if (attribute[childNodeId] > threshold) {
                    stack.push(childNodeId);
                } else {
                    writeSubtreeProperParts(tree, childNodeId, imgOutput, altitudeOf(view, nodeId));
                }
            }
        }
    }
 
    template <std::floating_point Real>
    static void filteringByPruningMaxAttributeImpl(AltitudeView view, const Real* attribute, Real threshold, ImagePtr<T> imgOutputPtr) {
        const char* context = "AttributeFilters::filteringByPruningMax";
        view.requireTopologyUnchanged(context);
        const MorphologicalTree& tree = view.topology();
        requireAttributePointer(attribute, context);
        requireOutputImage(tree, imgOutputPtr, context);
        // This internal marker records collapsible rejected subtrees, not a
        // direct keep decision. The post-order merge implements the universal
        // quantifier over descendants.
        std::vector<uint8_t> collapsibleRejectedSubtree(tree.numInternalNodeSlots(), false);
        detail::traversePostOrder(
            tree, tree.root(),
            [&collapsibleRejectedSubtree, attribute, threshold](NodeId nodeId) -> void {
                if (attribute[nodeId] <= threshold) {
                    collapsibleRejectedSubtree[nodeId] = true;
                }
            },
            [&collapsibleRejectedSubtree](NodeId parentNodeId, NodeId childNodeId) -> void {
                collapsibleRejectedSubtree[parentNodeId] =
                    (collapsibleRejectedSubtree[parentNodeId] & collapsibleRejectedSubtree[childNodeId]);
            },
            [](NodeId) -> void {});
 
        auto imgOutput = imgOutputPtr->rawData();
        std::stack<NodeId> stack;
        stack.push(tree.root());
        while (!stack.empty()) {
            const NodeId nodeId = stack.top();
            stack.pop();
            writeProperParts(tree, nodeId, imgOutput, altitudeOf(view, nodeId));
            for (NodeId childNodeId : detail::CommittedTreeAccess::children(tree, nodeId)) {
                if (!collapsibleRejectedSubtree[childNodeId]) {
                    stack.push(childNodeId);
                } else {
                    writeSubtreeProperParts(tree, childNodeId, imgOutput, altitudeOf(view, childNodeId));
                }
            }
        }
    }
 
 
  public:
    explicit AttributeFilters(AltitudeView view) : view_{view}, tree{view_.topology()}, treeMutationVersion_{tree.getMutationVersion()} {
        MMCFILTERS_CONTRACT_CHECKED_ONLY(view_.requireTopologyUnchanged("AttributeFilters"));
    }
 
    explicit AttributeFilters(const ValuedMorphologicalTree<T>& valuedTree) : AttributeFilters(valuedTree.asView()) { valuedTree_ = &valuedTree; }
 
    ~AttributeFilters() = default;
 
    template <std::floating_point Real> [[nodiscard]] ImagePtr<T> filteringByPruningMin(const std::shared_ptr<Real[]>& attr, Real threshold) {
        return filteringByPruningMin(attr.get(), threshold);
    }
 
    template <std::floating_point Real> [[nodiscard]] ImagePtr<T> filteringByPruningMin(const Real* attr, Real threshold) {
        ImagePtr<T> imgOutput = Image<T>::create(this->tree.numRows(), this->tree.numColumns());
        filteringByPruningMinAttributeImpl(view(), attr, threshold, imgOutput);
        return imgOutput;
    }
 
    template <std::floating_point Real> [[nodiscard]] ImagePtr<T> filteringByPruningMax(const std::shared_ptr<Real[]>& attr, Real threshold) {
        return filteringByPruningMax(attr.get(), threshold);
    }
 
    template <std::floating_point Real> [[nodiscard]] ImagePtr<T> filteringByPruningMax(const Real* attr, Real threshold) {
        ImagePtr<T> imgOutput = Image<T>::create(this->tree.numRows(), this->tree.numColumns());
        filteringByPruningMaxAttributeImpl(view(), attr, threshold, imgOutput);
        return imgOutput;
    }
 
    template <std::floating_point Real> [[nodiscard]] ImagePtr<T> filteringByViterbiRule(const Real* attr, Real threshold) {
        requireStableTree("AttributeFilters::filteringByViterbiRule");
        auto costs = detail::makeThresholdViterbiCosts(tree, attr, threshold);
        NodePreservationMask nodePreservationMask(detail::computeViterbiPreservationDecisions(tree, costs));
        return applyDirectAttributeFilter(view(), nodePreservationMask);
    }
 
    [[nodiscard]] ImagePtr<T> filteringByPruningMin(const NodePreservationMask& nodePreservationMask) {
        ImagePtr<T> imgOutput = Image<T>::create(this->tree.numRows(), this->tree.numColumns());
        filteringByPruningMinMaskImpl(view(), nodePreservationMask, imgOutput);
        return imgOutput;
    }
 
    [[nodiscard]] ImagePtr<T> filteringByPruningMax(const NodePreservationMask& nodePreservationMask) {
        ImagePtr<T> imgOutput = Image<T>::create(this->tree.numRows(), this->tree.numColumns());
        filteringByPruningMaxMaskImpl(view(), nodePreservationMask, imgOutput);
        return imgOutput;
    }
 
    static void filteringByPruningMin(const ValuedMorphologicalTreeView<T>& valuedTree, const NodePreservationMask& nodePreservationMask,
                                      ImagePtr<T> imgOutputPtr) {
        filteringByPruningMinMaskImpl(valuedTree, nodePreservationMask, imgOutputPtr);
    }
 
    static void filteringByPruningMin(const ValuedMorphologicalTree<T>& valuedTree, const NodePreservationMask& nodePreservationMask,
                                      ImagePtr<T> imgOutputPtr) {
        filteringByPruningMin(valuedTree.asView(), nodePreservationMask, imgOutputPtr);
    }
 
    static void filteringByPruningMax(const ValuedMorphologicalTreeView<T>& valuedTree, const NodePreservationMask& nodePreservationMask,
                                      ImagePtr<T> imgOutputPtr) {
        filteringByPruningMaxMaskImpl(valuedTree, nodePreservationMask, imgOutputPtr);
    }
 
    static void filteringByPruningMax(const ValuedMorphologicalTree<T>& valuedTree, const NodePreservationMask& nodePreservationMask,
                                      ImagePtr<T> imgOutputPtr) {
        filteringByPruningMax(valuedTree.asView(), nodePreservationMask, imgOutputPtr);
    }
 
    template <std::floating_point Real>
    static void filteringByPruningMin(const ValuedMorphologicalTreeView<T>& valuedTree, const std::shared_ptr<Real[]>& attribute, Real threshold, ImagePtr<T> imgOutputPtr) {
        filteringByPruningMin(valuedTree, attribute.get(), threshold, imgOutputPtr);
    }
 
    template <std::floating_point Real>
    static void filteringByPruningMin(const ValuedMorphologicalTree<T>& valuedTree, const std::shared_ptr<Real[]>& attribute, Real threshold,
                                      ImagePtr<T> imgOutputPtr) {
        filteringByPruningMin(valuedTree.asView(), attribute.get(), threshold, imgOutputPtr);
    }
 
    template <std::floating_point Real>
    static void filteringByPruningMin(const ValuedMorphologicalTreeView<T>& valuedTree, const Real* attribute, Real threshold, ImagePtr<T> imgOutputPtr) {
        filteringByPruningMinAttributeImpl(valuedTree, attribute, threshold, imgOutputPtr);
    }
 
    template <std::floating_point Real>
    static void filteringByPruningMin(const ValuedMorphologicalTree<T>& valuedTree, const Real* attribute, Real threshold, ImagePtr<T> imgOutputPtr) {
        filteringByPruningMin(valuedTree.asView(), attribute, threshold, imgOutputPtr);
    }
 
    template <std::floating_point Real>
    static void filteringByPruningMax(const ValuedMorphologicalTreeView<T>& valuedTree, const std::shared_ptr<Real[]>& attribute, Real threshold, ImagePtr<T> imgOutputPtr) {
        filteringByPruningMax(valuedTree, attribute.get(), threshold, imgOutputPtr);
    }
 
    template <std::floating_point Real>
    static void filteringByPruningMax(const ValuedMorphologicalTree<T>& valuedTree, const std::shared_ptr<Real[]>& attribute, Real threshold,
                                      ImagePtr<T> imgOutputPtr) {
        filteringByPruningMax(valuedTree.asView(), attribute.get(), threshold, imgOutputPtr);
    }
 
    template <std::floating_point Real>
    static void filteringByPruningMax(const ValuedMorphologicalTreeView<T>& valuedTree, const Real* attribute, Real threshold, ImagePtr<T> imgOutputPtr) {
        filteringByPruningMaxAttributeImpl(valuedTree, attribute, threshold, imgOutputPtr);
    }
 
    template <std::floating_point Real>
    static void filteringByPruningMax(const ValuedMorphologicalTree<T>& valuedTree, const Real* attribute, Real threshold, ImagePtr<T> imgOutputPtr) {
        filteringByPruningMax(valuedTree.asView(), attribute, threshold, imgOutputPtr);
    }
 
};
 
} // namespace mmcfilters