#!/usr/bin/env python3 """ Enhanced Rubik's Cube Visualizer This module provides improved visualization tools for the Rubik's cube environment, including progress tracking, move history visualization, and interactive elements. """ import base64 import datetime import os import random import re from io import BytesIO from typing import Dict, List import matplotlib.pyplot as plt def generate_progress_chart( move_history: List[str] = None, progress_history: List[float] = None, rewards_history: List[float] = None, solved_at_move: int = None, title: str = "Cube Solving Progress", ) -> str: """ Generate a chart showing progress over solving steps Args: move_history: List of moves applied progress_history: List of progress values (0.0-1.0) after each move rewards_history: Optional list of rewards for each move solved_at_move: Index of the move that solved the cube (if solved) title: Chart title Returns: Base64-encoded PNG image of the chart """ plt.figure(figsize=(12, 6)) # Ensure we have data to plot and initialize if None if move_history is None: move_history = [] if progress_history is None or not progress_history: plt.text( 0.5, 0.5, "No progress data available", horizontalalignment="center", verticalalignment="center", transform=plt.gca().transAxes, fontsize=14, ) # Save the figure to a base64 string buffer = BytesIO() plt.savefig(buffer, format="png", dpi=100, bbox_inches="tight") buffer.seek(0) image_png = buffer.getvalue() buffer.close() plt.close() return base64.b64encode(image_png).decode("utf-8") # Plot progress # Make sure move_indices and progress_history have the same length move_indices = list(range(len(progress_history))) plt.plot(move_indices, progress_history, "b-", linewidth=2, label="Progress") # Add markers for each move plt.plot(move_indices, progress_history, "bo", markersize=6) # Plot rewards if provided if rewards_history: # Ensure rewards_history has the same length as progress_history if len(rewards_history) != len(progress_history): # Truncate or extend rewards_history to match progress_history if len(rewards_history) > len(progress_history): rewards_history = rewards_history[: len(progress_history)] else: # Extend with the last value or 0 last_reward = rewards_history[-1] if rewards_history else 0 rewards_history.extend( [last_reward] * (len(progress_history) - len(rewards_history)) ) # Normalize rewards to 0-1 range for comparison if rewards_history: min_reward = min(rewards_history) max_reward = max(rewards_history) reward_range = max_reward - min_reward if reward_range > 0: normalized_rewards = [ (r - min_reward) / reward_range for r in rewards_history ] else: normalized_rewards = [0.5] * len(rewards_history) plt.plot( move_indices, normalized_rewards, "r--", linewidth=1.5, label="Reward" ) # Highlight the solving move if provided if solved_at_move is not None and 0 <= solved_at_move < len(progress_history): plt.axvline( x=solved_at_move, color="g", linestyle="--", label=f"Solved at move {solved_at_move+1}", ) plt.plot( [solved_at_move], [progress_history[solved_at_move]], "g*", markersize=15, label="Solution", ) # Add move labels (only if we have moves to label) if move_history: # Ensure move_history has the same length as progress_history if len(move_history) > len(progress_history): move_history = move_history[: len(progress_history)] elif len(move_history) < len(progress_history): # Extend with empty strings move_history.extend([""] * (len(progress_history) - len(move_history))) for i, move in enumerate(move_history): if move: # Only annotate non-empty moves plt.annotate( move, (i, progress_history[i]), textcoords="offset points", xytext=(0, 10), ha="center", ) # Add grid and labels plt.grid(True, linestyle="--", alpha=0.7) plt.xlabel("Move Number") plt.ylabel("Progress / Normalized Reward") plt.title(title) plt.ylim(-0.05, 1.05) # Ensure there's space for annotations plt.legend() # Save the figure to a base64 string buffer = BytesIO() plt.savefig(buffer, format="png", dpi=100, bbox_inches="tight") buffer.seek(0) image_png = buffer.getvalue() buffer.close() plt.close() return base64.b64encode(image_png).decode("utf-8") def parse_cube_state(cube_state: str) -> Dict[str, List[str]]: """ Parse the cube state string into a dictionary of faces. Args: cube_state: String representation of the cube Returns: Dictionary with keys 'up', 'down', 'left', 'right', 'front', 'back' Each containing a 3x3 grid represented as a flattened list of colors """ # Initialize faces dictionary faces = {"up": [], "down": [], "left": [], "right": [], "front": [], "back": []} # Extract face information current_face = None for line in cube_state.strip().split("\n"): line = line.strip() if line.startswith("U:"): current_face = "up" faces[current_face] = [c for c in line[2:].split() if c] elif line.startswith("D:"): current_face = "down" faces[current_face] = [c for c in line[2:].split() if c] elif line.startswith("L:"): current_face = "left" faces[current_face] = [c for c in line[2:].split() if c] elif line.startswith("R:"): current_face = "right" faces[current_face] = [c for c in line[2:].split() if c] elif line.startswith("F:"): current_face = "front" faces[current_face] = [c for c in line[2:].split() if c] elif line.startswith("B:"): current_face = "back" faces[current_face] = [c for c in line[2:].split() if c] elif current_face and line: # Continue accumulating colors for the current face colors = [c for c in line.split() if c] if colors: faces[current_face].extend(colors) # Ensure each face has exactly 9 elements (3x3 grid) for face in faces: # If we have too few colors, pad with a placeholder while len(faces[face]) < 9: faces[face].append("?") # If we have too many, truncate faces[face] = faces[face][:9] return faces def generate_enhanced_cube_html( cube_state: str, move_history: List[str] = None, progress_history: List[float] = None, rewards_history: List[float] = None, thinking_history: List[str] = None, scramble_sequence: List[str] = None, is_solved: bool = False, curriculum_level: int = None, curriculum_description: str = None, ) -> str: """ Generate enhanced HTML visualization of a Rubik's cube solve attempt Args: cube_state: String representation of the cube's current state move_history: List of moves applied during solving progress_history: List of progress values after each move rewards_history: List of rewards for each move thinking_history: List of thinking steps from the LLM scramble_sequence: List of moves used to scramble the cube is_solved: Whether the cube is solved curriculum_level: Current curriculum level (if using curriculum) curriculum_description: Description of the current curriculum level Returns: HTML string for visualization """ # Extract the colors from the cube state string faces = parse_cube_state(cube_state) # Generate progress chart if data is available progress_chart_base64 = None if progress_history: solved_at_move = None if is_solved and move_history: solved_at_move = len(move_history) - 1 # Make sure we have move_history even if it's empty if move_history is None: move_history = [] progress_chart_base64 = generate_progress_chart( move_history=move_history, progress_history=progress_history, rewards_history=rewards_history, solved_at_move=solved_at_move, title="Rubik's Cube Solving Progress", ) # Generate the HTML html = """ Enhanced Rubik's Cube Visualizer

Enhanced Rubik's Cube Visualization

""" # Add status badges if is_solved: html += 'SOLVED' else: html += 'UNSOLVED' if curriculum_level is not None: html += f'LEVEL {curriculum_level}' if curriculum_description: html += f"

{curriculum_description}

" html += """

Current State

""" # Create a 3D layout of the cube faces # Row 1: Empty, Up, Empty, Empty html += '
' html += '
' # Empty space html += generate_face_html("Up", faces["up"]) html += '
' # Empty space html += '
' # Empty space html += "
" # Row 2: Left, Front, Right, Back html += '
' html += generate_face_html("Left", faces["left"]) html += generate_face_html("Front", faces["front"]) html += generate_face_html("Right", faces["right"]) html += generate_face_html("Back", faces["back"]) html += "
" # Row 3: Empty, Down, Empty, Empty html += '
' html += '
' # Empty space html += generate_face_html("Down", faces["down"]) html += '
' # Empty space html += '
' # Empty space html += "
" html += """
Progress
Move History
Thinking Steps
""" # Progress Chart Tab html += '
' html += '
' html += "

Solving Progress

" if progress_chart_base64: img_src = f"data:image/png;base64,{progress_chart_base64}" html += f'Solving Progress Chart' else: html += "

No progress data available.

" html += "
" # Moves History Tab html += '
' html += '
' html += "

Move History

" # Add playback controls if moves exist if move_history: html += """
0 / 0
""" if scramble_sequence: html += "

Scramble Sequence

" html += '
' for move in scramble_sequence: move_class = move[0] if len(move) >= 1 else "" html += f'
{move}
' html += "
" if move_history: html += "

Solving Moves

" html += '
' for i, move in enumerate(move_history): move_class = move[0] if len(move) >= 1 else "" html += f'
{move}
' html += "
" else: html += "

No moves have been made yet.

" html += "
" # Thinking Steps Tab html += '
' html += '
' html += "

Thinking Process

" if thinking_history: for i, thinking in enumerate(thinking_history): html += f'
Step {i+1}: {thinking}
' else: html += "

No thinking steps recorded.

" html += "
" # Add JavaScript for interactivity html += """
""" return html def generate_face_html(face_name: str, face_colors: List[str]) -> str: """ Generate HTML for a single face of the cube Args: face_name: Name of the face (Up, Down, Left, Right, Front, Back) face_colors: List of 9 colors for the face Returns: HTML string for the face """ # Color mapping color_map = { "W": "#FFFFFF", # White "Y": "#FFFF00", # Yellow "R": "#FF0000", # Red "O": "#FFA500", # Orange "B": "#0000FF", # Blue "G": "#00FF00", # Green "?": "#CCCCCC", # Gray (unknown/placeholder) } html = f'
{face_name}
' for color in face_colors: bg_color = color_map.get( color[0], "#CCCCCC" ) # Get first char of color, default to gray html += ( f'
{color}
' ) html += "
" return html def extract_thinking_from_history(message_history: List[Dict]) -> List[str]: """ Extract thinking steps from LLM message history Args: message_history: List of message dictionaries with roles and content Returns: List of extracted thinking content """ thinking_steps = [] for message in message_history: if message.get("role") == "agent" and isinstance(message.get("content"), str): content = message["content"] thinking_match = re.search(r"(.*?)", content, re.DOTALL) if thinking_match: thinking_text = thinking_match.group(1).strip() if thinking_text: thinking_steps.append(thinking_text) return thinking_steps def save_enhanced_visualization( cube_state: str, move_history: List[str] = None, progress_history: List[float] = None, rewards_history: List[float] = None, thinking_history: List[str] = None, message_history: List[Dict] = None, scramble_sequence: List[str] = None, is_solved: bool = False, curriculum_level: int = None, curriculum_description: str = None, output_path: str = None, ) -> str: """ Generate and save an enhanced HTML visualization of a Rubik's cube solving attempt Args: cube_state: String representation of the cube's current state move_history: List of moves applied during solving progress_history: List of progress values after each move rewards_history: List of rewards for each move thinking_history: List of thinking steps extracted from LLM message_history: Optional full message history to extract thinking from scramble_sequence: List of moves used to scramble the cube is_solved: Whether the cube is solved curriculum_level: Current curriculum level (if using curriculum) curriculum_description: Description of the current curriculum level output_path: Optional file path to save the HTML Returns: Path to the saved HTML file """ # Extract thinking from message history if provided and not already extracted if message_history and not thinking_history: thinking_history = extract_thinking_from_history(message_history) # Generate the HTML html = generate_enhanced_cube_html( cube_state=cube_state, move_history=move_history, progress_history=progress_history, rewards_history=rewards_history, thinking_history=thinking_history, scramble_sequence=scramble_sequence, is_solved=is_solved, curriculum_level=curriculum_level, curriculum_description=curriculum_description, ) # Set default output path if not provided if output_path is None: timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") solved_status = "solved" if is_solved else "unsolved" output_path = f"rubiks_visualization_{solved_status}_{timestamp}.html" # Create directory if it doesn't exist os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) # Save the HTML to a file with open(output_path, "w") as f: f.write(html) return output_path if __name__ == "__main__": # Example usage example_cube_state = """ U: W W W W W W W W W D: Y Y Y Y Y Y Y Y Y L: O O O O O O O O O R: R R R R R R R R R F: G G G G G G G G G B: B B B B B B B B B """ # Create example data example_moves = [ "R", "U", "R'", "U'", "F", "R", "U", "R'", "U'", "F'", "U", "R", "U2", "R'", ] # Generate sample progress history progress_vals = [0.5] for i in range(1, len(example_moves)): # Random progress that generally increases next_progress = min(1.0, progress_vals[-1] + random.uniform(-0.05, 0.15)) progress_vals.append(next_progress) # Generate sample rewards rewards = [random.uniform(0.1, 0.9) for _ in range(len(example_moves))] # Sample thinking steps sample_thinking = [ "I see the cube has a white cross on top already. I'll focus on solving the first layer corners.", "I'll use the sequence R U R' U' to position the corner piece without disrupting the cross.", "Now I need to solve the middle layer edges. I'll use the appropriate algorithm based on the edge orientation.", "Looking at the last layer, I need to orient the yellow face first, then permute the corners and edges.", ] # Sample scramble sample_scramble = ["F", "R", "U'", "B", "L2", "D"] # Generate and save the visualization html_path = save_enhanced_visualization( cube_state=example_cube_state, move_history=example_moves, progress_history=progress_vals, rewards_history=rewards, thinking_history=sample_thinking, scramble_sequence=sample_scramble, is_solved=True, curriculum_level=2, curriculum_description="Level 2: Easy - Learn basic patterns and simple sequences", output_path="example_enhanced_rubiks.html", ) print(f"Enhanced visualization saved to {html_path}")