#! /usr/bin/python3

from math import sin, pi
import numpy as np
from numpy.testing import assert_allclose, assert_almost_equal
import unittest

class TestExpeyes(unittest.TestCase):
    
    def test_fft(self):
        import expeyes.eyemath as EM
        n = 4    # how many integer thousand samples
        k = 1000 # one thousand
        f = 20   # integer frequency in Hz
        ya = [sin(2*pi*t/k*f) for t in range(n*k)]
        si = 1 # sampling interval, in millisecond

        FFT = EM.fft(ya, si)

        # FFT[0] : abscissa to plot the Fourier transform, near np.arange(0, k/2, 1/n)
        assert_allclose(np.arange(0,k/2, 1/n), FFT[0])

        # FFT[1] should be the fast Fourier transform , with a peak at frequency f
        for i, amplitude in enumerate(FFT[1]):
            if i == n*f: # we are at the peak frequency
                assert_almost_equal(amplitude, 0.5)
            else:# we are not at the peak frequency
                assert_almost_equal(amplitude, 0.0)
        return

if __name__ == '__main__':
    unittest.main()    
