{-# LANGUAGE FlexibleContexts #-}

module Gather where

import Random

import System.Random.MWC
import Data.Array.Unboxed
import Data.Array.Accelerate as Acc hiding ((!))
import Prelude               as P


-- Tests
-- -----

gatherAcc :: Vector Int -> Vector Float -> Acc (Vector Float)
gatherAcc mapV inputV = Acc.gather (use mapV) (use inputV)

gatherIfAcc :: Vector Int -> Vector Int -> Vector Float -> Vector Float -> Acc (Vector Float)
gatherIfAcc mapV maskV defaultV inputV
 = Acc.gatherIf (use mapV) (use maskV) evenAcc (use defaultV) (use inputV)

evenAcc :: Exp Int -> Exp Bool
evenAcc v = (v `mod` 2) ==* 0


gatherRef :: UArray Int Int -> UArray Int Float -> UArray Int Float
gatherRef mapV inputV = amap (\ix -> inputV ! ix) mapV

gatherIfRef :: UArray Int Int -> UArray Int Int -> UArray Int Float -> UArray Int Float -> UArray Int Float
gatherIfRef mapV maskV defaultV inputV
  = listArray (bounds mapV)
  $ P.map (\(mIx, mV, dV) -> if evenRef mV then (inputV ! mIx) else dV)
  $ P.zip3 mapL maskL defaultL
  where
    mapL     = elems mapV
    maskL    = elems maskV
    defaultL = elems defaultV

evenRef :: Int -> Bool
evenRef = even


-- Main
-- ----
run :: String -> Int -> IO (() -> UArray Int Float, () -> Acc (Vector Float))
run alg n = withSystemRandom $ \gen -> do
  vec       <- randomUArrayR (-1, 1) gen n
  vec'      <- convertUArray vec

  mapV      <- randomUArrayR (0, n - 1) gen n
  mapV'     <- convertUArray mapV

  maskV     <- randomUArrayR (0, n) gen n
  maskV'    <- convertUArray maskV

  defaultV  <- randomUArrayR (-1, 1) gen n
  defaultV' <- convertUArray defaultV

  --
  let go f g = return (run_ref f vec, run_acc g vec')

  case alg of
    "gather"    -> go (gatherRef mapV) (gatherAcc mapV')
    "gather-if" -> go (gatherIfRef mapV maskV defaultV) (gatherIfAcc mapV' maskV' defaultV')
    x           -> error $ "unknown variant: " ++ x

  where
    {-# NOINLINE run_ref #-}
    run_ref f xs () = f xs
    run_acc f xs () = f xs