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

module Main ( main ) where

import           Control.Composition                   ((.$))
import qualified Data.Array.Accelerate                 as A
import           Data.Array.Accelerate.KullbackLiebler (hellinger, kullbackLiebler, scale)
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
    ]

commonData :: IO (A.Acc (A.Scalar Double))
commonData = do
    datums <- randDist (A.Z A.:. 1000000)
    pure $ kullbackLiebler .$ scale datums

hellingerData :: IO (A.Acc (A.Scalar Double))
hellingerData = do
    datums <- randDist (A.Z A.:. 1000000)
    pure $ hellinger .$ scale datums

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

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