semi_improved_fft.hs

semi_improved_fft.hs

data Complex = C Double Double deriving Show
data Nat     = E | O Nat | I Nat deriving Show
data Tree a  = L a | B (Tree a) (Tree a) deriving Show

cScale :: Double -> Complex -> Complex
cScale s (C ar ai) = C (ar * s) (ai * s)

cAdd :: Complex -> Complex -> Complex
cAdd (C ar ai) (C br bi) = C (ar + br) (ai + bi)

cSub :: Complex -> Complex -> Complex
cSub (C ar ai) (C br bi) = C (ar - br) (ai - bi)

cMul :: Complex -> Complex -> Complex
cMul (C ar ai) (C br bi) = C (ar * br - ai * bi) (ar * bi + ai * br)

cPol :: Double -> Complex
cPol ang = C (cos ang) (sin ang)

natVal :: Nat -> Double -> Double
natVal E     x = x
natVal (O n) x = natVal n (2 * x)
natVal (I n) x = natVal n (2 * x + 1)

natLen :: Nat -> Int
natLen E     = 0
natLen (O n) = 1 + natLen n
natLen (I n) = 1 + natLen n

fft :: Nat -> Tree Complex -> Tree Complex
fft ang (B eve odd) =
  let pt0 = fft (O ang) eve
      pt1 = fft (O ang) odd
  in mix ang pt0 pt1
fft ang (L x) = L x

mix :: Nat -> Tree Complex -> Tree Complex -> Tree Complex
mix ang (B ax ay) (B bx by) =
  let ptl = mix (O ang) ax bx
      ptr = mix (I ang) ay by
  in B ptl ptr
mix ang (L pt0) (L pt1) =
  let pt2 = cMul pt1 (cPol (2 * pi * natVal ang 0 / 2 ** fromIntegral (natLen ang + 1)))
      ptl = L (cAdd pt0 pt2)
      ptr = L (cSub pt0 pt2)
  in B ptl ptr

main :: IO ()
main = do
  let input = B (B (B (L (C 0 0)) (L (C 4 0))) (B (L (C 2 0)) (L (C 6 0)))) (B (B (L (C 1 0)) (L (C 5 0))) (B (L (C 3 0)) (L (C 7 0))))
  print $ fft (O E) input