import numpy as np

from typing import List
from numpy.typing import ArrayLike

from spatz.sensors import Sensor
from spatz.transforms import Transform
from spatz.dataset import Dataset
from spatz.logger import Logger


__all__=[
    'Accelerometer'
]


# Local definition of gravitation
g = 9.81


class Accelerometer(Sensor):
    def __init__(self, dataset: Dataset, logger: Logger, offset: float = 0, transforms: List[Transform] = []):
        super().__init__(dataset, logger, transforms)

        self._offset = np.array([offset, 0, 0])

    def _get_data(self) -> ArrayLike | float:
        acc = self._dataset.get_acceleration(frame='FL')
        acc += np.array([0, 0, g])

        self._logger.write('FL_x', acc[0], self._get_name())
        self._logger.write('FL_y', acc[1], self._get_name())
        self._logger.write('FL_z', acc[2], self._get_name())

        # Convert FL to body
        acc = self._dataset.launch_rail_to_body() @ acc

        self._logger.write('B_x', acc[0], self._get_name())
        self._logger.write('B_y', acc[1], self._get_name())
        self._logger.write('B_z', acc[2], self._get_name())

        # Flip axes to sensor's perspective.
        acc *= -1

        # Add the effects of the imu's offset.
        omega = self._dataset.get_angular_velocities()
        acc += (np.cross(omega, self._offset) + np.cross(omega, np.cross(omega, self._offset)))

        return acc