{-# LANGUAGE FlexibleContexts, TypeOperators #-}

module Fold where

import Random

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


-- one-dimension ah-ha-ha
-- ----------------------

toUA :: (IArray UArray a, IArray UArray b) => ([a] -> b) -> UArray Int a -> UArray () b
toUA f = listArray ((),()) . return . f . elems


sumAcc, prodAcc, maxAcc, minAcc :: Shape ix => Array (ix:.Int) Float -> Acc (Array ix Float)
sumAcc  = Acc.fold (+) 0 . Acc.use
prodAcc = Acc.fold (*) 1 . Acc.use
maxAcc  = Acc.fold1 Acc.max . Acc.use
minAcc  = Acc.fold1 Acc.min . Acc.use

sumRef, prodRef, maxRef, minRef :: UArray Int Float -> UArray () Float
sumRef  = toUA sum
prodRef = toUA product
maxRef  = toUA maximum
minRef  = toUA minimum


-- two-dimensions ah-ha-ha
-- -----------------------

foldU2D :: IArray UArray a => (a -> a -> a) -> a -> UArray (Int,Int) a -> UArray Int a
foldU2D f z arr =
  let (_,(m,_)) = bounds arr
  in  accumArray f z (0,m) [ (i,e) | ((i,_),e) <- assocs arr ]

sum2DRef, prod2DRef :: UArray (Int,Int) Float -> UArray Int Float
sum2DRef  = foldU2D (+) 0
prod2DRef = foldU2D (*) 1


-- Main
-- ----
run :: String -> Int -> IO (() -> UArray () Float, () -> Acc (Scalar Float))
run alg n = withSystemRandom $ \gen -> do
  vec  <- randomUArrayR (-1,1) gen n
  vec' <- convertUArray vec
  --
  let go  f g = return (run_ref f vec, run_acc g vec')
  case alg of
    "sum"        -> go sumRef sumAcc
    "product"    -> go prodRef prodAcc
    "maximum"    -> go maxRef maxAcc
    "minimum"    -> go minRef minAcc
    x            -> error $ "unknown variant: " ++ x
  where
    {-# NOINLINE run_ref #-}
    run_ref f xs () = f xs
    run_acc f xs () = f xs

run2d :: String -> Int -> IO (() -> UArray Int Float, () -> Acc (Vector Float))
run2d alg n = withSystemRandom $ \gen -> do
  let u = P.floor . sqrt $ (P.fromIntegral n :: Double)
      v = 2*u+1 :: Int
  mat  <- listArray ((0,0), (u-1,v-1)) `fmap` replicateM (u*v) (uniformR (-1,1) gen)
  mat' <- let m = fromIArray mat :: Array DIM2 Float
          in  evaluate (m `Acc.indexArray` (Z:.0:.0)) >> return m
  --
  let go f g = return (run_ref f mat, run_acc g mat')
  case alg of
    "sum-2d"     -> go sum2DRef sumAcc
    "product-2d" -> go prod2DRef prodAcc
    x            -> error $ "unknown variant: " ++ x
  where
    {-# NOINLINE run_ref #-}
    run_ref f xs () = f xs
    run_acc f xs () = f xs