#include "origincar_base/wall_fit_core.hpp"

#include <geometry_msgs/msg/transform_stamped.hpp>
#include <nav_msgs/msg/odometry.hpp>
#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/laser_scan.hpp>
#include <tf2_ros/transform_broadcaster.h>

#include <cmath>
#include <mutex>
#include <string>
#include <vector>

namespace
{
std::vector<origincar_wall::WallPoint> scanMsgToPoints(const sensor_msgs::msg::LaserScan &msg, int stride)
{
  std::vector<origincar_wall::WallPoint> points;
  stride = std::max(1, stride);
  points.reserve(msg.ranges.size() / static_cast<std::size_t>(stride) + 1);
  for (std::size_t index = 0; index < msg.ranges.size(); index += static_cast<std::size_t>(stride))
  {
    const auto distance = static_cast<double>(msg.ranges[index]);
    if (!std::isfinite(distance) || distance < msg.range_min || distance > msg.range_max)
    {
      continue;
    }
    const auto angle = static_cast<double>(msg.angle_min) + static_cast<double>(index) * msg.angle_increment;
    points.push_back({distance * std::cos(angle), distance * std::sin(angle)});
  }
  return points;
}

double quaternionToYaw(const geometry_msgs::msg::Quaternion &q)
{
  const auto siny_cosp = 2.0 * (q.w * q.z + q.x * q.y);
  const auto cosy_cosp = 1.0 - 2.0 * (q.y * q.y + q.z * q.z);
  return std::atan2(siny_cosp, cosy_cosp);
}

geometry_msgs::msg::Quaternion yawToQuaternion(double yaw)
{
  geometry_msgs::msg::Quaternion q;
  const auto half = yaw / 2.0;
  q.x = 0.0;
  q.y = 0.0;
  q.z = std::sin(half);
  q.w = std::cos(half);
  return q;
}

origincar_wall::Pose2D odomMsgToMapPose(const nav_msgs::msg::Odometry &msg,
                                        const origincar_wall::WallFitConfig &config)
{
  return {
      msg.pose.pose.position.x + config.coordinate_frame.odom_zero_map.x,
      msg.pose.pose.position.y + config.coordinate_frame.odom_zero_map.y,
      quaternionToYaw(msg.pose.pose.orientation),
  };
}
}  // namespace

class WallLocalizationNode : public rclcpp::Node
{
public:
  WallLocalizationNode()
      : Node("wall_localization_node"),
        tf_broadcaster_(std::make_unique<tf2_ros::TransformBroadcaster>(*this))
  {
    declare_parameter<std::string>("scan_topic", "/scan");
    declare_parameter<std::string>("odom_topic", "/odom_combined");
    declare_parameter<std::string>("wall_odom_topic", "/odom_wall");
    declare_parameter<std::string>("config_path", "/home/sunrise/yiliao_ws/src/origincar_base/config/wall_fit.json");
    declare_parameter<std::string>("odom_frame", "odom");
    declare_parameter<std::string>("base_frame", "base_footprint");
    declare_parameter<bool>("publish_tf", true);
    declare_parameter<double>("localization_hz", 10.0);
    declare_parameter<int>("scan_stride", 2);
    declare_parameter<double>("laser_x", 0.0);
    declare_parameter<double>("laser_y", 0.0);
    declare_parameter<double>("laser_yaw", 0.0);

    config_ = origincar_wall::loadWallFitConfig(get_parameter("config_path").as_string());
    wall_odom_topic_ = get_parameter("wall_odom_topic").as_string();
    odom_frame_ = get_parameter("odom_frame").as_string();
    base_frame_ = get_parameter("base_frame").as_string();
    publish_tf_ = get_parameter("publish_tf").as_bool();
    scan_stride_ = get_parameter("scan_stride").as_int();
    const auto localization_hz = std::max(1.0, get_parameter("localization_hz").as_double());
    min_scan_period_ = rclcpp::Duration::from_seconds(1.0 / localization_hz);
    laser_pose_ = {
        get_parameter("laser_x").as_double(),
        get_parameter("laser_y").as_double(),
        get_parameter("laser_yaw").as_double(),
    };

    wall_odom_pub_ = create_publisher<nav_msgs::msg::Odometry>(wall_odom_topic_, 10);
    scan_sub_ = create_subscription<sensor_msgs::msg::LaserScan>(
        get_parameter("scan_topic").as_string(), 10,
        [this](sensor_msgs::msg::LaserScan::SharedPtr msg) { onScan(*msg); });
    odom_sub_ = create_subscription<nav_msgs::msg::Odometry>(
        get_parameter("odom_topic").as_string(), 10,
        [this](nav_msgs::msg::Odometry::SharedPtr msg) { onOdom(*msg); });

    const auto bounds = origincar_wall::wallBounds(config_);
    RCLCPP_INFO(get_logger(), "Loaded wall fit; bounds x=[%.3f, %.3f] y=[%.3f, %.3f]",
                bounds.x_min, bounds.x_max, bounds.y_min, bounds.y_max);
    RCLCPP_INFO(get_logger(), "Publishing %s; publish_tf=%s", wall_odom_topic_.c_str(), publish_tf_ ? "true" : "false");
  }

private:
  void onOdom(const nav_msgs::msg::Odometry &msg)
  {
    std::lock_guard<std::mutex> lock(mutex_);
    latest_twist_ = msg.twist.twist;
    has_twist_ = true;
    const auto current_odom_pose = odomMsgToMapPose(msg, config_);
    if (!has_wall_pose_)
    {
      predicted_pose_ = current_odom_pose;
      last_wall_pose_ = current_odom_pose;
      has_wall_pose_ = true;
    }
    else if (has_odom_pose_)
    {
      predicted_pose_ = origincar_wall::predictPoseFromOdomDelta(last_wall_pose_, last_odom_pose_, current_odom_pose);
    }
    else
    {
      predicted_pose_ = last_wall_pose_;
    }
    last_odom_pose_ = current_odom_pose;
    has_odom_pose_ = true;
    has_prediction_ = true;
  }

  void onScan(const sensor_msgs::msg::LaserScan &msg)
  {
    const auto now = get_clock()->now();
    if (has_last_scan_time_ && (now - last_scan_time_) < min_scan_period_)
    {
      return;
    }
    last_scan_time_ = now;
    has_last_scan_time_ = true;

    origincar_wall::Pose2D predicted_pose;
    {
      std::lock_guard<std::mutex> lock(mutex_);
      if (!has_prediction_)
      {
        predicted_pose = {config_.coordinate_frame.odom_zero_map.x, config_.coordinate_frame.odom_zero_map.y, 0.0};
        predicted_pose_ = predicted_pose;
        has_prediction_ = true;
      }
      else
      {
        predicted_pose = predicted_pose_;
      }
    }

    const auto scan_points =
        origincar_wall::transformScanPointsToPoseFrame(scanMsgToPoints(msg, scan_stride_), laser_pose_);
    const auto result = origincar_wall::localizeFromScan(scan_points, config_, predicted_pose);

    origincar_wall::Pose2D pose_to_publish = result.ok ? result.pose : predicted_pose;
    {
      std::lock_guard<std::mutex> lock(mutex_);
      last_wall_pose_ = pose_to_publish;
      predicted_pose_ = pose_to_publish;
      has_wall_pose_ = true;
      has_prediction_ = true;
    }
    publishWallOdom(pose_to_publish, msg.header.stamp, result.ok);
    if (publish_tf_)
    {
      publishTf(pose_to_publish, msg.header.stamp);
    }
    if (!result.ok)
    {
      RCLCPP_WARN_THROTTLE(get_logger(), *get_clock(), 2000, "Wall localization degraded: %s",
                           result.message.c_str());
    }
  }

  void publishWallOdom(const origincar_wall::Pose2D &pose,
                       const builtin_interfaces::msg::Time &stamp,
                       bool valid)
  {
    nav_msgs::msg::Odometry odom;
    odom.header.stamp = stamp;
    odom.header.frame_id = odom_frame_;
    odom.child_frame_id = base_frame_;
    odom.pose.pose.position.x = pose.x;
    odom.pose.pose.position.y = pose.y;
    odom.pose.pose.position.z = 0.0;
    odom.pose.pose.orientation = yawToQuaternion(pose.theta);
    odom.pose.covariance[0] = valid ? 0.02 : 0.25;
    odom.pose.covariance[7] = valid ? 0.02 : 0.25;
    odom.pose.covariance[35] = valid ? 0.03 : 0.3;
    {
      std::lock_guard<std::mutex> lock(mutex_);
      if (has_twist_)
      {
        odom.twist.twist = latest_twist_;
      }
    }
    wall_odom_pub_->publish(odom);
  }

  void publishTf(const origincar_wall::Pose2D &pose,
                 const builtin_interfaces::msg::Time &stamp)
  {
    geometry_msgs::msg::TransformStamped transform;
    transform.header.stamp = stamp;
    transform.header.frame_id = odom_frame_;
    transform.child_frame_id = base_frame_;
    transform.transform.translation.x = pose.x;
    transform.transform.translation.y = pose.y;
    transform.transform.translation.z = 0.0;
    transform.transform.rotation = yawToQuaternion(pose.theta);
    tf_broadcaster_->sendTransform(transform);
  }

  std::mutex mutex_;
  origincar_wall::WallFitConfig config_;
  std::string wall_odom_topic_;
  std::string odom_frame_;
  std::string base_frame_;
  bool publish_tf_{true};
  int scan_stride_{2};
  rclcpp::Duration min_scan_period_{0, 100000000};
  bool has_last_scan_time_{false};
  rclcpp::Time last_scan_time_{0, 0, RCL_ROS_TIME};
  origincar_wall::Pose2D laser_pose_{0.0, 0.0, 0.0};
  origincar_wall::Pose2D last_odom_pose_{0.0, 0.0, 0.0};
  origincar_wall::Pose2D last_wall_pose_{0.0, 0.0, 0.0};
  origincar_wall::Pose2D predicted_pose_{0.0, 0.0, 0.0};
  bool has_odom_pose_{false};
  bool has_wall_pose_{false};
  bool has_prediction_{false};
  geometry_msgs::msg::Twist latest_twist_;
  bool has_twist_{false};
  rclcpp::Publisher<nav_msgs::msg::Odometry>::SharedPtr wall_odom_pub_;
  rclcpp::Subscription<sensor_msgs::msg::LaserScan>::SharedPtr scan_sub_;
  rclcpp::Subscription<nav_msgs::msg::Odometry>::SharedPtr odom_sub_;
  std::unique_ptr<tf2_ros::TransformBroadcaster> tf_broadcaster_;
};

int main(int argc, char **argv)
{
  rclcpp::init(argc, argv);
  rclcpp::spin(std::make_shared<WallLocalizationNode>());
  rclcpp::shutdown();
  return 0;
}
