import numpy as np

from typing import List, AnyStr

from numpy.typing import ArrayLike
from spatz.dataset import ArrayLike, Dataset
from spatz.logger import ArrayLike, Logger
from spatz.sensors import IMU, Accelerometer, Gyroscope, CoordSystem
from spatz.transforms import Transform, GaussianNoise


class BMI088(IMU):
    def __init__(self, dataset: Dataset, logger: Logger, orientation = np.identity(3), offset=0, transforms: List[Transform] = []):
        acc = BMI088Acc(dataset, logger, orientation, offset, transforms)
        gyro = BMI088Gyro(dataset, logger, offset, transforms)

        super().__init__(dataset, logger, acc, gyro, transforms)


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

    def _get_name(self) -> AnyStr:
        return 'BMI088'

    def _get_data(self) -> ArrayLike:
        rots = self._dataset.fetch_values(['roll_l', 'pitch_l', 'yaw_l'])

        return rots

    def _sensor_specific_effects(self, x: ArrayLike) -> ArrayLike:
        return x
    

class BMI088Acc(Accelerometer):
    def __init__(self, dataset: Dataset, logger: Logger, orientation = np.identity(3), offset: float = 0, transforms: List[Transform] = []):
        super().__init__(dataset, logger, CoordSystem.RIGHT_HANDED, orientation, offset, transforms)

        self.__noise = GaussianNoise(0, 0.05)

    def _get_name(self) -> AnyStr:
        return 'BMI088'
    
    def _sensor_specific_effects(self, x: ArrayLike) -> ArrayLike:
        return self.__noise(0, x)