{- 6.827 Problem Set 3, Problem 2.

Replace these lines
with the names of your
team members

Please do not remove this initial comment or alter the first or last line. -}

module Main(main) where

import Array
import PHArray
import Complex

{- 

Your main job in this problem will be to incrementally replace array
indexing and the array data types with indexing operations and types
of your own.  

-}

type ComplexVector = Array Int (Complex Float)

twoPi :: Float
twoPi = 2.0 * pi

rotate90 :: RealFloat a => Complex a -> Complex a
rotate90 (a:+b) = b:+(-a)

-- pick elements out of given vector for recursive ffts.
-- Normalizes bounds to 1.
shuffle :: ComplexVector -> (ComplexVector, ComplexVector)
shuffle v =
  let bnds@(l,u) = bounds v
      size = sizeRange bnds
      halfsize = size `div` 2
  in (array (1,halfsize) [ (i, v!(i*2+l-2)) | i <- [1..halfsize]],
      array (1,halfsize) [ (i, v!(i*2+l-1)) | i <- [1..halfsize]])

fft :: ComplexVector -> ComplexVector -> ComplexVector
fft v roU =
  let bnds = bounds v
      size = sizeRange bnds
  in  if (size == 4) then
        let [i1,i2,i3,i4] = range bnds
            l1 = v!i1 + v!i3
            l2 = v!i1 - v!i3
            r1 = v!i2 + v!i4
            r2 = v!i2 - v!i4
            r2' = rotate90 r2
        in array bnds 
             [(i1, l1 + r1),
              (i2, l2 + r2'),
              (i3, l1 - r1),
              (i4, l2 - r2')]
      else
        let (left_v, right_v) = shuffle v
            fft_left          = fft left_v  roU
            fft_right         = fft right_v roU
        in combine fft_left fft_right roU

-- combine recursively fft'd arrays.  Assumes their bounds came from "shuffle".
-- Note how the final comprehension generates two array elements at a time.
combine :: ComplexVector -> ComplexVector -> ComplexVector -> ComplexVector
combine u v roU =
  let bnds@(1,m) = bounds u
      (1,n)      = bounds roU
      index      = n `div` m
  in array (1, 2*m) [ element | i <- [1..m],
                                vprod <- [roU!((i-1)*index + 1) * v!i],
                                element <- [(i,   u!i + vprod), 
                                            (m+i, u!i - vprod)]]

-- computeRoots computes the nth roots of 1 and forms the array roU used above.
computeRoots :: Int -> ComplexVector
computeRoots n =
  let theta = -twoPi / (fromInt n)
      halfn = n `div` 2
      wprWpi = (-2.0) * (sin (0.5 * theta) ^ 2) :+ sin theta
      roU = array (1,halfn) ([(1, 1.0:+0.0)]
                             ++ [(i, wprWpi * roU!(i-1) + roU!(i-1)) | i <- [2..halfn]])
  in roU



-- Some possible inputs to FFT
test1FFT n =
  let c = array (1,n) [(i, fromInt i * 1.0 :+ 0.0) | i <- [1..n]]
  in fft c (computeRoots n)

constFFT n =
  let c = array (1,n) [(i,10.0) | i <- [1..n]]
  in fft c (computeRoots n)

rotFFT n m =
  let c = array (1,n) [(i, cos theta :+ sin theta) |
			i <- [1..n],
			theta <- [fromInt (i-1) * twoPi / fromInt m]]
  in fft c (computeRoots n)

-- throw away all but four decimal places in answer so that it's readable.
nicely a = array (bounds a) 
		[(i, a:+b) | (i,x:+y) <- assocs a, 
			     a <- [fourDec x],
			     b <- [fourDec y]]

fourDec x = (fromIntegral $ round $ x*10000.0) / 10000.0


main = (print . nicely) (test1FFT 64) >>
       (print . nicely) (constFFT 64) >>
       (print . nicely) (rotFFT 64 32)