import torch
import math
from typing import Optional, List

from . import Encoder, Decoder, STFTFB
from .stft_fb import perfect_synthesis_window
from . import transforms


def griffin_lim(mag_specgram, stft_enc, angles=None, istft_dec=None, n_iter=6, momentum=0.9):
    """Estimates matching phase from magnitude spectogram using the
    'fast' Griffin Lim algorithm [1].

    Args:
        mag_specgram (torch.Tensor): (any, dim, ension, freq, frames) as
            returned by `Encoder(STFTFB)`, the magnitude spectrogram to be
            inverted.
        stft_enc (Encoder[STFTFB]): The `Encoder(STFTFB())` object that was
            used to compute the input `mag_spec`.
        angles (None or Tensor): Angles to use to initialize the algorithm.
            If None (default), angles are init with uniform ditribution.
        istft_dec (None or Decoder[STFTFB]): Optional Decoder to use to get
            back to the time domain. If None (default), a perfect
            reconstruction Decoder is built from `stft_enc`.
        n_iter (int): Number of griffin-lim iterations to run.
        momentum (float): The momentum of fast Griffin-Lim. Original
            Griffin-Lim is obtained for momentum=0.

    Returns:
        torch.Tensor: estimated waveforms of shape (any, dim, ension, time).

    Examples
        >>> stft = Encoder(STFTFB(n_filters=256, kernel_size=256, stride=128))
        >>> wav = torch.randn(2, 1, 8000)
        >>> spec = stft(wav)
        >>> masked_spec = spec * torch.sigmoid(torch.randn_like(spec))
        >>> mag = transforms.mag(masked_spec, -2)
        >>> est_wav = griffin_lim(mag, stft, n_iter=32)

    References
        - [1] Perraudin et al. "A fast Griffin-Lim algorithm," WASPAA 2013.
        - [2] D. W. Griffin and J. S. Lim:  "Signal estimation from modified
        short-time Fourier transform," ASSP 1984.

    """
    # We can create perfect iSTFT from STFT Encoder
    if istft_dec is None:
        # Compute window for perfect resynthesis
        syn_win = perfect_synthesis_window(stft_enc.filterbank.window, stft_enc.stride)
        istft_dec = Decoder(STFTFB(**stft_enc.get_config(), window=syn_win))

    # If no intitial phase is provided initialize uniformly
    if angles is None:
        angles = 2 * math.pi * torch.rand_like(mag_specgram, device=mag_specgram.device)
    else:
        angles = angles.view(*mag_specgram.shape)

    # Initialize rebuilt (useful to use momentum)
    rebuilt = 0.0
    for _ in range(n_iter):
        prev_built = rebuilt
        # Go to the time domain
        complex_specgram = transforms.from_magphase(mag_specgram, angles)
        waveform = istft_dec(complex_specgram)
        # And back to TF domain
        rebuilt = stft_enc(waveform)
        # Update phase estimates (with momentum)
        diff = rebuilt - momentum / (1 + momentum) * prev_built
        angles = transforms.angle(diff)

    final_complex_spec = transforms.from_magphase(mag_specgram, angles)
    return istft_dec(final_complex_spec)


def misi(
    mixture_wav,
    mag_specgrams,
    stft_enc,
    angles=None,
    istft_dec=None,
    n_iter=6,
    momentum=0.0,
    src_weights=None,
    dim=1,
):
    """Jointly estimates matching phase from magnitude spectograms using the
    Multiple Input Spectrogram Inversion (MISI) algorithm [1].

    Args:
        mixture_wav (torch.Tensor): (batch, time)
        mag_specgrams (torch.Tensor): (batch, n_src, freq, frames) as
            returned by `Encoder(STFTFB)`, the magnitude spectrograms to be
            jointly inverted using MISI (modified or not).
        stft_enc (Encoder[STFTFB]): The `Encoder(STFTFB())` object that was
            used to compute the input `mag_spec`.
        angles (None or Tensor): Angles to use to initialize the algorithm.
            If None (default), angles are init with uniform ditribution.
        istft_dec (None or Decoder[STFTFB]): Optional Decoder to use to get
            back to the time domain. If None (default), a perfect
            reconstruction Decoder is built from `stft_enc`.
        n_iter (int): Number of MISI iterations to run.
        momentum (float): Momentum on updates (this argument comes from
            GriffinLim). Defaults to 0 as it was never proposed anywhere.
        src_weights (None or torch.Tensor): Consistency weight for each source.
            Shape needs to be broadcastable to `istft_dec(mag_specgrams)`.
            We make sure that the weights sum up to 1 along dim `dim`.
            If `src_weights` is None, compute them based on relative power.
        dim (int): Axis which contains the sources in `mag_specgrams`.
            Used for consistency constraint.

    Returns:
        torch.Tensor: estimated waveforms of shape (batch, n_src, time).

    Examples
        >>> stft = Encoder(STFTFB(n_filters=256, kernel_size=256, stride=128))
        >>> wav = torch.randn(2, 3, 8000)
        >>> specs = stft(wav)
        >>> masked_specs = specs * torch.sigmoid(torch.randn_like(specs))
        >>> mag = transforms.mag(masked_specs, -2)
        >>> est_wav = misi(wav.sum(1), mag, stft, n_iter=32)

    References
        [1] Gunawan and Sen, "Iterative Phase Estimation for the Synthesis of
        Separated Sources From Single-Channel Mixtures," in IEEE Signal
        Processing Letters, 2010.
        [2] Wang, LeRoux et al. “End-to-End Speech Separation with Unfolded
        Iterative Phase Reconstruction.” Interspeech 2018 (2018)
    """
    # We can create perfect iSTFT from STFT Encoder
    if istft_dec is None:
        # Compute window for perfect resynthesis
        syn_win = perfect_synthesis_window(stft_enc.filterbank.window, stft_enc.stride)
        istft_dec = Decoder(STFTFB(**stft_enc.get_config(), window=syn_win))

    # If no intitial phase is provided initialize uniformly
    if angles is None:
        angles = 2 * math.pi * torch.rand_like(mag_specgrams, device=mag_specgrams.device)
    # wav_dim is used in mixture_consistency.
    # Transform spec src dim to wav src dim for positive and negative dim
    wav_dim = dim if dim >= 0 else dim + 1

    # We forward/backward the mixture through STFT to have matching shapes
    # with the input spectrograms as well as  account for potential modulations
    # if the window were not chosen to enable perfect reconstruction.
    mixture_wav = istft_dec(stft_enc(mixture_wav))

    # Initialize rebuilt (useful to use momentum)
    rebuilt = 0.0
    for _ in range(n_iter):
        prev_built = rebuilt
        # Go to the time domain
        complex_specgram = transforms.from_magphase(mag_specgrams, angles)
        wavs = istft_dec(complex_specgram)
        # Make wavs sum up to the mixture
        consistent_wavs = mixture_consistency(
            mixture_wav, wavs, src_weights=src_weights, dim=wav_dim
        )
        # Back to TF domain
        rebuilt = stft_enc(consistent_wavs)
        # Update phase estimates (with momentum). Keep the momentum here
        # in case. Was shown useful in GF, might be here. We'll see.
        diff = rebuilt - momentum / (1 + momentum) * prev_built
        angles = transforms.angle(diff)
    # Final source estimates
    final_complex_spec = transforms.from_magphase(mag_specgrams, angles)
    return istft_dec(final_complex_spec)


def mixture_consistency(
    mixture: torch.Tensor,
    est_sources: torch.Tensor,
    src_weights: Optional[torch.Tensor] = None,
    dim: int = 1,
) -> torch.Tensor:
    """Applies mixture consistency to a tensor of estimated sources.

    Args
        mixture (torch.Tensor): Mixture waveform or TF representation.
        est_sources (torch.Tensor): Estimated sources waveforms or TF
            representations.
        src_weights (torch.Tensor): Consistency weight for each source.
            Shape needs to be broadcastable to `est_source`.
            We make sure that the weights sum up to 1 along dim `dim`.
            If `src_weights` is None, compute them based on relative power.
        dim (int): Axis which contains the sources in `est_sources`.

    Returns
        torch.Tensor with same shape as `est_sources`, after applying mixture
        consistency.

    Notes
        This method can be used only in 'complete' separation tasks, otherwise
        the residual error will contain unwanted sources. For example, this
        won't work with the task `sep_noisy` from WHAM.

    Examples
        >>> # Works on waveforms
        >>> mix = torch.randn(10, 16000)
        >>> est_sources = torch.randn(10, 2, 16000)
        >>> new_est_sources = mixture_consistency(mix, est_sources, dim=1)
        >>> # Also works on spectrograms
        >>> mix = torch.randn(10, 514, 400)
        >>> est_sources = torch.randn(10, 2, 514, 400)
        >>> new_est_sources = mixture_consistency(mix, est_sources, dim=1)

    References
        - Scott Wisdom, John R Hershey, Kevin Wilson, Jeremy Thorpe, Michael
        Chinen, Brian Patton, and Rif A Saurous. "Differentiable consistency
        constraints for improved deep speech enhancement", ICASSP 2019.
    """
    # If the source weights are not specified, the weights are the relative
    # power of each source to the sum. w_i = P_i / (P_all), P for power.
    if src_weights is None:
        all_dims: List[int] = torch.arange(est_sources.ndim).tolist()
        all_dims.pop(dim)  # Remove source axis
        all_dims.pop(0)  # Remove batch axis
        src_weights = torch.mean(est_sources ** 2, dim=all_dims, keepdim=True)
    # Make sure that the weights sum up to 1
    norm_weights = torch.sum(src_weights, dim=dim, keepdim=True) + 1e-8
    src_weights = src_weights / norm_weights

    # Compute residual mix - sum(est_sources)
    if mixture.ndim == est_sources.ndim - 1:
        # mixture (batch, *), est_sources (batch, n_src, *)
        residual = (mixture - est_sources.sum(dim=dim)).unsqueeze(dim)
    elif mixture.ndim == est_sources.ndim:
        # mixture (batch, 1, *), est_sources (batch, n_src, *)
        residual = mixture - est_sources.sum(dim=dim, keepdim=True)
    else:
        n, m = est_sources.ndim, mixture.ndim
        raise RuntimeError(
            f"The size of the mixture tensor should match the "
            f"size of the est_sources tensor. Expected mixture"
            f"tensor to have {n} or {n-1} dimension, found {m}."
        )
    # Compute remove
    new_sources = est_sources + src_weights * residual
    return new_sources
