#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from geometry_msgs.msg import Twist
from nav_msgs.msg import Odometry
from sensor_msgs.msg import LaserScan
import math
import time
import sys
import numpy as np

try:
    import pygame
except ImportError:
    print("Błąd: Brak biblioteki pygame. Zainstaluj wpisując: pip3 install pygame")
    sys.exit(1)

class HarvesterStateMachine(Node):
    def __init__(self):
        super().__init__('harvester_sm_node')

        odom_topic_name = '/odom' 
        cmd_vel_topic_name = '/cmd_vel'
        
        qos_profile = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=10
        )

        self.cmd_vel_pub = self.create_publisher(Twist, cmd_vel_topic_name, 10)
        self.odom_sub = self.create_subscription(Odometry, odom_topic_name, self.odom_callback, qos_profile)
        self.lidar_left_sub = self.create_subscription(LaserScan, '/lidar_left', self.lidar_left_callback, qos_profile)
        self.lidar_right_sub = self.create_subscription(LaserScan, '/lidar_right', self.lidar_right_callback, qos_profile)

        self.STATE_INITIALIZING = "INITIALIZING"
        self.STATE_GPS_DRIVE    = "GPS_DRIVE"
        self.STATE_LIDAR_DRIVE  = "LIDAR_DRIVE"
        self.STATE_HARVESTING   = "HARVESTING"
        self.STATE_MISSION_DONE = "MISSION_DONE"
        self.current_state = self.STATE_INITIALIZING

        self.robot_x, self.robot_y, self.robot_yaw = 0.0, 0.0, 0.0
        self.has_odom = False
        self.start_offset_x, self.start_offset_y = 0.0, 0.0

        self.pts_L = ([], [])
        self.pts_R = ([], [])
        
        self.avg_left = 0.0
        self.avg_right = 0.0
        self.end_of_row_detected = False

        self.latest_left_scan = None
        self.latest_right_scan = None

        self.lidar_start_x, self.lidar_start_y = 0.0, 0.0
        self.lidar_enter_time = 0.0
        self.harvesting_done_in_this_row = False
        self.harvest_trigger_distance = 6.0  

        #   PUNKTY DOCELOWE
        self.global_waypoints = [{"x": -3.0, "y": 0.8}, {"x": -5.0, "y": -5}]
        self.local_waypoints = []
        self.current_wp_idx = 0
        self.position_reached = False 

        #   PARAMETRY REGULACJI
        self.target_pos_tolerance = 0.4      
        self.target_heading_tolerance = 0.05  
        self.lidar_speed = 0.25               
        self.lidar_kp = 0.35                  
        self.deadzone = 0.08                  
        
        self.control_timer = self.create_timer(0.1, self.control_loop)

        #   WIZUALIZACJA
        pygame.init()
        self.win_size = 600
        self.screen = pygame.display.set_mode((self.win_size, self.win_size))
        pygame.display.set_caption("HARVESTER")

    def odom_callback(self, msg: Odometry):
        self.robot_x = msg.pose.pose.position.x
        self.robot_y = msg.pose.pose.position.y
        
        q = msg.pose.pose.orientation
        siny_cosp = 2 * (q.w * q.z + q.x * q.y)
        cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z)
        self.robot_yaw = math.atan2(siny_cosp, cosy_cosp)
        
        if not self.has_odom:
            self.start_offset_x, self.start_offset_y = self.robot_x, self.robot_y
            gazebo_spawn_x, gazebo_spawn_y = -10.0, -2.0

            for wp in self.global_waypoints:
                self.local_waypoints.append({
                    "x": self.start_offset_x + (wp["x"] - gazebo_spawn_x),
                    "y": self.start_offset_y + (wp["y"] - gazebo_spawn_y)
                })
            self.has_odom = True

    def lidar_left_callback(self, msg: LaserScan):
        self.latest_left_scan = msg

    def lidar_right_callback(self, msg: LaserScan):
        self.latest_right_scan = msg

    def process_lidar_fusion(self):
        if not self.latest_left_scan or not self.latest_right_scan:
            return

        x_left_wall, y_left_wall = [], []
        x_right_wall, y_right_wall = [], []

        l_msg = self.latest_left_scan
        l_x, l_y, l_yaw = -1.0, 0.5, -0.785398

        
        for i, distance in enumerate(l_msg.ranges):
            if math.isinf(distance) or math.isnan(distance) or distance < 3.0 or distance > l_msg.range_max: 
                continue
            angle = l_msg.angle_min + i * l_msg.angle_increment
            
            xs = distance * math.cos(angle)
            ys = distance * math.sin(angle)
            
            xr = l_x + (xs * math.cos(l_yaw) - ys * math.sin(l_yaw))
            yr = l_y + (xs * math.sin(l_yaw) + ys * math.cos(l_yaw))
            
            if yr > 0.0:
                x_left_wall.append(xr)
                y_left_wall.append(yr)
            else:
                x_right_wall.append(xr)
                y_right_wall.append(yr)

        r_msg = self.latest_right_scan
        r_x, r_y, r_yaw = -1.0, -0.5, 0.785398  # <--- TUTAJ ZMIANA
        
        for i, distance in enumerate(r_msg.ranges):
            if math.isinf(distance) or math.isnan(distance) or distance < 3.0 or distance > r_msg.range_max: 
                continue
            angle = r_msg.angle_min + i * r_msg.angle_increment
            
            xs = distance * math.cos(angle)
            ys = distance * math.sin(angle)
            
            xr = r_x + (xs * math.cos(r_yaw) - ys * math.sin(r_yaw))
            yr = r_y + (xs * math.sin(r_yaw) + ys * math.cos(r_yaw))
            
            if yr > 0.0:
                x_left_wall.append(xr)
                y_left_wall.append(yr)
            else:
                x_right_wall.append(xr)
                y_right_wall.append(yr)

        self.pts_L = (x_left_wall, y_left_wall)
        self.pts_R = (x_right_wall, y_right_wall)

        min_points = 4
        if len(y_left_wall) < min_points and len(y_right_wall) < min_points:
            self.end_of_row_detected = True
            self.avg_left, self.avg_right = 1.4, 1.4
        else:
            self.end_of_row_detected = False
            self.avg_left = sum(y_left_wall) / len(y_left_wall) if y_left_wall else 1.4
            self.avg_right = abs(sum(y_right_wall) / len(y_right_wall)) if y_right_wall else 1.4

    def control_loop(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

        if self.current_state == self.STATE_INITIALIZING:
            if self.has_odom and self.latest_left_scan is not None and self.latest_right_scan is not None:
                self.current_state = self.STATE_GPS_DRIVE
            self.draw_pygame_window()
            return

        self.process_lidar_fusion()

        if self.current_state == self.STATE_GPS_DRIVE: self.execute_gps_drive()
        elif self.current_state == self.STATE_LIDAR_DRIVE: self.execute_lidar_drive()
        elif self.current_state == self.STATE_HARVESTING: self.execute_harvesting()
        elif self.current_state == self.STATE_MISSION_DONE: self.stop_robot()

        self.draw_pygame_window()

    def execute_gps_drive(self):
        if self.current_wp_idx >= len(self.local_waypoints):
            self.current_state = self.STATE_MISSION_DONE
            return

        target = self.local_waypoints[self.current_wp_idx]
        cmd = Twist()

        if not self.position_reached:
            dx, dy = target["x"] - self.robot_x, target["y"] - self.robot_y
            distance = math.sqrt(dx**2 + dy**2)
            desired_yaw = math.atan2(dy, dx)
            yaw_error = math.atan2(math.sin(desired_yaw - self.robot_yaw), math.cos(desired_yaw - self.robot_yaw))

            if distance < self.target_pos_tolerance:
                self.stop_robot()
                self.position_reached = True
                return

            if abs(yaw_error) > 0.15:
                cmd.angular.z = max(min(0.6 * yaw_error, 0.5), -0.5)
            else:
                cmd.linear.x = 2.5  
                cmd.angular.z = 0.8 * yaw_error  
            self.cmd_vel_pub.publish(cmd)
        else:
            yaw_error = math.atan2(math.sin(0.0 - self.robot_yaw), math.cos(0.0 - self.robot_yaw))
            if abs(yaw_error) < self.target_heading_tolerance:
                self.stop_robot()
                self.position_reached = False
                self.lidar_start_x, self.lidar_start_y = self.robot_x, self.robot_y
                self.lidar_enter_time = time.time()
                self.harvesting_done_in_this_row = False
                self.current_state = self.STATE_LIDAR_DRIVE
                return

            cmd.angular.z = max(min(0.5 * yaw_error, 0.4), -0.4)
            if 0 < cmd.angular.z < 0.15: cmd.angular.z = 0.15
            if -0.15 < cmd.angular.z < 0: cmd.angular.z = -0.15
            self.cmd_vel_pub.publish(cmd)

    def execute_lidar_drive(self):
        cmd = Twist()
        time_in_lidar = time.time() - self.lidar_enter_time
        
        if self.end_of_row_detected and time_in_lidar > 6.0:
            self.stop_robot()
            self.get_logger().info("RETURNING TO GPS.")
            self.current_wp_idx += 1  
            self.current_state = self.STATE_GPS_DRIVE
            return

        dist_traveled = math.sqrt((self.robot_x - self.lidar_start_x)**2 + (self.robot_y - self.lidar_start_y)**2)
        if dist_traveled >= self.harvest_trigger_distance and not self.harvesting_done_in_this_row:
            self.stop_robot()
            self.current_state = self.STATE_HARVESTING
            return

        diff = self.avg_left - self.avg_right
        cmd.linear.x = self.lidar_speed

        if abs(diff) < self.deadzone:
            cmd.angular.z = 0.0  
        else:
            cmd.angular.z = self.lidar_kp * diff  

        cmd.angular.z = max(min(cmd.angular.z, 0.2), -0.2)
        self.cmd_vel_pub.publish(cmd)

    def execute_harvesting(self):
        self.stop_robot()
        time.sleep(2.0)
        self.harvesting_done_in_this_row = True
        self.current_state = self.STATE_LIDAR_DRIVE

    def stop_robot(self):
        cmd = Twist()
        cmd.linear.x, cmd.angular.z = 0.0, 0.0
        self.cmd_vel_pub.publish(cmd)

    def draw_pygame_window(self):
        self.screen.fill((12, 16, 24))
        
        cx, cy = self.win_size // 2, self.win_size // 2 + 50 
        scale = 55.0  

        # Siatka radarowa
        for r in range(1, 8):
            pygame.draw.circle(self.screen, (25, 35, 45), (cx, cy), int(r * scale), 1)

        # Rysowanie obrysu kombajnu
        robot_w = int(3.2 * scale)
        robot_h = int(4.0 * scale)
        rx = cx - robot_w // 2
        ry = cy - robot_h // 2
        pygame.draw.rect(self.screen, (55, 65, 80), pygame.Rect(rx, ry, robot_w, robot_h), 2)
        
        # Pozycja montażowa lidarow
        pygame.draw.circle(self.screen, (0, 255, 0), (cx - int(1.5*scale), cy + int(2.0*scale)), 5)   # Lewy
        pygame.draw.circle(self.screen, (255, 150, 0), (cx - int(-1.5*scale), cy + int(2.0*scale)), 5) # Prawy

        font = pygame.font.SysFont('Monospace', 13, bold=True)
        if self.current_state == self.STATE_INITIALIZING:
            status_text = "INITIALIZING..."
            if not self.has_odom: status_text += " [ODOM]"
            if self.latest_left_scan is None: status_text += " [LIDAR_L]"
            if self.latest_right_scan is None: status_text += " [LIDAR_R]"
            self.screen.blit(font.render(status_text, True, (255, 140, 0)), (15, 15))
            pygame.display.flip()
            return

        # Rysowanie przeliczonych chmur punktów
        for x, y in zip(self.pts_L[0], self.pts_L[1]):
            px = cx - int(y * scale)
            py = cy - int(x * scale)
            if 0 <= px < self.win_size and 0 <= py < self.win_size:
                pygame.draw.circle(self.screen, (100, 255, 100), (px, py), 2)

        for x, y in zip(self.pts_R[0], self.pts_R[1]):
            px = cx - int(y * scale)
            py = cy - int(x * scale)
            if 0 <= px < self.win_size and 0 <= py < self.win_size:
                pygame.draw.circle(self.screen, (255, 165, 0), (px, py), 2)

        # Wyświetlanie tekstów diagnostycznych
        self.screen.blit(font.render(f"STAN: {self.current_state}", True, (255, 255, 255)), (15, 15))
        self.screen.blit(font.render(f"DL:  {self.avg_left:.2f}m", True, (100, 255, 100)), (15, 35))
        self.screen.blit(font.render(f"DR:  {self.avg_right:.2f}m", True, (255, 165, 0)), (15, 55))
        pygame.display.flip()

def main(args=None):
    rclpy.init(args=args)
    node = HarvesterStateMachine()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.stop_robot()
        node.destroy_node()
        rclpy.shutdown()
        pygame.quit()

if __name__ == '__main__':
    main()