/*============================= FSM State =============================*/
/**
 * FSM State base class
 */

#include "FSM_State.h"

/**
 * Constructor for the FSM State class.
 *
 * @param _controlFSMData holds all of the relevant control data
 * @param stateNameIn the enumerated state name
 * @param stateStringIn the string name of the current FSM state
 */
template <typename T>
FSM_State<T>::FSM_State(ControlFSMData<T>* _controlFSMData,
                        FSM_StateName stateNameIn, std::string stateStringIn)
    : _data(_controlFSMData),
      stateName(stateNameIn),
      stateString(stateStringIn) {
  transitionData.zero();
  std::cout << "[FSM_State] Initialized FSM state: " << stateStringIn
            << std::endl;
}

/**
 * Cartesian impedance control for a given leg.
 *
 * @param leg the leg number to control
 * @param qDes desired joint position
 * @param dqDes desired joint velocity
 */
template <typename T>
void FSM_State<T>::jointPDControl(
    int leg, Vec3<T> qDes, Vec3<T> qdDes) {

  kpMat << 80, 0, 0, 0, 80, 0, 0, 0, 80;
  kdMat << 1, 0, 0, 0, 1, 0, 0, 0, 1;

  _data->_legController->commands[leg].kpJoint = kpMat;
  _data->_legController->commands[leg].kdJoint = kdMat;

  _data->_legController->commands[leg].qDes = qDes;
  _data->_legController->commands[leg].qdDes = qdDes;
}

/**
 * Cartesian impedance control for a given leg.
 *
 * @param leg the leg number to control
 * @param pDes desired foot position
 * @param vDes desired foot velocity
 * @param kp_cartesian P gains
 * @param kd_cartesian D gains
 */
template <typename T>
void FSM_State<T>::cartesianImpedanceControl(int leg, Vec3<T> pDes,
                                             Vec3<T> vDes,
                                             Vec3<double> kp_cartesian,
                                             Vec3<double> kd_cartesian) {
  _data->_legController->commands[leg].pDes = pDes;
  // Create the cartesian P gain matrix
  kpMat << kp_cartesian[0], 0, 0, 0,
      kp_cartesian[1], 0, 0, 0,
      kp_cartesian[2];
  _data->_legController->commands[leg].kpCartesian = kpMat;

  _data->_legController->commands[leg].vDes = vDes;
  // Create the cartesian D gain matrix
  kdMat << kd_cartesian[0], 0, 0, 0, kd_cartesian[1], 0, 0, 0, kd_cartesian[2];
  _data->_legController->commands[leg].kdCartesian = kdMat;
}

/**
 *
 */
/*
template <typename T>
void FSM_State<T>::footstepHeuristicPlacement(int leg) {

  // Create the projection matrix for the 2D foot placement components
  Mat23<float> projectionMatrix;
  projectionMatrix << 1, 0, 0,
                   0, 1, 0;

  Vec3<float> velDes = _data->_desiredStateCommand->data.stateDes.block<3, 1>(6,
0); Vec3<float> angVelDes = _data->_desiredStateCommand->data.stateDes.block<3,
1>(9, 0); Mat3<float> rBody = _data->_stateEstimate.rBody;

  // Find each of the footstep locations for the swing feet
  for (int leg = 0; leg < 4; leg++) {
    if (_data->_gaitScheduler->gaitData.contactStateScheduled(leg)) {
      // The leg is in contact so nothing to do here

    } else {
      if (_data->_gaitScheduler->gaitData._currentGait ==
GaitType::TRANSITION_TO_STAND) {
        // Position the legs under the hips to stand...
        // Could also get rid of this and simply send 0 velocity ang vel
commands
        // from the CoM desired planner...
        Vec3<float> posHip = _data->_quadruped.getHipLocation(leg);
        footstepLocations.col(leg) <<
projectionMatrix.transpose()*projectionMatrix*
                                   (_data->_stateEstimate.position + // rBody *
posHip); } else {
        // Pull out the approximate yaw rate component of the robot in the
world. Vec3<float> yaw_rate; yaw_rate << 0, 0, _stateEstimate.omegaWorld(3);

        Vec3<float> posHip = _data->_quadruped.getHipLocation(leg);

        float timeStance = _data->_gaitScheduler->gaitData.timeStance(leg);

        // Footstep heuristic composed of several parts in the world frame
        footstepLocations.col(leg) <<
projectionMatrix.transpose()*projectionMatrix*      // Ground projection
                                   (_stateEstimate.position + // rBody * posHip
+                                      // Foot under hips timeStance / 2 *
velDes +                             // Raibert Heuristic timeStance / 2 *
(angVelDes.cross(rBody * posHip)) +  // Turning Raibert Heuristic
                                    (_stateEstimate.vBody - velDes));
      }
    }

  }
}
*/

/**
 * Gait independent formulation for choosing appropriate GRF and step locations
 * as well as converting them to leg controller understandable values.
 */
template <typename T>
void FSM_State<T>::runControls() {
  // This option should be set from the user interface or autonomously
  // eventually
  int CONTROLLER_OPTION = 1;

  // Reset the forces and steps to 0
  footFeedForwardForces = Mat34<T>::Zero();
  footstepLocations = Mat34<T>::Zero();

  // Choose the controller to run for picking step locations and balance forces
  if (CONTROLLER_OPTION == 0) {
    // Test to make sure we can control the robot
    // Test to make sure we can control the robot these will be calculated by
    // the controllers
    for (int leg = 0; leg < 4; leg++) {
      footFeedForwardForces.col(leg) << 0.0, 0.0, 0;  //-220.36;
      // footFeedForwardForces.col(leg) = stateEstimate.rBody *
      // footFeedForwardForces.col(leg);

      footstepLocations.col(leg) << 0.0, 0.0,
          -_data->_quadruped->_maxLegLength / 2;
    }
  } else if (CONTROLLER_OPTION == 1) {
    // QP Balance Controller
    runBalanceController();

    // Swing foot landing positions are calculated with heuristics
    for (int leg = 0; leg < 4; leg++) {
      footstepLocations.col(leg) << 0.0, 0.0,
          -_data->_quadruped->_maxLegLength / 2;
    }  // footstepHeuristicPlacement();

  } else if (CONTROLLER_OPTION == 2) {
    // WBC
    runWholeBodyController();

  } else if (CONTROLLER_OPTION == 3) {
    // cMPC
    runConvexModelPredictiveController();

    // Swing foot landing positions are calculated with heuristics
    // footstepHeuristicPlacement();

  } else if (CONTROLLER_OPTION == 4) {
    // RPC
    runRegularizedPredictiveController();

  } else {
    // Zero out the commands if a controller was not selected
    jointFeedForwardTorques =
        Mat34<float>::Zero();                // feed forward joint torques
    jointPositions = Mat34<float>::Zero();   // joint angle positions
    jointVelocities = Mat34<float>::Zero();  // joint angular velocities
    footFeedForwardForces =
        Mat34<float>::Zero();              // feedforward forces at the feet
    footPositions = Mat34<float>::Zero();  // cartesian foot positions
    footVelocities = Mat34<float>::Zero();

    // Print an error message
    std::cout << "[FSM_State] ERROR: No known controller was selected: "
              << CONTROLLER_OPTION << std::endl;
  }
}

/**
 *
 */
template <typename T>
void FSM_State<T>::runBalanceController() {
  double minForce = 25;
  double maxForce = 500;
  double contactStateScheduled[4];  // = {1, 1, 1, 1};
  for (int i = 0; i < 4; i++) {
    contactStateScheduled[i] =
        _data->_gaitScheduler->gaitData.contactStateScheduled(i);
  }

  double minForces[4];  // = {minForce, minForce, minForce, minForce};
  double maxForces[4];  // = {maxForce, maxForce, maxForce, maxForce};
  for (int leg = 0; leg < 4; leg++) {
    minForces[leg] = contactStateScheduled[leg] * minForce;
    maxForces[leg] = contactStateScheduled[leg] * maxForce;
  }

  double COM_weights_stance[3] = {1, 1, 10};
  double Base_weights_stance[3] = {20, 10, 10};
  double pFeet[12], p_des[3], p_act[3], v_des[3], v_act[3], O_err[3], rpy[3],
      omegaDes[3];
  double se_xfb[13];
  double kpCOM[3], kdCOM[3], kpBase[3], kdBase[3];

  for (int i = 0; i < 4; i++) {
    se_xfb[i] = (double)_data->_stateEstimator->getResult().orientation(i);
  }
  // se_xfb[3] = 1.0;
  for (int i = 0; i < 3; i++) {
    rpy[i] = 0;  //(double)_data->_stateEstimator->getResult().rpy(i);
    p_des[i] = (double)_data->_stateEstimator->getResult().position(i);
    p_act[i] = (double)_data->_stateEstimator->getResult().position(i);
    omegaDes[i] =
        0;  //(double)_data->_stateEstimator->getResult().omegaBody(i);
    v_act[i] = (double)_data->_stateEstimator->getResult().vBody(i);
    v_des[i] = (double)_data->_stateEstimator->getResult().vBody(i);

    se_xfb[4 + i] = (double)_data->_stateEstimator->getResult().position(i);
    se_xfb[7 + i] = (double)_data->_stateEstimator->getResult().omegaBody(i);
    se_xfb[10 + i] = (double)_data->_stateEstimator->getResult().vBody(i);

    // Set the translational and orientation gains
    kpCOM[i] = (double)_data->controlParameters->kpCOM(i);
    kdCOM[i] = (double)_data->controlParameters->kdCOM(i);
    kpBase[i] = (double)_data->controlParameters->kpBase(i);
    kdBase[i] = (double)_data->controlParameters->kdBase(i);
  }

  Vec3<T> pFeetVec;
  Vec3<T> pFeetVecCOM;
  // Get the foot locations relative to COM
  for (int leg = 0; leg < 4; leg++) {
    computeLegJacobianAndPosition(**&_data->_quadruped,
                                  _data->_legController->datas[leg].q,
                                  (Mat3<T>*)nullptr, &pFeetVec, 1);
    //pFeetVecCOM = _data->_stateEstimator->getResult().rBody.transpose() *
                  //(_data->_quadruped->getHipLocation(leg) + pFeetVec);

    pFeetVecCOM = _data->_stateEstimator->getResult().rBody.transpose() *
                  (_data->_quadruped->getHipLocation(leg) + _data->_legController->datas[leg].p);


    pFeet[leg * 3] = (double)pFeetVecCOM[0];
    pFeet[leg * 3 + 1] = (double)pFeetVecCOM[1];
    pFeet[leg * 3 + 2] = (double)pFeetVecCOM[2];
  }
  
  balanceController.set_alpha_control(0.01);
  balanceController.set_friction(0.5);
  balanceController.set_mass(46.0);
  balanceController.set_wrench_weights(COM_weights_stance, Base_weights_stance);
  balanceController.set_PDgains(kpCOM, kdCOM, kpBase, kdBase);
  balanceController.set_desiredTrajectoryData(rpy, p_des, omegaDes, v_des);
  balanceController.SetContactData(contactStateScheduled, minForces, maxForces);
  balanceController.updateProblemData(se_xfb, pFeet, p_des, p_act, v_des, v_act,
                                      O_err, 0.0);

  double fOpt[12];
  balanceController.solveQP_nonThreaded(fOpt);

  // Publish the results over LCM
  balanceController.publish_data_lcm();

  // Copy the results to the feed forward forces
  for (int leg = 0; leg < 4; leg++) {
    footFeedForwardForces.col(leg) << (T)fOpt[leg * 3], (T)fOpt[leg * 3 + 1],
        (T)fOpt[leg * 3 + 2];
  }
}

/**
 * Gait independent formulation for choosing appropriate GRF and step locations
 * as well as converting them to leg controller understandable values.
 */
template <typename T>
void FSM_State<T>::turnOnAllSafetyChecks() {
  // Pre controls safety checks
  checkSafeOrientation = true;  // check roll and pitch

  // Post control safety checks
  checkPDesFoot = true;          // do not command footsetps too far
  checkForceFeedForward = true;  // do not command huge forces
  checkLegSingularity = true;    // do not let leg
}

/**
 *
 */
template <typename T>
void FSM_State<T>::turnOffAllSafetyChecks() {
  // Pre controls safety checks
  checkSafeOrientation = false;  // check roll and pitch

  // Post control safety checks
  checkPDesFoot = false;          // do not command footsetps too far
  checkForceFeedForward = false;  // do not command huge forces
  checkLegSingularity = false;    // do not let leg
}

// template class FSM_State<double>;
template class FSM_State<float>;