from aim_fsm import *
from aim_fsm.particle import *

class Lab3b(StateMachineProgram):
    def __init__(self):
        landmarks = {'ArucoMarker-1': Pose(-25, 160, 0, -pi/2),
                     'ArucoMarker-2': Pose( 25, 160, 0, -pi/2),
                     'ArucoMarker-3': Pose(160,  25, 0,  pi),
                     'ArucoMarker-4': Pose(160, -25, 0,  pi)}
        pf = ParticleFilter(robot,
                            num_particles=500, landmarks=landmarks,
                            #sensor_model = ArucoDistanceSensorModel,
                            sensor_model = ArucoBearingSensorModel,
                            #sensor_model = ArucoCombinedSensorModel
                            )
        super().__init__(particle_filter = pf,
                         wall_marker_dict = None,
                         speech = False,
                         launch_particle_viewer = True)

    def start(self):
        super().start()
        a1 = ArucoMarkerObj({'name':'ArucoMarker-1.a', 'id':1, 'marker':None}, x=-25, y=160, theta=-pi/2, is_fixed=True)
        a2 = ArucoMarkerObj({'name':'ArucoMarker-2.a', 'id':2, 'marker':None}, x= 25, y=160, theta=-pi/2, is_fixed=True)
        a3 = ArucoMarkerObj({'name':'ArucoMarker-3.a', 'id':3, 'marker':None}, x=160, y= 25, theta=pi, is_fixed=True)
        a4 = ArucoMarkerObj({'name':'ArucoMarker-4.a', 'id':4, 'marker':None}, x=160, y=-25, theta=pi, is_fixed=True)
        robot.world_map.objects['ArucoMarker-1.a'] = a1
        robot.world_map.objects['ArucoMarker-2.a'] = a2
        robot.world_map.objects['ArucoMarker-3.a'] = a3
        robot.world_map.objects['ArucoMarker-4.a'] = a4