"""
Provides Ops for FFT and DCT.

"""


# This module will soon be deprecated.
import warnings

import numpy as np
import numpy.fft

from theano import tensor
from theano.graph.basic import Apply
from theano.graph.op import Op
from theano.graph.type import generic


message = (
    "The module theano.sandbox.fourier will soon be deprecated."
    " Please use theano.tensor.fft, which supports gradients and "
    "automatic optimization transfers to the GPU ops."
)
warnings.warn(message)


class GradTodo(Op):
    # TODO : need description for class
    __props__ = ()

    def make_node(self, x):
        return Apply(self, [x], [x.type()])

    def perform(self, node, inputs, outputs):
        raise NotImplementedError("TODO")


grad_todo = GradTodo()


class FFT(Op):
    # TODO : need description for parameters
    """
    Fast Fourier Transform.

    .. TODO:
        The current implementation just works for matrix inputs, and permits
        taking a 1D FFT over either rows or columns. Add support for N-D FFTs
        as provided by either numpy or FFTW directly.

    .. TODO:
        Give the C code that uses FFTW.

    .. TODO:
        Unit tests.

    """

    default_output = 0
    # don't return the plan object in the 'buf' output

    half = False
    """Only return the first half (positive-valued) of the frequency
    components."""
    __props__ = ("half", "inverse")

    def __init__(self, half=False, inverse=False):
        self.half = half
        self.inverse = inverse

    def make_node(self, frames, n, axis):
        """
        Compute an n-point fft of frames along given axis.

        """
        _frames = tensor.as_tensor(frames, ndim=2)
        _n = tensor.as_tensor(n, ndim=0)
        _axis = tensor.as_tensor(axis, ndim=0)
        if self.half and _frames.type.dtype.startswith("complex"):
            raise TypeError("Argument to HalfFFT must not be complex", frames)
        spectrogram = tensor.zmatrix()
        buf = generic()
        # The `buf` output is present for future work
        # when we call FFTW directly and re-use the 'plan' that FFTW creates.
        # In that case, buf would store a CObject encapsulating the plan.
        rval = Apply(self, [_frames, _n, _axis], [spectrogram, buf])
        return rval

    def perform(self, node, inp, out):
        frames, n, axis = inp
        spectrogram, buf = out
        if self.inverse:
            fft_fn = numpy.fft.ifft
        else:
            fft_fn = numpy.fft.fft

        fft = fft_fn(frames, int(n), int(axis))
        if self.half:
            M, N = fft.shape
            if axis == 0:
                if M % 2:
                    raise ValueError("halfFFT on odd-length vectors is undefined")
                spectrogram[0] = fft[0 : M / 2, :]
            elif axis == 1:
                if N % 2:
                    raise ValueError("halfFFT on odd-length vectors is undefined")
                spectrogram[0] = fft[:, 0 : N / 2]
            else:
                raise NotImplementedError()
        else:
            spectrogram[0] = fft

    def grad(self, inp, out):
        frames, n, axis = inp
        g_spectrogram, g_buf = out
        return [grad_todo(frames), None, None]


fft = FFT(half=False, inverse=False)
half_fft = FFT(half=True, inverse=False)
ifft = FFT(half=False, inverse=True)
half_ifft = FFT(half=True, inverse=True)


def dct_matrix(rows, cols, unitary=True):
    # TODO : need description for parameters
    """
    Return a (rows x cols) matrix implementing a discrete cosine transform.

    This algorithm is adapted from Dan Ellis' Rastmat spec2cep.m, lines 15-20.

    """
    rval = np.zeros((rows, cols))
    col_range = np.arange(cols)
    scale = np.sqrt(2.0 / cols)
    for i in range(rows):
        rval[i] = np.cos(i * (col_range * 2 + 1) / (2.0 * cols) * np.pi) * scale

    if unitary:
        rval[0] *= np.sqrt(0.5)
    return rval
