import numpy as np

from typing import AnyStr, List
from numpy.typing import ArrayLike

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


class WSEN_ISDS(IMU):
    pass


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

        self.__variance = 0.05
        self.__noise = GaussianNoise(np.zeros(3), np.identity(3) * self.__variance)

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

    def _sensor_specific_effects(self, x: ArrayLike) -> ArrayLike:
        t = self._dataset.get_time()

        # Apply noise to the true values.
        y = self.__noise(t, x)
        noise = y - x

        # Log the chosen noise values.
        self._logger.write('acc_x_noise', noise[0], self._get_name())
        self._logger.write('acc_y_noise', noise[1], self._get_name())
        self._logger.write('acc_z_noise', noise[2], self._get_name())

        return y
    

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

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

    def _sensor_specific_effects(self, x: ArrayLike) -> ArrayLike:
        # Convert to degrees per second.
        x = (x / np.pi) * 180

        # TODO: Noise model.

        return x