content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE TypeFamilies #-}
  
  -- |
  -- Module      : Data.Massiv.Array.Delayed.Stream
  -- Copyright   : (c) Alexey Kuleshevich 2019-2022
  -- License     : BSD3
  -- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
  -- Stability   : experimental
  -- Portability : non-portable
  module Data.Massiv.Array.Delayed.Stream (
    DS (..),
    Array (..),
    toStreamArray,
    toStreamM,
    toStreamIxM,
    toSteps,
    fromSteps,
    fromStepsM,
  ) where
  
  import Control.Applicative
  import Control.Monad.ST
  import Data.Coerce
  import Data.Foldable
  import Data.Massiv.Array.Delayed.Pull
  import Data.Massiv.Core.Common
  import qualified Data.Massiv.Vector.Stream as S
  import GHC.Exts
  import Prelude hiding (drop, take)
  
  -- | Delayed stream array that represents a sequence of values that can be loaded
  -- sequentially. Important distinction from other arrays is that its size might no be
  -- known until it is computed.
  data DS = DS
  
  newtype instance Array DS Ix1 e = DSArray
    { dsArray :: S.Steps S.Id e
    }
  
  -- | /O(1)/ - Convert delayed stream array into `Steps`.
  --
  -- @since 0.4.1
  toSteps :: Vector DS e -> Steps Id e
  toSteps = coerce
  {-# INLINE toSteps #-}
  
  -- | /O(1)/ - Convert `Steps` into delayed stream array
  --
  -- @since 0.4.1
  fromSteps :: Steps Id e -> Vector DS e
  fromSteps = coerce
  {-# INLINE fromSteps #-}
  
  -- | /O(1)/ - Convert monadic `Steps` into delayed stream array
  --
  -- @since 0.5.0
  fromStepsM :: Monad m => Steps m e -> m (Vector DS e)
  fromStepsM = fmap DSArray . S.transSteps
  {-# INLINE fromStepsM #-}
  
  instance Shape DS Ix1 where
    linearSizeHint = stepsSize . dsArray
    {-# INLINE linearSizeHint #-}
  
    linearSize = SafeSz . unId . S.length . dsArray
    {-# INLINE linearSize #-}
  
    outerSize = linearSize
    {-# INLINE outerSize #-}
  
    isNull = S.unId . S.null . coerce
    {-# INLINE isNull #-}
  
  -- | For now only `Seq` strategy.
  instance Strategy DS where
    getComp _ = Seq
    setComp _ = id
    repr = DS
  
  instance Functor (Array DS Ix1) where
    fmap f = coerce . S.map f . dsArray
    {-# INLINE fmap #-}
    (<$) e = coerce . (e <$) . dsArray
    {-# INLINE (<$) #-}
  
  instance Applicative (Array DS Ix1) where
    pure = fromSteps . S.singleton
    {-# INLINE pure #-}
    (<*>) a1 a2 = fromSteps (S.zipWith ($) (coerce a1) (coerce a2))
    {-# INLINE (<*>) #-}
  
  #if MIN_VERSION_base(4,10,0)
    liftA2 f a1 a2 = fromSteps (S.zipWith f (coerce a1) (coerce a2))
    {-# INLINE liftA2 #-}
  #endif
  
  instance Monad (Array DS Ix1) where
    (>>=) arr f = coerce (S.concatMap (coerce . f) (dsArray arr))
    {-# INLINE (>>=) #-}
  
  instance Foldable (Array DS Ix1) where
    foldr f acc = S.unId . S.foldrLazy f acc . toSteps
    {-# INLINE foldr #-}
    foldl f acc = S.unId . S.foldlLazy f acc . toSteps
    {-# INLINE foldl #-}
    foldl' f acc = S.unId . S.foldl f acc . toSteps
    {-# INLINE foldl' #-}
    foldr1 f = S.unId . S.foldr1Lazy f . toSteps
    {-# INLINE foldr1 #-}
    foldl1 f = S.unId . S.foldl1Lazy f . toSteps
    {-# INLINE foldl1 #-}
    toList = S.toList . coerce
    {-# INLINE toList #-}
    length = S.unId . S.length . coerce
    {-# INLINE length #-}
    null = S.unId . S.null . coerce
    {-# INLINE null #-}
    sum = S.unId . S.foldl (+) 0 . toSteps
    {-# INLINE sum #-}
    product = S.unId . S.foldl (*) 1 . toSteps
    {-# INLINE product #-}
    maximum = S.unId . S.foldl1 max . toSteps
    {-# INLINE maximum #-}
    minimum = S.unId . S.foldl1 min . toSteps
    {-# INLINE minimum #-}
  
  instance Semigroup (Array DS Ix1 e) where
    (<>) a1 a2 = fromSteps (coerce a1 `S.append` coerce a2)
    {-# INLINE (<>) #-}
  
  instance Monoid (Array DS Ix1 e) where
    mempty = DSArray S.empty
    {-# INLINE mempty #-}
  #if !MIN_VERSION_base(4,11,0)
    mappend = (<>)
    {-# INLINE mappend #-}
  #endif
  
  instance IsList (Array DS Ix1 e) where
    type Item (Array DS Ix1 e) = e
    fromList = fromSteps . fromList
    {-# INLINE fromList #-}
    fromListN n = fromSteps . fromListN n
    {-# INLINE fromListN #-}
    toList = S.toList . coerce
    {-# INLINE toList #-}
  
  instance S.Stream DS Ix1 e where
    toStream = coerce
    {-# INLINE toStream #-}
    toStreamIx = S.indexed . coerce
    {-# INLINE toStreamIx #-}
  
  -- | Flatten an array into a stream of values.
  --
  -- @since 0.4.1
  toStreamArray :: (Index ix, Source r e) => Array r ix e -> Vector DS e
  toStreamArray = DSArray . S.steps
  {-# INLINE [1] toStreamArray #-}
  
  {-# RULES "toStreamArray/id" toStreamArray = id #-}
  
  -- | /O(1)/ - Convert an array into monadic `Steps`
  --
  -- @since 0.5.0
  toStreamM :: (Stream r ix e, Monad m) => Array r ix e -> Steps m e
  toStreamM = S.transStepsId . toStream
  {-# INLINE toStreamM #-}
  
  -- | /O(1)/ - Convert an array into monadic `Steps`
  --
  -- @since 0.5.0
  toStreamIxM :: (Stream r ix e, Monad m) => Array r ix e -> Steps m (ix, e)
  toStreamIxM = S.transStepsId . toStreamIx
  {-# INLINE toStreamIxM #-}
  
  -- | /O(n)/ - `size` implementation.
  instance Load DS Ix1 e where
    makeArrayLinear _ k = fromSteps . S.generate k
    {-# INLINE makeArrayLinear #-}
    replicate _ k = fromSteps . S.replicate k
    {-# INLINE replicate #-}
  
    iterArrayLinearST_ _scheduler arr uWrite =
      S.mapM_ (uncurry uWrite) $ S.indexed $ S.transStepsId (coerce arr)
    {-# INLINE iterArrayLinearST_ #-}
  
    unsafeLoadIntoST marr (DSArray sts) =
      S.unstreamIntoM marr (stepsSize sts) (stepsStream sts)
    {-# INLINE unsafeLoadIntoST #-}
  
    unsafeLoadIntoIO marr arr = stToIO $ unsafeLoadIntoST marr arr
    {-# INLINE unsafeLoadIntoIO #-}
  
  -- cons :: e -> Array DS Ix1 e -> Array DS Ix1 e
  -- cons e = coerce . S.cons e . dsArray
  -- {-# INLINE cons #-}
  
  -- uncons :: Array DS Ix1 e -> Maybe (e, Array DS Ix1 e)
  -- uncons = coerce . S.uncons . dsArray
  -- {-# INLINE uncons #-}
  
  -- snoc :: Array DS Ix1 e -> e -> Array DS Ix1 e
  -- snoc (DSArray sts) e = DSArray (S.snoc sts e)
  -- {-# INLINE snoc #-}
  
  -- TODO: skip the stride while loading
  -- instance StrideLoad DS Ix1 e where
  --   iterArrayLinearWithStrideST_ scheduler stride resultSize arr uWrite =
  --     let strideIx = unStride stride
  --         DIArray (DArray _ _ f) = arr
  --     in loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !start ->
  --           scheduleWork scheduler $
  --           iterLinearM_ resultSize start (totalElem resultSize) (numWorkers scheduler) (<) $
  --             \ !i ix -> uWrite i (f (liftIndex2 (*) strideIx ix))
  --   {-# INLINE iterArrayLinearWithStrideST_ #-}
  