content: {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE EmptyDataDecls #-}
  {-# LANGUAGE MagicHash #-}
  {-# LANGUAGE UnboxedTuples #-}
  module T17151a
    ( computeS
    , Stencil
    , P(..)
    , D(..)
    , makeConvolutionStencilFromKernel
    , mapStencil
    , Array
    , Construct(..)
    , Source(..)
    , Load(..)
    , Mutable(..)
    ) where
  
  import Control.Monad.ST
  import Data.Functor.Identity
  import GHC.STRef
  import GHC.ST
  import GHC.Exts
  import Unsafe.Coerce
  import Data.Kind
  
  ----  Hacked up stuff to simulate primitive package
  class Prim e where
    indexByteArray :: ByteArray -> Int -> e
    sizeOf :: e ->Int
  instance Prim Int where
    indexByteArray _ _ = 55
    sizeOf _ = 99
  
  data ByteArray = BA
  type MutableByteArray s = STRef s Int
  
  class Monad m => PrimMonad m where
    type PrimState m
    primitive :: (State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
  instance PrimMonad (ST s) where
    type PrimState (ST s) = s
    primitive = ST
  
  unsafeFreezeByteArray :: PrimMonad m => MutableByteArray (PrimState m) -> m ByteArray
  unsafeFreezeByteArray a = return (unsafeCoerce a)
  
  newByteArray :: PrimMonad m => Int -> m (MutableByteArray (PrimState m))
  newByteArray (I# n#)
    = primitive (\s# -> case newMutVar# 33 s# of
                          (# s'#, arr# #) -> (# s'#, STRef arr# #))
  
  writeByteArray :: PrimMonad m => MutableByteArray (PrimState m) -> Int -> e -> m ()
  writeByteArray _ _ _ = return ()
  
  -----  End of hacked up stuff
  
  --------------
  newtype Stencil ix e a =
    Stencil ((ix -> e) -> ix -> a)
  
  mapStencil :: Source r ix e => Stencil ix e a -> Array r ix e -> Array D ix a
  mapStencil (Stencil stencilF) arr = DArray (size arr) (stencilF (unsafeIndex arr))
  {-# INLINE mapStencil #-}
  
  makeConvolutionStencilFromKernel
    :: (Source r ix e, Num e)
    => Array r ix e
    -> Stencil ix e e
  makeConvolutionStencilFromKernel arr = Stencil stencil
    where
      sz = size arr
      sCenter = liftIndex (`quot` 2) sz
      stencil getVal ix =
        runIdentity $
        loopM 0 (< totalElem sz) (+ 1) 0 $ \i a ->
          pure $ accum a (fromLinearIndex sz i) (unsafeLinearIndex arr i)
        where
          ixOff = liftIndex2 (+) ix sCenter
          accum acc kIx kVal = getVal (liftIndex2 (-) ixOff kIx) * kVal + acc
          {-# INLINE accum #-}
      {-# INLINE stencil #-}
  {-# INLINE makeConvolutionStencilFromKernel #-}
  
  
  computeS :: (Mutable r ix e, Load r' ix e) => Array r' ix e -> Array r ix e
  computeS arr = runST $ do
    marr <- unsafeNew (size arr)
    unsafeLoadIntoS marr arr
    unsafeFreeze marr
  {-# INLINE computeS #-}
  
  
  data D = D deriving Show
  
  data instance  Array D ix e = DArray{dSize :: ix,
                                       dIndex :: ix -> e}
  
  instance Index ix => Construct D ix e where
    makeArray _ = DArray
    {-# INLINE makeArray #-}
  
  instance Index ix => Source D ix e where
    unsafeIndex = dIndex
    {-# INLINE unsafeIndex #-}
  
  instance Index ix => Load D ix e where
    size = dSize
    {-# INLINE size #-}
    loadArrayM arr = splitLinearlyWith_ (totalElem (size arr)) (unsafeLinearIndex arr)
    {-# INLINE loadArrayM #-}
  
  
  data P = P deriving Show
  
  data instance Array P ix e = PArray ix ByteArray
  
  instance (Prim e, Index ix) => Construct P ix e where
    makeArray _ sz f = computeS (makeArray D sz f)
    {-# INLINE makeArray #-}
  
  instance (Prim e, Index ix) => Source P ix e where
    unsafeIndex (PArray sz a) = indexByteArray a . toLinearIndex sz
    {-# INLINE unsafeIndex #-}
  
  instance (Prim e, Index ix) => Mutable P ix e where
    data MArray s P ix e = MPArray ix (MutableByteArray s)
    unsafeFreeze (MPArray sz a) = PArray sz <$> unsafeFreezeByteArray a
    {-# INLINE unsafeFreeze #-}
    unsafeNew sz = MPArray sz <$> newByteArray (totalElem sz * eSize)
      where
        eSize = sizeOf (undefined :: e)
    {-# INLINE unsafeNew #-}
    unsafeLinearWrite (MPArray _ ma) = writeByteArray ma
    {-# INLINE unsafeLinearWrite #-}
  
  
  instance (Prim e, Index ix) => Load P ix e where
    size (PArray sz _) = sz
    {-# INLINE size #-}
    loadArrayM arr = splitLinearlyWith_ (totalElem (size arr)) (unsafeLinearIndex arr)
    {-# INLINE loadArrayM #-}
  
  
  unsafeLinearIndex :: Source r ix e => Array r ix e -> Int -> e
  unsafeLinearIndex arr = unsafeIndex arr . fromLinearIndex (size arr)
  {-# INLINE unsafeLinearIndex #-}
  
  
  loopM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
  loopM init' condition increment initAcc f = go init' initAcc
    where
      go step acc
        | condition step = f step acc >>= go (increment step)
        | otherwise = return acc
  {-# INLINE loopM #-}
  
  splitLinearlyWith_ ::
       Monad m => Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
  splitLinearlyWith_ totalLength index write =
    loopM 0 (< totalLength) (+1) () $ \i () -> write i (index i)
  {-# INLINE splitLinearlyWith_ #-}
  
  
  data family Array r ix e :: Type
  
  class Index ix => Construct r ix e where
    makeArray :: r -> ix -> (ix -> e) -> Array r ix e
  
  class Load r ix e => Source r ix e where
    unsafeIndex :: Array r ix e -> ix -> e
  
  class Index ix => Load r ix e where
    size :: Array r ix e -> ix
    loadArrayM :: Monad m => Array r ix e -> (Int -> e -> m ()) -> m ()
    unsafeLoadIntoS ::
         (Mutable r' ix e, PrimMonad m) => MArray (PrimState m) r' ix e -> Array r ix e -> m ()
    unsafeLoadIntoS marr arr = loadArrayM arr (unsafeLinearWrite marr)
    {-# INLINE unsafeLoadIntoS #-}
  
  class (Construct r ix e, Source r ix e) => Mutable r ix e where
    data MArray s r ix e :: Type
    unsafeFreeze :: PrimMonad m => MArray (PrimState m) r ix e -> m (Array r ix e)
    unsafeNew :: PrimMonad m => ix -> m (MArray (PrimState m) r ix e)
    unsafeLinearWrite :: PrimMonad m => MArray (PrimState m) r ix e -> Int -> e -> m ()
  
  
  class (Eq ix, Ord ix, Show ix) =>
        Index ix
    where
    totalElem :: ix -> Int
    liftIndex2 :: (Int -> Int -> Int) -> ix -> ix -> ix
    liftIndex :: (Int -> Int) -> ix -> ix
    toLinearIndex :: ix -> ix -> Int
    fromLinearIndex :: ix -> Int -> ix
  
  instance Index Int where
    totalElem = id
    toLinearIndex _ = id
    fromLinearIndex _ = id
    liftIndex f = f
    liftIndex2 f = f
  