//===========================================================================


#include "Inverse_kinematics.h"

#include <graphene/geometry/Matrix3x3.h>

#if __APPLE__
#include "/opt/homebrew/Cellar/libomp/15.0.6/include/omp.h"
#else
#include <omp.h>
#endif

#include <iostream>
#include <cfloat>
#include <fstream>


//===========================================================================
namespace ik {
//===========================================================================


Inverse_kinematics::
Inverse_kinematics(double lambda)
    : lambda_(lambda),
      clamp_(false),
      skeleton_(NULL)
{
    map_effectors_ = std::map<std::string, std::vector<ik::Inverse_kinematics::Effector *> >();

    effectors_ = std::vector<ik::Inverse_kinematics::Effector *>();
}


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


Inverse_kinematics::
~Inverse_kinematics()
{
    clear_effectors();
}


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


bool
Inverse_kinematics::
set_targets_and_positions(const std::vector<graphene::Vec3f>& targets,
                          const std::vector<graphene::Vec3f>& positions,
                          const std::vector<double>& weights)
{
    bool weights_available = weights.size() > 0;

    if ((targets.size() != positions.size()) ||
        (weights_available && (targets.size() != weights.size())))
    {
        std::cerr << "Inverse_kinematics::set_targets_and_positions: [ERROR] #targets does not match #positions." << std::endl;
        return false;
    }

    clear_effectors();

    graphene::character::Joint* closest_ik_joint;

    for (size_t i = 0; i < positions.size(); ++i)
    {
        // find joint closest to i-th position
        closest_ik_joint = NULL;
        float closest_dist = FLT_MAX;

        for (size_t j = 0; j < skeleton_->joints_.size(); ++j)
        {
            graphene::character::Joint* ik_joint = skeleton_->joints_[j];

            const graphene::Vec3f& joint_position = ik_joint->get_global_translation();
            const float dist = norm(joint_position - positions[i]);

            if (dist < closest_dist)
            {
                closest_ik_joint = ik_joint;
                closest_dist = dist;
            }
        }

        if (closest_ik_joint != NULL)
        {
            double weight = weights_available ? weights[i] : 1.0;
            add_effector(closest_ik_joint, positions[i], targets[i], weight);
        }
    }

    return true;
}


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


bool
Inverse_kinematics::
solve_IK()
{
    // setup damped least squares system
    if (!setup())
    {
        std::cerr << "Inverse_kinematics::solve_IK: [ERROR] !setup()" << std::endl;
        return false;
    }

    // compute parameter update vector
    if (!compute_update())
    {
        std::cerr << "Inverse_kinematics::solve_IK: [ERROR] !compute_update()" << std::endl;
        return false;
    }

    // apply update
    if (!apply_update())
    {
        std::cerr << "Inverse_kinematics::solve_IK: [ERROR] !apply_update()" << std::endl;
        return false;
    }

    return true;
}


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


bool
Inverse_kinematics::
setup()
{
    // number of degrees of freedom
    const int n_dofs = current_.size();

    // how many constraints? three per effector (x,y,z)
    const int n_targets = 3 * effectors_.size();

    // if #constraints or #dofs is zero -> stop
    if (n_targets * n_dofs == 0)
    {
        std::cerr << "IK: matrix is degenerate!\n";
        return false;
    }

    // Jacobian matrix and vector of errors
    errors_   = Eigen::VectorXd::Zero(n_targets);
    jacobian_ = Eigen::MatrixXd::Zero(n_targets, n_dofs);

    // iterate over constraints
    const std::vector<ik::Inverse_kinematics::Effector*>& effectors = effectors_;

    for(int i = 0; i < effectors.size(); ++i)
    {
        const ik::Inverse_kinematics::Effector* e = effectors[i];

        const graphene::Vec3f& position = e->position_;
        const graphene::Vec3f& target   = e->target_;

        // compute errors
        errors_(3*i + 0) = e->weight_ * (target[0] - position[0]);
        errors_(3*i + 1) = e->weight_ * (target[1] - position[1]);
        errors_(3*i + 2) = e->weight_ * (target[2] - position[2]);

        // go up kinematic chain, add entries to Jacobian
        const graphene::character::Joint* joint = e->joint_;
        while (joint->parent_ != NULL)
        {
            // get all local axes defined for this joint
            const std::vector<ik::Inverse_kinematics::DoF>& joint_dofs = dofs_[joint->name_];

            // get rotation center of current joint
            const graphene::Vec3f& c = joint->get_global_translation();

            for (const auto& dof : joint_dofs)
            {
                // rotate rotation axis of joint by its global transformation
                const graphene::Vec3f axis = graphene::linear_transform(joint->global_, dof.axis_);

                // calculate gradient
                graphene::Vec3f grad = e->weight_ * cross(axis, position - c);

                // update Jacobian matrix
                jacobian_(3*i + 0, dof.index_) = grad[0];
                jacobian_(3*i + 1, dof.index_) = grad[1];
                jacobian_(3*i + 2, dof.index_) = grad[2];
            }

            joint = joint->parent_;
        }
    }

    return true;
}


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


Eigen::MatrixXd
Inverse_kinematics::
damping_matrix()
{
    const int n_dofs = current_.size();

    Eigen::MatrixXd invW = Eigen::MatrixXd::Identity(n_dofs, n_dofs);

    const double e = 1e-5;
    if (hprev_.size() != (size_t) n_dofs)
    {
        hprev_ = std::vector<double>(n_dofs, e);
    }
    for (int i = 0; i < n_dofs; ++i)
    {
        const ik::Inverse_kinematics::DoF& info = get_joint_info(i);
        if (info.index_ == -1)
        {
            std::cerr << "Inverse_kinematics::dampling_matrix: [ERROR] info.index_ == -1" << std::endl;
            exit(1);
        }

        if (!info.has_limits_)
            continue;

        double qmax = info.max_;
        double qmin = info.min_;

        double q = current_[i];

        // DEBUG
        //const float three_deg_in_radians = 3.0f * M_PI / 180.0f;
        //const float to_degrees = 180.0f / M_PI;
        //if (qmax != qmin &&
        //   qmax - qmin >= three_deg_in_radians &&
        //   (std::abs(qmax - q) < three_deg_in_radians ||
        //    std::abs(qmin - q) < three_deg_in_radians))
        //{
        //   std::cerr << "[WARNING] Joint " << info.joint_->name_ << " approaches limits on axis " << info.axis_ << std::endl;
        //   std::cerr << "[WARNING] Joint " << qmin * to_degrees << " < " << q * to_degrees << " < " << qmax * to_degrees << std::endl;
        //}

        double h = fabs( std::pow(qmax - qmin, 2) * (2 * q - qmax - qmin)
                    / ( 4 * std::pow(qmax - q, 2) * std::pow(q - qmin, 2) ) );

        double delta = h - hprev_[i];
        hprev_[i] = h;

        invW(i, i) = delta >= 0 ? (1.0 + h) : 1.0;
    }

    // apply damping
    for (int i = 0; i < jacobian_.cols(); ++i)
    {
        invW(i, i) *= lambda_;
    }

    return invW;
}


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


bool
Inverse_kinematics::
compute_update()
{
    dtheta_.setZero();

    const unsigned int n = jacobian_.cols();
    const Eigen::MatrixXd jacobianT = jacobian_.transpose();

    // joint limit avoidance
    if (clamp_)
    {
        // solve system with regularized damped pseudo-inverse:
        //  dt = (J^T * J + lambda*W)^-1 * J^T * e
        Eigen::MatrixXd DLS = jacobianT * jacobian_ + damping_matrix();
        dtheta_ = DLS.llt().solve( jacobianT * errors_ );
    }
    else // standard solution
    {
    	// solve system with damped pseudo-inverse:
        //  dt = (J^T * J + lambda*I)^-1 * J^T * e
        Eigen::MatrixXd DLS = jacobianT * jacobian_ + lambda_*Eigen::MatrixXd::Identity(n, n);
        dtheta_ = DLS.llt().solve( jacobianT * errors_ );
    }

    return true;
}


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


bool
Inverse_kinematics::
apply_update()
{
    const size_t n = dtheta_.rows();
    if (n == 0 || current_.size() == 0)
    {
        std::cerr << "Inverse_kinematics::apply_update: [ERROR] n: " << n << std::endl;
        std::cerr << "Inverse_kinematics::apply_update: [ERROR] current_.size(): " << current_.size() << std::endl;
        return false;
    }

    std::vector<double> ps(n);
    for(size_t i = 0; i < n; ++i)
        ps[i] = current_[i] + dtheta_(i);

    // update skeleton forward kinematics
    if (clamp_)
    {
        if (!set_parameters(constrain(ps)))
        {
            std::cerr << "Inverse_kinematics::apply_update: [ERROR] !set(constrain(ps))" << std::endl;
            return false;
        }
    }
    else
    {
        if (!set_parameters(ps))
        {
            std::cerr << "Inverse_kinematics::apply_update: [ERROR] !set(ps)" << std::endl;
            return false;
        }
    }

    return true;
}


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


void
Inverse_kinematics::
move_effectors()
{
    #if 0
    for (size_t i = 0; i < skeleton_->joints_.size(); ++i)
    {
        graphene::character::Joint* joint = skeleton_->joints_[i];

        if (has_effectors(joint->get_name()))
        {
            std::vector<ik::Inverse_kinematics::Effector *>& es = map_effectors_[joint->get_name()];

            const graphene::Mat4f& transf = joint->global_;

            for (size_t j = 0; j < es.size(); ++j)
            {
                const graphene::Vec3f& offset = es[j]->offset_;

                es[j]->position_ = graphene::linear_transform(transf, offset);
            }
        }
    }
    #endif
}


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


void
Inverse_kinematics::
add_effector(graphene::character::Joint* joint,
             const graphene::Vec3f& position,
             const graphene::Vec3f& target,
             const double weight)
{
    ik::Inverse_kinematics::Effector* effector = new ik::Inverse_kinematics::Effector(joint, position, target, weight);

    std::string name = effector->joint_->name_;

    map_effectors_[name].push_back(effector);
    effectors_.push_back(effector);
}


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


void
Inverse_kinematics::
clear_effectors()
{
    if (effectors_.size() == 0)
    {
        return;
    }

    for (size_t i = 0; i < effectors_.size(); ++i)
    {
        delete effectors_[i];
    }

    effectors_.clear();
    map_effectors_.clear();
}


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


bool
Inverse_kinematics::
has_effectors(const std::string &joint)
{
    return map_effectors_.find(joint) != map_effectors_.end();
}


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


bool
Inverse_kinematics::
set_parameters(const std::vector<double>& parameters)
{
    reset();

    if (!move(parameters))
    {
        std::cerr << "Inverse_kinematics::set: [ERROR] !move(parameters)" << std::endl;
        return false;
    }

    skeleton_->root_->update();

    // move effectors of the skeleton
    move_effectors();

    return true;
}


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


bool
Inverse_kinematics::
move(const std::vector<double>& parameters)
{
    if (parameters.size() != current_.size())
    {
        std::cerr << "Inverse_kinematics::move: [ERROR] parameters.size() != current_.size()" << std::endl;
        std::cerr << "parameters.size(): "<< parameters.size() << std::endl;
        std::cerr << "current_.size(): " << current_.size() << std::endl;
        return false;
    }

    const size_t n = current_.size();

    // separate rotation axes and pose parameters
    std::vector<double> rotateX(n, 0); // flexion
    std::vector<double> rotateZ(n, 0); // abduction
    std::vector<double> rotateY(n, 0); // twist

    for (size_t i = 0; i < n; ++i)
    {
        const ik::Inverse_kinematics::DoF& info = get_joint_info(i);
        if (info.index_ == -1)
        {
            std::cerr << "Inverse_kinematics::move: [ERROR] info.index_ == -1" << std::endl;
            return false;
        }

        const graphene::Vec3f& axis = info.axis_;
        if (axis == graphene::Vec3f(1,0,0))
            rotateX[i] = parameters[i];    // flexion
        else if (axis == graphene::Vec3f(0,1,0))
            rotateY[i] = parameters[i];    // abduction
        else if (axis == graphene::Vec3f(0,0,1))
            rotateZ[i] = parameters[i];    // twist
        else
            rotateY[i] = parameters[i];    // other
    }

    // transform joints separately
    transform_joints(rotateX); // flexion
    transform_joints(rotateZ); // abduction
    transform_joints(rotateY); // twist

    return true;
}


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


void
Inverse_kinematics::
reset()
{
    for (size_t i = 0; i < skeleton_->joints_.size(); ++i)
    {
        graphene::character::Joint* joint = skeleton_->joints_[i];

        if (joint != NULL)
        {
            joint->local_ = joint->bind_pose_local_;
        }
    }

    current_ = std::vector<double>(current_.size(), 0);
}


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


const std::vector<double>&
Inverse_kinematics::
get_current_parameters()
{
    return current_;
}


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


void
Inverse_kinematics::
transform_joints(const std::vector<double>& parameters)
{
    // transform joints specified in parameters
    for (size_t i = 0; i < parameters.size(); ++i)
    {
        const ik::Inverse_kinematics::DoF& info = get_joint_info(i);
        if (info.index_ == -1)
        {
            std::cerr << "Inverse_kinematics::transform_joints: [ERROR] info.index_ == -1" << std::endl;
            exit(1);
        }

        graphene::character::Joint* joint = info.joint_;
        if (joint == NULL)
            continue;

        const double p_rad = parameters[i];
        const double p_deg = p_rad * (180.0f/M_PI);

        joint->local_ = joint->local_ * graphene::Mat4f::rotate(info.axis_, p_deg);

        current_[i] += p_rad;
    }
}


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


const ik::Inverse_kinematics::DoF&
Inverse_kinematics::
get_joint_info(int index) const
{
    for (const auto& vec_dofs : dofs_)
    {
        for (const auto& dof : vec_dofs.second)
        {
            if (dof.index_ == index)
            {
                return dof;
            }
        }
    }

    std::cerr << "Inverse_kinematics::get_joint_info: [ERROR] Can't get joint info" << std::endl;
    exit(1);

    return invalid_dof;
}


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


std::vector<double>
Inverse_kinematics::
constrain(const std::vector<double>& param) const
{
    std::vector<double> result(param.size());

    for (size_t i = 0; i < param.size(); ++i)
    {
        const ik::Inverse_kinematics::DoF& info = get_joint_info(i);
        if (info.index_ == -1)
        {
            std::cerr << "Inverse_kinematics::constrain: [ERROR] info.index_ == -1" << std::endl;
            exit(1);
        }

        double qmax = info.max_;
        double qmin = info.min_;

        const double q = param[i];
        const double e = 0.001;

        if (q < qmin)
        {
            result[i] = qmin + e;
        }
        else if (q > qmax)
        {
            result[i] = qmax - e;
        }
        else
        {
            result[i] = q;
        }
    }

    return result;
}


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


bool
Inverse_kinematics::
load_dof(const std::string& filename)
{
    std::ifstream ifs(filename.c_str());
    if (!ifs)
    {
        std::cerr << "Inverse_kinematics::load_dof: [ERROR] Could not load DoF for IK. Filename: \"" << filename << "\"." << std::endl;
        return false;
    }

    free_joints_.clear();
    joint_limits_.clear();

    std::stringstream ss;
    while (ifs)
    {
        std::string line;
        std::getline(ifs, line);
        if (!ifs)
            break;

        if (line.find("#") == 0 || line.find("//") == 0)
            continue;

        ss.clear();
        ss.str(line);

        std::string joint_name;
        ss >> joint_name;

        free_joints_.push_back(joint_name);

        if (ss.str().size() > joint_name.size())
        {
            std::pair<graphene::Vec3f, graphene::Vec3f>& p = joint_limits_[joint_name];
            graphene::Vec3f& lmin = p.first;
            graphene::Vec3f& lmax = p.second;
            for (int i = 0; i < 3; ++i)
            {
                ss >> lmin[i];
            }
            for (int i = 0; i < 3; ++i)
            {
                ss >> lmax[i];
            }
            //convert from degree to radians
            lmax *= M_PI / 180.0f;
            lmin *= M_PI / 180.0f;
        }
    }

    ifs.close();

    return true;
}


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


bool
Inverse_kinematics::
collect_dofs()
{
    if (free_joints_.empty())
    {
        std::cerr << "Inverse_kinematics::collect_dofs: [ERROR] No free joints!" << std::endl;
        return false;
    }

    current_.clear();

    dofs_.clear();

    int dof_idx = 0;

    const graphene::Vec3f X_AXIS(1,0,0);
    const graphene::Vec3f Y_AXIS(0,1,0);
    const graphene::Vec3f Z_AXIS(0,0,1);

    graphene::Mat4f m1, m2;
    graphene::Vec3f current_joint_angles;

    for (size_t i = 0; i < free_joints_.size(); ++i)
    {
        graphene::character::Joint* joint = skeleton_->get_joint(free_joints_[i].c_str());

        m1 = joint->local_;
        m1[12] = m1[13] = m1[14] = 0.0f;
        m2 = joint->bind_pose_local_;
        m2[12] = m2[13] = m2[14] = 0.0f;

        current_joint_angles = graphene::Mat4f::to_euler_XYZ(inverse(m2) * m1);
        current_.push_back(current_joint_angles[0]);
        current_.push_back(current_joint_angles[1]);
        current_.push_back(current_joint_angles[2]);

        std::vector<ik::Inverse_kinematics::DoF>& joint_dofs = dofs_[free_joints_[i]];

        if (joint_limits_.count(free_joints_[i]) > 0) // there are joint limits
        {
            const std::pair<graphene::Vec3f, graphene::Vec3f>& limits = joint_limits_[free_joints_[i]];
            const graphene::Vec3f& min_angles = limits.first;
            const graphene::Vec3f& max_angles = limits.second;
            joint_dofs.push_back(ik::Inverse_kinematics::DoF(dof_idx++, joint, X_AXIS, min_angles[0], max_angles[0]));
            joint_dofs.push_back(ik::Inverse_kinematics::DoF(dof_idx++, joint, Y_AXIS, min_angles[1], max_angles[1]));
            joint_dofs.push_back(ik::Inverse_kinematics::DoF(dof_idx++, joint, Z_AXIS, min_angles[2], max_angles[2]));
        }
        else // or no joint limits
        {
            joint_dofs.push_back(ik::Inverse_kinematics::DoF(dof_idx++, joint, X_AXIS));
            joint_dofs.push_back(ik::Inverse_kinematics::DoF(dof_idx++, joint, Y_AXIS));
            joint_dofs.push_back(ik::Inverse_kinematics::DoF(dof_idx++, joint, Z_AXIS));
        }
    }

    //current_.resize(dof_idx, 0);

    return true;
}


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


void
Inverse_kinematics::
set_skeleton(graphene::character::Skeleton* skeleton)
{
    skeleton_ = skeleton;
}


//===========================================================================
} // namespace ik
//===========================================================================
