--------------------------------------------------------------------------------
-- |
-- Module      :  Data.IndexedDoublyLinkedList
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--------------------------------------------------------------------------------
module Data.IndexedDoublyLinkedList( DLList(..)
                                   , Cell(..), emptyCell
                                   , DLListMonad, runDLListMonad
                                   , Index

                                   , singletons
                                   , writeList
                                   , valueAt, getNext, getPrev
                                   , toListFrom, toListFromR, toListContains
                                   , toListFromK, toListFromRK
                                   , insertAfter, insertBefore
                                   , delete
                                   , dump
                                   ) where

import           Control.Monad.Primitive (PrimMonad(..))
import           Control.Monad.Reader (ReaderT, runReaderT)
import           Control.Monad.Reader.Class
import           Control.Monad.ST
import           Data.Foldable (forM_)
import           Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NonEmpty
import           Data.Util
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV

--------------------------------------------------------------------------------

-- | Cell indices. Must be non-negative.
type Index = Int

-- TODO: Switch to unobxed sums for these!

-- | Cells in the Linked List
data Cell = Cell { prev :: Maybe Index
                 , next :: Maybe Index
                 } deriving (Show,Eq)

-- | Empty cell with no next or prev cells.
emptyCell :: Cell
emptyCell = Cell Nothing Nothing

-- | Doubly linked list implemented by a mutable vector. So actually
-- this data type can represent a collection of Linked Lists that can
-- efficiently be concatenated and split.
--
-- Supports O(1) indexing, and O(1) insertions, deletions
data DLList s a = DLList { values :: !(V.Vector a)
                         , llist  :: !(MV.MVector s Cell)
                         }

instance Functor (DLList s) where
  fmap f (DLList v l) = DLList (fmap f v) l


--------------------------------------------------------------------------------

-- | Monad in which we can use the IndexedDoublyLinkedList.
newtype DLListMonad s b a = DLListMonad { runDLListMonad' :: ReaderT (DLList s b) (ST s) a }
                         deriving (Functor,Applicative,Monad)

instance PrimMonad (DLListMonad s b) where
  type PrimState (DLListMonad s b) = s
  primitive = DLListMonad . primitive

instance MonadReader (DLList s b) (DLListMonad s b) where
  local f = DLListMonad . local f . runDLListMonad'
  ask = DLListMonad ask

-- | Runs a DLList Computation, starting with singleton values, crated
-- from the input vector.
runDLListMonad         :: V.Vector b -> (forall s. DLListMonad s b a) -> a
runDLListMonad vs comp = runST $ singletons vs >>= runReaderT (runDLListMonad' comp)

----------------------------------------

-- | Constructs a new DoublyLinkedList. Every element is its own singleton list
singletons    :: (PrimMonad m, s ~ PrimState m) => V.Vector b -> m (DLList s b)
singletons vs = DLList vs <$> MV.replicate (V.length vs) emptyCell

-- | Sets the DoublyLinkedList to the given List.
--
-- Indices that do not occur in the list are not touched.
writeList   :: NonEmpty Index -> DLListMonad s b ()
writeList h = do v <- asks llist
                 forM_ (withNeighs h) $ \(STR p i s) ->
                   modify v i $ \c -> c { prev = p , next = s }
  where
    withNeighs (x:|xs) = let l = x:xs
                         in zipWith3 STR (Nothing : map Just l) l (map Just xs ++ [Nothing])

----------------------------------------
-- * Queries

-- | Gets the value at Index i
valueAt    :: Index -> DLListMonad s b b
valueAt  i = asks ((V.! i) . values)

-- | Next element in the List
getNext   :: Index -> DLListMonad s b (Maybe Index)
getNext i = do v <- asks llist
               next <$> MV.read v i

-- | Previous Element in the List
getPrev   :: Index -> DLListMonad s b (Maybe Index)
getPrev i = do v <- asks llist
               prev <$> MV.read v i

-- | Computes a maximal length list starting from the Given index
--
-- running time: \(O(k)\), where \(k\) is the length of the output list
toListFrom   :: Index -> DLListMonad s b (NonEmpty Index)
toListFrom i = (i :|) <$> iterateM getNext i

-- | Takes the current element and its k next's
toListFromK     :: Index -> Int -> DLListMonad s b (NonEmpty Index)
toListFromK i k = (i :|) <$> replicateM k getNext i

-- | Computes a maximal length list by walking backwards in the
-- DoublyLinkedList, starting from the Given index
--
-- running time: \(O(k)\), where \(k\) is the length of the output list
toListFromR :: Index -> DLListMonad s b (NonEmpty Index)
toListFromR i = (i :|) <$> iterateM getPrev i

-- | Takes the current element and its k prev's
toListFromRK     :: Index -> Int -> DLListMonad s b (NonEmpty Index)
toListFromRK i k = (i :|) <$> replicateM k getPrev i

-- | Computes a maximal length list that contains the element i.
--
-- running time: \(O(k)\), where \(k\) is the length of the output
-- list
toListContains   :: Index -> DLListMonad s b (NonEmpty Index)
toListContains i = f <$> toListFromR i <*> toListFrom i
  where
    f l r = NonEmpty.fromList $ reverse (NonEmpty.toList l) <> NonEmpty.tail r


----------------------------------------
-- * Updates

-- | Inserts the second argument after the first one into the linked list
insertAfter     :: Index -> Index -> DLListMonad s b ()
insertAfter i j = do v  <- asks llist
                     mr <- getNext i
                     modify  v i  $ \c -> c { next = Just j }
                     modify  v j  $ \c -> c { prev = Just i , next = mr }
                     mModify v mr $ \c -> c { prev = Just j }

-- | Inserts the second argument before the first one into the linked list
insertBefore     :: Index -> Index -> DLListMonad s b ()
insertBefore i h = do v <- asks llist
                      ml <- getPrev i
                      mModify v ml $ \c -> c { next = Just h }
                      modify  v h  $ \c -> c { prev = ml , next = Just i }
                      modify  v i  $ \c -> c { prev = Just h }

-- | Deletes the element from the linked list. This element thus
-- essentially becomes a singleton list. Returns the pair of indices
-- that now have become neighbours (i.e. the predecessor and successor
-- of j just before we deleted j).
delete   :: Index -> DLListMonad s b (Maybe Index, Maybe Index)
delete j = do v <- asks llist
              ml <- getPrev j
              mr <- getNext j
              modify  v j  $ \c -> c { prev = Nothing, next = Nothing }
              mModify v ml $ \c -> c { next = mr }
              mModify v mr $ \c -> c { prev = ml }
              pure (ml,mr)




----------------------------------------
-- * Helper functions

-- | Applies the action at most n times.
replicateM     :: Monad m => Int -> (a -> m (Maybe a)) -> a -> m [a]
replicateM n f = go n
  where
    go 0 _ = pure []
    go k x = f x >>= \case
               Nothing -> pure []
               Just y  -> (y:) <$> go (k-1) y

iterateM  :: Monad m => (a -> m (Maybe a)) -> a -> m [a]
iterateM f = go
  where
    go x = f x >>= \case
             Nothing -> pure []
             Just y  -> (y:) <$> go y

mModify   :: PrimMonad m => MV.MVector (PrimState m) a -> Maybe Int -> (a -> a) -> m ()
mModify v mi f = case mi of
                   Nothing -> pure ()
                   Just i  -> modify v i f

modify        :: PrimMonad m => MV.MVector (PrimState m) a -> Int -> (a -> a) -> m ()
modify v i f = MV.modify v f i


--------------------------------------------------------------------------------

-- | For debugging purposes, dump the values and the cells
dump :: DLListMonad s a (V.Vector a, V.Vector Cell)
dump = do DLList v cs <- ask
          cs' <- V.freeze cs
          pure (v,cs')