#include <prim_gen.hpp>
#include <helper.hpp>
#include <algorithm/find_mismatch_permutation.hpp>

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

void contiguous_hash(size_t& seed, uint32_t value)
{
    seed = seed ^ (std::hash<uint32_t>()(value) + 0x9e3779b9 + (seed << 6) + (seed >> 2));
}

void contiguous_hash(size_t& seed, size_t value)
{
    seed = seed ^ (std::hash<size_t>()(value) + 0x9e3779b9 + (seed << 6) + (seed >> 2));
}

void extract_iso_mesh(const std::array<uint32_t, 3>&              tet_active_subface_counts,
                      const std::vector<arrangement_t>&           tetrahedron_arrangements,
                      const scene_bg_mesh_info_t&                 scene_bg_mesh_info,
                      const std::vector<std::array<uint32_t, 4>>& tetrahedrons,
                      const std::vector<uint32_t>&                tetrahedron_active_subface_start_index,
                      const std::vector<uint32_t>&                active_subface_indices,
                      const Eigen::Ref<const Eigen::MatrixXd>     vertex_infos,
                      const flat_hash_map<uint32_t, uint32_t>&    vertex_indices_mapping,
                      std::vector<Eigen::Vector3d>&               iso_pts,
                      std::vector<iso_vertex_t>&                  iso_verts,
                      std::vector<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<uint32_t, uint32_t> vert_on_tetVert{};
    flat_hash_map<size_t, uint32_t>   vert_on_tetEdge{};
    vert_on_tetEdge.reserve(num_1_func + 3 * num_2_func + num_more_func);
    flat_hash_map<size_t, 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<pod_key_t<3>, uint32_t> face_on_tetFace{};

    //
    flat_hash_set<uint32_t> tet_iso_verts{};
    flat_hash_set<uint32_t> tet_iso_faces{};
    // map: local vert index --> iso-vert index
    std::vector<uint32_t>   iso_vId_of_vert{};
    std::vector<uint32_t>   face_verts{};
    face_verts.reserve(4);
    pod_key_t<3>            key3;
    pod_key_t<5>            key5;
    std::array<bool, 4>     used_pId;
    std::array<uint32_t, 2> vIds2;
    std::array<uint32_t, 3> vIds3;
    std::array<uint32_t, 3> implicit_pIds;
    std::array<uint32_t, 3> boundary_pIds;
    std::array<uint32_t, 4> not_boundary_vIds{};

    //
    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);
                    tet_iso_verts.insert(face.vertices.begin(), face.vertices.end());
                }
            }
        } 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];
                // plane 0,1,2,3 are tet boundaries
                auto        implicit_plane_iter = std::find_if(arrangement.unique_planes[uId].begin(),
                                                        arrangement.unique_planes[uId].end(),
                                                        [](uint32_t plane_index) { return plane_index > 3; });
                if (implicit_plane_iter != arrangement.unique_planes[uId].end()) {
                    tet_iso_faces.emplace(*implicit_plane_iter);
                    tet_iso_verts.insert(face.vertices.begin(), face.vertices.end());
                }
            }
        }

        iso_vId_of_vert.clear();
        iso_vId_of_vert.resize(vertices.size());
        const auto& tet_vertices = tetrahedrons[tet_index];
        // 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[vertex_index];

            auto boundary_plane_end = boundary_pIds.begin();
            auto impl_plane_end     = implicit_pIds.begin();
            for (const auto& vertex_plane : vertex) {
                if (vertex_plane >= 4) { // plane 0,1,2,3 are tet boundaries
                    const auto impl_plane = active_subface_indices[vertex_plane - 4 + start_index];
                    *impl_plane_end       = impl_plane;
                    ++impl_plane_end;
                } else {
                    *boundary_plane_end = vertex_plane;
                    ++boundary_plane_end;
                }
            }
            not_boundary_vIds.fill(invalid_index);
            auto end_iter = algorithm::find_mismatch_permutation(boundary_pIds.begin(),
                                                                 boundary_plane_end,
                                                                 size_t{4},
                                                                 not_boundary_vIds.begin());
            std::transform(not_boundary_vIds.begin(), end_iter, not_boundary_vIds.begin(), [&](uint32_t local_index) {
                return tet_vertices[local_index];
            });
            size_t hash_key =
                XXH3_64bits(implicit_pIds.data(), sizeof(uint32_t) * std::distance(implicit_pIds.begin(), impl_plane_end));
            for (auto iter = not_boundary_vIds.begin(); iter != end_iter - 1; ++iter) {
                assert(*iter < *(iter + 1));
                contiguous_hash(hash_key, *iter);
            }
            contiguous_hash(hash_key, *(end_iter - 1));

            uint32_t              num_boundary_planes = std::distance(boundary_pIds.begin(), boundary_plane_end);
            const vertex_header_t vertex_header{tet_index, vertex_index, 4 - num_boundary_planes};
            auto                  mapped_not_boundary_vIds = not_boundary_vIds;
            std::transform(mapped_not_boundary_vIds.begin(),
                           mapped_not_boundary_vIds.begin() + (end_iter - not_boundary_vIds.begin()),
                           mapped_not_boundary_vIds.begin(),
                           [&](uint32_t vId) { return vertex_indices_mapping.at(vId); });

            auto generate_new_iso_vert = [&](auto&& vert_hash, auto&& f) {
                const auto& [_, iso_vert_index] = *vert_hash.lazy_emplace(hash_key, [&](auto&& ctor) {
                    ctor(hash_key, static_cast<uint32_t>(iso_verts.size()));

                    auto& iso_vert                  = iso_verts.emplace_back();
                    iso_vert.header                 = vertex_header;
                    iso_vert.simplex_vertex_indices = not_boundary_vIds;

                    f();
                });
                local_vert_to_iso_vert          = iso_vert_index;
            };

            if (num_boundary_planes == 0) { // in tet cell
                local_vert_to_iso_vert          = static_cast<uint32_t>(iso_verts.size());
                auto& iso_vert                  = iso_verts.emplace_back();
                iso_vert.header                 = vertex_header;
                iso_vert.simplex_vertex_indices = not_boundary_vIds;

                const auto &impl0 = vertex_infos.col(implicit_pIds[0]), &impl1 = vertex_infos.col(implicit_pIds[1]),
                           &impl2   = vertex_infos.col(implicit_pIds[2]);
                const std::array f1 = {impl0[mapped_not_boundary_vIds[0]],
                                       impl0[mapped_not_boundary_vIds[1]],
                                       impl0[mapped_not_boundary_vIds[2]],
                                       impl0[mapped_not_boundary_vIds[3]]},
                                 f2 = {impl1[mapped_not_boundary_vIds[0]],
                                       impl1[mapped_not_boundary_vIds[1]],
                                       impl1[mapped_not_boundary_vIds[2]],
                                       impl1[mapped_not_boundary_vIds[3]]},
                                 f3 = {impl2[mapped_not_boundary_vIds[0]],
                                       impl2[mapped_not_boundary_vIds[1]],
                                       impl2[mapped_not_boundary_vIds[2]],
                                       impl2[mapped_not_boundary_vIds[3]]};
                const auto coord    = compute_barycentric_coords(f1, f2, f3);
                iso_pts.emplace_back(coord[0] * scene_bg_mesh_info.get_vert_pos(not_boundary_vIds[0])
                                     + coord[1] * scene_bg_mesh_info.get_vert_pos(not_boundary_vIds[1])
                                     + coord[2] * scene_bg_mesh_info.get_vert_pos(not_boundary_vIds[2])
                                     + coord[3] * scene_bg_mesh_info.get_vert_pos(not_boundary_vIds[3]));
            } else if (num_boundary_planes == 1) { // on tet face
                generate_new_iso_vert(vert_on_tetFace, [&] {
                    const auto &     impl0 = vertex_infos.col(implicit_pIds[0]), &impl1 = vertex_infos.col(implicit_pIds[1]);
                    const std::array f1 = {impl0[mapped_not_boundary_vIds[0]],
                                           impl0[mapped_not_boundary_vIds[1]],
                                           impl0[mapped_not_boundary_vIds[2]]},
                                     f2 = {impl1[mapped_not_boundary_vIds[0]],
                                           impl1[mapped_not_boundary_vIds[1]],
                                           impl1[mapped_not_boundary_vIds[2]]};
                    const auto coord    = compute_barycentric_coords(f1, f2);
                    iso_pts.emplace_back(coord[0] * scene_bg_mesh_info.get_vert_pos(not_boundary_vIds[0])
                                         + coord[1] * scene_bg_mesh_info.get_vert_pos(not_boundary_vIds[1])
                                         + coord[2] * scene_bg_mesh_info.get_vert_pos(not_boundary_vIds[2]));
                });
            } else if (num_boundary_planes == 2) { // on tet edge
                generate_new_iso_vert(vert_on_tetEdge, [&] {
                    const auto f1    = vertex_infos(vertex_indices_mapping.at(not_boundary_vIds[0]), implicit_pIds[0]);
                    const auto f2    = vertex_infos(vertex_indices_mapping.at(not_boundary_vIds[1]), implicit_pIds[0]);
                    const auto coord = compute_barycentric_coords(f1, f2);
                    iso_pts.emplace_back(coord[0] * scene_bg_mesh_info.get_vert_pos(not_boundary_vIds[0])
                                         + coord[1] * scene_bg_mesh_info.get_vert_pos(not_boundary_vIds[1]));
                });
            } else { // on tet vertex
                generate_new_iso_vert(vert_on_tetVert,
                                      [&] { iso_pts.emplace_back(scene_bg_mesh_info.get_vert_pos(not_boundary_vIds[0])); });
            }
        }
        // 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{true};
            if (face.negative_cell == invalid_index) {
                compute_iso_face_key(face_verts, key3);
                const auto& [_, iso_face_index] = *face_on_tetFace.lazy_emplace(key3, [&](auto&& ctor) {
                    ctor(key3, static_cast<uint32_t>(iso_faces.size()));

                    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);
            } else { // face not on tet boundary
                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_face.headers.emplace_back(face_header);
            }
        }
    }

    iso_pts.shrink_to_fit();
    iso_verts.shrink_to_fit();
    iso_faces.shrink_to_fit();
}