#include <pnkx_local_planner/transforms.h>
#include <robot_nav_2d_utils/path_ops.h>
#include <robot_nav_2d_utils/tf_help.h>
#include <robot/robot.h>
#include <data_convert/data_convert.h>
#include <robot_tf3_geometry_msgs/tf3_geometry_msgs.h>

bool pnkx_local_planner::getPose(TFListenerPtr tf, std::string base_frame_id, std::string global_frame, robot_nav_2d_msgs::Pose2DStamped& global_pose, double transform_tolerance)
{
  if (!tf)
    return false;

  robot_geometry_msgs::PoseStamped global_pose_stamped;
  tf3::toMsg(tf3::Transform::getIdentity(), global_pose_stamped.pose);
  robot_geometry_msgs::PoseStamped robot_pose;
  tf3::toMsg(tf3::Transform::getIdentity(), robot_pose.pose);
  robot_pose.header.frame_id = base_frame_id;
  robot_pose.header.stamp = robot::Time();
  robot::Time current_time = robot::Time::now(); // save time for checking tf delay later
  
  // Convert robot::Time to tf3::Time
  tf3::Time tf3_current_time = data_convert::convertTime(current_time);
  tf3::Time tf3_zero_time = data_convert::convertTime(robot::Time());
  
  // get the global pose of the robot
  try
  {
    // use current time if possible (makes sure it's not in the future)
    std::string error_msg;
    if (tf->canTransform(global_frame, base_frame_id, tf3_current_time, &error_msg))
    {
      // Transform is available at current time
      tf3::TransformStampedMsg transform = tf->lookupTransform(global_frame, base_frame_id, tf3_current_time);
      tf3::doTransform(robot_pose, global_pose_stamped, transform);
    }
    // use the latest otherwise (tf3::Time() means latest available)
    else
    {
      // Try to get latest transform
      tf3::TransformStampedMsg transform = tf->lookupTransform(global_frame, base_frame_id, tf3_zero_time);
      tf3::doTransform(robot_pose, global_pose_stamped, transform);
    }
    global_pose = robot_nav_2d_utils::poseStampedToPose2D(global_pose_stamped);
  }
  catch (tf3::LookupException& ex)
  {
    robot::log_warning_throttle(1.0, "No Transform available Error looking up robot pose: %s\n", ex.what());
    return false;
    // throw robot_nav_core2::LocalPlannerException(ex.what());
  }
  catch (tf3::ConnectivityException& ex)
  {
    robot::log_warning_throttle(1.0, "Connectivity Error looking up robot pose: %s\n", ex.what());
    return false;
    // throw robot_nav_core2::LocalPlannerException(ex.what());
  }
  catch (tf3::ExtrapolationException& ex)
  {
    robot::log_warning_throttle(1.0, "Extrapolation Error looking up robot pose: %s\n", ex.what());
    return false;
    // throw robot_nav_core2::LocalPlannerException(ex.what());
  }
  // check global_pose timeout
  if (!global_pose_stamped.header.stamp.isZero() && current_time.toSec() - global_pose_stamped.header.stamp.toSec() > transform_tolerance)
  {
    robot::log_warning_throttle(1.0, "Local Planner transform %s %s timeout. Current time: %.4f, global_pose stamp: %.4f, tolerance: %.4f", base_frame_id.c_str(), global_frame.c_str(),
      current_time.toSec(), global_pose_stamped.header.stamp.toSec(), transform_tolerance);
    return false;
  }
  return true;
}

bool pnkx_local_planner::transformGlobalPose(TFListenerPtr tf, const std::string& global_frame,
  const robot_nav_2d_msgs::Pose2DStamped& stamped_pose, robot_nav_2d_msgs::Pose2DStamped& newer_pose)
{
  if (tf == nullptr)
    return false;
  bool result = robot_nav_2d_utils::transformPose(tf, global_frame, stamped_pose, newer_pose);
  newer_pose.header.seq = stamped_pose.header.seq;
  return result;
}

double pnkx_local_planner::getSquareDistance(const robot_geometry_msgs::Pose2D& pose_a, const robot_geometry_msgs::Pose2D& pose_b)
{
  double x_diff = pose_a.x - pose_b.x;
  double y_diff = pose_a.y - pose_b.y;
  return x_diff * x_diff + y_diff * y_diff;
}

bool pnkx_local_planner::transformGlobalPlan(
  TFListenerPtr tf, const robot_nav_2d_msgs::Path2D& global_plan, const robot_nav_2d_msgs::Pose2DStamped& pose,
  const robot_costmap_2d::Costmap2DROBOT* costmap, const std::string& global_frame, double max_plan_length,
  robot_nav_2d_msgs::Path2D& transformed_plan, const bool from_global_frame)
{
  if (global_plan.poses.size() == 0)
  {
    robot::log_warning("Received plan with zero length");
    throw robot_nav_core2::PlannerException("Received plan with zero length");
  }

  transformed_plan.poses.clear();
  if (tf == nullptr)
    throw robot_nav_core2::PlannerException("TFListenerPtr is null");

  if (global_plan.poses.empty())
    throw robot_nav_core2::PlannerException("Received plan with zero length");

  try
  {
    // let's get the pose of the robot in the frame of the plan
    robot_nav_2d_msgs::Pose2DStamped robot_pose;
    if (!robot_nav_2d_utils::transformPose(tf, global_plan.header.frame_id, pose, robot_pose))
    {
      throw robot_nav_core2::PlannerTFException("transformGlobalPlan: Unable to transform robot pose into global plan's frame");
    }

    transformed_plan.header.frame_id = costmap->getGlobalFrameID();
    transformed_plan.header.stamp = pose.header.stamp;

    robot_nav_2d_msgs::Pose2DStamped new_pose;
    robot_nav_2d_msgs::Pose2DStamped stamped_pose;
    stamped_pose.header.frame_id = global_plan.header.frame_id;

    for (unsigned int i = 0; i < global_plan.poses.size(); i++)
    {
      // now we'll transform until points are outside of our distance threshold
      stamped_pose = global_plan.poses[i];
      if (pnkx_local_planner::transformGlobalPose(tf, global_frame, global_plan.poses[i], new_pose))
      {
        transformed_plan.poses.push_back(new_pose);
      }
    }

    // Prune both plans based on robot position
    // Note that this part of the algorithm assumes that the global plan starts near the robot (at one point)
    // Otherwise it may take a few iterations to converge to the proper behavior

      // let's get the pose of the robot in the frame of the transformed_plan/costmap
    robot_nav_2d_msgs::Pose2DStamped costmap_pose;
    if (!robot_nav_2d_utils::transformPose(tf, transformed_plan.header.frame_id, pose, costmap_pose))
    {
      throw robot_nav_core2::PlannerTFException("Unable to transform robot pose into costmap's frame");
    }

    robot::log_assert(global_plan.poses.size() >= transformed_plan.poses.size(), "Global plan size is less than transformed plan size");

    std::vector<robot_nav_2d_msgs::Pose2DStamped>::iterator it = transformed_plan.poses.begin();

    double sq_prune_dist = 0.1;

    std::vector<robot_geometry_msgs::Point> footprint = costmap->getRobotFootprint();
    size_t n = footprint.size();
    if (n > 1)
    {
      double max_length = 0.0;
      for (size_t i = 0; i < n; ++i)
      {
        const robot_geometry_msgs::Point& p1 = footprint[i];
        const robot_geometry_msgs::Point& p2 = footprint[(i + 1) % n]; // wrap around to first point
        double len = std::hypot(p2.x - p1.x, p2.y - p1.y);
        if (len > max_length)
        {
          max_length = len;
        }
      }
      sq_prune_dist += max_length;
    }
    // if(from_global_frame) sq_prune_dist = 0;
    while (it != transformed_plan.poses.end())
    {
      const robot_nav_2d_msgs::Pose2DStamped& w = *it;
      // Fixed error bound of 1 meter for now. Can reduce to a portion of the map size or based on the resolution
      if (pnkx_local_planner::getSquareDistance(costmap_pose.pose, w.pose) < sq_prune_dist)
      {
        // robot::log_debug_at(__FILE__, __LINE__, "Nearest waypoint to <%f, %f> is <%f, %f>\n",
        //   costmap_pose.pose.x, costmap_pose.pose.y, w.pose.x, w.pose.y);
        break;
      }

      it = transformed_plan.poses.erase(it);
    }

    if (transformed_plan.poses.size() == 0)
    {
      throw robot_nav_core2::PlannerTFException("Resulting plan has 0 poses in it.");
    }
  }
  catch (tf3::LookupException& ex)
  {
    robot::log_error_at(__FILE__, __LINE__, "No Transform available Error: %s\n", ex.what());
    return false;
  }
  catch (tf3::ConnectivityException& ex)
  {
    robot::log_error("Connectivity Error: %s\n", ex.what());
    return false;
  }
  catch (tf3::ExtrapolationException& ex)
  {
    robot::log_error("Extrapolation Error: %s\n", ex.what());
    if (global_plan.poses.size() > 0)
      robot::log_error("[%s:%d]\n Robot Frame: %s Plan Frame size %d: %s\n", __FILE__, __LINE__, global_frame.c_str(), (unsigned int)global_plan.poses.size(), global_plan.header.frame_id.c_str());
    return false;
  }
  return true;
}