import math
import numpy as np

from typing import List

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


class PressureSensor(Sensor):
    def __init__(self, dataset: Dataset, logger: Logger, transforms: List[Transform] = [], ts_effects=True):
        """
        Args:
            dataset (Dataset): A dataset instance.
            transforms (List[Transform], optional): Transforms to apply to the sensor outputs. Defaults to [].
            ts_effects (bool, optional): If True, models transsonic effects. Defaults to True.
        """
        super(PressureSensor, self).__init__(dataset, logger, transforms)

        self._ts_effects = ts_effects

    def _get_data(self) -> float:
        x = self._dataset.get_pressure()

        if self._ts_effects:
            # Pre-defined constants.
            _p = 3e6
            sigma = 40

            # How far away from transsonic speed (mach 1) are we?
            vvec = self._dataset.get_velocity()
            dv = np.abs(np.linalg.norm(vvec) - self._dataset.get_speed_of_sound())

            # Model transsonic effects by a peak at mach 1 which decays the further we are away from it.
            ts_eff = _p * math.exp(-0.5* (dv / sigma)**2 ) / (sigma * math.sqrt(2*math.pi))

            # Log the values for the transsonic effect.
            self._logger.write('ts_effects', ts_eff, domain=self._get_name())
            self._logger.write('mach_no', self._dataset.get_mach_number(), domain='mach')
            self._logger.write('speedofsound', self._dataset.get_speed_of_sound(), domain='mach')

            x = x + ts_eff

        return x