```-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Transform.Fourier.PFA
--
-- Stability   :  experimental
-- Portability :  portable
--
-- Prime Factor Algorithm
--
-----------------------------------------------------------------------------

module Numeric.Transform.Fourier.PFA (fft_pfa) where

import Data.List
import Data.Array
import Data.Complex

{-# specialize fft_pfa :: Array Int (Complex Float) -> Int -> Int -> (Array Int (Complex Float) -> Array Int (Complex Float)) -> Array Int (Complex Float) #-}
{-# specialize fft_pfa :: Array Int (Complex Double) -> Int -> Int -> (Array Int (Complex Double) -> Array Int (Complex Double)) -> Array Int (Complex Double) #-}

-- | Prime Factor Algorithm doing row FFT's then column FFT's

fft_pfa :: (Ix a, Integral a, RealFloat b) => Array a (Complex b) -- ^ x[n]
-> a -- ^ nrows
-> a -- ^ ncols
-> (Array a (Complex b) -> Array a (Complex b)) -- ^ FFT function
-> Array a (Complex b) -- ^ X[k]

fft_pfa a l m fft = array (0,n-1) \$ zip ks (elems x')
where x = listArray ((0,0),(l-1,m-1)) [ a!i | i <- xs ]
f = listArray ((0,0),(l-1,m-1)) (flatten_rows \$ map fft \$ rows x)
x' = listArray ((0,0),(l-1,m-1)) (flatten_cols \$ map fft \$ cols f)
(xs,ks) = pfa_index_map l m
n = l * m

{-# specialize pfa_index_map :: Int -> Int -> ([Int],[Int]) #-}

pfa_index_map :: (Integral a) => a -> a -> ([a],[a])
pfa_index_map l m = (ns,ks)
where ns = [ (m * n1 + l * n2) `mod` n | n1 <- [0..(l-1)], n2 <- [0..(m-1)] ]
ks = [ (c * m * k1 + d * l * k2) `mod` n | k1 <- [0..(l-1)], k2 <- [0..(m-1)] ]
c = find_inverse m l
d = find_inverse l m
n = l * m

{-# specialize find_inverse :: Int -> Int -> Int #-}

find_inverse :: (Integral a) => a -> a -> a
find_inverse a0 n0 = find_inverse' a0 n0 1
where find_inverse' a n a' | (a*a') `mod` n == 1 = a'
| otherwise = find_inverse' a n (a'+1)

{-# specialize rows :: Array (Int,Int) (Complex Float) -> [Array Int (Complex Float)] #-}
{-# specialize rows :: Array (Int,Int) (Complex Double) -> [Array Int (Complex Double)] #-}

rows :: (Ix a, Integral a, RealFloat b) => Array (a,a) (Complex b) -> [Array a (Complex b)]
rows x = [ listArray (0,m) [ x!(i,j) | j <- [0..m] ] | i <- [0..l] ]
where ((_,_),(l,m)) = bounds x

{-# specialize cols :: Array (Int,Int) (Complex Float) -> [Array Int (Complex Float)] #-}
{-# specialize cols :: Array (Int,Int) (Complex Double) -> [Array Int (Complex Double)] #-}

cols :: (Ix a, Integral a, RealFloat b) => Array (a,a) (Complex b) -> [Array a (Complex b)]
cols x = [ listArray (0,l) [ x!(i,j) | i <- [0..l] ] | j <- [0..m] ]
where ((_,_),(l,m)) = bounds x

{-# specialize flatten_rows :: [Array Int (Complex Float)] -> [(Complex Float)] #-}
{-# specialize flatten_rows :: [Array Int (Complex Double)] -> [(Complex Double)] #-}

flatten_rows :: (Ix a, Integral a, RealFloat b) => [Array a (Complex b)] -> [(Complex b)]
flatten_rows a = foldr (++) [] \$ map elems a

{-# specialize flatten_cols :: [Array Int (Complex Float)] -> [(Complex Float)] #-}
{-# specialize flatten_cols :: [Array Int (Complex Double)] -> [(Complex Double)] #-}

flatten_cols :: (Ix a, Integral a, RealFloat b) => [Array a (Complex b)] -> [(Complex b)]
flatten_cols a = foldr (++) [] \$ transpose \$ map elems a
```