"""Beamformer for multi-mic processing.

Authors
 * Nauman Dawalatabad
"""

import torch

from speechbrain.processing.features import ISTFT, STFT
from speechbrain.processing.multi_mic import Covariance, DelaySum, GccPhat


class DelaySum_Beamformer(torch.nn.Module):
    """Generate beamformed signal from multi-mic data using DelaySum beamforming.

    Arguments
    ---------
    sampling_rate : int (default: 16000)
        Sampling rate of audio signals.
    """

    def __init__(self, sampling_rate=16000):
        super().__init__()
        self.fs = sampling_rate
        self.stft = STFT(sample_rate=self.fs)
        self.cov = Covariance()
        self.gccphat = GccPhat()
        self.delaysum = DelaySum()
        self.istft = ISTFT(sample_rate=self.fs)

    def forward(self, mics_signals):
        """Returns beamformed signal using multi-mic data.

        Arguments
        ---------
        mics_signals : torch.Tensor
            Set of audio signals to be transformed.

        Returns
        -------
        sig : torch.Tensor
        """
        with torch.no_grad():

            Xs = self.stft(mics_signals)
            XXs = self.cov(Xs)
            tdoas = self.gccphat(XXs)
            Ys_ds = self.delaysum(Xs, tdoas)
            sig = self.istft(Ys_ds)

        return sig