ExtinctionValues.hpp

#pragma once
 
#include "../trees/TreeAltitudeAlgorithms.hpp"
#include "../trees/ValuedMorphologicalTree.hpp"
#include "../trees/ValuedMorphologicalTreeView.hpp"
#include "../trees/detail/CommittedTreeAccess.hpp"
#include "../trees/detail/HierarchyCapabilityValidation.hpp"
#include "../trees/saliency/HierarchySaliencyMapValidation.hpp"
#include "../trees/saliency/HierarchySaliencyMap.hpp"
#include "../trees/saliency/HierarchicalWatershedSaliency.hpp"
#include "../utils/Image.hpp"
#include "../utils/Common.hpp"
#include "../utils/Contract.hpp"
#include "../contours/ContourComputation.hpp"
 
#include <algorithm>
#include <cmath>
#include <concepts>
#include <limits>
#include <memory>
#include <span>
#include <stack>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
 
namespace mmcfilters {
 
template <std::floating_point Real = float> struct RegionalExtremaNode {
    NodeId leaf;
 
    NodeId cutoffNode;
 
    Real extinction;
 
    RegionalExtremaNode(NodeId leaf, NodeId cutoffNode, Real extinction) : leaf(leaf), cutoffNode(cutoffNode), extinction(extinction) {}
};
 
template <std::floating_point Real = float> class ExtinctionSelectionPolicy {
    enum class Mode { TopK, MinimumExtinction };
 
    Mode mode_ = Mode::TopK;
    int extremaToKeep_ = 0;
    Real threshold_ = Real{0};
 
  public:
    [[nodiscard]] static ExtinctionSelectionPolicy byTopK(int extremaToKeep) noexcept {
        ExtinctionSelectionPolicy policy;
        policy.mode_ = Mode::TopK;
        policy.extremaToKeep_ = extremaToKeep;
        return policy;
    }
 
    [[nodiscard]] static ExtinctionSelectionPolicy byThreshold(Real threshold) noexcept {
        ExtinctionSelectionPolicy policy;
        policy.mode_ = Mode::MinimumExtinction;
        policy.threshold_ = threshold;
        return policy;
    }
 
    [[nodiscard]] bool selectsTopK() const noexcept { return mode_ == Mode::TopK; }
 
    [[nodiscard]] int extremaToKeep() const noexcept { return extremaToKeep_; }
 
    [[nodiscard]] Real minimumExtinction() const noexcept { return threshold_; }
};
 
enum class ExtinctionContourScorePolicy {
    RankScore,
    ExtinctionValue
};
 
template <AltitudeValue T, std::floating_point Real = float> class ExtinctionValues {
  protected:
 
    using AltitudeView = ValuedMorphologicalTreeView<T>;
 
    std::vector<RegionalExtremaNode<Real>> regionalExtremaNodes_;
    std::vector<Real> extinctionValueAttribute_;
    AltitudeView view_;
    const ValuedMorphologicalTree<T>* valuedTree_ = nullptr;
    const MorphologicalTree& tree;
    std::size_t treeMutationVersion_ = 0;
 
    struct SelectedExtremum {
        const RegionalExtremaNode<Real>* record = nullptr;
        int rankScore = 0;
    };
 
    [[nodiscard]] AltitudeView view() const { return valuedTree_ != nullptr ? valuedTree_->asView() : view_; }
 
    void requireStableTree(const char* context) const { tree.requireMutationVersion(treeMutationVersion_, context); }
 
    static void requireAttributePointer(const Real* attr, const char* context) {
        MMCFILTERS_CONTRACT_REQUIRE(attr != nullptr,
                                    throw std::invalid_argument(std::string(context) + " requires a non-null attribute buffer."));
    }
 
    static const Real* requireAttributeBuffer(const MorphologicalTree& tree, const std::vector<Real>& attr, const char* context) {
        MMCFILTERS_CONTRACT_REQUIRE(attr.size() == static_cast<std::size_t>(tree.numInternalNodeSlots()),
                                    throw std::invalid_argument(std::string(context) + " attribute size must match the internal node slot count."));
        return attr.data();
    }
 
    static T altitudeOf(const AltitudeView& view, NodeId nodeId) noexcept {
        return view.nodeAltitudes()[static_cast<std::size_t>(nodeId)];
    }
 
    std::vector<NodeId> collectComponentTreeExtrema() const { return this->tree.leaves(); }
 
    static Real dominantExtremumSentinel() noexcept { return std::numeric_limits<Real>::max(); }
 
    static void requireNonNegativeExtremaToKeep(int extremaToKeep, const char* context) {
        MMCFILTERS_CONTRACT_REQUIRE(extremaToKeep >= 0,
                                    throw std::invalid_argument(std::string(context) + " requires a non-negative extremaToKeep value."));
    }
 
    static void requireFiniteThreshold(Real threshold, const char* context) {
        MMCFILTERS_CONTRACT_REQUIRE(std::isfinite(threshold),
                                    throw std::invalid_argument(std::string(context) + " requires a finite extinction threshold."));
    }
 
    static void requireValidSelectionPolicy(const ExtinctionSelectionPolicy<Real>& policy, const char* context) {
        if (policy.selectsTopK()) {
            requireNonNegativeExtremaToKeep(policy.extremaToKeep(), context);
        } else {
            requireFiniteThreshold(policy.minimumExtinction(), context);
        }
    }
 
    static void requireFormalExtinctionValue(Real value, NodeId nodeId, const char* context) {
        if (!std::isfinite(value)) {
            throw std::invalid_argument(std::string(context) + " requires finite extinction values.");
        }
        if (value < Real{0}) {
            throw std::invalid_argument(std::string(context) + " requires non-negative extinction values.");
        }
        (void)nodeId;
    }
 
    void buildExtinctionValueAttribute() {
        extinctionValueAttribute_.assign(static_cast<std::size_t>(tree.numInternalNodeSlots()), Real{0});
        for (const RegionalExtremaNode<Real>& record : regionalExtremaNodes_) {
            Real& localValue = extinctionValueAttribute_[static_cast<std::size_t>(record.leaf)];
            if (localValue < record.extinction) {
                localValue = record.extinction;
            }
        }
 
        for (NodeId nodeId : tree.postOrder()) {
            Real& nodeValue = extinctionValueAttribute_[static_cast<std::size_t>(nodeId)];
            for (NodeId childId : detail::CommittedTreeAccess::children(tree, nodeId)) {
                const Real childValue = extinctionValueAttribute_[static_cast<std::size_t>(childId)];
                if (nodeValue < childValue) {
                    nodeValue = childValue;
                }
            }
        }
    }
 
    void validateExtinctionValueAttribute(const char* context) const {
        for (const RegionalExtremaNode<Real>& record : regionalExtremaNodes_) {
            requireFormalExtinctionValue(record.extinction, record.leaf, context);
        }
        HierarchySaliencyMapValidation::validateHierarchyValuation(tree, std::span<const Real>(extinctionValueAttribute_),
                                                                   HierarchyValuationPolicy::AllowLevelCollapse,
                                                                   HierarchyValuationRangePolicy::RequireNonNegative, context);
    }
 
    [[nodiscard]] std::vector<SelectedExtremum> selectExtrema(const ExtinctionSelectionPolicy<Real>& policy, const char* context) const {
        requireValidSelectionPolicy(policy, context);
 
        std::vector<const RegionalExtremaNode<Real>*> records;
        if (policy.selectsTopK()) {
            const int extremaToSelect = std::min(policy.extremaToKeep(), static_cast<int>(regionalExtremaNodes_.size()));
            records.reserve(static_cast<std::size_t>(extremaToSelect));
            for (int i = 0; i < extremaToSelect; ++i) {
                records.push_back(&regionalExtremaNodes_[static_cast<std::size_t>(i)]);
            }
        } else {
            records.reserve(regionalExtremaNodes_.size());
            for (const RegionalExtremaNode<Real>& record : regionalExtremaNodes_) {
                if (record.extinction >= policy.minimumExtinction()) {
                    records.push_back(&record);
                }
            }
        }
 
        std::vector<SelectedExtremum> selected;
        selected.reserve(records.size());
        const int selectedCount = static_cast<int>(records.size());
        for (int i = 0; i < selectedCount; ++i) {
            selected.push_back(SelectedExtremum{records[static_cast<std::size_t>(i)], selectedCount - i});
        }
        return selected;
    }
 
    [[nodiscard]] ImagePtr<T> filteringFromExtremumPreservationMask(std::vector<uint8_t> extremumPreservationMask, const char* context) const {
        requireStableTree(context);
        const AltitudeView altitudeView = view();
        for (NodeId nodeId : tree.postOrder()) {
            if (!tree.isRoot(nodeId) && extremumPreservationMask[static_cast<std::size_t>(nodeId)]) {
                extremumPreservationMask[static_cast<std::size_t>(detail::CommittedTreeAccess::nodeParent(tree, nodeId))] = true;
            }
        }
 
        ImagePtr<T> imgOutputPtr = Image<T>::create(tree.numRows(), tree.numColumns(), T{});
        auto imgOutput = imgOutputPtr->rawData();
        std::stack<NodeId> stack;
        stack.push(tree.root());
        while (!stack.empty()) {
            const NodeId nodeId = stack.top();
            stack.pop();
            const T level = altitudeOf(altitudeView, nodeId);
            for (PixelId pixel : detail::CommittedTreeAccess::properParts(tree, nodeId)) {
                imgOutput[pixel] = level;
            }
            for (NodeId childNodeId : detail::CommittedTreeAccess::children(tree, nodeId)) {
                if (extremumPreservationMask[static_cast<std::size_t>(childNodeId)]) {
                    stack.push(childNodeId);
                } else {
                    for (NodeId subtreeNodeId : detail::CommittedTreeAccess::subtree(tree, childNodeId)) {
                        for (PixelId pixel : detail::CommittedTreeAccess::properParts(tree, subtreeNodeId)) {
                            imgOutput[pixel] = level;
                        }
                    }
                }
            }
        }
        return imgOutputPtr;
    }
 
    [[nodiscard]] ImagePtr<T> filteringFromSelectedExtrema(const std::vector<SelectedExtremum>& selected, const char* context) const {
        std::vector<uint8_t> extremumPreservationMask(tree.numInternalNodeSlots(), false);
        for (const SelectedExtremum& item : selected) {
            extremumPreservationMask[static_cast<std::size_t>(item.record->leaf)] = true;
        }
        return filteringFromExtremumPreservationMask(std::move(extremumPreservationMask), context);
    }
 
    void initialize(const Real* attr) {
        requireAttributePointer(attr, "ExtinctionValues");
        std::vector<NodeId> leaves = collectComponentTreeExtrema();
        regionalExtremaNodes_.reserve(leaves.size());
        std::vector<uint8_t> visited(this->tree.numInternalNodeSlots(), false);
        for (NodeId leafNodeId : leaves) {
            Real extinction = dominantExtremumSentinel();
            NodeId cutoffNodeId = leafNodeId;
            NodeId parentNodeId = detail::CommittedTreeAccess::nodeParent(this->tree, cutoffNodeId);
            bool flag = true;
            while (flag && !this->tree.isRoot(cutoffNodeId)) {
                if (detail::CommittedTreeAccess::numChildren(this->tree, parentNodeId) > 1) {
                    for (NodeId sonNodeId : detail::CommittedTreeAccess::children(this->tree, parentNodeId)) {
                        if (flag) {
                            if (visited[sonNodeId] && sonNodeId != cutoffNodeId && attr[sonNodeId] == attr[cutoffNodeId]) {
                                flag = false;
                            } else if (sonNodeId != cutoffNodeId && attr[sonNodeId] > attr[cutoffNodeId]) {
                                flag = false;
                            }
                            visited[sonNodeId] = true;
                        }
                    }
                }
                if (flag) {
                    cutoffNodeId = parentNodeId;
                    parentNodeId = detail::CommittedTreeAccess::nodeParent(this->tree, cutoffNodeId);
                }
            }
            if (!this->tree.isRoot(cutoffNodeId)) {
                extinction = attr[cutoffNodeId];
            }
            regionalExtremaNodes_.emplace_back(leafNodeId, cutoffNodeId, extinction);
        }
 
        std::sort(regionalExtremaNodes_.begin(), regionalExtremaNodes_.end(), [](const auto& a, const auto& b) {
            if (a.extinction != b.extinction) {
                return a.extinction > b.extinction;
            }
            if (a.cutoffNode != b.cutoffNode) {
                return a.cutoffNode < b.cutoffNode;
            }
            return a.leaf < b.leaf;
        });
        buildExtinctionValueAttribute();
    }
 
  public:
    using value_type = Real;
 
    ExtinctionValues(const AltitudeView& view, const std::shared_ptr<Real[]>& attr) : ExtinctionValues(view, attr.get()) {}
 
    ExtinctionValues(const AltitudeView& view, const std::vector<Real>& attr)
        : ExtinctionValues(view, requireAttributeBuffer(view.topology(), attr, "ExtinctionValues")) {}
 
    ExtinctionValues(const AltitudeView& view, const Real* attr) : view_(view), tree(view_.topology()), treeMutationVersion_(tree.getMutationVersion()) {
        view_.requireTopologyUnchanged("ExtinctionValues");
        detail::validateGlobalMonotoneAltitudeOrder(this->tree, "ExtinctionValues");
        initialize(attr);
    }
 
    ExtinctionValues(const ValuedMorphologicalTree<T>& valuedTree, const std::shared_ptr<Real[]>& attr) : ExtinctionValues(valuedTree.asView(), attr.get()) {
        valuedTree_ = &valuedTree;
    }
 
    ExtinctionValues(const ValuedMorphologicalTree<T>& valuedTree, const std::vector<Real>& attr) : ExtinctionValues(valuedTree.asView(), attr) {
        valuedTree_ = &valuedTree;
    }
 
    ExtinctionValues(const ValuedMorphologicalTree<T>& valuedTree, const Real* attr) : ExtinctionValues(valuedTree.asView(), attr) { valuedTree_ = &valuedTree; }
 
    [[nodiscard]] ImagePtr<Real> contourMap(const ExtinctionSelectionPolicy<Real>& selection, ExtinctionContourScorePolicy scorePolicy) const {
        constexpr const char* context = "ExtinctionValues::contourMap";
        requireStableTree(context);
        const std::vector<SelectedExtremum> selected = selectExtrema(selection, context);
 
        std::vector<uint8_t> keep(tree.numInternalNodeSlots(), false);
        std::vector<Real> scoreByNode(tree.numInternalNodeSlots(), Real{0});
        std::vector<NodeId> keptNodes;
        for (const SelectedExtremum& item : selected) {
            const NodeId cutoffNode = item.record->cutoffNode;
            Real score = Real{0};
            switch (scorePolicy) {
            case ExtinctionContourScorePolicy::RankScore:
                score = static_cast<Real>(item.rankScore);
                break;
            case ExtinctionContourScorePolicy::ExtinctionValue:
                score = item.record->extinction;
                break;
            default:
                throw std::invalid_argument(std::string(context) + " received an unknown extinction contour score policy.");
            }
            if (!keep[static_cast<std::size_t>(cutoffNode)]) {
                keptNodes.push_back(cutoffNode);
                keep[static_cast<std::size_t>(cutoffNode)] = true;
                scoreByNode[static_cast<std::size_t>(cutoffNode)] = score;
            } else if (scoreByNode[static_cast<std::size_t>(cutoffNode)] < score) {
                scoreByNode[static_cast<std::size_t>(cutoffNode)] = score;
            }
        }
 
        ImagePtr<Real> imgOutputPtr = Image<Real>::create(tree.numRows(), tree.numColumns(), Real{0});
        auto contourOutput = imgOutputPtr->rawData();
 
        if (keptNodes.empty()) {
            return imgOutputPtr;
        }
        // Preserve the selected-node overwrite priority independently of the
        // contour traversal's post-order (shared boundary pixels can overlap).
        std::vector<std::size_t> priorityByNode(tree.numInternalNodeSlots(), 0);
        for (std::size_t selectionIndex = 0; selectionIndex < keptNodes.size(); ++selectionIndex) {
            priorityByNode[static_cast<std::size_t>(keptNodes[selectionIndex])] = selectionIndex + 1;
        }
        std::vector<std::size_t> pixelPriority(tree.numPixels(), 0);
        ContourComputation(tree).forEachContour([&](NodeId node, std::span<const PixelId> pixels) {
            const auto selectionPriority = priorityByNode[static_cast<std::size_t>(node)];
            if (selectionPriority == 0) {
                return;
            }
            for (PixelId pixel : pixels) {
                if (selectionPriority > pixelPriority[static_cast<std::size_t>(pixel)]) {
                    contourOutput[pixel] = scoreByNode[static_cast<std::size_t>(node)];
                    pixelPriority[static_cast<std::size_t>(pixel)] = selectionPriority;
                }
            }
        });
 
        return imgOutputPtr;
    }
 
    [[nodiscard]] const std::vector<Real>& getExtinctionValueAttribute() const {
        constexpr const char* context = "ExtinctionValues::getExtinctionValueAttribute";
        requireStableTree(context);
        validateExtinctionValueAttribute(context);
        return extinctionValueAttribute_;
    }
 
    [[nodiscard]] std::vector<int> computeRankedExtinctionValueAttribute() const {
        const std::vector<Real>& valuation = getExtinctionValueAttribute();
        return HierarchySaliencyMapValidation::rankHierarchyValuation(tree, std::span<const Real>(valuation), HierarchyValuationPolicy::AllowLevelCollapse);
    }
 
    [[nodiscard]] EdgeSaliencyMap<Real> computeFormalSaliencyEdgeMap(const RegularGridAdjacency2D& adjacency) const {
        const std::vector<Real>& valuation = getExtinctionValueAttribute();
        return HierarchicalWatershedSaliency::compute(view(), std::span<const Real>(valuation), adjacency);
    }
 
    [[nodiscard]] EdgeSaliencyMap<Real> computeFormalSaliencyEdgeMap() const {
        return computeFormalSaliencyEdgeMap(HierarchySaliencyMap::requireProjectionAdjacency(tree, "ExtinctionValues::computeFormalSaliencyEdgeMap"));
    }
 
    [[nodiscard]] EdgeSaliencyMap<int> computeRankedFormalSaliencyEdgeMap(const RegularGridAdjacency2D& adjacency) const {
        const std::vector<Real>& valuation = getExtinctionValueAttribute();
        return HierarchicalWatershedSaliency::computeRanked(view(), std::span<const Real>(valuation), adjacency);
    }
 
    [[nodiscard]] EdgeSaliencyMap<int> computeRankedFormalSaliencyEdgeMap() const {
        return computeRankedFormalSaliencyEdgeMap(
            HierarchySaliencyMap::requireProjectionAdjacency(tree, "ExtinctionValues::computeRankedFormalSaliencyEdgeMap"));
    }
 
    [[nodiscard]] EdgeSaliencyMap<Real> computeMonotoneExtinctionProjection(const RegularGridAdjacency2D& adjacency) const {
        const std::vector<Real>& valuation = getExtinctionValueAttribute();
        return HierarchySaliencyMap::computeSaliencyEdgeMap(tree, adjacency, std::span<const Real>(valuation), HierarchyValuationPolicy::AllowLevelCollapse);
    }
 
    [[nodiscard]] EdgeSaliencyMap<Real> computeMonotoneExtinctionProjection() const {
        return computeMonotoneExtinctionProjection(
            HierarchySaliencyMap::requireProjectionAdjacency(tree, "ExtinctionValues::computeMonotoneExtinctionProjection"));
    }
 
    [[nodiscard]] EdgeSaliencyMap<int> computeRankedMonotoneExtinctionProjection(const RegularGridAdjacency2D& adjacency) const {
        const std::vector<Real>& valuation = getExtinctionValueAttribute();
        return HierarchySaliencyMap::computeCanonicalRankedSaliencyEdgeMap(tree, adjacency, std::span<const Real>(valuation),
                                                                           HierarchyValuationPolicy::AllowLevelCollapse);
    }
 
    [[nodiscard]] EdgeSaliencyMap<int> computeRankedMonotoneExtinctionProjection() const {
        return computeRankedMonotoneExtinctionProjection(
            HierarchySaliencyMap::requireProjectionAdjacency(tree, "ExtinctionValues::computeRankedMonotoneExtinctionProjection"));
    }
 
    [[nodiscard]] ImagePtr<T> filtering(const ExtinctionSelectionPolicy<Real>& selection) const {
        constexpr const char* context = "ExtinctionValues::filtering";
        const std::vector<SelectedExtremum> selected = selectExtrema(selection, context);
        return filteringFromSelectedExtrema(selected, context);
    }
 
    [[nodiscard]] const std::vector<RegionalExtremaNode<Real>>& getRegionalExtrema() const {
        requireStableTree("ExtinctionValues::getRegionalExtrema");
        return regionalExtremaNodes_;
    }
};
 
} // namespace mmcfilters