#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Computes the spectrogram of a test signal using cupy and cuFFT.

Author: Jan Schlüter
"""

import sys
import os
import timeit

import numpy as np
import cupy as cp

INPUT_ON_GPU = True
OUTPUT_ON_GPU = True

from testfile import make_test_signal


def spectrogram(signal, sample_rate=22050, frame_len=1024, fps=70):
    """
    Computes a magnitude spectrogram at a given sample rate (in Hz), frame
    length (in samples) and frame rate (in Hz), on CUDA using cupy.
    """
    if not INPUT_ON_GPU:
        signal = cp.array(signal.astype(np.float32))  # already blown up to a list of frames
    win = cp.hanning(frame_len).astype(cp.float32)
    # apply window function
    #signal *= win  # this doesn't work correctly for some reason.
    signal = signal * win
    # perform FFT
    spect = cp.fft.rfft(signal)
    # convert into magnitude spectrogram
    spect = cp.abs(spect)
    # return
    if OUTPUT_ON_GPU:
        cp.cuda.get_current_stream().synchronize()
    else:
        return spect.get()


def main():
    # load input
    global x, spectrogram
    x = make_test_signal()
    # we do the following here because cupy cannot do stride tricks
    # the actual copying work is included in the benchmark unless INPUT_ON_GPU
    hop_size = 22050 // 70
    frame_len = 1024
    frames = len(x) - frame_len + 1
    x = np.lib.stride_tricks.as_strided(
            x, (frames, frame_len), (x.strides[0], x.strides[0]))[::hop_size]
    if INPUT_ON_GPU:
        x = cp.array(x.astype(np.float32))

    # benchmark
    times = timeit.repeat(
            setup='from __main__ import x, spectrogram',
            stmt='spectrogram(x)',
            repeat=5, number=32)
    print("Took %.3fs." % (min(times) / 32))

    # save result
    #assert not OUTPUT_ON_GPU
    #np.save(sys.argv[0][:-2] + 'npy', spectrogram(x))


if __name__=="__main__":
    main()
[((845, 864), 'cupy.fft.rfft', 'cp.fft.rfft', (['signal'], {}), '(signal)\n', (856, 864), True, 'import cupy as cp\n'), ((918, 931), 'cupy.abs', 'cp.abs', (['spect'], {}), '(spect)\n', (924, 931), True, 'import cupy as cp\n'), ((1120, 1138), 'testfile.make_test_signal', 'make_test_signal', ([], {}), '()\n', (1136, 1138), False, 'from testfile import make_test_signal\n'), ((1582, 1689), 'timeit.repeat', 'timeit.repeat', ([], {'setup': '"""from __main__ import x, spectrogram"""', 'stmt': '"""spectrogram(x)"""', 'repeat': '(5)', 'number': '(32)'}), "(setup='from __main__ import x, spectrogram', stmt=\n    'spectrogram(x)', repeat=5, number=32)\n", (1595, 1689), False, 'import timeit\n'), ((1378, 1468), 'numpy.lib.stride_tricks.as_strided', 'np.lib.stride_tricks.as_strided', (['x', '(frames, frame_len)', '(x.strides[0], x.strides[0])'], {}), '(x, (frames, frame_len), (x.strides[0], x.\n    strides[0]))\n', (1409, 1468), True, 'import numpy as np\n'), ((653, 674), 'cupy.hanning', 'cp.hanning', (['frame_len'], {}), '(frame_len)\n', (663, 674), True, 'import cupy as cp\n'), ((975, 1003), 'cupy.cuda.get_current_stream', 'cp.cuda.get_current_stream', ([], {}), '()\n', (1001, 1003), True, 'import cupy as cp\n')]