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
-}

bW9kdWxlIE1haW4gd2hlcmUKIAppbXBvcnQgRGF0YS5BcnJheS5SZXBhIGhpZGluZyAobWFwKQppbXBvcnQgQ3JpdGVyaW9uLk1haW4KaW1wb3J0IENyaXRlcmlvbi5UeXBlcwppbXBvcnQgU3lzdGVtLklPLlVuc2FmZQoKCnR5cGUgRElNMV9VX0QgPSBBcnJheSBVIERJTTEgRG91YmxlCnR5cGUgRElNMV9EX0QgPSBBcnJheSBEIERJTTEgRG91YmxlCnR5cGUgRElNMl9EX0QgPSBBcnJheSBEIERJTTIgRG91YmxlCgpjYWxjIDo6IERJTTFfVV9EIC0+IERJTTFfVV9EIC0+IERJTTFfRF9ECmNhbGMgayBiID0gbGV0IHNpemVCID0gc2l6ZSAoZXh0ZW50IGIpCiAgICAgICAgICAgICAgIHNpemVLID0gc2l6ZSAoZXh0ZW50IGspCiAgICAgICAgICAgICAgIHAga3AgYnAgKFo6LmljKSA9IHN1bSAkIG1hcAogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAoXGkgLT4ga3AgKFo6LmkpICogYnAgKFo6LmljK2kpKQogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBbMC4uc2l6ZUstMV0KICAgICAgICAgICBpbiB0cmF2ZXJzZTIgayBiIChcXyBfIC0+IFo6LnNpemVCLXNpemVLKzEpIHAKCmNhbGMwLCBjYWxjMSwgY2FsYzIgOjogRElNMV9VX0QgLT4gRElNMV9VX0QgLT4gRElNMV9VX0QKY2FsYzAgayBiID0gdW5zYWZlUGVyZm9ybUlPICQgY29tcHV0ZVAgJCBjYWxjIGsgYiAKCmNhbGMxIGsgYiA9IGxldAogICAgWiA6LiBrcyA9IGV4dGVudCBrCiAgICBaIDouIGJzID0gZXh0ZW50IGIKICAgIHNoID0gWiA6LiBicyAtIGtzICsgMSA6LiBrcwogICAgaycsIGInIDo6IERJTTJfRF9ECiAgICBrJyA9IHRyYXZlcnNlIGsgKGNvbnN0IHNoKSAoXGYgKFogOi4gXyA6LiBqKSAtPiBmIChaIDouIGopKQogICAgYicgPSB0cmF2ZXJzZSBiIChjb25zdCBzaCkgKFxmIChaIDouIGkgOi4gaikgLT4gZiAoWiA6LiAoaStqKSkpCiAgaW4gdW5zYWZlUGVyZm9ybUlPICQgc3VtUCAkIGsnICpeIGInCgpjYWxjMiBrIGIgPSBsZXQKICAgIHMgKFogOi4ga3MpIChaIDouIGJzKSA9IFogOi4gYnMgLSBrcyArIDEgOi4ga3MKICAgIGYga2YgYmYgKFogOi4gaSA6LiBqKSA9IGtmIChaIDouIGopICogYmYgKFogOi4gKGkraikpCiAgaW4gdW5zYWZlUGVyZm9ybUlPICQgc3VtUCAkIHRyYXZlcnNlMiBrIGIgcyBmCgprZXJuZWwsIGJhc2UgOjogW0RvdWJsZV0KLS0gKGtlcm5lbCwgYmFzZSkgPSAoWzYsIDMsIDVdLCBbMywgNSwgMiwgOSwgNF0pCi0tIChrZXJuZWwsIGJhc2UpID0gKFsxLi4xMF0sIFsxLi4xMDBdKQooa2VybmVsLCBiYXNlKSA9IChbMS4uMjAwXSwgWzEuLjQwMDAwXSkKCgoKbWFpbiA9IGRvCiAgayA8LSByZXR1cm4gJCEgZnJvbUxpc3RVbmJveGVkIChaIDouIGxlbmd0aCBrZXJuZWwpIGtlcm5lbAogIGIgPC0gcmV0dXJuICQhIGZyb21MaXN0VW5ib3hlZCAoWiA6LiBsZW5ndGggYmFzZSkgYmFzZQotLSAgdGVzdDEgayBiCiAgdGVzdDIgayBiCgp0ZXN0MSBrIGIgPSBkbwogIHByaW50ICQgY2FsYzAgayBiCiAgcHJpbnQgJCBjYWxjMSBrIGIKICBwcmludCAkIGNhbGMyIGsgYgoKdGVzdDIgayBiID0KICBkZWZhdWx0TWFpbiBbYmVuY2ggImNhbGMwIiAkIHdobmYgKGNhbGMwIGspIGIsCiAgICAgICAgICAgICAgIGJlbmNoICJjYWxjMSIgJCB3aG5mIChjYWxjMSBrKSBiLAogICAgICAgICAgICAgICBiZW5jaCAiY2FsYzIiICQgd2huZiAoY2FsYzIgaykgYl0KCnstIENPTU1BTkQKZ2hjIC1PMiAtcnRzb3B0cyAtdGhyZWFkZWQgLS1tYWtlIDI2NS5ocwpzbGVlcCAxMDAKLi8yNjUgK1JUUyAtTgotfQoKey0gUkVTVUxUCi19