    {-# 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 arr1 :: Array r1 (sh :. Int) a
    arr1 arr2 :: 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 dim1 :: tail :. head
    dim1 dim2 :: tail :. head
    dim2
             | tail :. head
    dim1 (tail :. head) -> (tail :. head) -> Bool
    forall a. Eq a => a -> a -> Bool
    == tail :. head
    dim2
             , sh :: tail
    sh :. len :: 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
    * 2)

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

            elemFn :: ((tail :. head) -> p)
    -> ((tail :. head) -> p) -> (tail :. head) -> p
    elemFn get1 :: (tail :. head) -> p
    get1 get2 :: (tail :. head) -> p
    get2 (sh :: tail
    sh :. ix :: head
    ix)
             = case head
    ix head -> head -> head
    forall a. Integral a => a -> a -> a
    `mod` 2 of
                    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` 2)
                    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` 2)
                    _       -> [Char] -> p
    forall a. HasCallStack => [Char] -> a
    error "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 arr1 :: Array r1 (sh :. Int) a
    arr1 arr2 :: Array r2 (sh :. Int) a
    arr2 arr3 :: 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 dim1 :: tail :. head
    dim1 dim2 :: tail :. head
    dim2 dim3 :: 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
             , sh :: tail
    sh :. len :: 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
    * 3)

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

            elemFn :: ((tail :. head) -> p)
    -> ((tail :. head) -> p)
    -> ((tail :. head) -> p)
    -> (tail :. head)
    -> p
    elemFn get1 :: (tail :. head) -> p
    get1 get2 :: (tail :. head) -> p
    get2 get3 :: (tail :. head) -> p
    get3 (sh :: tail
    sh :. ix :: head
    ix)
             = case head
    ix head -> head -> head
    forall a. Integral a => a -> a -> a
    `mod` 3 of
                    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` 3)
                    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` 3)
                    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` 3)
                    _       -> [Char] -> p
    forall a. HasCallStack => [Char] -> a
    error "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 arr1 :: Array r1 (sh :. Int) a
    arr1 arr2 :: Array r2 (sh :. Int) a
    arr2 arr3 :: Array r3 (sh :. Int) a
    arr3 arr4 :: 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 dim1 :: tail :. head
    dim1 dim2 :: tail :. head
    dim2 dim3 :: tail :. head
    dim3 dim4 :: 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
             , sh :: tail
    sh :. len :: 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
    * 4)

             | Bool
    otherwise
             = [Char] -> tail :. head
    forall a. HasCallStack => [Char] -> a
    error "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 get1 :: (tail :. head) -> p
    get1 get2 :: (tail :. head) -> p
    get2 get3 :: (tail :. head) -> p
    get3 get4 :: (tail :. head) -> p
    get4 (sh :: tail
    sh :. ix :: head
    ix)
             = case head
    ix head -> head -> head
    forall a. Integral a => a -> a -> a
    `mod` 4 of
                    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` 4)
                    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` 4)
                    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` 4)
                    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` 4)
                    _       -> [Char] -> p
    forall a. HasCallStack => [Char] -> a
    error "Data.Array.Repa.interleave4: this never happens :-P"