import numpy as np

from typing import List
from enum import Enum
from numpy.typing import ArrayLike, NDArray

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


__all__=[
    'Accelerometer'
]


# Local definition of gravitation
g = 9.81


class CoordSystem:
    LEFT_HANDED = 0,
    RIGHT_HANDED = 0,


class Accelerometer(Sensor):
    def __init__(
            self, dataset: DataSource, 
            logger: Logger, 
            orientation: NDArray = np.identity(3), 
            offset: float = 0, 
            transforms: List[Transform] = []):
        """Accelerometer sensor base class.

        Args:
            dataset (Dataset): A dataset object to fetch data from.
            logger (Logger): A logger object to write data to.
            coord_system (CoordSystem, optional): The type of coordinate system to use. Defaults to CoordSystem.RIGHT_HANDED.
            orientation (NDArray, optional): The orientation of the sensor inside the spacecraft. Defaults to np.identity(3).
            offset (float, optional): The offset of the sensor from the origin around which the object rotates. Defaults to 0.
            transforms (List[Transform], optional): A list of transformations applied to the sensor measurements. Defaults to [].
        """
        super().__init__(dataset, logger, transforms)

        assert orientation.shape == (3, 3), 'Orientation has to be a 3x3 matrix.'

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

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

        return acc