main.py

import sys
import time
import sim
import keyboard
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm

from slam.occupancy_grid import OccupancyGrid
from slam.localization     import Localization
from slam.mapping          import Mapping
from robot.controller      import RobotController
from robot.sensors         import RobotSensors

def connect_to_coppeliasim():
    sim.simxFinish(-1)
    clientID = sim.simxStart('127.0.0.1', 19999, True, True, 5000, 5)
    if clientID == -1:
        sys.exit("[ERROR] Failed to connect to CoppeliaSim")
    return clientID

def get_handles(clientID):
    _, left_motor   = sim.simxGetObjectHandle(
        clientID, '/PioneerP3DX/leftMotor', sim.simx_opmode_oneshot_wait
    )
    _, right_motor  = sim.simxGetObjectHandle(
        clientID, '/PioneerP3DX/rightMotor', sim.simx_opmode_oneshot_wait
    )
    _, robot_handle = sim.simxGetObjectHandle(
        clientID, '/PioneerP3DX',           sim.simx_opmode_oneshot_wait
    )
    err, lidar      = sim.simxGetObjectHandle(
        clientID, 'Hokuyo',                 sim.simx_opmode_blocking
    )
    if err != 0:
        sys.exit("[ERROR] Failed to get LiDAR sensor handle")
    return left_motor, right_motor, robot_handle, lidar

def get_robot_pose(clientID, handle):
    err_p, pos = sim.simxGetObjectPosition(
        clientID, handle, -1, sim.simx_opmode_buffer
    )
    err_o, ori = sim.simxGetObjectOrientation(
        clientID, handle, -1, sim.simx_opmode_buffer
    )
    if err_p == 0 and err_o == 0:
        return pos[0], pos[1], ori[2]
    return None

def main():
    global occ, mapping
    WHEEL_BASE   = 0.381
    dt           = 0.1
    grid_m       = 10
    cell_m       = 0.1
    lidar_fov    = 270
    max_range    = 10.0
    W = int(grid_m / cell_m)
    H = int(grid_m / cell_m)
    origin = (-5, -5)

    occ     = OccupancyGrid(W, H, cell_m, origin, max_range)
    mapping = Mapping(occ)

    print("Program Started")
    clientID = connect_to_coppeliasim()
    left, right, robot_handle, lidar = get_handles(clientID)
    print("[OK] Connected to CoppeliaSim")

    err_p, pos = sim.simxGetObjectPosition(
        clientID, robot_handle, -1, sim.simx_opmode_blocking
    )
    err_o, ori = sim.simxGetObjectOrientation(
        clientID, robot_handle, -1, sim.simx_opmode_blocking
    )
    if err_p != 0 or err_o != 0:
        sys.exit("[ERROR] Failed to get initial pose")
    odom_x, odom_y, odom_theta = pos[0], pos[1], ori[2]

    controller   = RobotController(clientID, left, right, WHEEL_BASE)
    sensors      = RobotSensors(clientID, lidar)
    localization = Localization()

    packed_data = sim.simxGetStringSignal(
        clientID, 'laserScanData', sim.simx_opmode_buffer
    )[1]
    if packed_data:
        float_data = sim.simxUnpackFloats(packed_data)
        sensor_data = [(float_data[i], float_data[i+1]) for i in range(0, len(float_data), 2)]
        occ.update(odom_x, odom_y, odom_theta, sensor_data)

    sim.simxGetObjectPosition(clientID, robot_handle, -1, sim.simx_opmode_streaming)
    sim.simxGetObjectOrientation(clientID, robot_handle, -1, sim.simx_opmode_streaming)
    sim.simxGetStringSignal(clientID, 'laserScanData', sim.simx_opmode_streaming)
    time.sleep(1)
    print("[INFO] Streaming started.")

    plt.ion()
    fig_map, ax_map = plt.subplots(figsize=(6,6), dpi=100)
    ax_map.set_xlim(-5, W)
    ax_map.set_ylim(-5, H)
    ax_map.set_aspect('equal', 'box')

    cmap = ListedColormap(['white', 'grey', 'black'])
    norm = BoundaryNorm([0.0, 0.4, 0.6, 1.0], cmap.N)

    im = ax_map.imshow(occ.grid_to_prob(), cmap=cmap, norm=norm, origin='lower')
    cbar = fig_map.colorbar(im, ax=ax_map, ticks=[0.0, 0.5, 1.0])
    cbar.ax.set_yticklabels(['free (<0.4)', 'unknown (0.4–0.6)', 'occ (>0.6)'])

    ax_map.set_title("Occupancy Grid with Robot")
    ax_map.set_xlabel("Grid X (cells)")
    ax_map.set_ylabel("Grid Y (cells)")

    robot_dot, = ax_map.plot([-1], [-1], 'bo', markersize=6)
    quiv = ax_map.quiver([0], [0], [1], [0],
                         color='red', angles='xy', scale_units='xy', scale=1)

    # Trajectory
    trajectory_x = []
    trajectory_y = []
    trajectory_plot, = ax_map.plot([], [], 'r-', linewidth=1)

    # Initialize frontier‐selection variables
    current_goal   = None
    goal_tolerance = 0.5   
    last_goal      = None   
    cached_num_ranges = None
    cached_angles     = None

    stuck_counter = 0
    last_pose = (odom_x, odom_y)

    try:
        frame_count = 0
        while True:
            print(f"[Current_Goal]: current_goal = {current_goal}")
            teleop = False
            vL = vR = 0.0
            if   keyboard.is_pressed('up'):
                vL = vR =  1.0; teleop = True
            elif keyboard.is_pressed('down'):
                vL = vR = -1.0; teleop = True
            elif keyboard.is_pressed('left'):
                vL, vR = -1.0, 1.0; teleop = True
            elif keyboard.is_pressed('right'):
                vL, vR = 1.0, -1.0; teleop = True
            elif keyboard.is_pressed('esc'):
                break
            if teleop:
                controller.set_speed(vL, vR)
                continue

            pose = get_robot_pose(clientID, robot_handle)
            if pose is not None:
                odom_x, odom_y, odom_theta = pose
                mx, my = occ.world_to_grid(odom_x, odom_y)
                print(f"[DEBUG] Robot world pose: ({odom_x:.2f}, {odom_y:.2f}), grid cell: ({mx}, {my})")

                
                trajectory_x.append(mx)
                trajectory_y.append(my)
                trajectory_plot.set_data(trajectory_x, trajectory_y)

            raw = sensors.get_lidar_data()
            if raw is None:
                time.sleep(dt)
                continue

            ranges = sensors.process_lidar_data(raw)
            if cached_num_ranges != len(ranges):
                cached_num_ranges = len(ranges)
                cached_angles     = np.deg2rad(
                    np.linspace(-lidar_fov/2, lidar_fov/2, len(ranges))
                )
            angles = cached_angles

            sensor_data = list(zip(ranges, angles))
            occ.update(odom_x, odom_y, odom_theta, sensor_data)
            prob = occ.grid_to_prob()
            print(
                "[DEBUG] Free:", np.sum(prob < 0.4),
                "Occupied:", np.sum(prob > 0.6),
                "Unknown:", np.sum((prob >= 0.4) & (prob <= 0.6))
            )

            pose_est = localization.estimate(ranges, angles, odom=(odom_x, odom_y, odom_theta))
            
            #loop closure detection
            loop_point = localization.detect_loop_closure()
            if loop_point:
                lx, ly = occ.world_to_grid(*loop_point)
                ax_map.plot(lx, ly, 'go', markersize=8)  # green dot for loop closure
                ax_map.annotate('Loop', (lx, ly), color='green', fontsize=8)

                # Optional pose correction (snap to matched loop pose)
                localization.x, localization.y = loop_point
                print("[LOOP] Corrected current pose using loop closure.")

            mapping.update(pose_est, ranges, angles)

            if current_goal is None:
                current_goal = mapping.frontier((odom_x, odom_y))
                print(f"[DEBUG] mapping.frontier() returned: {current_goal}")
                if current_goal is not None and current_goal != last_goal:
                    print(f"[GOAL] New goal set: ({current_goal[0]:.2f}, {current_goal[1]:.2f})")
                    last_goal = current_goal
            else:
                if not mapping.is_frontier(current_goal[0], current_goal[1]):
                    print("[INFO] Current goal is no longer a frontier. Selecting new goal.")
                    current_goal = mapping.frontier((odom_x, odom_y))

            if current_goal is not None:
                gx, gy = current_goal
                dist = np.hypot(gx - odom_x, gy - odom_y)
                print(f"[***] Distance to goal: {dist:.2f}")

                if dist < goal_tolerance:
                    print("[GOAL] Reached frontier goal!")
                    controller.compute_velocity(pose_est, current_goal)
                    current_goal = None
                else:
                    vL2, vR2 = controller.compute_velocity(pose_est, current_goal)
                    print(f"[WHEELS_VELOCITIES] Controller output: vL2={vL2:.2f}, vR2={vR2:.2f}")
                    vL2 = np.clip(vL2, -1.0, 1.0)
                    vR2 = np.clip(vR2, -1.0, 1.0)
                    controller.set_speed(vL2, vR2)
            else:
                controller.stop()

            if np.hypot(odom_x - last_pose[0], odom_y - last_pose[1]) < 0.05:
                stuck_counter += 1
            else:
                stuck_counter = 0
            last_pose = (odom_x, odom_y)

            if stuck_counter > 20:
                print("[RECOVERY] Robot appears stuck. Selecting new frontier.")
                current_goal = mapping.frontier((odom_x, odom_y))
                stuck_counter = 0

            im.set_data(prob)
            fig_map.canvas.draw_idle()
            plt.pause(0.001)

            mx, my = occ.world_to_grid(odom_x, odom_y)
            robot_dot.set_data([mx], [my])
            dx = 5 * np.cos(odom_theta)
            dy = 5 * np.sin(odom_theta)
            quiv.set_offsets([[mx, my]])
            quiv.set_UVC([dx], [dy])

            fig_map.canvas.draw_idle()
            plt.pause(0.001)

            frame_count += 1
            time.sleep(dt)

    except KeyboardInterrupt:
        print("[INFO] Interrupted by user.")
    finally:
        controller.stop()
        sim.simxFinish(clientID)
        plt.ioff()
        plt.show()
        print("[INFO] Shut down cleanly.")

if __name__ == "__main__":
    main()