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


#include "Energy_term_landmarks_skinned_multiple_scans.h"
#include <graphene/surface_mesh/algorithms/templatefitting/utility/my_helper.h>
#include <graphene/geometry/Blend_skinning_matrices.h>

#include "../settings.h"


//== NAMESPACE ================================================================


namespace graphene {
namespace surface_mesh {


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


Energy_term_landmarks_skinned_multiple_scans::
Energy_term_landmarks_skinned_multiple_scans(const Energy_term_data& energy_term_data, double weight) :
    Energy_term(energy_term_data, weight, "Landmarks Skinned Multiple Scans")
{
}


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


Energy_term_landmarks_skinned_multiple_scans::
~Energy_term_landmarks_skinned_multiple_scans()
{}


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


void
Energy_term_landmarks_skinned_multiple_scans::
pre_minimize_action()
{
    if (!energy_term_data_.character)
    {
        std::cerr << "[ERROR] - Energy_term_landmarks_skinned_multiple_scans::!energy_term_data_.character - EXIT!" << std::endl;
        exit(1);
    }

    Surface_mesh::Vertex_property<Vec4f> vdepends = energy_term_data_.template_mesh->get_vertex_property<Vec4f>("v:skin_depend");
    Surface_mesh::Vertex_property<Vec4f> vweights = energy_term_data_.template_mesh->get_vertex_property<Vec4f>("v:skin_weight");
    Surface_mesh::Vertex_property<Vec4f> vdepends2 = energy_term_data_.template_mesh->get_vertex_property<Vec4f>("v:skin_depend2");
    Surface_mesh::Vertex_property<Vec4f> vweights2 = energy_term_data_.template_mesh->get_vertex_property<Vec4f>("v:skin_weight2");

    if (!(vdepends && vweights && vdepends2 && vweights2))
    {
        std::cerr << "[ERROR] - Energy_term_landmarks_skinned_multiple_scans::pre_minimize_action() - EXIT!" << std::endl;
        exit(1);
    }

    const auto& fitting_targets = energy_term_data_.fitting_targets;
    size_t num_fitting_targets = fitting_targets.size();
    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        const std::vector<Mat4f>* pose_vector = fitting_targets[i].pose_vector;
        if (!pose_vector)
        {
            std::cerr << "[ERROR] - Energy_term_landmarks_skinned_multiple_scans::pre_minimize_action() - pose_vector not available" << std::endl;
            exit(1);
        }
    }
}


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


double
Energy_term_landmarks_skinned_multiple_scans::
evaluate()
{
    const Landmarks_manager& landmarks_manager = energy_term_data_.landmarks_manager;
    const Surface_mesh* template_mesh = energy_term_data_.template_mesh;

    Surface_mesh::Vertex_property<Vec4f> vdepends  = template_mesh->get_vertex_property<Vec4f>("v:skin_depend");
    Surface_mesh::Vertex_property<Vec4f> vweights  = template_mesh->get_vertex_property<Vec4f>("v:skin_weight");
    Surface_mesh::Vertex_property<Vec4f> vdepends2 = template_mesh->get_vertex_property<Vec4f>("v:skin_depend2");
    Surface_mesh::Vertex_property<Vec4f> vweights2 = template_mesh->get_vertex_property<Vec4f>("v:skin_weight2");

    const auto& fitting_targets = energy_term_data_.fitting_targets;
    size_t num_fitting_targets  = fitting_targets.size();

    // Check num landmarks
    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        if (!landmarks_manager.check_dimension(fitting_targets[i].point_set, template_mesh, fitting_targets[i].landmark_property_name))
        {
            std::cerr << "Energy_term_landmarks_skinned_multiple_scans::evaluate(): [ERROR] Number of landmarks do not match for front_body and/or back_body." << std::endl;
            return DBL_MAX;
        }
    }

    // Gather data per fitting target
    std::vector<Per_fitting_target_data> per_fitting_target_data;

    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        Per_fitting_target_data temp;

        per_fitting_target_data.push_back(temp);

        auto& data = per_fitting_target_data.back();
        auto& fitting_target = fitting_targets[i];

        // Landmark position and vertices on pointset and template
        landmarks_manager.get_landmarks_point_set_stdvec(fitting_target.point_set, data.lm_ps);
        landmarks_manager.get_landmarks_template_stdvec(template_mesh, data.lm_templ, fitting_target.landmark_property_name);
        landmarks_manager.get_indices_template(template_mesh, data.lm_templ_vertices, fitting_target.landmark_property_name);

        // Landmark weight
        data.local_landmark_weight = template_mesh->get_vertex_property<double>(fitting_target.landmark_weight_property_name);

        // Set pose and get skinning matrices
        surface_mesh::set_pose_vector(*fitting_target.pose_vector, energy_term_data_.character->skeleton());
        data.skinning_tpose_to_scanpose = energy_term_data_.character->skeleton().skinning_matrices_;
        if (data.skinning_tpose_to_scanpose.size() == 0)
        {
            std::cerr << "[ERROR] - Energy_term_landmarks_skinned_multiple_scans::evaluate() - skinning_tpose_to_scanpose.size(): " << data.skinning_tpose_to_scanpose.size() << std::endl;
            exit(1);
        }
    }

    double error_match_landmarks = 0.0;

    unsigned int L_total = 0;
    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        L_total += per_fitting_target_data[i].lm_ps.size();
    }

    // Get error per fitting target
    for (size_t fitting_target_index = 0;
         fitting_target_index < num_fitting_targets;
         ++fitting_target_index)
    {
        auto& data = per_fitting_target_data[fitting_target_index];
        unsigned int num_landmarks = data.lm_ps.size();

        for (unsigned int i = 0; i < num_landmarks; ++i)
        {
            auto v = data.lm_templ_vertices[i];

            double local_weight = 1.0;
            if (data.local_landmark_weight)
            {
                local_weight = data.local_landmark_weight[v];
                if (local_weight < 1e-5)
                    continue;
                local_weight = sqrt(local_weight);
            }

            Mat4f blend_mat = blend_matrices(vweights[v], vweights2[v],
                                             vdepends[v], vdepends2[v],
                                             data.skinning_tpose_to_scanpose);

            Point the_point = affine_transform(blend_mat, data.lm_templ[i]);

            // Distance between skinned point and pointset
            double u = norm( data.lm_ps[i] - the_point );
            if ( robustness_parameters_.solver_type == SOLVER_TYPE_L2 )
            {
                error_match_landmarks += u * u * local_weight;
            }
            else if ( robustness_parameters_.solver_type == SOLVER_TYPE_L1 )
            {
                error_match_landmarks += std::pow( u, robustness_parameters_.parameter ) * local_weight;
            }
            else if ( robustness_parameters_.solver_type == SOLVER_TYPE_HUBER )
            {
                std::cerr << "ERSTMAL NICHT " << std::endl; exit(1);
                double huber_thr = robustness_parameters_.parameter;
                if ( u > huber_thr )
                {
                    error_match_landmarks += ( huber_thr*u - huber_thr*huber_thr/2.0 ) * local_weight;
                }
                else
                {
                    error_match_landmarks += ( u*u/2.0 ) * local_weight;
                }
            }
            else if ( robustness_parameters_.solver_type == SOLVER_TYPE_GEMANMCCLURE )
            {
                std::cerr << "ERSTMAL NICHT " << std::endl; exit(1);
                double scale_parameter = robustness_parameters_.parameter;
    //                error_match_landmarks += ( u*u / (scale_parameter*scale_parameter + u*u) ); // VERSION 2
                error_match_landmarks += ( ( u*u / ( 2.0 * ( 1.0 + (u*u/(scale_parameter*scale_parameter)) ) ) ) ) * local_weight; // VERSION 3
            }
            else // default to L2
            {
                error_match_landmarks += u * u * local_weight;
            }
        }
    }

    error_match_landmarks *= (weight_ / L_total);

    return error_match_landmarks;
}


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


void
Energy_term_landmarks_skinned_multiple_scans::
add_rows_big_lse(std::vector<Tripl>& coeffs,
                 Eigen::VectorXd& b,
                 unsigned int& r)
{
    //std::cout << "[INFO] Energy_term_landmarks_skinned_multiple_scans::add_rows_big_lse(): weight_: " << weight_ << std::endl;

    const Landmarks_manager& landmarks_manager = energy_term_data_.landmarks_manager;
    const Column_offsets&  column_offsets = energy_term_data_.column_offsets;
    const Surface_mesh* template_mesh = energy_term_data_.template_mesh;

    Surface_mesh::Vertex_property<Vec4f> vdepends  = template_mesh->get_vertex_property<Vec4f>("v:skin_depend");
    Surface_mesh::Vertex_property<Vec4f> vweights  = template_mesh->get_vertex_property<Vec4f>("v:skin_weight");
    Surface_mesh::Vertex_property<Vec4f> vdepends2 = template_mesh->get_vertex_property<Vec4f>("v:skin_depend2");
    Surface_mesh::Vertex_property<Vec4f> vweights2 = template_mesh->get_vertex_property<Vec4f>("v:skin_weight2");

    const auto& fitting_targets = energy_term_data_.fitting_targets;
    size_t num_fitting_targets  = fitting_targets.size();

    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        if (!landmarks_manager.check_dimension(fitting_targets[i].point_set, template_mesh, fitting_targets[i].landmark_property_name))
        {
            std::cerr << "Energy_term_landmarks_skinned_multiple_scans::evaluate(): [ERROR] Number of landmarks does not match." << std::endl;
            return;
        }
    }

    // Gather per fitting target data
    std::vector<Per_fitting_target_data> per_fitting_target_data;

    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        Per_fitting_target_data temp;
        per_fitting_target_data.push_back(temp);

        auto& data = per_fitting_target_data.back();

        // Landmark position and vertices on pointset and template
        landmarks_manager.get_indices_template(template_mesh, data.lm_templ_vertices, fitting_targets[i].landmark_property_name);
        landmarks_manager.get_landmarks_point_set_stdvec(fitting_targets[i].point_set, data.lm_ps);

        // Landmark weight
        data.local_landmark_weight = template_mesh->get_vertex_property<double>(fitting_targets[i].landmark_weight_property_name);

        // Skinning matrices
        surface_mesh::set_pose_vector(*fitting_targets[i].pose_vector, energy_term_data_.character->skeleton());
        data.skinning_tpose_to_scanpose = energy_term_data_.character->skeleton().skinning_matrices_;
        if (data.skinning_tpose_to_scanpose.size() == 0)
        {
            std::cerr << "[ERROR] - Energy_term_landmarks_skinned_multiple_scans::evaluate() - skinning_tpose_to_scanpose.size(): " << data.skinning_tpose_to_scanpose.size() << std::endl;
            exit(1);
        }
    }

    double local_weight = 1.0;

    unsigned int L_total = 0;
    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        L_total += per_fitting_target_data[i].lm_ps.size();
    }

    const double weight_match_landmarks = sqrt(weight_ / L_total);

    for (size_t fitting_target_index = 0;
         fitting_target_index < num_fitting_targets;
         ++fitting_target_index)
    {
        auto& data = per_fitting_target_data[fitting_target_index];
        size_t num_landmarks = data.lm_ps.size();

        for (unsigned int i = 0; i < num_landmarks; ++i)
        {
            auto v = data.lm_templ_vertices[i];
            const int idx = v.idx();
            const Point t = data.lm_ps[i];

            if (data.local_landmark_weight)
            {
                local_weight = data.local_landmark_weight[v];
                if (local_weight < 1e-5)
                {
                    continue;
                }
                local_weight = sqrt(local_weight);
            }

            Mat4f blend_mat = blend_matrices(vweights[v], vweights2[v],
                                             vdepends[v], vdepends2[v],
                                             data.skinning_tpose_to_scanpose);

            // right hand side
            b(r) = local_weight * weight_match_landmarks * (t[0] - blend_mat(0, 3));

            // begin_row
            coeffs.push_back( Tripl(r, column_offsets.offset_match + 3*idx + 0, local_weight * weight_match_landmarks * blend_mat(0, 0)) );
            coeffs.push_back( Tripl(r, column_offsets.offset_match + 3*idx + 1, local_weight * weight_match_landmarks * blend_mat(0, 1)) );
            coeffs.push_back( Tripl(r, column_offsets.offset_match + 3*idx + 2, local_weight * weight_match_landmarks * blend_mat(0, 2)) );
            ++r;

            // right hand side
            b(r) = local_weight * weight_match_landmarks * (t[1] - blend_mat(1, 3));

            // begin_row
            coeffs.push_back( Tripl(r, column_offsets.offset_match + 3*idx + 0, local_weight * weight_match_landmarks * blend_mat(1, 0)) );
            coeffs.push_back( Tripl(r, column_offsets.offset_match + 3*idx + 1, local_weight * weight_match_landmarks * blend_mat(1, 1)) );
            coeffs.push_back( Tripl(r, column_offsets.offset_match + 3*idx + 2, local_weight * weight_match_landmarks * blend_mat(1, 2)) );
            ++r;

            // right hand side
            b(r) = local_weight * weight_match_landmarks * (t[2] - blend_mat(2, 3));

            // begin_row
            coeffs.push_back( Tripl(r, column_offsets.offset_match + 3*idx + 0, local_weight * weight_match_landmarks * blend_mat(2, 0)) );
            coeffs.push_back( Tripl(r, column_offsets.offset_match + 3*idx + 1, local_weight * weight_match_landmarks * blend_mat(2, 1)) );
            coeffs.push_back( Tripl(r, column_offsets.offset_match + 3*idx + 2, local_weight * weight_match_landmarks * blend_mat(2, 2)) );
            ++r;
        }
    }
}


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


unsigned int
Energy_term_landmarks_skinned_multiple_scans::
n_rows_big_sle()
{
    //std::cout << "[INFO] Energy_term_landmarks_skinned_multiple_scans::n_rows_big_sle()" << std::endl;

    unsigned int c = 0;

    const Surface_mesh* template_mesh = energy_term_data_.template_mesh;

    const auto& fitting_targets = energy_term_data_.fitting_targets;
    size_t num_fitting_targets  = fitting_targets.size();

    // Count valid landmarks
    for (size_t fitting_target_index = 0;
         fitting_target_index < num_fitting_targets;
         ++fitting_target_index)
    {
        const auto& fitting_target = fitting_targets[fitting_target_index];
        auto mesh_property_template_landmarks = template_mesh->get_mesh_property<std::vector<Surface_mesh::Vertex>>(fitting_target.landmark_property_name);
        if (mesh_property_template_landmarks)
        {
            auto template_landmarks    = mesh_property_template_landmarks[0];
            auto local_landmark_weight = template_mesh->get_vertex_property<double>(fitting_target.landmark_weight_property_name);

            if (local_landmark_weight)
            {
                for (size_t i = 0; i < template_landmarks.size(); ++i)
                {
                    if (local_landmark_weight[template_landmarks[i]] >= 1e-5)
                    {
                        ++c;
                    }
                }
            }
            else
            {
                c += template_landmarks.size();
            }
        }
    }

    return c * 3;
}


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