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


#include "Landmarks_manager.h"

#include <fstream>
#include <limits>
#include <cfloat>

#include <graphene/surface_mesh/algorithms/surface_mesh_tools/misc.h>

#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>

#include "settings.h"


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


namespace graphene {
namespace surface_mesh {


//== IMPLEMENTATION ============================================================



Landmarks_manager::
Landmarks_manager()
{
}


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


bool
Landmarks_manager::
load(Point_set *pointset, const std::string &filename)
{

    std::vector<unsigned int> indices;
    if (read(filename, indices))
    {
        pointset->landmarks_ = indices;
        pointset->selected_indices_ = indices;
        return true;
    }

    return false;
}


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


bool
Landmarks_manager::
load(Surface_mesh *template_mesh,
     const std::string &filename,
     const std::string& prop_name,
     const std::string& prop_lmvw)
{
    std::vector<unsigned int> indices;
    Surface_mesh::Vertex v;

    std::string ext;
    ext = filename.substr(filename.rfind('.') + 1);
    if (ext == "sel")
    {
        if (read(filename, indices))
        {
            Surface_mesh::Mesh_property< std::vector<Surface_mesh::Vertex> > landmarks
                    = template_mesh->mesh_property< std::vector<Surface_mesh::Vertex> >(prop_name);
            std::vector<Surface_mesh::Vertex>& landmark_vec = landmarks[0];

            landmark_vec.resize(indices.size());
            for (size_t i = 0; i < indices.size(); ++i)
            {
                v = Surface_mesh::Vertex((int)indices[i]);
                landmark_vec[i] = v;
            }

            return true;
        }
    }
    else if (ext == "lmvw")
    {
        std::vector<double> local_weights;
        if (read(filename, indices, local_weights))
        {
            Surface_mesh::Mesh_property< std::vector<Surface_mesh::Vertex> > landmarks
                = template_mesh->mesh_property< std::vector<Surface_mesh::Vertex> >(prop_name);
            std::vector<Surface_mesh::Vertex> &landmark_vec = landmarks[0];

            Surface_mesh::Vertex_property< double > landmarks_lw
                    = template_mesh->vertex_property< double >(prop_lmvw);
            for (auto v : template_mesh->vertices())
            {
                landmarks_lw[v] = 1.0;
            }

            landmark_vec.resize(indices.size());
            for (size_t i = 0; i < indices.size(); ++i)
            {
                v = Surface_mesh::Vertex((int)indices[i]);
                landmark_vec[i] = v;

                landmarks_lw[v] = local_weights[i];
            }

            return true;
        }
    }

    return false;
}


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


bool
Landmarks_manager::
load_with_mapping(Surface_mesh *template_mesh,
                  const std::string &filename,
                  const std::string& filename_mapping,
                  const std::string& prop_name,
                  const std::string& prop_lmvw)
{
    std::cout << "[INFO] load landmarks on template w.r.t. a mapping" << std::endl;

    // load mapping
    std::vector<int> mapping_vec;
    if (!graphene::surface_mesh::load_selection_from_file(mapping_vec, filename_mapping))
    {
        std::cerr << "[ERROR] While loading mapping." << std::endl;
        return false;
    }

    std::vector<unsigned int> indices;
    Surface_mesh::Vertex v;

    std::string ext;
    ext = filename.substr(filename.rfind('.') + 1);
    if (ext == "sel")
    {
        if (read(filename, indices))
        {
            Surface_mesh::Mesh_property< std::vector<Surface_mesh::Vertex> > landmarks
                    = template_mesh->mesh_property< std::vector<Surface_mesh::Vertex> >(prop_name);
            std::vector<Surface_mesh::Vertex>& landmark_vec = landmarks[0];

            landmark_vec.resize(indices.size());
            for (size_t i = 0; i < indices.size(); ++i)
            {
                const int map = mapping_vec[ indices[i] ];

                if (map == -1)
                {
                    std::cerr << "[ERROR] map == -1" << std::endl;
                    return false;
                }

                v = Surface_mesh::Vertex(map);
                landmark_vec[i] = v;
            }

            return true;
        }
    }
    else if (ext == "lmvw")
    {
        std::vector<double> local_weights;
        if (read(filename, indices, local_weights))
        {
            Surface_mesh::Mesh_property< std::vector<Surface_mesh::Vertex> > landmarks
                    = template_mesh->mesh_property< std::vector<Surface_mesh::Vertex> >(prop_name);
            std::vector<Surface_mesh::Vertex>& landmark_vec = landmarks[0];

            Surface_mesh::Vertex_property< double > landmarks_lw
                = template_mesh->vertex_property< double >(prop_lmvw);
            for (auto v : template_mesh->vertices())
            {
                landmarks_lw[v] = 1.0;
            }

            landmark_vec.resize(indices.size());
            for (size_t i = 0; i < indices.size(); ++i)
            {
                const int map = mapping_vec[ indices[i] ];

                if (map == -1)
                {
                    std::cerr << "[ERROR] map == -1" << std::endl;
                    return false;
                }

                v = Surface_mesh::Vertex(map);
                landmark_vec[i] = v;

                landmarks_lw[v] = local_weights[i];
            }

            return true;
        }
    }

    return false;
}


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


bool
Landmarks_manager::
load_landmark_descriptions(const std::string& filename,
                           const std::string& prop_name,
                           Surface_mesh* mesh)
{
    auto prop = mesh->mesh_property<std::vector<Landmark_description>>(prop_name);

    return ::load_landmark_descriptions(filename, prop[0]);
}


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


void
Landmarks_manager::
save(const Point_set *pointset, const std::string &filename)
{
    write(filename, pointset->landmarks_);
}


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


void
Landmarks_manager::
save(const Surface_mesh *template_mesh,
     const std::string &filename,
     const std::string& prop_name)
{
    std::vector<unsigned int> indices;
    Surface_mesh::Mesh_property< std::vector<Surface_mesh::Vertex> > landmarks
            = template_mesh->get_mesh_property< std::vector<Surface_mesh::Vertex> >(prop_name);

    if (landmarks)
    {
        const std::vector<Surface_mesh::Vertex> &landmark_vec = landmarks[0];
        indices.resize(landmark_vec.size());

        for (size_t i=0; i < landmark_vec.size(); ++i)
        {
            const Surface_mesh::Vertex& v = landmark_vec[i];
            indices[i] = (unsigned int)v.idx();
        }

        write(filename, indices);
    }
}


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


void
Landmarks_manager::
clear()
{
}



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

size_t
Landmarks_manager::
n_landmarks_template(const Surface_mesh *template_mesh,
                     const std::string& prop_name)
{
    auto landmarks = template_mesh->get_mesh_property< std::vector<Surface_mesh::Vertex> >(prop_name);

    if (landmarks)
    {
        const std::vector<Surface_mesh::Vertex> &lms = landmarks[0];
        return lms.size();
    }
    else
    {
        return 0;
    }
}


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


size_t
Landmarks_manager::
n_landmarks_pointset(const geometry::Point_set *point_set) const
{
    return point_set->landmarks_.size();
}


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


void
Landmarks_manager::
get_landmarks_point_set_stdvec(const Point_set* point_set, std::vector<Point>& landm_ps) const
{
//    std::cerr << "in: 'get_landmarks_point_set_stdvec(...)'" << std::endl;

    landm_ps.resize(point_set->landmarks_.size());

    unsigned int i = 0;
    for (auto v : point_set->landmarks_)
    {
        landm_ps[i] = point_set->points_[v];
        ++i;
    }
}



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


void
Landmarks_manager::
get_landmarks_template_stdvec(const Surface_mesh* template_mesh,
                              std::vector<Point>& landm_templ,
                              const std::string& prop_name) const
{
//    std::cerr << "in: 'get_landmarks_template_stdvec(...)'" << std::endl;

    auto points = template_mesh->get_vertex_property<Point>("v:point");
    auto landmarks = template_mesh->get_mesh_property<std::vector<Surface_mesh::Vertex> >(prop_name);

    if (landmarks)
    {
        const std::vector<Surface_mesh::Vertex>& lm_template = landmarks[0];

        landm_templ.resize(lm_template.size());

        unsigned int i = 0;
        for (auto v : lm_template)
        {
            landm_templ[i] = points[v];
            ++i;
        }
    }
}

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

void
Landmarks_manager::
get_landmarks_combined_stdvec(Surface_mesh* template_mesh, const Point_set* point_set,
                              std::vector<Point>& landm_tm,
                              std::vector<Point>& landm_ps,
                              const std::string& prop_name,
                              const std::string& prop_lmvw) const
{
    auto points = template_mesh->get_vertex_property<Point>("v:point");
    auto template_landmarks = template_mesh->get_mesh_property<std::vector<Surface_mesh::Vertex> >(prop_name);
    auto template_landmark_weights = template_mesh->get_vertex_property< double >(prop_lmvw);
    auto pointset_landmarks = point_set->landmarks_;

    if (!template_landmarks)
    {
        std::cerr << "[ERROR] Landmarks_manger: cannot get template landmarks from property '"<<  prop_name << "'" << std::endl;
        return;
    }

    size_t index = 0;
    for (auto v : template_landmarks[0])
    {
        if (template_landmark_weights)
        {
            if (template_landmark_weights[v] > 1e-5)
            {
                landm_tm.push_back(points[v]);
                landm_ps.push_back(point_set->points()[pointset_landmarks[index]]);
            }
        }
        else
        {
            landm_tm.push_back(points[v]);
            landm_ps.push_back(point_set->points()[pointset_landmarks[index]]);
        }

        index++;
    }


    auto lm_ears_prop = template_mesh->mesh_property<std::vector<Surface_mesh::Vertex>>("m:ears_lm");
    if (lm_ears_prop && (point_set->lm_ears_.size() > 0))
    {
        const auto lm_ears_template = lm_ears_prop[0];

        if (lm_ears_template.size() == point_set->lm_ears_.size())
        {
            for (auto lm : lm_ears_template)
                landm_tm.push_back(points[lm]);

            for (auto lm : point_set->lm_ears_)
                landm_ps.push_back(point_set->points()[lm]);
        }
    }
}


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

void
Landmarks_manager::
maybe_invalidate_landmarks(const std::string& fn_landmark_validity,
                           Surface_mesh* template_mesh,
                           Point_set* point_set,
                           const std::string& prop_name,
                           const std::string& description_prop_name)
{
    if (fn_landmark_validity.empty())
    {
        printf("[WARNING] Landmarks_manager::maybe_invalidate_landmarks: No filename given for landmark validity. All landmarks are assumed to be valid\n");
        return;
    }

    if (!check_dimension(point_set, template_mesh, prop_name))
    {
        return;
    }

    std::ifstream landmark_validity_file(fn_landmark_validity);
    if (!landmark_validity_file)
    {
        printf("[WARNING] Landmarks_manager::maybe_invalidate_landmarks: Cannot load %s. All landmarks are assumed to be valid\n",
               fn_landmark_validity.c_str());
        return;
    }

    //
    // Read validity from file
    //
    size_t num_landmarks = n_landmarks_template(template_mesh, prop_name);
    std::vector<bool> landmark_validity(num_landmarks);
    bool landmark_i_valid;
    for (size_t i = 0; i < num_landmarks; ++i)
    {
        landmark_validity_file >> landmark_i_valid;
        landmark_validity[i] = landmark_i_valid;

        if (landmark_validity_file.fail())
        {
            printf("[ERROR] Landmarks_manager::maybe_invalidate_landmarks: Could not parse landmark validity. All landmarks are assumed to be valid\n");
            landmark_validity = std::vector<bool>(num_landmarks, true);
            break;
        }
    }

    //
    // Only keep valid landmarks
    //
    auto template_landmarks = template_mesh->get_mesh_property<std::vector<Surface_mesh::Vertex> >(prop_name);
    auto pointset_landmarks = point_set->landmarks_;
    std::vector<Surface_mesh::Vertex> template_landmarks_to_keep;
    std::vector<unsigned int>         pointset_landmarks_to_keep;

    auto description_prop = template_mesh->get_mesh_property<std::vector<Landmark_description>>(description_prop_name);

    std::vector<Landmark_description> description_to_keep;

    if (description_prop)
    {
        if (description_prop[0].size() != pointset_landmarks.size())
        {
            fprintf(stderr, "[ERROR] Landmark description size does not match landmark size (%lu vs. %lu)\n",
                    description_prop[0].size(), pointset_landmarks.size());
        }
    }


    for (size_t i = 0; i < num_landmarks; ++i)
    {
        if (landmark_validity[i])
        {
            template_landmarks_to_keep.push_back(template_landmarks[0][i]);
            pointset_landmarks_to_keep.push_back(pointset_landmarks[i]);

            if (description_prop)
            {
                description_to_keep.push_back(description_prop[0][i]);
            }

        }
    }

    point_set->landmarks_ = pointset_landmarks_to_keep;
    point_set->selected_indices_ = pointset_landmarks_to_keep;
    template_landmarks[0] = template_landmarks_to_keep;

    if (description_prop)
    {
        description_prop[0] = description_to_keep;
    }
}

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

void
Landmarks_manager::
get_indices_pointset(const Point_set *point_set, std::vector<unsigned int>& indices) const
{
    indices = point_set->landmarks_;
}


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


void
Landmarks_manager::
get_indices_template(const Surface_mesh* template_mesh,
                     std::vector<Surface_mesh::Vertex>& indices,
                     const std::string& prop_name) const
{
    indices.clear();

    Surface_mesh::Mesh_property< std::vector<Surface_mesh::Vertex> > landmarks
        = template_mesh->get_mesh_property<std::vector<Surface_mesh::Vertex> >(prop_name);

    if (landmarks)
    {
        indices = landmarks[0];
    }
}


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


void
Landmarks_manager::
get_indices_template(const Surface_mesh* template_mesh,
                     std::vector<size_t>& indices,
                     const std::string& prop_name) const
{
    indices.clear();

    Surface_mesh::Mesh_property< std::vector<Surface_mesh::Vertex> > landmarks
        = template_mesh->get_mesh_property<std::vector<Surface_mesh::Vertex> >(prop_name);

    if (landmarks)
    {
        for (size_t i = 0; i < landmarks[0].size(); ++i)
        {
            indices.push_back(landmarks[0][i].idx());
        }
    }
}


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


bool
Landmarks_manager::
get_matching_landmarks(Surface_mesh* mesh, std::vector<Landmark_match>& matching_landmarks)
{
    auto body_prop = mesh->get_mesh_property<std::vector<Landmark_description>>("m:body_landmark_descriptions");
    auto face_prop = mesh->get_mesh_property<std::vector<Landmark_description>>("m:face_landmark_descriptions");

    if (!body_prop || !face_prop)
    {
        printf("[WARNING] No semantic landmark descriptions loaded for mesh\n");
        return false;
    }

    return ::get_matching_landmarks(body_prop[0], face_prop[0], matching_landmarks);
}


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


bool
Landmarks_manager::
check_dimension(const geometry::Point_set *point_set,
                const Surface_mesh *template_mesh,
                const std::string& prop_name) const
{
    Surface_mesh::Mesh_property< std::vector<Surface_mesh::Vertex> > landmarks
        = template_mesh->get_mesh_property< std::vector<Surface_mesh::Vertex> >(prop_name);

    if (landmarks)
    {
        const std::vector<Surface_mesh::Vertex>& tm_landmarks = landmarks[0];
        if (( tm_landmarks.size() == point_set->landmarks_.size() ) && ! tm_landmarks.empty())
        {
            return true;
        }
        else
        {
            std::cerr << "Landmarks_manager::check_dimension: [ERROR] Landmarks dimensions do not fit: "
                      << point_set->landmarks_.size()
                      << " on pointset vs "
                      << tm_landmarks.size()
                      << " on template."
                      << std::endl;
            return false;
        }
    }
    else
    {
        std::cerr << "Landmarks_manger::check_dimension: [ERROR] Cannot get template landmarks from: "
                  << "'" << prop_name << "'"
                  << std::endl;
        return false;
    }
}


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


bool
Landmarks_manager::
read(const std::string &filename, std::vector<unsigned int> &indices)
{
    std::string ext;
    ext = filename.substr(filename.rfind('.') + 1);
    if (ext == "sel")
    {
        return load_selection_from_file(indices, filename);
    }
    else
    {
        std::cerr << "Landmarks_manager::read: [WARNING] Cannot read landmarks from \"" << filename << "\"" << std::endl;
        return false;
    }
}


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


bool
Landmarks_manager::
read(const std::string &filename, std::vector<unsigned int> &indices, std::vector<double>& local_weights)
{
    std::string ext;
    ext = filename.substr(filename.rfind('.') + 1);
    if (ext == "lmvw")
    {
        return load_selection_from_file_with_weights(indices, local_weights, filename);
    }
    else
    {
        std::cerr << "Landmarks_manager::read: [WARNING] Cannot read landmarks with local weights from \"" << filename << "\"" << std::endl;
        return false;
    }
}


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


bool
Landmarks_manager::
write(const std::string &filename, const std::vector<unsigned int> &indices)
{
    std::string ext;
    ext = filename.substr(filename.rfind('.') + 1);
    if (ext == "sel")
    {
        return write_sel(filename, indices);
    }
    else
    {
        std::cerr << "Landmarks_manager::read: [WARNING] Cannot read landmarks from \"" << filename << "\"" << std::endl;
        return false;
    }
}


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


bool
Landmarks_manager::
write_sel(const std::string &filename, const std::vector<unsigned int> &indices)
{
    std::ofstream ofs;
    ofs.open(filename.c_str());

    if (!ofs)
        return false;


    for (size_t i=0; i < indices.size();++i)
    {
        ofs << indices[i] << "\n";
    }

    return true;
}


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