{-# LANGUAGE CPP #-}

-- | 
-- A random-access list implementation based on Chris Okasaki's approach
-- on his book \"Purely Functional Data Structures\", Cambridge University
-- Press, 1998, chapter 9.3.
--
-- It provides random-access ('lookup', 'update', etc.) in /O(log n)/ while
-- the list functionality 'head', 'tail' and 'cons' still works in /O(1)/.
--
-- A 'RandomAccessList' uses 'Int's for effective indexing. The valid index
-- range of a 'RandomAccessList' of size @n@ is @[0 .. n-1]@. If an index is
-- out of range, an 'error' is raised.
--
-- This module is best imported @qualified@ in order to prevent name clashes
-- with other modules.

module Data.RandomAccessList
	(
	-- * The RandomAccessList type
	  RandomAccessList
	-- * Operators
	, (.!.), (.:.), (.+.)
	-- * Construction
	, empty, singleton, replicate
	-- * Query
	, null, isEmpty, size, member, index
	-- * List specific operations
	, head, tail, extractHead, cons, append, zip, zipWith
	-- * Random-access specific operations
	, lookup, update, adjust, adjustLookup
	-- * Conversion
	-- ** List
	, fromList, toList, toIndexedList
	-- ** Map
	, toMap, toIntMap
	-- ** Array
	, fromArray, toArray
	) where

-- TODO
--
-- error handling, messages
--
-- strict tuples, maybe using {-# UNPACK #-}?
-- langage pragmas fromList/toList, etc.
--
-- filter, subranges (improve zip when done), insert, delete
-- sort, fold, zip, reverse, ...!

import Data.Array.IArray (IArray, Ix, array, elems)
import Data.Foldable (Foldable(foldr), foldr')
import Data.IntMap as IMap (IntMap, fromDistinctAscList)
import Data.Map as Map (Map, fromDistinctAscList)
import Data.Maybe
import Data.Monoid
import qualified Prelude (replicate, zip, zipWith)
import Prelude hiding (head, lookup, null, replicate, tail, zip, zipWith)
import Text.Read

-- |
-- Random-access lists allowing /O(1)/ list operations and /O(log n)/
-- indexed access.
data RandomAccessList a
	= RandomAccessList {-# UNPACK #-} !Int ![(Int, CBTree a)]

instance (Show a) => Show (RandomAccessList a) where
	show list = "fromList " ++ show (toList list)

instance (Read a) => Read (RandomAccessList a) where
#ifdef __GLASGOW_HASKELL__
	readPrec = parens $ prec 10 $ do
		Ident "fromList" <- lexP
		xs <- readPrec
		return (fromList xs)
	readListPrec = readListPrecDefault
#else
	readsPrec p = readParen (p > 10) $ \r -> do
		("fromList", s) <- lex r
		(xs, t) <- reads s
		return (fromList xs, t)
#endif

instance (Eq a) => Eq (RandomAccessList a) where
	(RandomAccessList size1 list1) == (RandomAccessList size2 list2)
		= size1 == size2   -- no need to compare weights after this
		&& (map snd list1) == (map snd list2)

instance (Ord a) => Ord (RandomAccessList a) where
	compare list1 list2 -- TODO: efficiency?
		| isEmpty list1 = if isEmpty list2 then EQ else LT
		| isEmpty list2 = GT
		| otherwise     = let
			(h1, t1) = extractHead list1
			(h2, t2) = extractHead list2
			in
			case compare h1 h2 of
				EQ -> compare t1 t2
				x  -> x

instance Monoid (RandomAccessList a) where
	mempty  = empty
	mappend = append

instance Functor RandomAccessList where
	fmap f (RandomAccessList s list) = RandomAccessList s
		(fmap (\(weight, tree) -> (weight, fmap f tree)) list)

instance Foldable RandomAccessList where
	foldr f zero list = foldr' f zero (toList list)

-- |
-- Complete binary trees. It is not possible to ensure completeness via
-- the type system, so make sure you do!
data CBTree a
	= Leaf a
	| Node a !(CBTree a) !(CBTree a)
	deriving (Eq)

instance Functor CBTree where
	fmap f (Leaf x)            = Leaf (f x)
	fmap f (Node x left right) = Node (f x) (fmap f left) (fmap f right)

-- |
-- Generate a readable error message for out-of-range index access.
indexErrorMessage :: Maybe (Int, Int) -- ^ Accessed index and list size if available.
	-> String                     -- ^ The error message.
indexErrorMessage = maybe
	"RandomAccessList: index out of range"
	(\(i, s) -> "RandomAccessList: index " ++ show i
		++ " out of range [0.." ++ show (s - 1) ++ "]")

-- |
-- /O(log n)/. Retrieve the /i/th element of the list. Unless
-- /0 <= i < n/, an 'error' is raised, see 'lookup'.
(.!.) :: RandomAccessList a -> Int -> a
(.!.) = flip lookup
infixr 9 .!.

-- |
-- /O(1)/. Prepend an element to the 'RandomAccessList', see 'cons'.
(.:.) :: a -> RandomAccessList a -> RandomAccessList a
(.:.) = cons
infixr 5 .:.

-- |
-- /O(n)/ where /n/ is the size of the first list. Appends two
-- 'RandomAccessList's, see 'append'.
(.+.) :: RandomAccessList a -> RandomAccessList a -> RandomAccessList a
(.+.) = append
infixr 5 .+.

-- |
-- /O(1)/. Builds an empty 'RandomAccessList'.
empty :: RandomAccessList a
empty = RandomAccessList 0 []

-- |
-- /O(1)/. Builds a singleton 'RandomAccessList'.
singleton :: a -> RandomAccessList a
singleton x = RandomAccessList 1 [(1, Leaf x)]

-- |
-- /O(n)/. @'replicate' n x@ constructs a 'RandomAccessList' that
-- contains the same element @x@ @n@ times.
replicate :: Int -> a -> RandomAccessList a
replicate n = fromList . (Prelude.replicate n)

-- |
-- /O(1)/. Is the 'RandomAccessList' empty?
null :: RandomAccessList a -> Bool
null = isEmpty

-- |
-- /O(1)/. Is the 'RandomAccessList' empty?
isEmpty :: RandomAccessList a -> Bool
isEmpty (RandomAccessList 0 _) = True
isEmpty _                      = False

-- |
-- /O(1)/. The number of elements contained in a 'RandomAccessList'.
size :: RandomAccessList a -> Int
size (RandomAccessList s _) = s

-- |
-- /O(n)/. Is the given element a member of the 'RandomAccessList'?
member :: (Eq a) => a -> RandomAccessList a -> Bool
member x = isJust . (index x)

-- |
-- /O(n)/. Find the index of a given element. If the element is not
-- a member of the 'RandomAccessList', this function will 'fail' or
-- @'return' index@ otherwise.
index :: (Eq a, Monad m) => a -> RandomAccessList a -> m Int
index x (RandomAccessList _ list) = index' 0 list
	where
	index' _     []                     = fail "RandomAccessList: element not found"
	index' accum ((weight, tree):trees) = maybe
		(index' (accum + weight) trees)
		(return . (+accum))
		(indexTree weight x tree)

-- |
-- Find the index of a given element in a 'CBTree'. Returns @'Just'
-- index@ on success and 'Nothing' on failure.
indexTree :: (Eq a) => Int -> a -> CBTree a -> Maybe Int
indexTree _ x (Leaf e) = if x == e then Just 0 else Nothing
indexTree weight x (Node e left right) = if x == e
	then Just 0
	else case indexTree halfWeight x left of
		Just i  -> Just (i + 1)
		Nothing -> case indexTree halfWeight x right of
			Just index' -> Just (index' + 1 + halfWeight)
			Nothing     -> Nothing
	where
	halfWeight = weight `div` 2

-- |
-- /O(1)/. Returns the head of a 'RandomAccessList'.
head :: RandomAccessList a -> a
head = fst . extractHead

-- |
-- /O(1)/. Retrieve the tail of a 'RandomAccessList'.
tail :: RandomAccessList a -> RandomAccessList a
tail = snd . extractHead

-- |
-- /O(1)/. Retrieve both, 'head' and 'tail' of a 'RandomAccessList'.
extractHead :: RandomAccessList a -> (a, RandomAccessList a)
extractHead (RandomAccessList s list) = case list of
	[]                                     -> error "RandomAccessList: empty list"
	((_, Leaf x) : trees')                 -> (x, RandomAccessList (s - 1) trees')
	((weight, Node x left right) : trees') -> let
		halfWeight = weight `div` 2
		trees''    = (halfWeight, left) : (halfWeight, right) : trees'
		in
		(x, RandomAccessList (s - 1) trees'')

-- |
-- /O(1)/. Prepend an element to the 'RandomAccessList'.
cons :: a -> RandomAccessList a -> RandomAccessList a
cons x (RandomAccessList s list) = RandomAccessList (s + 1) (case list of
	xss@((size1, tree1) : (size2, tree2) : xs) -> if size1 == size2
		then (1 + 2 * size1, Node x tree1 tree2) : xs
		else (1,             Leaf x)             : xss
	xs                                         -> (1, Leaf x) : xs)

-- |
-- /O(n)/ where /n/ is the 'size' of the first list. Appends the second
-- list to the first list.
append :: RandomAccessList a -> RandomAccessList a -> RandomAccessList a
append = (flip (foldr' cons))

-- |
-- /O(max(n, m))/. List-like 'Prelude.zip'. This function is faster
-- when called with two 'RandomAccessList's of equal size.
zip :: RandomAccessList a -> RandomAccessList b -> RandomAccessList (a, b)
zip = zipWith (,)

-- |
-- /O(max(n, m))/. List-like 'Prelude.zipWith'. This function is faster
-- when called with two 'RandomAccessList's of equal size.
zipWith :: (a -> b -> c) -> RandomAccessList a -> RandomAccessList b -> RandomAccessList c
zipWith f l1@(RandomAccessList size1 as) l2@(RandomAccessList size2 bs) = if size1 == size2
	then RandomAccessList size1 $ zip' as bs
	else fromList (Prelude.zipWith f (toList l1) (toList l2))
	where
	zip' []                          []                          = []
	zip' ((weightA, treeA) : treesA) ((weightB, treeB) : treesB) = if weightA == weightB
		then (weightA, zipWithTree f treeA treeB) : zip' treesA treesB
		else error skewError
	zip' _                           _                           = error skewError
	skewError = "RandomAccessList: invalid skew binary number representation"

-- |
-- Performs 'zipWith' for complete binary trees.
zipWithTree :: (a -> b -> c) -> CBTree a -> CBTree b -> CBTree c
zipWithTree f (Leaf x1)              (Leaf x2)              = Leaf (f x1 x2)
zipWithTree f (Node x1 left1 right1) (Node x2 left2 right2) = Node
	(f x1 x2)
	(zipWithTree f left1 left2)
	(zipWithTree f right1 right2)
zipWithTree _ _                      _                      = error
	"RandomAccessList: invalid skew binary number representation"

-- |
-- /O(log n)/. Retrieve the /i/th element of the list. Unless
-- /0 <= i < n/, an 'error' is raised.
lookup :: Int -> RandomAccessList a -> a
lookup i = fst . (adjustLookup undefined i)

-- |
-- /O(log n)/. Set the /i/th element of the list. Unless
-- /0 <= i < n/, an 'error' is raised.
update :: Int -> a -> RandomAccessList a -> RandomAccessList a
update i x = snd . (adjustLookup (const x) i)

-- |
-- /O(log n)/. Adjust /i/th element of the list according to the
-- given function. Unless /0 <= i < n/, an 'error' is raised.
adjust :: (a -> a) -> Int -> RandomAccessList a -> RandomAccessList a
adjust f i = snd . (adjustLookup f i)

-- |
-- /O(log n)/. Find the /i/th element of the list and change it. This
-- function returns the element that is at index /i/ in the original
-- 'RandomAccessList' and a new 'RandomAccessList' with the /i/th
-- element replaced according to the given function:
--
-- @
--    lookup   index list === fst (adjustLookup undefined index list)
--    adjust f index list === snd (adjustLookup f         index list)
-- @
--
-- Unless /0 <= i < n/, an 'error' is raised.
adjustLookup :: (a -> a)           -- ^ Modifying element function.
	-> Int                     -- ^ Index of the affected element.
	-> RandomAccessList a      -- ^ 'RandomAccessList' to be modified.
	-> (a, RandomAccessList a) -- ^ Original element and modified 'RandomAccessList'.
adjustLookup f i (RandomAccessList s list) = let
	(x, list') = updateLookup' i list
	in
	(x, RandomAccessList s list')
	where
	updateLookup' _      []                    =
		error $ indexErrorMessage (Just (i, s))
	updateLookup' i' ((weight, tree) : xs) = if i' < weight
		then let
			(x, tree') = adjustLookupTree f weight i' tree
			in
			(x, (weight, tree') : xs)
		else let (x, trees) = updateLookup' (i' - weight) xs
			in
			(x, (weight, tree) : trees)

-- |
-- Performs 'adjustLookup' on a complete binary tree. Note that it
-- is crucial, that the 'CBTree' really is complete!
adjustLookupTree :: (a -> a) -- ^ Function to change a given element.
	-> Int               -- ^ Size of the complete binary tree.
	-> Int               -- ^ Index to be replaced.
	-> CBTree a          -- ^ The tree.
	-> (a, CBTree a)
adjustLookupTree f 1      0 (Leaf x) = (x, Leaf (f x))
adjustLookupTree _ _      _ (Leaf _) = error (indexErrorMessage Nothing)
adjustLookupTree f weight i (Node x left right)
	| i == 0          = (x, Node (f x) left right)
	| i <= halfWeight = let
		(x', left') = adjustLookupTree f halfWeight (i - 1) left
		in
		(x', Node x left' right)
	| otherwise       = let
		(x', right') = adjustLookupTree f halfWeight (i - 1 - halfWeight) right
		in
		(x', Node x left right')
	where
	halfWeight = weight `div` 2

-- |
-- /O(n)/. Build a 'RandomAccessList' from a list.
fromList :: [a] -> RandomAccessList a
fromList = foldr' cons empty -- just as fast as building the trees by hand!

-- |
-- /O(n)/. Convert a 'RandomAccessList' to a list.
toList :: RandomAccessList a -> [a]
toList list
	| isEmpty list = []
	| otherwise    = let (x, xs) = extractHead list in x : toList xs

-- |
-- /O(n)/. Convert a 'RandomAccessList' to a list of tuples each holding
-- an element and its index. The list is ordered ascending regarding the
-- indices.
toIndexedList :: RandomAccessList a -> [(Int, a)]
toIndexedList list = Prelude.zip [0..] (toList list)

-- |
-- /O(n)/. Build a 'Map.Map' from a 'RandomAccessList'. The keys in the
-- 'Map.Map' are the indices of the elements in the 'RandomAccessList'.
toMap :: RandomAccessList a -> Map Int a
toMap = Map.fromDistinctAscList . toIndexedList

-- |
-- /O(n)/. Build an 'IMap.IntMap' from a 'RandomAccessList'. The keys in the
-- 'IMap.IntMap' are the indices of the elements in the 'RandomAccessList'.
toIntMap :: RandomAccessList a -> IntMap a
toIntMap = IMap.fromDistinctAscList . toIndexedList

-- |
-- /O(n)/. Given an 'IArray', generate a 'RandomAccessList'. The elements'
-- order will be preserved.
fromArray :: (IArray a e, Ix i) => a i e -> RandomAccessList e
fromArray = fromList . elems

-- |
-- /O(n)/. Build an 'IArray' from the 'RandomAccessList'. It will have
-- an index range from @[0 .. n-1]@, where @n@ is the 'RandomAccessList's
-- 'size'.
toArray :: (IArray a e) => RandomAccessList e -> a Int e
toArray list = array (0, size list - 1) (toIndexedList list)