#!/usr/bin/env python

import math

import sys

import os

import random

import struct

import popen2

import getopt

import numpy

pi=math.pi

e=math.e

j=complex(0,1)

doreal=0

datatype = os.environ.get('DATATYPE','float')

util = '../tools/fft_' + datatype

minsnr=90

if datatype == 'double':

    fmt='d'

elif datatype=='int16_t':

    fmt='h'

    minsnr=10

elif datatype=='int32_t':

    fmt='i'

elif datatype=='simd':

    fmt='4f'

    sys.stderr.write('testkiss.py does not yet test simd')

    sys.exit(0)

elif datatype=='float':

    fmt='f'

else:

    sys.stderr.write('unrecognized datatype %s\n' % datatype)

    sys.exit(1)

def dopack(x,cpx=1):

    x = numpy.reshape( x, ( numpy.size(x),) )

    if cpx:

        s = ''.join( [ struct.pack(fmt*2,c.real,c.imag) for c in x ] )

    else:

        s = ''.join( [ struct.pack(fmt,c.real) for c in x ] )

    return s

def dounpack(x,cpx):

    uf = fmt * ( len(x) / struct.calcsize(fmt) )

    s = struct.unpack(uf,x)

    if cpx:

        return numpy.array(s[::2]) + numpy.array( s[1::2] )*j

    else:

        return numpy.array(s )

def make_random(dims=[1]):

    res = []

    for i in range(dims[0]):

        if len(dims)==1:

            r=random.uniform(-1,1)

            if doreal:

                res.append( r )

            else:

                i=random.uniform(-1,1)

                res.append( complex(r,i) )

        else:

            res.append( make_random( dims[1:] ) )

    return numpy.array(res)

def flatten(x):

    ntotal = numpy.size(x)

    return numpy.reshape(x,(ntotal,))

def randmat( ndims ):

    dims=[]

    for i in range( ndims ):

        curdim = int( random.uniform(2,5) )

        if doreal and i==(ndims-1):

            curdim = int(curdim/2)*2 # force even last dimension if real

        dims.append( curdim )

    return make_random(dims )

def test_fft(ndims):

    x=randmat( ndims )

    if doreal:

        xver = numpy.fft.rfftn(x)

    else:

        xver = numpy.fft.fftn(x)

    open('/tmp/fftexp.dat','w').write(dopack( flatten(xver) , True ) )

    x2=dofft(x,doreal)

    err = xver - x2

    errf = flatten(err)

    xverf = flatten(xver)

    errpow = numpy.vdot(errf,errf)+1e-10

    sigpow = numpy.vdot(xverf,xverf)+1e-10

    snr = 10*math.log10(abs(sigpow/errpow) )

    print 'SNR (compared to NumPy) : %.1fdB' % float(snr)

    if snr

        print 'xver=',xver

        print 'x2=',x2

        print 'err',err

        sys.exit(1)

def dofft(x,isreal):

    dims=list( numpy.shape(x) )

    x = flatten(x)

    scale=1

    if datatype=='int16_t':

        x = 32767 * x

        scale = len(x) / 32767.0

    elif datatype=='int32_t':

        x = 2147483647.0 * x

        scale = len(x) / 2147483647.0

    cmd='%s -n ' % util

    cmd += ','.join([str(d) for d in dims])

    if doreal:

        cmd += ' -R '

    print cmd

    p = popen2.Popen3(cmd )

    open('/tmp/fftin.dat','w').write(dopack( x , isreal==False ) )

    p.tochild.write( dopack( x , isreal==False ) )

    p.tochild.close()

    res = dounpack( p.fromchild.read() , 1 )

    open('/tmp/fftout.dat','w').write(dopack( flatten(res) , True ) )

    if doreal:

        dims[-1] = int( dims[-1]/2 ) + 1

    res = scale * res

    p.wait()

    return numpy.reshape(res,dims)

def main():

    opts,args = getopt.getopt(sys.argv[1:],'r')

    opts=dict(opts)

    global doreal

    doreal = opts.has_key('-r')

    if doreal:

        print 'Testing multi-dimensional real FFTs'

    else:

        print 'Testing multi-dimensional FFTs'

    for dim in range(1,4):

        test_fft( dim )

if __name__ == "__main__":

    main()