module Examples.Math.Fft where

import qualified Prelude as P
import Feldspar
import Feldspar.Wrap
import Feldspar.Vector
import Feldspar.Matrix
import Feldspar.Compiler

-- | Wrapper to define the size of vectors in 'fft'
fftInstance :: Data [Complex Float] -> Data [Complex Float]
fftInstance = desugar . fft . thawVector' 256

-- | Wrapper to define the size of vectors in 'fft'
fft_wrapped ::Data' D256 [Complex Float] ->  Data [Complex Float]
fft_wrapped = wrap fft

-- | Wrapper to define the size of vectors in 'ifft'
ifftInstance :: Data [Complex Float] -> Data [Complex Float]
ifftInstance = desugar . ifft . thawVector' 256

-- | Wrapper to define the size of vectors in 'ifft'
ifft_wrapped ::Data' D256 [Complex Float] ->  Data [Complex Float]
ifft_wrapped = wrap ifft


-- =================== INTERFACE ==================================
-- | Radix-2 Decimation-In-Frequeny 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.)
fft :: Vector1 (Complex Float) -> Vector1 (Complex Float)
fft v = bitRev (loglen-1) $ fftCore (loglen-1) v
    where loglen = f2i $ logBase 2 $ i2f $  length v

-- | Radix-2 Decimation-In-Frequeny 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.)
ifft :: Vector1 (Complex Float) -> Vector1 (Complex Float)
ifft v = bitRev (loglen-1) $ ifftCore (loglen-1) v
    where loglen = f2i $ logBase 2 $ i2f $ length v
-- ================================================================




-- | fftCore function uses 2^(n+1) input vector.
--   Output from the last stage needs to be bit reversed using bitRev function (if required)
fftCore ::  Data Index ->  Vector1 (Complex Float) -> Vector1 (Complex Float)
fftCore n v = composeOn stage (reverse (0...n)) v

stage k (Indexed l ixf Empty) = (Indexed l ixf' Empty)
  where
    k2 = 1 .<<. k
    ixf' i = condition (testBit i k)   (twid * (b-a))   (a+b)
      where
        a = ixf i
        b = ixf (i `xor` k2)
        twid = cis (-pi*(i2f (lsbs k i)) / i2f k2)


-- | ifftCore function uses 2^(n+1) input vector.
--   Output from the last stage needs to be bit reversed using bitRev function (if required)
ifftCore ::  Data Index ->  Vector1 (Complex Float) -> Vector1 (Complex Float)
ifftCore n v = map (/ (complex (i2f (2^(n+1))) 0)) $ composeOn istage (reverse (0...n)) v

istage k (Indexed l ixf Empty) = (Indexed l ixf' Empty)
  where
    k2 = 1 .<<. k
    ixf' i = condition (testBit i k)   (twid * (b-a))   (a+b)
      where
        a = ixf i
        b = ixf (i `xor` k2)
        twid = cis (pi*(i2f (lsbs k i)) / i2f k2)


-- | bitRev function transforms the given vector to bitreversal order
--   parameter n is the size of the input vector
bitRev :: Type a => Data Index -> Vector (Data a) -> Vector (Data a)
bitRev n = pipe riffle (1...n)


-- | Helper functions for fftCore and ifftCore and bitRev
pipe :: (Syntax a) => (Data Index -> a -> a) -> Vector (Data Index) -> a -> a
pipe = flip.fold.flip

composeOn f is as = fold (flip f) as is

allOnes = complement 0

oneBits n = complement (allOnes .<<. n)

lsbs k i = i .&. oneBits k

par m n f = mat2Vec m n . map f . vec2Mat m n

-- k at least 1
rotBit :: Data Index -> Data Index -> Data Index
rotBit 0 _ = P.error "k should be at least 1"
rotBit k i = lefts .|. rights
  where
    ir = i .>>. 1
    rights = ir .&. (oneBits k)
    lefts  = (((ir .>>. k) .<<. 1) .|. (i .&. 1)) .<<. k

riffle k (Indexed l ixf Empty) = indexed l (ixf.rotBit k)

vec2Mat :: Data Index -> Data Index -> Vector (Data a) -> Matrix a
vec2Mat m n (Indexed l ixf Empty) = indexedMat (1 .<<. m) (1 .<<. n) ixf'
  where
    ixf' i j = ixf $ (i .<<. n) `xor` j

mat2Vec :: Type a => Data Index -> Data Index -> Matrix a -> Vector (Data a)
mat2Vec m n matr = Indexed (1 .<<. m .<<. n) ixf Empty
  where
    ixf i = matr ! y ! x
      where
        y = i .>>. n
        x = i .&. (oneBits n)



-- Ad-hoc function to generate a power-of-two sequence in order to test fft and ifft
pow2Seq :: Data WordN -> Vector1 (Complex Float)
pow2Seq n = indexed (2 ^ n) (\i ->complex (i2f i) 0 )