#include "Leg_InvDyn_Controller.hpp"
#include <iostream>

void Leg_InvDyn_Controller::runController(){
  static int iter = 0;
  iter ++;

  

  // This section eventually won't be needed after state estimate
  // is auto propagated to the model
  // {
  
  FBModelState<float> state;
  
  state.bodyOrientation = _stateEstimate->orientation;
  state.bodyPosition    = _stateEstimate->position;
  state.bodyVelocity.head(3) = _stateEstimate->omegaBody;
  state.bodyVelocity.tail(3) = _stateEstimate->vBody;

  state.q.setZero(12);
  state.qd.setZero(12);

  for (int i = 0; i < 4; ++i) {
    state.q(3*i+0) = _legController->datas[i].q[0];
    state.q(3*i+1) = _legController->datas[i].q[1];
    state.q(3*i+2) = _legController->datas[i].q[2];
    state.qd(3*i+0)= _legController->datas[i].qd[0];
    state.qd(3*i+1)= _legController->datas[i].qd[1];
    state.qd(3*i+2)= _legController->datas[i].qd[2];
  }
  _model->setState(state);
  // }

  float t = _controlParameters->controller_dt*iter;

  // Desired trajectory parameters for a simple joint-space trajectory
  float freq_Hz = 1;
  float freq_rad = freq_Hz * 2* 3.14159;
  float amplitude = 3.1415/3;


  // Desired angles, angular velocities, and angular accelerations. 
  // We'll be lazy and use the same desired trajectory for each moving joint.
  Vec12<float> qDes, qdDes, qddDes;
  qDes.setZero();
  qdDes.setZero();
  qddDes.setZero();

  // Desired trajecotry for each joint is a sin wave
  float desired_angle = sin(t*freq_rad)*amplitude;
  float desired_rate = freq_rad*cos(t*freq_rad)*amplitude;
  float desired_acceleration = freq_rad*freq_rad*sin(t*freq_rad)*amplitude;

  // Set desired for a subset of legs
  for( int i = 0 ; i < (int) userParameters.num_moving_legs ; i++) {
    qDes.segment(3*i,3).setConstant(desired_angle);
    qdDes.segment(3*i,3).setConstant(desired_rate);
    qddDes.segment(3*i,i).setConstant(desired_acceleration);
  }

  // Opposite Ab/ad for L & R Legs
  qDes(0)*=-1; qdDes(0)*=-1; qddDes(0)*=-1;
  qDes(6)*=-1; qdDes(6)*=-1; qddDes(6)*=-1;


  // Construct commanded acceleration
  FBModelStateDerivative<float> commandedAccleration;
  commandedAccleration.dBodyVelocity.setZero();
  commandedAccleration.qdd = qddDes + 25*(qdDes - state.qd) + 150*(qDes - state.q);
 
  // Run RNEA inverse dynamics
  Vec18<float> generalizedForce = _model->inverseDynamics(commandedAccleration); // [base_force ; joint_torques]
  Vec12<float> jointTorques = generalizedForce.tail(12); // prune away base force


  // Alternate strategy: Assemble equations of motion
  Mat18<float> H;
  Vec18<float> Cqd, tau_grav;
  
  H = _model->massMatrix();
  Cqd = _model->generalizedCoriolisForce();
  tau_grav = _model->generalizedGravityForce();

  Vec18<float> generalizedAcceleration;
  generalizedAcceleration.head(6)  = commandedAccleration.dBodyVelocity;
  generalizedAcceleration.tail(12) = commandedAccleration.qdd;
  Vec18<float> generalizedForce2 = H*generalizedAcceleration + Cqd + tau_grav;

  // Make sure they match
  Vec18<float> err = generalizedForce - generalizedForce2;
  assert( err.norm() < 1e-4 );

  // Send joint torques to leg controllers
  int dof = 0;
  for(int leg(0); leg<4; ++leg){
    for(int jidx(0); jidx<3; ++jidx){
      _legController->commands[leg].qDes[jidx] = 0;
      _legController->commands[leg].qdDes[jidx] = 0.;
      _legController->commands[leg].tauFeedForward[jidx] = jointTorques(dof);
      dof++;
    }
    _legController->commands[leg].kpJoint = Mat3<float>::Zero();
    _legController->commands[leg].kdJoint = Mat3<float>::Zero();
  }
}