    -----------------------------------------------------------------------------
    -- |
    -- Module      : Data.Array.Parallel.Unlifted.Sequential.Segmented.Basics
    -- Copyright   : (c) [2001..2002] Manuel M T Chakravarty & Gabriele Keller
    --		 (c) 2006         Manuel M T Chakravarty & Roman Leshchinskiy
    -- License     : see libraries/ndp/LICENSE
    --
    -- Maintainer  : Roman Leshchinskiy <rl@cse.unsw.edu.au>
    -- Stability   : internal
    -- Portability : portable
    --
    -- Description ---------------------------------------------------------------
    --
    --  Basics operations on unlifted segmented arrays.
    --
    -- Todo ----------------------------------------------------------------------
    --

    {-# LANGUAGE CPP #-}

    #include "fusion-phases.h"

    module Data.Array.Parallel.Unlifted.Sequential.Segmented.Basics (
      lengthSU, singletonSU, replicateSU, replicateCU, (!:^),
      flattenSU, (>:), segmentU, segmentArrU, concatSU, (^+:+^),
      sliceIndexSU, extractIndexSU, indexedSU,
      fstSU, sndSU, zipSU,
      enumFromToSU, enumFromThenToSU,
      toSU, fromSU,(+:+^)
    ) where

    import Data.Array.Parallel.Base (
      (:*:)(..))
    import Data.Array.Parallel.Stream (
      Step(..), Stream(..),SStream(..),(+++), (^+++^))
    import Data.Array.Parallel.Unlifted.Sequential.Flat

    import Data.Array.Parallel.Unlifted.Sequential.Segmented.Stream (streamSU,unstreamSU)
    import Data.Array.Parallel.Unlifted.Sequential.Segmented.SUArr (
      SUArr, lengthSU, (>:), flattenSU, segdSU, lengthsSU, indicesSU,
      lengthsToUSegd, singletonUSegd, toUSegd)

    -- TODO: Remove
    import Debug.Trace
    import Data.Array.Parallel.Unlifted.Sequential.Segmented.USegd (lengthsUSegd, indicesUSegd)

    -- lengthSU reexported from SUArr

    -- |Segmentation
    -- -------------

    -- flattenSU and (>:) reexported from SUArr

    singletonSU :: UA e => UArr e -> SUArr e
    {-# INLINE_U singletonSU #-}
    singletonSU es = singletonUSegd (lengthU es) >: es

    replicateSU :: UA e => UArr Int -> UArr e -> SUArr e
    {-# INLINE_U replicateSU #-}
    replicateSU ns es = lengthsToUSegd ns >: replicateEachU (sumU ns) ns es

    -- |Yield a segmented array, where each element contains the same array value
    --
    replicateCU:: (UA e) => Int -> UArr e -> SUArr e
    {-# INLINE_U replicateCU #-}
    replicateCU n arr = segmentArrU (replicateU n rLen) $ repeatU n arr
      where
        rLen    = lengthU arr

    -- |Array indexing
    --
    (!:^) :: (UA e) => SUArr e -> UArr Int -> UArr e
    {-# INLINE (!:^) #-}
    (!:^) sArr inds = bpermuteU (flattenSU sArr) newInds
      where
        xsLens  = lengthsSU sArr
        newInds = zipWithU (+) inds $ scanU (+) 0 xsLens

    -- |Segment an array according to the segmentation of the first argument
    --
    segmentU :: (UA e', UA e) => SUArr e' -> UArr e -> SUArr e
    {-# INLINE_U segmentU #-}
    segmentU template arr = segdSU template >: arr

    -- |Segment an array according to the segmentation of the first argument
    --
    segmentArrU :: (UA e) => UArr Int -> UArr e -> SUArr e
    {-# INLINE_U segmentArrU #-}
    segmentArrU lengths arr = (lengthsToUSegd lengths) >: arr

    -- |Concatenate the subarrays of an array of arrays
    --
    concatSU :: UA e => SUArr e -> UArr e
    concatSU = flattenSU

    -- |Indexing
    -- ---------

    -- |Extract the segment at the given index using the given extraction function
    -- (either 'sliceU' or 'extractU').
    --
    indexSU :: UA e => (UArr e -> Int -> Int -> UArr e) -> SUArr e -> Int -> UArr e
    indexSU copy sa i = copy (concatSU sa) (indicesSU sa !: i)
                                           (lengthsSU sa !: i)

    -- |Extract the segment at the given index without copying the elements.
    --
    sliceIndexSU :: UA e => SUArr e -> Int -> UArr e
    sliceIndexSU = indexSU sliceU

    -- |Extract the segment at the given index (elements are copied).
    --
    extractIndexSU :: UA e => SUArr e -> Int -> UArr e
    extractIndexSU = indexSU extractU

    -- |Associate each data element with its index
    --
    indexedSU :: UA e => SUArr e -> SUArr (Int :*: e)
    {-# INLINE_U indexedSU #-}
    indexedSU xss = segdSU xss >: zipU is xs
      where
        xs = concatSU xss

        is = enumFromToEachU (lengthU xs)
           . zipU (replicateU (lengthSU xss) 0)
           . mapU (subtract 1)
           $ lengthsSU xss

    -- |Concatenation
    -- --------------

    infixr 5 ^+:+^

    (^+:+^) :: UA a => SUArr a -> SUArr a -> SUArr a
    {-# INLINE_U (^+:+^) #-}
    xss ^+:+^ yss = toUSegd (zipU lens idxs)
                    >: unstreamU (streamSU xss ^+++^ streamSU yss)
      where
        lens = zipWithU (+) (lengthsSU xss) (lengthsSU yss)
        idxs = zipWithU (+) (indicesSU xss) (indicesSU yss)

    -- |Zipping
    -- --------

    fstSU :: (UA a, UA b) => SUArr (a :*: b) -> SUArr a
    {-# INLINE_U fstSU #-}
    fstSU sa = segdSU sa >: fstU (concatSU sa)

    sndSU :: (UA a, UA b) => SUArr (a :*: b) -> SUArr b
    {-# INLINE_U sndSU #-}
    sndSU sa = segdSU sa >: sndU (concatSU sa)

    zipSU :: (UA a, UA b) => SUArr a -> SUArr b -> SUArr (a :*: b)
    {-# INLINE_U zipSU #-}
    zipSU sa sb = segdSU sa >: zipU (concatSU sa) (concatSU sb)

    -- |Enumeration functions
    -- ----------------------

    -- |Yield a segmented enumerated array
    --
    enumFromToSU :: (Enum e, UA e) => UArr e -> UArr e -> SUArr e
    {-# INLINE_U enumFromToSU #-}
    enumFromToSU starts = enumFromThenToSU starts (mapU succ starts)

    -- |Yield a segmented enumerated array using a specific step
    --
    enumFromThenToSU :: (Enum e, UA e)
    		 => UArr e -> UArr e -> UArr e -> SUArr e
    {-# INLINE_U enumFromThenToSU #-}
    enumFromThenToSU starts nexts ends =
      segd >: unstreamU (enumFromThenToEachS len
                        (streamU (zipU (zipU lens starts) nexts)))
      where
        lens    = zipWith3U calcLen starts nexts ends
    	      where
    		calcLen start next end =
    		  abs (end' - start' + delta) `div` (abs delta)
    		  where
    		    start' = fromEnum start
    		    next'  = fromEnum next
    		    end'   = fromEnum end
    		    delta  = next' - start'
        len     = sumU lens
        segd    = lengthsToUSegd lens

    enumFromThenToEachS :: Enum a => Int -> Stream (Int :*: a :*: a) -> Stream a
    {-# INLINE_STREAM enumFromThenToEachS #-}
    enumFromThenToEachS n (Stream next s _) =
      Stream next' (0 :*: 0 :*: 0 :*: s) n
      where
        next' (0 :*: start :*: delta :*: s) =
          case next s of
            Done    -> Done
            Skip s' -> Skip (0 :*: start :*: delta :*: s')
            Yield (len :*: i :*: k) s'
                    -> let start' = fromEnum i
                       in Skip (len :*: start' :*: fromEnum k - start' :*: s')
        next' (n :*: start :*: delta :*: s) =
          Yield (toEnum start) (n-1 :*: start+delta :*: delta :*: s)

    -- |Concatenate two arrays
    --
    -- FIXME: rename
    --
    (+:+^) :: UA e => SUArr e -> SUArr e -> SUArr e
    {-# INLINE_U (+:+^) #-}
    a1 +:+^ a2 = unstreamSU $ SStream (segs1 +++ segs2) (vals1 +++ vals2)
      where
        (SStream segs1 vals1) = streamSU a1
        (SStream segs2 vals2) = streamSU a2

    -- |Conversion
    -- -----------

    -- |Turn a nested list into a segmented parallel array
    --
    toSU :: UA e => [[e]] -> SUArr e
    {-# INLINE_U toSU #-}
    toSU ls = let lens = toU $ map length ls
    	      a    = toU $ concat ls
              in
    	  lengthsToUSegd lens >: a

    -- |Turn a segmented array into a nested list
    --
    fromSU :: UA e => SUArr e -> [[e]]
    {-# INLINE_U fromSU #-}
    fromSU sa = let a       = concatSU sa
                    lens    = fromU $ lengthsSU sa
                    starts  = fromU $ indicesSU sa
                in
                [[a !: i | i <- [start .. start + len - 1]]
                                | (start, len) <- zip starts lens]