//== INCLUDES ===================================================================


#include "Teeth_correction_helper.h"

#include <graphene/surface_mesh/algorithms/surface_mesh_tools/misc.h>
#include <graphene/geometry/registration.h>
#include <graphene/surface_mesh/algorithms/templatefitting/utility/my_helper.h>

//== NAMESPACES ================================================================


namespace graphene {
namespace teeth_correction_helper {

using namespace surface_mesh;

bool correct_teeth(character::Character &inout_character,
                   const std::string &s_bodymesh,
                   const std::string &s_teeth_up,
                   const std::string &s_teeth_down,
                   const std::string &fn_sel_mouth_region)
{

    // prepare skins ------------------------------------------------------

    const std::vector<graphene::character::Surface_mesh_skin*>& skins = inout_character.skins();

    graphene::character::Surface_mesh_skin* base_mesh       = 0;
    graphene::character::Surface_mesh_skin* teeth_up_mesh   = 0;
    graphene::character::Surface_mesh_skin* teeth_down_mesh = 0;

    for (size_t i = 0; i < skins.size(); ++i)
    {
        auto skin_id = skins[i]->mesh_property<std::string>("m:id");

        if ( skin_id[0] == s_bodymesh )
        {
            base_mesh = skins[i];
        }
        else if ( skin_id[0] == s_teeth_up )
        {
            teeth_up_mesh = skins[i];
        }
        else if ( skin_id[0] == s_teeth_down )
        {
            teeth_down_mesh = skins[i];
        }
    }

    if (!base_mesh || !teeth_up_mesh || !teeth_down_mesh)
    {
        std::cerr << "correct_teeth: [ERROR] Cannot correct teeth. No base mesh or teeth meshes found!" << std::endl;
        return false;
    }


    // prepare source undeformed ------------------------------------------

    Surface_mesh::Vertex_property<Point> undefmesh_prop =
            inout_character.get_selected_skin()->get_vertex_property<Point>("v:undef_point");

    if (!undefmesh_prop)
    {
        std::cerr << "correct_teeth: [ERROR] No undeformed mesh vertices found." << std::endl;
        return false;
    }


    // update teeth -------------------------------------------------------

    std::vector<unsigned int> mouth_vec;
    if (!graphene::surface_mesh::load_selection_from_file(mouth_vec, fn_sel_mouth_region))
    {
        std::cerr << "correct_teeth: [ERROR] Could not load selection from file: \"" << fn_sel_mouth_region << "\"." << std::endl;
        return false;
    }

    std::vector< Point > src_points;
    std::vector< Point > tar_points;
    for ( unsigned int i = 0; i < mouth_vec.size(); ++i)
    {
        Surface_mesh::Vertex v(mouth_vec[i]);
        src_points.push_back( undefmesh_prop[v] );
        tar_points.push_back( base_mesh->position(v) );
    }

    // similarity transformation
    //        const Mat4f M_teeth = graphene::geometry::registration(src_points, tar_points, graphene::geometry::CONFORMAL_REGISTRATION);
    Mat4f slightly_push_back = Mat4f::translate(Vec3f(0.0f,0.0f,-0.0009f));
    const Mat4f M_teeth = slightly_push_back *
                          graphene::geometry::registration(src_points, tar_points,
                                                           graphene::geometry::R_TRANS_SCALE_ANISO_XYZ);

    // const Mat4f M_teeth = graphene::geometry::registration(src_points, tar_points, graphene::geometry::R_TRANS_SCALE_ANISO_XYZ);
    transform_mesh(*teeth_up_mesh, M_teeth);
    teeth_up_mesh->copy_unskinned_properties();
    transform_mesh(*teeth_down_mesh, M_teeth);
    teeth_down_mesh->copy_unskinned_properties();

    // update teeth blendshape targets

    for (graphene::character::Blendshapes* blendshape_set : inout_character.blendshapes())
    {
        if (blendshape_set->base() == teeth_up_mesh ||
            blendshape_set->base() == teeth_down_mesh)
        {
            auto& targets = blendshape_set->targets();
            for (unsigned int i = 0; i < targets.size(); ++i)
            {
                transform_mesh(*targets[i], M_teeth);
            }
        }
    }

    return true;
}

//-----------------------------------------------------------------------------

bool delete_teeth_vertices(scene_graph::Character_node* character_node,
                           const std::string& s_bodymesh,
                           const std::string& s_teeth_up,
                           const std::string& fn_sel_teeth_up_vertices,
                           const std::string& fn_sel_oralcavity_indices)
{
    //
    // Setup
    //
    int teeth_up_idx = character_node->character().get_skin_idx(s_teeth_up);
    if (teeth_up_idx == -1)
    {
        printf("[ERROR] delete_teeth_vertices: Cannot find teeth up skin under name %s\n", s_teeth_up.c_str());
        return false;
    }

    int base_mesh_idx = character_node->character().get_skin_idx(s_bodymesh);
    if (base_mesh_idx == -1)
    {
        printf("[ERROR] delete_teeth_vertices: Cannot find basemesh skin under name %s\n", s_bodymesh.c_str());
        return false;
    }

    std::vector<size_t> teeth_indices;
    if (!graphene::surface_mesh::load_selection_from_file(teeth_indices, fn_sel_teeth_up_vertices))
    {
        printf("[ERROR] delete_teeth_vertices: Cannot read teeth_up selection from %s\n", fn_sel_teeth_up_vertices.c_str());
        return false;
    }

    std::vector<size_t> oralcavity_indices;
    if (!graphene::surface_mesh::load_selection_from_file(oralcavity_indices, fn_sel_oralcavity_indices))
    {
        printf("[ERROR] delete_teeth_vertices: Cannot read oralcavity selection from %s\n", fn_sel_oralcavity_indices.c_str());
        return false;
    }

    //
    // Select, delete and reset selections
    //
    character_node->set_selected_skin(teeth_up_idx);
    character_node->set_selection(teeth_indices);
    character_node->delete_selected();
    character_node->clear_selection(teeth_indices);

    character_node->set_selected_skin(base_mesh_idx);
    character_node->set_selection(oralcavity_indices);
    character_node->delete_selected();
    character_node->clear_selection(oralcavity_indices);

    // NOTE(swenninger): Basemesh stays selected skin here

    return true;
}

//-----------------------------------------------------------------------------

bool squash_teeth_up_vertex_ring(character::Character& character,
                                 const std::string& s_teeth_up,
                                 const std::string& fn_sel_teeth_up_vertices_to_squash)
{

    int teeth_up_idx = character.get_skin_idx(s_teeth_up);
    if (teeth_up_idx < 0)
    {
        std::cerr << "squash_teeth_up_vertex_ring: [ERROR] Cannot get mesh " << s_teeth_up << std::endl;
        return false;
    }

    graphene::character::Surface_mesh_skin* teeth_up = character.skins()[teeth_up_idx];

    std::vector<unsigned int> indices_to_squash;
    if (!graphene::surface_mesh::load_selection_from_file(indices_to_squash, fn_sel_teeth_up_vertices_to_squash))
    {
        std::cerr << "squash_teeth_up_vertex_ring: [ERROR] Could not load selection from file: \""
                  << fn_sel_teeth_up_vertices_to_squash << "\"." << std::endl;
        return false;
    }

    if (indices_to_squash.size() % 2 != 0)
    {
        std::cerr << "squash_teeth_up_vertex_ring: [ERROR] Uneven # of vertices to squash: "
                  << indices_to_squash.size() << std::endl;
        return false;
    }

    for (int i = 0; i < indices_to_squash.size(); i += 2)
    {
        Surface_mesh::Vertex v1(indices_to_squash[i]);
        Surface_mesh::Vertex v2(indices_to_squash[i + 1]);

        Vec3f p1 = teeth_up->position(v1);
        Vec3f p2 = teeth_up->position(v2);

        teeth_up->position(v1) += 0.9f * (p2 - p1);
    }

    teeth_up->update_vertex_normals();

    return true;
}

//=============================================================================
} // namespace surface_mesh
} // namespace graphene
//=============================================================================

