import matplotlib.pyplot as plt
import matplotlib.patches as patches

# Create a figure and axis
fig, ax = plt.subplots(figsize=(8, 10))

# Define the grid size
rows, cols = 5, 4

# Define the colors for clear and not clear tiles
clear_color = 'lightgreen'
not_clear_color = 'salmon'

# Define the robot colors
robot1_color = 'white'
robot2_color = 'black'

# Define the tile status
tile_status = {
    'tile_0-1': 'clear', 'tile_0-2': 'clear', 'tile_0-3': 'clear', 'tile_0-4': 'clear',
    'tile_1-1': 'not clear', 'tile_1-2': 'not clear', 'tile_1-3': 'clear', 'tile_1-4': 'clear',
    'tile_2-1': 'clear', 'tile_2-2': 'clear', 'tile_2-3': 'clear', 'tile_2-4': 'clear',
    'tile_3-1': 'clear', 'tile_3-2': 'clear', 'tile_3-3': 'clear', 'tile_3-4': 'clear',
    'tile_4-1': 'clear', 'tile_4-2': 'clear', 'tile_4-3': 'clear', 'tile_4-4': 'clear',
}

# Define robot positions
robot_positions = {
    'robot1': (1, 0),
    'robot2': (1, 1),
}

# Draw the tiles
for row in range(rows):
    for col in range(cols):
        tile_id = f'tile_{row}-{col+1}'
        status = tile_status[tile_id]
        color = clear_color if status == 'clear' else not_clear_color
        rect = patches.Rectangle((col, rows-row-1), 1, 1, linewidth=1, edgecolor='black', facecolor=color)
        ax.add_patch(rect)
        ax.text(col + 0.5, rows-row-0.5, tile_id, ha='center', va='center', fontsize=8, color='black')

# Draw the robots
for robot, (row, col) in robot_positions.items():
    color = robot1_color if robot == 'robot1' else robot2_color
    circle = patches.Circle((col + 0.5, rows-row-0.5), 0.3, edgecolor='black', facecolor=color)
    ax.add_patch(circle)
    ax.text(col + 0.5, rows-row-0.7, robot, ha='center', va='center', fontsize=8, color='black' if color == 'white' else 'white')

# Set limits and labels
ax.set_xlim(0, cols)
ax.set_ylim(0, rows)
ax.set_xticks([])
ax.set_yticks([])

# Add legend
legend_elements = [
    patches.Patch(facecolor=clear_color, edgecolor='black', label='Clear Tile'),
    patches.Patch(facecolor=not_clear_color, edgecolor='black', label='Not Clear Tile'),
    patches.Patch(facecolor=robot1_color, edgecolor='black', label='Robot1 (White)'),
    patches.Patch(facecolor=robot2_color, edgecolor='black', label='Robot2 (Black)'),
]
ax.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(1.3, 1))

# Save the figure
plt.savefig('tiles/one_shot/ini_diagram_code/attempts/attempt_19.png', bbox_inches='tight')
plt.show()