"""PyTorch compatible cwt code.

This module is based on pywt's cwt implementation.
"""
from typing import Any, Tuple, Union

import numpy as np
import torch
from pywt import ContinuousWavelet, DiscreteContinuousWavelet, Wavelet
from pywt._functions import scale2frequency
from torch.fft import fft, ifft


def _next_fast_len(n: int) -> int:
    """Round up size to the nearest power of two.

    Given a number of samples `n`, returns the next power of two
    following this number to take advantage of FFT speedup.
    This fallback is less efficient than `scipy.fftpack.next_fast_len`
    """
    return int(2 ** np.ceil(np.log2(n)))


def cwt(
    data: torch.Tensor,
    scales: Union[np.ndarray, torch.Tensor],  # type: ignore
    wavelet: Union[ContinuousWavelet, str],
    sampling_period: float = 1.0,
) -> Tuple[torch.Tensor, np.ndarray]:  # type: ignore
    """Compute the single dimensional continuous wavelet transform.

    This function is a PyTorch port of pywt.cwt as found at:
    https://github.com/PyWavelets/pywt/blob/master/pywt/_cwt.py

    Args:
        data (torch.Tensor): The input tensor of shape [batch_size, time].
        scales (torch.Tensor or np.array):
            The wavelet scales to use. One can use
            ``f = pywt.scale2frequency(wavelet, scale)/sampling_period`` to determine
            what physical frequency, ``f``. Here, ``f`` is in hertz when the
            ``sampling_period`` is given in seconds.
        wavelet (ContinuousWavelet or str): The continuous wavelet to work with.
        sampling_period (float): Sampling period for the frequencies output (optional).
            The values computed for ``coefs`` are independent of the choice of
            ``sampling_period`` (i.e. ``scales`` is not scaled by the sampling
            period).

    Raises:
        ValueError: If a scale is too small for the input signal.

    Returns:
        Tuple[torch.Tensor, np.ndarray]: The first tuple-element contains
            the transformation matrix of shape [scales, batch, time].
            The second element contains an array with frequency information.

    Example:
        >>> import torch, ptwt
        >>> import numpy as np
        >>> import scipy.signal as signal
        >>> t = np.linspace(-2, 2, 800, endpoint=False)
        >>> sig = signal.chirp(t, f0=1, f1=12, t1=2, method="linear")
        >>> widths = np.arange(1, 31)
        >>> cwtmatr, freqs = ptwt.cwt(
        >>>     torch.from_numpy(sig), widths, "mexh", sampling_period=(4 / 800) * np.pi
        >>> )
    """
    # accept array_like input; make a copy to ensure a contiguous array
    if not isinstance(
        wavelet, (ContinuousWavelet, Wavelet, _DifferentiableContinuousWavelet)
    ):
        wavelet = DiscreteContinuousWavelet(wavelet)
    if type(scales) is torch.Tensor:
        scales = scales.numpy()
    elif np.isscalar(scales):
        scales = np.array([scales])

    if isinstance(wavelet, torch.nn.Module):
        if data.is_cuda:
            wavelet.cuda()

    precision = 10
    int_psi, x = _integrate_wavelet(wavelet, precision=precision)
    if type(wavelet) is ContinuousWavelet:
        int_psi = np.conj(int_psi) if wavelet.complex_cwt else int_psi
        int_psi = torch.tensor(int_psi, device=data.device)
    elif isinstance(wavelet, torch.nn.Module):
        int_psi = torch.conj(int_psi) if wavelet.complex_cwt else int_psi
    else:
        int_psi = torch.tensor(int_psi, device=data.device)
        x = torch.tensor(x, device=data.device)

    # convert int_psi, x to the same precision as the data
    # x = np.asarray(x, dtype=data.cpu().numpy().real.dtype)

    size_scale0 = -1
    fft_data = None

    out = []
    for scale in scales:
        step = x[1] - x[0]
        j = torch.arange(
            scale * (x[-1] - x[0]) + 1, device=data.device, dtype=data.dtype
        ) / (scale * step)
        j = torch.floor(j).type(torch.long)
        if j[-1] >= len(int_psi):
            # j = np.extract(j < len(int_psi), j)
            j = torch.masked_select(j, j < len(int_psi))
        int_psi_scale = int_psi[j].flip(0)

        # The padding is selected for:
        # - optimal FFT complexity
        # - to be larger than the two signals length to avoid circular
        #   convolution
        size_scale = _next_fast_len(data.shape[-1] + len(int_psi_scale) - 1)
        if size_scale != size_scale0:
            # Must recompute fft_data when the padding size changes.
            fft_data = fft(data, size_scale, dim=-1)
        size_scale0 = size_scale
        fft_wav = fft(int_psi_scale, size_scale, dim=-1)
        conv = ifft(fft_wav * fft_data, dim=-1)
        conv = conv[..., : data.shape[-1] + len(int_psi_scale) - 1]

        coef = -np.sqrt(scale) * torch.diff(conv, dim=-1)

        # transform axis is always -1
        d = (coef.shape[-1] - data.shape[-1]) / 2.0
        if d > 0:
            coef = coef[..., int(np.floor(d)) : -int(np.ceil(d))]
        elif d < 0:
            raise ValueError("Selected scale of {} too small.".format(scale))

        out.append(coef)
    out_tensor = torch.stack(out)
    if type(wavelet) is Wavelet:
        out_tensor = out_tensor.real
    elif isinstance(wavelet, _DifferentiableContinuousWavelet):
        out_tensor = out_tensor if wavelet.complex_cwt else out_tensor.real
        wavelet.cpu()
    else:
        out_tensor = out_tensor if wavelet.complex_cwt else out_tensor.real

    with torch.no_grad():
        frequencies = scale2frequency(wavelet, scales, precision)
        if np.isscalar(frequencies):
            frequencies = np.array([frequencies])
    frequencies /= sampling_period

    if isinstance(wavelet, _DifferentiableContinuousWavelet):
        if data.is_cuda:
            wavelet.cuda()

    return out_tensor, frequencies


def _integrate_wavelet(
    wavelet: Union[ContinuousWavelet, str], precision: int = 8
) -> Any:
    """
    Integrate `psi` wavelet function from -Inf to x using rectangle integration.

    Modified to enable gradient flow through the cwt.

    Ported from:
    https://github.com/PyWavelets/pywt/blob/cef09e7f419aaf4c39b9f778bdc2d54b32fd7337/pywt/_functions.py#L60


    Parameters
    ----------
    wavelet : Wavelet instance or str
        Wavelet to integrate.  If a string, should be the name of a wavelet.
    precision : int, optional
        Precision that will be used for wavelet function
        approximation computed with the wavefun(level=precision)
        Wavelet's method (default: 8).
    Returns
    -------
    [int_psi, x] :
        for orthogonal wavelets
    [int_psi_d, int_psi_r, x] :
        for other wavelets
    Examples
    --------
    >>> from pywt import Wavelet, _integrate_wavelet
    >>> wavelet1 = Wavelet('db2')
    >>> [int_psi, x] = _integrate_wavelet(wavelet1, precision=5)
    >>> wavelet2 = Wavelet('bior1.3')
    >>> [int_psi_d, int_psi_r, x] = _integrate_wavelet(wavelet2, precision=5)
    """

    def _integrate(
        arr: Union[np.ndarray, torch.Tensor],  # type: ignore
        step: Union[np.ndarray, torch.Tensor],  # type: ignore
    ) -> Union[np.ndarray, torch.Tensor]:  # type: ignore
        if type(arr) is np.ndarray:
            integral = np.cumsum(arr)
        elif type(arr) is torch.Tensor:
            integral = torch.cumsum(arr, -1)
        else:
            raise TypeError("Only ndarrays or tensors are integratable.")
        integral *= step
        return integral

    if type(wavelet) is str:
        wavelet = DiscreteContinuousWavelet(wavelet)
    elif not isinstance(
        wavelet, (Wavelet, ContinuousWavelet, _DifferentiableContinuousWavelet)
    ):
        wavelet = DiscreteContinuousWavelet(wavelet)

    functions_approximations = wavelet.wavefun(precision)

    if len(functions_approximations) == 2:  # continuous wavelet
        psi, x = functions_approximations
        step = x[1] - x[0]
        return _integrate(psi, step), x

    elif len(functions_approximations) == 3:  # orthogonal wavelet
        _, psi, x = functions_approximations
        step = x[1] - x[0]
        return _integrate(psi, step), x

    else:  # biorthogonal wavelet
        _, psi_d, _, psi_r, x = functions_approximations
        step = x[1] - x[0]
        return _integrate(psi_d, step), _integrate(psi_r, step), x


class _WaveletParameter(torch.nn.Parameter):
    pass


class _DifferentiableContinuousWavelet(
    torch.nn.Module, ContinuousWavelet  # type: ignore
):
    """A base class for learnable Continuous Wavelets."""

    def __init__(self, name: str):
        """Create a trainable shannon wavelet."""
        super().__init__()
        super(ContinuousWavelet, self).__init__()

        self.dtype = torch.float64
        # Use torch nn parameter
        self.bandwidth_par = _WaveletParameter(
            torch.sqrt(torch.tensor(self.bandwidth_frequency, dtype=self.dtype)),
            requires_grad=True,
        )
        self.center_par = _WaveletParameter(
            torch.sqrt(torch.tensor(self.center_frequency, dtype=self.dtype)),
            requires_grad=True,
        )

    def __call__(self, grid_values: torch.Tensor) -> torch.Tensor:
        """Return numerical values for the wavelet on a grid."""
        raise NotImplementedError

    @property
    def bandwidth(self) -> torch.Tensor:
        """Square the bandwith parameter to ensure positive values."""
        return self.bandwidth_par * self.bandwidth_par

    @property
    def center(self) -> torch.Tensor:
        """Square the bandwith parameter to ensure positive values."""
        return self.center_par * self.center_par

    def wavefun(
        self, precision: int, dtype: torch.dtype = torch.float64
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Define a grid and evaluate the wavelet on it."""
        length = 2**precision
        # load the bounds from untyped pywt code.
        lower_bound: float = float(self.lower_bound)  # type: ignore
        upper_bound: float = float(self.upper_bound)  # type: ignore
        grid = torch.linspace(
            lower_bound,
            upper_bound,
            length,
            dtype=dtype,
            device=self.bandwidth_par.device,
        )
        return self(grid), grid


class _ShannonWavelet(_DifferentiableContinuousWavelet):
    """A differentiable Shannon wavelet."""

    def __call__(self, grid_values: torch.Tensor) -> torch.Tensor:
        """Return numerical values for the wavelet on a grid."""
        shannon = (
            torch.sqrt(self.bandwidth)
            * (
                torch.sin(torch.pi * self.bandwidth * grid_values)
                / (torch.pi * self.bandwidth * grid_values)
            )
            * torch.exp(1j * 2 * torch.pi * self.center * grid_values)
        )
        return shannon


class _ComplexMorletWavelet(_DifferentiableContinuousWavelet):
    """A differentiable Shannon wavelet."""

    def __call__(self, grid_values: torch.Tensor) -> torch.Tensor:
        """Return numerical values for the wavelet on a grid."""
        morlet = (
            1.0
            / torch.sqrt(torch.pi * self.bandwidth)
            * torch.exp(-(grid_values**2) / self.bandwidth)
            * torch.exp(1j * 2 * torch.pi * self.center * grid_values)
        )
        return morlet
