{-# 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 arr1 arr2
 = unsafeTraverse2 arr1 arr2 shapeFn elemFn
 where
	shapeFn dim1 dim2
	 | dim1 == dim2
	 , sh :. len	<- dim1
	 = sh :. (len * 2)

	 | otherwise
	 = error "Data.Array.Repa.interleave2: arrays must have same extent"

	elemFn get1 get2 (sh :. ix)
	 = case ix `mod` 2 of
		0	-> get1 (sh :. ix `div` 2)
		1	-> get2 (sh :. ix `div` 2)
		_	-> 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 arr1 arr2 arr3
 = unsafeTraverse3 arr1 arr2 arr3 shapeFn elemFn
 where
	shapeFn dim1 dim2 dim3
	 | dim1 == dim2
	 , dim1 == dim3
	 , sh :. len	<- dim1
	 = sh :. (len * 3)

	 | otherwise
	 = error "Data.Array.Repa.interleave3: arrays must have same extent"

	elemFn get1 get2 get3 (sh :. ix)
	 = case ix `mod` 3 of
		0	-> get1 (sh :. ix `div` 3)
		1	-> get2 (sh :. ix `div` 3)
		2	-> get3 (sh :. ix `div` 3)
		_	-> 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 arr1 arr2 arr3 arr4
 = unsafeTraverse4 arr1 arr2 arr3 arr4 shapeFn elemFn
 where
	shapeFn dim1 dim2 dim3 dim4
	 | dim1 == dim2
	 , dim1 == dim3
	 , dim1 == dim4
	 , sh :. len	<- dim1
	 = sh :. (len * 4)

	 | otherwise
	 = error "Data.Array.Repa.interleave4: arrays must have same extent"

	elemFn get1 get2 get3 get4 (sh :. ix)
	 = case ix `mod` 4 of
		0	-> get1 (sh :. ix `div` 4)
		1	-> get2 (sh :. ix `div` 4)
		2	-> get3 (sh :. ix `div` 4)
		3	-> get4 (sh :. ix `div` 4)
		_	-> error "Data.Array.Repa.interleave4: this never happens :-P"