import Data.Complex
import Data.Array
import Data.Bits
import System.Random
import Data.Random.Normal
import Criterion.Main
import Control.Parallel
import Control.Monad.Par  
import Control.Parallel.Strategies

-- file given.hs for use with Lab 1 Part 1 of the Chalmers PFP Course
-- Please write your names in the file if submitting it


-- generating input for FFT or DFT. Borrowed from Simon Marlow I believe.
mX, mY, sdX, sdY :: Float
mX = 0
mY = 0
sdX = 0.5
sdY = 1.5    

generate2DSamplesList :: Int           -- number of samples to generate
                  -> Float -> Float    -- X and Y mean
                  -> Float -> Float    -- X and Y standard deviations
                  -> IO [Complex Float]
generate2DSamplesList n mx my sdx sdy = do
  gen <- getStdGen
  let (genx, geny) = split gen
      xsamples = normals' (mx,sdx) genx
      ysamples = normals' (my,sdy) geny
  return $ zipWith (:+) (take n xsamples) ysamples




-- Task 1
divConq :: (prob -> Bool)              -- is the problem indivisible?
            -> (prob -> [prob])        -- split
            -> ([sol] -> sol)          -- join
            -> (prob -> sol)           -- solve a sub-problem
            -> (prob -> sol)

divConq indiv split join f prob = undefined



-- Task 2
-- simple recursive parallel prefix or scan (the Sklansky algorithm)

halve as = splitAt n' as
  where
    n' = div (length as + 1) 2

-- for input [a1,a2,a3, ...], it returns [a1, a1 `op` a2, a1 `op` a2 `op` a3, ...]
dpscan :: (a -> a -> a) -> [a] -> [a]
dpscan _ []  = []
dpscan _ [a] = [a]
dpscan op as = ls ++ [op (last ls) b | b <- rs]
  where (ls0,rs0) = halve as
        (ls,rs)   = (dpscan op ls0, dpscan op rs0)

{--
*Main> dpscan (+) [1..20]
[1,3,6,10,15,21,28,36,45,55,66,78,91,105,120,136,153,171,190,210]
--}


-- suggested type for pscan made using divConq. The first param. is a threshold
-- revert to scanl1 from the prelude for inputs of length less than this
pscan :: Int -> (a -> a -> a) -> [a] -> [a]
pscan t op = undefined  -- divConq indiv split join f   where ...
                        

-- Task 4

-- twiddle factors
tw :: Int -> Int -> Complex Float
tw n k = cis (-2 * pi * fromIntegral k / fromIntegral n)

dft :: [Complex Float] -> [Complex Float]
dft xs = [ sum [ xs!!j * tw n (j*k) | j <- [0..n']] | k <- [0..n']]
  where
    n = length xs
    n' = n-1


-- the outputs of this FFT should be have the bit-reversal permutation applied
--  to get them into the same order as in the dft spec
-- see the bitRev function that is provided below for completeness
-- Feel free to skip it when benchmarking your FFT
-- In case you are wondering, this is the Decimation in Frequency (DIF) 
-- radix 2 Cooley-Tukey FFT

fft :: [Complex Float] -> [Complex Float]
fft [a] = [a]
fft as = ls ++ rs
  where
    (cs,ds) = bflyS as
    ls = fft cs
    rs = fft ds

bflyS :: [Complex Float] -> ([Complex Float], [Complex Float])
bflyS as = (los,rts)
  where
    (ls,rs) = halve as
    los = zipWith (+) ls rs
    ros = zipWith (-) ls rs
    rts = zipWith (*) ros [tw (length as) i | i <- [0..(length ros) - 1]]





-- bit reversal (for completeness, borrowed from Feldspar)

oneBitsN :: Int -> Int
oneBitsN  k = complement (shiftL (complement 0) k)

bitr :: Int -> Int -> Int
bitr n a = let mask = (oneBitsN n) in
    (complement mask .&. a) .|. rotateL (reverseBits (mask .&. a)) n

bitRevA :: Int  -> Array Int a -> Array Int a
bitRevA n xs = ixmap (bounds xs) (bitr n) xs

-- bit reverse sublists of length 2^n. Assume input list has length 2^(n+k) for k >= 0
bitRev :: Int -> [a] -> [a]
bitRev n xs = elems . bitRevA n . listArray (0, length xs - 1) $ xs

reverseBits :: Bits b => b -> b
reverseBits b = revLoop b 0 (0 `asTypeOf` b)
  where
    bSz = bitSize b
    revLoop b i n | i >= bSz    = n
    revLoop b i n | testBit b i = revLoop b (i+1) (setBit n (bSz - i - 1))
    revLoop b i n | otherwise   = revLoop b (i+1) n