/**
   Sample walking motion controller for the SR1 robot model.
   This program was ported from the "SamplePD" sample of OpenHRP3.
   @author Shin'ichiro Nakaoka
*/

#include <cnoid/SimpleController>
#include <cnoid/Link>
#include <cnoid/BodyMotion>
#include <cnoid/ExecutablePath>
#include <cnoid/FileUtil>

using namespace std;
using namespace cnoid;

static double pgain[] = {
    8000.0, 8000.0, 8000.0, 8000.0, 8000.0, 8000.0,
    3000.0, 3000.0, 3000.0, 3000.0, 3000.0, 3000.0, 3000.0, 
    8000.0, 8000.0, 8000.0, 8000.0, 8000.0, 8000.0,
    3000.0, 3000.0, 3000.0, 3000.0, 3000.0, 3000.0, 3000.0, 
    8000.0, 8000.0, 8000.0 };

static double dgain[] = {
    100.0, 100.0, 100.0, 100.0, 100.0, 100.0,
    100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0,
    100.0, 100.0, 100.0, 100.0, 100.0, 100.0,
    100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0,
    100.0, 100.0, 100.0 };


class SR1WalkPatternController : public cnoid::SimpleController
{
    Body* ioBody;
    int currentFrameIndex;
    int lastFrameIndex;
    std::shared_ptr<MultiValueSeq> qseq;
    MultiValueSeq::Frame qref0;
    MultiValueSeq::Frame qref1;
    vector<double> q0;
    double dt;
    int actuationMode;
    
public:

    virtual bool initialize(SimpleControllerIO* io) override
    {
        ioBody = io->body();
        string patternFile;
        string opt = io->optionString();
        
        if(opt == "position"){
            actuationMode = Link::JointAngle;
            patternFile = "SR1WalkPattern2.seq";
            io->os() << "SR1WalkPatternController: position control mode." << endl;
        } else if(opt == "velocity"){
            actuationMode = Link::JointVelocity;
            patternFile = "SR1WalkPattern2.seq";
            io->os() << "SR1WalkPatternController: velocity control mode." << endl;
        } else {
            actuationMode = Link::JointTorque;
            patternFile = "SR1WalkPattern1.seq";
            io->os() << "SR1WalkPatternController: torque control mode." << endl;
        }

        for(auto joint : ioBody->joints()){
            joint->setActuationMode(actuationMode);
            io->enableIO(joint);
        }

        auto path = shareDirPath() / "motion" / "SR1" / patternFile;
        auto filename = path.make_preferred().string();
        
        BodyMotion motion;
        if(!motion.load(filename)){
            io->os() << motion.seqMessage() << endl;
            return false;
        }
        motion.updateJointPosSeqWithBodyStateSeq();
        qseq = motion.jointPosSeq();
        if(qseq->numFrames() == 0){
            io->os() << "Empty motion data." << endl;
            return false;
        }

        if(ioBody->numJoints() != qseq->numParts()){
            io->os() << "The number of joints must be " << qseq->numParts() << endl;
            return false;
        }

        dt = io->timeStep();
        if(fabs(dt - qseq->timeStep()) > 1.0e-6){
            io->os() << "Warning: the simulation time step is different from that of the motion data " << endl;
        }
        
        currentFrameIndex = 0;
        lastFrameIndex = std::max(0, qseq->numFrames() - 1);
        qref1 = qseq->frame(0);

        q0.resize(qseq->numParts());
        for(int i=0; i < ioBody->numJoints(); ++i){
            q0[i] = ioBody->joint(i)->q();
        }

        return true;
    }

    virtual bool control() override
    {
        switch(actuationMode){

        case Link::JointTorque:
            qref0 = qref1;
            qref1 = qseq->frame(currentFrameIndex);
            for(int i=0; i < ioBody->numJoints(); ++i){
                Link* joint = ioBody->joint(i);
                double q_ref = qref1[i];
                double q = joint->q();
                double dq_ref = (q_ref - qref0[i]) / dt;
                double dq = (q - q0[i]) / dt;
                joint->u() = (q_ref - q) * pgain[i] + (dq_ref - dq) * dgain[i];
                q0[i] = q;
            }
            break;

        case Link::JointAngle:
            qref0 = qseq->frame(currentFrameIndex);
            for(int i=0; i < ioBody->numJoints(); ++i){
                ioBody->joint(i)->q_target() = qref0[i];
            }
            break;

        case Link::JointVelocity:
            qref0 = qref1;
            qref1 = qseq->clampedFrame(currentFrameIndex + 1);
            for(int i=0; i < ioBody->numJoints(); ++i){
                ioBody->joint(i)->dq_target() = (qref1[i] - qref0[i]) / dt;
            }
            break;

        default:
            break;
        }

        if(currentFrameIndex < lastFrameIndex){
            ++currentFrameIndex;
            return true;
        } else {
            return false;
        }
    }
};

CNOID_IMPLEMENT_SIMPLE_CONTROLLER_FACTORY(SR1WalkPatternController)
