{-# LANGUAGE FlexibleInstances #-}
module FFT(fftRadix2) where

import Data.Complex
import Data.List

-- TOOLS
trunc x = if (x < 0.00001) then 0 else x
ctrunc (x :+ y) = (trunc x) :+ (trunc y)
  
ones :: [Complex Double]
ones = (1 :+ 0) : ones 


-- The FFT in a nutshell:

-- 1. init: create hierarchy of size 1 DFTs == decimated signal
-- 2. update: recursively combine size n to size n*2.


-- Intermediate result of a FFT operation, from fully decimated signal
-- tree to single DFT spectrum list.
data Decim a = Decim Int (Decim a) (Decim a)
             | DFT [a]   -- list of DFTs
             deriving (Show, Eq)

-- 1. Decimate signal.
--
-- Annotate each subdivision node with level number l.  The number 2^l
-- is the decimation factor for the intermediate DFT result at that
-- level.  It is used for computing the "phase shift" modulation for
-- the odd spectrum.

fftDecim levels = dec 0 where  
  dec l a = if (l >= levels)
            then DFT $ [head a] 
            else Decim l (d $ pair fst a) (d $ pair snd a) where
              d = dec (l + 1)
              pair p (x0:x1:xs) = p (x0,x1) : pair p xs
              pair _ [] = []

-- 2. Recursive combination.
--
--               
fftRecurse nbLevels = fft where
  
  -- Phase shift the Odd spectrum for subsapling factor 2^level, which
  -- aligns Odd and Even spectra so they can be combined.
  shift = shiftSpectrum (wPowers nbLevels)
  
  -- Recursive combination of decimated tree.
  fft (DFT dft) = DFT dft
  fft (Decim l sigEven sigOdd ) = DFT dft where
    -- Sub-FFT on decimated signals + shift Odd spectrum
    DFT dftE = fft sigEven
    DFT dftO = shift l $ fft sigOdd
    -- Compute butterfly.
    dft = zipWith (+) dftE dftO ++
          zipWith (-) dftE dftO
  

-- Perform Radix 2 fft by taking the first 2^level elements from s.
fftRadix2 nbLevels sig = sig' where
  -- Step 1: decimate the signal.
  decimSig = fftDecim nbLevels sig
  
  -- Step 2: recursively compute FFTs
  DFT sig' = fftRecurse nbLevels decimSig
  
                          
-- Phase shift spectrum by multiplying with [1,w^{2^l},...]
shiftSpectrum wPows level (DFT fs) = DFT fs' where
  fs' = zipWith (*) fs wsl
  wsl = map (wPows !!) [0,2^level..]


-- Powers of w
wPowers levels = map w [0..(2^levels)-1] where
  w n = (cos t :+ sin t) where
    t = 2 * n * pi / (2^levels)