#include "ros/ros.h"
#include "sensor_msgs/Image.h"
#include "geometry_msgs/Twist.h"
#include "cv_bridge/cv_bridge.h"
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/core/core.hpp>
#include <iostream>
#include <cmath>

class RedBallChaser {
public:
    RedBallChaser() {
        cmd_vel_pub = nh.advertise<geometry_msgs::Twist>("/cmd_vel", 1);
        image_sub = nh.subscribe("/camera/image", 1, &RedBallChaser::imageCallback, this);

        cv::namedWindow("Red Ball Detection", cv::WINDOW_NORMAL);
    }

    ~RedBallChaser() {
        cv::destroyAllWindows();
    }

    void imageCallback(const sensor_msgs::ImageConstPtr& msg) {
        cv_bridge::CvImagePtr cv_ptr;
        try {
            cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("cv_bridge exception: %s", e.what());
            return;
        }

        cv::Mat hsv_image;
        cv::cvtColor(cv_ptr->image, hsv_image, cv::COLOR_BGR2HSV);

        cv::Mat red_mask1, red_mask2, red_mask;
        cv::inRange(hsv_image, cv::Scalar(0, 120, 70), cv::Scalar(10, 255, 255), red_mask1);
        cv::inRange(hsv_image, cv::Scalar(170, 120, 70), cv::Scalar(180, 255, 255), red_mask2);
        red_mask = red_mask1 | red_mask2;

        std::vector<std::vector<cv::Point>> contours;
        cv::findContours(red_mask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

        bool ball_found = false;
        for (auto &contour : contours) {
            double area = cv::contourArea(contour);
            if (area > 0) {
                float radius;
                cv::Point2f center;
                cv::minEnclosingCircle(contour, center, radius);

                if (radius > 5 && isCircular(contour, radius)) {
                    ball_found = true;
                    if (radius >= 40) {
                        executeHalfCircle();
                        return;
                    }
                    // Adjust robot's position and orientation
                    adjustRobotPosition(center, cv_ptr->image.cols);
                    break;
                }
            }
        }

        if (!ball_found) {
            executeSearch();
        }
    }

    void executeHalfCircle() {
        geometry_msgs::Twist twist_msg;
        twist_msg.angular.z = 0.5; // Adjust for half-circle
        twist_msg.linear.x = -0.05; // Adjust for slight forward movement

        ros::Time start = ros::Time::now();
        double duration = 4.0; // Adjust for proper half-circle based on your robot
        while (ros::Time::now() - start < ros::Duration(duration)) {
            cmd_vel_pub.publish(twist_msg);
            ros::Duration(0.1).sleep();
        }

        twist_msg.angular.z = 0;
        twist_msg.linear.x = 0;
        cmd_vel_pub.publish(twist_msg);
    }

    void executeSearch() {
        rotate360();
        moveRandomDirection();
        rotate360();
    }

private:
    ros::NodeHandle nh;
    ros::Publisher cmd_vel_pub;
    ros::Subscriber image_sub;

    bool isCircular(std::vector<cv::Point> &contour, float radius) {
        double area = cv::contourArea(contour);
        double circularity = 4 * M_PI * area / (2 * M_PI * radius) * (2 * M_PI * radius);
        return circularity >= 0.7 && circularity <= 1.3;
    }

    void adjustRobotPosition(cv::Point2f center, int frame_width) {
        float err = center.x - frame_width / 2;
        geometry_msgs::Twist twist_msg;

        if (std::abs(err) > 100) {
            twist_msg.angular.z = -0.005 * err;
        } else {
            twist_msg.linear.x = -0.2;
        }
        cmd_vel_pub.publish(twist_msg);
    }

    void rotate360() {
        geometry_msgs::Twist twist_msg;
        twist_msg.angular.z = 0.5; // Adjust as needed for the full rotation

        ros::Time start = ros::Time::now();
        while (ros::Time::now() - start < ros::Duration(2 * M_PI / twist_msg.angular.z)) {
            cmd_vel_pub.publish(twist_msg);
            ros::Duration(0.1).sleep();
        }

        twist_msg.angular.z = 0;
        cmd_vel_pub.publish(twist_msg);
    }

    void moveRandomDirection() {
        geometry_msgs::Twist twist_msg;
        twist_msg.linear.x = -0.2; // Move forward

        double move_duration = (rand() % 20 + 10) / 10.0; // Random duration between 1 and 3 seconds
        ros::Time start = ros::Time::now();
        while (ros::Time::now() - start < ros::Duration(move_duration)) {
            cmd_vel_pub.publish(twist_msg);
            ros::Duration(0.1).sleep();
        }

        twist_msg.linear.x = 0;
        cmd_vel_pub.publish(twist_msg);
    }
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "redball_chase_cpp");
    RedBallChaser chaser;
    ros::spin();
    return 0;
}