"""
In the particle viewer window:
   the WASD keys move the robot
  'e' forces an evaluation step
  'r' forces a resampling
  'v' displays the weight statistics
  'z' re-randomizes the particles.
"""

from cozmo_fsm import *
from cozmo.util import degrees, Pose

ARUCO_MARKER_SIZE = 50  # mm

class Lab3(StateMachineProgram):
    def __init__(self):
        landmarks = {
            'Aruco-0' : Pose(-55, 160, 0, angle_z=degrees(90)),
            'Aruco-1' : Pose( 55, 160, 0, angle_z=degrees(90)),
            'Aruco-2' : Pose(160,  55, 0, angle_z=degrees( 0)),
            'Aruco-3' : Pose(160, -55, 0, angle_z=degrees( 0))
        }

        pf = ParticleFilter(robot,
                            landmarks = landmarks,
                            sensor_model = ArucoDistanceSensorModel(robot))

        super().__init__(aruco_marker_size=50, # millimeters
                         particle_filter=pf,
                         particle_viewer=True)