{-# OPTIONS_GHC -O2 #-}

{- |
  Module      :  Numeric.FFT
  Copyright   :  (c) Matt Morrow 2008
  License     :  BSD3
  Maintainer  :  Matt Morrow <mjm2002@gmail.com>
  Stability   :  experimental
  Portability :  portable


  A radix-2 DIT version of
  the Cooley-Tukey FFT algorithm.
-}

module Numeric.FFT (
    fft, ifft
  , dft, idft
) where

import Data.List(foldl')
import Data.Complex(Complex(..))

-----------------------------------------------------------------------------

-- | /O(n lg n)/. A radix-2 DIT
--  (decimation-in-time) version of the
--  Cooley-Tukey FFT algorithm.
--  The length of the input list /must/
--  be a power of two, or only the prefix
--  of the list of length equal to the largest
--  power of two less than the length of the list
--  will be transformed.
fft :: [Complex Double] -> [Complex Double]
fft []    = []
fft [x,y] = dft [x,y]
fft xs    = fmap (go len ys zs) [0..len*2-1]
  where (ys,zs) = (fft***fft)
            . deinterleave $ xs
        len = length ys
        go len xs ys k
          | k <  len  = (xs!!k) + (ys!!k) * f k (len*2)
          | otherwise = let k' = k - len
              in (xs!!k') - (ys!!k') * f k' (len*2)
          where i = 0 :+ 1
                fi = fromIntegral
                f k n = exp (negate(2*pi*i*fi k)/fi n)
        (***) f g (x,y) = (f x, g y)
        deinterleave :: [a] -> ([a],[a])
        deinterleave = unzip . pairs
        pairs :: [a] -> [(a,a)]
        pairs []       = []
        pairs (x:y:zs) = (x,y) : pairs zs
        pairs _        = []


ifft :: [Complex Double] -> [Complex Double]
ifft xs = let n = (fromIntegral . length) xs in fmap (/n) (fft xs)

-----------------------------------------------------------------------------

-- | /O(n^2)/. The Discrete Fourier Transform.
dft :: [Complex Double] -> [Complex Double]
dft xs = let len = length xs
          in zipWith (go len) [0..len-1] (repeat xs)
  where i = 0 :+ 1
        fi = fromIntegral
        go len k xs = foldl' (+) 0 . flip fmap [0..len-1]
          $ \n -> (xs!!n) * exp (negate(2*pi*i*fi n*fi k)/fi len)

idft :: [Complex Double] -> [Complex Double]
idft xs = let n = (fromIntegral . length) xs in fmap (/n) (dft xs)

-----------------------------------------------------------------------------

{-
"A radix-2 decimation-in-time (DIT) FFT is
the simplest and most common form of the
Cooley-Tukey algorithm, although highly
optimized Cooley-Tukey implementations
typically use other forms of the algorithm"

"Radix-2 DIT divides a DFT of size N into
two interleaved DFTs (hence the name "radix-2")
of size N/2 with each recursive stage."

<http://en.wikipedia.org/wiki/Cooley-Tukey_FFT_algorithm>
-}

-----------------------------------------------------------------------------

{-
[m@ganon SPHODRA]$ time ./test_fft
(-32767.999999213338) :+ (-6.835652750528328e8)

real    0m0.757s
user    0m0.729s
sys     0m0.019s


[m@ganon SPHODRA]$ time ./test_dft
^C

real    0m21.221s
user    0m20.938s
sys     0m0.017s

import Math.FFT(fft)
main = print . last . fft . fmap fromIntegral $ [0..2^16-1]

import Math.FFT(dft)
main = print . last . dft . fmap fromIntegral $ [0..2^16-1]
-}

-----------------------------------------------------------------------------