from __future__ import annotations

import numpy as np

from utils.math_ops import MathOps


class FieldLandmarks:
    def __init__(self, field):
        self.field = field
        self.world = field.world
        self.landmarks: dict[str, np.ndarray] = {}
        self.canonical_landmarks: dict[str, np.ndarray] = field.get_canonical_landmarks()

    def update_from_perception(self, landmark_id: str, landmark_pos: np.ndarray) -> None:
        """
        Calculates the global position of all currently visible landmarks.
        """

        world = self.world
        local_cart_3d = MathOps.deg_sph2cart(landmark_pos)

        global_pos_3d = MathOps.rel_to_global_3d(
            local_pos_3d=local_cart_3d,
            global_pos_3d=world.global_position,
            global_orientation_quat=world.agent.robot.global_orientation_quat,
        )

        self.landmarks[landmark_id] = global_pos_3d

    def get_landmark_position(
        self, landmark_id: str, use_canonical: bool = False
    ) -> np.ndarray | None:
        """
        Returns the current perceived or canonical global position for a landmark.
        """
        source = self.canonical_landmarks if use_canonical else self.landmarks
        return source.get(landmark_id)

    def get_canonical_landmark_position(self, landmark_id: str) -> np.ndarray | None:
        return self.canonical_landmarks.get(landmark_id)