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


#include "Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans.h"

#include "../settings.h"

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

#include <graphene/geometry/bary_coord.h>

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


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


namespace graphene {
namespace surface_mesh {


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


Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans::
Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans(const Energy_term_data& energy_term_data, double weight)
: Energy_term(energy_term_data, weight, "PCA CPC PointSet-to-Mesh Skinned Multiple Scans")
{
}


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


Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans::
~Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans()
{}

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

void
Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans::
pre_minimize_action()
{
    if (!energy_term_data_.character)
    {
        std::cerr << "[ERROR] - Energy_term_pca_cpc_ps2mesh_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_cpc_ps2mesh_skinned_multiple_scans::pre_minimize_action() - EXIT!" << std::endl;
        exit(1);
    }
}

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


double
Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans::
evaluate()
{
    const auto& fitting_targets = energy_term_data_.fitting_targets;
    size_t num_fitting_targets  = fitting_targets.size();

    const PCA& pca = energy_term_data_.pca;
    const Surface_mesh* template_mesh = energy_term_data_.template_mesh;
    auto* character = energy_term_data_.character;

    auto template_points = template_mesh->get_vertex_property<Point>("v:point");

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

    double error_reg_pca_cpc_ps2mesh = 0.0;

    unsigned int C_total = 0;
    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        C_total += fitting_targets[i].correspondences.size();
    }

    for (size_t fitting_target_index = 0;
         fitting_target_index < num_fitting_targets;
         ++fitting_target_index)
    {
        auto& fitting_target = fitting_targets[fitting_target_index];

        const auto& correspondences = fitting_target.correspondences;
        unsigned int num_correspondences = correspondences.size();

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

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

        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::barycentric_interpolation(
            correspondences,
            pca.mapping_full_to_reduced_,
            *template_mesh,
            B);

        /*
            Objective:
                template correspondence of skinned pca result should equal
                pointset correspondence

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

        Eigen::VectorXd resulting_correspondences = B * (R * (P * pca.d_pca_ + mu) + t);
        std::vector<Point> resulting_points;
        resulting_points.reserve(num_correspondences);
        for (unsigned int i = 0; i < num_correspondences; ++i)
        {
            resulting_points.push_back(
                Point(resulting_correspondences[3 * i + 0],
                      resulting_correspondences[3 * i + 1],
                      resulting_correspondences[3 * i + 2])
            );
        }

        for (unsigned int i = 0; i < num_correspondences; ++i)
        {

            const auto& c = correspondences[i];

            double u = norm( resulting_points[i] - c.on_ps );
            if ( robustness_parameters_.solver_type == SOLVER_TYPE_L2 )
            {
                error_reg_pca_cpc_ps2mesh += u * u;
            }
            else
            {
                std::cerr << "[ERROR] in 'Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans::evaluate()' robust function not supported! Aborting..." << std::endl;
                exit(1);
            }
        }
    }

    error_reg_pca_cpc_ps2mesh *= (weight_ / C_total);

    return error_reg_pca_cpc_ps2mesh;
}


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


void
Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans::
add_rows_big_lse( std::vector<Tripl>& coeffs,
                  Eigen::VectorXd& b,
                  unsigned int& r)
{
    //std::cerr << "[INFO] in: Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans::add_rows_big_lse(): weight_: " << weight_ << std::endl;
    const auto& fitting_targets = energy_term_data_.fitting_targets;
    size_t num_fitting_targets  = fitting_targets.size();

    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 Surface_mesh* template_mesh = energy_term_data_.template_mesh;
    auto* character = energy_term_data_.character;

    auto template_points = template_mesh->get_vertex_property<Point>("v:point");

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

    unsigned int C_total = 0;
    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        C_total += fitting_targets[i].correspondences.size();
    }

    const double weight_model = sqrt(weight_ / C_total);

    //std::cout << "[INFO] Energy_term_pca_cpc_ps2mesh_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)
    {
        const auto& fitting_target = fitting_targets[fitting_target_index];

        const auto& correspondences = fitting_target.correspondences;
        const unsigned int num_correspondences = correspondences.size();

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

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

        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::barycentric_interpolation(
            correspondences,
            pca.mapping_full_to_reduced_,
            *template_mesh,
            B);

        Eigen::VectorXd target(3 * num_correspondences);
        for (size_t i = 0; i < num_correspondences; ++i)
        {
            const auto& c = correspondences[i];
            target(3 * i + 0) = c.on_ps[0];
            target(3 * i + 1) = c.on_ps[1];
            target(3 * i + 2) = c.on_ps[2];
        }

        /*
            Objective:
                template correspondence of skinned pca result should equal
                pointset correspondence

            B * (R * (Pd + mu) + t) = target
            BRPd + BRmu + Bt = target
            BRPd = target - BRmu - Bt

            -> lhs = BRP
            -> rhs = target - BRmu - Bt
        */

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

        for (unsigned int i = 0; i < num_correspondences; ++i)
        {
            const auto& c = correspondences[i];

            // right hand side
            b(r) = 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) * weight_model) ); // d's
            }
            ++r;

            // right hand side
            b(r) = 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) * weight_model) ); // d's
            }
            ++r;

            // right hand side
            b(r) = 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) * weight_model) ); // d's
            }
            ++r;
        }
    }
}

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


unsigned int
Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans::
n_rows_big_sle()
{
    //std::cout << "[INFO] Energy_term_pca_cpc_ps2mesh_skinned_multiple_scans::n_rows_big_sle()" << std::endl;
    const auto& fitting_targets = energy_term_data_.fitting_targets;
    size_t num_fitting_targets  = fitting_targets.size();

    unsigned int c = 0;

    for (size_t i = 0; i < num_fitting_targets; ++i)
    {
        c += fitting_targets[i].correspondences.size();
    }

    return c*3;
}

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