{-# LANGUAGE Rank2Types,  UndecidableInstances, MultiParamTypeClasses, TypeFamilies, FlexibleInstances, GeneralizedNewtypeDeriving #-}

module Control.Monad.Array.Unboxed where

import Control.Monad.ST.Class
import Control.Monad.Array.Class

import Data.Array.Vector

import Control.Monad.ST
import Control.Monad.State.Strict
import Control.Monad.Reader
import Control.Monad.Writer.Class

type UArrayM s e = UArrayT e (ST s)
newtype UArrayT e m a = UArrayT {runUArrayT :: StateT (MUArr e (StateThread m)) m a} deriving (Monad, MonadReader r, MonadWriter w, MonadIO, MonadFix)

evalUArrayT :: (UA e, MonadST m, Monad m) => Int -> UArrayT e m a -> m a
evalUArrayT n m = liftST (newMU n) >>= evalStateT (runUArrayT m)

defaultUA :: UA e => e
defaultUA = runST $ do	arr <- newMU 1
			readMU arr 0

evalUArrayM :: UA e => Int -> (forall s . UArrayM s e a) -> a
evalUArrayM n m = runST (evalUArrayT n m)

instance MonadTrans (UArrayT e) where
	lift m = UArrayT (lift m)

instance MonadState s m => MonadState s (UArrayT e m) where
	get = lift get
	put = lift . put

-- instance MonadReader r m => MonadReader r (UArrayT e m) where
-- 	ask = lift ask
-- 	local f (UArrayT m) = UArrayT $ ReaderT $ \ r -> local f (runReaderT m r)

instance (UA e, MonadST m, Monad m) => MonadArray (UArrayT e m) where
	{-# SPECIALIZE instance UA e => MonadArray (UArrayM s e) #-}
	type ArrayElem (UArrayT e m) = e
	askSize = UArrayT $ gets lengthMU
	readAt i = UArrayT $ gets (`readMU` i) >>= liftST
	writeAt i x = UArrayT $ gets (\ arr -> writeMU arr i x) >>= liftST
	resize n' = UArrayT $ do	n <- gets lengthMU
					ice <- gets unsafeFreezeAllMU
					put =<< liftST (do
						noob <- newMU n'
						copyMU noob 0 =<< ice
						return noob)
	askElems = UArrayT $ liftST =<< gets (liftM fromU . unsafeFreezeAllMU)