content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# OPTIONS_GHC -fno-warn-redundant-constraints #-}
  
  -- |
  -- Module      : Data.Massiv.Array.Ops.Transform
  -- Copyright   : (c) Alexey Kuleshevich 2018-2022
  -- License     : BSD3
  -- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
  -- Stability   : experimental
  -- Portability : non-portable
  module Data.Massiv.Array.Ops.Transform (
    -- ** Transpose
    transpose,
    transposeInner,
    transposeOuter,
  
    -- ** Reverse
    reverse,
    reverse',
    reverseM,
  
    -- ** Backpermute
    backpermuteM,
    backpermute',
  
    -- ** Resize
    resizeM,
    resize',
    flatten,
  
    -- ** Extract
    extractM,
    extract',
    extractFromToM,
    extractFromTo',
    deleteRowsM,
    deleteColumnsM,
    deleteRegionM,
  
    -- ** Append/Split
    appendOuterM,
    appendM,
    append',
    concatOuterM,
    concatM,
    concat',
    stackSlicesM,
    stackOuterSlicesM,
    stackInnerSlicesM,
    splitAtM,
    splitAt',
    splitExtractM,
    replaceSlice,
    replaceOuterSlice,
  
    -- ** Upsample/Downsample
    upsample,
    downsample,
  
    -- ** Zoom
    zoom,
    zoomWithGrid,
  
    -- ** Transform
    transformM,
    transform',
    transform2M,
    transform2',
  ) where
  
  import Control.Monad as M (foldM_, forM_, unless)
  import Control.Monad.ST
  import Control.Scheduler (traverse_)
  import Data.Bifunctor (bimap)
  import Data.Foldable as F (foldl', foldrM, length, toList)
  import qualified Data.List as L (uncons)
  import Data.Massiv.Array.Delayed.Pull
  import Data.Massiv.Array.Delayed.Push
  import Data.Massiv.Array.Mutable
  import Data.Massiv.Array.Ops.Construct
  import Data.Massiv.Array.Ops.Map
  import Data.Massiv.Core
  import Data.Massiv.Core.Common -- (size, unsafeIndex, unsafeResize, evaluate', evaluateM)
  import Prelude as P hiding (
    concat,
    drop,
    mapM_,
    reverse,
    splitAt,
    take,
    traverse,
   )
  
  -- | Extract a sub-array from within a larger source array. Array that is being extracted must be
  -- fully encapsulated in a source array, otherwise `SizeSubregionException` will be thrown.
  --
  -- ====__Examples__
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> m <- resizeM (Sz (3 :. 3)) $ Ix1 1 ... 9
  -- >>> m
  -- Array D Seq (Sz (3 :. 3))
  --   [ [ 1, 2, 3 ]
  --   , [ 4, 5, 6 ]
  --   , [ 7, 8, 9 ]
  --   ]
  -- >>> extractM (0 :. 1) (Sz (2 :. 2)) m
  -- Array D Seq (Sz (2 :. 2))
  --   [ [ 2, 3 ]
  --   , [ 5, 6 ]
  --   ]
  -- >>> a <- resizeM (Sz (3 :> 2 :. 4)) $ Ix1 11 ... 34
  -- >>> a
  -- Array D Seq (Sz (3 :> 2 :. 4))
  --   [ [ [ 11, 12, 13, 14 ]
  --     , [ 15, 16, 17, 18 ]
  --     ]
  --   , [ [ 19, 20, 21, 22 ]
  --     , [ 23, 24, 25, 26 ]
  --     ]
  --   , [ [ 27, 28, 29, 30 ]
  --     , [ 31, 32, 33, 34 ]
  --     ]
  --   ]
  -- >>> extractM (0 :> 1 :. 1) (Sz (3 :> 1 :. 2)) a
  -- Array D Seq (Sz (3 :> 1 :. 2))
  --   [ [ [ 16, 17 ]
  --     ]
  --   , [ [ 24, 25 ]
  --     ]
  --   , [ [ 32, 33 ]
  --     ]
  --   ]
  --
  -- @since 0.3.0
  extractM
    :: forall r ix e m
     . (MonadThrow m, Index ix, Source r e)
    => ix
    -- ^ Starting index
    -> Sz ix
    -- ^ Size of the resulting array
    -> Array r ix e
    -- ^ Source array
    -> m (Array D ix e)
  extractM !sIx !newSz !arr
    | isSafeIndex sz1 sIx && isSafeIndex eIx1 sIx && isSafeIndex sz1 eIx =
        pure $ unsafeExtract sIx newSz arr
    | otherwise = throwM $ SizeSubregionException (size arr) sIx newSz
    where
      sz1 = Sz (liftIndex (+ 1) (unSz (size arr)))
      eIx1 = Sz (liftIndex (+ 1) eIx)
      eIx = liftIndex2 (+) sIx $ unSz newSz
  {-# INLINE extractM #-}
  
  -- | Same as `extractM`, but will throw a runtime exception from pure code if supplied dimensions
  -- are incorrect.
  --
  -- @since 0.1.0
  extract'
    :: forall r ix e
     . (HasCallStack, Index ix, Source r e)
    => ix
    -- ^ Starting index
    -> Sz ix
    -- ^ Size of the resulting array
    -> Array r ix e
    -- ^ Source array
    -> Array D ix e
  extract' sIx newSz = throwEither . extractM sIx newSz
  {-# INLINE extract' #-}
  
  -- | Similar to `extractM`, except it takes starting and ending index. Result array will not include
  -- the ending index.
  --
  -- ====__Examples__
  --
  -- >>> a <- resizeM (Sz (3 :> 2 :. 4)) $ Ix1 11 ... 34
  -- >>> a
  -- Array D Seq (Sz (3 :> 2 :. 4))
  --   [ [ [ 11, 12, 13, 14 ]
  --     , [ 15, 16, 17, 18 ]
  --     ]
  --   , [ [ 19, 20, 21, 22 ]
  --     , [ 23, 24, 25, 26 ]
  --     ]
  --   , [ [ 27, 28, 29, 30 ]
  --     , [ 31, 32, 33, 34 ]
  --     ]
  --   ]
  -- >>> extractFromToM (1 :> 0 :. 1) (3 :> 2 :. 4) a
  -- Array D Seq (Sz (2 :> 2 :. 3))
  --   [ [ [ 20, 21, 22 ]
  --     , [ 24, 25, 26 ]
  --     ]
  --   , [ [ 28, 29, 30 ]
  --     , [ 32, 33, 34 ]
  --     ]
  --   ]
  --
  -- @since 0.3.0
  extractFromToM
    :: forall r ix e m
     . (MonadThrow m, Index ix, Source r e)
    => ix
    -- ^ Starting index
    -> ix
    -- ^ Index up to which elements should be extracted.
    -> Array r ix e
    -- ^ Source array.
    -> m (Array D ix e)
  extractFromToM sIx eIx = extractM sIx (Sz (liftIndex2 (-) eIx sIx))
  {-# INLINE extractFromToM #-}
  
  -- | Same as `extractFromToM`, but throws an error on invalid indices.
  --
  -- @since 0.2.4
  extractFromTo'
    :: forall r ix e
     . (HasCallStack, Index ix, Source r e)
    => ix
    -- ^ Starting index
    -> ix
    -- ^ Index up to which elmenets should be extracted.
    -> Array r ix e
    -- ^ Source array.
    -> Array D ix e
  extractFromTo' sIx eIx = extract' sIx $ Sz (liftIndex2 (-) eIx sIx)
  {-# INLINE extractFromTo' #-}
  
  -- | /O(1)/ - Change the size of an array. Throws
  -- `SizeElementsMismatchException` if total number of elements does not match
  -- the supplied array.
  --
  -- @since 0.3.0
  resizeM
    :: forall r ix ix' e m
     . (MonadThrow m, Index ix', Index ix, Size r)
    => Sz ix'
    -> Array r ix e
    -> m (Array r ix' e)
  resizeM sz arr = guardNumberOfElements (size arr) sz >> pure (unsafeResize sz arr)
  {-# INLINE resizeM #-}
  
  -- | Same as `resizeM`, but will throw an error if supplied dimensions are incorrect.
  --
  -- @since 0.1.0
  resize'
    :: forall r ix ix' e
     . (HasCallStack, Index ix', Index ix, Size r)
    => Sz ix'
    -> Array r ix e
    -> Array r ix' e
  resize' sz = throwEither . resizeM sz
  {-# INLINE resize' #-}
  
  -- | /O(1)/ - Reduce a multi-dimensional array into a flat vector
  --
  -- @since 0.3.1
  flatten :: forall r ix e. (Index ix, Size r) => Array r ix e -> Vector r e
  flatten arr = unsafeResize (SafeSz (totalElem (size arr))) arr
  {-# INLINE flatten #-}
  
  -- | Transpose a 2-dimensional array
  --
  -- ==== __Examples__
  --
  -- >>> import Data.Massiv.Array
  -- >>> arr = makeArrayLinearR D Seq (Sz (2 :. 3)) id
  -- >>> arr
  -- Array D Seq (Sz (2 :. 3))
  --   [ [ 0, 1, 2 ]
  --   , [ 3, 4, 5 ]
  --   ]
  -- >>> transpose arr
  -- Array D Seq (Sz (3 :. 2))
  --   [ [ 0, 3 ]
  --   , [ 1, 4 ]
  --   , [ 2, 5 ]
  --   ]
  --
  -- @since 0.1.0
  transpose :: forall r e. Source r e => Matrix r e -> Matrix D e
  transpose = transposeInner
  {-# INLINE [1] transpose #-}
  
  {-# RULES
  "transpose . transpose" [~1] forall arr. transpose (transpose arr) = delay arr
  "transposeInner . transposeInner" [~1] forall arr. transposeInner (transposeInner arr) = delay arr
  "transposeOuter . transposeOuter" [~1] forall arr. transposeOuter (transposeOuter arr) = delay arr
    #-}
  
  -- | Transpose inner two dimensions of at least rank-2 array.
  --
  -- ===__Examples__
  --
  -- >>> import Data.Massiv.Array
  -- >>> arr = makeArrayLinearR U Seq (Sz (2 :> 3 :. 4)) id
  -- >>> arr
  -- Array U Seq (Sz (2 :> 3 :. 4))
  --   [ [ [ 0, 1, 2, 3 ]
  --     , [ 4, 5, 6, 7 ]
  --     , [ 8, 9, 10, 11 ]
  --     ]
  --   , [ [ 12, 13, 14, 15 ]
  --     , [ 16, 17, 18, 19 ]
  --     , [ 20, 21, 22, 23 ]
  --     ]
  --   ]
  -- >>> transposeInner arr
  -- Array D Seq (Sz (3 :> 2 :. 4))
  --   [ [ [ 0, 1, 2, 3 ]
  --     , [ 12, 13, 14, 15 ]
  --     ]
  --   , [ [ 4, 5, 6, 7 ]
  --     , [ 16, 17, 18, 19 ]
  --     ]
  --   , [ [ 8, 9, 10, 11 ]
  --     , [ 20, 21, 22, 23 ]
  --     ]
  --   ]
  --
  -- @since 0.1.0
  transposeInner
    :: forall r ix e
     . (Index (Lower ix), Index ix, Source r e)
    => Array r ix e
    -> Array D ix e
  transposeInner !arr = makeArray (getComp arr) newsz newVal
    where
      transInner !ix =
        either throwImpossible id $ do
          n <- getDimM ix dix
          m <- getDimM ix (dix - 1)
          ix' <- setDimM ix dix m
          setDimM ix' (dix - 1) n
      {-# INLINE transInner #-}
      newVal = unsafeIndex arr . transInner
      {-# INLINE newVal #-}
      !newsz = Sz (transInner (unSz (size arr)))
      !dix = dimensions newsz
  {-# INLINE [1] transposeInner #-}
  
  -- | Transpose outer two dimensions of at least rank-2 array.
  --
  -- ====__Examples__
  --
  -- >>> import Data.Massiv.Array
  -- >>> :seti -XTypeApplications
  -- >>> arr = makeArrayLinear @U Seq (Sz (2 :> 3 :. 4)) id
  -- >>> arr
  -- Array U Seq (Sz (2 :> 3 :. 4))
  --   [ [ [ 0, 1, 2, 3 ]
  --     , [ 4, 5, 6, 7 ]
  --     , [ 8, 9, 10, 11 ]
  --     ]
  --   , [ [ 12, 13, 14, 15 ]
  --     , [ 16, 17, 18, 19 ]
  --     , [ 20, 21, 22, 23 ]
  --     ]
  --   ]
  -- >>> transposeOuter arr
  -- Array D Seq (Sz (2 :> 4 :. 3))
  --   [ [ [ 0, 4, 8 ]
  --     , [ 1, 5, 9 ]
  --     , [ 2, 6, 10 ]
  --     , [ 3, 7, 11 ]
  --     ]
  --   , [ [ 12, 16, 20 ]
  --     , [ 13, 17, 21 ]
  --     , [ 14, 18, 22 ]
  --     , [ 15, 19, 23 ]
  --     ]
  --   ]
  --
  --
  -- @since 0.1.0
  transposeOuter
    :: forall r ix e
     . (Index (Lower ix), Index ix, Source r e)
    => Array r ix e
    -> Array D ix e
  transposeOuter !arr = makeArray (getComp arr) newsz newVal
    where
      transOuter !ix =
        either throwImpossible id $ do
          n <- getDimM ix 1
          m <- getDimM ix 2
          ix' <- setDimM ix 1 m
          setDimM ix' 2 n
      {-# INLINE transOuter #-}
      newVal = unsafeIndex arr . transOuter
      {-# INLINE newVal #-}
      !newsz = Sz (transOuter (unSz (size arr)))
  {-# INLINE [1] transposeOuter #-}
  
  -- | Reverse an array along some dimension. Dimension supplied is checked at compile time.
  --
  -- ==== __Example__
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> arr = makeArrayLinear Seq (Sz2 4 5) (+10) :: Array D Ix2 Int
  -- >>> arr
  -- Array D Seq (Sz (4 :. 5))
  --   [ [ 10, 11, 12, 13, 14 ]
  --   , [ 15, 16, 17, 18, 19 ]
  --   , [ 20, 21, 22, 23, 24 ]
  --   , [ 25, 26, 27, 28, 29 ]
  --   ]
  -- >>> A.reverse Dim1 arr
  -- Array D Seq (Sz (4 :. 5))
  --   [ [ 14, 13, 12, 11, 10 ]
  --   , [ 19, 18, 17, 16, 15 ]
  --   , [ 24, 23, 22, 21, 20 ]
  --   , [ 29, 28, 27, 26, 25 ]
  --   ]
  -- >>> A.reverse Dim2 arr
  -- Array D Seq (Sz (4 :. 5))
  --   [ [ 25, 26, 27, 28, 29 ]
  --   , [ 20, 21, 22, 23, 24 ]
  --   , [ 15, 16, 17, 18, 19 ]
  --   , [ 10, 11, 12, 13, 14 ]
  --   ]
  --
  -- @since 0.4.1
  reverse
    :: forall n r ix e
     . (IsIndexDimension ix n, Index ix, Source r e)
    => Dimension n
    -> Array r ix e
    -> Array D ix e
  reverse dim = reverse' (fromDimension dim)
  {-# INLINE reverse #-}
  
  -- | Similarly to `reverse`, flip an array along a particular dimension, but throws
  -- `IndexDimensionException` for an incorrect dimension.
  --
  -- @since 0.4.1
  reverseM
    :: forall r ix e m
     . (MonadThrow m, Index ix, Source r e)
    => Dim
    -> Array r ix e
    -> m (Array D ix e)
  reverseM dim arr = do
    let sz = size arr
    k <- getDimM (unSz sz) dim
    pure $ makeArray (getComp arr) sz $ \ix ->
      unsafeIndex arr (snd $ modifyDim' ix dim (\i -> k - i - 1))
  {-# INLINE reverseM #-}
  
  -- | Reverse an array along some dimension. Same as `reverseM`, but throws the
  -- `IndexDimensionException` from pure code.
  --
  -- @since 0.4.1
  reverse'
    :: forall r ix e
     . (HasCallStack, Index ix, Source r e)
    => Dim
    -> Array r ix e
    -> Array D ix e
  reverse' dim = throwEither . reverseM dim
  {-# INLINE reverse' #-}
  
  -- | Rearrange elements of an array into a new one by using a function that maps indices of the
  -- newly created one into the old one. This function can throw `IndexOutOfBoundsException`.
  --
  -- ===__Examples__
  --
  -- >>> import Data.Massiv.Array
  -- >>> :seti -XTypeApplications
  -- >>> arr = makeArrayLinear @D Seq (Sz (2 :> 3 :. 4)) id
  -- >>> arr
  -- Array D Seq (Sz (2 :> 3 :. 4))
  --   [ [ [ 0, 1, 2, 3 ]
  --     , [ 4, 5, 6, 7 ]
  --     , [ 8, 9, 10, 11 ]
  --     ]
  --   , [ [ 12, 13, 14, 15 ]
  --     , [ 16, 17, 18, 19 ]
  --     , [ 20, 21, 22, 23 ]
  --     ]
  --   ]
  -- >>> backpermuteM @U (Sz (4 :. 2)) (\(i :. j) -> j :> j :. i) arr
  -- Array U Seq (Sz (4 :. 2))
  --   [ [ 0, 16 ]
  --   , [ 1, 17 ]
  --   , [ 2, 18 ]
  --   , [ 3, 19 ]
  --   ]
  --
  -- @since 0.3.0
  backpermuteM
    :: forall r ix e r' ix' m
     . (Manifest r e, Index ix, Source r' e, Index ix', MonadUnliftIO m, PrimMonad m, MonadThrow m)
    => Sz ix
    -- ^ Size of the result array
    -> (ix -> ix')
    -- ^ A function that maps indices of the new array into the source one.
    -> Array r' ix' e
    -- ^ Source array.
    -> m (Array r ix e)
  backpermuteM sz ixF !arr = generateArray (getComp arr) sz (evaluateM arr . ixF)
  {-# INLINE backpermuteM #-}
  
  -- | Similar to `backpermuteM`, with a few notable differences:
  --
  -- * Creates a delayed array, instead of manifest, therefore it can be fused
  -- * Respects computation strategy, so it can be parallelized
  -- * Throws a runtime `IndexOutOfBoundsException` from pure code.
  --
  -- @since 0.3.0
  backpermute'
    :: forall r ix ix' e
     . (HasCallStack, Source r e, Index ix, Index ix')
    => Sz ix'
    -- ^ Size of the result array
    -> (ix' -> ix)
    -- ^ A function that maps indices of the new array into the source one.
    -> Array r ix e
    -- ^ Source array.
    -> Array D ix' e
  backpermute' sz ixF !arr = makeArray (getComp arr) sz (evaluate' arr . ixF)
  {-# INLINE backpermute' #-}
  
  -- | Append two arrays together along a particular dimension. Sizes of both arrays must match, with
  -- an allowed exception of the dimension they are being appended along, otherwise `Nothing` is
  -- returned.
  --
  -- ====__Examples__
  --
  -- Append two 2D arrays along both dimensions. Note that they do agree on inner dimensions.
  --
  -- >>> import Data.Massiv.Array
  -- >>> arrA = makeArrayR U Seq (Sz2 2 3) (\(i :. j) -> ('A', i, j))
  -- >>> arrB = makeArrayR U Seq (Sz2 2 3) (\(i :. j) -> ('B', i, j))
  -- >>> appendM 1 arrA arrB
  -- Array DL Seq (Sz (2 :. 6))
  --   [ [ ('A',0,0), ('A',0,1), ('A',0,2), ('B',0,0), ('B',0,1), ('B',0,2) ]
  --   , [ ('A',1,0), ('A',1,1), ('A',1,2), ('B',1,0), ('B',1,1), ('B',1,2) ]
  --   ]
  -- >>> appendM 2 arrA arrB
  -- Array DL Seq (Sz (4 :. 3))
  --   [ [ ('A',0,0), ('A',0,1), ('A',0,2) ]
  --   , [ ('A',1,0), ('A',1,1), ('A',1,2) ]
  --   , [ ('B',0,0), ('B',0,1), ('B',0,2) ]
  --   , [ ('B',1,0), ('B',1,1), ('B',1,2) ]
  --   ]
  --
  -- Now appending arrays with different sizes:
  --
  -- >>> arrC = makeArrayR U Seq (Sz (2 :. 4)) (\(i :. j) -> ('C', i, j))
  -- >>> appendM 1 arrA arrC
  -- Array DL Seq (Sz (2 :. 7))
  --   [ [ ('A',0,0), ('A',0,1), ('A',0,2), ('C',0,0), ('C',0,1), ('C',0,2), ('C',0,3) ]
  --   , [ ('A',1,0), ('A',1,1), ('A',1,2), ('C',1,0), ('C',1,1), ('C',1,2), ('C',1,3) ]
  --   ]
  -- >>> appendM 2 arrA arrC
  -- *** Exception: SizeMismatchException: (Sz (2 :. 3)) vs (Sz (2 :. 4))
  -- ...
  --
  -- @since 0.3.0
  appendM
    :: forall r1 r2 ix e m
     . (MonadThrow m, Index ix, Source r1 e, Source r2 e)
    => Dim
    -> Array r1 ix e
    -> Array r2 ix e
    -> m (Array DL ix e)
  appendM n !arr1 !arr2 = do
    let !sz1 = size arr1
        !sz2 = size arr2
    (k1, szl1) <- pullOutSzM sz1 n
    (k2, szl2) <- pullOutSzM sz2 n
    unless (szl1 == szl2) $ throwM $ SizeMismatchException sz1 sz2
    let !k1' = unSz k1
    newSz <- insertSzM szl1 n (SafeSz (k1' + unSz k2))
    let load :: Loader e
        load scheduler !startAt dlWrite _dlSet = do
          scheduleWork scheduler $
            iterA_ zeroIndex (unSz sz1) (pureIndex 1) (<) $ \ix ->
              dlWrite (startAt + toLinearIndex newSz ix) (unsafeIndex arr1 ix)
          scheduleWork scheduler $
            iterA_ zeroIndex (unSz sz2) (pureIndex 1) (<) $ \ix ->
              let i = getDim' ix n
                  ix' = setDim' ix n (i + k1')
               in dlWrite (startAt + toLinearIndex newSz ix') (unsafeIndex arr2 ix)
        {-# INLINE load #-}
    return $
      DLArray
        { dlComp = getComp arr1 <> getComp arr2
        , dlSize = newSz
        , dlLoad = load
        }
  {-# INLINE appendM #-}
  
  -- | Same as `appendM`, but will throw an exception in pure code on mismatched sizes.
  --
  -- @since 0.3.0
  append'
    :: forall r1 r2 ix e
     . (HasCallStack, Index ix, Source r1 e, Source r2 e)
    => Dim
    -> Array r1 ix e
    -> Array r2 ix e
    -> Array DL ix e
  append' dim arr1 arr2 = throwEither $ appendM dim arr1 arr2
  {-# INLINE append' #-}
  
  -- | Concat many arrays together along some dimension.
  --
  -- @since 0.3.0
  concat'
    :: forall f r ix e
     . (HasCallStack, Foldable f, Index ix, Source r e)
    => Dim
    -> f (Array r ix e)
    -> Array DL ix e
  concat' n = throwEither . concatM n
  {-# INLINE concat' #-}
  
  -- | Concatenate many arrays together along some dimension. It is important that all sizes are
  -- equal, with an exception of the dimensions along which concatenation happens.
  --
  -- /__Exceptions__/: `IndexDimensionException`, `SizeMismatchException`
  --
  -- @since 0.3.0
  concatM
    :: forall r ix e f m
     . (MonadThrow m, Foldable f, Index ix, Source r e)
    => Dim
    -> f (Array r ix e)
    -> m (Array DL ix e)
  concatM n arrsF =
    case L.uncons (F.toList arrsF) of
      Nothing -> pure empty
      Just (a, arrs) -> do
        let sz = unSz (size a)
            szs = unSz . size <$> arrs
        (k, szl) <- pullOutDimM sz n
        -- / remove the dimension out of all sizes along which concatenation will happen
        (ks, szls) <-
          F.foldrM (\ !csz (ks, szls) -> bimap (: ks) (: szls) <$> pullOutDimM csz n) ([], []) szs
        -- / make sure to fail as soon as at least one of the arrays has a mismatching inner size
        traverse_
          (\(sz', _) -> throwM (SizeMismatchException (SafeSz sz) (SafeSz sz')))
          (dropWhile ((== szl) . snd) $ P.zip szs szls)
        let kTotal = SafeSz $ F.foldl' (+) k ks
        newSz <- insertSzM (SafeSz szl) n kTotal
        let load :: Loader e
            load scheduler startAt dlWrite _dlSet =
              let arrayLoader !kAcc (!kCur, arr) = do
                    scheduleWork scheduler $
                      iforM_ arr $ \ix e -> do
                        i <- getDimM ix n
                        ix' <- setDimM ix n (i + kAcc)
                        dlWrite (startAt + toLinearIndex newSz ix') e
                    pure $! kAcc + kCur
                  {-# INLINE arrayLoader #-}
               in M.foldM_ arrayLoader 0 $ (k, a) : P.zip ks arrs
            {-# INLINE load #-}
        return $
          DLArray{dlComp = getComp a <> foldMap getComp arrs, dlSize = newSz, dlLoad = load}
  {-# INLINE concatM #-}
  
  -- | Stack slices on top of each other along the specified dimension.
  --
  -- /__Exceptions__/: `IndexDimensionException`, `SizeMismatchException`
  --
  -- ====__Examples__
  --
  -- Here are the three different ways to stack up two 2D Matrix pages into a 3D array.
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> x = compute (iterateN 3 succ 0) :: Matrix P Int
  -- >>> y = compute (iterateN 3 succ 9) :: Matrix P Int
  -- >>> x
  -- Array P Seq (Sz (3 :. 3))
  --   [ [ 1, 2, 3 ]
  --   , [ 4, 5, 6 ]
  --   , [ 7, 8, 9 ]
  --   ]
  -- >>> y
  -- Array P Seq (Sz (3 :. 3))
  --   [ [ 10, 11, 12 ]
  --   , [ 13, 14, 15 ]
  --   , [ 16, 17, 18 ]
  --   ]
  -- >>> stackSlicesM 1 [x, y] :: IO (Array DL Ix3 Int)
  -- Array DL Seq (Sz (3 :> 3 :. 2))
  --   [ [ [ 1, 10 ]
  --     , [ 2, 11 ]
  --     , [ 3, 12 ]
  --     ]
  --   , [ [ 4, 13 ]
  --     , [ 5, 14 ]
  --     , [ 6, 15 ]
  --     ]
  --   , [ [ 7, 16 ]
  --     , [ 8, 17 ]
  --     , [ 9, 18 ]
  --     ]
  --   ]
  -- >>> stackSlicesM 2 [x, y] :: IO (Array DL Ix3 Int)
  -- Array DL Seq (Sz (3 :> 2 :. 3))
  --   [ [ [ 1, 2, 3 ]
  --     , [ 10, 11, 12 ]
  --     ]
  --   , [ [ 4, 5, 6 ]
  --     , [ 13, 14, 15 ]
  --     ]
  --   , [ [ 7, 8, 9 ]
  --     , [ 16, 17, 18 ]
  --     ]
  --   ]
  -- >>> stackSlicesM 3 [x, y] :: IO (Array DL Ix3 Int)
  -- Array DL Seq (Sz (2 :> 3 :. 3))
  --   [ [ [ 1, 2, 3 ]
  --     , [ 4, 5, 6 ]
  --     , [ 7, 8, 9 ]
  --     ]
  --   , [ [ 10, 11, 12 ]
  --     , [ 13, 14, 15 ]
  --     , [ 16, 17, 18 ]
  --     ]
  --   ]
  --
  -- @since 0.5.4
  stackSlicesM
    :: forall r ix e f m
     . (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix)
    => Dim
    -> f (Array r (Lower ix) e)
    -> m (Array DL ix e)
  stackSlicesM dim !arrsF = do
    case L.uncons (F.toList arrsF) of
      Nothing -> pure empty
      Just (a, arrs) -> do
        let sz = size a
            len = SafeSz (F.length arrsF)
        -- / make sure all arrays have the same size
        M.forM_ arrsF $ \arr ->
          unless (sz == size arr) $ throwM (SizeMismatchException sz (size arr))
        newSz <- insertSzM sz dim len
        let load :: Loader e
            load scheduler startAt dlWrite _dlSet =
              let loadIndex k ix = dlWrite (toLinearIndex newSz (insertDim' ix dim k) + startAt)
                  arrayLoader !k arr = (k + 1) <$ scheduleWork scheduler (imapM_ (loadIndex k) arr)
                  {-# INLINE arrayLoader #-}
               in M.foldM_ arrayLoader 0 arrsF
            {-# INLINE load #-}
        return $
          DLArray{dlComp = foldMap getComp arrs, dlSize = newSz, dlLoad = load}
  {-# INLINE stackSlicesM #-}
  
  -- | Specialized `stackSlicesM` to handling stacking from the outside. It is the inverse of
  -- `Data.Massiv.Array.outerSlices`.
  --
  -- /__Exceptions__/: `SizeMismatchException`
  --
  -- ====__Examples__
  --
  -- In this example we stack vectors as row of a matrix from top to bottom:
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> x = compute (iterateN 3 succ 0) :: Matrix P Int
  -- >>> x
  -- Array P Seq (Sz (3 :. 3))
  --   [ [ 1, 2, 3 ]
  --   , [ 4, 5, 6 ]
  --   , [ 7, 8, 9 ]
  --   ]
  -- >>> rows = outerSlices x
  -- >>> A.mapM_ print rows
  -- Array P Seq (Sz1 3)
  --   [ 1, 2, 3 ]
  -- Array P Seq (Sz1 3)
  --   [ 4, 5, 6 ]
  -- Array P Seq (Sz1 3)
  --   [ 7, 8, 9 ]
  -- >>> stackOuterSlicesM rows :: IO (Matrix DL Int)
  -- Array DL Seq (Sz (3 :. 3))
  --   [ [ 1, 2, 3 ]
  --   , [ 4, 5, 6 ]
  --   , [ 7, 8, 9 ]
  --   ]
  --
  -- @since 0.5.4
  stackOuterSlicesM
    :: forall r ix e f m
     . (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix)
    => f (Array r (Lower ix) e)
    -> m (Array DL ix e)
  stackOuterSlicesM = stackSlicesM (dimensions (Proxy :: Proxy ix))
  {-# INLINE stackOuterSlicesM #-}
  
  -- | Specialized `stackSlicesM` to handling stacking from the inside. It is the inverse of
  -- `Data.Massiv.Array.innerSlices`.
  --
  -- /__Exceptions__/: `SizeMismatchException`
  --
  -- ====__Examples__
  --
  -- In this example we stack vectors as columns of a matrix from left to right:
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> x = compute (iterateN 3 succ 0) :: Matrix P Int
  -- >>> x
  -- Array P Seq (Sz (3 :. 3))
  --   [ [ 1, 2, 3 ]
  --   , [ 4, 5, 6 ]
  --   , [ 7, 8, 9 ]
  --   ]
  -- >>> columns = innerSlices x
  -- >>> A.mapM_ print columns
  -- Array D Seq (Sz1 3)
  --   [ 1, 4, 7 ]
  -- Array D Seq (Sz1 3)
  --   [ 2, 5, 8 ]
  -- Array D Seq (Sz1 3)
  --   [ 3, 6, 9 ]
  -- >>> stackInnerSlicesM columns :: IO (Matrix DL Int)
  -- Array DL Seq (Sz (3 :. 3))
  --   [ [ 1, 2, 3 ]
  --   , [ 4, 5, 6 ]
  --   , [ 7, 8, 9 ]
  --   ]
  --
  -- @since 0.5.4
  stackInnerSlicesM
    :: forall r ix e f m
     . (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix)
    => f (Array r (Lower ix) e)
    -> m (Array DL ix e)
  stackInnerSlicesM = stackSlicesM 1
  {-# INLINE stackInnerSlicesM #-}
  
  -- | /O(1)/ - Split an array into two at an index along a specified dimension.
  --
  -- /Related/: 'splitAt'', `splitExtractM`, 'Data.Massiv.Vector.sliceAt'', `Data.Massiv.Vector.sliceAtM`
  --
  -- /__Exceptions__/: `IndexDimensionException`, `SizeSubregionException`
  --
  -- @since 0.3.0
  splitAtM
    :: forall r ix e m
     . (MonadThrow m, Index ix, Source r e)
    => Dim
    -- ^ Dimension along which to split
    -> Int
    -- ^ Index along the dimension to split at
    -> Array r ix e
    -- ^ Source array
    -> m (Array D ix e, Array D ix e)
  splitAtM dim i arr = do
    let Sz sz = size arr
    eIx <- setDimM sz dim i
    sIx <- setDimM zeroIndex dim i
    arr1 <- extractFromToM zeroIndex eIx arr
    arr2 <- extractFromToM sIx sz arr
    return (arr1, arr2)
  {-# INLINE splitAtM #-}
  
  -- | /O(1)/ - Split an array into two at an index along a specified dimension. Throws an
  -- error for a wrong dimension or incorrect indices.
  --
  -- /Related/: `splitAtM`, `splitExtractM`, 'Data.Massiv.Vector.sliceAt'', `Data.Massiv.Vector.sliceAtM`
  --
  -- ==== __Examples__
  --
  --
  -- @since 0.1.0
  splitAt'
    :: forall r ix e
     . (HasCallStack, Index ix, Source r e)
    => Dim
    -> Int
    -> Array r ix e
    -> (Array D ix e, Array D ix e)
  splitAt' dim i = throwEither . splitAtM dim i
  {-# INLINE splitAt' #-}
  
  -- | Split an array in three parts across some dimension
  --
  -- @since 0.3.5
  splitExtractM
    :: forall r ix e m
     . (MonadThrow m, Index ix, Source r e)
    => Dim
    -- ^ Dimension along which to do the extraction
    -> Ix1
    -- ^ Start index along the dimension that needs to be extracted
    -> Sz Ix1
    -- ^ Size of the extracted array along the dimension that it will be extracted
    -> Array r ix e
    -> m (Array D ix e, Array D ix e, Array D ix e)
  splitExtractM dim startIx1 (Sz extractSzIx1) arr = do
    let Sz szIx = size arr
    midStartIx <- setDimM zeroIndex dim startIx1
    midExtractSzIx <- setDimM szIx dim extractSzIx1
    midArr <- extractM midStartIx (Sz midExtractSzIx) arr
    leftArrSzIx <- setDimM szIx dim startIx1
    leftArr <- extractM zeroIndex (Sz leftArrSzIx) arr
    rightArrStartIx <- setDimM zeroIndex dim (startIx1 + extractSzIx1)
    rightArr <- extractFromToM rightArrStartIx szIx arr
    pure (leftArr, midArr, rightArr)
  {-# INLINE splitExtractM #-}
  
  -- | Replace a slice of an array with another one
  --
  -- ==== __Example__
  --
  -- >>> import Data.Massiv.Array
  -- >>> arr = makeArrayR U Seq (Sz3 3 4 5) fromIx3
  -- >>> arr' = makeArrayR U Seq (Sz3 3 4 5) (fromIx3 . liftIndex (* 100))
  -- >>> replaceSlice 2 1 (arr' <!> (2, 3)) arr
  -- Array DL Seq (Sz (3 :> 4 :. 5))
  --   [ [ [ (0,0,0), (0,0,1), (0,0,2), (0,0,3), (0,0,4) ]
  --     , [ (0,300,0), (0,300,100), (0,300,200), (0,300,300), (0,300,400) ]
  --     , [ (0,2,0), (0,2,1), (0,2,2), (0,2,3), (0,2,4) ]
  --     , [ (0,3,0), (0,3,1), (0,3,2), (0,3,3), (0,3,4) ]
  --     ]
  --   , [ [ (1,0,0), (1,0,1), (1,0,2), (1,0,3), (1,0,4) ]
  --     , [ (100,300,0), (100,300,100), (100,300,200), (100,300,300), (100,300,400) ]
  --     , [ (1,2,0), (1,2,1), (1,2,2), (1,2,3), (1,2,4) ]
  --     , [ (1,3,0), (1,3,1), (1,3,2), (1,3,3), (1,3,4) ]
  --     ]
  --   , [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3), (2,0,4) ]
  --     , [ (200,300,0), (200,300,100), (200,300,200), (200,300,300), (200,300,400) ]
  --     , [ (2,2,0), (2,2,1), (2,2,2), (2,2,3), (2,2,4) ]
  --     , [ (2,3,0), (2,3,1), (2,3,2), (2,3,3), (2,3,4) ]
  --     ]
  --   ]
  --
  -- @since 0.6.1
  replaceSlice
    :: forall r r' ix e m
     . (MonadThrow m, Source r e, Source r' e, Index ix, Index (Lower ix))
    => Dim
    -> Ix1
    -> Array r' (Lower ix) e
    -> Array r ix e
    -> m (Array DL ix e)
  replaceSlice dim i sl arr = do
    (l, m, r) <- splitExtractM dim i (SafeSz 1) arr
    m' <- resizeM (size m) sl
    concatM dim [l, delay m', r]
  {-# INLINE replaceSlice #-}
  
  -- | Replace an outer slice of an array with another one
  --
  -- ==== __Example__
  --
  -- >>> import Data.Massiv.Array
  -- >>> arr = makeArrayR U Seq (Sz3 3 4 5) fromIx3
  -- >>> arr' = makeArrayR U Seq (Sz3 3 4 5) (fromIx3 . liftIndex (* 100))
  -- >>> replaceOuterSlice 1 (arr' !> 2) arr
  -- Array DL Seq (Sz (3 :> 4 :. 5))
  --   [ [ [ (0,0,0), (0,0,1), (0,0,2), (0,0,3), (0,0,4) ]
  --     , [ (0,1,0), (0,1,1), (0,1,2), (0,1,3), (0,1,4) ]
  --     , [ (0,2,0), (0,2,1), (0,2,2), (0,2,3), (0,2,4) ]
  --     , [ (0,3,0), (0,3,1), (0,3,2), (0,3,3), (0,3,4) ]
  --     ]
  --   , [ [ (200,0,0), (200,0,100), (200,0,200), (200,0,300), (200,0,400) ]
  --     , [ (200,100,0), (200,100,100), (200,100,200), (200,100,300), (200,100,400) ]
  --     , [ (200,200,0), (200,200,100), (200,200,200), (200,200,300), (200,200,400) ]
  --     , [ (200,300,0), (200,300,100), (200,300,200), (200,300,300), (200,300,400) ]
  --     ]
  --   , [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3), (2,0,4) ]
  --     , [ (2,1,0), (2,1,1), (2,1,2), (2,1,3), (2,1,4) ]
  --     , [ (2,2,0), (2,2,1), (2,2,2), (2,2,3), (2,2,4) ]
  --     , [ (2,3,0), (2,3,1), (2,3,2), (2,3,3), (2,3,4) ]
  --     ]
  --   ]
  --
  -- @since 0.6.1
  replaceOuterSlice
    :: forall r ix e m
     . (MonadThrow m, Index ix, Source r e, Load r (Lower ix) e)
    => Ix1
    -> Array r (Lower ix) e
    -> Array r ix e
    -> m (Array DL ix e)
  replaceOuterSlice i sl arr = replaceSlice (dimensions (size arr)) i sl arr
  {-# INLINE replaceOuterSlice #-}
  
  -- | Delete a region from an array along the specified dimension.
  --
  -- ==== __Examples__
  --
  -- >>> import Data.Massiv.Array
  -- >>> arr = fromIx3 <$> (0 :> 0 :. 0 ..: 3 :> 2 :. 6)
  -- >>> deleteRegionM 1 2 3 arr
  -- Array DL Seq (Sz (3 :> 2 :. 3))
  --   [ [ [ (0,0,0), (0,0,1), (0,0,5) ]
  --     , [ (0,1,0), (0,1,1), (0,1,5) ]
  --     ]
  --   , [ [ (1,0,0), (1,0,1), (1,0,5) ]
  --     , [ (1,1,0), (1,1,1), (1,1,5) ]
  --     ]
  --   , [ [ (2,0,0), (2,0,1), (2,0,5) ]
  --     , [ (2,1,0), (2,1,1), (2,1,5) ]
  --     ]
  --   ]
  -- >>> v = Ix1 0 ... 10
  -- >>> deleteRegionM 1 3 5 v
  -- Array DL Seq (Sz1 6)
  --   [ 0, 1, 2, 8, 9, 10 ]
  --
  -- @since 0.3.5
  deleteRegionM
    :: forall r ix e m
     . (MonadThrow m, Index ix, Source r e)
    => Dim
    -- ^ Along which axis should the removal happen
    -> Ix1
    -- ^ At which index to start dropping slices
    -> Sz Ix1
    -- ^ Number of slices to drop
    -> Array r ix e
    -- ^ Array that will have it's subarray removed
    -> m (Array DL ix e)
  deleteRegionM dim ix sz arr = do
    (leftArr, _, rightArr) <- splitExtractM dim ix sz arr
    appendM dim leftArr rightArr
  {-# INLINE deleteRegionM #-}
  
  -- | Similar to `deleteRegionM`, but drop a specified number of rows from an array that
  -- has at least 2 dimensions.
  --
  -- ====__Example__
  --
  -- >>> import Data.Massiv.Array
  -- >>> arr = fromIx2 <$> (0 :. 0 ..: 3 :. 6)
  -- >>> arr
  -- Array D Seq (Sz (3 :. 6))
  --   [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ]
  --   , [ (1,0), (1,1), (1,2), (1,3), (1,4), (1,5) ]
  --   , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ]
  --   ]
  -- >>> deleteRowsM 1 1 arr
  -- Array DL Seq (Sz (2 :. 6))
  --   [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ]
  --   , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ]
  --   ]
  --
  -- @since 0.3.5
  deleteRowsM
    :: forall r ix e m
     . (MonadThrow m, Index ix, Index (Lower ix), Source r e)
    => Ix1
    -> Sz Ix1
    -> Array r ix e
    -> m (Array DL ix e)
  deleteRowsM = deleteRegionM 2
  {-# INLINE deleteRowsM #-}
  
  -- | Similar to `deleteRegionM`, but drop a specified number of columns an array.
  --
  -- ====__Example__
  --
  -- >>> import Data.Massiv.Array
  -- >>> arr = fromIx2 <$> (0 :. 0 ..: 3 :. 6)
  -- >>> arr
  -- Array D Seq (Sz (3 :. 6))
  --   [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ]
  --   , [ (1,0), (1,1), (1,2), (1,3), (1,4), (1,5) ]
  --   , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ]
  --   ]
  -- >>> deleteColumnsM 2 3 arr
  -- Array DL Seq (Sz (3 :. 3))
  --   [ [ (0,0), (0,1), (0,5) ]
  --   , [ (1,0), (1,1), (1,5) ]
  --   , [ (2,0), (2,1), (2,5) ]
  --   ]
  --
  -- @since 0.3.5
  deleteColumnsM
    :: forall r ix e m
     . (MonadThrow m, Index ix, Source r e)
    => Ix1
    -> Sz Ix1
    -> Array r ix e
    -> m (Array DL ix e)
  deleteColumnsM = deleteRegionM 1
  {-# INLINE deleteColumnsM #-}
  
  -- | Discard elements from the source array according to the stride.
  --
  -- @since 0.3.0
  downsample
    :: forall r ix e
     . (Source r e, Load r ix e)
    => Stride ix
    -> Array r ix e
    -> Array DL ix e
  downsample stride arr =
    DLArray{dlComp = getComp arr, dlSize = resultSize, dlLoad = load}
    where
      resultSize = strideSize stride (size arr)
      strideIx = unStride stride
      unsafeLinearWriteWithStride =
        unsafeIndex arr . liftIndex2 (*) strideIx . fromLinearIndex resultSize
      {-# INLINE unsafeLinearWriteWithStride #-}
      load :: Loader e
      load scheduler startAt dlWrite _ =
        splitLinearlyWithStartAtM_
          scheduler
          startAt
          (totalElem resultSize)
          (pure . unsafeLinearWriteWithStride)
          dlWrite
      {-# INLINE load #-}
  {-# INLINE downsample #-}
  
  -- | Insert the same element into a `Load`able array according to the supplied stride.
  --
  -- ====__Examples__
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> arr = iterateN (Sz2 3 2) succ (0 :: Int)
  -- >>> arr
  -- Array DL Seq (Sz (3 :. 2))
  --   [ [ 1, 2 ]
  --   , [ 3, 4 ]
  --   , [ 5, 6 ]
  --   ]
  -- >>> upsample 0 (Stride (2 :. 3)) arr
  -- Array DL Seq (Sz (6 :. 6))
  --   [ [ 1, 0, 0, 2, 0, 0 ]
  --   , [ 0, 0, 0, 0, 0, 0 ]
  --   , [ 3, 0, 0, 4, 0, 0 ]
  --   , [ 0, 0, 0, 0, 0, 0 ]
  --   , [ 5, 0, 0, 6, 0, 0 ]
  --   , [ 0, 0, 0, 0, 0, 0 ]
  --   ]
  -- >>> upsample 9 (Stride (1 :. 2)) arr
  -- Array DL Seq (Sz (3 :. 4))
  --   [ [ 1, 9, 2, 9 ]
  --   , [ 3, 9, 4, 9 ]
  --   , [ 5, 9, 6, 9 ]
  --   ]
  --
  -- @since 0.3.0
  upsample
    :: forall r ix e
     . Load r ix e
    => e
    -- ^ Element to use for filling the newly added cells
    -> Stride ix
    -- ^ Fill cells according to this stride
    -> Array r ix e
    -- ^ Array that will have cells added to
    -> Array DL ix e
  upsample !fillWith safeStride arr =
    DLArray
      { dlComp = getComp arr
      , dlSize = newsz
      , dlLoad = load
      }
    where
      load :: Loader e
      load scheduler startAt uWrite uSet = do
        uSet startAt (toLinearSz newsz) fillWith
        iterArrayLinearST_ scheduler arr (\i -> uWrite (adjustLinearStride (i + startAt)))
      {-# INLINE load #-}
      adjustLinearStride = toLinearIndex newsz . timesStride . fromLinearIndex sz
      {-# INLINE adjustLinearStride #-}
      timesStride !ix = liftIndex2 (*) stride ix
      {-# INLINE timesStride #-}
      !stride = unStride safeStride
      ~sz = outerSize arr -- intentionally lazy in case it is used with DS
      !newsz = SafeSz (timesStride $ unSz sz)
  {-# INLINE upsample #-}
  
  -- | General array transformation, that forces computation and produces a manifest array.
  --
  -- @since 0.3.0
  transformM
    :: forall r ix e r' ix' e' a m
     . (Manifest r e, Index ix, Source r' e', Index ix', MonadUnliftIO m, PrimMonad m, MonadThrow m)
    => (Sz ix' -> m (Sz ix, a))
    -> (a -> (ix' -> m e') -> ix -> m e)
    -> Array r' ix' e'
    -> m (Array r ix e)
  transformM getSzM getM arr = do
    (sz, a) <- getSzM (size arr)
    generateArray (getComp arr) sz (getM a (evaluateM arr))
  {-# INLINE transformM #-}
  
  -- | General array transformation
  --
  -- @since 0.3.0
  transform'
    :: forall ix e r' ix' e' a
     . (HasCallStack, Source r' e', Index ix', Index ix)
    => (Sz ix' -> (Sz ix, a))
    -> (a -> (ix' -> e') -> ix -> e)
    -> Array r' ix' e'
    -> Array D ix e
  transform' getSz get arr = makeArray (getComp arr) sz (get a (evaluate' arr))
    where
      (sz, a) = getSz (size arr)
  {-# INLINE transform' #-}
  
  -- | Same as `transformM`, but operates on two arrays
  --
  -- @since 0.3.0
  transform2M
    :: ( Manifest r e
       , Index ix
       , Source r1 e1
       , Source r2 e2
       , Index ix1
       , Index ix2
       , MonadUnliftIO m
       , PrimMonad m
       , MonadThrow m
       )
    => (Sz ix1 -> Sz ix2 -> m (Sz ix, a))
    -> (a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e)
    -> Array r1 ix1 e1
    -> Array r2 ix2 e2
    -> m (Array r ix e)
  transform2M getSzM getM arr1 arr2 = do
    (sz, a) <- getSzM (size arr1) (size arr2)
    generateArray (getComp arr1 <> getComp arr2) sz (getM a (evaluateM arr1) (evaluateM arr2))
  {-# INLINE transform2M #-}
  
  -- | Same as 'transform'', but operates on two arrays
  --
  -- @since 0.3.0
  transform2'
    :: (HasCallStack, Source r1 e1, Source r2 e2, Index ix, Index ix1, Index ix2)
    => (Sz ix1 -> Sz ix2 -> (Sz ix, a))
    -> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e)
    -> Array r1 ix1 e1
    -> Array r2 ix2 e2
    -> Array D ix e
  transform2' getSz get arr1 arr2 =
    makeArray (getComp arr1 <> getComp arr2) sz (get a (evaluate' arr1) (evaluate' arr2))
    where
      (sz, a) = getSz (size arr1) (size arr2)
  {-# INLINE transform2' #-}
  
  -- | Replicate each element of the array by a factor in stride along each dimension and surround each
  -- such group with a box of supplied grid value. It will essentially zoom up an array and create a
  -- grid around each element from the original array. Very useful for zooming up images to inspect
  -- individual pixels.
  --
  -- ==== __Example__
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> arr = resize' (Sz2 3 2) (Ix1 1 ... 6)
  -- >>> arr
  -- Array D Seq (Sz (3 :. 2))
  --   [ [ 1, 2 ]
  --   , [ 3, 4 ]
  --   , [ 5, 6 ]
  --   ]
  -- >>> zoomWithGrid 0 (Stride (2 :. 3)) arr
  -- Array DL Seq (Sz (10 :. 9))
  --   [ [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
  --   , [ 0, 1, 1, 1, 0, 2, 2, 2, 0 ]
  --   , [ 0, 1, 1, 1, 0, 2, 2, 2, 0 ]
  --   , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
  --   , [ 0, 3, 3, 3, 0, 4, 4, 4, 0 ]
  --   , [ 0, 3, 3, 3, 0, 4, 4, 4, 0 ]
  --   , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
  --   , [ 0, 5, 5, 5, 0, 6, 6, 6, 0 ]
  --   , [ 0, 5, 5, 5, 0, 6, 6, 6, 0 ]
  --   , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
  --   ]
  --
  -- @since 0.3.1
  zoomWithGrid
    :: forall r ix e
     . (Index ix, Source r e)
    => e
    -- ^ Value to use for the grid
    -> Stride ix
    -- ^ Scaling factor
    -> Array r ix e
    -- ^ Source array
    -> Array DL ix e
  zoomWithGrid gridVal (Stride zoomFactor) arr = unsafeMakeLoadArray Seq newSz (Just gridVal) load
    where
      !kx = liftIndex (+ 1) zoomFactor
      !lastNewIx = liftIndex2 (*) kx $ unSz (size arr)
      !newSz = Sz (liftIndex (+ 1) lastNewIx)
      load :: forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ()
      load scheduler _ writeElement =
        iforSchedulerM_ scheduler arr $ \ !ix !e ->
          let !kix = liftIndex2 (*) ix kx
           in mapM_ (\ !ix' -> writeElement (toLinearIndex newSz ix') e) $
                range Seq (liftIndex (+ 1) kix) (liftIndex2 (+) kix kx)
      {-# INLINE load #-}
  {-# INLINE zoomWithGrid #-}
  
  -- | Increaze the size of the array accoridng to the stride multiplier while replicating
  -- the same element to fill the neighbors. It is exactly the same as `zoomWithGrid`, but
  -- without the grid.
  --
  -- ==== __Example__
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> arr = resize' (Sz3 1 3 2) (Ix1 1 ... 6)
  -- >>> arr
  -- Array D Seq (Sz (1 :> 3 :. 2))
  --   [ [ [ 1, 2 ]
  --     , [ 3, 4 ]
  --     , [ 5, 6 ]
  --     ]
  --   ]
  -- >>> zoom (Stride (2 :> 2 :. 3)) arr
  -- Array DL Seq (Sz (2 :> 6 :. 6))
  --   [ [ [ 1, 1, 1, 2, 2, 2 ]
  --     , [ 1, 1, 1, 2, 2, 2 ]
  --     , [ 3, 3, 3, 4, 4, 4 ]
  --     , [ 3, 3, 3, 4, 4, 4 ]
  --     , [ 5, 5, 5, 6, 6, 6 ]
  --     , [ 5, 5, 5, 6, 6, 6 ]
  --     ]
  --   , [ [ 1, 1, 1, 2, 2, 2 ]
  --     , [ 1, 1, 1, 2, 2, 2 ]
  --     , [ 3, 3, 3, 4, 4, 4 ]
  --     , [ 3, 3, 3, 4, 4, 4 ]
  --     , [ 5, 5, 5, 6, 6, 6 ]
  --     , [ 5, 5, 5, 6, 6, 6 ]
  --     ]
  --   ]
  --
  -- @since 0.4.4
  zoom
    :: forall r ix e
     . (Index ix, Source r e)
    => Stride ix
    -- ^ Scaling factor
    -> Array r ix e
    -- ^ Source array
    -> Array DL ix e
  zoom (Stride zoomFactor) arr = unsafeMakeLoadArray Seq newSz Nothing load
    where
      !lastNewIx = liftIndex2 (*) zoomFactor $ unSz (size arr)
      !newSz = Sz lastNewIx
      load :: forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ()
      load scheduler _ writeElement =
        iforSchedulerM_ scheduler arr $ \ !ix !e ->
          let !kix = liftIndex2 (*) ix zoomFactor
           in mapM_ (\ !ix' -> writeElement (toLinearIndex newSz ix') e) $
                range Seq kix (liftIndex2 (+) kix zoomFactor)
      {-# INLINE load #-}
  {-# INLINE zoom #-}
  