#################################
## complex fourier transformation
## simple discrete dft and idft
## fast fft and ifft
## transform:  x --> X
##
## Wolfgang.Urban@schule.at, 7/2004
#################################


############################################################
# discrete fourier transform
############################################################

def dft(x, sign=-1):
    from cmath import pi, exp
    N = len(x)
    omega = [0]*N
    indices = range(N)
    phi = sign*2j*pi/N
    for i in indices: omega[i] = exp(phi*i)
    
    X = [0]*N
    for n in indices:
        sum = 0
        for k in indices:
            sum = sum+omega[(n*k)%N]*x[k]
        X[n] = sum
    return X


def idft(X):
    N = len(X)
    x = dft(X,1)
    for i in range(N):
        x[i] = x[i]/float(N)
    return x

############################################################
# fast fourier transform
# number of points has to be power of 2
############################################################

def nextpower2(i):
    n = 2
    while n<i: n = n*2
    return n

def bitreverse(x):
    N = len(x)
    if N != nextpower2(N): raise ValueError, 'N is not power of 2'

    x = x[:]
    for i in range(N):
        k, b, a = 0, N>>1, 1
        while b >= a:
            if b&i: k = k|a
            if a&i: k = k|b
            b, a = b>>1, a<<1
        if i < k:               # swap each pair only once, avoid i==k
            x[i], x[k] = x[k], x[i]
    return x
   

def fft(x, sign=-1):
    from cmath import pi, exp
    N = len(x)
    omega = [0]*N
    for i in range(N):
        omega[i] = exp(sign*2j*pi*i/N)
    x = bitreverse(x)
    m = 2
    while m <= N:
        for s in range(0, N, m):
            for i in range(m/2):
                n = i * N/m
                a, b = s + i, s + i + m/2
                x[a], x[b] = x[a]+omega[n%N]*x[b], x[a]-omega[n%N]*x[b]
        m = m * 2
    return x

def ifft(X):
    N = len(X)
    x = fft(X, sign=1)   
    for i in range(N):
        x[i] = x[i]/float(N)
    return x


############################################################
# some functions for testing purposes
############################################################

def rect(r, w, deg=0):          # degrees or radians
    from math import cos, sin, pi
    if deg:
        w = pi*w/180.0
    return r*cos(w), r*sin(w)

def polar(x, y, deg=0):
    from math import hypot, atan2, pi
    if deg: return hypot(x,y), 180.0*atan2(y,x)/pi
    else: return hypot(x,y), atan2(y,x)

############################################################

# calculate function values for test routine

def f_saw_rise(n):
    dt = 1.0/n
    f = [0]*n
    for t in range(n): f[t] = complex(2.0*t*dt-1.0)
    return f
    
def f_saw_fall(n):
    dt = 1.0/n
    f = [0]*n
    for t in range(n): f[t] = complex(1.0-2.0*t*dt)
    return f

############################################################

# show complex function values nicely
def show(a):
    for n in range(len(a)):
        print "%3d: %6.3f %6.3fi" %(n,a[n].real,a[n].imag)

def test():
    print "---f-------------------------------"
    f = f_saw_fall(16)
    show(f)
    print "---fft-----------------------------"
    show(fft(f))
    print "---ifft(fft)-----------------------"
    show(ifft(fft(f)))
    
test()