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

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

# Define grid dimensions
rows, cols = 6, 4
grid_size = 1

# Define colors
colors = {
    'rightl0': 'blue',
    'straight0': 'green',
    'straight1': 'red',
    'square0': 'yellow',
    'square1': 'orange',
    'clear': 'white'
}

# Define piece positions for the child state
pieces = {
    'rightl0': [(3, 2), (4, 2), (4, 3)],
    'straight0': [(0, 1), (1, 1)],
    'straight1': [(0, 2), (1, 2)],
    'square0': [(1, 0)],
    'square1': [(3, 0)]
}

# Draw grid
for row in range(rows):
    for col in range(cols):
        # Determine if the cell is occupied and by which piece
        occupied = False
        piece_name = ''
        for name, positions in pieces.items():
            if (row, col) in positions:
                occupied = True
                piece_name = name
                break
        
        # Set color and text based on occupancy
        if occupied:
            color = colors[piece_name]
            status = 'occupied'
        else:
            color = colors['clear']
            status = 'clear'
        
        # Draw the rectangle
        rect = patches.Rectangle((col * grid_size, (rows - 1 - row) * grid_size), grid_size, grid_size, linewidth=1, edgecolor='black', facecolor=color)
        ax.add_patch(rect)

        # Add text
        text_color = 'black' if color in ['white', 'yellow', 'orange'] else 'white'
        ax.text(col * grid_size + 0.5 * grid_size, (rows - 1 - row) * grid_size + 0.7 * grid_size, piece_name if occupied else '', ha='center', va='center', color=text_color, fontsize=8, fontweight='bold')
        ax.text(col * grid_size + 0.5 * grid_size, (rows - 1 - row) * grid_size + 0.3 * grid_size, f"f{row}-{col}f\n{status}", ha='center', va='center', color=text_color, fontsize=6)

# Bold outer edges of pieces
for name, positions in pieces.items():
    for pos in positions:
        row, col = pos
        # Draw bold rectangle
        rect = patches.Rectangle((col * grid_size, (rows - 1 - row) * grid_size), grid_size, grid_size, linewidth=3, edgecolor='black', facecolor='none')
        ax.add_patch(rect)

# Set axis limits and aspect
ax.set_xlim(0, cols * grid_size)
ax.set_ylim(0, rows * grid_size)
ax.set_aspect('equal')

# Remove axes
ax.axis('off')

# Add legend
legend_elements = [
    patches.Patch(facecolor='blue', edgecolor='black', label='rightl0: occupied'),
    patches.Patch(facecolor='green', edgecolor='black', label='straight0: occupied'),
    patches.Patch(facecolor='red', edgecolor='black', label='straight1: occupied'),
    patches.Patch(facecolor='yellow', edgecolor='black', label='square0: occupied'),
    patches.Patch(facecolor='orange', edgecolor='black', label='square1: occupied'),
    patches.Patch(facecolor='white', edgecolor='black', label='clear: clear')
]
ax.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(1.2, 1))

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