{-# LANGUAGE TypeFamilies #-}
module Data.Array.Comfort.Storable.Mutable.Internal where

import qualified Data.Array.Comfort.Storable.Internal as Imm
import qualified Data.Array.Comfort.Shape as Shape

import qualified Foreign.Marshal.Array.Guarded as Alloc
import Foreign.Marshal.Array (pokeArray, peekArray, )
import Foreign.Storable (Storable, pokeElemOff, peekElemOff, )
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, )
import Foreign.Ptr (Ptr, )

import System.IO.Unsafe (unsafePerformIO, )

import Control.Applicative ((<$>))
import Control.Monad ((<=<))

import Data.Tuple.HT (mapFst)

import Prelude hiding (map, read, )


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

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


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

createWithSize ::
   (Shape.C sh, Storable a) => sh -> (Int -> Ptr a -> IO ()) -> IO (Array sh a)
createWithSize sh f =
   fst <$> createWithSizeAndResult sh f

createWithSizeAndResult ::
   (Shape.C sh, Storable a) =>
   sh -> (Int -> Ptr a -> IO b) -> IO (Array sh a, b)
createWithSizeAndResult sh f =
   let size = Shape.size sh
   in fmap (mapFst (Array sh)) $ Alloc.create size $ f size

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

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

write ::
   (Shape.Indexed sh, Storable a) => Array sh a -> Shape.Index sh -> a -> IO ()
write (Array sh fptr) ix a =
   withForeignPtr fptr $ \ptr -> pokeElemOff ptr (Shape.offset sh ix) a

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

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

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


freeze :: (Shape.Indexed sh, Storable a) => Array sh a -> IO (Imm.Array sh a)
freeze = Imm.copyIO <=< unsafeFreeze

thaw :: (Shape.Indexed sh, Storable a) => Imm.Array sh a -> IO (Array sh a)
thaw = unsafeThaw <=< Imm.copyIO

unsafeFreeze ::
   (Shape.Indexed sh, Storable a) => Array sh a -> IO (Imm.Array sh a)
unsafeFreeze (Array sh fptr) = return (Imm.Array sh fptr)

unsafeThaw ::
   (Shape.Indexed sh, Storable a) => Imm.Array sh a -> IO (Array sh a)
unsafeThaw (Imm.Array sh fptr) = return (Array sh fptr)