{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PatternSynonyms #-}

-- |
-- Module      : Data.Massiv.Core.Index
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <alexey@kuleshevi.ch>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Core.Index (
  Ix0 (..),
  type Ix1,
  pattern Ix1,
  type Ix2 (Ix2, (:.)),
  IxN ((:>), Ix3, Ix4, Ix5),
  HighIxN,
  type Ix3,
  type Ix4,
  type Ix5,
  Ix,

  -- ** Size
  type Sz1,
  type Sz2,
  type Sz3,
  type Sz4,
  type Sz5,
  Sz (Sz, Sz1, Sz2, Sz3, Sz4, Sz5),
  unSz,
  zeroSz,
  oneSz,
  liftSz,
  liftSz2,
  consSz,
  unconsSz,
  snocSz,
  unsnocSz,
  setSzM,
  insertSzM,
  pullOutSzM,
  toLinearSz,
  mkSzM,

  -- ** Dimension
  Dim (..),
  Dimension (Dim1, Dim2, Dim3, Dim4, Dim5, DimN),
  IsIndexDimension,
  IsDimValid,
  ReportInvalidDim,

  -- ** Stride
  Stride (Stride),
  unStride,
  toLinearIndexStride,
  strideStart,
  strideSize,
  oneStride,

  -- ** Border
  Border (..),
  handleBorderIndex,

  -- ** Index functions
  Lower,
  Index (..),
  zeroIndex,
  oneIndex,
  isZeroSz,
  isNotZeroSz,
  headDim,
  tailDim,
  lastDim,
  initDim,
  getDim',
  setDim',
  modifyDim',
  dropDimM,
  dropDim',
  pullOutDim',
  insertDim',
  fromDimension,
  getDimension,
  setDimension,
  modifyDimension,
  dropDimension,
  pullOutDimension,
  insertDimension,

  -- * Iterators
  iter,
  iterA_,
  iterM_,
  iterLinearM,
  iterLinearM_,
  module Data.Massiv.Core.Loop,
  module Data.Massiv.Core.Index.Iterator,
  module Data.Massiv.Core.Index.Tuple,

  -- * Exceptions
  IndexException (..),
  SizeException (..),
  ShapeException (..),
  guardNumberOfElements,
  indexAssert,
  indexWith,
) where

import Control.DeepSeq
import Control.Monad.Catch (MonadThrow (..))
import Data.Coerce
import Data.Functor.Identity (runIdentity)
import Data.Massiv.Core.Exception
import Data.Massiv.Core.Index.Internal
import Data.Massiv.Core.Index.Iterator
import Data.Massiv.Core.Index.Ix
import Data.Massiv.Core.Index.Stride
import Data.Massiv.Core.Index.Tuple
import Data.Massiv.Core.Loop
import GHC.TypeLits

#include "massiv.h"

-- | 1-dimensional type synonym for size.
--
-- @since 0.3.0
type Sz1 = Sz Ix1

-- | 2-dimensional size type synonym.
--
-- @since 0.3.0
type Sz2 = Sz Ix2

-- | 3-dimensional size type synonym.
--
-- @since 0.3.0
type Sz3 = Sz Ix3

-- | 4-dimensional size type synonym.
--
-- @since 0.3.0
type Sz4 = Sz Ix4

-- | 5-dimensional size type synonym.
--
-- @since 0.3.0
type Sz5 = Sz Ix5

-- | Approach to be used near the borders during various transformations.
-- Whenever a function needs information not only about an element of interest, but
-- also about it's neighbors, it will go out of bounds near the array edges,
-- hence is this set of approaches that specify how to handle such situation.
data Border e
  = -- | Fill in a constant element.
    --
    -- @
    --            outside |  Array  | outside
    -- ('Fill' 0) : 0 0 0 0 | 1 2 3 4 | 0 0 0 0
    -- @
    Fill e
  | -- | Wrap around from the opposite border of the array.
    --
    -- @
    --            outside |  Array  | outside
    -- 'Wrap' :     1 2 3 4 | 1 2 3 4 | 1 2 3 4
    -- @
    Wrap
  | -- | Replicate the element at the edge.
    --
    -- @
    --            outside |  Array  | outside
    -- 'Edge' :     1 1 1 1 | 1 2 3 4 | 4 4 4 4
    -- @
    Edge
  | -- | Mirror like reflection.
    --
    -- @
    --            outside |  Array  | outside
    -- 'Reflect' :  4 3 2 1 | 1 2 3 4 | 4 3 2 1
    -- @
    Reflect
  | -- | Also mirror like reflection, but without repeating the edge element.
    --
    -- @
    --            outside |  Array  | outside
    -- 'Continue' : 1 4 3 2 | 1 2 3 4 | 3 2 1 4
    -- @
    Continue
  deriving (Eq, Show)

instance NFData e => NFData (Border e) where
  rnf b = case b of
    Fill e -> rnf e
    Wrap -> ()
    Edge -> ()
    Reflect -> ()
    Continue -> ()

-- | Apply a border resolution technique to an index
--
-- ==== __Examples__
--
-- >>> handleBorderIndex (Fill 100) (Sz (2 :. 3)) id (2 :. 3)
-- 100 :. 100
-- >>> handleBorderIndex Wrap (Sz (2 :. 3)) id (2 :. 3)
-- 0 :. 0
-- >>> handleBorderIndex Edge (Sz (2 :. 3)) id (2 :. 3)
-- 1 :. 2
--
-- @since 0.1.0
handleBorderIndex
  :: Index ix
  => Border e
  -- ^ Broder resolution technique
  -> Sz ix
  -- ^ Size
  -> (ix -> e)
  -- ^ Index function that produces an element
  -> ix
  -- ^ Index
  -> e
handleBorderIndex border !sz getVal !ix =
  case border of
    Fill val -> if isSafeIndex sz ix then getVal ix else val
    Wrap -> getVal (repairIndex sz ix wrap wrap)
    Edge -> getVal (repairIndex sz ix (const (const 0)) (\(SafeSz k) _ -> k - 1))
    Reflect ->
      getVal
        ( repairIndex
            sz
            ix
            (\(SafeSz k) !i -> (abs i - 1) `mod` k)
            (\(SafeSz k) !i -> (-i - 1) `mod` k)
        )
    Continue ->
      getVal
        ( repairIndex
            sz
            ix
            (\(SafeSz k) !i -> abs i `mod` k)
            (\(SafeSz k) !i -> (-i - 2) `mod` k)
        )
  where
    wrap (SafeSz k) i = i `mod` k
    {-# INLINE [1] wrap #-}
{-# INLINE [1] handleBorderIndex #-}

-- | Index with all zeros
--
-- ==== __Examples__
--
-- >>> zeroIndex :: Ix4
-- 0 :> 0 :> 0 :. 0
--
-- @since 0.1.0
zeroIndex :: Index ix => ix
zeroIndex = pureIndex 0
{-# INLINE [1] zeroIndex #-}

-- | Index with all ones
--
-- @since 0.3.0
oneIndex :: Index ix => ix
oneIndex = pureIndex 1
{-# INLINE [1] oneIndex #-}

-- | Checks whether size can hold at least one element.
--
-- ==== __Examples__
--
-- >>> isNotZeroSz (Sz3 1 0 2)
-- False
--
-- @since 1.0.0
isNotZeroSz :: Index ix => Sz ix -> Bool
isNotZeroSz !sz = isSafeIndex sz zeroIndex
{-# INLINE [1] isNotZeroSz #-}

-- TODO: benchmark against (also adjust `isEmpty` with fastest):
-- - foldlIndex (*) 1 (unSz sz) /= 0
-- - foldlIndex (\a x -> a && x /= 0) True (unSz sz)
-- - totalElem sz == 0

-- | Checks whether size can hold at least one element.
--
-- ==== __Examples__
--
-- >>> isZeroSz (Sz3 1 0 2)
-- True
--
-- @since 1.0.0
isZeroSz :: Index ix => Sz ix -> Bool
isZeroSz = not . isNotZeroSz
{-# INLINE [1] isZeroSz #-}

-- | Convert a size to a linear size.
--
-- @since 0.5.8
toLinearSz :: Index ix => Sz ix -> Sz1
toLinearSz = coerce . totalElem
{-# INLINE [1] toLinearSz #-}

-- | Get the outmost dimension of the index.
--
-- ==== __Examples__
--
-- >>> headDim (2 :> 3 :> 4 :. 5)
-- 2
--
-- @since 0.1.0
headDim :: Index ix => ix -> Int
headDim = fst . unconsDim
{-# INLINE [1] headDim #-}

-- | Drop the outmost dimension from the index
--
-- ==== __Examples__
--
-- >>> tailDim (2 :> 3 :> 4 :. 5)
-- 3 :> 4 :. 5
--
-- @since 0.1.0
tailDim :: Index ix => ix -> Lower ix
tailDim = snd . unconsDim
{-# INLINE [1] tailDim #-}

-- | Get the innermost dimension from the index
--
-- ==== __Examples__
--
-- >>> lastDim (2 :> 3 :> 4 :. 5)
-- 5
--
-- @since 0.1.0
lastDim :: Index ix => ix -> Int
lastDim = snd . unsnocDim
{-# INLINE [1] lastDim #-}

-- | Drop the innermost dimension from the index
--
-- ==== __Examples__
--
-- >>> initDim (2 :> 3 :> 4 :. 5)
-- 2 :> 3 :. 4
--
-- @since 0.1.0
initDim :: Index ix => ix -> Lower ix
initDim = fst . unsnocDim
{-# INLINE [1] initDim #-}

-- | Change the value of a specific dimension within the index. See `setDimM` for a safer
-- version and `setDimension` for a type safe version.
--
-- ==== __Examples__
--
-- >>> setDim' (2 :> 3 :> 4 :. 5) 3 10
-- 2 :> 10 :> 4 :. 5
--
-- @since 0.2.4
setDim' :: (HasCallStack, Index ix) => ix -> Dim -> Int -> ix
setDim' ix dim = throwEither . setDimM ix dim
{-# INLINE [1] setDim' #-}

-- | Change the value from a specific dimension within the index. See
-- `getDimM` for a safer version and `getDimension` for a type safe version.
--
-- ==== __Examples__
--
-- >>> getDim' (2 :> 3 :> 4 :. 5) 3
-- 3
--
-- @since 0.2.4
getDim' :: (HasCallStack, Index ix) => ix -> Dim -> Int
getDim' ix = throwEither . getDimM ix
{-# INLINE [1] getDim' #-}

-- | Update the value of a specific dimension within the index. See
-- `modifyDimM` for a safer version and `modifyDimension` for a type safe version.
--
-- ==== __Examples__
--
-- >>> modifyDim' (2 :> 3 :> 4 :. 5) 2 (+ 10)
-- (4,2 :> 3 :> 14 :. 5)
--
-- @since 0.4.1
modifyDim' :: (HasCallStack, Index ix) => ix -> Dim -> (Int -> Int) -> (Int, ix)
modifyDim' ix dim = throwEither . modifyDimM ix dim
{-# INLINE [1] modifyDim' #-}

-- | Remove a dimension from the index.
--
-- ==== __Examples__
--
-- >>> dropDimM (2 :> 3 :> 4 :. 5) 3 :: Maybe Ix3
-- Just (2 :> 4 :. 5)
-- >>> dropDimM (2 :> 3 :> 4 :. 5) 6 :: Maybe Ix3
-- Nothing
--
-- @since 0.3.0
dropDimM :: (MonadThrow m, Index ix) => ix -> Dim -> m (Lower ix)
dropDimM ix = fmap snd . pullOutDimM ix
{-# INLINE [1] dropDimM #-}

-- | Remove a dimension from the index.
--
-- ==== __Examples__
--
-- >>> dropDim' (2 :> 3 :> 4 :. 5) 3
-- 2 :> 4 :. 5
--
-- @since 0.2.4
dropDim' :: (HasCallStack, Index ix) => ix -> Dim -> Lower ix
dropDim' ix = throwEither . dropDimM ix
{-# INLINE [1] dropDim' #-}

-- | Lower the dimension of the index by pulling the specified dimension. See
-- `pullOutDimM` for a safer version and `pullOutDimension` for a type safe version.
--
-- ==== __Examples__
--
-- >>> pullOutDim' (2 :> 3 :> 4 :. 5) 3
-- (3,2 :> 4 :. 5)
--
-- @since 0.2.4
pullOutDim' :: (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix)
pullOutDim' ix = throwEither . pullOutDimM ix
{-# INLINE [1] pullOutDim' #-}

-- | Raise the dimension of the index by inserting one in the specified dimension. See
-- `insertDimM` for a safer version and `insertDimension` for a type safe version.
--
-- ==== __Examples__
--
-- >>> insertDim' (2 :> 3 :> 4 :. 5) 3 10 :: Ix5
-- 2 :> 3 :> 10 :> 4 :. 5
--
-- @since 0.2.4
insertDim' :: (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' ix dim = throwEither . insertDimM ix dim
{-# INLINE [1] insertDim' #-}

-- | Get the value level `Dim` from the type level equivalent.
--
-- ==== __Examples__
--
-- >>> fromDimension Dim4
-- (Dim 4)
-- >>> :set -XDataKinds
-- >>> fromDimension (DimN :: Dimension 10)
-- (Dim 10)
--
-- @since 0.2.4
fromDimension :: KnownNat n => Dimension n -> Dim
fromDimension = fromIntegral . natVal
{-# INLINE [1] fromDimension #-}

-- | Type safe way to set value of index at a particular dimension.
--
-- ==== __Examples__
--
-- >>> setDimension (2 :> 3 :> 4 :. 5) Dim4 10
-- 10 :> 3 :> 4 :. 5
--
-- @since 0.2.4
setDimension :: IsIndexDimension ix n => ix -> Dimension n -> Int -> ix
setDimension ix = setDim' ix . fromDimension
{-# INLINE [1] setDimension #-}

-- | Type safe way to set value of index at a particular dimension.
--
-- ==== __Examples__
--
-- >>> modifyDimension (2 :> 3 :> 4 :. 5) Dim3 (+ 2)
-- (3,2 :> 5 :> 4 :. 5)
--
-- @since 0.4.1
modifyDimension :: IsIndexDimension ix n => ix -> Dimension n -> (Int -> Int) -> (Int, ix)
modifyDimension ix = modifyDim' ix . fromDimension
{-# INLINE [1] modifyDimension #-}

-- | Type safe way to extract value of index at a particular dimension.
--
-- ==== __Examples__
--
-- >>> getDimension (2 :> 3 :> 4 :. 5) Dim2
-- 4
--
-- @since 0.2.4
getDimension :: IsIndexDimension ix n => ix -> Dimension n -> Int
getDimension ix = getDim' ix . fromDimension
{-# INLINE [1] getDimension #-}

-- | Type safe way of dropping a particular dimension, thus lowering index
-- dimensionality.
--
-- ==== __Examples__
--
-- >>> dropDimension (2 :> 3 :> 4 :. 5) Dim2
-- 2 :> 3 :. 5
--
-- @since 0.2.4
dropDimension :: IsIndexDimension ix n => ix -> Dimension n -> Lower ix
dropDimension ix = dropDim' ix . fromDimension
{-# INLINE [1] dropDimension #-}

-- | Type safe way of pulling out a particular dimension, thus lowering index
-- dimensionality and returning the value at specified dimension.
--
-- ==== __Examples__
--
-- >>> pullOutDimension (2 :> 3 :> 4 :. 5) Dim2
-- (4,2 :> 3 :. 5)
--
-- @since 0.2.4
pullOutDimension :: IsIndexDimension ix n => ix -> Dimension n -> (Int, Lower ix)
pullOutDimension ix = pullOutDim' ix . fromDimension
{-# INLINE [1] pullOutDimension #-}

-- | Type safe way of inserting a particular dimension, thus raising index dimensionality.
--
-- ==== __Examples__
--
-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim5 10 :: Ix5
-- 10 :> 2 :> 3 :> 4 :. 5
-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim4 10 :: Ix5
-- 2 :> 10 :> 3 :> 4 :. 5
-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim3 10 :: Ix5
-- 2 :> 3 :> 10 :> 4 :. 5
-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim2 10 :: Ix5
-- 2 :> 3 :> 4 :> 10 :. 5
-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim1 10 :: Ix5
-- 2 :> 3 :> 4 :> 5 :. 10
--
-- @since 0.2.5
insertDimension :: IsIndexDimension ix n => Lower ix -> Dimension n -> Int -> ix
insertDimension ix = insertDim' ix . fromDimension
{-# INLINE [1] insertDimension #-}

-- | Row-major iterator for the index. Same as `iterM`, but pure.
--
-- ==== __Examples__
--
-- >>> iter (Ix1 0) 1000 1 (<) 0 (+)
-- 499500
-- >>> iter (0 :. 0) (2 :. 3) oneIndex (<) 100 $ \ (i :. j) acc -> (acc + i) * (j + 1)
-- 3615
--
-- @since 0.1.0
iter
  :: Index ix
  => ix
  -- ^ Start index
  -> ix
  -- ^ End index
  -> ix
  -- ^ Increment
  -> (Int -> Int -> Bool)
  -- ^ Continuation condition
  -> a
  -- ^ Accumulator
  -> (ix -> a -> a)
  -- ^ Iterating function
  -> a
iter sIx eIx incIx cond acc f =
  runIdentity $ iterM sIx eIx incIx cond acc (\ix -> return . f ix)
{-# INLINE iter #-}

-- | Iterate over N-dimensional space linearly from start to end in row-major fashion with an
-- accumulator
--
-- ==== __Examples__
--
-- >>> sz = Sz2 3 4
-- >>> iterLinearM sz 0 3 1 (<) 100 $ \ k ix acc -> (acc + k) <$ print (fromLinearIndex sz k == ix)
-- True
-- True
-- True
-- 103
--
-- @since 0.1.0
iterLinearM
  :: (Index ix, Monad m)
  => Sz ix
  -- ^ Size
  -> Int
  -- ^ Linear start (must be non-negative)
  -> Int
  -- ^ Linear end (must be less than or equal to @`totalElem` sz@)
  -> Int
  -- ^ Increment (must not be zero)
  -> (Int -> Int -> Bool)
  -- ^ Continuation condition (continue if @True@)
  -> a
  -- ^ Accumulator
  -> (Int -> ix -> a -> m a)
  -> m a
iterLinearM !sz !k0 !k1 !inc cond !acc f =
  loopM k0 (`cond` k1) (+ inc) acc $ \ !i !acc0 -> f i (fromLinearIndex sz i) acc0
{-# INLINE iterLinearM #-}

-- | Same as `iterLinearM`, except without an accumulator.
--
-- ==== __Examples__
--
-- >>> sz = Sz2 3 4
-- >>> iterLinearM_ sz 0 3 1 (<) $ \ k ix -> print (toLinearIndex sz ix == k)
-- True
-- True
-- True
--
-- @since 0.1.0
iterLinearM_
  :: (Index ix, Monad m)
  => Sz ix
  -- ^ Size
  -> Int
  -- ^ Start (must be non-negative)
  -> Int
  -- ^ End
  -> Int
  -- ^ Increment (must not be zero)
  -> (Int -> Int -> Bool)
  -- ^ Continuation condition (continue if @True@)
  -> (Int -> ix -> m ())
  -- ^ Monadic action that takes index in both forms
  -> m ()
iterLinearM_ sz !k0 !k1 !inc cond f =
  loopA_ k0 (`cond` k1) (+ inc) $ \ !i -> f i (fromLinearIndex sz i)
{-# INLINE iterLinearM_ #-}

-- | This is used by the @unsafe-checks@ cabal flag.
--
-- @since 1.1.0
#ifdef MASSIV_UNSAFE_CHECKS
indexAssert :: (HasCallStack, Index ix) => String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert funName getSize f arr ix
  | isSafeIndex sz ix = f arr ix
  | otherwise = _errorIx ("<" ++ funName ++ ">") sz ix
  where
    sz = getSize arr
#else
indexAssert :: String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert _funName _getSize f arr ix = f arr ix
#endif
{-# INLINE indexAssert #-}

-- | This is used by @INDEX_CHECK@ macro and thus used whenever the @unsafe-checks@ cabal
-- flag is on.
--
-- @since 0.4.0
indexWith
  :: Index ix
  => String
  -- ^ Source file name, eg. __FILE__
  -> Int
  -- ^ Line number in th source file, eg. __LINE__
  -> String
  -> (arr -> Sz ix)
  -- ^ Get size of the array
  -> (arr -> ix -> e)
  -- ^ Indexing function
  -> arr
  -- ^ Array
  -> ix
  -- ^ Index
  -> e
indexWith fileName lineNo funName getSize f arr ix
  | isSafeIndex sz ix = f arr ix
  | otherwise = _errorIx ("<" ++ fileName ++ ":" ++ show lineNo ++ "> " ++ funName) sz ix
  where
    sz = getSize arr
{-# DEPRECATED indexWith "In favor of `indexAssert` that uses HasCallStack" #-}

-- | Helper function for throwing out of bounds error. Used by `indexAssert`
_errorIx :: (HasCallStack, Show ix, Show ix') => String -> ix -> ix' -> a
_errorIx fName sz ix =
  error $
    fName
      ++ ": Index out of bounds: ("
      ++ show ix
      ++ ") for Array of size: ("
      ++ show sz
      ++ ")"
{-# NOINLINE _errorIx #-}