content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  -- |
  -- Module      : Data.Massiv.Array.Delayed.Windowed
  -- Copyright   : (c) Alexey Kuleshevich 2018-2025
  -- License     : BSD3
  -- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
  -- Stability   : experimental
  -- Portability : non-portable
  module Data.Massiv.Array.Delayed.Windowed (
    DW (..),
    Array (..),
    Window (..),
    insertWindow,
    getWindow,
    dropWindow,
    makeWindowedArray,
  ) where
  
  import Control.Monad (when)
  import Control.Scheduler (trivialScheduler_)
  import Data.Massiv.Array.Delayed.Pull
  import Data.Massiv.Array.Manifest.Boxed
  import Data.Massiv.Array.Manifest.Internal
  import Data.Massiv.Core
  import Data.Massiv.Core.Common
  import Data.Massiv.Core.List (showArrayList, showsArrayPrec)
  import Data.Maybe (fromMaybe)
  import GHC.Base (modInt)
  import GHC.TypeLits
  
  -- | Delayed Windowed Array representation.
  data DW = DW
  
  data Window ix e = Window
    { windowStart :: !ix
    -- ^ Index of where window will start at.
    , windowSize :: !(Sz ix)
    -- ^ Size of the window
    , windowIndex :: ix -> e
    -- ^ Indexing function for the window
    , windowUnrollIx2 :: !(Maybe Int)
    -- ^ Setting this value during stencil application improves cache
    -- utilization by unrolling the loop for Ix2 and higher dimensions.
    -- Has no affect on arrays with one dimension.
    }
  
  instance Functor (Window ix) where
    fmap f arr@Window{windowIndex} = arr{windowIndex = f . windowIndex}
  
  data instance Array DW ix e = DWArray
    { dwComp :: !Comp
    , dwSize :: !(Sz ix)
    , dwIndex :: ix -> e
    , dwWindow :: !(Maybe (Window ix e))
    }
  
  instance (Ragged L ix e, Load DW ix e, Show e) => Show (Array DW ix e) where
    showsPrec = showsArrayPrec (computeAs B)
    showList = showArrayList
  
  instance Strategy DW where
    setComp c arr = arr{dwComp = c}
    {-# INLINE setComp #-}
    getComp = dwComp
    {-# INLINE getComp #-}
    repr = DW
  
  instance Functor (Array DW ix) where
    fmap f arr@DWArray{dwIndex, dwWindow} =
      arr
        { dwIndex = f . dwIndex
        , dwWindow = fmap f <$> dwWindow
        }
    {-# INLINE fmap #-}
  
  --
  --
  -- @since 0.3.0
  -- _makeWindowedArrayM
  --   :: Source r ix e
  --   => Array r ix e -- ^ Source array that will have a window inserted into it
  --   -> ix -- ^ Start index for the window
  --   -> Sz ix -- ^ Size of the window
  --   -> (ix -> e) -- ^ Inside window indexing function
  --   -> Array DW ix e
  -- _makeWindowedArrayM !arr !windowStart !windowSize windowIndex
  --   | not (isSafeIndex sz windowStart) =
  --     error $
  --     "makeWindowedArray: Incorrect window starting index: (" ++
  --     show windowStart ++ ") for array size: (" ++ show (size arr) ++ ")"
  --   | totalElem windowSize == 0 =
  --     error $
  --     "makeWindowedArray: Window can't hold any elements with this size: (" ++ show windowSize ++ ")"
  --   | not
  --      (isSafeIndex
  --         (Sz (liftIndex (+ 1) (unSz sz)))
  --         (liftIndex2 (+) windowStart (unSz windowSize))) =
  --     error $
  --     "makeWindowedArray: Incorrect window size: (" ++
  --     show windowSize ++
  --     ") and/or starting index: (" ++
  --     show windowStart ++ ") for array size: (" ++ show (size arr) ++ ")"
  --   | otherwise =
  --     DWArray {dwArray = delay arr, dwWindow = Just $! Window {..}}
  --   where
  --     windowUnrollIx2 = Nothing
  --     sz = size arr
  -- {-# INLINE _makeWindowedArrayM #-}
  
  -- | Construct a delayed windowed array by supply a separate element producing function for the
  -- interior of an array. This is very usful for stencil mapping, where interior function does not
  -- perform boundary checks, thus significantly speeding up computation process.
  --
  -- @since 0.1.3
  makeWindowedArray
    :: (Index ix, Source r e)
    => Array r ix e
    -- ^ Source array that will have a window inserted into it
    -> ix
    -- ^ Start index for the window
    -> Sz ix
    -- ^ Size of the window
    -> (ix -> e)
    -- ^ Indexing function foto use inside window
    -> Array DW ix e
  makeWindowedArray !arr wStart wSize wIndex =
    insertWindow (delay arr) $
      Window{windowStart = wStart, windowSize = wSize, windowIndex = wIndex, windowUnrollIx2 = Nothing}
  {-# INLINE makeWindowedArray #-}
  
  -- | Inserts a `Window` into a delayed array while scaling the window down if it doesn't fit inside
  -- that array.
  --
  -- @since 0.3.0
  insertWindow
    :: Index ix
    => Array D ix e
    -- ^ Source array that will have a window inserted into it
    -> Window ix e
    -- ^ Window to place inside the delayed array
    -> Array DW ix e
  insertWindow !arr !window =
    DWArray
      { dwComp = getComp arr
      , dwSize = arrSize
      , dwIndex = unsafeIndex arr
      , dwWindow =
          Just $!
            Window
              { windowStart = wStart'
              , windowSize = Sz (liftIndex2 min wSize (liftIndex2 (-) sz wStart'))
              , windowIndex = wIndex
              , windowUnrollIx2 = wUnrollIx2
              }
      }
    where
      wStart' = unSz (Sz (liftIndex2 min wStart (liftIndex (subtract 1) sz)))
      arrSize@(Sz sz) = size arr
      Window
        { windowStart = wStart
        , windowSize = Sz wSize
        , windowIndex = wIndex
        , windowUnrollIx2 = wUnrollIx2
        } = window
  {-# INLINE insertWindow #-}
  
  -- | Get the `Window` from a windowed array.
  --
  -- @since 0.2.1
  getWindow :: Array DW ix e -> Maybe (Window ix e)
  getWindow = dwWindow
  {-# INLINE getWindow #-}
  
  -- | Drop the `Window` from a windowed array.
  --
  -- @since 0.3.0
  dropWindow :: Array DW ix e -> Array D ix e
  dropWindow DWArray{..} =
    DArray
      { dComp = dwComp
      , dSize = dwSize
      , dPrefIndex = PrefIndex dwIndex
      }
  {-# INLINE dropWindow #-}
  
  zeroWindow :: Index ix => Window ix e
  zeroWindow = Window zeroIndex zeroSz windowError Nothing
  {-# INLINE zeroWindow #-}
  
  data EmptyWindowException = EmptyWindowException deriving (Eq, Show)
  
  instance Exception EmptyWindowException where
    displayException _ = "Index of zero size Window"
  
  windowError :: ix -> a
  windowError _ = throwImpossible EmptyWindowException
  {-# NOINLINE windowError #-}
  
  loadWithIx1
    :: Monad m
    => (m () -> m ())
    -> Array DW Ix1 e
    -> (Ix1 -> e -> m a)
    -> m (Ix1 -> Ix1 -> m (), Ix1, Ix1)
  loadWithIx1 with (DWArray _ sz uIndex mWindow) uWrite = do
    let Window it wk uwIndex _ = fromMaybe zeroWindow mWindow
        wEnd = it + unSz wk
    with $ iterA_ 0 it 1 (<) $ \ !i -> uWrite i (uIndex i)
    with $ iterA_ wEnd (unSz sz) 1 (<) $ \ !i -> uWrite i (uIndex i)
    return (\from to -> with $ iterA_ from to 1 (<) $ \ !i -> uWrite i (uwIndex i), it, wEnd)
  {-# INLINE loadWithIx1 #-}
  
  instance Index ix => Shape DW ix where
    maxLinearSize = Just . linearSize
    {-# INLINE maxLinearSize #-}
    linearSize = SafeSz . totalElem . dwSize
    {-# INLINE linearSize #-}
    outerSize = dwSize
    {-# INLINE outerSize #-}
  
  instance Load DW Ix1 e where
    makeArray c sz f = DWArray c sz f Nothing
    {-# INLINE makeArray #-}
    iterArrayLinearST_ scheduler arr uWrite = do
      (loadWindow, wStart, wEnd) <- loadWithIx1 (scheduleWork scheduler) arr uWrite
      let (chunkWidth, slackWidth) = (wEnd - wStart) `quotRem` numWorkers scheduler
      loopA_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid ->
        let !it' = wid * chunkWidth + wStart
         in loadWindow it' (it' + chunkWidth)
      when (slackWidth > 0) $
        let !itSlack = numWorkers scheduler * chunkWidth + wStart
         in loadWindow itSlack (itSlack + slackWidth)
    {-# INLINE iterArrayLinearST_ #-}
  
  instance StrideLoad DW Ix1 e where
    iterArrayLinearWithStrideST_ scheduler stride sz arr uWrite = do
      (loadWindow, (wStart, wEnd)) <- loadArrayWithIx1 (scheduleWork scheduler) arr stride sz uWrite
      let (chunkWidth, slackWidth) = (wEnd - wStart) `quotRem` numWorkers scheduler
      loopA_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid ->
        let !it' = wid * chunkWidth + wStart
         in loadWindow (it', it' + chunkWidth)
      when (slackWidth > 0) $
        let !itSlack = numWorkers scheduler * chunkWidth + wStart
         in loadWindow (itSlack, itSlack + slackWidth)
    {-# INLINE iterArrayLinearWithStrideST_ #-}
  
  loadArrayWithIx1
    :: Monad m
    => (m () -> m ())
    -> Array DW Ix1 e
    -> Stride Ix1
    -> Sz1
    -> (Ix1 -> e -> m a)
    -> m ((Ix1, Ix1) -> m (), (Ix1, Ix1))
  loadArrayWithIx1 with (DWArray _ arrSz uIndex mWindow) stride _ uWrite = do
    let Window it wk uwIndex _ = fromMaybe zeroWindow mWindow
        wEnd = it + unSz wk
        strideIx = unStride stride
    with $ iterA_ 0 it strideIx (<) $ \ !i -> uWrite (i `div` strideIx) (uIndex i)
    with $
      iterA_ (strideStart stride wEnd) (unSz arrSz) strideIx (<) $ \ !i ->
        uWrite (i `div` strideIx) (uIndex i)
    return
      ( \(from, to) ->
          with $
            iterA_ (strideStart stride from) to strideIx (<) $ \ !i ->
              uWrite (i `div` strideIx) (uwIndex i)
      , (it, wEnd)
      )
  {-# INLINE loadArrayWithIx1 #-}
  
  loadWithIx2
    :: Monad m
    => (m () -> m ())
    -> Array DW Ix2 t1
    -> (Int -> t1 -> m ())
    -> m (Ix2 -> m (), Ix2)
  loadWithIx2 with arr uWrite = do
    let DWArray _ (Sz (m :. n)) uIndex window = arr
        Window (it :. jt) (Sz (wm :. wn)) uwIndex mUnrollHeight = fromMaybe zeroWindow window
        ib :. jb = (wm + it) :. (wn + jt)
        !blockHeight = maybe 1 (min 7 . max 1) mUnrollHeight
        !sz = strideSize oneStride $ outerSize arr
        writeB !ix = uWrite (toLinearIndex sz ix) (uIndex ix)
        {-# INLINE writeB #-}
        writeW !ix = uWrite (toLinearIndex sz ix) (uwIndex ix)
        {-# INLINE writeW #-}
    with $ iterA_ (0 :. 0) (it :. n) (1 :. 1) (<) writeB
    with $ iterA_ (ib :. 0) (m :. n) (1 :. 1) (<) writeB
    with $ iterA_ (it :. 0) (ib :. jt) (1 :. 1) (<) writeB
    with $ iterA_ (it :. jb) (ib :. n) (1 :. 1) (<) writeB
    let f (it' :. ib') = with $ unrollAndJam blockHeight (it' :. jt) (ib' :. jb) 1 writeW
        {-# INLINE f #-}
    return (f, it :. ib)
  {-# INLINE loadWithIx2 #-}
  
  loadArrayWithIx2
    :: Monad m
    => (m () -> m ())
    -> Array DW Ix2 e
    -> Stride Ix2
    -> Sz2
    -> (Int -> e -> m ())
    -> m (Ix2 -> m (), Ix2)
  loadArrayWithIx2 with arr stride sz uWrite = do
    let DWArray _ (Sz (m :. n)) uIndex window = arr
        Window (it :. jt) (Sz (wm :. wn)) uwIndex mUnrollHeight = fromMaybe zeroWindow window
        ib :. jb = (wm + it) :. (wn + jt)
        !blockHeight = maybe 1 (min 7 . max 1) mUnrollHeight
        strideIx@(is :. js) = unStride stride
        writeB !ix = uWrite (toLinearIndexStride stride sz ix) (uIndex ix)
        {-# INLINE writeB #-}
        writeW !ix = uWrite (toLinearIndexStride stride sz ix) (uwIndex ix)
        {-# INLINE writeW #-}
    with $ iterA_ (0 :. 0) (it :. n) strideIx (<) writeB
    with $ iterA_ (strideStart stride (ib :. 0)) (m :. n) strideIx (<) writeB
    with $ iterA_ (strideStart stride (it :. 0)) (ib :. jt) strideIx (<) writeB
    with $ iterA_ (strideStart stride (it :. jb)) (ib :. n) strideIx (<) writeB
    let f (it' :. ib')
          | is > 1 || blockHeight <= 1 =
              -- Turn off unrolling for vertical strides
              iterA_ (strideStart stride (it' :. jt)) (ib' :. jb) strideIx (<) writeW
          | otherwise =
              unrollAndJam blockHeight (strideStart stride (it' :. jt)) (ib' :. jb) js writeW
        {-# INLINE f #-}
    return (with . f, it :. ib)
  {-# INLINE loadArrayWithIx2 #-}
  
  loadWindowIx2 :: Monad m => Int -> (Ix2 -> m ()) -> Ix2 -> m ()
  loadWindowIx2 nWorkers loadWindow (it :. ib) = do
    let !(chunkHeight, slackHeight) = (ib - it) `quotRem` nWorkers
    loopA_ 0 (< nWorkers) (+ 1) $ \ !wid ->
      let !it' = wid * chunkHeight + it
       in loadWindow (it' :. (it' + chunkHeight))
    when (slackHeight > 0) $
      let !itSlack = nWorkers * chunkHeight + it
       in loadWindow (itSlack :. (itSlack + slackHeight))
  {-# INLINE loadWindowIx2 #-}
  
  instance Load DW Ix2 e where
    makeArray c sz f = DWArray c sz f Nothing
    {-# INLINE makeArray #-}
    iterArrayLinearST_ scheduler arr uWrite =
      loadWithIx2 (scheduleWork scheduler) arr uWrite
        >>= uncurry (loadWindowIx2 (numWorkers scheduler))
    {-# INLINE iterArrayLinearST_ #-}
  
  instance StrideLoad DW Ix2 e where
    iterArrayLinearWithStrideST_ scheduler stride sz arr uWrite =
      loadArrayWithIx2 (scheduleWork scheduler) arr stride sz uWrite
        >>= uncurry (loadWindowIx2 (numWorkers scheduler))
    {-# INLINE iterArrayLinearWithStrideST_ #-}
  
  instance (Index (IxN n), Load DW (Ix (n - 1)) e) => Load DW (IxN n) e where
    makeArray c sz f = DWArray c sz f Nothing
    {-# INLINE makeArray #-}
    iterArrayLinearST_ = loadWithIxN
    {-# INLINE iterArrayLinearST_ #-}
  
  instance (Index (IxN n), StrideLoad DW (Ix (n - 1)) e) => StrideLoad DW (IxN n) e where
    iterArrayLinearWithStrideST_ = loadArrayWithIxN
    {-# INLINE iterArrayLinearWithStrideST_ #-}
  
  loadArrayWithIxN
    :: (Index ix, StrideLoad DW (Lower ix) e)
    => Scheduler s ()
    -> Stride ix
    -> Sz ix
    -> Array DW ix e
    -> (Int -> e -> ST s ())
    -> ST s ()
  loadArrayWithIxN scheduler stride szResult arr uWrite = do
    let DWArray _ sz uIndex window = arr
        Window{windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window
        !(!headSourceSize, !lowerSourceSize) = unconsSz sz
        !lowerSize = snd $ unconsSz szResult
        !(!s, !lowerStrideIx) = unconsDim $ unStride stride
        !(!curWindowStart, lowerWindowStart) = unconsDim windowStart
        !(!headWindowSz, tailWindowSz) = unconsSz windowSize
        !curWindowEnd = curWindowStart + unSz headWindowSz
        !pageElements = totalElem lowerSize
        lowerWindow =
          Window
            { windowStart = lowerWindowStart
            , windowSize = tailWindowSz
            , windowIndex = \_ -> error "Window index uninitialized"
            , windowUnrollIx2 = windowUnrollIx2
            }
        mkLowerWindow !i =
          lowerWindow
            { windowIndex = windowIndex . consDim i
            }
        loadLower mkWindow !i =
          let !lowerArray =
                DWArray
                  { dwComp = Seq
                  , dwSize = lowerSourceSize
                  , dwIndex = uIndex . consDim i
                  , dwWindow = mkWindow i
                  }
              !innerScheduler =
                if numWorkers scheduler <= unSz (strideSize (Stride s) headSourceSize)
                  then trivialScheduler_
                  else scheduler
           in scheduleWork_ scheduler $
                iterArrayLinearWithStrideST_ innerScheduler (Stride lowerStrideIx) lowerSize lowerArray $ \k ->
                  uWrite (k + pageElements * (i `div` s))
        {-# INLINE loadLower #-}
    loopA_ 0 (< headDim windowStart) (+ s) (loadLower (const Nothing))
    loopA_
      (strideStart (Stride s) curWindowStart)
      (< curWindowEnd)
      (+ s)
      (loadLower (Just . mkLowerWindow))
    loopA_
      (strideStart (Stride s) curWindowEnd)
      (< unSz headSourceSize)
      (+ s)
      (loadLower (const Nothing))
  {-# INLINE loadArrayWithIxN #-}
  
  loadWithIxN
    :: (Index ix, Load DW (Lower ix) e)
    => Scheduler s ()
    -> Array DW ix e
    -> (Int -> e -> ST s ())
    -> ST s ()
  loadWithIxN scheduler arr uWrite = do
    let DWArray _ sz uIndex window = arr
        Window{windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window
        !(!si, !szL) = unconsSz sz
        !windowEnd = liftIndex2 (+) windowStart (unSz windowSize)
        !(!t, windowStartL) = unconsDim windowStart
        !pageElements = totalElem szL
        lowerWindow =
          Window
            { windowStart = windowStartL
            , windowSize = snd $ unconsSz windowSize
            , windowIndex = \_ -> error "Window index uninitialized"
            , windowUnrollIx2 = windowUnrollIx2
            }
        mkLowerWindow !i =
          lowerWindow
            { windowIndex = windowIndex . consDim i
            }
        loadLower mkWindow !i =
          let !lowerArray =
                DWArray
                  { dwComp = Seq
                  , dwSize = szL
                  , dwIndex = uIndex . consDim i
                  , dwWindow = mkWindow i
                  }
              !innerScheduler =
                if numWorkers scheduler <= unSz si
                  then trivialScheduler_
                  else scheduler
           in scheduleWork_ scheduler $
                iterArrayLinearST_ innerScheduler lowerArray (\k -> uWrite (k + pageElements * i))
        {-# INLINE loadLower #-}
    loopA_ 0 (< headDim windowStart) (+ 1) (loadLower (const Nothing))
    loopA_ t (< headDim windowEnd) (+ 1) (loadLower (Just . mkLowerWindow))
    loopA_ (headDim windowEnd) (< unSz si) (+ 1) (loadLower (const Nothing))
  {-# INLINE loadWithIxN #-}
  
  unrollAndJam
    :: Monad m
    => Int
    -- ^ Block height. Must not be zero.
    -> Ix2
    -- ^ Top corner
    -> Ix2
    -- ^ Bottom corner
    -> Int
    -- ^ Column Stride
    -> (Ix2 -> m ())
    -- ^ Writing function
    -> m ()
  unrollAndJam !bH (it :. jt) (ib :. jb) js f = do
    let
      f2 (i :. j) = f (i :. j) >> f ((i + 1) :. j)
      f3 (i :. j) = f (i :. j) >> f2 ((i + 1) :. j)
      f4 (i :. j) = f (i :. j) >> f3 ((i + 1) :. j)
      f5 (i :. j) = f (i :. j) >> f4 ((i + 1) :. j)
      f6 (i :. j) = f (i :. j) >> f5 ((i + 1) :. j)
      f7 (i :. j) = f (i :. j) >> f6 ((i + 1) :. j)
      f' = case bH of
        1 -> f
        2 -> f2
        3 -> f3
        4 -> f4
        5 -> f5
        6 -> f6
        _ -> f7
      !ibS = ib - ((ib - it) `modInt` bH)
    loopA_ it (< ibS) (+ bH) $ \ !i ->
      loopA_ jt (< jb) (+ js) $ \ !j ->
        f' (i :. j)
    loopA_ ibS (< ib) (+ 1) $ \ !i ->
      loopA_ jt (< jb) (+ js) $ \ !j ->
        f (i :. j)
  {-# INLINE unrollAndJam #-}
  
  -- TODO: Implement Hilbert curve
  