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

module Data.Array.Repa.Operators.IndexSpace
        ( reshape
        , append, (++)
        , transpose
        , extract
        , backpermute,         unsafeBackpermute
        , backpermuteDft,      unsafeBackpermuteDft
        , extend,              unsafeExtend 
        , slice,               unsafeSlice)
where
import Data.Array.Repa.Index
import Data.Array.Repa.Slice
import Data.Array.Repa.Base
import Data.Array.Repa.Repr.Delayed
import Data.Array.Repa.Operators.Traversal
import Data.Array.Repa.Shape            as S
import Prelude                          hiding ((++), traverse)
import qualified Prelude                as P 

stage   = "Data.Array.Repa.Operators.IndexSpace"

-- Index space transformations ------------------------------------------------
-- | Impose a new shape on the elements of an array.
--   The new extent must be the same size as the original, else `error`.
reshape :: ( Shape sh1, Shape sh2
           , Source r1 e)
        => sh2
        -> Array r1 sh1 e
        -> Array D  sh2 e

reshape sh2 arr
        | not $ S.size sh2 == S.size (extent arr)
        = error 
        $ stage P.++ ".reshape: reshaped array will not match size of the original"

reshape sh2 arr
        = fromFunction sh2 
        $ unsafeIndex arr . fromIndex (extent arr) . toIndex sh2
{-# INLINE [2] reshape #-}
 

-- | Append two arrays.
append, (++)
        :: ( Shape sh
           , Source r1 e, Source r2 e)
        => Array r1 (sh :. Int) e
        -> Array r2 (sh :. Int) e
        -> Array D  (sh :. Int) e

append arr1 arr2
 = unsafeTraverse2 arr1 arr2 fnExtent fnElem
 where
        (_ :. n)        = extent arr1

        fnExtent (sh :. i) (_  :. j)
                = sh :. (i + j)

        fnElem f1 f2 (sh :. i)
                | i < n         = f1 (sh :. i)
                | otherwise     = f2 (sh :. (i - n))
{-# INLINE [2] append #-}


(++) arr1 arr2 = append arr1 arr2
{-# INLINE (++) #-}


-- | Transpose the lowest two dimensions of an array.
--      Transposing an array twice yields the original.
transpose
        :: (Shape sh, Source r e)
        => Array  r (sh :. Int :. Int) e
        -> Array  D (sh :. Int :. Int) e

transpose arr
 = unsafeTraverse arr
        (\(sh :. m :. n)        -> (sh :. n :.m))
        (\f -> \(sh :. i :. j)  -> f (sh :. j :. i))
{-# INLINE [2] transpose #-}


-- | Extract a sub-range of elements from an array.
extract :: (Shape sh, Source r e)
        => sh                   -- ^ Starting index.
        -> sh                   -- ^ Size of result.
        -> Array r sh e 
        -> Array D sh e
extract start sz arr
        = fromFunction sz (\ix -> arr `unsafeIndex` (addDim start ix))
{-# INLINE [2] extract #-}


-- | Backwards permutation of an array's elements.
backpermute, unsafeBackpermute
        :: forall r sh1 sh2 e
        .  ( Shape sh1, Shape sh2
           , Source r e)
        => sh2                  -- ^ Extent of result array.
        -> (sh2 -> sh1)         -- ^ Function mapping each index in the result array
                                --      to an index of the source array.
        -> Array r  sh1 e       -- ^ Source array.
        -> Array D  sh2 e

backpermute newExtent perm arr
        = traverse arr (const newExtent) (. perm)
{-# INLINE [2] backpermute #-}

unsafeBackpermute newExtent perm arr
        = unsafeTraverse arr (const newExtent) (. perm)
{-# INLINE [2] unsafeBackpermute #-}


-- | Default backwards permutation of an array's elements.
--      If the function returns `Nothing` then the value at that index is taken
--      from the default array (@arrDft@)
backpermuteDft, unsafeBackpermuteDft
        :: forall r1 r2 sh1 sh2 e
        .  ( Shape sh1,   Shape sh2
           , Source r1 e, Source r2 e)
        => Array r2 sh2 e       -- ^ Default values (@arrDft@)
        -> (sh2 -> Maybe sh1)   -- ^ Function mapping each index in the result array
                                --      to an index in the source array.
        -> Array r1 sh1 e       -- ^ Source array.
        -> Array D  sh2 e

backpermuteDft arrDft fnIndex arrSrc
        = fromFunction (extent arrDft) fnElem
        where   fnElem ix
                 = case fnIndex ix of
                        Just ix'        -> arrSrc `index` ix'
                        Nothing         -> arrDft `index` ix
{-# INLINE [2] backpermuteDft #-}

unsafeBackpermuteDft arrDft fnIndex arrSrc
        = fromFunction (extent arrDft) fnElem
        where   fnElem ix
                 = case fnIndex ix of
                        Just ix'        -> arrSrc `unsafeIndex` ix'
                        Nothing         -> arrDft `unsafeIndex` ix
{-# INLINE [2] unsafeBackpermuteDft #-}



-- | Extend an array, according to a given slice specification.
--
--   For example, to replicate the rows of an array use the following:
--
--   @extend (Any :. (5::Int) :. All) arr@
--
extend, unsafeExtend
        :: ( Slice sl
           , Shape (SliceShape sl)
           , Shape (FullShape sl)
           , Source r e)
        => sl
        -> Array r (SliceShape sl) e
        -> Array D (FullShape sl)  e

extend sl arr
        = backpermute
                (fullOfSlice sl (extent arr))
                (sliceOfFull sl)
                arr
{-# INLINE [2] extend #-}

unsafeExtend sl arr
        = unsafeBackpermute
                (fullOfSlice sl (extent arr))
                (sliceOfFull sl)
                arr
{-# INLINE [2] unsafeExtend #-}



-- | Take a slice from an array, according to a given specification.
--
--   For example, to take a row from a matrix use the following:
--
--   @slice arr (Any :. (5::Int) :. All)@
--
--   To take a column use:
--
--   @slice arr (Any :. (5::Int))@
--
slice, unsafeSlice
        :: ( Slice sl
           , Shape (FullShape sl)
           , Shape (SliceShape sl)
           , Source r e)
        => Array r (FullShape sl) e
        -> sl
        -> Array D (SliceShape sl) e

slice arr sl
        = backpermute
                (sliceOfFull sl (extent arr))
                (fullOfSlice sl)
                arr
{-# INLINE [2] slice #-}


unsafeSlice arr sl
        = unsafeBackpermute
                (sliceOfFull sl (extent arr))
                (fullOfSlice sl)
                arr
{-# INLINE [2] unsafeSlice #-}