module Main where
 
import Data.Array.Repa hiding (map)
import Criterion.Main
import Criterion.Types
import System.IO.Unsafe


type DIM1_U_D = Array U DIM1 Double
type DIM1_D_D = Array D DIM1 Double
type DIM2_D_D = Array D DIM2 Double

calc :: DIM1_U_D -> DIM1_U_D -> DIM1_D_D
calc k b = let sizeB = size (extent b)
               sizeK = size (extent k)
               p kp bp (Z:.ic) = sum $ map
                                 (\i -> kp (Z:.i) * bp (Z:.ic+i))
                                 [0..sizeK-1]
           in traverse2 k b (\_ _ -> Z:.sizeB-sizeK+1) p

calc0, calc1, calc2 :: DIM1_U_D -> DIM1_U_D -> DIM1_U_D
calc0 k b = unsafePerformIO $ computeP $ calc k b 

calc1 k b = let
    Z :. ks = extent k
    Z :. bs = extent b
    sh = Z :. bs - ks + 1 :. ks
    k', b' :: DIM2_D_D
    k' = traverse k (const sh) (\f (Z :. _ :. j) -> f (Z :. j))
    b' = traverse b (const sh) (\f (Z :. i :. j) -> f (Z :. (i+j)))
  in unsafePerformIO $ sumP $ k' *^ b'

calc2 k b = let
    s (Z :. ks) (Z :. bs) = Z :. bs - ks + 1 :. ks
    f kf bf (Z :. i :. j) = kf (Z :. j) * bf (Z :. (i+j))
  in unsafePerformIO $ sumP $ traverse2 k b s f

kernel, base :: [Double]
-- (kernel, base) = ([6, 3, 5], [3, 5, 2, 9, 4])
-- (kernel, base) = ([1..10], [1..100])
(kernel, base) = ([1..200], [1..40000])



main = do
  k <- return $! fromListUnboxed (Z :. length kernel) kernel
  b <- return $! fromListUnboxed (Z :. length base) base
--  test1 k b
  test2 k b

test1 k b = do
  print $ calc0 k b
  print $ calc1 k b
  print $ calc2 k b

test2 k b =
  defaultMain [bench "calc0" $ whnf (calc0 k) b,
               bench "calc1" $ whnf (calc1 k) b,
               bench "calc2" $ whnf (calc2 k) b]

{- COMMAND
ghc -O2 -rtsopts -threaded --make 265.hs
sleep 100
./265 +RTS -N
-}

{- RESULT
-}