{-# OPTIONS -ffi #-}
------------------------------------------------------------------------
-- |
-- Module      : FFTW3
-- Copyright   : (c) Aaron Denney 2004
-- License     : GPL v. 2
-- 
-- Maintainer  : wnoise-haskell@ofb.net
-- Stability   : experimental
-- Portability : GHC.  (-#include extension).
--
-- Provides a fairly simple interface to the FFTW (version 3), an FFT
-- library whose "codelets" are generated and combined by an O'Caml
-- optimizer, so is fairly efficient for a large range of composite
-- numbers, not just powers of two.
-- 
-- Only exposes the "basic" interface, and not all of that yet.
------------------------------------------------------------------------
#include "fftw3.h"

module FFTW3 (fft, ifft) where
import Complex
import CComplex
import Foreign
import CForeign

data PlanStruct = PS
type Plan = Ptr PlanStruct
type Flags = CInt
type Sign = CInt
type FFTWComplex = CComplex CDouble

fftw_measure, fftw_destroy_input, fftw_unaligned, fftw_conserve_memory :: Flags
fftw_exhaustive, fftw_preserve_input, fftw_patient, fftw_estimate :: Flags

fftw_measure         = #const FFTW_MEASURE 
fftw_destroy_input   = #const FFTW_DESTROY_INPUT
fftw_unaligned       = #const FFTW_UNALIGNED
fftw_conserve_memory = #const FFTW_CONSERVE_MEMORY
fftw_exhaustive      = #const FFTW_EXHAUSTIVE
fftw_preserve_input  = #const FFTW_PRESERVE_INPUT
fftw_patient         = #const FFTW_PATIENT
fftw_estimate        = #const FFTW_ESTIMATE

-- preserving plans safel is a pain, so we don't expose them.  This
-- means only the fftw_estimate method for construction will be useful.
-- Since we've just marshalled the data in, we shouldn't care if it gets
-- destroyed.
use_flags :: CInt
use_flags = fftw_estimate .|. fftw_destroy_input

foreign import ccall safe fftw_execute :: Plan -> IO ()
foreign import ccall unsafe fftw_destroy_plan :: Plan -> IO ()
foreign import ccall unsafe fftw_cleanup :: IO ()

-- Basic interface, complex double -> complex double
foreign import ccall unsafe fftw_plan_dft_1d :: CInt -> Ptr FFTWComplex -> Ptr FFTWComplex -> Sign -> Flags -> IO Plan;


fft  :: RealFloat a => [Complex a] -> [Complex a]
fft = map c2c . unsafePerformIO . fft_dir (-1) . map c2c
ifft :: RealFloat a => [Complex a] -> [Complex a]
ifft = map c2c . unsafePerformIO . fft_dir (1) . map c2c

-- Because CComplex is a type synonym, we can get away with
-- only messing about with the double inside.
c2c :: (RealFloat a, RealFloat b) => Complex a -> Complex b
c2c (a :+ b) = (realToFrac a :+ realToFrac b)

fft_wrap l sign out inarr = fftw_plan_dft_1d l inarr out sign use_flags

fft_dir dir x = do let l = length x
                   out <- mallocArray l
                   plan <- withArray x (fft_wrap (fromIntegral l) dir out)
                   fftw_execute plan
                   fftw_destroy_plan plan
                   ret <- peekArray l out
                   free out
                   return ret