#include "JointPath.h"
#include "Jacobian.h"
#include "Body.h"
#include "CustomJointPathHandler.h"
#include "BodyCustomizerInterface.h"
#include <cnoid/EigenUtil>
#include <cnoid/TruncatedSVD>
#include <algorithm>
#include <numeric>

using namespace std;
using namespace cnoid;


double JointPath::numericalIkDefaultDeltaScale()
{
    return 0.9;
}

int JointPath::numericalIkDefaultMaxIterations()
{
    return 50;
}

double JointPath::numericalIkDefaultMaxIkError()
{
    return 1.0e-7;
}

double JointPath::numericalIkDefaultDampingConstant()
{
    return 1.0e-6;
}

int JointPath::numericalIkDefaultMaxStagnationCount()
{
    return 3;
}

//! \deprecated
double JointPath::numericalIkDefaultTruncateRatio()
{
    return TruncatedSVD<MatrixXd>::defaultTruncateRatio();
}

namespace cnoid {

class NumericalIK
{
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    
    bool isBestEffortIkMode;
    double deltaScale;
    int maxIterations;
    int iteration;
    double maxIkErrorSqr;
    double dampingConstantSqr;
    int maxStagnationCount;
    vector<double> q0;
    Isometry3 T0;
    vector<double> qBest;  // Best solution found during iterations (for best-effort mode)
    double minErrorSqr;     // Minimum error squared achieved (for best-effort mode)
    MatrixXd J;
    VectorXd dTask;
    VectorXd dq;
    MatrixXd JJ;
    Eigen::ColPivHouseholderQR<MatrixXd> QR;
    TruncatedSVD<MatrixXd> svd;
    std::function<double(VectorXd& out_error)> errorFunc;
    std::function<void(MatrixXd& out_Jacobian)> jacobianFunc;

    NumericalIK() {
        deltaScale = JointPath::numericalIkDefaultDeltaScale();
        maxIterations = JointPath::numericalIkDefaultMaxIterations();
        iteration = 0;
        dTask.resize(6);
        isBestEffortIkMode = false;
        double e = JointPath::numericalIkDefaultMaxIkError();
        maxIkErrorSqr = e * e;
        double d = JointPath::numericalIkDefaultDampingConstant();
        dampingConstantSqr = d * d;
        maxStagnationCount = JointPath::numericalIkDefaultMaxStagnationCount();
    }

    void resize(int numJoints){
        J.resize(dTask.size(), numJoints);
        dq.resize(numJoints);
    }
};

}


JointPath::JointPath()
{
    initialize();
}


JointPath::JointPath(Link* base, Link* end)
    : linkPath_(base, end)
{
    initialize();
    extractJoints();
}


JointPath::JointPath(Link* end)
    : linkPath_(end)
{
    initialize();
    extractJoints();
}


bool JointPath::setPath(Link* base, Link* end)
{
    remainingLinkTraverse.reset();
    if(!linkPath_.setPath(base, end)){
        return false;
    }
    initialize();
    extractJoints();
    return true;
}


void JointPath::initialize()
{
    needForwardKinematicsBeforeIK = false;
    numericalIK = nullptr;
    isIkJointLimitEnabled_ = false;
    isCustomIkDisabled_ = false;
}    


void JointPath::extractJoints()
{
    numUpwardJointConnections = 0;
    idOrderIndices_.clear();

    int n = linkPath_.size();
    if(n <= 1){
        joints_.clear();
    } else {
        int i = 0;
        if(linkPath_.isDownward(i)){
            i++;
        }
        joints_.clear();
        joints_.reserve(n); // reserve size n buffer
        int m = n - 1;
        while(i < m){
            Link* link = linkPath_[i];
            if(link->jointId() >= 0){
                if(link->hasActualJoint()){
                    joints_.push_back(link);
                    if(!linkPath_.isDownward(i)){
                        numUpwardJointConnections++;
                    }
                }
            }
            ++i;
        }
        if(linkPath_.isDownward(m-1)){
            Link* link = linkPath_[m];
            if(link->jointId() >= 0){
                if(link->hasActualJoint()){
                    joints_.push_back(link);
                }
            }
        }
    }
}


int JointPath::indexOf(const Link* link) const
{
    for(size_t i=0; i < joints_.size(); ++i){
        if(joints_[i] == link){
            return i;
        }
    }
    return -1;
}


void JointPath::updateIdOrderIndices() const
{
    if(idOrderIndices_.empty() && !joints_.empty()){
        int n = joints_.size();
        idOrderIndices_.resize(n);
        std::iota(idOrderIndices_.begin(), idOrderIndices_.end(), 0);
        std::sort(idOrderIndices_.begin(), idOrderIndices_.end(),
                  [this](int a, int b){
                      return joints_[a]->jointId() < joints_[b]->jointId();
                  });
    }
}


Link* JointPath::jointAtIdOrder(int index) const
{
    updateIdOrderIndices();
    return joints_[idOrderIndices_[index]];
}


int JointPath::jointIndexAtIdOrder(int index) const
{
    updateIdOrderIndices();
    return idOrderIndices_[index];
}


void JointPath::calcJacobian(Eigen::MatrixXd& out_J) const
{
    const int n = joints_.size();
    out_J.resize(6, n);
	
    if(n > 0){

        //! \todo compare the efficiency for the following codes
        if(false){
            setJacobian<0x3f, 0, 0>(*this, linkPath_.endLink(), out_J);

        } else {
        
            Link* targetLink = linkPath_.endLink();
		
            for(int i=0; i < n; ++i){
			
                Link* link = joints_[i];
			
                switch(link->jointType()){
				
                case Link::RevoluteJoint:
                {
                    Vector3 omega = link->R() * link->a();
                    const Vector3 arm = targetLink->p() - link->p();
                    if(!isJointDownward(i)){
                        omega = -omega;
                    }
                    out_J.col(i) << omega.cross(arm), omega;
                }
                break;
				
                case Link::PrismaticJoint:
                {
                    Vector3 dp = link->R() * link->d();
                    if(!isJointDownward(i)){
                        dp = -dp;
                    }
                    out_J.col(i) << dp, Vector3::Zero();
                }
                break;
				
                default:
                    out_J.col(i).setZero();
                }
            }
        }
    }
}


NumericalIK* JointPath::getOrCreateNumericalIK()
{
    if(!numericalIK){
        numericalIK = new NumericalIK;
    }
    return numericalIK;
}


bool JointPath::isBestEffortIkAvailable() const
{
    return true;  // JointPath always supports best-effort IK
}


bool JointPath::isBestEffortIkEnabled() const
{
    return numericalIK ? numericalIK->isBestEffortIkMode : false;
}


void JointPath::setBestEffortIkEnabled(bool on)
{
    getOrCreateNumericalIK()->isBestEffortIkMode = on;
}


int JointPath::getDOF() const
{
    return numJoints();
}


void JointPath::setNumericalIkMaxIkError(double e)
{
    getOrCreateNumericalIK()->maxIkErrorSqr = e * e;
}


void JointPath::setNumericalIkDeltaScale(double s)
{
    getOrCreateNumericalIK()->deltaScale = s;
}


void JointPath::setNumericalIkMaxIterations(int n)
{
    getOrCreateNumericalIK()->maxIterations = n;
}


void JointPath::setNumericalIkDampingConstant(double lambda)
{
    getOrCreateNumericalIK()->dampingConstantSqr = lambda * lambda;
}


void JointPath::setNumericalIkMaxStagnationCount(int n)
{
    getOrCreateNumericalIK()->maxStagnationCount = n;
}


/**
   \deprecated
   This parameter is used when SVD is used to solve a numerical IK, but the current
   implementation uses the damped least square method instead of SVD by default.
   Therefore this parameter is meaningless in the current implementation.
*/
void JointPath::setNumericalIkTruncateRatio(double r)
{
    getOrCreateNumericalIK()->svd.setTruncateRatio(r);
}

    
void JointPath::customizeTarget
(int numTargetElements,
 std::function<double(VectorXd& out_error)> errorFunc,
 std::function<void(MatrixXd& out_Jacobian)> jacobianFunc)
{
    auto nuIK = getOrCreateNumericalIK();
    nuIK->dTask.resize(numTargetElements);
    nuIK->errorFunc = errorFunc;
    nuIK->jacobianFunc = jacobianFunc;
}


JointPath& JointPath::setBaseLinkGoal(const Isometry3& T)
{
    linkPath_.baseLink()->setPosition(T);
    needForwardKinematicsBeforeIK = true;
    return *this;
}


bool JointPath::calcInverseKinematics()
{
    if(numericalIK && numericalIK->errorFunc){
        Isometry3 T; // dummy
        return calcInverseKinematics(T);
    }
    return false;
}


/**
 * Calculates inverse kinematics using numerical iteration.
 *
 * Best-effort mode behavior:
 * - Tracks and returns the solution with minimum error across all iterations
 * - Useful for singularities and joint limit configurations
 * - For high-precision exploration, set maxIkError=0 and increase maxIterations
 *   to explore more configurations and return the best one found
 */
bool JointPath::calcInverseKinematics(const Isometry3& T)
{
    const bool USE_USUAL_INVERSE_SOLUTION_FOR_6x6_NON_BEST_EFFORT_PROBLEM = false;
    const bool USE_SVD_FOR_BEST_EFFORT_IK = false;

    if(joints_.empty()){
        if(linkPath_.empty()){
            return false;
        }
        if(baseLink() == endLink()){
            endLink()->setPosition(T);
            if(endLink()->isFreeJoint() && !endLink()->isRoot()){
                endLink()->setOffsetPosition(endLink()->parent()->T().inverse() * T);
            }
            return true;
        } else {
            // \todo implement here
            return false;
        }
    }
    const int n = numJoints();

    auto nuIK = getOrCreateNumericalIK();
    
    if(!nuIK->jacobianFunc){
        nuIK->jacobianFunc = [&](MatrixXd& out_Jacobian){ setJacobian<0x3f, 0, 0>(*this, endLink(), out_Jacobian); };
    }
    nuIK->resize(n);
    
    Link* target = linkPath_.endLink();

    if(needForwardKinematicsBeforeIK){
        calcForwardKinematics();
        needForwardKinematicsBeforeIK = false;
    }

    nuIK->q0.resize(n);
    if(!nuIK->isBestEffortIkMode){
        for(int i=0; i < n; ++i){
            nuIK->q0[i] = joints_[i]->q();
        }
        nuIK->T0 = target->T();
    } else {
        // Initialize best solution tracking for best-effort mode
        nuIK->minErrorSqr = std::numeric_limits<double>::max();
        nuIK->qBest.resize(n);
    }

    double prevErrsqr = std::numeric_limits<double>::max();
    bool completed = false;
    int stagnationCount = 0;

    bool useUsualInverseSolution = false;
    if(USE_USUAL_INVERSE_SOLUTION_FOR_6x6_NON_BEST_EFFORT_PROBLEM){
        if(!nuIK->isBestEffortIkMode && (n == 6) && (nuIK->dTask.size() == 6)){
            useUsualInverseSolution = true;
        }
    }
        
    if(USE_SVD_FOR_BEST_EFFORT_IK && !useUsualInverseSolution && !nuIK->isBestEffortIkMode){
        // disable truncation
        nuIK->svd.setTruncateRatio(std::numeric_limits<double>::max());
    }

    for(nuIK->iteration = 0; nuIK->iteration < nuIK->maxIterations; ++nuIK->iteration){

        double errorSqr;
        if(nuIK->errorFunc){
            errorSqr = nuIK->errorFunc(nuIK->dTask);
        } else {
            nuIK->dTask.head<3>() = T.translation() - target->p();
            nuIK->dTask.segment<3>(3) = target->R() * omegaFromRot(target->R().transpose() * T.linear());
            errorSqr = nuIK->dTask.squaredNorm();
        }
        // Track the best solution in best-effort mode
        if(nuIK->isBestEffortIkMode && errorSqr < nuIK->minErrorSqr){
            nuIK->minErrorSqr = errorSqr;
            for(int j=0; j < n; ++j){
                nuIK->qBest[j] = joints_[j]->q();
            }
            stagnationCount = 0;  // Reset stagnation counter when best solution improves
        }

        if(errorSqr < nuIK->maxIkErrorSqr){
            completed = true;
            target->T() = T;
            break;
        }

        // Check for convergence stagnation or divergence
        double improvement = prevErrsqr - errorSqr;
        if(improvement < nuIK->maxIkErrorSqr){
            // No significant improvement or error increased
            stagnationCount++;
            if(stagnationCount >= nuIK->maxStagnationCount){
                break;
            }
        } else {
            // Good improvement, reset stagnation counter (unless already reset by best solution update)
            if(!nuIK->isBestEffortIkMode){
                stagnationCount = 0;
            }
        }
        prevErrsqr = errorSqr;

        nuIK->jacobianFunc(nuIK->J);

        if(useUsualInverseSolution){
            nuIK->dq = nuIK->QR.compute(nuIK->J).solve(nuIK->dTask);
        } else {
            if(USE_SVD_FOR_BEST_EFFORT_IK){
                nuIK->svd.compute(nuIK->J).solve(nuIK->dTask, nuIK->dq);
            } else {
                // The damped least squares (singurality robust inverse) method
                nuIK->JJ = nuIK->J * nuIK->J.transpose() + nuIK->dampingConstantSqr * MatrixXd::Identity(nuIK->J.rows(), nuIK->J.rows());
                nuIK->dq = nuIK->J.transpose() * nuIK->QR.compute(nuIK->JJ).solve(nuIK->dTask);
            }
        }

        // Update joint angles
        for(int j=0; j < n; ++j){
            joints_[j]->q() += nuIK->deltaScale * nuIK->dq(j);
        }

        // Apply joint limits
        if(isIkJointLimitEnabled_){
            for(int j=0; j < n; ++j){
                auto joint = joints_[j];
                double& q = joint->q();
                if(q > joint->q_upper()){
                    q = joint->q_upper();
                } else if(q < joint->q_lower()){
                    q = joint->q_lower();
                }
            }
        }

        calcForwardKinematics();
    }

    if(!completed){
        if(nuIK->isBestEffortIkMode){
            // Restore the best solution found during iterations
            for(int i=0; i < n; ++i){
                joints_[i]->q() = nuIK->qBest[i];
            }
            calcForwardKinematics();
        } else {
            // Normal mode: restore initial state on failure
            for(int i=0; i < n; ++i){
                joints_[i]->q() = nuIK->q0[i];
            }
            calcForwardKinematics();
            target->T() = nuIK->T0;
        }
    }

    return completed;
}


bool JointPath::calcRemainingPartForwardKinematicsForInverseKinematics()
{
    if(!remainingLinkTraverse){
        remainingLinkTraverse = make_shared<LinkTraverse>();
        Link* remainingBaseLink = nullptr;
        if(empty()){
            remainingLinkTraverse->find(baseLink(), true, true);
            remainingBaseLink = baseLink();
        } else if(isJointDownward(0)){
            remainingLinkTraverse->find(linkPath_[1], false, true);
            remainingLinkTraverse->insertRoot(baseLink());
            remainingBaseLink = baseLink();
        } else {
            remainingLinkTraverse->find(baseLink(), true, false);
            while(!remainingLinkTraverse->empty()){
                auto nextLink = remainingLinkTraverse->link(1);
                remainingBaseLink = nextLink;
                if(remainingLinkTraverse->isDownward(1)){
                    break;
                }
                if(nextLink->child()->sibling()){ // has multiple children
                    break;
                }
                remainingLinkTraverse->remove(0);
            }
        }
        int n = linkPath_.numLinks();
        for(int i=0; i < n; ++i){
            auto link = linkPath_[i];
            if(link != remainingBaseLink){
                remainingLinkTraverse->remove(link);
            }
        }
    }
    remainingLinkTraverse->calcForwardKinematics();
    return true;
}


int JointPath::numIterations() const
{
    return numericalIK ? numericalIK->iteration : 0;
}


bool JointPath::hasCustomIK() const
{
    return false;
}


bool JointPath::hasAnalyticalIK() const
{
    return hasCustomIK();
}


bool JointPath::calcInverseKinematics
(const Vector3& base_p, const Matrix3& base_R, const Vector3& end_p, const Matrix3& end_R)
{
    Isometry3 T_base;
    T_base.linear() = base_R;
    T_base.translation() = base_p;

    Isometry3 T_end;
    T_end.linear() = end_R;
    T_end.translation() = end_p;
    
    return setBaseLinkGoal(T_base).calcInverseKinematics(T_end);
}


std::ostream& operator<<(std::ostream& os, JointPath& path)
{
    int n = path.numJoints();
    for(int i=0; i < n; ++i){
        Link* link = path.joint(i);
        os << link->name();
        if(i != n){
            os << (path.isJointDownward(i) ? " => " : " <= ");
        }
    }
    os << std::endl;
    return os;
}


namespace {

// deprecated
class JointPathWithCustomizerIk : public JointPath
{
    BodyPtr body;
    int ikTypeId;
    bool isCustomizedIkPathReversed;
    
public:
    JointPathWithCustomizerIk(const BodyPtr& body, Link* baseLink, Link* endLink)
        : JointPath(baseLink, endLink),
          body(body)
    {
        ikTypeId = body->customizerInterface()->initializeAnalyticIk(
            body->customizerHandle(), baseLink->index(), endLink->index());
        if(ikTypeId){
            isCustomizedIkPathReversed = false;
        } else {
            // try reversed path
            ikTypeId = body->customizerInterface()->initializeAnalyticIk(
                body->customizerHandle(), endLink->index(), baseLink->index());
            if(ikTypeId){
                isCustomizedIkPathReversed = true;
            }
        }
    }
        
    virtual bool hasCustomIK() const override
    {
        return (ikTypeId != 0);
    }
    
    virtual bool calcInverseKinematics(const Isometry3& T) override
    {
        if(isCustomIkDisabled() || ikTypeId == 0){
            return JointPath::calcInverseKinematics(T);
        }
        
        const Link* baseLink_ = baseLink();
        Vector3 p;
        Matrix3 R;
        
        if(!isCustomizedIkPathReversed){
            p = baseLink_->R().transpose() * (T.translation() - baseLink_->p());
            R.noalias() = baseLink_->R().transpose() * T.linear();
        } else {
            p = T.linear().transpose() * (baseLink_->p() - T.translation());
            R.noalias() = T.linear().transpose() * baseLink_->R();
        }
        
        bool solved = body->customizerInterface()->
            calcAnalyticIk(body->customizerHandle(), ikTypeId, p, R);

        if(solved){
            calcForwardKinematics();
        }
        
        return solved;
    }
};

}


std::shared_ptr<JointPath> JointPath::getCustomPath(Link* baseLink, Link* endLink)
{
    auto body = baseLink->body();
    auto customJointPathHandler = body->findHandler<CustomJointPathHandler>();
    if(customJointPathHandler){
        auto customPath = customJointPathHandler->getCustomJointPath(baseLink, endLink);
        if(customPath){
            return customPath;
        }
    }

    // deprecated
    if(body->customizerInterface() && body->customizerInterface()->initializeAnalyticIk){
        return make_shared<JointPathWithCustomizerIk>(body, baseLink, endLink);
    }

    return make_shared<JointPath>(baseLink, endLink);
}


std::shared_ptr<JointPath> cnoid::getCustomJointPath(Body* /* body */, Link* baseLink, Link* endLink)
{
    return JointPath::getCustomPath(baseLink, endLink);
}
