{-# LANGUAGE TypeOperators, ExplicitForAll, FlexibleContexts #-}

module Data.Array.Repa.Operators.Interleave
        ( interleave2
        , interleave3
        , interleave4)
where
import Data.Array.Repa.Shape
import Data.Array.Repa.Index
import Data.Array.Repa.Base
import Data.Array.Repa.Repr.Delayed
import Data.Array.Repa.Operators.Traversal
import Prelude                          hiding ((++))


-- Interleave -----------------------------------------------------------------
-- | Interleave the elements of two arrays.
--   All the input arrays must have the same extent, else `error`.
--   The lowest dimension of the result array is twice the size of the inputs.
--
-- @
--  interleave2 a1 a2   b1 b2  =>  a1 b1 a2 b2
--              a3 a4   b3 b4      a3 b3 a4 b4
-- @
--
interleave2
        :: ( Shape sh
           , Source r1 a, Source r2 a)
        => Array r1 (sh :. Int) a
        -> Array r2 (sh :. Int) a
        -> Array D  (sh :. Int) a

{-# INLINE [2] interleave2 #-}
interleave2 :: Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a -> Array D (sh :. Int) a
interleave2 Array r1 (sh :. Int) a
arr1 Array r2 (sh :. Int) a
arr2
 = Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a
-> ((sh :. Int) -> (sh :. Int) -> sh :. Int)
-> (((sh :. Int) -> a) -> ((sh :. Int) -> a) -> (sh :. Int) -> a)
-> Array D (sh :. Int) a
forall r1 r2 sh sh' sh'' a b c.
(Source r1 a, Source r2 b, Shape sh, Shape sh') =>
Array r1 sh a
-> Array r2 sh' b
-> (sh -> sh' -> sh'')
-> ((sh -> a) -> (sh' -> b) -> sh'' -> c)
-> Array D sh'' c
unsafeTraverse2 Array r1 (sh :. Int) a
arr1 Array r2 (sh :. Int) a
arr2 (sh :. Int) -> (sh :. Int) -> sh :. Int
forall tail head.
(Eq tail, Eq head, Num head) =>
(tail :. head) -> (tail :. head) -> tail :. head
shapeFn ((sh :. Int) -> a) -> ((sh :. Int) -> a) -> (sh :. Int) -> a
forall head tail p.
Integral head =>
((tail :. head) -> p)
-> ((tail :. head) -> p) -> (tail :. head) -> p
elemFn
 where
        shapeFn :: (tail :. head) -> (tail :. head) -> tail :. head
shapeFn tail :. head
dim1 tail :. head
dim2
         | tail :. head
dim1 (tail :. head) -> (tail :. head) -> Bool
forall a. Eq a => a -> a -> Bool
== tail :. head
dim2
         , tail
sh :. head
len    <- tail :. head
dim1
         = tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. (head
len head -> head -> head
forall a. Num a => a -> a -> a
* head
2)

         | Bool
otherwise
         = [Char] -> tail :. head
forall a. HasCallStack => [Char] -> a
error [Char]
"Data.Array.Repa.interleave2: arrays must have same extent"

        elemFn :: ((tail :. head) -> p)
-> ((tail :. head) -> p) -> (tail :. head) -> p
elemFn (tail :. head) -> p
get1 (tail :. head) -> p
get2 (tail
sh :. head
ix)
         = case head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`mod` head
2 of
                head
0       -> (tail :. head) -> p
get1 (tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`div` head
2)
                head
1       -> (tail :. head) -> p
get2 (tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`div` head
2)
                head
_       -> [Char] -> p
forall a. HasCallStack => [Char] -> a
error [Char]
"Data.Array.Repa.interleave2: this never happens :-P"


-- | Interleave the elements of three arrays.
interleave3
        :: ( Shape sh
           , Source r1 a, Source r2 a, Source r3 a)
        => Array r1 (sh :. Int) a
        -> Array r2 (sh :. Int) a
        -> Array r3 (sh :. Int) a
        -> Array D  (sh :. Int) a

{-# INLINE [2] interleave3 #-}
interleave3 :: Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a
-> Array r3 (sh :. Int) a
-> Array D (sh :. Int) a
interleave3 Array r1 (sh :. Int) a
arr1 Array r2 (sh :. Int) a
arr2 Array r3 (sh :. Int) a
arr3
 = Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a
-> Array r3 (sh :. Int) a
-> ((sh :. Int) -> (sh :. Int) -> (sh :. Int) -> sh :. Int)
-> (((sh :. Int) -> a)
    -> ((sh :. Int) -> a) -> ((sh :. Int) -> a) -> (sh :. Int) -> a)
-> Array D (sh :. Int) a
forall r1 r2 r3 sh1 sh2 sh3 sh4 a b c d.
(Source r1 a, Source r2 b, Source r3 c, Shape sh1, Shape sh2,
 Shape sh3) =>
Array r1 sh1 a
-> Array r2 sh2 b
-> Array r3 sh3 c
-> (sh1 -> sh2 -> sh3 -> sh4)
-> ((sh1 -> a) -> (sh2 -> b) -> (sh3 -> c) -> sh4 -> d)
-> Array D sh4 d
unsafeTraverse3 Array r1 (sh :. Int) a
arr1 Array r2 (sh :. Int) a
arr2 Array r3 (sh :. Int) a
arr3 (sh :. Int) -> (sh :. Int) -> (sh :. Int) -> sh :. Int
forall tail head.
(Eq tail, Eq head, Num head) =>
(tail :. head) -> (tail :. head) -> (tail :. head) -> tail :. head
shapeFn ((sh :. Int) -> a)
-> ((sh :. Int) -> a) -> ((sh :. Int) -> a) -> (sh :. Int) -> a
forall head tail p.
Integral head =>
((tail :. head) -> p)
-> ((tail :. head) -> p)
-> ((tail :. head) -> p)
-> (tail :. head)
-> p
elemFn
 where
        shapeFn :: (tail :. head) -> (tail :. head) -> (tail :. head) -> tail :. head
shapeFn tail :. head
dim1 tail :. head
dim2 tail :. head
dim3
         | tail :. head
dim1 (tail :. head) -> (tail :. head) -> Bool
forall a. Eq a => a -> a -> Bool
== tail :. head
dim2
         , tail :. head
dim1 (tail :. head) -> (tail :. head) -> Bool
forall a. Eq a => a -> a -> Bool
== tail :. head
dim3
         , tail
sh :. head
len    <- tail :. head
dim1
         = tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. (head
len head -> head -> head
forall a. Num a => a -> a -> a
* head
3)

         | Bool
otherwise
         = [Char] -> tail :. head
forall a. HasCallStack => [Char] -> a
error [Char]
"Data.Array.Repa.interleave3: arrays must have same extent"

        elemFn :: ((tail :. head) -> p)
-> ((tail :. head) -> p)
-> ((tail :. head) -> p)
-> (tail :. head)
-> p
elemFn (tail :. head) -> p
get1 (tail :. head) -> p
get2 (tail :. head) -> p
get3 (tail
sh :. head
ix)
         = case head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`mod` head
3 of
                head
0       -> (tail :. head) -> p
get1 (tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`div` head
3)
                head
1       -> (tail :. head) -> p
get2 (tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`div` head
3)
                head
2       -> (tail :. head) -> p
get3 (tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`div` head
3)
                head
_       -> [Char] -> p
forall a. HasCallStack => [Char] -> a
error [Char]
"Data.Array.Repa.interleave3: this never happens :-P"


-- | Interleave the elements of four arrays.
interleave4
        :: ( Shape sh
           , Source r1 a, Source r2 a, Source r3 a, Source r4 a)
        => Array r1 (sh :. Int) a
        -> Array r2 (sh :. Int) a
        -> Array r3 (sh :. Int) a
        -> Array r4 (sh :. Int) a
        -> Array D  (sh :. Int) a

{-# INLINE [2] interleave4 #-}
interleave4 :: Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a
-> Array r3 (sh :. Int) a
-> Array r4 (sh :. Int) a
-> Array D (sh :. Int) a
interleave4 Array r1 (sh :. Int) a
arr1 Array r2 (sh :. Int) a
arr2 Array r3 (sh :. Int) a
arr3 Array r4 (sh :. Int) a
arr4
 = Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a
-> Array r3 (sh :. Int) a
-> Array r4 (sh :. Int) a
-> ((sh :. Int)
    -> (sh :. Int) -> (sh :. Int) -> (sh :. Int) -> sh :. Int)
-> (((sh :. Int) -> a)
    -> ((sh :. Int) -> a)
    -> ((sh :. Int) -> a)
    -> ((sh :. Int) -> a)
    -> (sh :. Int)
    -> a)
-> Array D (sh :. Int) a
forall r1 r2 r3 r4 sh1 sh2 sh3 sh4 sh5 a b c d e.
(Source r1 a, Source r2 b, Source r3 c, Source r4 d, Shape sh1,
 Shape sh2, Shape sh3, Shape sh4) =>
Array r1 sh1 a
-> Array r2 sh2 b
-> Array r3 sh3 c
-> Array r4 sh4 d
-> (sh1 -> sh2 -> sh3 -> sh4 -> sh5)
-> ((sh1 -> a)
    -> (sh2 -> b) -> (sh3 -> c) -> (sh4 -> d) -> sh5 -> e)
-> Array D sh5 e
unsafeTraverse4 Array r1 (sh :. Int) a
arr1 Array r2 (sh :. Int) a
arr2 Array r3 (sh :. Int) a
arr3 Array r4 (sh :. Int) a
arr4 (sh :. Int)
-> (sh :. Int) -> (sh :. Int) -> (sh :. Int) -> sh :. Int
forall tail head.
(Eq tail, Eq head, Num head) =>
(tail :. head)
-> (tail :. head)
-> (tail :. head)
-> (tail :. head)
-> tail :. head
shapeFn ((sh :. Int) -> a)
-> ((sh :. Int) -> a)
-> ((sh :. Int) -> a)
-> ((sh :. Int) -> a)
-> (sh :. Int)
-> a
forall head tail p.
Integral head =>
((tail :. head) -> p)
-> ((tail :. head) -> p)
-> ((tail :. head) -> p)
-> ((tail :. head) -> p)
-> (tail :. head)
-> p
elemFn
 where
        shapeFn :: (tail :. head)
-> (tail :. head)
-> (tail :. head)
-> (tail :. head)
-> tail :. head
shapeFn tail :. head
dim1 tail :. head
dim2 tail :. head
dim3 tail :. head
dim4
         | tail :. head
dim1 (tail :. head) -> (tail :. head) -> Bool
forall a. Eq a => a -> a -> Bool
== tail :. head
dim2
         , tail :. head
dim1 (tail :. head) -> (tail :. head) -> Bool
forall a. Eq a => a -> a -> Bool
== tail :. head
dim3
         , tail :. head
dim1 (tail :. head) -> (tail :. head) -> Bool
forall a. Eq a => a -> a -> Bool
== tail :. head
dim4
         , tail
sh :. head
len    <- tail :. head
dim1
         = tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. (head
len head -> head -> head
forall a. Num a => a -> a -> a
* head
4)

         | Bool
otherwise
         = [Char] -> tail :. head
forall a. HasCallStack => [Char] -> a
error [Char]
"Data.Array.Repa.interleave4: arrays must have same extent"

        elemFn :: ((tail :. head) -> p)
-> ((tail :. head) -> p)
-> ((tail :. head) -> p)
-> ((tail :. head) -> p)
-> (tail :. head)
-> p
elemFn (tail :. head) -> p
get1 (tail :. head) -> p
get2 (tail :. head) -> p
get3 (tail :. head) -> p
get4 (tail
sh :. head
ix)
         = case head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`mod` head
4 of
                head
0       -> (tail :. head) -> p
get1 (tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`div` head
4)
                head
1       -> (tail :. head) -> p
get2 (tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`div` head
4)
                head
2       -> (tail :. head) -> p
get3 (tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`div` head
4)
                head
3       -> (tail :. head) -> p
get4 (tail
sh tail -> head -> tail :. head
forall tail head. tail -> head -> tail :. head
:. head
ix head -> head -> head
forall a. Integral a => a -> a -> a
`div` head
4)
                head
_       -> [Char] -> p
forall a. HasCallStack => [Char] -> a
error [Char]
"Data.Array.Repa.interleave4: this never happens :-P"