#include <memory>
#include <mutex>
#include <cmath>
#include <vector>

#include "rclcpp/rclcpp.hpp"
#include "std_msgs/msg/float32_multi_array.hpp"
#include "nav_msgs/msg/odometry.hpp"
#include "visualization_msgs/msg/marker_array.hpp"
#include "geometry_msgs/msg/point.hpp"
#include "std_msgs/msg/header.hpp"
#include "tf2/LinearMath/Quaternion.h"
#include "tf2/LinearMath/Matrix3x3.h"

#include "past_control/msg/obstacle_array.hpp"
#include "past_control/msg/obstacle.hpp"

// DNN detection format (per the existing hobot_dnn convention)
// Typically a custom array of detections: [class_id, x1, y1, x2, y2, confidence, ...]
// We'll use Float32MultiArray for input

class ObstacleDetectorNode : public rclcpp::Node
{
public:
  ObstacleDetectorNode()
  : Node("obstacle_detector_node"), has_odom_(false)
  {
    this->declare_parameter("confidence_threshold", 0.5);
    this->declare_parameter("processing_latency", 0.1);  // seconds
    this->declare_parameter("camera_hfov", 1.0472);      // 60 degrees
    this->declare_parameter("camera_height", 0.3);       // meters above ground
    this->declare_parameter("camera_pitch", 0.0);        // radians
    this->declare_parameter("max_detection_range", 3.0); // meters
    this->declare_parameter("publish_rate", 10.0);

    confidence_threshold_ = this->get_parameter("confidence_threshold").as_double();
    processing_latency_ = this->get_parameter("processing_latency").as_double();
    camera_hfov_ = this->get_parameter("camera_hfov").as_double();
    camera_height_ = this->get_parameter("camera_height").as_double();
    camera_pitch_ = this->get_parameter("camera_pitch").as_double();
    max_detection_range_ = this->get_parameter("max_detection_range").as_double();

    // Subscribers
    dnn_sub_ = this->create_subscription<std_msgs::msg::Float32MultiArray>(
      "hobot_dnn_detection", 10,
      std::bind(&ObstacleDetectorNode::dnnCallback, this, std::placeholders::_1));

    odom_sub_ = this->create_subscription<nav_msgs::msg::Odometry>(
      "/odom", 10,
      std::bind(&ObstacleDetectorNode::odomCallback, this, std::placeholders::_1));

    // Publishers
    obstacles_pub_ = this->create_publisher<past_control::msg::ObstacleArray>(
      "/obstacles", 10);

    marker_pub_ = this->create_publisher<visualization_msgs::msg::MarkerArray>(
      "/obstacle_markers", 10);

    double rate = this->get_parameter("publish_rate").as_double();
    int period_ms = static_cast<int>(1000.0 / rate);
    timer_ = this->create_wall_timer(
      std::chrono::milliseconds(period_ms),
      std::bind(&ObstacleDetectorNode::publishLoop, this));

    RCLCPP_INFO(this->get_logger(), "ObstacleDetectorNode started");
  }

private:
  void dnnCallback(const std_msgs::msg::Float32MultiArray::SharedPtr msg)
  {
    std::lock_guard<std::mutex> lock(mutex_);

    // Parse DNN detections
    // Expected format: interleaved [class_id, x_center, y_center, width, height, conf]
    // Normalized coordinates (0-1) within image frame
    const int fields_per_detection = 6;
    int num_detections = msg->data.size() / fields_per_detection;

    latest_detections_.clear();
    for (int i = 0; i < num_detections; ++i) {
      int base = i * fields_per_detection;
      Detection det;
      det.class_id  = static_cast<int>(msg->data[base + 0]);
      det.x_center  = msg->data[base + 1];  // normalized [0,1]
      det.y_center  = msg->data[base + 2];  // normalized [0,1]
      det.width     = msg->data[base + 3];  // normalized
      det.height    = msg->data[base + 4];  // normalized
      det.confidence = msg->data[base + 5];

      if (det.confidence >= confidence_threshold_) {
        latest_detections_.push_back(det);
      }
    }

    has_detections_ = true;
  }

  void odomCallback(const nav_msgs::msg::Odometry::SharedPtr msg)
  {
    std::lock_guard<std::mutex> lock(mutex_);
    cur_x_ = msg->pose.pose.position.x;
    cur_y_ = msg->pose.pose.position.y;

    // Extract yaw from quaternion
    double qx = msg->pose.pose.orientation.x;
    double qy = msg->pose.pose.orientation.y;
    double qz = msg->pose.pose.orientation.z;
    double qw = msg->pose.pose.orientation.w;
    cur_yaw_ = std::atan2(2.0*(qw*qz + qx*qy), 1.0 - 2.0*(qy*qy + qz*qz));
    has_odom_ = true;
  }

  void publishLoop()
  {
    std::lock_guard<std::mutex> lock(mutex_);

    past_control::msg::ObstacleArray obs_array;
    visualization_msgs::msg::MarkerArray marker_array;

    if (!has_odom_ || !has_detections_) {
      // Publish empty
      obstacles_pub_->publish(obs_array);
      marker_pub_->publish(marker_array);
      return;
    }

    int marker_id = 0;
    for (const auto& det : latest_detections_) {
      // Convert pixel coords to world coords using simple pinhole model
      // x_center normalized [0,1] → angle offset from camera optical axis
      double pixel_offset = det.x_center - 0.5;  // [-0.5, 0.5]
      double angle_offset = pixel_offset * camera_hfov_;  // radians

      double obstacle_angle = cur_yaw_ + angle_offset;

      // Estimate distance from bounding box height
      // Larger bbox = closer object (simple inverse relationship)
      double est_distance = (1.0 - det.height) * max_detection_range_ + 0.5;

      // World coordinates of obstacle
      double obs_x = cur_x_ + est_distance * std::cos(obstacle_angle);
      double obs_y = cur_y_ + est_distance * std::sin(obstacle_angle);
      double obs_radius = 0.15;  // default obstacle radius

      // Add to obstacle array
      past_control::msg::Obstacle obs;
      obs.header.stamp = this->now();
      obs.header.frame_id = "odom";
      obs.x = obs_x;
      obs.y = obs_y;
      obs.radius = obs_radius;
      obs.type = "circle";
      obs_array.obstacles.push_back(obs);

      // Add marker
      visualization_msgs::msg::Marker marker;
      marker.header.stamp = this->now();
      marker.header.frame_id = "odom";
      marker.ns = "obstacles";
      marker.id = marker_id++;
      marker.type = visualization_msgs::msg::Marker::CYLINDER;
      marker.action = visualization_msgs::msg::Marker::ADD;
      marker.pose.position.x = obs_x;
      marker.pose.position.y = obs_y;
      marker.pose.position.z = 0.0;
      marker.pose.orientation.w = 1.0;
      marker.scale.x = obs_radius * 2;
      marker.scale.y = obs_radius * 2;
      marker.scale.z = 0.3;
      marker.color.r = 1.0f;
      marker.color.g = 0.0f;
      marker.color.b = 0.0f;
      marker.color.a = 0.8f;
      marker.lifetime = rclcpp::Duration::from_seconds(0.5);
      marker_array.markers.push_back(marker);
    }

    obstacles_pub_->publish(obs_array);
    marker_pub_->publish(marker_array);
  }

  struct Detection
  {
    int class_id;
    double x_center, y_center;
    double width, height;
    double confidence;
  };

  // Subscribers
  rclcpp::Subscription<std_msgs::msg::Float32MultiArray>::SharedPtr dnn_sub_;
  rclcpp::Subscription<nav_msgs::msg::Odometry>::SharedPtr odom_sub_;

  // Publishers
  rclcpp::Publisher<past_control::msg::ObstacleArray>::SharedPtr obstacles_pub_;
  rclcpp::Publisher<visualization_msgs::msg::MarkerArray>::SharedPtr marker_pub_;

  // Timer
  rclcpp::TimerBase::SharedPtr timer_;

  // State
  std::vector<Detection> latest_detections_;
  double cur_x_, cur_y_, cur_yaw_;
  bool has_odom_, has_detections_;
  std::mutex mutex_;

  // Parameters
  double confidence_threshold_;
  double processing_latency_;
  double camera_hfov_;
  double camera_height_;
  double camera_pitch_;
  double max_detection_range_;
};

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