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

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

# Define the grid size
rows, cols = 3, 3

# Define the colors for clear and painted tiles
clear_color = 'lightgreen'
white_paint_color = 'white'
black_paint_color = 'black'

# Define the tile status for the goal state
tile_status = {
    'tile_0-1': 'no constraints', 'tile_0-2': 'no constraints', 'tile_0-3': 'no constraints',
    'tile_1-1': 'painted white', 'tile_1-2': 'painted black', 'tile_1-3': 'painted white',
    'tile_2-1': 'painted black', 'tile_2-2': 'painted white', 'tile_2-3': 'painted black',
}

# 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]
        if 'painted white' in status:
            color = white_paint_color
        elif 'painted black' in status:
            color = black_paint_color
        else:
            color = 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, f'{tile_id}\n{status}', ha='center', va='center', fontsize=8, color='black' if color != black_paint_color else 'white')

# Draw the robot with no constraints
robot_position = (3.5, 1.5)
circle = patches.Circle(robot_position, 0.3, edgecolor='black', facecolor='white')
ax.add_patch(circle)
ax.text(robot_position[0], robot_position[1] - 0.2, 'robot1\n(no constraints)', ha='center', va='center', fontsize=8, color='black')

# Draw available colors
color_white_position = (3.5, 0.5)
color_black_position = (3.5, 0.0)

rect_white = patches.Rectangle((color_white_position[0] - 0.2, color_white_position[1] - 0.2), 0.4, 0.4, linewidth=1, edgecolor='black', facecolor='white')
ax.add_patch(rect_white)
ax.text(color_white_position[0], color_white_position[1], 'color_white\navailable', ha='center', va='center', fontsize=8, color='black')

rect_black = patches.Rectangle((color_black_position[0] - 0.2, color_black_position[1] - 0.2), 0.4, 0.4, linewidth=1, edgecolor='black', facecolor='black')
ax.add_patch(rect_black)
ax.text(color_black_position[0], color_black_position[1], 'color_black\navailable', ha='center', va='center', fontsize=8, color='white')

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

# Add legend
legend_elements = [
    patches.Patch(facecolor=clear_color, edgecolor='black', label='Clear Tile'),
    patches.Patch(facecolor=white_paint_color, edgecolor='black', label='Painted White'),
    patches.Patch(facecolor=black_paint_color, edgecolor='black', label='Painted Black'),
    patches.Patch(facecolor='white', edgecolor='black', label='Robot1 (White)'),
]
ax.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(1.3, 1))

# Save the figure
plt.savefig('<PATH_REMOVED>', bbox_inches='tight')
plt.show()