/*============================ Vision =============================*/
/**
 * FSM State for robot locomotion. Manages the contact specific logic
 * and handles calling the interfaces to the controllers. This state
 * should be independent of controller, gait, and desired trajectory.
 */

#include "FSM_State_Vision.h"
#include <Utilities/Timer.h>
#include <Controllers/WBC_Ctrl/LocomotionCtrl/LocomotionCtrl.hpp>

/**
 * Constructor for the FSM State that passes in state specific info to
 * the generic FSM State constructor.
 *
 * @param _controlFSMData holds all of the relevant control data
 */
template <typename T>
FSM_State_Vision<T>::FSM_State_Vision(
    ControlFSMData<T>* _controlFSMData)
    : FSM_State<T>(_controlFSMData, FSM_StateName::VISION, "VISION"),
        vision_MPC(_controlFSMData->controlParameters->controller_dt,
                30 / (1000. * _controlFSMData->controlParameters->controller_dt),
                _controlFSMData->userParameters),
        cMPCOld(_controlFSMData->controlParameters->controller_dt,
                30 / (1000. * _controlFSMData->controlParameters->controller_dt),
                _controlFSMData->userParameters),
         _visionLCM(getLcmUrl(255))
{
  // Set the safety checks
  this->turnOnAllSafetyChecks();
  // Turn off Foot pos command since it is set in WBC as operational task
  this->checkPDesFoot = false;

  // Initialize GRF and footstep locations to 0s
  this->footFeedForwardForces = Mat34<T>::Zero();
  this->footstepLocations = Mat34<T>::Zero();
  _wbc_ctrl = new LocomotionCtrl<T>(_controlFSMData->_quadruped->buildModel());
  _wbc_data = new LocomotionCtrlData<T>();
  zero_vec3.setZero();
  _global_robot_loc.setZero();
  _robot_rpy.setZero();

  _visionLCM.subscribe("local_heightmap", &FSM_State_Vision<T>::handleHeightmapLCM, this);
  _visionLCM.subscribe("traversability", &FSM_State_Vision<T>::handleIndexmapLCM, this);
  _visionLCM.subscribe("global_to_robot", &FSM_State_Vision<T>::handleLocalization, this);
  _visionLCMThread = std::thread(&FSM_State_Vision<T>::visionLCMThread, this);

  _height_map = DMat<T>::Zero(x_size, y_size);
  _idx_map = DMat<int>::Zero(x_size, y_size);
}
template<typename T>
void FSM_State_Vision<T>::handleLocalization(const lcm::ReceiveBuffer* rbuf, 
    const std::string& chan, const localization_lcmt* msg){
  (void)rbuf;
  (void)chan;

  for(size_t i(0); i<3; ++i){
    _robot_rpy[i] = msg->rpy[i];
    _global_robot_loc[i] = msg->xyz[i];
  }
  _b_localization_data = true;
}

template<typename T>
void FSM_State_Vision<T>::handleHeightmapLCM(const lcm::ReceiveBuffer *rbuf,
                                      const std::string &chan,
                                      const heightmap_t *msg) {
  (void)rbuf;
  (void)chan;

  for(size_t i(0); i<x_size; ++i){
    for(size_t j(0); j<y_size; ++j){
      _height_map(i,j) = msg->map[i][j];
    }
  }
}


template<typename T>
void FSM_State_Vision<T>::handleIndexmapLCM(const lcm::ReceiveBuffer *rbuf,
    const std::string &chan,
    const traversability_map_t *msg) {
  (void)rbuf;
  (void)chan;

  for(size_t i(0); i<x_size; ++i){
    for(size_t j(0); j<y_size; ++j){
      _idx_map(i,j) = msg->map[i][j];
    }
  }
}


template <typename T>
void FSM_State_Vision<T>::onEnter() {
  // Default is to not transition
  this->nextStateName = this->stateName;

  // Reset the transition data
  this->transitionData.zero();
  vision_MPC.initialize();

  if(_b_localization_data){
    _updateStateEstimator();
    _ini_body_pos = _global_robot_loc;
    _ini_body_ori_rpy = _robot_rpy;
  }else{
    _ini_body_pos = (this->_data->_stateEstimator->getResult()).position;
    _ini_body_ori_rpy = (this->_data->_stateEstimator->getResult()).rpy;
  }
  printf("[FSM VISION] On Enter\n");
}

/**
 * Calls the functions to be executed on each control loop iteration.
 */
template <typename T>
void FSM_State_Vision<T>::run() {
  if(_b_localization_data){
    _updateStateEstimator();
  }
  // Call the locomotion control logic for this iteration
  Vec3<T> des_vel; // x,y, yaw
  _UpdateObstacle();

  // Vision Walking
  _UpdateVelCommand(des_vel);
  _LocomotionControlStep(des_vel);

  // Convex Locomotion
  //_RCLocomotionControl();
  

  // Stand still
  //_JPosStand(); 
  _Visualization(des_vel);
}

template <typename T>
void FSM_State_Vision<T>::_RCLocomotionControl() {
  cMPCOld.run<T>(*this->_data);

  if(this->_data->userParameters->use_wbc > 0.9){
    _wbc_data->pBody_des = cMPCOld.pBody_des;
    _wbc_data->vBody_des = cMPCOld.vBody_des;
    _wbc_data->aBody_des = cMPCOld.aBody_des;

    _wbc_data->pBody_RPY_des = cMPCOld.pBody_RPY_des;
    _wbc_data->vBody_Ori_des = cMPCOld.vBody_Ori_des;
    
    for(size_t i(0); i<4; ++i){
      _wbc_data->pFoot_des[i] = cMPCOld.pFoot_des[i];
      _wbc_data->vFoot_des[i] = cMPCOld.vFoot_des[i];
      _wbc_data->aFoot_des[i] = cMPCOld.aFoot_des[i];
      _wbc_data->Fr_des[i] = cMPCOld.Fr_des[i]; 
    }
    _wbc_data->contact_state = cMPCOld.contact_state;

    _wbc_ctrl->run(_wbc_data, *this->_data);

  }
}

template<typename T>
void FSM_State_Vision<T>::_updateStateEstimator(){
  StateEstimate<T> * estimate_handle = this->_data->_stateEstimator->getResultHandle();

  //printf("global robot:%f,%f,%f\n", _global_robot_loc[0], _global_robot_loc[1], _global_robot_loc[2]);
  estimate_handle->position = _global_robot_loc;
  estimate_handle->rpy = _robot_rpy;
  // 
  estimate_handle->rBody = rpyToRotMat(_robot_rpy);
  estimate_handle->orientation = rpyToQuat(_robot_rpy);
}

template<typename T>
void FSM_State_Vision<T>::_JPosStand(){
  Vec3<T> stand_jpos;
  stand_jpos << 0.f, -.8f, 1.6f;
  for(int leg(0); leg<4;++leg){
    this->jointPDControl(leg, stand_jpos, zero_vec3);
  }
}
template <typename T>
void FSM_State_Vision<T>::_UpdateObstacle(){
  _obs_list.clear();
  // TEST
  //Vec3<T> test1; 
  //test1<< 0.7, 0.20, 0.4;
  //_obs_list.push_back(test1);
  //test1<< 0.7, -0.20, 0.4; //_obs_list.push_back(test1); 
  T obstacle_height = 0.15; 
  T threshold_gap = 0.08; 
  bool add_obs(true); 
  Vec3<T> robot_loc; 
  if(_b_localization_data){ robot_loc = _global_robot_loc; } 
  else{ robot_loc = (this->_data->_stateEstimator->getResult()).position; } 
  Vec3<T> obs; obs[2] = 0.27; 

  T dist = 0;
  for(size_t i(0); i<x_size; ++i){
    for(size_t j(0); j<y_size; ++j){
      if( (_height_map(i, j) > obstacle_height) ){ // if too high point
        add_obs = true;
        obs[0] = i*grid_size - 50*grid_size + robot_loc[0];
        obs[1] = j*grid_size - 50*grid_size + robot_loc[1];

        // check with already added obstacle
        for(size_t idx_obs(0); idx_obs < _obs_list.size(); ++idx_obs){
          dist = sqrt(  (obs[0] - _obs_list[idx_obs][0])*(obs[0] - _obs_list[idx_obs][0]) + 
              (obs[1] - _obs_list[idx_obs][1])*(obs[1] - _obs_list[idx_obs][1]) ); 
          if(dist < threshold_gap){
            add_obs = false;
          }
        }// distance check
        if(add_obs){
            _obs_list.push_back(obs);
        }
      }
    } // y loop
  } // x loop

  //_print_obstacle_list();
}

template <typename T>
void FSM_State_Vision<T>::_print_obstacle_list(){
  for(size_t i(0); i < _obs_list.size(); ++i){
    printf("%lu th obstacle: %f, %f, %f\n", i, _obs_list[i][0], _obs_list[i][1], _obs_list[i][2] );
  }
}

template <typename T>
void FSM_State_Vision<T>::_Visualization(const Vec3<T> & des_vel){
  velocity_visual_t vel_visual;
  for(size_t i(0); i<3; ++i){
    vel_visual.vel_cmd[i] = des_vel[i];
    vel_visual.base_position[i] = (this->_data->_stateEstimator->getResult()).position[i];
  }
  _visionLCM.publish("velocity_cmd", &vel_visual);
  

  _obs_visual_lcm.num_obs = _obs_list.size();
  for(size_t i(0); i<_obs_list.size(); ++i){
    _obs_visual_lcm.location[i][0] = _obs_list[i][0];
    _obs_visual_lcm.location[i][1] = _obs_list[i][1];
    _obs_visual_lcm.location[i][2] = _obs_list[i][2];
  }
  _obs_visual_lcm.sigma = 0.15;
  _obs_visual_lcm.height = 0.5;

  _visionLCM.publish("obstacle_visual", &_obs_visual_lcm);
}

template <typename T>
void FSM_State_Vision<T>::_UpdateVelCommand(Vec3<T> & des_vel) {
  des_vel.setZero();

  Vec3<T> target_pos, curr_pos, curr_ori_rpy;

  target_pos.setZero();

  T moving_time = 25.0;
  T curr_time = (T)iter * 0.002;
  //if(curr_time > moving_time){
    //curr_time = moving_time;
  //}
  //target_pos[0] += 2.0 * curr_time/moving_time;
  //target_pos[0] += 2.5 * curr_time/moving_time;
  //target_pos[0] = 0.7 * (1-cos(2*M_PI*curr_time/moving_time));
  target_pos[0] = 1.0 * (1-cos(2*M_PI*curr_time/moving_time));

  if(_b_localization_data){
    curr_pos = _global_robot_loc;
    curr_pos -= _ini_body_pos;
    curr_ori_rpy = _robot_rpy;
  }else{
    curr_pos = (this->_data->_stateEstimator->getResult()).position;
    curr_pos -= _ini_body_pos;
    curr_ori_rpy = (this->_data->_stateEstimator->getResult()).rpy;

  }
  target_pos = rpyToRotMat(_ini_body_ori_rpy).transpose() * target_pos;
  des_vel[0] = 0.7 * (target_pos[0] - curr_pos[0]);
  des_vel[1] = 0.7 * (target_pos[1] - curr_pos[1]);
  des_vel[2] = 0.7 * (_ini_body_ori_rpy[2] - curr_ori_rpy[2]);

  //des_vel[0] = 0.5;
  T inerproduct;
  T x, y, x_obs, y_obs;
  T vel_x, vel_y;
  T sigma(0.15);
  T h(0.5);
  for(size_t i(0); i<_obs_list.size(); ++i){
    x = curr_pos[0];
    y = curr_pos[1];
    x_obs = _obs_list[i][0];
    y_obs = _obs_list[i][1];
    h = _obs_list[i][2];

    inerproduct = (x-x_obs)*(x-x_obs) + (y-y_obs)*(y-y_obs);
    vel_x = h*((x-x_obs)/ sigma/sigma)*exp(-inerproduct/(2*sigma*sigma));
    vel_y = h*((y-y_obs)/ sigma/sigma)*exp(-inerproduct/(2*sigma*sigma));
    //printf("vel_x, y: %f, %f\n", vel_x, vel_y);

    des_vel[0] += vel_x;
    des_vel[1] += vel_y;
  }
  //printf("des vel_x, y: %f, %f\n", des_vel[0], des_vel[1]);

  des_vel[0] = fminf(fmaxf(des_vel[0], -1.), 1.);
  des_vel[1] = fminf(fmaxf(des_vel[1], -1.), 1.);

}
 
/**
 * Manages which states can be transitioned into either by the user
 * commands or state event triggers.
 *
 * @return the enumerated FSM state name to transition into
 */
template <typename T>
FSM_StateName FSM_State_Vision<T>::checkTransition() {
  // Get the next state
  iter++;

  // Switch FSM control mode
  switch ((int)this->_data->controlParameters->control_mode) {
    case K_VISION:
      break;

    case K_LOCOMOTION:
      this->nextStateName = FSM_StateName::LOCOMOTION;
      this->transitionDuration = 0.0;
      break;

    case K_BALANCE_STAND:
      // Requested change to BALANCE_STAND
      this->nextStateName = FSM_StateName::BALANCE_STAND;

      // Transition time is immediate
      this->transitionDuration = 0.0;

      break;

    case K_PASSIVE:
      // Requested change to BALANCE_STAND
      this->nextStateName = FSM_StateName::PASSIVE;

      // Transition time is immediate
      this->transitionDuration = 0.0;

      break;

    case K_RECOVERY_STAND:
      this->nextStateName = FSM_StateName::RECOVERY_STAND;
      this->transitionDuration = 0.;
      break;

    default:
      std::cout << "[CONTROL FSM] Bad Request: Cannot transition from "
                << K_VISION << " to "
                << this->_data->controlParameters->control_mode << std::endl;
  }

  // Return the next state name to the FSM
  return this->nextStateName;
}

/**
 * Handles the actual transition for the robot between states.
 * Returns true when the transition is completed.
 *
 * @return true if transition is complete
 */
template <typename T>
TransitionData<T> FSM_State_Vision<T>::transition() {
  // Switch FSM control mode
  switch (this->nextStateName) {
    case FSM_StateName::BALANCE_STAND:
      _LocomotionControlStep(zero_vec3);

      iter++;
      if (iter >= this->transitionDuration * 1000) {
        this->transitionData.done = true;
      } else {
        this->transitionData.done = false;
      }

      break;

    case FSM_StateName::PASSIVE:
      this->turnOffAllSafetyChecks();
      this->transitionData.done = true;

      break;

    case FSM_StateName::RECOVERY_STAND:
      this->transitionData.done = true;
      break;

    case FSM_StateName::LOCOMOTION:
      this->transitionData.done = true;
      break;


    default:
      std::cout << "[CONTROL FSM] Something went wrong in transition"
                << std::endl;
  }

  // Return the transition data to the FSM
  return this->transitionData;
}

/**
 * Cleans up the state information on exiting the state.
 */
template <typename T>
void FSM_State_Vision<T>::onExit() {
  // Nothing to clean up when exiting
  iter = 0;
}

/**
 * Calculate the commands for the leg controllers for each of the feet by
 * calling the appropriate balance controller and parsing the results for
 * each stance or swing leg.
 */
template <typename T>
void FSM_State_Vision<T>::_LocomotionControlStep(const Vec3<T> & des_vel) {
  // StateEstimate<T> stateEstimate = this->_data->_stateEstimator->getResult();

  // Contact state logic
  vision_MPC.run<T>(*this->_data, des_vel, _height_map, _idx_map);

  if(this->_data->userParameters->use_wbc > 0.9){
    _wbc_data->pBody_des = vision_MPC.pBody_des;
    _wbc_data->vBody_des = vision_MPC.vBody_des;
    _wbc_data->aBody_des = vision_MPC.aBody_des;

    _wbc_data->pBody_RPY_des = vision_MPC.pBody_RPY_des;
    _wbc_data->vBody_Ori_des = vision_MPC.vBody_Ori_des;
    
    for(size_t i(0); i<4; ++i){
      _wbc_data->pFoot_des[i] = vision_MPC.pFoot_des[i];
      _wbc_data->vFoot_des[i] = vision_MPC.vFoot_des[i];
      _wbc_data->aFoot_des[i] = vision_MPC.aFoot_des[i];
      _wbc_data->Fr_des[i] = vision_MPC.Fr_des[i]; 
    }
    _wbc_data->contact_state = vision_MPC.contact_state;
    _wbc_ctrl->run(_wbc_data, *this->_data);

  }
}

template class FSM_State_Vision<float>;