{-# LANGUAGE FlexibleContexts #-}

-- | FFT implementation inspired by the paper "Feldspar: Application and
-- implementation":
--
-- <http://publications.lib.chalmers.se/records/fulltext/local_156271.pdf>
--
-- There are a few differences, partly due to the paper using a different
-- Feldspar implementation. But regardless, the best way to understand the
-- definitions in this file is by reading the paper.

module FFT
  ( tw  -- Exported to allow pre-computation
  , fftCore
  , fft
  , ifft
  ) where



import Prelude ()

import Data.Bool (bool)

import Feldspar.Run
import Feldspar.Data.Vector
import Feldspar.Data.Buffered



----------------------------------------
-- * Helper functions
----------------------------------------

rotBit :: Data Index -> Data Index -> Data Index
rotBit k i = lefts .|. rights
  where
    k'     = i2n k
    ir     = i .>>. 1
    rights = ir .&. oneBits k'
    lefts  = (((ir .>>. k') .<<. 1) .|. (i .&. 1)) .<<. k'

riffle :: (Pully pull a, Syntax a) => Data Index -> pull -> Pull a
riffle = backPermute . const . rotBit

testBit :: (Bits a, Integral a, PrimType a) => Data a -> Data Index -> Data Bool
testBit a i = i2b (a .&. (1 .<<. i2n i))

-- | @2^n@
twoTo :: (Num a, Bits a, PrimType a) => Data Index -> Data a
twoTo n = 1 .<<. i2n n

flipBit :: (Num a, Bits a, PrimType a) => Data a -> Data Index -> Data a
flipBit i k = i `xor` twoTo k



----------------------------------------
-- * Bit-reversal permutation
----------------------------------------

bitRev :: (Manifestable Run vec a, Finite vec, Syntax a)
    => Store a
    -> Length  -- ^ Unrolling steps in inner loops (1 means no unrolling)
    -> Data Length
    -> vec
    -> Run (Manifest a)
bitRev st u n = loopStore st (1,1,Excl n) $ \i -> return . unroll u . riffle i



----------------------------------------
-- * FFT
----------------------------------------

tw :: (Floating a, PrimType a, PrimType (Complex a))
    => Bool  -- ^ Inverse FFT?
    -> Data Index
    -> Data Index
    -> Data (Complex a)
tw inv n k = polar 1 (bool (-2) 2 inv * π * i2n k / i2n n)

twids
    :: ( Pully ts (Data (Complex a))
       , RealFloat a
       , PrimType a
       , PrimType (Complex a)
       , Pully vec (Data (Complex a))
       )
    => ts
    -> Data Index
    -> Data Index
    -> Data Length
    -> vec
    -> DPull (Complex a)
twids ts n k l vec = Pull l $ \i ->
    let j = (lsbs (i2n k) i) .<<. (n'-1-k')
    in  (testBit i k) ? ((ts!j) * (vec!i)) $ (vec!i)
  where
    n' = i2n n
    k' = i2n k

bfly
    :: ( Pully vec (Data (Complex a))
       , RealFloat a
       , PrimType a
       , PrimType (Complex a)
       )
    => Data Index -> vec -> DPull (Complex a)
bfly k as = Pull (length as) $ \i ->
    let a = as ! i
        b = as ! flipBit i k
    in  (testBit i k) ? (b-a) $ (a+b)

-- | Core of the FFT
--
-- It is normally better to use 'fft' or 'ifft' than this functon; however, for
-- doing repeated FFT on vectors of the same size, 'fftCore' can be used to
-- avoid recomputing the twiddle factors and the number of stages.
fftCore
    :: ( Pully ts (Data (Complex a))
       , Manifestable Run vec (Data (Complex a))
       , Finite vec
       , RealFloat a
       , PrimType a
       , PrimType (Complex a)
       )
    => Store (Data (Complex a))
    -> Length       -- ^ Unrolling steps in inner loops (1 means no unrolling)
    -> ts           -- ^ Twiddle factors
    -> Data Length  -- ^ Number of stages
    -> vec
    -> Run (DManifest (Complex a))
fftCore st u ts n vec = do
    let step i = return . unroll u . twids ts n i (length vec) . bfly i
    loopStore st ((i2n n :: Data Int32)-1,-1,Incl 0) (step . i2n) vec
      >>= bitRev st u n
      -- `i2n` is used to make the loop index a signed number. Otherwise the
      -- index will wrap to maxBound before the loop test after the final
      -- iteration.
      --
      -- An alternative is to use:
      --
      --     loopStore st (n,-1,Excl 0) (step . subtract 1) vec

-- | Radix-2 Decimation-In-Frequency Fast Fourier Transformation of the given
-- complex vector. The given vector must be power-of-two sized, (for example 2,
-- 4, 8, 16, 32, etc.) The output is non-normalized.
--
-- The length of the vector must be divisible by the number of unrolling steps.
--
-- The optimal amount of unrolling depends on the target architecture, but a
-- value of 2 might be a reasonable default that gives some performance
-- improvements on many systems and doesn't lead to too much code size increase.
fft
    :: ( Manifestable Run vec (Data (Complex a))
       , Finite vec
       , RealFloat a
       , PrimType a
       , PrimType (Complex a)
       )
    => Store (Data (Complex a))
    -> Length  -- ^ Unrolling steps in inner loops (1 means no unrolling)
    -> vec
    -> Run (DManifest (Complex a))
fft st u vec = do
    n  <- shareM (ilog2 (length vec))
    ts <- manifestFresh $ Pull (twoTo (n-1)) (tw False (twoTo n))
      -- Change `manifestFresh` to `return` to avoid pre-computing twiddle
      -- factors
    fftCore st u ts n vec

-- | Radix-2 Decimation-In-Frequency Inverse Fast Fourier Transformation of the
-- given complex vector. The given vector must be power-of-two sized, (for
-- example 2, 4, 8, 16, 32, etc.) The output is divided with the input size,
-- thus giving @`ifft` . `fft` == id@.
--
-- The length of the vector must be divisible by the number of unrolling steps.
--
-- The optimal amount of unrolling depends on the target architecture, but a
-- value of 2 might be a reasonable default that gives some performance
-- improvements on many systems and doesn't lead to too much code size increase.
ifft
    :: ( Manifestable Run vec (Data (Complex a))
       , Finite vec
       , RealFloat a
       , PrimType a
       , PrimType (Complex a)
       )
    => Store (Data (Complex a))
    -> Length  -- ^ Unrolling steps in inner loops (1 means no unrolling)
    -> vec
    -> Run (DPull (Complex a))
ifft st u vec = do
    n  <- shareM (ilog2 (length vec))
    ts <- manifestFresh $ Pull (twoTo (n-1)) (tw True (twoTo n))
      -- Change `manifestFresh` to `return` to avoid pre-computing twiddle
      -- factors
    normalize <$> fftCore st u ts n vec
  where
    normalize = map (/ (i2n $ length vec))