from __future__ import division
import numpy
import scipy
import scipy.signal


def print_hi(name):
    # Use a breakpoint in the code line below to debug your script.
    print(f'Hi, {name}')  # Press Ctrl+F8 to toggle the breakpoint.


"""
Module to calculate the fast 2D fractional fourier transform.
"""


def frft2d(matrix, a=[0.7, 0.7]):
    """
    Calculate the 2D fast fractional fourier transform.
    Parameters
    ----------
    matrix : 2D numpy array
        The signal to be transformed.
    a : list / array
        fractional powers
    Returns
    -------
    data : numpy array
        The transformed 2D signal.
    """
    temp = numpy.zeros((matrix.shape[0], matrix.shape[1]), dtype=numpy.complex)
    for k in range(0, matrix.shape[0]):
        temp[k, :] = frft(matrix[k, :], a[0])
    out = numpy.zeros((temp.shape[0], temp.shape[1]), dtype=numpy.complex)
    for m in range(0, temp.shape[1]):
        out[:, m] = frft(temp[:, m], a[1])

    return out


def frft(f, a):
    """
    Calculate the fast fractional fourier transform.
    Parameters
    ----------
    f : numpy array
        The signal to be transformed.
    a : float
        fractional power
    Returns
    -------
    data : numpy array
        The transformed signal.
    References
    ---------
     .. [1] This algorithm implements `frft.m` from
        https://nalag.cs.kuleuven.be/research/software/FRFT/
    """
    ret = numpy.zeros_like(f, dtype=numpy.complex)
    f = f.copy().astype(numpy.complex)
    N = len(f)
    shft = numpy.fmod(numpy.arange(N) + numpy.fix(N / 2), N).astype(int)
    sN = numpy.sqrt(N)
    a = numpy.remainder(a, 4.0)

    # Special cases
    if a == 0.0:
        return f
    if a == 2.0:
        return numpy.flipud(f)
    if a == 1.0:
        ret[shft] = numpy.fft.fft(f[shft]) / sN
        return ret
    if a == 3.0:
        ret[shft] = numpy.fft.ifft(f[shft]) * sN
        return ret

    # reduce to interval 0.5 < a < 1.5
    if a > 2.0:
        a = a - 2.0
        f = numpy.flipud(f)
    if a > 1.5:
        a = a - 1
        f[shft] = numpy.fft.fft(f[shft]) / sN
    if a < 0.5:
        a = a + 1
        f[shft] = numpy.fft.ifft(f[shft]) * sN

    # the general case for 0.5 < a < 1.5
    alpha = a * numpy.pi / 2
    tana2 = numpy.tan(alpha / 2)
    sina = numpy.sin(alpha)
    f = numpy.hstack((numpy.zeros(N - 1), sincinterp(f), numpy.zeros(N - 1))).T

    # chirp premultiplication
    chrp = numpy.exp(-1j * numpy.pi / N * tana2 / 4 *
                     numpy.arange(-2 * N + 2, 2 * N - 1).T ** 2)
    f = chrp * f

    # chirp convolution
    c = numpy.pi / N / sina / 4
    tmp = numpy.exp(1j * c * numpy.arange(-(4 * N - 4), 4 * N - 3).T ** 2)
    ret = scipy.signal.fftconvolve(
        numpy.exp(1j * c * numpy.arange(-(4 * N - 4), 4 * N - 3).T ** 2),
        f
    )
    ret = ret[4 * N - 4:8 * N - 7] * numpy.sqrt(c / numpy.pi)

    # chirp post multiplication
    ret = chrp * ret

    # normalizing constant
    ret = numpy.exp(-1j * (1 - a) * numpy.pi / 4) * ret[N - 1:-N + 1:2]

    return ret


def ifrft(f, a):
    """
    Calculate the inverse fast fractional fourier transform.
    Parameters
    ----------
    f : numpy array
        The signal to be transformed.
    a : float
        fractional power
    Returns
    -------
    data : numpy array
        The transformed signal.
    """
    return frft(f, -a)


def sincinterp(x):
    N = len(x)
    y = numpy.zeros(2 * N - 1, dtype=x.dtype)
    y[:2 * N:2] = x
    xint = scipy.signal.fftconvolve(
        y[:2 * N],
        numpy.sinc(numpy.arange(-(2 * N - 3), (2 * N - 2)).T / 2),
    )

    return xint[2 * N - 3: -2 * N + 3]


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    print_hi('PyCharm')
    f = numpy.array([10.8392, -39.9709, -44.5170, -28.4400, -4.7966, 6.5892, 26.0688, 25.6156])
    a = 0.6
    result = frft(f, a)
    # print(type(result))
    # result = frft2d(f)
    for ans in result:
        print(ans)
