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


#include "Energy_term_pca_landmarks_skinned_multiple_scans.h"
#include <graphene/surface_mesh/algorithms/templatefitting/utility/my_helper.h>

#include "../settings.h"
#include "../utility/Energy_matrices.h"


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


namespace graphene {
namespace surface_mesh {


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


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


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


Energy_term_pca_landmarks_skinned_multiple_scans::
~Energy_term_pca_landmarks_skinned_multiple_scans()
{}


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


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

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

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


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


double
Energy_term_pca_landmarks_skinned_multiple_scans::
evaluate()
{
    const Landmarks_manager& landmarks_manager = energy_term_data_.landmarks_manager;
    const PCA& pca = energy_term_data_.pca;
    const Surface_mesh* template_mesh = energy_term_data_.template_mesh;
    auto* character = energy_term_data_.character;

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

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

    // 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();

        // Pointset landmarks
        const auto* point_set = fitting_targets[i].point_set;
        data.lm_ps.resize(point_set->landmarks_.size());
        for (size_t j = 0; j < data.lm_ps.size(); ++j)
        {
            data.lm_ps[j] = point_set->points_[point_set->landmarks_[j]];
        }

        // Template landmarks
        auto lm_prop = template_mesh->get_mesh_property<std::vector<Surface_mesh::Vertex>>(fitting_targets[i].landmark_property_name);
        if (!lm_prop)
        {
            std::cerr << "Energy_term_pca_landmarks_skinned_multiple_scans::evaluate(): [ERROR] lm_prop not available." << std::endl;
            return DBL_MAX;
        }
        data.lm_templ_vertices = lm_prop[0];

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

        // check dimensions
        if (data.lm_ps.size() != data.lm_templ_vertices.size())
        {
            std::cerr << "Energy_term_pca_landmarks_skinned_multiple_scans::evaluate(): [ERROR] Number of landmarks does not match." << std::endl;
            return DBL_MAX;
        }
    }

    double error_reg_pca_ff = 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();
    }

    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];

        // Set skeleton to current pose for fitting target in order
        // to correctly retrieve skinning matrices
        surface_mesh::set_pose_vector(*fitting_targets[fitting_target_index].pose_vector,
                                      character->skeleton());

        // Gather part matrices
        Eigen::MatrixXd P;
        Eigen::VectorXd mu;
        SpMat R;
        Eigen::VectorXd t;
        SpMat L;

        energy_matrix::pca_matrix(pca, P);
        energy_matrix::pca_mean(pca, mu);
        energy_matrix::skinning(*character, *template_mesh,
                                pca.mapping_full_to_reduced_,
                                R, t);
        energy_matrix::selection(data.lm_templ_vertices,
                                 template_mesh->n_vertices(),
                                 pca.mapping_full_to_reduced_,
                                 L);

        /*
            Objective:
                Landmarks of skinned pca result should equal pointset landmarks

            L * (R * (Pd + mu) + t) = target
        */

        Eigen::VectorXd resulting_landmarks = L * (R * (P * pca.d_pca_ + mu) + t);
        std::vector<Point> resulting_points;
        unsigned int num_landmarks = data.lm_ps.size();
        for (unsigned int i = 0; i < num_landmarks; ++i)
        {
            resulting_points.push_back(
                Point(resulting_landmarks[3 * i + 0],
                      resulting_landmarks[3 * i + 1],
                      resulting_landmarks[3 * i + 2])
            );
        }

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

            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);
            }

            double u = norm( resulting_points[i] - data.lm_ps[i] );
            if ( robustness_parameters_.solver_type == SOLVER_TYPE_L2 )
            {
                error_reg_pca_ff += u * u * local_weight;
            }
            else if ( robustness_parameters_.solver_type == SOLVER_TYPE_L1 )
            {
                error_reg_pca_ff += std::pow( u, robustness_parameters_.parameter) * local_weight;
            }
            else if ( robustness_parameters_.solver_type == SOLVER_TYPE_HUBER )
            {
                double huber_thr = robustness_parameters_.parameter;
                if ( u > huber_thr )
                {
                    error_reg_pca_ff += ( huber_thr*u - huber_thr*huber_thr/2.0 ) * local_weight;
                }
                else
                {
                    error_reg_pca_ff += ( u*u/2.0 ) * local_weight;
                }
            }
            else if ( robustness_parameters_.solver_type == SOLVER_TYPE_GEMANMCCLURE )
            {
                double scale_parameter = robustness_parameters_.parameter;
    //                error_reg_pca_ff += ( u*u / (scale_parameter*scale_parameter + u*u) ); // VERSION 2
                  error_reg_pca_ff += ( ( u*u / ( 2.0 * ( 1.0 + (u*u/(scale_parameter*scale_parameter)) ) ) ) ) * local_weight; // VERSION 3
            }
            else // default to L2
            {
                error_reg_pca_ff += u * u * local_weight;
            }
        }
    }

    error_reg_pca_ff *= (weight_ / L_total);

    return error_reg_pca_ff;
}


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


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

    const Column_offsets& column_offsets = energy_term_data_.column_offsets;
    const unsigned int offset_nrd_w_pca = column_offsets.offset_nrd_w_pca;
    const PCA& pca = energy_term_data_.pca;
    const Landmarks_manager& landmarks_manager = energy_term_data_.landmarks_manager;
    const Surface_mesh* template_mesh = energy_term_data_.template_mesh;
    auto* character = energy_term_data_.character;

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

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

    // 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();

        // Pointset landmarks
        const auto* point_set = fitting_targets[i].point_set;
        data.lm_ps.resize(point_set->landmarks_.size());
        for (size_t j = 0; j < data.lm_ps.size(); ++j)
        {
            data.lm_ps[j] = point_set->points_[point_set->landmarks_[j]];
        }

        // Template landmarks
        auto lm_prop = template_mesh->get_mesh_property<std::vector<Surface_mesh::Vertex>>(fitting_targets[i].landmark_property_name);
        if (!lm_prop)
        {
            std::cerr << "Energy_term_pca_landmarks_skinned_multiple_scans::evaluate(): [ERROR] lm_prop not available." << std::endl;
            return;
        }
        data.lm_templ_vertices = lm_prop[0];

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

        // check dimensions
        if (data.lm_ps.size() != data.lm_templ_vertices.size())
        {
            std::cerr << "Energy_term_pca_landmarks_skinned_multiple_scans::evaluate(): [ERROR] Number of landmarks does not match." << std::endl;
            return;
        }
    }

    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_model = sqrt(weight_ / L_total);

    //std::cout << "[INFO] Energy_term_pca_landmarks_skinned_multiple_scans::add_rows_big_lse(): weight_model: " << weight_model << std::endl;

    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];

        // Set skeleton to current pose for fitting target in order
        // to correctly retrieve skinning matrices
        surface_mesh::set_pose_vector(*fitting_targets[fitting_target_index].pose_vector,
                                      character->skeleton());

        // Gather part matrices
        Eigen::MatrixXd P;
        Eigen::VectorXd mu;
        SpMat R;
        Eigen::VectorXd t;
        SpMat L;

        energy_matrix::pca_matrix(pca, P);
        energy_matrix::pca_mean(pca, mu);
        energy_matrix::skinning(*character, *template_mesh,
                                pca.mapping_full_to_reduced_,
                                R, t);
        energy_matrix::selection(data.lm_templ_vertices,
                                 template_mesh->n_vertices(),
                                 pca.mapping_full_to_reduced_,
                                 L);

        unsigned int num_landmarks = data.lm_ps.size();

        Eigen::VectorXd target(3 * num_landmarks);
        for (unsigned int i = 0; i < num_landmarks; ++i)
        {
            const auto& p = data.lm_ps[i];
            target(3 * i + 0) = p[0];
            target(3 * i + 1) = p[1];
            target(3 * i + 2) = p[2];
        }

        /*
            Objective:
                Landmarks of skinned pca result should equal pointset landmarks

            L * (R * (Pd + mu) + t) = target
            L * (RPd + Rmu + t) = target
            LRPd + LRmu + Lt = target
            LRPd = target - LRmu - Lt

            -> lhs = LRP
            -> rhs = target - LRmu - Lt
        */

        Eigen::MatrixXd lhs = L * R * P;
        Eigen::VectorXd rhs = target - L * R * mu - L * t;

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

            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);
            }

            // right hand side
            b(r) = local_weight * weight_model * rhs[3 * i + 0];

            // begin_row
            for (unsigned int dcount = 0; dcount < pca.dim_pca_model_; ++dcount)
            {
                coeffs.push_back( Tripl(r, offset_nrd_w_pca + dcount, lhs(3 * i + 0, dcount) * local_weight * weight_model) ); // d's
            }
            ++r;

            // right hand side
            b(r) = local_weight * weight_model * rhs[3 * i + 1];

            // begin_row
            for (unsigned int dcount = 0; dcount < pca.dim_pca_model_; ++dcount)
            {
                coeffs.push_back( Tripl(r, offset_nrd_w_pca + dcount, lhs(3 * i + 1, dcount) * local_weight * weight_model) ); // d's
            }
            ++r;

            // right hand side
            b(r) = local_weight * weight_model * rhs[3 * i + 2];

            // begin_row
            for (unsigned int dcount = 0; dcount < pca.dim_pca_model_; ++dcount)
            {
                coeffs.push_back( Tripl(r, offset_nrd_w_pca + dcount, lhs(3 * i + 2, dcount) * local_weight * weight_model) ); // d's
            }
            ++r;
        }
    }
}

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


unsigned int
Energy_term_pca_landmarks_skinned_multiple_scans::
n_rows_big_sle()
{
    //std::cout << "[INFO] Energy_term_pca_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();

    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        auto lm_prop = template_mesh->get_mesh_property<std::vector<Surface_mesh::Vertex>>(fitting_targets[i].landmark_property_name);
        if (!lm_prop)
            continue;

        auto lm_templ_vertices = lm_prop[0];
        auto local_landmark_weight = template_mesh->get_vertex_property<double>(fitting_targets[i].landmark_weight_property_name);

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

    return c*3;
}


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