implement reset

src/components/cube-piece.tsx

@@ -1,7 +1,7 @@
 'use client';
 
 import { RoundedBox } from '@react-three/drei';
-import { useEffect, useRef, useState } from 'react';
+import { forwardRef, useEffect, useImperativeHandle, useRef } from 'react';
 import { Mesh } from 'three';
 
 import { FacingDirection, Rotations } from './consts';
@@ -17,14 +17,17 @@ const CUBE_COLORS = {
   bottom: '#ffffff', // White
 };
 
+export type CubePieceRef = {
+  resetPosition: () => void;
+};
+
 type CubePieceProps = {
   roughness: number;
   initialPosition: [number, number, number];
 };
 
-export const CubePiece = ({ roughness, initialPosition }: CubePieceProps) => {
+export const CubePiece = forwardRef<CubePieceRef, CubePieceProps>(({ roughness, initialPosition }, ref) => {
   const [x, y, z] = initialPosition;
-  const [position] = useState<[number, number, number]>([x, y, z]);
 
   const meshRef = useRef<Mesh | null>(null);
   useEffect(() => {
@@ -33,6 +36,13 @@ export const CubePiece = ({ roughness, initialPosition }: CubePieceProps) => {
     }
   }, [meshRef]);
 
+  useImperativeHandle(ref, () => ({
+    resetPosition: () => {
+      meshRef.current?.position.set(x, y, z);
+      meshRef.current?.rotation.set(0, 0, 0);
+    },
+  }));
+
   const visibleFaces: Record<FacingDirection, boolean> = {
     front: z > 0,
     back: z < 0,
@@ -51,7 +61,7 @@ export const CubePiece = ({ roughness, initialPosition }: CubePieceProps) => {
   };
 
   return (
-    <mesh position={position} ref={meshRef}>
+    <mesh position={[x, y, z]} ref={meshRef}>
       <RoundedBox args={[0.95, 0.95, 0.95]} radius={0.05} smoothness={4}>
         <meshStandardMaterial color="#2a2a2a" metalness={1} roughness={roughness} />
       </RoundedBox>
@@ -68,4 +78,6 @@ export const CubePiece = ({ roughness, initialPosition }: CubePieceProps) => {
       })}
     </mesh>
   );
-};
+});
+
+CubePiece.displayName = 'CubePiece';

src/components/rotation-controller.ts

@@ -65,6 +65,16 @@ export class RotationController {
     return RotationController.instance;
   }
 
+  stopRotation(cb: () => void) {
+    this.rotationSteps = [];
+    const cancel = setInterval(() => {
+      if (!this.rotatingStep) {
+        clearInterval(cancel);
+        cb();
+      }
+    }, 50);
+  }
+
   setCubeGroup(cubeGroup: Group) {
     this.cubeGroup = cubeGroup;
   }

src/components/rubiks-cube.tsx

@@ -2,7 +2,7 @@ import { forwardRef, useEffect, useImperativeHandle, useRef } from 'react';
 import { Group } from 'three';
 
 import { RotationStep } from './consts';
-import { CubePiece } from './cube-piece';
+import { CubePiece, CubePieceRef } from './cube-piece';
 import { rotationController } from './rotation-controller';
 import { Rotator, RotatorRef } from './rotator';
 
@@ -17,6 +17,7 @@ for (let x = -0.5; x <= 0.5; x += 1) {
 
 export type RubiksCubeRef = {
   rotate: (steps: Array<RotationStep>) => void;
+  reset: () => void;
 };
 
 type RubiksCubeProps = {
@@ -27,9 +28,17 @@ type RubiksCubeProps = {
 export const RubiksCube = forwardRef<RubiksCubeRef, RubiksCubeProps>(({ cubeRoughness, cubeSpeed }, ref) => {
   const cubeGroupRef = useRef<Group | null>(null);
   const rotatorRef = useRef<RotatorRef | null>(null);
+  const cubePieceRefs = useRef<Map<string, CubePieceRef>>(new Map());
 
   useImperativeHandle(ref, () => ({
     rotate: (steps: Array<RotationStep>) => rotatorRef.current?.rotate(steps),
+    reset: () => {
+      rotationController.stopRotation(() => {
+        cubePieceRefs.current.forEach((cubePieceRef) => {
+          cubePieceRef.resetPosition();
+        });
+      });
+    },
   }));
 
   useEffect(() => {
@@ -39,9 +48,23 @@ export const RubiksCube = forwardRef<RubiksCubeRef, RubiksCubeProps>(({ cubeRoug
   return (
     <>
       <group ref={cubeGroupRef}>
-        {CUBE_POSITIONS.map((position) => (
-          <CubePiece key={position.join(',')} initialPosition={position} roughness={cubeRoughness} />
-        ))}
+        {CUBE_POSITIONS.map((position) => {
+          const positionKey = position.join(',');
+          return (
+            <CubePiece
+              key={positionKey}
+              ref={(ref) => {
+                if (ref) {
+                  cubePieceRefs.current.set(positionKey, ref);
+                } else {
+                  cubePieceRefs.current.delete(positionKey);
+                }
+              }}
+              initialPosition={position}
+              roughness={cubeRoughness}
+            />
+          );
+        })}
       </group>
       <Rotator ref={rotatorRef} cubeSpeed={cubeSpeed} />
     </>

src/components/ui-controls.tsx

@@ -16,7 +16,7 @@ export const UIControls = () => {
   };
 
   const reset = () => {
-    alert('Working on it!');
+    rubiksCubeRef?.current?.reset();
   };
 
   const solve = () => {

src/components/ui-title.tsx

@@ -1,10 +1,14 @@
 export const UITitle = () => {
   return (
     <div className="z-10 pointer-events-none">
-      <div className="text-6xl font-bold">Rubik&apos;s Cube Solver</div>
-      <div className="text-2xl text-gray-800">with Reinforcement Learning</div>
-      <div className="text-gray-700">
-        <a href="https://cross-entropy.ai">https://cross-entropy.ai</a>
+      <div className="pointer-events-auto w-fit">
+        <div className="text-6xl font-bold">Rubik&apos;s Cube Solver</div>
+        <div className="text-2xl text-gray-800">with Reinforcement Learning</div>
+        <div className="text-gray-700">
+          <a href="https://cross-entropy.ai" target="_blank" className="underline text-primary font-bold">
+            https://cross-entropy.ai
+          </a>
+        </div>
       </div>
     </div>
   );