#include <mkt_algorithm/diff/diff_go_straight.h>
#include <boost/dll/alias.hpp>
#include <robot/robot.h>

mkt_algorithm::diff::GoStraight::GoStraight()
  : PredictiveTrajectory()
{
}

mkt_algorithm::diff::GoStraight::~GoStraight() {}

void mkt_algorithm::diff::GoStraight::initialize(
    robot::NodeHandle &nh, const std::string &name, TFListenerPtr tf, robot_costmap_2d::Costmap2DROBOT *costmap_robot, const score_algorithm::TrajectoryGenerator::Ptr &traj)
{
  if (!initialized_)
  {
    nh_ = nh;
    nh_priv_ = robot::NodeHandle(nh, name);
    name_ = name;
    tf_ = tf;
    traj_ = traj;
    costmap_robot_ = costmap_robot;
    nh_.param<double>("min_approach_linear_velocity", min_approach_linear_velocity_, 0.0);
    this->getParams();

    std::vector<robot_geometry_msgs::Point> footprint = costmap_robot_? costmap_robot_->getRobotFootprint() : std::vector<robot_geometry_msgs::Point>();
    if (footprint.size() > 1)
    {
      double min_length = 1e10;
      double max_length = 0.0;
      for (size_t i = 0; i < footprint.size(); ++i)
      {
        const robot_geometry_msgs::Point &p1 = footprint[i];
        const robot_geometry_msgs::Point &p2 = footprint[(i + 1) % footprint.size()]; // wrap around to first point
        double len = std::hypot(p2.x - p1.x, p2.y - p1.y) / 2.0;
        if (len > max_length)
        {
          max_length = len;
        }
        if (len < min_length)
        {
          min_length = len;
        }
      }
      this->min_path_distance_ = min_length > 0.1 ? min_length : 0.1;
      this->max_path_distance_ = max_length > 0.1 ? max_length : 0.1;
    }

    // kalman
    last_actuator_update_ = robot::Time::now();
    n_ = 6; // [x, vx, ax, y, vy, ay, theta, vtheta, atheta]
    m_ = 2; // measurements: x, y, theta
    double dt = control_duration_;

    // Khởi tạo ma trận
    A = Eigen::MatrixXd::Identity(n_, n_);
    C = Eigen::MatrixXd::Zero(m_, n_);
    Q = Eigen::MatrixXd::Zero(n_, n_);
    R = Eigen::MatrixXd::Identity(m_, m_);
    P = Eigen::MatrixXd::Identity(n_, n_);

    for (int i = 0; i < n_; i += 3)
    {
      A(i, i + 1) = dt;
      A(i, i + 2) = 0.5 * dt * dt;
      A(i + 1, i + 2) = dt;
    }

    C(0, 0) = 1;
    C(1, 3) = 1;
    Q(2, 2) = 0.1;
    Q(5, 5) = 0.6;

    R(0, 0) = 0.1;
    R(1, 1) = 0.2;

    P(3, 3) = 0.4;
    P(4, 4) = 0.4;
    P(5, 5) = 0.4;

    kf_ = boost::make_shared<KalmanFilter>(dt, A, C, Q, R, P);
    Eigen::VectorXd x0(n_);
    x0 << 0, 0, 0, 0, 0, 0;
    kf_->init(robot::Time::now().toSec(), x0);

    x_direction_ = y_direction_ = theta_direction_ = 0;
    initialized_ = true;
    robot::log_info("GoStraight Initialized successfully");
  }
}

mkt_msgs::Trajectory2D mkt_algorithm::diff::GoStraight::calculator(
    const robot_nav_2d_msgs::Pose2DStamped &pose, const robot_nav_2d_msgs::Twist2D &velocity)
{
  // robot::log_info("x %f y %f theta %f", pose.pose.x, pose.pose.y, pose.pose.theta);
  mkt_msgs::Trajectory2D result;
  result.velocity.x = result.velocity.y = result.velocity.theta = 0.0;
  if (!traj_)
    return result;
  if (compute_plan_.poses.size() < 2)
  {
    robot::log_warning("Local compute plan is available");
    return result;
  }

  robot::Time current_time = robot::Time::now();
  const double dt = (current_time - last_actuator_update_).toSec();
  last_actuator_update_ = current_time;

  robot_nav_2d_msgs::Twist2D drive_cmd;
  double sign_x = x_direction_;

  robot_nav_2d_msgs::Twist2D twist;
  traj_->startNewIteration(velocity); // Nhận tốc độ hiện tại từ odom
  while (robot::ok() && traj_->hasMoreTwists())
  {
    twist = traj_->nextTwist();
  }
  double v_max = sign_x > 0 ? traj_->getTwistLinear(true).x : traj_->getTwistLinear(false).x;
  drive_cmd.x = std::min(sqrt(twist.x * twist.x),  fabs(v_max));
  robot_nav_2d_msgs::Path2D transformed_plan = this->transform_plan_;
  if (transformed_plan.poses.empty())
  {
    robot::log_warning("[%s:%d]\n Transformed plan is empty. Cannot determine a localglobal_plan.", __FILE__, __LINE__);
    stopWithAccLimits(pose, velocity, drive_cmd, dt);
    result.velocity = drive_cmd;
    prevous_drive_cmd_ = drive_cmd;
    return result;
  }
  double lookahead_dist = getLookAheadDistance(velocity);
  auto carrot_pose = *getLookAheadPoint(velocity, lookahead_dist, transformed_plan);
  robot_geometry_msgs::PoseStamped carrot_pose_stamped = robot_nav_2d_utils::pose2DToPoseStamped(carrot_pose);
  
//   // === Final Heading Alignment Check ===
//   double xy_error = 0.0, heading_error = 0.0;
//   if (shouldAlignToFinalHeading(transformed_plan, carrot_pose, velocity, xy_error, heading_error, sign_x))
//   {
//     // Use Arc Motion controller for final heading alignment
//     alignToFinalHeading(xy_error, heading_error, velocity, sign_x, dt, drive_cmd);
// #ifdef BUILD_WITH_ROS
//     ROS_INFO("xy_err=%.3f, heading_err=%.3f deg, v=%.3f, w_current=%.3f, w_target=%.3f",
//                       xy_error, heading_error * 180.0 / M_PI, drive_cmd.x, velocity.theta, drive_cmd.theta);
// #endif
//   }
//   else
//   {
    // robot::log_info_at(__FILE__, __LINE__, "journey : %f lookahead_dist : %f", 
    //   journey(transformed_plan.poses, 0, transformed_plan.poses.size() - 1), lookahead_dist);
    if(fabs(carrot_pose.pose.y) > 0.2)
    {
      lookahead_dist = sqrt(carrot_pose.pose.y *carrot_pose.pose.y + lookahead_dist * lookahead_dist);
    }
    robot_nav_2d_msgs::Twist2D drive_target;
    transformed_plan = this->generateTrajectory(transformed_plan, drive_cmd, velocity, sign_x, drive_target, dt);
    carrot_pose = *getLookAheadPoint(velocity, lookahead_dist, transformed_plan);

    // Normal Pure Pursuit
    this->computePurePursuit(
    carrot_pose, 
    drive_target,
    velocity, 
    transformed_plan,
    min_approach_linear_velocity_, 
    sign_x, 
    dt, 
    drive_cmd);
  // }
  applyDistanceSpeedScaling(compute_plan_, velocity, drive_cmd, sign_x, dt);

  if (this->nav_stop_)
  {
    if (!stopped(velocity, rot_stopped_velocity_, trans_stopped_velocity_))
    {
      if (!stopWithAccLimits(pose, velocity, drive_cmd, dt))
        return result;
    }
    else
      drive_cmd = {};
    result.velocity = drive_cmd;
    return result;
  }

  Eigen::VectorXd y(2);
  y << drive_cmd.x, drive_cmd.theta;

  // Cập nhật lại A nếu dt thay đổi
  for (int i = 0; i < n_; ++i)
    for (int j = 0; j < n_; ++j)
      A(i, j) = (i == j ? 1.0 : 0.0);
  for (int i = 0; i < n_; i += 3)
  {
    A(i, i + 1) = dt;
    A(i, i + 2) = 0.5 * dt * dt;
    A(i + 1, i + 2) = dt;
  }
  kf_->update(y, dt, A);
  double v_min = min_approach_linear_velocity_;
  drive_cmd.x = std::clamp(kf_->state()[0], -fabs(v_max), fabs(v_max));
  drive_cmd.x = fabs(drive_cmd.x) >= v_min ? drive_cmd.x : std::copysign(v_min, sign_x);
  drive_cmd.theta = std::clamp(kf_->state()[3], -max_vel_theta_, max_vel_theta_);
  

  result.poses.clear();
  result.poses.reserve(transformed_plan.poses.size());
  for (const auto &pose_stamped : transformed_plan.poses)
  {
    result.poses.push_back(pose_stamped.pose);
    if(fabs(pose_stamped.pose.x - carrot_pose.pose.x) < 1e-6 && 
      fabs(pose_stamped.pose.y - carrot_pose.pose.y) < 1e-6 && 
      fabs(pose_stamped.pose.theta - carrot_pose.pose.theta) < 1e-6)
    break;
  }
  drive_cmd.theta = std::clamp(drive_cmd.theta, -fabs(0.02), fabs(0.02));
  result.velocity = drive_cmd;
  prevous_drive_cmd_ = drive_cmd;
  return result;
}

score_algorithm::ScoreAlgorithm::Ptr mkt_algorithm::diff::GoStraight::create()
{
  return std::make_shared<mkt_algorithm::diff::GoStraight>();
}

BOOST_DLL_ALIAS(mkt_algorithm::diff::GoStraight::create, MKTAlgorithmDiffGoStraight)