from dataclasses import dataclass, field
from typing import List, Optional

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
@dataclass
class FastRandomRIRGenerator:
    sr: int = 16000
    direct_range: List = field(default_factory=lambda: [-6, 50])
    max_T60: float = 0.8
    alpha: float = 0.25
    a: float = -2.0
    b: float = 2.0
    tau: float = 0.2
    room_seed: Optional[int] = None
    source_seed: Optional[int] = None

    def __post_init__(self):
        self.room_rng = (
            np.random.default_rng(self.room_seed)
            if self.room_seed is not None
            else np.random.default_rng()
        )
        self.source_rng = (
            np.random.default_rng(self.source_seed)
            if self.source_seed is not None
            else np.random.default_rng()
        )

    def __call__(self, nsource: int = 1) -> np.ndarray:
        """
        :param nsource: number of sources (RIR filters) to simulate. Default: 1.
        :return: simulated RIR filter for all sources, shape: (nsource, nsample)
        """
        from lhotse.augmentation.torchaudio import (
            check_for_torchaudio,
            get_or_create_resampler,
        )

        check_for_torchaudio()

        from torchaudio.functional import highpass_biquad

        # the sample rate at which the original RIR filter is generated
        ratio = 64
        sample_sr = self.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)), self.sr)

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

        # sample T60 of the room
        T60 = torch.from_numpy(self.room_rng.uniform(0.1, self.max_T60, size=(1,)))[
            0
        ].data

        # sample room-related statistics for calculating the reflection coefficient R
        R = torch.from_numpy(self.room_rng.uniform(0.1, 1.2, size=(1,)))[0].data

        # sample distance between the sound sources and the receiver (d_0) if not given
        direct_dist = torch.from_numpy(
            self.source_rng.uniform(0.2, 12.0, size=(nsource,))
        )

        # number of virtual sound sources
        image = self.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(self.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.from_numpy(
            self.source_rng.uniform(self.a, self.b, size=(nsource, image))
        ) * (dist_ratio.pow(self.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 * self.direct_range[0] // 1000, 0) : min(
                    direct_idx[i] + sample_sr * self.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()
