import numpy as np
import scipy.signal
import math
from collections import deque

from .delayLine import DelayLine
from .material import Material

class SimulatorManager:
    """
    This class defines the simulation logic. It contains all the methods and components that are needed
    to run the acoustic simulation in the defined scenario, and compute the output signals. The
    acoustic scene defined in the `Environment` module contains a single sound source, an array of microphones, 
    and some additional parameters (e.g. a sampling frequency, the atmospheric conditions, the road surface properties,
    a background noise). The `SimulatorManager` simulates the propagation of sound from the moving source to a 
    single, static microphone, exploiting the scene parameters that it receives at the time of its initialization.

    The received signal will contain two different components: the direct sound, propagating directly from 
    the source (assumed to be omnidirectional and point-like) to the microphone (assumed
    to be omnidirectional and point-like too), and the sound reflected by the road surface. In order to simulate the
    reflection, the simulator exploits the Image-Source Method (ISM): a virtual source is generated by reflecting the
    position of the real sound source w.r.t the road surface, and the reflected sound is assumed to be emitted by
    this source as if it was emitting in the free-field. The sound produced the image source is then filtered by a 
    set of FIR filters that simulate the asphalt surface absorption and the attenuation due to the propagation of
    the sound through the air.
    
    The simulation is divided in two separate phases:
    * During the `initialization` phase, the `DelayLine` components used to simulate the sound propagation delay
    are instantiated, and the initial delays between the microphone, the sound source and the reflection point are
    computed.
    * During the `update` phase, first the signal emitted by the source is written on the two `DelayLine` instances 
    controlling the direct and reflected sound propagation delays. Then, the delays are updated according to the 
    new position of the moving source, together with the distances source -> receiver, image source -> road surface 
    and road surface -> receiver. The corresponding delayed signal samples are read from the `DelayLines` and filtered
    via a set of filters simulating the (angle-dependent) asphalt reflection and the air absorption. Finally, the
    signal received by the microphone is computed by summing the direct and the reflected components of the simulated
    sound field.

    Attributes
    -------
    c: float
        Speed of sound in the air
    fs: int
        Sampling frequency used in the simulations
    Z0: float
        Acoustic impedance of air
    road_material: Material
        Asphalt type of the road surface, defined by its absorption and reflection properties
    primaryDelayLine: DelayLine
        Primary `DelayLine` component, used to compute delays corresponding to propagation of sound between source
        and receiver on direct path, and between image source and road surface, on reflected path
    secondaryDelayLine: DelayLine
        Secondary `DelayLine` component, used to compute delay corresponding to the propagation of sound on the 
        reflected path from the reflection point on the road surface to the microphone position
    airAbsorptionCoefficients: np.ndarray
        1D array containing the air absorption coefficients computed at a set of frequency bands covering the 
        range from 0 to fs/2
    asphaltReflectionFilterTable: np.ndarray
        2D array containing the asphalt reflection filters computed at different reflection angles. The filters
        have a fixed number of taps ntaps = 10, and their impulse response at a set of angles from -89 to + 89 degrees
        is pre-computed and stored in this table in order to speed up the execution.
    simulation_params: dict
        Dictionary containing the keys:
        * `interp_method`: str, defines the interpolation method used in the delay line interpolated reads
        * `include_reflected_path`: bool, if `True` the received sample includes both the direct and the reflected sound,
        otherwise it includes just the direct sound
        * `include_air_abosrption`: bool, if `True` the air absorption is included in the simulation

    Methods
    ------
    initialize(src_pos: np.ndarray, mic_pos: np.ndarray):
        Sets the initial delays in the primary and secondary delay lines, according to the (initial) position of the
        source and the position of the microphone used for the current simulation
    update(src_pos: np.ndarray, mic_pos: np.ndarray, signal_sample: float):
        Updates delays on the primary and secondary delay lines given the new position of the source, and produces
        a new sample received by the microphone
    """

    def __init__(
            self,
            c: float,                       # speed of sound in air 
            fs: int,                        # sampling frequency
            Z0: int,                        # characteristic impedance of air
            road_material: Material,        # Absorption properties of road surface
            airAbsorptionCoefficients: np.ndarray,
            simulation_params: dict = {
                "interp_method": "Sinc",
                "include_reflected_path": True,
                "include_air_absorption": True,
            },
        ) -> None:
        """
        Creates a `SimulatorManager` object containing all the methods and attributes required to run the acoustic
        simulation. When the constructor is called, the environmental parameters previously defined in the object
        `Environment` are stored into the simulator manager (c, fs, Z0, road_material, airAbsorptionCoefficients). 
        In addition, the two `DelayLine` objects required for the simulation are instantiated as empty circular 
        arrays with 48000 entries. The `asphaltReflectionFilterTable` is pre-computed and stored as an instance 
        attribute containing the FIR filter coefficients that describe asphalt reflection at different angles of
        incidence of the sound waves. This will result in a reduction of the computational load,
        since the computation of a new filter depending on the instantaneous incidence angle is turned into a 
        table lookup. The set of considered incidence angles is stored in `_theta_vector` in degrees. In order to
        reduce the computational complexity, also the pseudo inverse matrix, needed by the least squares filter design
        method in the air absorption filters evaluation, is pre-computed. Finally, a set of auxiliary buffers are 
        instantiated and will be used for filtering purpose.

        Parameters
        ----------
        c : float
            Speed of sound in the air
        fs: int
            Sampling frequency used in the simulation
        Z0: float
            Characteristic impedance of the air at the chosen environmental conditions (i.e. T, p)
        road_material: Material
            Material object defining the absorption and reflection properties of the road surface
        airAbsorptionCoefficients: np.ndarray
            1D Array containing the air absorption coefficients computed in the defined frequency bands at the chosen 
            environmental conditions (i.e. T, p, rel. humidity)
        
        Raises
        ------
        KeyError:
            If simulation_params does not contain keys `interp_method`, `include_reflected_path` and 
            `include_air_absorption`.
        """
        self.c = c
        self.fs = fs
        self.Z0 = Z0
        self.road_material = road_material
        self.airAbsorptionCoefficients = airAbsorptionCoefficients

        # Define simulation parameters
        if ("interp_method" not in simulation_params or "include_reflected_path" 
            not in simulation_params or "include_air_absorption" not in simulation_params):
            raise KeyError("simulation_params must be a dict with keys 'interp_method' (string),"
              "'include_reflection_path' (bool) and 'include_air_absorption (bool'")
        self.simulation_params = simulation_params

        # Instantiation of two delay lines used for the simulations
        self.primaryDelLine = DelayLine(N = 48000, num_read_ptrs = 2, 
            interpolation = self.simulation_params["interp_method"]) 
        self.secondaryDelLine = DelayLine(N = 48000, num_read_ptrs = 1, 
            interpolation = self.simulation_params["interp_method"])

        # Array containing all possible incidence angles for pre-computed reflection filter table
        self._theta_vector = np.arange(-89,90,1)

        # Instantiation of table containing asphalt reflection filters
        self.asphaltReflectionFilterTable = self._compute_angle_reflection_table(ntaps = 11)
       
        # Buffers to store previous data read from delay lines for filtering purpose
        self._read1Buf = deque(np.zeros(11), maxlen=11)
        self._read2Buf = deque(np.zeros(11), maxlen=11)
        self._read3Buf = deque(np.zeros(11), maxlen=11)
        self._read4Buf = deque(np.zeros(11), maxlen=11)

        # Compute air absorption filter 
        freqs = np.linspace(0, self.fs / 2, len(airAbsorptionCoefficients))
        w = 2 * np.pi * freqs / (fs / 2)
        
        # Filter order
        L = 10

        # Matrix A
        A = np.zeros((len(w),int(L/2 + 1)))
        A[:,0] = 1
        for i in range(len(w)):
            for j in range(1,int(L/2 + 1)):
                A[i,j] = 2 * math.cos(w[i] * j)
        
        self._pseud_A = np.linalg.inv(A.T.dot(A)).dot(A.T)
        

    def initialize(self, src_pos: np.ndarray, mic_pos: np.ndarray) -> None:
        """
        Computes the propagation delays along the paths:
        * Direct path from source initial position to microphone
        * Path from image source position to incidence point on asphalt surface, along reflected path
        * Path from incidence point on asphalt surface to microphone, along reflected path
        Sets the corresponding initial delays on the two delay lines.

        This method must be called at the start of every simulation to set the initial conditions that
        ensure no discontinuity is generated at the first simulation instant.

        Parameters
        ----------
        src_pos : np.ndarray
            1D array containing the cartesian coordinates [x,y,z] denoting the initial position of the sound source
        mic_pos : np.ndarray
            1D array containing the cartesian coordinates [x,y,z] denoting the position of the microphone
        """

        # Compute direct distance and delay
        _, tau = self._compute_delay(src_pos, mic_pos)
        
        # Compute incidence angle
        theta = self._compute_angle(src_pos, mic_pos)

        # Compute distance and delay between src and reflection point
        a = src_pos[2] / math.sin(theta)
        tau_1 = a / self.c
        
        # Compute distance and delay between reflection point and microphone
        b = src_pos[2] / math.sin(theta)
        tau_2 = b / self.c

        # Set initial delays
        self.primaryDelLine.set_delays(np.array([tau * self.fs, tau_1 * self.fs]))
        self.secondaryDelLine.set_delays(np.array([tau_2 * self.fs]))
        
    def update(self, src_pos: np.ndarray, mic_pos: np.ndarray, signal_sample: float) -> float:
        """
        This method is used to compute the output of one instant of the simulation. It takes as inputs the instantaneous
        positions of the source and of the microphone used in the current simulation, and the sample of the source
        signal emitted at the considered instant. It then performs the update of the `DelayLine` objects, writing the new 
        sample `signal_sample` on the primary line, and computes the relative distances and delays between the source,
        the road surface at the incidence point, and the microphone, together with the incidence angle. According to
        these delays, the `DelayLine` read pointers positions are updated. Then, the output sample received at the
        microphone position is computed as a sum of two components:
        * the direct component of the sound field, filtered with by air absorption filter ,depending on the distance
        between the source and the microphone, and attenuated depending on the travelled distance;
        * the reflected component: the reflected sound is analyzed using a cascade of 3 filters and two multiplications.
        The reflected sound wave travels from the source to the road surface, along a trajectory computed 
        using the principles of geometrical acoustics. This path is represented by a propagation delay, a filter that
        models the air absorption depending on the distance between the source and the incidence point, and an
        attenuation factor depending on the same travelled distance. At the incidence point, the sound wave is 
        filtered with a second FIR filter representing the asphalt reflection, depending on the road surface material 
        and on the incidence angle. After the reflection, the sound wave travels from the reflection point to the 
        microphone, along a path modeled by a second delay element, a second air absorption filter depending on the
        distance between reflection point and microphone, and a second attenuation factor depending on the same
        distance. 

        The two components are summed to produce the output sample, that is returned by the update function. All the 
        filtering operations are performed in the time domain, as the convolution between the considered signals 
        and the FIR filters.

        Parameters
        ----------
        src_pos : np.ndarray
            1D array containing the cartesian coordinates [x,y,z] denoting the position of the sound source
        mic_pos : np.ndarray
            1D array containing the cartesian coordinates [x,y,z] denoting the position of the microphone
        signal_sample : float
            Sample produced by the source at the considered instant

        Returns
        -------
        float
            Sample of the signal recorded by the cosidered microphone at the given time instant
        """

        # Compute direct distance and delay
        d, tau = self._compute_delay(src_pos, mic_pos)

        # Compute incidence angle
        theta = self._compute_angle(src_pos, mic_pos)

        # Compute distance and delay between src and reflection point
        a = src_pos[2] / math.sin(theta)
        tau_1 = a / self.c
        
        # Compute distance between reflection point and microphone
        b = mic_pos[2] / math.sin(theta)
        tau_2 = b / self.c

        # Update delays and get new sample reads
        y_primary = self.primaryDelLine.update_delay_line(signal_sample, np.array([tau * self.fs, 
            tau_1 * self.fs]))

        # Store the read samples in a circular array for filtering with air absorption and asphalt reflection
        self._read1Buf.appendleft(y_primary[0])
        self._read2Buf.appendleft(y_primary[1])
        
        y_received = 0
        ### DIRECT PATH ###

        if self.simulation_params["include_air_absorption"]:
            # Attenuation due to air absorption
            filt_coeffs = self._compute_air_absorption_filter(d, numtaps = 11)
            sample_eval = filt_coeffs.dot(list(self._read1Buf))
        else:
            sample_eval = self._read1Buf[0]
        
        # Attenuation due to distance
        att = self._compute_sound_attenuation(d)

        # Direct Path Output Sample
        y_dir =  att * sample_eval

        ### REFLECTED PATH ###
        if self.simulation_params["include_reflected_path"]:
        # 1. From Source to Road Surface

            if self.simulation_params["include_air_absorption"]:

                # Attenuation due to air absorption
                filt_coeffs = self._compute_air_absorption_filter(a, numtaps = 11)
                sample_eval = filt_coeffs.dot(list(self._read2Buf))
            else:
                sample_eval = self._read2Buf[0]
            
            # Attenuation due to distance
            att = self._compute_sound_attenuation(a)
            sample_eval = att * sample_eval
            
            # Write output of first filter in buffer, to be able to apply second filter in cascade
            self._read3Buf.appendleft(sample_eval)

            # 2. Asphalt Absorption
            asphalt_filter_coeffs = self._get_asphalt_reflection_filter(90 - math.degrees(theta))
            
            sample_eval = asphalt_filter_coeffs.dot(list(self._read3Buf))
            
            
            # 3. Second path in air --> Secondary Delay Line
            y_secondary = self.secondaryDelLine.update_delay_line(sample_eval, np.array([tau_2 * self.fs]))

            # Store output of secondary delay line in buffer for cascade air abs filter
            self._read4Buf.appendleft(y_secondary[0])

            # 4. From Road Surface to Receiver
            if self.simulation_params["include_air_absorption"]:
                # Attenuation due to air absorption
                filt_coeffs = self._compute_air_absorption_filter(b, numtaps = 11)
                sample_eval = filt_coeffs.dot(list(self._read4Buf))
            else:
                sample_eval = self._read4Buf[0]
            
            # Attenuation due to distance
            att = self._compute_sound_attenuation(b)

            y_refl = att * sample_eval

            # Compute received sample as sum of direct and reflected path outputs
            y_received = y_refl
        
        y_received = y_received + y_dir

        return y_received
    
    def _compute_air_absorption_filter(self, distance: float, numtaps: int) -> np.ndarray:
        """
        Computes air absorption filter as a FIR filter with `numtaps` coefficients. The filter depends
        on the distance between the source and the receiver.

        Parameters
        ----------
        distance : float
            Distance travelled by the sound wave, in meters
        numtaps : int
            Number of coefficients of the FIR filter

        Returns
        -------
        np.ndarray
            1D array containing `numtaps` FIR filter coefficients modelling air absorption
        """
        filt_coeffs = np.empty(numtaps, np.float64)
        alpha = 10 ** (-self.airAbsorptionCoefficients * distance / 20)     # Convert coeffs in dB to linear scale
        filt_coeffs[int((numtaps+1)/2)-1:] = self._pseud_A.dot(alpha)
        filt_coeffs[0:int((numtaps+1)/2)-1] = np.flip(filt_coeffs[int((numtaps+1)/2):])
        
        return filt_coeffs

    def _compute_angle_reflection_table(self, ntaps: int = 11) -> np.ndarray:
        """
        Computes the coefficients of the FIR filters modelling asphalt reflection, at a set of different
        incidence angles, and stores them in a table. Each row of the table contains the `ntaps` FIR filter 
        corresponding to a different incidence angle. The set of angles is defined in `self._theta_vector`.

        The filter coefficients are computed using the `scipy.signal.firwin2` method, given the reflection coefficients
        computed in the frequency bands `center_freqs`, defined in the instance attribute `road_material`. 
        The reflection coefficients are derived from the  absorption coefficients of the road surface, defined
        in the instance attribute `road_material`. The road surface is assumed to be  locally reactive, meaning 
        that its impedance Z does not depend on the incidence angle.

        Parameters
        ----------
        ntaps : int
            Number of FIR filter coefficients, must be odd

        Returns
        -------
        np.ndarray
            2D Array with len(self._theta_vector) rows, each containing `ntaps` FIR filter coefficients corresponding
            to each angle defined in `self._theta_vector`
        """
        
        Z0 = self.Z0
        refl = np.sqrt(1 - np.array(self.road_material.absorption["coeffs"]))
        
        # The impedance does not depend on the angle of incidence -> locally reactive surface
        Z = - Z0 * (refl + 1) / (refl-1)
        
        b_fir = np.zeros((len(self._theta_vector), ntaps))
        
        # Compute filters coefficients for all thetas
        for idx, theta in enumerate(self._theta_vector):
            R = np.abs((Z * np.cos(math.radians(theta)) - Z0) / (Z * np.cos(math.radians(theta)) + Z0))
            b_fir[idx] = scipy.signal.firwin2(ntaps, np.array(self.road_material.absorption["center_freqs"]) 
                / (self.fs/2), R)
        
        return b_fir
    
    def _get_asphalt_reflection_filter(self, theta: float) -> np.ndarray:
        """
        Select the asphalt reflection filter from the pre-computed table given the incidence angle of the 
        sound wave hitting the road surface, in degrees. For the table lookup, the incidence angle is 
        rounded to the closest integer in the range defined in `self._theta_vector` attribute.

        Parameters
        ----------
        theta : float
            Incidence angle of the sound wave hitting the road surface, in degrees

        Returns
        -------
        np.ndarray
            1D Array containing asphalt reflection filter coefficients computed at angle theta
        """

        idx = np.where(self._theta_vector == round(theta))
        idx = idx[0][0]
        return self.asphaltReflectionFilterTable[idx]

    def _compute_sound_attenuation(self, distance: float) -> float:
        """
        Computes the attenuation of sound due to its propagation in the empty space. This attenuation depends
        only on the distance between the source and the receiver and is computed assuming a spherical propagation
        model.

        Parameters
        ----------
        distance : float
            Relative distance between the sound source and the receiver

        Returns
        -------
        float
            Attenuation factor
        """
        return 1 / (distance)
    
    def _compute_delay(self, src_pos: np.ndarray, mic_pos: np.ndarray) -> tuple[float, float]:
        """
        Computes the distance between the source and the microphone, and the time needed by the sound to travel 
        from source to microphone, given their distance and the speed of sound in air.

        Parameters
        ----------
        src_pos : np.ndarray
            1D array containing the cartesian coordinates [x,y,z] denoting the position of the source
        mic_pos : np.ndarray
            1D array containing the cartesian coordinates [x,y,z] denoting the position of the microphone

        Returns
        -------
        Tuple[float, float]
            Tuple containing the distance between source and microphone (in meters) and the time delay (in seconds)
        """

        d = math.sqrt(np.sum((src_pos - mic_pos) ** 2))
        tau = d / self.c
        return d, tau
    
    def _compute_angle(self, src_pos: np.ndarray, mic_pos: np.ndarray) -> float:
        """
        Computes the incidence angle of the sound wave travelling from the source to the road surface, along the
        reflected path connecting the source to the microphone with a single reflection generated by the road.

        Parameters
        ----------
        src_pos : np.ndarray
            1D array containing the cartesian coordinates [x,y,z] denoting the position of the source
        mic_pos : np.ndarray
            1D array containing the cartesian coordinates [x,y,z] denoting the position of the microphone

        Returns
        -------
        float
            Incidence angle of sound wave travelling from source to road surface along reflected path, in radians
        """

        # Distance between image and microphone
        dist = math.sqrt(np.sum(((src_pos - np.array([0, 0 , 2*src_pos[2]])) - mic_pos) ** 2))
        # Incidence angle
        theta = math.asin((src_pos[2] + mic_pos[2]) / dist)
        return theta