import numpy as np
from torch.utils.data import Dataset

from ..transforms import functional as F
from .utils import sampling_rate_after_transform


__doctest_skip__ = ['*']


class WhiteNoise(Dataset):
    r"""White noise data set.

    The white noise is generated by numpy.random.standard_normal.

    * :attr:`duration` controls the duration of the noise signal
    * :attr:`sampling_rate` holds the sampling rate of the returned data
    * :attr:`mean` controls the mean of the underlying distribution
    * :attr:`stdev` controls the standard deviation of the underlying
      distribution
    * :attr:`transform` controls the input transform
    * :attr:`target_transform` controls the target transform

    As white noise has not really a sampling rate you can use the following
    attribute to change it instead of resampling:

    * :attr:`original_sampling_rate` controls the sampling rate of the data set

    Args:
        duration (float): duration of the noise signal in seconds
        sampling_rate (int, optional): sampling rate in Hz. Default: `44100`
        mean (float, optional): mean of underlying distribution. Default: `0`
        stdev (float, optional): standard deviation of underlying distribution.
            Default: `1`
        transform (callable, optional): function/transform applied on the
            signal. Default: `None`
        target_transform (callable, optional): function/transform applied on
            the target. Default: `None`

    Note:
        Even `WhiteNoise` has an infintely number of entries, its length is
        `1` as repeated calls of the same index return different signals.

    Example:
        >>> import sounddevice as sd
        >>> data = WhiteNoise(duration=1, sampling_rate=44100)
        >>> print(data)
        Dataset WhiteNoise
            Number of data points: Inf
            Signal length: 1s
            Sampling Rate: 44100Hz
            Label (str): noise type
        >>> signal, target = data[0]
        >>> target
        'white noise'
        >>> sd.play(signal.transpose(), data.sampling_rate)

    """

    def __init__(
            self,
            *,
            duration=1,
            sampling_rate=44100,
            mean=0,
            stdev=1,
            transform=None,
            target_transform=None,
    ):
        super().__init__()
        self.duration = duration
        self.mean = mean
        self.stdev = stdev
        self.transform = transform
        self.target_transform = target_transform
        self.original_sampling_rate = sampling_rate

    @property
    def sampling_rate(self):
        return sampling_rate_after_transform(self)

    def __len__(self):
        # This has no meaningful __len__ as its actual length is `Inf`.
        # Return `1` to make it work with `random.choice` like operations.
        return 1

    def __getitem__(self, item):
        samples = int(np.ceil(self.duration * self.sampling_rate))
        signal = np.random.normal(loc=self.mean, scale=self.stdev,
                                  size=samples)
        signal = F.normalize(np.expand_dims(signal, axis=0))
        target = 'white noise'

        if self.transform is not None:
            signal = self.transform(signal)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return signal, target

    def __repr__(self):
        fmt_str = f'Dataset {self.__class__.__name__}\n'
        fmt_str += '    Number of data points: Inf\n'
        fmt_str += f'    Signal length: {self.duration}s\n'
        if self.sampling_rate == self.original_sampling_rate:
            fmt_str += f'    Sampling Rate: {self.sampling_rate}Hz\n'
        else:
            fmt_str += (
                f'    Sampling Rate: {self.sampling_rate}Hz '
                f'(original: {self.original_sampling_rate}Hz)\n'
            )
        fmt_str += '    Label (str): noise type\n'
        tmp1 = '    Transform: '
        tmp2 = self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp1))
        if self.transform:
            fmt_str += f'{tmp1}{tmp2}\n'
        tmp1 = '    Target Transform: '
        tmp2 = self.target_transform.__repr__().replace('\n',
                                                        '\n' + ' ' * len(tmp1))
        if self.target_transform:
            fmt_str += f'{tmp1}{tmp2}'
        return fmt_str
