import numpy as np

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

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


class WSEN_ISDS(IMU):
    pass


class WSEN_ISDS_ACC(Accelerometer):
    def __init__(self, dataset: DataSource, logger: Logger, orientation=np.identity(3), offset=0, transforms: List[Transform] = []):
        super().__init__(dataset, logger, CoordSystem.LEFT_HANDED, orientation, 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_ACC'

    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())

        # The WSEN-ISDS accelerometer only measures acceleration between -16g and 16g. 
        g = 9.81
        y = np.clip(y, -16*g, +16*g)

        return y
    

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

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

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

        # TODO: Noise model.

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

        return x