{-# LANGUAGE TypeOperators, PatternGuards #-}

-- | Computation of Fast Fourier Transforms using the Cooley-Tuckey algorithm.
module FFT 
	( fft
	, ifft
	, fft2d
	, fftWithRoots )
where
import Data.Array.Repa		as A
import Data.Ratio
import StrictComplex
import Roots

-- Vector Transform -------------------------------------------------------------------------------
-- | Compute the (fast) Discrete Fourier Transform of a vector.
fft	:: Shape sh
	=> Array (sh :. Int) Complex
	-> Array (sh :. Int) Complex

fft v
 = let	rofu	= calcRofu (extent v)
   in	force $ fftWithRoots rofu v


-- | Compute the (fast) Inverse Discrete Fourier Transform of a vector.
ifft	:: Shape sh
	=> Array (sh :. Int) Complex
	-> Array (sh :. Int) Complex

ifft v
 = let	_ :. len	= extent v
	scale		= fromIntegral len :*: 0
	rofu		= calcInverseRofu (extent v)
   in	force $ A.map (/ scale) $ fftWithRoots rofu v


-- Matrix Transform -------------------------------------------------------------------------------
-- | Compute the (fast) Discrete Fourier Transform of a square matrix.
fft2d 	:: Array DIM2 Complex
	-> Array DIM2 Complex

fft2d arr
 = let	rofu		= calcRofu (extent arr)
  	fftTrans	= transpose . fftWithRoots rofu
   in	fftTrans $ fftTrans arr


-- Cube Transform ---------------------------------------------------------------------------------
-- | Compute the (fast) Discrete Fourier Transform of a 3d array.
fft3d 	:: Array DIM3 Complex
	-> Array DIM3 Complex

fft3d arr
 = let	rofu		= calcRofu (extent arr)

	transpose3 arr
	 = traverse arr 
        	(\(Z :. k :. l :. m)   -> (Z :. l :. m :. k)) 
            	(\f (Z :. l :. m :. k) -> f (Z :. k :. l :. m)) 

	fftTrans	= transpose3 . fftWithRoots rofu
	
  in	fftTrans $ fftTrans $ fftTrans arr

	
-- Worker -----------------------------------------------------------------------------------------
-- | Generic function for computation of forward or inverse Discrete Fourier Transforms.
--	The length of the roots vector must be the same as the values vector.
--	The length of these vectors must be a power of two.
fftWithRoots	
	:: Shape sh
	=> Array (sh :. Int) Complex		-- ^ Roots of unity.
	-> Array (sh :. Int) Complex		-- ^ Input values.
        -> Array (sh :. Int) Complex

fftWithRoots rofu v
	| not $ (denominator $ toRational (logBase 2 $ fromIntegral vLen)) == 1
	= error $ "fft: vector length of " ++ show vLen ++ " is not a power of 2"
	
	| rLen /= vLen
	= error $ "fft: length of vector is not the length of the roots"
	
	| otherwise
	= fftWithRoots' rofu v

	where	_ :. rLen	= extent rofu
		_ :. vLen	= extent v


fftWithRoots'
	:: Shape sh
	=> Array (sh :. Int) Complex
	-> Array (sh :. Int) Complex
        -> Array (sh :. Int) Complex

{-# INLINE fftWithRoots' #-}
fftWithRoots' rofu v
 = case extent v of
	_ :. 2	-> fft_two   rofu v
	dim	-> fft_split rofu v dim

{-# INLINE fft_two #-}
fft_two rofu v
 = let	vFn' vFn (sh :. 0)  = vFn (sh :. 0) + vFn (sh :. 1)
	vFn' vFn (sh :. 1)  = vFn (sh :. 0) - vFn (sh :. 1)
   in	traverse v id vFn'
	
{-# INLINE fft_split #-}
fft_split rofu v vLen
 = let 	fft_lr = force $ fftWithRoots' (splitRofu rofu) (splitVector v)

	fft_l  = traverse2 fft_lr rofu 
 		   (\(sh :. 2 :. n) _ -> sh :. n)
		   (\f r (sh :. i)    -> f (sh :. 0 :. i) + r (sh :. i) * f (sh :. 1 :. i))

	fft_r  = traverse2 fft_lr rofu 
		   (\(sh :. 2 :. n) _ -> sh :. n)
		   (\f r (sh :. i)    -> f (sh :. 0 :. i) - r (sh :. i) * f (sh :. 1 :. i))

   in	fft_l +:+ fft_r

{-# INLINE splitRofu #-}
splitRofu rofu
 = traverse rofu
	(\(rSh :. rLen) 	-> rSh :. (2::Int) :. (rLen `div` 2))
	(\rFn (sh :. _ :. i) 	-> rFn (sh :. 2*i))

{-# INLINE splitVector #-}
splitVector v 
 = let	vFn' vFn (sh :. 0 :. i) = vFn (sh :. 2*i)
	vFn' vFn (sh :. 1 :. i) = vFn (sh :. 2*i+1)
   in	traverse v
		(\(vSh :. vLen)    -> vSh :. 2 :. (vLen `div` 2)) 
		vFn'