#include "qt_pi_serial/ros_node.h"

sensor_msgs::Imu imu_data;
unsigned char in_imu_data_buf[14];

qt_pi_msgs::Encoder prev_enc_data;
nav_msgs::Odometry odom_data;
int32_t in_enc_data_buf[2];
float prev_vel_data[2];
float theta;

void init_odom();

void report_error(int ret){
  switch(ret){
    case 1:
       ROS_ERROR("Wait for reply timed out.");
       break;
    case -1:
       ROS_ERROR("Invalid packet received");
       break;
  }
}

bool GET_ENC_SRV_CB(qt_pi_msgs::GetEncoderData::Request &req,
                    qt_pi_msgs::GetEncoderData::Response &res) {

  com_buf[COMB] = GET_ENCODERS;
  com_buf[CIDB] = 1;
  com_buf[DATASTRB] = 1;

  com_buf_len = com_preamble_len + 1;

  int ret = sendCom();

  if (ret)
  {
    report_error(ret);
    return false;
  }

  memcpy(&res.data.LEFT, com_buf + DATASTRB, 4);
  memcpy(&res.data.RIGHT, com_buf + DATASTRB + 4, 4);

  return true;
}

bool GET_BATT_SRV_CB(qt_pi_msgs::GetBatteryInfo::Request &req,
                     qt_pi_msgs::GetBatteryInfo::Response &res) {
  com_buf[COMB] = GET_BATTERY_LVL;
  com_buf[CIDB] = 1;
  com_buf[DATASTRB] = 1;

  com_buf_len = com_preamble_len + 1;

  int ret = sendCom();

  if (ret){
    report_error(ret);
    return false;
  }

  uint16_t volts;

  memcpy(&volts, com_buf + DATASTRB, 2);

  res.volts = volts * 11.14 / 876.0;

  return true;

}

bool GET_PID_SRV_CB(qt_pi_msgs::GetPIDData::Request &req,
                    qt_pi_msgs::GetPIDData::Response &res) {

  com_buf[COMB] = GET_PID_P;
  com_buf[CIDB] = 1;
  com_buf[DATASTRB] = 1;

  com_buf_len = com_preamble_len + 1;

  int ret = sendCom();

  if (ret){
    report_error(ret);
    return false;
  }

  res.LEFT.kp = com_buf[DATASTRB];
  res.LEFT.kd = com_buf[DATASTRB + 1];
  res.LEFT.ki = com_buf[DATASTRB + 2];
  res.LEFT.ko = com_buf[DATASTRB + 3];
  res.RIGHT.kp = com_buf[DATASTRB + 4];
  res.RIGHT.kd = com_buf[DATASTRB + 5];
  res.RIGHT.ki = com_buf[DATASTRB + 6];
  res.RIGHT.ko = com_buf[DATASTRB + 7];

  return true;
}

bool GET_BAUD_SRV_CB(qt_pi_msgs::GetBaudRate::Request &req,
                     qt_pi_msgs::GetBaudRate::Response &res) {
  com_buf[COMB] = GET_BAUDRATE;
  com_buf[CIDB] = 1;
  com_buf[DATASTRB] = 1;

  com_buf_len = com_preamble_len + 1;
  int ret = sendCom();

  if (ret){
    report_error(ret);
    return false;
  }

  memcpy(&(res.BAUDRATE), com_buf + DATASTRB, 4);

  return true;
}

bool SET_PID_SRV_CB(qt_pi_msgs::SetPIDData::Request &req,
                    qt_pi_msgs::SetPIDData::Response &res) {
  if(req.motor == LEFT_WHEEL)
    com_buf[COMB] = SET_PID_L;
  else 
    com_buf[COMB] = SET_PID_R;

  com_buf[CIDB] = 1;

  com_buf[DATASTRB] = req.data.kp;
  com_buf[DATASTRB + 1] = req.data.kd;
  com_buf[DATASTRB + 2] = req.data.ki;
  com_buf[DATASTRB + 3] = req.data.ko;

  com_buf_len = 6;

  int ret = sendCom();

  if (ret){
    report_error(ret);
    return false;
  }

  return true;
}

bool EN_IMU_SRV_CB(std_srvs::SetBool::Request &req,
                   std_srvs::SetBool::Response &res) {
  com_buf[COMB] = EN_IMU;
  com_buf[CIDB] = 1;
  com_buf[DATASTRB] = req.data;

  com_buf_len = com_preamble_len + 1;

  res.success = 0;

  int ret = sendCom();

  if (ret){
    report_error(ret);
    return false;
  }
  
  res.success = 1;

  return true;
}

bool EN_ODOM_SRV_CB(std_srvs::SetBool::Request &req,
                    std_srvs::SetBool::Response &res) {
  com_buf[COMB] = EN_ODOM;
  com_buf[CIDB] = 1;
  com_buf[DATASTRB] = req.data;

  com_buf_len = com_preamble_len + 1;

  res.success = 0;

  int ret = sendCom();

  if (ret){
    report_error(ret);
    return false;
  }

  res.success = 1;

  return true;
}

bool RESET_ODOM_SRV_CB(std_srvs::Trigger::Request &req,
                       std_srvs::Trigger::Response &res) {
  com_buf[COMB] = RESET_ENCODERS;
  com_buf[CIDB] = 1;
  com_buf[DATASTRB] = 1;

  com_buf_len = com_preamble_len + 1;

  int ret = sendCom();

  res.success = 0;

  if (ret){
    report_error(ret);
    return false;
  }

  init_odom();

  res.success = 1;

  return true;
}

void init_odom(){

  memset(in_enc_data_buf, 0, 8);
  memset(&prev_enc_data, 0, sizeof(qt_pi_msgs::Encoder));
  memset(&odom_data, 0, sizeof(nav_msgs::Odometry));
  memset(prev_vel_data, 0 , sizeof(float) * 2);

  theta = 0;
  
   odom_data.pose.pose.position.x    = 0;
   odom_data.pose.pose.position.y    = 0;
   odom_data.pose.pose.position.z    = 0;

   double cov_pose[36] = {0.05, 0, 0, 0, 0, 0,
			    0, 0.05, 0, 0, 0, 0,
			    0, 0, 0, 0, 0, 0,
			    0, 0, 0, 0, 0, 0,
			    0, 0, 0, 0, 0, 0,
			    0, 0, 0, 0, 0, 0.3};

   double cov_twist[36] = {0.01, 0, 0, 0, 0, 0,
			    0, 0, 0, 0, 0, 0,
			    0, 0, 0, 0, 0, 0,
			    0, 0, 0, 0, 0, 0,
			    0, 0, 0, 0, 0, 0,
			    0, 0, 0, 0, 0, 0.03};

   memcpy(&(odom_data.pose.covariance),  cov_pose,  36 * sizeof(double));
   memcpy(&(odom_data.twist.covariance), cov_twist, 36 * sizeof(double));

   odom_data.pose.pose.orientation.x = 0;
   odom_data.pose.pose.orientation.y = 0;
   odom_data.pose.pose.orientation.z = 0;
   odom_data.pose.pose.orientation.w = 1;

   odom_data.twist.twist.linear.x  = 0;
   odom_data.twist.twist.linear.y  = 0;
   odom_data.twist.twist.linear.z  = 0;

   odom_data.twist.twist.angular.x = 0;
   odom_data.twist.twist.angular.y = 0;
   odom_data.twist.twist.angular.z = 0;

   odom_data.header.stamp = ros::Time::now();
   odom_data.header.frame_id = "odom";
   odom_data.child_frame_id = "base_link";
}

void ODOM_DATA_PUBLISH(unsigned char *buf) {
   
   static tf::TransformBroadcaster br;   

   memcpy(in_enc_data_buf, buf, 8);
   
   float diff_disp[2];
   float dx, dy, dtheta, ddisp, vl, vr;
   
   diff_disp[LEFT_WHEEL]  = M_PI * wheel_dia * (in_enc_data_buf[LEFT_WHEEL]  - prev_enc_data.LEFT ) / (gear_ratio * counts_per_rot);
   diff_disp[RIGHT_WHEEL] = M_PI * wheel_dia * (in_enc_data_buf[RIGHT_WHEEL] - prev_enc_data.RIGHT) / (gear_ratio * counts_per_rot);

   ddisp  = ( diff_disp[LEFT_WHEEL]  + diff_disp[RIGHT_WHEEL] ) / 2.0;
   dtheta = ( diff_disp[RIGHT_WHEEL] - diff_disp[LEFT_WHEEL]  ) / wheel_sep;

   dx = ddisp * cos(theta + dtheta);    //Check if dtheta will be halved or not
   dy = ddisp * sin(theta + dtheta);    

   theta += dtheta;

   odom_data.pose.pose.position.x   += dx;
   odom_data.pose.pose.position.y   += dy;

   odom_data.pose.pose.orientation.z = sin( theta / 2.0 );
   odom_data.pose.pose.orientation.w = cos( theta / 2.0 );

   prev_enc_data.LEFT  = in_enc_data_buf[LEFT_WHEEL];
   prev_enc_data.RIGHT = in_enc_data_buf[RIGHT_WHEEL];

   vl = diff_disp[LEFT_WHEEL]  * ODOM_RATE;
   vr = diff_disp[RIGHT_WHEEL] * ODOM_RATE;
   
   vl = vl * (1 - v_alpha) + prev_vel_data[LEFT_WHEEL]  * v_alpha;
   vr = vr * (1 - v_alpha) + prev_vel_data[RIGHT_WHEEL] * v_alpha;

   prev_vel_data[LEFT_WHEEL]  = vl;
   prev_vel_data[RIGHT_WHEEL] = vr;

   odom_data.twist.twist.linear.x  = ( vr + vl ) / 2.0;
   odom_data.twist.twist.angular.z = ( vr - vl ) / 2.0;

   odom_data.header.stamp = ros::Time::now();

   odom_pub.publish(odom_data);

   br.sendTransform(
                    tf::StampedTransform(
                                         tf::Transform(
                                                       tf::Quaternion(
                                                                     odom_data.pose.pose.orientation.x,
                                                                     odom_data.pose.pose.orientation.y,
                                                                     odom_data.pose.pose.orientation.z,
                                                                     odom_data.pose.pose.orientation.w
                                                                     ), 
                                                       tf::Vector3(
                                                                  odom_data.pose.pose.position.x, 
                                                                  odom_data.pose.pose.position.y, 
                                                                  odom_data.pose.pose.position.z
                                                                  )
                                                      ),
                                         ros::Time::now(),
                                         "odom",
                                         "base_link"
                                         )
                   );

   br.sendTransform(
                    tf::StampedTransform(
                                         tf::Transform(
                                                       tf::Quaternion(
                                                                     0,
                                                                     0,
                                                                     0,
                                                                     1
                                                                     ), 
                                                       tf::Vector3(
                                                                    0.076, 
                                                                   -0.09 , 
                                                                    0.15
                                                                  )
                                                      ),
                                         ros::Time::now(),
                                         "base_link",
                                         "camera_link"
                                         )
                   );

   br.sendTransform(
                    tf::StampedTransform(
                                         tf::Transform(
                                                       tf::Quaternion(
                                                                     0,
                                                                     0,
                                                                    -0.70710678119,
                                                                     0.70710678119
                                                                     ), 
                                                       tf::Vector3(
                                                                   -0.025,
                                                                   -0.05 ,
                                                                    0.085
                                                                  )
                                                      ),
                                         ros::Time::now(),
                                         "base_link",
                                         "imu_link"
                                         )
                   );
}

void TARE_CALLBACK(unsigned char *buf){
  
  static int count = 0;
  static int32_t imu_off_ave[] = {0, 0, 0, 0, 0, 0};

  memcpy(in_imu_data_buf, buf, 14);

  int16_t imu_raw[6];
  
  imu_raw[0] = ((int16_t)in_imu_data_buf[0] << 8 ) | in_imu_data_buf[1];
  imu_raw[1] = ((int16_t)in_imu_data_buf[2] << 8 ) | in_imu_data_buf[3];
  imu_raw[2] = ((int16_t)in_imu_data_buf[4] << 8 ) | in_imu_data_buf[5];

  imu_raw[3] = ((int16_t)in_imu_data_buf[ 8] << 8) | in_imu_data_buf[ 9];
  imu_raw[4] = ((int16_t)in_imu_data_buf[10] << 8) | in_imu_data_buf[11];
  imu_raw[5] = ((int16_t)in_imu_data_buf[12] << 8) | in_imu_data_buf[13];
  
  if(count == 200)
  {

     for(int i = 0; i < 6; i++)
     {
       imu_off_ave[i] /= count;
       imu_offset[i] = (int16_t)(imu_off_ave[i]);
     }

     imu_offset[2] -= 16384;

    set_CB(&IMU_DATA_PUBLISH);
    return;
  }
  
  for(int i = 0; i < 6; i++)
  {
    imu_off_ave[i] += imu_raw[i];
  }

  count++;
}

void tare_routine(){
  
    com_buf[COMB] = EN_IMU;
    com_buf[CIDB] = 1;
    com_buf[DATASTRB] = 0;

    com_buf_len = com_preamble_len + 1;

    int ret = sendCom();

    if (ret){
      report_error(ret);
      ROS_ERROR("Could not disable IMU. Returning.");
      return;
    }

  set_CB(&TARE_CALLBACK);

  com_buf[COMB] = EN_IMU;
  com_buf[CIDB] = 1;
  com_buf[DATASTRB] = 1;

  com_buf_len = com_preamble_len + 1;

  ret = sendCom();

  if (ret){
    report_error(ret);
    ROS_ERROR("Could not enable IMU. Setting normal imu CB and returning.");
    set_CB(&IMU_DATA_PUBLISH);
    return;
  }
}

void IMU_TARE_CMD_CB(const std_msgs::Bool::ConstPtr& msg ){
  tare_routine();
}

void IMU_DATA_PUBLISH(unsigned char *buf) {

  imu_data.header.stamp = ros::Time::now();
  imu_data.header.frame_id = "imu_link";

  memcpy(in_imu_data_buf, buf, 14);  

  int16_t imu_raw[6];
  
  imu_raw[0] = ((int16_t)in_imu_data_buf[0] << 8 ) | in_imu_data_buf[1];
  imu_raw[1] = ((int16_t)in_imu_data_buf[2] << 8 ) | in_imu_data_buf[3];
  imu_raw[2] = ((int16_t)in_imu_data_buf[4] << 8 ) | in_imu_data_buf[5];

  imu_raw[3] = ((int16_t)in_imu_data_buf[ 8] << 8) | in_imu_data_buf[ 9];
  imu_raw[4] = ((int16_t)in_imu_data_buf[10] << 8) | in_imu_data_buf[11];
  imu_raw[5] = ((int16_t)in_imu_data_buf[12] << 8) | in_imu_data_buf[13];

  imu_data.orientation_covariance[0] = -1;
  
  imu_data.linear_acceleration.x =  (imu_raw[0] - imu_offset[0]) * 9.8 / 16384.0;
  imu_data.linear_acceleration.y =  (imu_raw[1] - imu_offset[1]) * 9.8 / 16384.0;
  imu_data.linear_acceleration.z =  (imu_raw[2] - imu_offset[2]) * 9.8 / 16384.0;

  imu_data.linear_acceleration_covariance[0] = imu_data.linear_acceleration_covariance[4] = imu_data.linear_acceleration_covariance[8] = 0.0392266;
  imu_data.linear_acceleration_covariance[1] = imu_data.linear_acceleration_covariance[3] = imu_data.linear_acceleration_covariance[6] = 0.0;
  imu_data.linear_acceleration_covariance[2] = imu_data.linear_acceleration_covariance[5] = imu_data.linear_acceleration_covariance[7] = 0.0;

  imu_data.angular_velocity.x = (imu_raw[3] - imu_offset[3]) * (M_PI / 180.0) / 131.0;
  imu_data.angular_velocity.y = (imu_raw[4] - imu_offset[4]) * (M_PI / 180.0) / 131.0;
  imu_data.angular_velocity.z = (imu_raw[5] - imu_offset[5]) * (M_PI / 180.0) / 131.0;
  
  imu_data.angular_velocity_covariance[0] = imu_data.angular_velocity_covariance[4] = imu_data.angular_velocity_covariance[8] = 0.0008726646;
  imu_data.angular_velocity_covariance[1] = imu_data.angular_velocity_covariance[3] = imu_data.angular_velocity_covariance[6] = 0.0;
  imu_data.angular_velocity_covariance[2] = imu_data.angular_velocity_covariance[5] = imu_data.angular_velocity_covariance[7] = 0.0;
  
  imu_pub.publish(imu_data);
}

void CMD_VEL_CB(const geometry_msgs::Twist::ConstPtr &msg){
    
  com_buf[COMB] = SET_SPEEDS;
  com_buf[CIDB] = 1;

  double left  = (msg->linear.x - msg->angular.z * wheel_sep / 2.0);
  double right = (msg->linear.x + msg->angular.z * wheel_sep / 2.0);

  double ticks_per_frame_l = left  * ( gear_ratio * counts_per_rot / ( M_PI * wheel_dia ) ) / ODOM_RATE;
  double ticks_per_frame_r = right * ( gear_ratio * counts_per_rot / ( M_PI * wheel_dia ) ) / ODOM_RATE;

  int32_t l = ticks_per_frame_l * 1000;
  int32_t r = ticks_per_frame_r * 1000;

  memcpy(com_buf + DATASTRB, &l, 4);
  memcpy(com_buf + DATASTRB + 4, &r, 4);

  com_buf_len = com_preamble_len + 8;  

  int ret = sendCom();
  
  ROS_INFO("%f\t%f\n", ticks_per_frame_l, ticks_per_frame_r);  

  if (ret){
    report_error(ret);
    ROS_ERROR("Could not set motor speeds. Returning.");
    return;
  }
  
}

void init_node(ros::NodeHandle n){

  ROS_INFO("Opening serial port.");

  if (init_comm("/dev/qt_pi_serial_port") < 0) {
    ROS_FATAL("Unable to open port");
    exit(0);
  }
  
  ROS_INFO("Port opened successfully. Initializing services.");  

  pthread_yield();

  usleep(2000000);

  imu_pub = n.advertise<sensor_msgs::Imu>("imu", 1);
  odom_pub = n.advertise<nav_msgs::Odometry>("odom", 1);

  get_baudrate = n.advertiseService("get_baudrate", GET_BAUD_SRV_CB);
  get_batt_state = n.advertiseService("get_battery_state", GET_BATT_SRV_CB);

  get_encoder = n.advertiseService("get_encoder_state", GET_ENC_SRV_CB);
  reset_odom = n.advertiseService("reset_odom", RESET_ODOM_SRV_CB);

  get_pid_val = n.advertiseService("get_pid_values", GET_PID_SRV_CB);
  set_pid_val = n.advertiseService("set_pid_values", SET_PID_SRV_CB);

  imu_tare = n.subscribe("imu_tare", 1, IMU_TARE_CMD_CB);
  en_imu = n.advertiseService("set_imu_state", EN_IMU_SRV_CB);
  en_odom = n.advertiseService("set_odom_state", EN_ODOM_SRV_CB);

  move_cmd = n.subscribe("cmd_vel", 1, CMD_VEL_CB);

  ROS_INFO("Services initialized");  
  
  tare_routine();
  init_odom();
}

void shutdown_node(){

    com_buf[COMB] = EN_IMU;
    com_buf[CIDB] = 1;
    com_buf[DATASTRB] = 0;

    com_buf_len = com_preamble_len + 1;

    int ret = sendCom();

    if (ret){
      report_error(ret);
      ROS_ERROR("Could not disable IMU. Shutting down.");
      return;
    }

    
    com_buf[COMB] = EN_ODOM;
    com_buf[CIDB] = 1;
    com_buf[DATASTRB] = 0;

    com_buf_len = com_preamble_len + 1;

    ret = sendCom();

    if (ret){
      report_error(ret);
      ROS_ERROR("Could not disable ODOM. Shutting down.");
      return;
    }
}
