{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedLists  #-}

module Main ( main ) where

import           Control.Composition                   ((.$))
import qualified Data.Array.Accelerate                 as A
import           Data.Array.Accelerate.KullbackLiebler (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            (randDist, the)
import           Test.Tasty
import           Test.Tasty.HUnit

main :: IO ()
main = defaultMain accTests

accTests :: TestTree
accTests = testGroup "Unit tests"
    [
#ifdef ACCELERATE_GPU
      testCase "Setup (GPU)" $
        assertBool "simple unit test" =<< selfGPU
    , testCase "Setup (CPU)" $
#else
      testCase "Setup (CPU)" $
#endif
        assertBool "simple unit test" =<< selfCPU
    ]

--TODO: random gen/normalization?
commonData :: IO (A.Acc (A.Scalar Double))
commonData = do
    datums <- randDist (A.Z A.:. 1000000)
    pure $ kullbackLiebler .$ datums

#ifdef ACCELERATE_GPU
selfGPU :: IO Bool
selfGPU = do
    datums <- commonData
    let res = the $ GPU.run datums
    pure $ res == 0
#endif

selfCPU :: IO Bool
selfCPU = do
    datums <- commonData
    let res = the $ CPU.run datums
    pure $ res == 0