TopoOffsetTriMesh.h

#pragma once
#include <wmtk/TriMesh.h>
#include <wmtk/TriOptimizerMesh.h>
#include <algorithm>
#include <atomic>
#include <functional>
#include <set>
#include <wmtk/threading/enumerable_thread_specific.hpp>
#include "Parameters.h"
#include "SimplicialComplexBVH.hpp"
 
using CellTag = std::set<int64_t>;
 
 
namespace wmtk::components::topological_offset {
 
 
const int64_t TEMP_OFFSET_TRI_TAG = -1;
const CellTag TEMP_OFFSET_TRI_TAG_SET{TEMP_OFFSET_TRI_TAG};
 
constexpr double TOPOLOGY_BLOCK_AVG_FRAC = 0.05;
 
class VertexExtra2d
{
public:
    int label = 0;
    bool m_is_on_input = false; // on the input complex
    bool m_is_on_offset = false; // on the offset boundary itself
    bool m_is_on_region = false; // on some OTHER tag region's boundary
};
 
 
class EdgeExtra2d
{
public:
    int label = 0;
};
 
 
class FaceExtra2d
{
public:
    int label = 0;
};
 
 
class TopoOffsetTriMesh : public wmtk::TriOptimizerMesh
{
public: // mode for splitting in marching tets
    enum class EdgeSplitMode {
        Midpoint = 0, // used for simplicial embedding steps
        Initial = 1, // this is used to initialize the complex. Its a little hacky
        Optimization = 2 // the optimization phase; the shared engine places the vertex
    };
 
public:
    int m_vtu_counter = 0;
    std::array<size_t, 3> m_init_counts = {{0, 0, 0}};
    size_t m_tags_count;
    SimplicialComplexBVH m_input_complex_bvh;
    EdgeSplitMode m_edge_split_mode = EdgeSplitMode::Midpoint;
 
    // tag name maps
    std::map<std::string, int64_t> m_tag_name_to_id;
    std::map<int64_t, std::string> m_tag_id_to_name;
    CellTag m_offset_output_tag_ids;
 
    // if in 'singlebody' mode
    bool m_singlebody = false;
    int64_t m_single_tag;
 
    // just for retaining in output. dont actually use
    bool m_has_envelope = false;
    MatrixXd m_V_envelope;
    MatrixXi m_F_envelope;
 
    static constexpr int INPUT_SURFACE_CLASS = 0;
    static constexpr int OFFSET_SURFACE_CLASS = 1;
    static constexpr int REGION_SURFACE_CLASS = 2;
 
    Parameters& m_offset_params;
 
    using VertexExtraCol = wmtk::AttributeCollection<VertexExtra2d>;
    using EdgeExtraCol = wmtk::AttributeCollection<EdgeExtra2d>;
    using FaceExtraCol = wmtk::AttributeCollection<FaceExtra2d>;
    // m_vertex_attribute, m_edge_attribute and m_face_attribute are the base's; these three
    // are registered alongside them in its attribute groups.
    VertexExtraCol m_vertex_extra;
    EdgeExtraCol m_edge_extra;
    FaceExtraCol m_face_extra;
 
    TopoOffsetTriMesh(Parameters& _m_offset_params, int _num_threads = 0)
        : wmtk::TriOptimizerMesh(_m_offset_params)
        , m_offset_params(_m_offset_params)
    {
        NUM_THREADS = _num_threads;
        m_vertex_attr_group.add(&m_vertex_extra);
        m_edge_attr_group.add(&m_edge_extra);
        m_face_attr_group.add(&m_face_extra);
    }
 
    ~TopoOffsetTriMesh() override = default;
 
    void set_vertex_position(const size_t vid, const Vector2d& p)
    {
        m_vertex_attribute[vid].m_posf = p;
        m_vertex_attribute[vid].m_pos = to_rational(p);
        m_vertex_attribute[vid].m_is_rounded = true;
    }
 
    bool edge_is_offset(const size_t eid) const
    {
        return m_edge_attribute[eid].m_is_surface_fs &&
               m_edge_attribute[eid].m_surface_class == OFFSET_SURFACE_CLASS;
    }
    bool edge_is_region_boundary(const size_t eid) const
    {
        return m_edge_attribute[eid].m_is_surface_fs &&
               m_edge_attribute[eid].m_surface_class == REGION_SURFACE_CLASS;
    }
    struct EdgeSnapshot2d
    {
        EdgeAttributes tags;
        EdgeExtra2d extra;
    };
    struct FaceSnapshot2d
    {
        FaceAttributes attrs;
        FaceExtra2d extra;
    };
    EdgeSnapshot2d edge_snapshot(const size_t eid) const
    {
        return EdgeSnapshot2d{m_edge_attribute[eid], m_edge_extra[eid]};
    }
    void restore_edge(const size_t eid, const EdgeSnapshot2d& s)
    {
        m_edge_attribute[eid] = s.tags;
        m_edge_extra[eid] = s.extra;
    }
    FaceSnapshot2d face_snapshot(const size_t fid) const
    {
        return FaceSnapshot2d{m_face_attribute[fid], m_face_extra[fid]};
    }
    void restore_face(const size_t fid, const FaceSnapshot2d& s)
    {
        m_face_attribute[fid] = s.attrs;
        m_face_extra[fid] = s.extra;
    }
 
    void label_offset_boundary();
 
    bool face_in_region(const size_t fid) const;
 
    bool face_is_input_complex(const size_t fid) const;
 
    bool vertex_is_on_surface(const size_t vid) const override;
    bool edge_is_on_surface(const std::array<size_t, 2>& vids) const override;
 
    bool project_offset_vertex(const Tuple& t);
 
    bool region_boundary_tangent(const Tuple& t, Vector2d& tang) const;
 
    double minimize_distance_along_tangent(
        const size_t vid,
        const Vector2d& p0,
        const Vector2d& tang,
        const double cap,
        const std::function<bool()>& rejected);
 
    void optimize_offset(const std::filesystem::path& output_file);
 
    bool face_is_offset_band(const size_t fid) const;
 
    bool edge_is_region_boundary_live(const Tuple& t) const
    {
        const std::optional<Tuple> opp = t.switch_face(*this);
        if (!opp) return false; // domain boundary: frozen, never a region boundary
        const size_t fa = t.fid(*this), fb = opp->fid(*this);
        if (face_is_input_complex(fa) != face_is_input_complex(fb)) {
            return false; // input-complex boundary: frozen, not the envelope's job
        }
        if (face_is_offset_band(fa) != face_is_offset_band(fb)) {
            return false; // offset boundary: exempt by design, it is what the optimization moves
        }
        return m_face_attribute[fa].tags != m_face_attribute[fb].tags;
    }
 
    std::pair<double, double> compute_distance_deviation() const;
 
    mutable size_t m_worst_dist_vid = static_cast<size_t>(-1);
    void log_worst_dist_vertex() const;
 
    std::pair<double, double> compute_normal_deviation() const;
 
    double edge_normal_deviation(const Tuple& e) const;
 
    bool offset_field_normal(const Vector2d& p, Vector2d& n) const;
 
    bool edge_is_offset_surface_live(const Tuple& e) const;
 
    double max_offset_normal_deviation_at_vertex(const size_t vid) const;
 
    void init_offset_sizing_field();
 
    std::vector<std::array<double, 4>> optimization_metrics;
 
    std::vector<std::array<int, 3>> op_counts;
    std::atomic<int> iter_cnt_split = 0, iter_cnt_collapse = 0, iter_cnt_swap = 0;
    std::atomic<int> iter_cnt_collapse_nd_reject = 0;
    std::atomic<int> iter_cnt_swap_nd_reject = 0;
 
    mutable wmtk::threading::enumerable_thread_specific<double> m_collapse_nd_before;
 
    mutable wmtk::threading::enumerable_thread_specific<double> m_swap_nd_before;
    mutable wmtk::threading::enumerable_thread_specific<std::array<size_t, 4>> m_swap_verts;
 
    struct OptSplitCache2d
    {
        std::map<size_t, int> face_label;
        size_t v1_id = 0;
        size_t v2_id = 0;
    };
    wmtk::threading::enumerable_thread_specific<OptSplitCache2d> m_opt_split_cache;
 
    bool marching_split_edge_before(const Tuple& t);
    bool marching_split_edge_after(const Tuple& t);
 
    bool collapse_edge_before(const Tuple& t) override;
 
    bool swap_edge_before(const Tuple& t) override;
 
    bool swap_edge_after(const Tuple& t) override;
    bool collapse_before_vertex(size_t v1, size_t v2) override;
    void collapse_after_vertex(size_t v1, size_t v2) override;
    bool collapse_edge_after(const Tuple& t) override;
    void split_after_vertex(size_t v_new) override;
 
    bool split_adjust_position(size_t v_new, const std::vector<Tuple>& children) override;
 
    bool smooth_before(const Tuple& t) override;
    bool smooth_after(const Tuple& t) override;
 
    bool edge_is_input(const size_t eid) const
    {
        return m_edge_attribute[eid].m_is_surface_fs &&
               m_edge_attribute[eid].m_surface_class == INPUT_SURFACE_CLASS;
    }
 
    bool vertex_is_frozen(const size_t vid) const
    {
        return m_vertex_extra[vid].m_is_on_input || !m_vertex_attribute[vid].on_bbox_faces.empty();
    }
    bool edge_is_frozen(const size_t eid) const
    {
        return edge_is_input(eid) || m_edge_attribute[eid].m_is_bbox_fs >= 0;
    }
 
    void init_region_boundary_envelope();
 
    void init_region_boundary_envelope_from_input();
 
    void warn_if_offset_reaches_domain_boundary() const;
 
    bool region_boundary_is_outside_envelope(const size_t vid) const;
 
    struct SmoothTrace
    {
        std::atomic<int> attempted{0}; 
        std::atomic<int> before_bbox{0}; 
        std::atomic<int> before_unrounded{0}; 
        std::atomic<int> before_on_input{0}; 
        std::atomic<int> offset_attempted{0}; 
        std::atomic<int> offset_no_neighbours{0}; 
        std::atomic<int> offset_on_complex{0}; 
        std::atomic<int> offset_inverted{0}; 
        std::atomic<int> offset_envelope{0}; 
        std::atomic<int> offset_accepted{0};
        std::atomic<int> offset_clamped{0}; 
        std::atomic<int> offset_clamp_env{0}; 
        std::atomic<int> offset_clamp_inv{0}; 
        std::atomic<int> offset_slid{0};
        std::atomic<long long> offset_err_before_nano{0};
        std::atomic<long long> offset_err_after_nano{0};
        std::atomic<int> offset_err_max_before_nano{0};
        std::atomic<int> offset_err_max_after_nano{0};
        std::atomic<int> interior_attempted{0}; 
        std::atomic<int> region_attempted{0}; 
 
        void reset()
        {
            offset_err_before_nano.store(0);
            offset_err_after_nano.store(0);
            for (std::atomic<int>* c :
                 {&attempted,
                  &before_bbox,
                  &before_unrounded,
                  &before_on_input,
                  &offset_attempted,
                  &offset_no_neighbours,
                  &offset_on_complex,
                  &offset_inverted,
                  &offset_envelope,
                  &offset_accepted,
                  &offset_clamped,
                  &offset_clamp_env,
                  &offset_clamp_inv,
                  &offset_slid,
                  &offset_err_max_before_nano,
                  &offset_err_max_after_nano,
                  &interior_attempted,
                  &region_attempted}) {
                c->store(0);
            }
        }
    };
    SmoothTrace m_smooth_trace;
    void log_smooth_trace() const;
 
 
    std::shared_ptr<SampleEnvelope> surface_envelope_for_edge(
        const std::array<size_t, 2>& vids) const override
    {
        const auto [t, eid] = tuple_from_edge(vids);
        if (eid == static_cast<size_t>(-1) || !t.is_valid(*this)) {
            // The segment does not exist yet -- an operation asking about one it is about to
            // create. Answer with the envelope: containing a segment that turns out not to be a
            // region boundary only ever costs a rejected operation, while missing one that is
            // lets a region drift.
            return m_envelope;
        }
        // LIVE, not edge_is_region_boundary(eid): this is called from inside the very
        // split/collapse/swap passes that create the segment being asked about, well before the
        // next label_offset_boundary() would ever revisit its cached class. See
        // edge_is_region_boundary_live()'s comment for the collapse case that motivated this.
        return edge_is_region_boundary_live(t) ? m_envelope : nullptr;
    }
 
    bool smoothing_position_is_allowed(size_t vid, const Vector2d&) const override
    {
        return !m_vertex_extra[vid].m_is_on_input;
    }
 
    size_t refine_sizing_around_worst(double) override { return update_sizing_field(); }
    size_t update_sizing_field();
    void write_smoothing_debug_output(const std::string& path) const override
    {
        const_cast<TopoOffsetTriMesh*>(this)->write_vtu(path);
    }
 
 
    void init_from_image(
        const MatrixXd& V,
        const MatrixXi& F,
        const MatrixSi& F_tags,
        const MatrixXd& V_env,
        const MatrixXi& F_env,
        const std::vector<std::string>& tag_names);
 
    bool ambient_assert();
 
    void label_input_complex();
 
    bool empty_input_complex();
 
    void init_input_complex_bvh();
 
    bool split_edge_before(const Tuple& t) override;
    bool split_edge_after(const Tuple& t) override;
    bool split_face_before(const Tuple& t) override;
    bool split_face_after(const Tuple& t) override;
    bool invariants(const std::vector<Tuple>& tris) override;
 
    void execute_offset(const std::filesystem::path& output_file);
 
    void marching_tris();
 
 
 
    bool is_simplicially_embedded() const;
 
    bool tri_is_simp_emb(const Tuple& t) const;
 
    void simplicial_embedding();
 
    // variable offset stuff
 
    bool tag_tri_consistent_topology(size_t f_id, int64_t tag) const;
 
    bool offset_tri_consistent_topology(const size_t f_id) const;
 
    bool tri_is_in_offset_conservative(const size_t f_id, const double threshold_r) const;
 
    void grow_offset_conservative();
 
    void set_offset_tri_tags();
 
    bool offset_is_manifold();
 
    void write_input_complex(const std::string& path);
    void write_vtu(const std::string& path);
    // void write_msh(const std::string& file);
    void write_msh_groups(const std::string& file);
 
private:
    struct EdgeSplitCache
    {
        size_t v1_id;
        size_t v2_id;
        Vector2d new_v_pos;
        VertexExtra2d new_v_extra;
 
        // cache edge attributes
        EdgeSnapshot2d split_eattr;
        std::map<simplex::Edge, EdgeSnapshot2d> existing_eattr;
 
        // cache face attributes
        std::map<size_t, FaceSnapshot2d> opp_v_fattr;
    };
    wmtk::threading::enumerable_thread_specific<EdgeSplitCache> edge_split_cache;
 
    struct FaceSplitCache
    {
        size_t v1_id;
        size_t v2_id;
        size_t v3_id;
        Vector2d new_v_pos;
        VertexExtra2d new_v_extra;
 
        std::map<simplex::Edge, EdgeSnapshot2d> existing_eattr; // 3 orig edges
        FaceSnapshot2d split_fattr; // split face attributes
    };
    wmtk::threading::enumerable_thread_specific<FaceSplitCache> face_split_cache;
 
private: // helpers
    bool any_tag_present(const CellTag& tag1, const CellTag& tag2) const
    {
        for (const int64_t& i : tag1) {
            if (tag2.find(i) != tag2.end()) {
                return true;
            }
        }
        return false;
 
        // if (tag2.empty()) {
        //     return tag1.empty();
        // }
        // if (tag1.empty()) { // tag1 is ambient and tag2 is not
        //     return false;
        // }
 
        // for (const int64_t& i : tag1) {
        //     if (tag2.find(i) != tag2.end()) {
        //         return true;
        //     }
        // }
        // return false;
    }
 
    void sort_edges_by_length(std::vector<simplex::Edge>& edges)
    {
        std::sort(
            edges.begin(),
            edges.end(),
            [this](const simplex::Edge& e1, const simplex::Edge& e2) {
                double len1 = (m_vertex_attribute[e1.vertices()[0]].m_posf -
                               m_vertex_attribute[e1.vertices()[1]].m_posf)
                                  .squaredNorm();
                double len2 = (m_vertex_attribute[e2.vertices()[0]].m_posf -
                               m_vertex_attribute[e2.vertices()[1]].m_posf)
                                  .squaredNorm();
                return len1 > len2;
            });
    }
 
public: // helpers
    size_t edge_id_from_simplex(const simplex::Edge& e) const
    {
        const auto& verts = e.vertices();
        const auto incident = simplex_incident_triangles(e);
        const auto& faces = incident.faces();
 
        assert(!faces.empty()); // throw error here otherwise
 
        const size_t f_id = tuple_from_simplex(faces.front()).fid(*this);
        const Tuple t_edge = tuple_from_edge(verts[0], verts[1], f_id);
        return t_edge.eid(*this);
    }
 
    Tuple get_tuple_from_edge(const simplex::Edge& e) const
    {
        const auto& v = e.vertices();
        const auto faces = simplex_incident_triangles(e).faces();
        assert(!faces.empty());
        const size_t fid = tuple_from_simplex(faces.front()).fid(*this);
        return tuple_from_edge(v[0], v[1], fid);
    }
 
    std::vector<Tuple> get_edge_adjacent_faces(const Tuple& f) const
    {
        std::vector<Tuple> adj_tris;
        auto tri_1 = f.switch_face(*this);
        if (tri_1) {
            adj_tris.push_back(tri_1.value());
        }
        auto tri_2 = f.switch_edge(*this).switch_face(*this);
        if (tri_2) {
            adj_tris.push_back(tri_2.value());
        }
        auto tri_3 = f.switch_vertex(*this).switch_edge(*this).switch_face(*this);
        if (tri_3) {
            adj_tris.push_back(tri_3.value());
        }
        return adj_tris;
    }
};
 
 
} // namespace wmtk::components::topological_offset