#!/usr/bin/env python

import sys
import numpy

def array_print(name, data):
	if type_scalar == numpy.float16:
		a = ['{:#04x}'.format(h.astype(int)) for h in data.view(dtype=numpy.uint16)]
	else:
		a = [repr(x) for x in data]
	print 'fftval_t {}[{}] = {{ {} }};'.format(name, len(a), ', '.join(a))

def array_permute(data, radix = 2):
	logradix = int(numpy.log2(radix))
	term_mask = int(radix - 1)
#	num_term = int(numpy.log2(len(data)) / logradix)
	for i in xrange(0, len(data)):
		# Obtain permuted address
		i_left = i
		permuted = 0
		cur_fft_size = radix
		while cur_fft_size <= len(data):
			permuted = (permuted << logradix) | (i_left & term_mask)
			i_left >>= logradix
			cur_fft_size <<= logradix
		# Permute only once and when addresses are different
		if i < permuted:
			tmp = data[i]
			data[i] = data[permuted]
			data[permuted] = tmp

fft_size = 1024
type_scalar = numpy.float16
permute = True

if len(sys.argv) > 1:
	type_scalar = {
		'16' : numpy.float16,
		'32' : numpy.float32,
		'64' : numpy.float64 }.get(sys.argv[1].strip(), None)
	if type_scalar == None:
		sys.exit('Invalid datatype')

if len(sys.argv) > 2:
	fft_size = sys.argv[2].strip()
	if not fft_size.isdigit():
		sys.exit('Invalid FFT size')
	fft_size = int(fft_size)

numpy.random.seed(seed=0)

print '#include "fft_const.h"'

# Generate input data
input_real = numpy.random.uniform(size=fft_size).astype(type_scalar)
input_imag = numpy.random.uniform(size=fft_size).astype(type_scalar)

# Compute reference FFT output
type_cmplx = numpy.complex128 if (type_scalar == numpy.float64) else numpy.complex64
input_cmplx = input_real.astype(type_cmplx)
input_cmplx.imag = input_imag
output_cmplx = numpy.fft.fft(input_cmplx)

if permute:
	array_permute(input_real)
	array_permute(input_imag)

array_print('input_data_real', input_real)
array_print('input_data_imag', input_imag)

output_real = output_cmplx.real.astype(type_scalar)
output_imag = output_cmplx.imag.astype(type_scalar)

array_print('output_data_real', output_real)
array_print('output_data_imag', output_imag)

# Generate twiddle factors (TFs)
# Negate sine since twiddle factor angles are generated *clockwise* from 0
rad = [(2.0 * numpy.pi * (float(i) / fft_size)) for i in xrange(fft_size)]
tf_real = numpy.cos(rad).astype(type_scalar)
tf_imag = -numpy.sin(rad).astype(type_scalar)

array_print('tf_real', tf_real)
array_print('tf_imag', tf_imag)