{-# LANGUAGE BangPatterns #-}

-- | Functions specialised for arrays of dimension 2.
module Data.Array.Repa.Specialised.Dim2
	( isInside2
	, isOutside2
	, clampToBorder2
	, makeBordered2)
where
import Data.Array.Repa.Index
import Data.Array.Repa.Base
import Data.Array.Repa.Repr.Partitioned
import Data.Array.Repa.Repr.Undefined


-- | Check if an index lies inside the given extent.
--   As opposed to `inRange` from "Data.Array.Repa.Index",
--   this is a short-circuited test that checks that lowest dimension first.
isInside2
	:: DIM2 	-- ^ Extent of array.
	-> DIM2 	-- ^ Index to check.
	-> Bool

{-# INLINE isInside2 #-}
isInside2 ex 	= not . isOutside2 ex


-- | Check if an index lies outside the given extent.
--   As opposed to `inRange` from "Data.Array.Repa.Index",
--   this is a short-circuited test that checks the lowest dimension first.
isOutside2
	:: DIM2		-- ^ Extent of array.
	-> DIM2		-- ^ Index to check.
	-> Bool

{-# INLINE isOutside2 #-}
isOutside2 (_ :. yLen :. xLen) (_ :. yy :. xx)
	| xx < 0	= True
	| xx >= xLen	= True
	| yy < 0	= True
	| yy >= yLen	= True
	| otherwise	= False


-- | Given the extent of an array, clamp the components of an index so they
--   lie within the given array. Outlying indices are clamped to the index
--   of the nearest border element.
clampToBorder2
	:: DIM2 	-- ^ Extent of array.
	-> DIM2		-- ^ Index to clamp.
	-> DIM2

{-# INLINE clampToBorder2 #-}
clampToBorder2 (_ :. yLen :. xLen) (sh :. j :. i)
 = clampX j i
 where 	{-# INLINE clampX #-}
	clampX !y !x
	  | x < 0	= clampY y 0
	  | x >= xLen	= clampY y (xLen - 1)
	  | otherwise	= clampY y x

	{-# INLINE clampY #-}
	clampY !y !x
	  | y < 0	= sh :. 0	   :. x
	  | y >= yLen	= sh :. (yLen - 1) :. x
	  | otherwise	= sh :. y	   :. x



-- | Make a 2D partitioned array from two others, one to produce the elements
--   in the internal region, and one to produce elements in the border region.
--   The two arrays must have the same extent.
--   The border must be the same width on all sides.
--
makeBordered2
	:: (Source r1 a, Source r2 a)
        => DIM2			-- ^ Extent of array.
	-> Int			-- ^ Width of border.
	-> Array r1 DIM2 a	-- ^ Array for internal elements.
	-> Array r2 DIM2 a	-- ^ Array for border elements.
	-> Array (P r1 (P r2 (P r2 (P r2 (P r2 X))))) DIM2 a

{-# INLINE makeBordered2 #-}
makeBordered2 sh@(_ :. aHeight :. aWidth) bWidth arrInternal arrBorder
 = checkDims `seq` 
   let
	-- minimum and maximum indicies of values in the inner part of the image.
	!inX		= bWidth
	!inY		= bWidth
        !inW            = aWidth  - 2 * bWidth 
        !inH            = aHeight - 2 * bWidth

	inInternal (Z :. y :. x)
		=  x >= inX && x < (inX + inW)
		&& y >= inY && y < (inY + inH)
        {-# INLINE inInternal #-}

	inBorder 	= not . inInternal
        {-# INLINE inBorder #-}

   in	
    --  internal region
        APart sh (Range (Z :. inY     :. inX)       (Z :. inH :. inW )    inInternal) arrInternal

    --  border regions
    $   APart sh (Range (Z :. 0         :. 0)         (Z :. bWidth :. aWidth) inBorder) arrBorder
    $   APart sh (Range (Z :. inY + inH :. 0)         (Z :. bWidth :. aWidth) inBorder) arrBorder
    $   APart sh (Range (Z :. inY       :. 0)         (Z :. inH    :. bWidth) inBorder) arrBorder
    $   APart sh (Range (Z :. inY       :. inX + inW) (Z :. inH    :. bWidth) inBorder) arrBorder
    $   AUndefined sh

 where
        checkDims
         = if (extent arrInternal) == (extent arrBorder)
                then ()
                else error "makeBordered2: internal and border arrays have different extents"
        {-# NOINLINE checkDims #-}
        --  NOINLINE because we don't want the branch in the core code.