{-# LINE 1 "Numeric/FFTW.hsc" #-}
{-| Bindings to FFTW.

Example usage:

> import Foreign.Marshal.Array
> import Data.Complex
> import Foreign.Storable.Complex
> import FFTW
> 
> main = do
>     inA  <- fftwAllocComplex 1024
>     outA <- fftwAllocComplex 1024
> 
>     plan <- planDFT1d 1024 inA outA Forward fftwEstimate
> 
>     pokeArray inA $ map (:+ 0) [0..1023]
>     execute plan
>     res <- peekArray 1024 outA
> 
>     fftwFree inA
>     fftwFree outA
> 
>     print res
-}

module Numeric.FFTW (
    -- * Memory allocation functions
    fftwMalloc,
    fftwFree,
    fftwFreePtr,
    fftwAllocReal,
    fftwAllocComplex,

    -- * FFT planning flags
    Direction(..),
    Flag(),
    -- ** Planning rigor flags
    fftwMeasure,
    fftwExhaustive,
    fftwPatient,
    fftwEstimate,
    fftwWisdomOnly,
    -- ** Algorithm restriction flags
    fftwDestroyInput,
    fftwUnaligned,
    fftwPreserveInput,

    -- * FFT planning
    FFTWPlan,
    planDFT1d,
    planDFTR2C1d,

    -- * FFT execution
    execute,
    executeDFT,
    executeDFTR2C
    ) where

import Foreign.C.Types
import Foreign.Ptr
import Data.Word
import Data.Complex
import Control.Monad
import Data.Bits
import Data.Monoid
import Data.Semigroup as Sem



foreign import ccall unsafe "fftw_malloc"
    c_fftwMalloc :: CUInt -> IO (Ptr a)

-- | Like malloc, but ensures that the pointer obeys the alignment restrictions of FFTW (e.g. for SIMD acceleration). You probably want to use 'fftwAllocReal' or 'fftwAllocComplex' instead.
fftwMalloc :: Word32 -- ^ size
           -> IO (Ptr a)
fftwMalloc = c_fftwMalloc . fromIntegral

-- | Free a pointer returned by 'fftwMalloc', 'fftwAllocReal', or 'fftwAllocComplex'
foreign import ccall unsafe "fftw_free"
    fftwFree :: Ptr a -- ^ the pointer to be freed
             -> IO ()

-- | A function pointer to @fftwFree@.
foreign import ccall unsafe "&fftw_free"
    fftwFreePtr :: FunPtr (Ptr a -> IO ())

foreign import ccall unsafe "fftw_alloc_real"
    c_fftwAllocReal :: CUInt -> IO (Ptr CDouble)

-- | Allocates an array of Doubles. It ensures that the pointer obeys the alignment restrictions of FFTW (e.g. for SIMD acceleration).
fftwAllocReal :: Word32 -- ^ size
              -> IO (Ptr CDouble)
fftwAllocReal = c_fftwAllocReal . fromIntegral

foreign import ccall unsafe "fftw_alloc_complex"
    c_fftwAllocComplex :: CUInt -> IO (Ptr (Complex CDouble))

-- | Allocates an array of complex Doubles (i.e. the c type "double complex"). It ensures that the pointer obeys the alignment restrictions of FFTW (e.g. for SIMD acceleration).
fftwAllocComplex :: Word32 -- ^ size
                 -> IO (Ptr (Complex CDouble))
fftwAllocComplex = c_fftwAllocComplex . fromIntegral

-- | The direction of the transform: Forward for a normal transform, Backward for an inverse transform
data Direction = Forward
               | Backward

dirToInt :: Direction -> CInt
dirToInt Forward  = -1
{-# LINE 109 "Numeric/FFTW.hsc" #-}
dirToInt Backward = 1
{-# LINE 110 "Numeric/FFTW.hsc" #-}

-- | FFTW planner flags. These flags affect the planning process. They can be combined using the 'Monoid' instance. See the FFTW flag documentation: <http://www.fftw.org/doc/Planner-Flags.html>.
newtype Flag = Flag {unFlag :: CUInt}

instance Sem.Semigroup Flag where
  (Flag x) <> (Flag y) = Flag (x .|. y)

instance Monoid Flag where
    mempty   = Flag 0

{-# LINE 122 "Numeric/FFTW.hsc" #-}

fftwMeasure, fftwExhaustive, fftwPatient, fftwEstimate, fftwWisdomOnly :: Flag
fftwEstimate       = Flag 64
{-# LINE 125 "Numeric/FFTW.hsc" #-}
fftwMeasure        = Flag 0
{-# LINE 126 "Numeric/FFTW.hsc" #-}
fftwPatient        = Flag 32
{-# LINE 127 "Numeric/FFTW.hsc" #-}
fftwExhaustive     = Flag 8
{-# LINE 128 "Numeric/FFTW.hsc" #-}
fftwWisdomOnly     = Flag 2097152
{-# LINE 129 "Numeric/FFTW.hsc" #-}

fftwDestroyInput, fftwUnaligned, fftwPreserveInput :: Flag
fftwDestroyInput   = Flag 1
{-# LINE 132 "Numeric/FFTW.hsc" #-}
fftwUnaligned      = Flag 2
{-# LINE 133 "Numeric/FFTW.hsc" #-}
fftwPreserveInput  = Flag 16
{-# LINE 134 "Numeric/FFTW.hsc" #-}

data CFFTWPlan

-- | A @FFTWPlan i o@ contains all of the information necessary to perform a transform from an input array of type @i@ to an output array of type @o@, including pointers to the input and output arrays.
newtype FFTWPlan i o = FFTWPlan (Ptr CFFTWPlan)

foreign import ccall unsafe "fftw_plan_dft_1d"
    c_planDFT1d :: CInt -> Ptr (Complex CDouble) -> Ptr (Complex CDouble) -> CInt -> CUInt -> IO (Ptr CFFTWPlan)

--This appears to be missing from the fft package on Hackage
-- | Create a plan for a 1 dimensional complex to complex DFT. The plan stores pointers to the input and output arrays, and these will be used if you 'execute' the plan in the future. They are required even if you intend to specify different input and output arrays in the future (i.e. using 'executeDFT')
planDFT1d :: Int                   -- ^ size
          -> Ptr (Complex CDouble) -- ^ input pointer
          -> Ptr (Complex CDouble) -- ^ output pointer
          -> Direction             -- ^ direction
          -> Flag                  -- ^ planner flags
          -> IO (FFTWPlan (Complex CDouble) (Complex CDouble))
planDFT1d n inp out sign flags = liftM FFTWPlan $ c_planDFT1d (fromIntegral n) inp out (dirToInt sign) (unFlag flags)

foreign import ccall unsafe "fftw_plan_dft_r2c_1d"
    c_planDFTR2C1d :: CInt -> Ptr CDouble -> Ptr (Complex CDouble) -> CUInt -> IO (Ptr CFFTWPlan)

--This appears to be missing from the fft package on Hackage
-- | Create a plan for a 1 dimensional real to complex DFT. The plan stores pointers to the input and output arrays, and these will be used if you 'execute' the plan in the future. They are required even if you intend to specify different input and output arrays in the future (i.e. using 'executeDFTR2C')
planDFTR2C1d :: Int                   -- ^ size
             -> Ptr CDouble           -- ^ input pointer
             -> Ptr (Complex CDouble) -- ^ output pointer
             -> Flag                  -- ^ planner flags
             -> IO (FFTWPlan CDouble (Complex CDouble))
planDFTR2C1d n inp out flags = liftM FFTWPlan $ c_planDFTR2C1d (fromIntegral n) inp out (unFlag flags)

foreign import ccall unsafe "fftw_execute"
    c_execute :: Ptr CFFTWPlan -> IO ()

-- | Execute a plan. Performs an FFT. The input and output arrays are stored within the plan so do not need to be given.
execute :: FFTWPlan i o -- ^ the plan to execute
        -> IO ()
execute (FFTWPlan p) = c_execute p

foreign import ccall unsafe "fftw_execute_dft"
    c_executeDFT :: Ptr CFFTWPlan -> Ptr (Complex CDouble) -> Ptr (Complex CDouble) -> IO ()

-- | Execute a complex to complex DFT but on different input and output arrays to those specified when the plan was created.
executeDFT :: FFTWPlan (Complex CDouble) (Complex CDouble) -- ^ the plan to execute
           -> Ptr (Complex CDouble)                        -- ^ input pointer
           -> Ptr (Complex CDouble)                        -- ^ output pointer
           -> IO ()
executeDFT (FFTWPlan p) inp out = c_executeDFT p inp out

foreign import ccall unsafe "fftw_execute_dft_r2c"
    c_executeDFTR2C :: Ptr CFFTWPlan -> Ptr CDouble -> Ptr (Complex CDouble) -> IO ()

-- | Execute a real to complex DFT but on different input and output arrays to those specified when the plan was created.
executeDFTR2C :: FFTWPlan CDouble (Complex CDouble) -- ^ the plan to execute
              -> Ptr CDouble                        -- ^ input pointer
              -> Ptr (Complex CDouble)              -- ^ output pointer
              -> IO ()
executeDFTR2C (FFTWPlan p) inp out = c_executeDFTR2C p inp out