    {-# LANGUAGE FunctionalDependencies, UndecidableInstances #-}

    module Data.Array.Repa.Operators.Mapping
            ( -- * Generic maps
              map
            , zipWith
            , (+^), (-^), (*^), (/^)

              -- * Structured maps
            , Structured(..))
    where
    import Data.Array.Repa.Shape
    import Data.Array.Repa.Base
    import Data.Array.Repa.Repr.ByteString
    import Data.Array.Repa.Repr.Cursored
    import Data.Array.Repa.Repr.Delayed
    import Data.Array.Repa.Repr.ForeignPtr
    import Data.Array.Repa.Repr.HintSmall
    import Data.Array.Repa.Repr.HintInterleave
    import Data.Array.Repa.Repr.Partitioned
    import Data.Array.Repa.Repr.Unboxed
    import Data.Array.Repa.Repr.Undefined
    import Prelude hiding (map, zipWith)
    import Foreign.Storable
    import Data.Word

    -- | Apply a worker function to each element of an array,
    --   yielding a new array with the same extent.
    --
    map     :: (Shape sh, Source r a)
            => (a -> b) -> Array r sh a -> Array D sh b
    map f arr
     = case delay arr of
            ADelayed sh g -> ADelayed sh (f . g)
    {-# INLINE [3] map #-}

    -- ZipWith --------------------------------------------------------------------
    -- | Combine two arrays, element-wise, with a binary operator.
    --	If the extent of the two array arguments differ,
    --	then the resulting array's extent is their intersection.
    --
    zipWith :: (Shape sh, Source r1 a, Source r2 b)
            => (a -> b -> c)
            -> Array r1 sh a -> Array r2 sh b
            -> Array D sh c
    zipWith f arr1 arr2
     = let  get ix  = f (arr1 `unsafeIndex` ix) (arr2 `unsafeIndex` ix)
            {-# INLINE get #-}

       in   fromFunction
                    (intersectDim (extent arr1) (extent arr2))
                    get
    {-# INLINE [2] zipWith #-}

    infixl 7  *^, /^
    infixl 6  +^, -^

    (+^)	= zipWith (+)
    {-# INLINE (+^) #-}

    (-^)	= zipWith (-)
    {-# INLINE (-^) #-}

    (*^)	= zipWith (*)
    {-# INLINE (*^) #-}

    (/^)	= zipWith (/)
    {-# INLINE (/^) #-}

    -- Structured -------------------------------------------------------------------
    -- | Structured versions of @map@ and @zipWith@ that preserve the representation
    --   of cursored and partitioned arrays.
    --
    --   For cursored (@C@) arrays, the cursoring of the source array is preserved.
    --
    --   For partitioned (@P@) arrays, the worker function is fused with each array
    --   partition separately, instead of treating the whole array as a single
    --   bulk object.
    --
    --   Preserving the cursored and\/or paritioned representation of an array
    --   is will make follow-on computation more efficient than if the array was
    --   converted to a vanilla Delayed (@D@) array as with plain `map` and `zipWith`.
    --
    --   If the source array is not cursored or partitioned then `smap` and
    --   `szipWith` are identical to the plain functions.
    --
    class Structured r1 a b where
     -- | The target result representation.
     type TR r1

     -- | Structured @map@.
     smap   :: Shape sh
            => (a -> b)
            -> Array r1     sh a
            -> Array (TR r1) sh b

     -- | Structured @zipWith@.
     --   If you have a cursored or partitioned source array then use that as
     --   the third argument (corresponding to @r1@ here)
     szipWith
            :: (Shape sh, Source r c)
            => (c -> a -> b)
            -> Array r      sh c
            -> Array r1     sh a
            -> Array (TR r1) sh b

    -- ByteString -------------------------
    instance Structured B Word8 b where
     type TR B = D
     smap           = map
     szipWith       = zipWith

    -- Cursored ---------------------------
    instance Structured C a b where
     type TR C = C

     smap f (ACursored sh makec shiftc loadc)
            = ACursored sh makec shiftc (f . loadc)
     {-# INLINE [3] smap #-}

     szipWith f arr1 (ACursored sh makec shiftc loadc)
      = let makec' ix               = (ix, makec ix)
            {-# INLINE makec' #-}

            shiftc' off (ix, cur)   = (addDim off ix, shiftc off cur)
            {-# INLINE shiftc' #-}

            load' (ix, cur)         = f (arr1 `unsafeIndex` ix) (loadc cur)
            {-# INLINE load' #-}

        in  ACursored
                    (intersectDim (extent arr1) sh)
                    makec' shiftc' load'
     {-# INLINE [2] szipWith #-}

    -- Delayed ----------------------------
    instance Structured D a b where
     type TR D = D
     smap           = map
     szipWith       = zipWith

    -- ForeignPtr -------------------------
    instance Storable a => Structured F a b where
     type TR F = D
     smap           = map
     szipWith       = zipWith

    -- Partitioned ------------------------
    instance (Structured r1 a b
            , Structured r2 a b)
           => Structured (P r1 r2) a b where
     type TR (P r1 r2) = P (TR r1) (TR r2)

     smap f (APart sh range arr1 arr2)
            = APart sh range (smap f arr1) (smap f arr2)
     {-# INLINE [3] smap #-}

     szipWith f arr1 (APart sh range arr21 arr22)
            = APart sh range (szipWith f arr1 arr21)
                             (szipWith f arr1 arr22)
     {-# INLINE [2] szipWith #-}

    -- Small ------------------------------
    instance   Structured r1 a b
            => Structured (S r1) a b where
     type TR (S r1) = S (TR r1)

     smap f (ASmall arr1)
            = ASmall (smap f arr1)
     {-# INLINE [3] smap #-}

     szipWith f arr1 (ASmall arr2)
            = ASmall (szipWith f arr1 arr2)
     {-# INLINE [3] szipWith #-}

    -- Interleaved ------------------------
    instance   Structured r1 a b
            => Structured (I r1) a b where
     type TR (I r1) = I (TR r1)

     smap f (AInterleave arr1)
            = AInterleave (smap f arr1)
     {-# INLINE [3] smap #-}

     szipWith f arr1 (AInterleave arr2)
            = AInterleave (szipWith f arr1 arr2)
     {-# INLINE [3] szipWith #-}

    -- Unboxed ----------------------------
    instance Unbox a => Structured U a b where
     type TR U = D
     smap           = map
     szipWith       = zipWith

    -- Undefined --------------------------
    instance Structured X a b where
     type TR X = X
     smap     _   (AUndefined sh) = AUndefined sh
     szipWith _ _ (AUndefined sh) = AUndefined sh