import matplotlib.

pyplot as plt
import numpy as np
import time

# Constants
road_length = 300
car_length = 5
car_height = 0.5
lane_width = 2
desired_speed = 10 # Desired speed for IDM (m/s)
min_gap = 2 # Minimum gap to the leading vehicle for IDM (m)

# Initialize Matplotlib
plt.axis([0, road_length, 0, 10])
plt.gca().invert_yaxis()
plt.xlabel("Road")
plt.ylabel("Position")

# Create a list of cars (dictionaries)

cars = [{'lane':1,'car_position': 10, 'car_speed':25},
{'lane':1,'car_position': 50, 'car_speed':2},
{'lane':1,'car_position': 0, 'car_speed':25},
{'lane':1,'car_position': 10, 'car_speed':2},
{'lane':1,'car_position': 5, 'car_speed':20},
{'lane':1,'car_position': 6, 'car_speed':20},
{'lane':1,'car_position': 13, 'car_speed':20},
{'lane':1,'car_position': 1, 'car_speed':20},
{'lane':1,'car_position': 2, 'car_speed':20},
{'lane':0,'car_position': 2, 'car_speed':20},
{'lane':0,'car_position': 1, 'car_speed':20},
{'lane':0,'car_position': 5, 'car_speed':20},
{'lane':0,'car_position': 7, 'car_speed':20},
{'lane':0,'car_position': 10, 'car_speed':20},
]

lanes = {0:{'x0':0,'y0':1,'x1':road_length,'y1':1,'car_color':'red'},
1:{'x0':0,'y0':2,'x1':road_length,'y1':2,'car_color':'blue'}}

#plot lane centerline

for k in lanes.keys():
l = lanes[k]
plt.plot([l['x0'], l['x1']], [l['y0'], l['y1']], color="gray", linestyle="--",
linewidth=2, label="Road "+str(k))

# Create the car (rectangle)

car_patches = []
car_text = []
for i, car_info in enumerate(cars):
car_patch = plt.Rectangle((car_info['car_position'], lanes[car_info['lane']]
['y0']-car_height/2), car_length, car_height, fc=lanes[car_info['lane']]
['car_color'], label="Car")
car_patches.append(car_patch)
plt.gca().add_patch(car_patch)
car_label = plt.text(car_info['car_position']+car_length/2,
lanes[car_info['lane']]['y0'], str(i), fontsize=10, color='black', ha='center',
va='center')
car_text.append(car_label)

# Update car position and redraw

plt.legend(loc="upper right")
plt.draw()

# Function to calculate IDM acceleration. #r_time = reaction time #b_max =

comfortable deceleration for vehicle
def idm_acceleration(v_lead, v, desired_speed, dist, desired_dist, a_max=0.73,
delta=4,r_time=1.5,b_max=1.67):
term1 = a_max * (1 - (v/desired_speed)**delta)
term2 = a_max*((desired_dist+v*r_time)/dist +
v*(v-v_lead)/(2*dist*(a_max*b_max)**0.5))**2
acceleration = term1 - term2
return acceleration

# Simulation time step

dt = 0.1 # Time step (in seconds)

# Simulate car movements

simulation_time = 0
num_cars = len(cars)
while all(car['car_position'] < road_length - car_length for car in cars):
print("-----")
for i, car_info in enumerate(cars):
# Calculate gap to the leading vehicle
cars[i]['gap_to_leading'] = float('inf')
gap_to_leading = float('inf') # Initialize to infinity

cars[i]['lane_change_lead'] = 1000
cars[i]['lane_change_follow'] = 1000

for j in range(0, num_cars):

#only consider cars in the same lane for gap_to_leading calculation
if cars[j]['lane']==cars[i]['lane']:
gap_candidate = cars[j]['car_position'] - cars[i]['car_position']
if gap_candidate>0 and gap_candidate<cars[i]['gap_to_leading']:
cars[i]['gap_to_leading'] = gap_candidate
cars[i]['leading_car_speed'] = cars[j]['car_speed']

if cars[j]['lane']!=cars[i]['lane']:
if cars[j]['car_position']>=cars[i]['car_position']:
cars[i]['lane_change_lead'] = min(cars[j]['car_position'] -
cars[i]['car_position'], cars[i]['lane_change_lead'])
if cars[j]['car_position']<cars[i]['car_position']:
cars[i]['lane_change_follow'] = min(cars[i]['car_position'] -
cars[j]['car_position'], cars[i]['lane_change_follow'])

#lane change
if cars[i]['lane_change_follow']>5 and cars[i]['lane_change_lead']>cars[i]
['gap_to_leading'] and simulation_time>20:
cars[i]['lane'] = 1- cars[i]['lane']
car_patches[i].set_y(lanes[cars[i]['lane']]['y0']-car_height/2)
car_text[i].set_y(lanes[cars[i]['lane']]['y0'])
cars[i]['gap_to_leading'] = cars[i]['lane_change_lead']
car_patches[i].set_facecolor(lanes[cars[i]['lane']]['car_color'])

# If there are no cars ahead, set gap to the road length

if cars[i]['gap_to_leading'] == float('inf'):
cars[i]['gap_to_leading'] = road_length - cars[i]['car_position']
cars[i]['leading_car_speed'] = 25
 # Calculate IDM acceleration
acceleration = idm_acceleration(cars[i]['leading_car_speed'],cars[i]
['car_speed'], desired_speed, cars[i]['gap_to_leading'], min_gap)
cars[i]['acceleration'] = acceleration
cars[i]['car_speed'] = max(cars[i]['car_speed'] + cars[i]['acceleration'] *
dt, 0)
cars[i]['car_position'] = cars[i]['car_position'] + cars[i]['car_speed'] *
dt

car_patches[i].set_x(cars[i]['car_position'])
car_text[i].set_x(cars[i]['car_position']+car_length/2)
simulation_time += dt

for c in cars:
print(c)

plt.pause(0.1)

# Keep the plot window open

plt.show()