#include <prim_gen.hpp>
#include <helper.hpp>

const std::array tet_boundary_faces = {0u, 1u, 2u, 3u};

void extract_iso_mesh(const std::array<uint32_t, 3>&                tet_active_subface_counts,
                      const stl_vector_mp<arrangement_t>&           tetrahedron_arrangements,
                      const scene_bg_mesh_info_t&                   scene_bg_mesh_info,
                      const stl_vector_mp<std::array<uint32_t, 4>>& tetrahedrons,
                      const stl_vector_mp<uint32_t>&                tetrahedron_active_subface_start_index,
                      const stl_vector_mp<uint32_t>&                active_subface_indices,
                      const Eigen::Ref<const Eigen::MatrixXd>       vertex_infos,
                      const flat_hash_map_mp<uint32_t, uint32_t>&   vertex_indices_mapping,
                      stl_vector_mp<Eigen::Vector3d>&               iso_pts,
                      stl_vector_mp<iso_vertex_t>&                  iso_verts,
                      stl_vector_mp<polygon_face_t>&                iso_faces)
{
    // estimate number of iso-verts and iso-faces
    const auto& [num_1_func, num_2_func, num_more_func] = tet_active_subface_counts;
    uint32_t max_num_face                               = num_1_func + 4 * num_2_func + 8 * num_more_func;
    uint32_t max_num_vert                               = max_num_face;
    iso_pts.reserve(max_num_vert);
    iso_verts.reserve(max_num_vert);
    iso_faces.reserve(max_num_face);

    // hash table for vertices on the boundary of tetrahedron
    flat_hash_map_mp<uint32_t, uint32_t>     vert_on_tetVert{};
    flat_hash_map_mp<pod_key_t<3>, uint32_t> vert_on_tetEdge{};
    vert_on_tetEdge.reserve(num_1_func + 3 * num_2_func + num_more_func);
    flat_hash_map_mp<pod_key_t<5>, uint32_t> vert_on_tetFace{};
    vert_on_tetFace.reserve(num_2_func + 6 * num_more_func);

    // hash table for faces on the boundary of tetrahedron
    flat_hash_map_mp<pod_key_t<3>, uint32_t> face_on_tetFace{};

    //
    flat_hash_set_mp<uint32_t> tet_iso_verts{};
    flat_hash_set_mp<uint32_t> tet_iso_faces{};
    // map: local vert index --> iso-vert index
    stl_vector_mp<uint32_t>    iso_vId_of_vert{};
    stl_vector_mp<uint32_t>    face_verts{};
    face_verts.reserve(4);
    pod_key_t<3>            key3;
    pod_key_t<5>            key5;
    std::array<uint16_t, 3> implicit_pIds;
    std::array<uint32_t, 3> boundary_pIds;
    std::array<uint32_t, 3> not_boundary_pIds;

    //
    for (uint32_t tet_index = 0; tet_index < tetrahedrons.size(); tet_index++) {
        const auto& arrangement = tetrahedron_arrangements[tet_index];
        const auto& vertices    = arrangement.vertices;
        const auto& faces       = arrangement.faces;
        auto        start_index = tetrahedron_active_subface_start_index[tet_index];

        // find vertices and faces on isosurface
        tet_iso_verts.clear();
        tet_iso_verts.reserve(vertices.size());
        tet_iso_faces.clear();
        tet_iso_faces.reserve(faces.size());
        if (arrangement.unique_planes.empty()) { // all planes are unique
            for (size_t i = 0; i < faces.size(); ++i) {
                const auto& face = faces[i];
                if (face.supporting_plane > 3) { // plane 0,1,2,3 are tet boundaries
                    tet_iso_faces.emplace(i);
                    for (const auto& vId : face.vertices) tet_iso_verts.emplace(vId);
                }
            }
        } else {
            for (size_t i = 0; i < faces.size(); ++i) {
                const auto& face = faces[i];
                const auto  pId  = face.supporting_plane;
                const auto  uId  = arrangement.unique_plane_indices[pId];
                for (const auto& plane_index : arrangement.unique_planes[uId]) {
                    if (plane_index > 3) { // plane 0,1,2,3 are tet boundaries
                        tet_iso_faces.emplace(i);
                        for (const auto& vId : face.vertices) tet_iso_verts.emplace(vId);
                    }
                }
            }
        }

        iso_vId_of_vert.clear();
        iso_vId_of_vert.reserve(vertices.size());
        // create iso-vertices
        for (const auto& vertex_index : tet_iso_verts) {
            const auto& vertex                 = vertices[vertex_index];
            auto&       local_vert_to_iso_vert = iso_vId_of_vert.emplace_back(invalid_index);
            auto&       iso_vert               = iso_verts.emplace_back();

            uint32_t num_boundary_planes{};
            uint32_t num_impl_planes{};
            for (const auto& vertex_plane : vertex) {
                if (vertex_plane >= 4) { // plane 0,1,2,3 are tet boundaries
                    implicit_pIds[num_impl_planes++] =
                        static_cast<uint16_t>(active_subface_indices[vertex_plane - 4 + start_index]);
                } else {
                    boundary_pIds[num_boundary_planes++] = vertex_plane;
                }
            }
            std::sort(boundary_pIds.begin(), boundary_pIds.begin() + num_boundary_planes);
            std::set_difference(tet_boundary_faces.begin(),
                                tet_boundary_faces.end(),
                                boundary_pIds.begin(),
                                boundary_pIds.begin() + num_impl_planes,
                                not_boundary_pIds.begin());

            const auto& tet_vertices = tetrahedrons[tet_index];
            iso_vert.header          = {tet_index, vertex_index, 4 - num_boundary_planes};
            if (num_boundary_planes == 0) { // in tet cell
                local_vert_to_iso_vert          = static_cast<uint32_t>(iso_verts.size());
                iso_vert.simplex_vertex_indices = tet_vertices;
                iso_vert.subface_indices        = implicit_pIds;

                const auto &vert_face0 = vertex_infos.row(vertex_indices_mapping.at(tet_vertices[0])),
                           vert_face1  = vertex_infos.row(vertex_indices_mapping.at(tet_vertices[1])),
                           vert_face2  = vertex_infos.row(vertex_indices_mapping.at(tet_vertices[2])),
                           vert_face3  = vertex_infos.row(vertex_indices_mapping.at(tet_vertices[3]));
                const auto f1          = std::array{vert_face0(implicit_pIds[0]),
                                           vert_face1(implicit_pIds[0]),
                                           vert_face2(implicit_pIds[0]),
                                           vert_face3(implicit_pIds[0])};
                const auto f2          = std::array{vert_face0(implicit_pIds[1]),
                                           vert_face1(implicit_pIds[1]),
                                           vert_face2(implicit_pIds[1]),
                                           vert_face3(implicit_pIds[1])};
                const auto f3          = std::array{vert_face0(implicit_pIds[2]),
                                           vert_face1(implicit_pIds[2]),
                                           vert_face2(implicit_pIds[2]),
                                           vert_face3(implicit_pIds[2])};
                const auto coord       = compute_barycentric_coords(f1, f2, f3);
                iso_pts.emplace_back(coord[0] * get_vert_pos(tet_vertices[0], scene_bg_mesh_info)
                                     + coord[1] * get_vert_pos(tet_vertices[1], scene_bg_mesh_info)
                                     + coord[2] * get_vert_pos(tet_vertices[2], scene_bg_mesh_info)
                                     + coord[3] * get_vert_pos(tet_vertices[3], scene_bg_mesh_info));
            } else if (num_boundary_planes == 1) { // on tet face
                key5               = {tet_vertices[not_boundary_pIds[0]],
                                      tet_vertices[not_boundary_pIds[1]],
                                      tet_vertices[not_boundary_pIds[2]],
                                      implicit_pIds[0],
                                      implicit_pIds[1]};
                auto iter_inserted = vert_on_tetFace.try_emplace(key5, static_cast<uint32_t>(iso_verts.size()));
                if (iter_inserted.second) {
                    iso_vert.simplex_vertex_indices = {key5[0], key5[1], key5[2]};
                    iso_vert.subface_indices        = {implicit_pIds[0], implicit_pIds[1]};

                    //
                    const auto &vert_face0 = vertex_infos.row(vertex_indices_mapping.at(key5[0])),
                               vert_face1  = vertex_infos.row(vertex_indices_mapping.at(key5[1])),
                               vert_face2  = vertex_infos.row(vertex_indices_mapping.at(key5[2]));
                    const auto f1 =
                        std::array{vert_face0(implicit_pIds[0]), vert_face1(implicit_pIds[0]), vert_face2(implicit_pIds[0])};
                    const auto f2 =
                        std::array{vert_face0(implicit_pIds[1]), vert_face1(implicit_pIds[1]), vert_face2(implicit_pIds[1])};
                    const auto coord = compute_barycentric_coords(f1, f2);
                    iso_pts.emplace_back(coord[0] * get_vert_pos(key5[0], scene_bg_mesh_info)
                                         + coord[1] * get_vert_pos(key5[1], scene_bg_mesh_info)
                                         + coord[2] * get_vert_pos(key5[2], scene_bg_mesh_info));
                }
                local_vert_to_iso_vert = iter_inserted.first->second;
            } else if (num_boundary_planes == 2) { // on tet edge
                key3               = {tet_vertices[not_boundary_pIds[0]], tet_vertices[not_boundary_pIds[1]], implicit_pIds[0]};
                auto iter_inserted = vert_on_tetEdge.try_emplace(key3, static_cast<uint32_t>(iso_verts.size()));
                if (iter_inserted.second) {
                    iso_vert.simplex_vertex_indices = {key3[0], key3[1]};
                    iso_vert.subface_indices        = {implicit_pIds[0]};

                    const auto f1    = vertex_infos(vertex_indices_mapping.at(key3[0]), implicit_pIds[0]);
                    const auto f2    = vertex_infos(vertex_indices_mapping.at(key3[1]), implicit_pIds[0]);
                    const auto coord = compute_barycentric_coords(f1, f2);
                    iso_pts.emplace_back(coord[0] * get_vert_pos(key3[0], scene_bg_mesh_info)
                                         + coord[1] * get_vert_pos(key3[1], scene_bg_mesh_info));
                }
                local_vert_to_iso_vert = iter_inserted.first->second;
            } else if (num_boundary_planes == 3) { // on tet vertex
                auto key           = tet_vertices[not_boundary_pIds[0]];
                auto iter_inserted = vert_on_tetVert.try_emplace(key, static_cast<uint32_t>(iso_verts.size()));
                if (iter_inserted.second) {
                    iso_vert.subface_indices = {static_cast<uint16_t>(key)};

                    iso_pts.emplace_back(get_vert_pos(iso_vert.simplex_vertex_indices[0], scene_bg_mesh_info));
                }
                local_vert_to_iso_vert = iter_inserted.first->second;
            }
        }
        // create iso-faces
        for (const auto& face_index : tet_iso_faces) {
            const auto& face = faces[face_index];
            face_verts.clear();

            for (unsigned long vId : face.vertices) { face_verts.emplace_back(iso_vId_of_vert[vId]); }
            const face_header_t face_header{tet_index, face_index};
            // face is on tet boundary if face.negative_cell is NONE
            uint32_t            iso_face_index{static_cast<uint32_t>(iso_faces.size())};
            bool                is_new_iso_face{};
            if (face.negative_cell == invalid_index) {
                compute_iso_face_key(face_verts, key3);
                auto iter_inserted = face_on_tetFace.try_emplace(key3, static_cast<uint32_t>(iso_faces.size()));
                if (iter_inserted.second) {
                    iso_face_index  = (iter_inserted.first)->second;
                    is_new_iso_face = true;
                }
            }
            if (is_new_iso_face) {
                auto& iso_face          = iso_faces.emplace_back();
                iso_face.vertex_indices = face_verts;
                iso_face.subface_index  = active_subface_indices[face.supporting_plane - 4 + start_index];
            }
            iso_faces[iso_face_index].headers.emplace_back(face_header);
        }
    }
}