{-# LANGUAGE TypeFamilies #-}
module Data.Array.Comfort.Storable.Internal (
   Array(Array, shape, buffer),
   reshape,
   mapShape,

   (!),
   unsafeCreate,
   toList,
   fromList,
   vectorFromList,

   createIO,
   createWithSizeIO,
   showIO,
   readIO,
   toListIO,
   fromListIO,
   vectorFromListIO,
   ) where

import qualified Data.Array.Comfort.Shape as Shape

import Foreign.Marshal.Array (pokeArray, peekArray, )
import Foreign.Storable (Storable, peekElemOff, )
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, mallocForeignPtrArray, )
import Foreign.Ptr (Ptr, )

import System.IO.Unsafe (unsafePerformIO, )

import Prelude hiding (readIO, )


data Array sh a =
   Array {
      shape :: sh,
      buffer :: ForeignPtr a
   }

instance (Shape.C sh, Show sh, Storable a, Show a) => Show (Array sh a) where
   show = unsafePerformIO . showIO

-- add assertion, at least in an exposed version
reshape :: sh1 -> Array sh0 a -> Array sh1 a
reshape sh (Array _ fptr) = Array sh fptr

mapShape :: (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape f (Array sh fptr) = Array (f sh) fptr


infixl 9 !

unsafeCreate ::
   (Shape.C sh, Storable a) => sh -> (Ptr a -> IO ()) -> Array sh a
unsafeCreate sh = unsafePerformIO . createIO sh

(!) :: (Shape.C sh, Storable a) => Array sh a -> Shape.Index sh -> a
(!) arr = unsafePerformIO . readIO arr

toList :: (Shape.C sh, Storable a) => Array sh a -> [a]
toList = unsafePerformIO . toListIO

fromList :: (Shape.C sh, Storable a) => sh -> [a] -> Array sh a
fromList sh = unsafePerformIO . fromListIO sh

vectorFromList :: (Storable a) => [a] -> Array (Shape.ZeroBased Int) a
vectorFromList = unsafePerformIO . vectorFromListIO


createIO ::
   (Shape.C sh, Storable a) => sh -> (Ptr a -> IO ()) -> IO (Array sh a)
createIO sh f = createWithSizeIO sh $ const f

createWithSizeIO ::
   (Shape.C sh, Storable a) => sh -> (Int -> Ptr a -> IO ()) -> IO (Array sh a)
createWithSizeIO sh f = do
   let size = Shape.size sh
   fptr <- mallocForeignPtrArray size
   withForeignPtr fptr $ f size
   return $ Array sh fptr

showIO :: (Shape.C sh, Show sh, Storable a, Show a) => Array sh a -> IO String
showIO arr = do
   xs <- toListIO arr
   return $ "fromList " ++ showsPrec 11 (shape arr) (' ' : show xs)

readIO :: (Shape.C sh, Storable a) => Array sh a -> Shape.Index sh -> IO a
readIO (Array sh fptr) ix =
   withForeignPtr fptr $ flip peekElemOff (Shape.offset sh ix)

toListIO :: (Shape.C sh, Storable a) => Array sh a -> IO [a]
toListIO (Array sh fptr) =
   withForeignPtr fptr $ peekArray (Shape.size sh)

fromListIO ::
   (Shape.C sh, Storable a) =>
   sh -> [a] -> IO (Array sh a)
fromListIO sh xs =
   createWithSizeIO sh $ \size ptr ->
      pokeArray ptr $ take size $
      xs ++
      repeat (error "Array.Comfort.Storable.fromList: list too short for shape")

vectorFromListIO :: (Storable a) => [a] -> IO (Array (Shape.ZeroBased Int) a)
vectorFromListIO xs =
   createIO (Shape.ZeroBased $ length xs) $ flip pokeArray xs