#!/usr/pkg/bin/python
"""package of variants on fourier transform, including power spectrum

   see descriptions of individual functions"""

Copyright = """
    power - variants on fourier transform
    Copyright (C) 2004  John Comeau <jc@jcomeau.com>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
    """
errormessage = "Not all needed libraries found, upgrade or check path"
try:
 True # not defined in older Python releases
except:
 True, False = 1, 0
try:
 import sys, os, types, re
except exception:
 try:
  sys.stderr.write("%s\n" % errormessage)
 except:
  print errormessage
 raise exception

# get name this program was called as
self = sys.argv[0].split(os.sep)[-1]
# now get name we gave it when we wrote it
originalself = Copyright.split()[0]

# globals and routines that should be in every program
# (yes, you could import them, but there are problems in that approach too)

debugging = True # set to False for production release

def DebugPrint(*whatever):
 global debugging
 if debugging:
  sys.stderr.write("%s\n" % repr(whatever))

def test(*args):
 "check if lowercase functions actually work as programs"
 print "this is a test. it is only a test. passed arguments follow..."
 DebugPrint(args)

# other globals, specific to this program

from Numeric import *
import FFT, types

power = "this string has to be here, see template.py for reason why"

def getarray(args):
 """return a numeric array and size...

    whether it was passed AS an array, or another sequence type,
    or even as a list of strings like .12, .22, .65, etc.
    if the first element is an array, then the second element will
    be taken as the size parameter to be passed to fft()"""
 size = len(args) # just in case args is just a list, as from commandline
 while type(args[0]) is types.ListType or \
  type(args[0]) is types.TupleType or \
  type(args[0]) is ArrayType:
  #DebugPrint('data is already in sequence format')
  size = len(args[0])
  if len(args) > 1:
   size = int(args[1])
  args = args[0]
 data = list(args)
 #DebugPrint('getarray(): ', 'data', data, 'size', size)
 data = map(float, data) # just in case it's strings
 return array(data), size # return as Numeric array (q.v.)

def Fft(*args):
 "for command-line testing"
 return fft(args)

def fft(*args):
 "pydoc FFT.fft, this is just an interface into that"
 #DebugPrint('calling fft()')
 data, size = getarray(args)
 #DebugPrint('data', data, 'size', size)
 return FFT.fft(data, size)

def PowerSpectrum(*args):
 "for command-line testing"
 return powerspectrum(args)

def powerspectrum(*args):
 """Power spectrum code from University of Cape Town website

 http://www.phy.uct.ac.za/courses/phy400w/cp/pyfrags.htm"""
 data, size = getarray(args)
 transform = fft(data, size)
 spectrum = abs(transform * conjugate(transform)) / len(transform) ** 2
 return spectrum

def main():
 """main routine, only used for command-line testing

    let's say this program is called 'template.py'
    don't, by the way, name the program with the name
     of an included function if you're using this system...
    now if the program is invoked as './template.py', it will simply
     output a usage message.
    if invoked as './template.py MyFunction 1 2 3'
     we will attempt to call MyFunction(1, 2, 3) and
     show whatever is found
    if invoked as './template.py myfunction 1 2 3'
     we attempt to call myfunction(1, 2, 3) as a program and
     let it output whatever it decides
    likewise if you symlink like so: 'ln -s template.py myfunction'
     then invoke as './myfunction 1 2 3'"""
 program = sys.argv[0].split(os.sep)[-1]
 function = program.split('.')[0] # in case of .py extension
 try:
  if eval("type(%s)" % function) == types.FunctionType:
   if re.compile('^[a-z0-9]+$').match(function) != None:
    eval("%s(sys.argv[1:])" % function)
   else:
    try:
     print repr(eval("%s(sys.argv[1:])" % function))
    except:
     sys.stderr.write("Error occurred testing routine %s(%s)\n" % (
      function, repr(sys.argv[1:])))
     raise
  else:
   raise NameError, '%s not a function' % function
 except NameError, instance:
  try:
   message = instance.args[0]
  except:
   message = 'unknown NameError'
  if message != '%s not a function' % function:
   raise NameError, message
  # called name didn't match anything defined
  if len(sys.argv) > 1:
   # hopefully the next arg is the function name we want
   sys.argv.pop(0)
   main()
  else:
   # called name didn't match any routine, so output usage message
   sys.stderr.write("""Usage: %s [OPTIONS] INPUT[...]
    or for more detailed help, read source or:
    pydoc %s
    """ % (self, self.split('.')[0]))

if __name__ == '__main__':
 # if this program was imported by another, the above test will fail,
 # and this following code won't be used...
 main()