module Main where

import qualified Data.Array.Accelerate.CUFFT.Batched as Batched
import qualified Data.Array.Accelerate.CUFFT.Single as Single
import qualified Data.Array.Accelerate.LLVM.PTX as CUDA
import qualified Data.Array.Accelerate as A
import Data.Array.Accelerate (Z(Z), (:.)((:.)))


mainSingle :: IO ()
mainSingle = do
   target <- Batched.getBestTarget
   let dim = Z:.7
   hf <- Single.plan1D target Single.forwardReal dim
   let spec =
          CUDA.run1With target (Single.transform hf) $
          A.fromList dim $ 0 : 1 : repeat (0 :: Float)
   print spec

   hb <- Single.plan1D target Single.inverseReal dim
   print $ CUDA.run1With target (Single.transform hb) spec


mainBatched :: IO ()
mainBatched = do
   let count, width :: Int
       count = 3; width = 7
       dim :: A.DIM2
       dim = Z:.count:.width

   target <- Batched.getBestTarget
   hf <- Batched.plan1D target Batched.forwardReal dim
   let spec =
          CUDA.run1With target (Batched.transform hf) $
          A.fromList dim $ concat $
          take count $ map (take width) $
          iterate (0:) $ 1 : repeat (0 :: Float)
   print spec

   hb <- Batched.plan1D target Batched.inverseReal dim
   print $ CUDA.run1With target (Batched.transform hb) spec


main :: IO ()
main = mainSingle >> mainBatched