
import numpy as np
from featureflow import Node
from scipy.fftpack import dct
from scipy.stats.mstats import gmean

from .functional import fft, mdct
from .frequencyscale import LinearScale, ChromaScale, BarkScale
from .weighting import AWeighting
from .tfrepresentation import FrequencyDimension
from .frequencyadaptive import FrequencyAdaptive
from zounds.core import ArrayWithUnits, IdentityDimension
from zounds.nputil import safe_log
from zounds.timeseries import audio_sample_rate
from .sliding_window import HanningWindowingFunc


class FrequencyWeighting(Node):
    """
    `FrequencyWeighting` is a processing node that expects to be passed an
    :class:`~zounds.core.ArrayWithUnits` instance whose last dimension is a
    :class:`~zounds.spectral.FrequencyDimension`

    Args:
        weighting (FrequencyWeighting): the frequency weighting to apply
        needs (Node): a processing node on which this node depends whose last
            dimension is a :class:`~zounds.spectral.FrequencyDimension`
    """

    def __init__(self, weighting=None, needs=None):
        super(FrequencyWeighting, self).__init__(needs=needs)
        self.weighting = weighting

    def _process(self, data):
        yield data * self.weighting


class FFT(Node):
    """
    A processing node that performs an FFT of a real-valued signal

    Args:
        axis (int): The axis over which the FFT should be computed
        padding_samples (int): number of zero samples to pad each window with
            before applying the FFT
        needs (Node): a processing node on which this one depends

    See Also:
        :class:`~zounds.synthesize.FFTSynthesizer`
    """

    def __init__(self, needs=None, axis=-1, padding_samples=0):
        super(FFT, self).__init__(needs=needs)
        self._axis = axis
        self._padding_samples = padding_samples

    def _process(self, data):
        yield fft(data, axis=self._axis, padding_samples=self._padding_samples)


class DCT(Node):
    """
    A processing node that performs a Type II Discrete Cosine Transform
    (https://en.wikipedia.org/wiki/Discrete_cosine_transform#DCT-II) of the
    input

    Args:
        axis (int): The axis over which to perform the DCT transform
        needs (Node): a processing node on which this one depends

    See Also:
        :class:`~zounds.synthesize.DctSynthesizer`
    """

    def __init__(self, axis=-1, scale_always_even=False, needs=None):
        super(DCT, self).__init__(needs=needs)
        self.scale_always_even = scale_always_even
        self._axis = axis

    def _process(self, data):
        transformed = dct(data, norm='ortho', axis=self._axis)

        sr = audio_sample_rate(
            int(data.shape[1] / data.dimensions[0].duration_in_seconds))
        scale = LinearScale.from_sample_rate(
            sr, transformed.shape[-1], always_even=self.scale_always_even)

        yield ArrayWithUnits(
            transformed, [data.dimensions[0], FrequencyDimension(scale)])


class DCTIV(Node):
    """
    A processing node that performs a Type IV Discrete Cosine Transform
    (https://en.wikipedia.org/wiki/Discrete_cosine_transform#DCT-IV) of the
    input

    Args:
        needs (Node): a processing node on which this one depends

    See Also:
        :class:`~zounds.synthesize.DCTIVSynthesizer`
    """

    def __init__(self, scale_always_even=False, needs=None):
        super(DCTIV, self).__init__(needs=needs)
        self.scale_always_even = scale_always_even

    def _process_raw(self, data):
        l = data.shape[1]
        tf = np.arange(0, l)
        z = np.zeros((len(data), l * 2))
        z[:, :l] = (data * np.exp(-1j * np.pi * tf / 2 / l)).real
        z = np.fft.fft(z)[:, :l]
        raw = np.sqrt(2 / l) * \
              (z * np.exp(-1j * np.pi * (tf + 0.5) / 2 / l)).real
        return raw

    def _process(self, data):
        raw = self._process_raw(data)
        sr = audio_sample_rate(
            int(data.shape[1] / data.dimensions[0].duration_in_seconds))