import numpy as np

from numpy.typing import ArrayLike
from typing import List

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


class Gyroscope(Sensor):
    def __init__(self, dataset: DataSource, logger: Logger, orientation=np.identity(3), transforms: List[Transform] = []):
        super().__init__(dataset, logger, transforms)

        self._orientation = orientation

    def calibrate(self, n_samples):
        return np.sum(self() for i in range(n_samples)) / n_samples

    def _get_data(self) -> ArrayLike | float:
        # Rotation in rad/sec
        omegas = self._dataset.get_angular_velocity()
        omegas = self._orientation @ omegas

        self._log('ox', omegas[0])
        self._log('oy', omegas[1])
        self._log('oz', omegas[2])

        return omegas