{-# LANGUAGE FlexibleContexts #-}

module Main (main) where

import           Criterion.Main
import qualified Data.Array.Accelerate                 as A
import           Data.Array.Accelerate.KullbackLiebler
import qualified Data.Array.Accelerate.LLVM.Native     as CPU
#ifdef ACCELERATE_GPU
import qualified Data.Array.Accelerate.LLVM.PTX        as GPU
#endif
import           Data.Array.Accelerate.Test

main :: IO ()
main = do
    ps <- distEnv
    qs <- distEnv
#ifdef ACCELERATE_GPU
    let res = kullbackLieblerGPU ps qs -- TODO: better way to force this
#else
    let res = kullbackLieblerCPU ps qs -- TODO: better way to force this
#endif
    putStrLn (res `seq` "Setup complete")
    defaultMain [
#ifdef ACCELERATE_GPU
                  bgroup "GPU"
                      [ bench "K-L (1e7)" $ nf (kullbackLieblerGPU ps) qs
                      , bench "entropy (1e7)" $ nf entropyGPU ps
                      , bench "K-L (1e7) (scaled)" $ nf (kullbackLieblerGPUScaled ps) qs
                      , bench "entropy (1e7) (scaled)" $ nf entropyGPUScaled ps
                      ]
                , bgroup "CPU"
#else
                  bgroup "CPU"
#endif
                      [ bench "K-L (1e7)" $ nf (kullbackLieblerCPU ps) qs
                      , bench "K-L (1e7) (scaled)" $ nf (kullbackLieblerCPUScaled ps) qs
                      , bench "entropy (1e7)" $ nf entropyCPU ps
                      , bench "entropy (1e7) (scaled)" $ nf entropyCPUScaled ps
                      ]
                ]

    where distEnv :: IO (A.Acc (A.Vector Double))
          distEnv = randDist (A.Z A.:. 10000000)
#ifdef ACCELERATE_GPU
          kullbackLieblerGPU ps qs = the $ GPU.run $ kullbackLiebler ps qs
          kullbackLieblerGPUScaled ps qs = the $ GPU.run $ kullbackLiebler (scale ps) (scale qs)
          entropyGPU = the . GPU.run . entropy
          entropyGPUScaled = the . GPU.run . entropy . scale
#endif
          -- TODO: runQ
          kullbackLieblerCPU ps qs = the $ CPU.run $ kullbackLiebler ps qs
          kullbackLieblerCPUScaled ps qs = the $ CPU.run $ kullbackLiebler (scale ps) (scale qs)
          entropyCPU = the . CPU.run . entropy
          entropyCPUScaled = the . CPU.run . entropy . scale