from typing import List

import numpy as np
import torch

try:
    # Pytorch >= 1.7
    from torch.fft import irfft, rfft
except ImportError:
    from torch import irfft, rfft

# Implementation based on torch-audiomentations:
# https://github.com/asteroid-team/torch-audiomentations/blob/master/torch_audiomentations/utils/convolution.py

_NEXT_FAST_LEN = {}


def next_fast_len(size):
    """
    Returns the next largest number ``n >= size`` whose prime factors are all
    2, 3, or 5. These sizes are efficient for fast fourier transforms.
    Equivalent to :func:`scipy.fftpack.next_fast_len`.
    Note: This function was originally copied from the https://github.com/pyro-ppl/pyro
    repository, where the license was Apache 2.0. Any modifications to the original code can be
    found at https://github.com/asteroid-team/torch-audiomentations/commits
    :param int size: A positive number.
    :returns: A possibly larger number.
    :rtype int:
    """
    try:
        return _NEXT_FAST_LEN[size]
    except KeyError:
        pass

    assert isinstance(size, int) and size > 0
    next_size = size
    while True:
        remaining = next_size
        for n in (2, 3, 5):
            while remaining % n == 0:
                remaining //= n
        if remaining == 1:
            _NEXT_FAST_LEN[size] = next_size
            return next_size
        next_size += 1


def convolve1d(signal: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor:
    """
    Computes the 1-d convolution of signal by kernel using FFTs.
    Both signal and kernel must be 1-dimensional.
    :param torch.Tensor signal: A signal to convolve.
    :param torch.Tensor kernel: A convolution kernel.
    :param str mode: One of: 'full', 'valid', 'same'.
    :return: torch.Tensor Convolution of signal with kernel. Returns the full convolution, i.e.,
        the output tensor will have size m + n - 1, where m is the length of the
        signal and n is the length of the kernel.
    """
    assert (
        signal.ndim == 1 and kernel.ndim == 1
    ), "signal and kernel must be 1-dimensional"
    m = signal.size(-1)
    n = kernel.size(-1)

    # Compute convolution using fft.
    padded_size = m + n - 1
    # Round up for cheaper fft.
    fast_ftt_size = next_fast_len(padded_size)
    f_signal = rfft(signal, n=fast_ftt_size)
    f_kernel = rfft(kernel, n=fast_ftt_size)
    f_result = f_signal * f_kernel
    result = irfft(f_result, n=fast_ftt_size)

    return result[:padded_size]


# The following is based on: https://github.com/yluo42/FRA-RIR/blob/main/FRA-RIR.py
def generate_fast_random_rir(
    nsource: int = 1,
    sr: int = 16000,
    direct_range: List = [-6, 50],
    max_T60: float = 0.8,
    alpha: float = 0.25,
    a: float = -2.0,
    b: float = 2.0,
    tau: float = 0.2,
) -> np.ndarray:
    """
    The fast random approximation of room impulse response (FRA-RIR) method.
    :param nsource: number of sources (RIR filters) to simulate. Default: 1.
    :param sr: target sample rate. Default: 16000.
    :param direct_range: the context range (at milliseconds) at the first peak of the RIR filter to define the direct-path RIR. Default: [-6, 50] ms.
    :param max_T60: the maximum range of T60 to sample from. Default: 0.8.
    :param alpha: controlling the probability distribution to sample the distance of the virtual sound sources from. Default: 0.25.
    :param a, b: controlling the random pertubation added to each virtual sound source. Default: -2, 2.
    :param tau: controlling the relationship between the distance and the number of reflections of each virtual sound source. Default: 0.25.
    :return: simulated RIR filter for all sources, shape: (nsource, nsample)
    """
    from torchaudio.functional import highpass_biquad

    from lhotse.augmentation.torchaudio import get_or_create_resampler

    # the sample rate at which the original RIR filter is generated
    ratio = 64
    sample_sr = sr * ratio

    # two resampling operations
    resample1 = get_or_create_resampler(sample_sr, sample_sr // int(np.sqrt(ratio)))
    resample2 = get_or_create_resampler(sample_sr // int(np.sqrt(ratio)), sr)

    eps = np.finfo(np.float16).eps

    # sample distance between the sound sources and the receiver (d_0)
    direct_dist = torch.FloatTensor(nsource).uniform_(0.2, 12)

    # sample T60 of the room
    T60 = torch.FloatTensor(1).uniform_(0.1, max_T60)[0].data

    # sample room-related statistics for calculating the reflection coefficient R
    R = torch.FloatTensor(1).uniform_(0.1, 1.2)[0].data

    # number of virtual sound sources
    image = sr * 2

    # sound velocity
    velocity = 340.0

    # indices of direct-path signals based on the sampled d_0
    direct_idx = torch.ceil(direct_dist * sample_sr / velocity).long()

    # length of the RIR filter based on the sampled T60
    rir_length = int(np.ceil(sample_sr * T60))

    # calculate the reflection coefficient based on the Eyring's empirical equation
    reflect_coef = (1 - (1 - torch.exp(-0.16 * R / T60)).pow(2)).sqrt()

    # randomly sample the propagation distance for all the virtual sound sources
    dist_range = [
        torch.linspace(1.0, velocity * T60 / direct_dist[i] - 1, image)
        for i in range(nsource)
    ]
    # a simple quadratic function
    dist_prob = torch.linspace(alpha, 1.0, image).pow(2)
    dist_prob = dist_prob / dist_prob.sum()
    dist_select_idx = dist_prob.multinomial(
        num_samples=image * nsource, replacement=True
    ).view(nsource, image)
    # the distance is sampled as a ratio between d_0 and each virtual sound sources
    dist_ratio = torch.stack(
        [dist_range[i][dist_select_idx[i]] for i in range(nsource)], 0
    )
    dist = direct_dist.view(-1, 1) * dist_ratio

    # sample the number of reflections (can be nonintegers)
    # calculate the maximum number of reflections
    reflect_max = (
        torch.log10(velocity * T60) - torch.log10(direct_dist) - 3
    ) / torch.log10(reflect_coef + eps)
    # calculate the number of reflections based on the assumption that
    # virtual sound sources which have longer propagation distances may reflect more frequently
    reflect_ratio = (dist / (velocity * T60)).pow(2) * (
        reflect_max.view(nsource, -1) - 1
    ) + 1
    # add a random pertubation based on the assumption that
    # virtual sound sources which have similar propagation distances can have different routes and reflection patterns
    reflect_pertub = torch.FloatTensor(nsource, image).uniform_(a, b) * dist_ratio.pow(
        tau
    )
    # all virtual sound sources should reflect for at least once
    reflect_ratio = torch.maximum(reflect_ratio + reflect_pertub, torch.ones(1))

    # calculate the rescaled dirac comb as RIR filter
    dist = torch.cat([direct_dist.reshape(-1, 1), dist], 1)
    reflect_ratio = torch.cat([torch.zeros(nsource, 1), reflect_ratio], 1)
    rir = torch.zeros(nsource, rir_length)
    delta_idx = torch.minimum(
        torch.ceil(dist * sample_sr / velocity), torch.ones(1) * rir_length - 1
    ).long()
    delta_decay = reflect_coef.pow(reflect_ratio) / dist
    for i in range(nsource):
        rir[i][delta_idx[i]] += delta_decay[i]

    # a binary mask for direct-path RIR
    direct_mask = torch.zeros(nsource, rir_length).float()
    for i in range(nsource):
        direct_mask[
            i,
            max(direct_idx[i] + sample_sr * direct_range[0] // 1000, 0) : min(
                direct_idx[i] + sample_sr * direct_range[1] // 1000, rir_length
            ),
        ] = 1.0
    rir_direct = rir * direct_mask

    # downsample
    all_rir = torch.stack([rir, rir_direct], 1).view(nsource * 2, -1)
    rir_downsample = resample1(all_rir)

    # apply high-pass filter
    rir_hp = highpass_biquad(rir_downsample, sample_sr // int(np.sqrt(ratio)), 80.0)

    # downsample again
    rir = resample2(rir_hp).float().view(nsource, 2, -1)

    # RIR filter and direct-path RIR filter at target sample rate
    rir_filter = rir[:, 0]  # nsource, nsample

    # convert to numpy array
    return rir_filter.numpy()
