content: {-# LANGUAGE FlexibleContexts #-}
  
  module FFT
    where
  
  import qualified Prelude                                as P
  import Data.Array.Accelerate                            as A
  import Data.Array.Accelerate.IO.Codec.BMP               as A
  import Data.Array.Accelerate.Data.Colour.RGBA           as A
  import Data.Array.Accelerate.Data.Complex               as A
  import Data.Array.Accelerate.Math.FFT                   as A
  import Data.Array.Accelerate.Math.DFT.Centre            as A
  
  
  imageFFT :: Int -> Acc (Array DIM2 RGBA32) -> Acc (Array DIM2 RGBA32, Array DIM2 RGBA32)
  imageFFT cutoff img = lift (arrMag, arrPhase)
    where
      -- Load in the image luminance
      arrComplex :: Acc (Array DIM2 (Complex Float))
      arrComplex  = A.map (\r -> lift (r :+ constant 0))
                  $ A.map (luminance . unpackRGBA) img
  
      -- Apply the centering transform so that the output has the zero frequency
      -- in the middle of the image
      arrCentered = centre2D arrComplex
  
      -- Do the transform
      arrFreq     = fft2D Forward arrCentered
  
      -- Clip the magnitude of the transformed array
      clipMag     = the (unit (constant (P.fromIntegral cutoff)))
      clip x      = x > clipMag ? ( 1 , x / clipMag )
      arrMag      = A.map (packRGBA . grey . clip . magnitude) arrFreq
  
      -- Get the phase of the transformed array
      scale x     = (phase x + pi) / (2 * pi)
      arrPhase    = A.map (packRGBA . grey . scale) arrFreq
  
      grey x      = rgba x x x 1
  
  