simulation_env_enhanced.py

import pybullet as p
import pybullet_data
import time
import numpy as np
from PIL import Image
import os
import math

# Global variables to store simulation state
obstacleId = None
planeId = None

def setup_pybullet_environment():
    """Setup PyBullet environment with ground plane and obstacle"""
    global obstacleId, planeId
    
    # Connect to PyBullet physics server
    p.connect(p.GUI)
    
    # Set additional search path for URDF files
    p.setAdditionalSearchPath(pybullet_data.getDataPath())
    
    # Set gravity
    p.setGravity(0, 0, -9.81)
    
    # Load ground plane
    planeId = p.loadURDF("plane.urdf")
    
    # Create obstacle - Box: width=0.5m, depth=0.1m, height=0.05m at x=0.75m
    obstacle_collision_shape = p.createCollisionShape(
        p.GEOM_BOX, 
        halfExtents=[0.25, 0.05, 0.025]  # Half extents for box
    )
    obstacle_visual_shape = p.createVisualShape(
        p.GEOM_BOX,
        halfExtents=[0.25, 0.05, 0.025],
        rgbaColor=[1, 0, 0, 1]  # Red color
    )
    
    obstacleId = p.createMultiBody(
        baseMass=0,  # Static obstacle
        baseCollisionShapeIndex=obstacle_collision_shape,
        baseVisualShapeIndex=obstacle_visual_shape,
        basePosition=[0.75, 0, 0.025]  # Centered at x=0.75m, base on ground
    )
    
    # Set camera view
    p.resetDebugVisualizerCamera(
        cameraDistance=2.0,
        cameraYaw=0,
        cameraPitch=-30,
        cameraTargetPosition=[0.5, 0, 0]
    )
    
    return obstacleId, planeId

def create_vehicle(vehicle_specs):
    """Create vehicle (robot or drone) based on specifications"""
    vehicle_type = vehicle_specs.get("vehicle_type", "robot")
    
    if vehicle_type == "robot":
        return create_robot(vehicle_specs)
    elif vehicle_type == "drone":
        return create_drone(vehicle_specs)
    else:
        raise ValueError(f"Unknown vehicle type: {vehicle_type}")

def create_robot(robot_specs):
    """Create robot based on specifications with corrected wheel orientation and joint axis."""
    # Extract specifications with defaults
    wheel_type = robot_specs.get("wheel_type", "large_smooth")
    body_clearance_cm = robot_specs.get("body_clearance_cm", 7)
    # approach_sensor_enabled = robot_specs.get("approach_sensor_enabled", True) # Not used directly in physics
    main_material = robot_specs.get("main_material", "light_plastic")
    
    # Wheel parameters
    wheel_params = {
        "small_high_grip": {"radius": 0.06, "friction": 1.5, "width": 0.03},
        "large_smooth": {"radius": 0.07, "friction": 0.8, "width": 0.04},
        "tracked_base": {"radius": 0.065, "friction": 2.0, "width": 0.05} # Simplified as wide wheels
    }
    
    wheel_config = wheel_params.get(wheel_type, wheel_params["large_smooth"])
    wheel_radius = wheel_config["radius"]
    wheel_friction = wheel_config["friction"]
    wheel_width = wheel_config["width"] # This is the thickness of the wheel
    
    # Body parameters
    body_length = 0.25 # Along X
    body_width = 0.20  # Along Y
    body_height = 0.04 # Along Z
    
    obstacle_height = 0.05
    min_clearance = obstacle_height + 0.01
    body_clearance = max(body_clearance_cm / 100.0, min_clearance) # This is vertical clearance
    
    # Calculate Z position for the body's center
    # The wheels should touch the ground (Z=0), so:
    # wheel_center_z = wheel_radius (center of wheel above ground)
    # body_bottom_z = wheel_center_z + body_clearance
    # body_center_z = body_bottom_z + body_height/2
    body_center_z_pos = wheel_radius + body_clearance + (body_height / 2.0)

    # Material properties
    material_mass = {
        "light_plastic": 2.0,
        "sturdy_metal_alloy": 3.0
    }
    body_mass = material_mass.get(main_material, 2.0)
    wheel_mass = 0.3
    
    # Create collision shapes
    body_collision_shape = p.createCollisionShape(
        p.GEOM_BOX,
        halfExtents=[body_length/2, body_width/2, body_height/2]
    )
    # For GEOM_CYLINDER, the 'height' is along the cylinder's local Z-axis.
    # 'radius' is in its local XY plane.
    # We want wheels whose flat sides are perpendicular to the Y-axis of the robot.
    # So, the cylinder's length (PyBullet's 'height' param) should be wheel_width.
    wheel_collision_shape = p.createCollisionShape(
        p.GEOM_CYLINDER,
        radius=wheel_radius,
        height=wheel_width # This is the thickness of the wheel disk
    )
    
    # Create visual shapes
    body_visual_shape = p.createVisualShape(
        p.GEOM_BOX,
        halfExtents=[body_length/2, body_width/2, body_height/2],
        rgbaColor=[0, 0, 1, 1]
    )
    # For visual shape, 'length' corresponds to cylinder's Z-axis length
    wheel_visual_shape = p.createVisualShape(
        p.GEOM_CYLINDER,
        radius=wheel_radius,
        length=wheel_width, # Thickness of the wheel disk
        rgbaColor=[0.2, 0.2, 0.2, 1]
    )
    
    # Link properties for two wheels
    link_masses = [wheel_mass, wheel_mass]
    link_collision_shape_indices = [wheel_collision_shape, wheel_collision_shape]
    link_visual_shape_indices = [wheel_visual_shape, wheel_visual_shape]
    
    # Position of wheel links relative to the robot's base (body) center
    # Wheels should be at ground level (Z=0), so their centers are at Z=wheel_radius
    # Relative to body center: wheel_z = wheel_radius - body_center_z_pos
    wheel_link_z_offset = wheel_radius - body_center_z_pos # This should be negative

    link_positions = [
        [0, body_width/2 + wheel_width/2, wheel_link_z_offset], # Right wheel
        [0, -(body_width/2 + wheel_width/2), wheel_link_z_offset] # Left wheel
    ]
    
    # Orientation of wheel links:
    # Use identity orientation (no rotation) and set joint axis to X directly
    wheel_link_orientation_quat = p.getQuaternionFromEuler([0, 0, 0])  # No rotation
    link_orientations = [wheel_link_orientation_quat, wheel_link_orientation_quat]
    
    link_inertial_frame_positions = [[0, 0, 0], [0, 0, 0]] # Relative to link frame
    link_inertial_frame_orientations = [[0,0,0,1], [0,0,0,1]] # Identity quaternion
    link_parent_indices = [0, 0] # Both wheels attached to the base (index 0 for links)
    link_joint_types = [p.JOINT_REVOLUTE, p.JOINT_REVOLUTE]
    
    # Joint Axis: Use X-axis directly for forward motion
    link_joint_axis = [[1, 0, 0], [1, 0, 0]]
    
    robotId = p.createMultiBody(
        baseMass=body_mass,
        baseCollisionShapeIndex=body_collision_shape,
        baseVisualShapeIndex=body_visual_shape,
        basePosition=[0, 0, body_center_z_pos], # Initial position of the base
        baseOrientation=[0,0,0,1],
        linkMasses=link_masses,
        linkCollisionShapeIndices=link_collision_shape_indices,
        linkVisualShapeIndices=link_visual_shape_indices,
        linkPositions=link_positions,
        linkOrientations=link_orientations,
        linkInertialFramePositions=link_inertial_frame_positions,
        linkInertialFrameOrientations=link_inertial_frame_orientations,
        linkParentIndices=link_parent_indices,
        linkJointTypes=link_joint_types,
        linkJointAxis=link_joint_axis
    )
    
    # Set dynamics properties for body
    p.changeDynamics(robotId, -1, 
                    lateralFriction=0.8,
                    spinningFriction=0.1,
                    rollingFriction=0.05, # Rolling friction for the body itself if it contacts
                    linearDamping=0.1, 
                    angularDamping=0.3)
    
    # Set dynamics properties for wheels
    # Joint indices for createMultiBody start from 0 for the first link.
    wheel_joint_indices = [0, 1] 
    for joint_idx in wheel_joint_indices:
        p.changeDynamics(robotId, joint_idx,
                        lateralFriction=wheel_friction, # Friction against sideways slip
                        spinningFriction=0.05, 
                        rollingFriction=0.001, # Low rolling friction for the wheel itself
                        linearDamping=0.05,
                        angularDamping=0.1)
        # Enable motor for velocity control
        p.setJointMotorControl2(robotId, joint_idx, p.VELOCITY_CONTROL, force=0)


    print(f"Created robot: body_z_pos={body_center_z_pos:.3f}m, wheel_radius={wheel_radius:.3f}m, actual_clearance_under_body={(body_center_z_pos - body_height/2 - wheel_radius):.3f}m")
    
    return robotId, wheel_joint_indices, "robot"

def create_drone(drone_specs):
    """Create drone based on specifications"""
    # Extract specifications with defaults
    propeller_size = drone_specs.get("propeller_size", "medium")
    flight_height_cm = drone_specs.get("flight_height_cm", 20)
    stability_mode = drone_specs.get("stability_mode", "auto_hover")
    main_material = drone_specs.get("main_material", "light_carbon_fiber")
    
    # Propeller parameters
    propeller_params = {
        "small_agile": {"radius": 0.05, "thrust_coeff": 1.2, "mass": 0.1},
        "medium": {"radius": 0.08, "thrust_coeff": 1.5, "mass": 0.15},
        "large_stable": {"radius": 0.12, "thrust_coeff": 2.0, "mass": 0.2}
    }
    
    prop_config = propeller_params.get(propeller_size, propeller_params["medium"])
    prop_radius = prop_config["radius"]
    thrust_coeff = prop_config["thrust_coeff"]
    prop_mass = prop_config["mass"]
    
    # Body parameters
    body_length = 0.20
    body_width = 0.20
    body_height = 0.05
    flight_height = max(flight_height_cm / 100.0, 0.15)  # Minimum 15cm flight height
    
    # Material properties
    material_mass = {
        "light_carbon_fiber": 0.8,
        "sturdy_aluminum": 1.2
    }
    body_mass = material_mass.get(main_material, 0.8)
    
    # Calculate starting position - above obstacle
    body_z_pos = flight_height
    prop_offset = body_length / 2 + prop_radius / 2
    
    # Create collision shapes
    body_collision_shape = p.createCollisionShape(
        p.GEOM_BOX,
        halfExtents=[body_length/2, body_width/2, body_height/2]
    )
    
    prop_collision_shape = p.createCollisionShape(
        p.GEOM_CYLINDER,
        radius=prop_radius,
        height=0.01  # Very thin propellers
    )
    
    # Create visual shapes
    body_visual_shape = p.createVisualShape(
        p.GEOM_BOX,
        halfExtents=[body_length/2, body_width/2, body_height/2],
        rgbaColor=[0, 1, 0, 1]  # Green body for drone
    )
    
    prop_visual_shape = p.createVisualShape(
        p.GEOM_CYLINDER,
        radius=prop_radius,
        length=0.01,
        rgbaColor=[0.1, 0.1, 0.1, 0.8]  # Semi-transparent dark propellers
    )
    
    # Create the drone with body and 4 propellers
    link_masses = [prop_mass] * 4  # Four propellers
    link_collision_shape_indices = [prop_collision_shape] * 4
    link_visual_shape_indices = [prop_visual_shape] * 4
    link_positions = [
        [prop_offset, prop_offset, 0.03],    # Front right
        [-prop_offset, prop_offset, 0.03],   # Front left  
        [-prop_offset, -prop_offset, 0.03],  # Rear left
        [prop_offset, -prop_offset, 0.03]    # Rear right
    ]
    link_orientations = [[0, 0, 0, 1]] * 4
    link_inertial_frame_positions = [[0, 0, 0]] * 4
    link_inertial_frame_orientations = [[0, 0, 0, 1]] * 4
    link_parent_indices = [0, 0, 0, 0]  # All propellers connected to base
    link_joint_types = [p.JOINT_REVOLUTE] * 4  # Revolute joints for propellers
    link_joint_axis = [[0, 0, 1]] * 4  # Rotate around Z-axis (vertical)
    
    droneId = p.createMultiBody(
        baseMass=body_mass,
        baseCollisionShapeIndex=body_collision_shape,
        baseVisualShapeIndex=body_visual_shape,
        basePosition=[0, 0, body_z_pos],
        linkMasses=link_masses,
        linkCollisionShapeIndices=link_collision_shape_indices,
        linkVisualShapeIndices=link_visual_shape_indices,
        linkPositions=link_positions,
        linkOrientations=link_orientations,
        linkInertialFramePositions=link_inertial_frame_positions,
        linkInertialFrameOrientations=link_inertial_frame_orientations,
        linkParentIndices=link_parent_indices,
        linkJointTypes=link_joint_types,
        linkJointAxis=link_joint_axis
    )
    
    # Set dynamics properties for body
    p.changeDynamics(droneId, -1,
                    lateralFriction=0.1,
                    spinningFriction=0.1,
                    rollingFriction=0.1,
                    linearDamping=0.3,
                    angularDamping=0.5)
    
    # Set dynamics properties for propellers
    for prop_idx in range(4):
        p.changeDynamics(droneId, prop_idx,
                        lateralFriction=0.1,
                        spinningFriction=0.05,
                        rollingFriction=0.01,
                        linearDamping=0.2,
                        angularDamping=0.3)
    
    # Store thrust coefficient for flight control
    drone_props = {
        "thrust_coeff": thrust_coeff,
        "target_height": flight_height,
        "stability_mode": stability_mode
    }
    
    propeller_joint_indices = [0, 1, 2, 3]  # Joint indices for the four propellers
    
    print(f"Created drone: flight_height={flight_height:.3f}m, prop_radius={prop_radius:.3f}m, thrust_coeff={thrust_coeff}")
    
    return droneId, propeller_joint_indices, "drone", drone_props

def run_simulation_step(vehicleId, joint_indices, control_params, vehicle_type="robot", vehicle_props=None):
    """Run one simulation step with vehicle control"""
    if vehicle_type == "robot":
        run_robot_simulation_step(vehicleId, joint_indices, control_params)
    elif vehicle_type == "drone":
        run_drone_simulation_step(vehicleId, joint_indices, control_params, vehicle_props)
    
    # Step simulation
    p.stepSimulation()
    time.sleep(1./240.)  # 240 Hz simulation

def run_robot_simulation_step(robotId, wheel_joint_indices, control_params):
    """Run robot simulation step with wheel control"""
    if wheel_joint_indices:  # Robot has wheel joints
        # Set target velocity for forward motion - TRY POSITIVE direction
        target_velocity = 5.0  # rad/s - positive for forward motion
        max_force = 50.0  # Nm - much more torque for climbing
        
        # Apply velocity control to both wheels for forward motion
        for joint_idx in wheel_joint_indices:
            p.setJointMotorControl2(
                robotId,
                joint_idx,
                p.VELOCITY_CONTROL,
                targetVelocity=target_velocity,
                force=max_force
            )
    else:
        # Fallback: apply direct force
        force_magnitude = 5.0  # Newtons
        p.applyExternalForce(robotId, -1, [force_magnitude, 0, 0], [0, 0, 0], p.WORLD_FRAME)

def run_drone_simulation_step(droneId, propeller_joint_indices, control_params, drone_props):
    """Run drone simulation step with flight control"""
    # Get current drone state
    drone_pos, drone_orn = p.getBasePositionAndOrientation(droneId)
    drone_vel, drone_ang_vel = p.getBaseVelocity(droneId)
    
    target_height = drone_props.get("target_height", 0.2)
    thrust_coeff = drone_props.get("thrust_coeff", 1.5)
    stability_mode = drone_props.get("stability_mode", "auto_hover")
    
    # Calculate thrust needed for hovering
    drone_mass = p.getDynamicsInfo(droneId, -1)[0]
    gravity_force = drone_mass * 9.81
    base_thrust_per_prop = gravity_force / 4  # Four propellers
    
    # Height control (PID-like)
    height_error = target_height - drone_pos[2]
    height_thrust_correction = height_error * 10.0  # Higher proportional gain
    
    # Forward motion - apply body force instead of individual propeller forces
    forward_force = 3.0  # Newtons - direct forward force
    
    # Apply main thrust for hovering
    total_thrust = (base_thrust_per_prop + height_thrust_correction) * thrust_coeff
    
    # Apply upward thrust at drone center
    if total_thrust > 0:
        p.applyExternalForce(
            droneId, -1,
            [0, 0, total_thrust * 4],  # Total upward force
            [0, 0, 0],  # At center of mass
            p.WORLD_FRAME
        )
    
    # Apply forward force directly to drone body
    p.applyExternalForce(
        droneId, -1,
        [forward_force, 0, 0],  # Forward force
        [0, 0, 0],  # At center of mass
        p.WORLD_FRAME
    )
    
    # Add slight damping to prevent oscillations
    linear_damping = -0.1
    p.applyExternalForce(
        droneId, -1,
        [drone_vel[0] * linear_damping, drone_vel[1] * linear_damping, 0],
        [0, 0, 0],
        p.WORLD_FRAME
    )
    
    # Spin propellers for visual effect
    if propeller_joint_indices:
        for prop_idx in propeller_joint_indices:
            p.setJointMotorControl2(
                droneId,
                prop_idx,
                p.VELOCITY_CONTROL,
                targetVelocity=20.0,  # Fast spinning for visual effect
                force=0.1
            )

def get_simulation_feedback(vehicleId, obstacleId, start_time, current_sim_time, vehicle_type="robot"):
    """Get feedback from current simulation state"""
    # Get vehicle position and orientation
    vehicle_pos, vehicle_orn = p.getBasePositionAndOrientation(vehicleId)
    
    # Check if vehicle is upright/stable
    euler_angles = p.getEulerFromQuaternion(vehicle_orn)
    roll, pitch, yaw = euler_angles
    
    if vehicle_type == "robot":
        # Robot is considered upright if roll and pitch are small
        is_stable = abs(roll) < 0.5 and abs(pitch) < 0.5
    elif vehicle_type == "drone":
        # Drone is considered stable if not completely inverted and at reasonable height
        is_stable = abs(roll) < 1.0 and abs(pitch) < 1.0 and vehicle_pos[2] > 0.05
    else:
        is_stable = True
    
    # Check for contacts with obstacle
    contact_points = p.getContactPoints(vehicleId, obstacleId)
    obstacle_contacts_exist = len(contact_points) > 0
    
    feedback = {
        "robot_position": list(vehicle_pos),  # Keep "robot_position" for compatibility
        "robot_orientation_quaternion": list(vehicle_orn),
        "obstacle_contacts_exist": obstacle_contacts_exist,
        "is_robot_upright": is_stable,  # Keep "is_robot_upright" for compatibility
        "current_sim_time_sec": current_sim_time,
        "vehicle_type": vehicle_type
    }
    
    return feedback

def reset_simulation():
    """Reset and disconnect from PyBullet simulation"""
    p.resetSimulation()
    p.disconnect()

def capture_frame():
    """Capture current frame from simulation"""
    # Get camera image
    width, height, rgb_img, depth_img, seg_img = p.getCameraImage(
        width=640,
        height=480,
        viewMatrix=p.computeViewMatrixFromYawPitchRoll(
            cameraTargetPosition=[0.5, 0, 0],
            distance=2.0,
            yaw=0,
            pitch=-30,
            roll=0,
            upAxisIndex=2
        ),
        projectionMatrix=p.computeProjectionMatrixFOV(
            fov=60,
            aspect=640/480,
            nearVal=0.1,
            farVal=100.0
        )
    )
    
    # Convert to PIL Image
    rgb_array = np.array(rgb_img).reshape(height, width, 4)[:, :, :3]  # Remove alpha channel
    image = Image.fromarray(rgb_array, 'RGB')
    
    return image

def get_obstacle_info():
    """Get information about the obstacle"""
    return {
        "width_m": 0.5,
        "depth_m": 0.1,
        "height_m": 0.05,
        "position_x_m": 0.75,
        "position_y_m": 0.0,
        "position_z_m": 0.025,
        "material": "static_red_box",
        "success_threshold_x_m": 0.8
    }