{- Copyright (c) 2011 Robert Henderson
This source file is distributed under the terms of a BSD3-style
license, which can be found in the file LICENSE at the root of
this package. -}

-- | 'MultiMap' data structure: similar to a map ("Data.Map"), but allows
-- multiple values with the same key.
--
module Data.MultiMap.Rosso1
    (MultiMap

    ,fromList
    ,toList

    ,empty
    ,singleton
    ,insert
    ,insertMany
    ,insertList
    ,insertManyLists

    ,null
    ,lookup
    ,deleteList
    ,extractList
    ,extractEachListWithKey
    ,alter

    ,maxView
    ,elems
    ,descElems
    ,assocs
    ,descAssocs

    ) where

{-- Notes --

The only other multimap implementation that I know of on Hackage is
part of the 'Holumbus' system. The Holumbus multimap has different
semantics to this one, as it maps each key to a /set/ of values rather
than a list.

Todo:
 () Add Big-O time complexities to the documentation.

-}

-----------------------------------------------------------------
import Prelude hiding (null, lookup)
import qualified Prelude

import Data.Maybe.Rosso1
import Data.Tuple.Rosso1
import Data.Map.Rosso1 (Map)
import qualified Data.Map.Rosso1 as Map



data MultiMap k a = MultiMap (Map k [a])
-- In a valid multimap, the list mapped to by a key is always non-empty.



-- Helper function: converts from 'toList' form to 'assocs' form. Preserves
-- the order of the values.
--
listToAssocs :: [(k, [a])] -> [(k, a)]
listToAssocs lst = do (k, vs) <- lst
                      v <- vs
                      return (k, v)



-- | Converts an association list into a multimap. If the association
-- list contains duplicate keys, then the corresponding lists of
-- values become concatenated.
--
-- >>> fromList [(4, "dca"), (1, "aba"), (2, "b"), (1, "ac"), (3, "")]
-- fromList [(1,"abaac"),(2,"b"),(4,"dca")]
--
fromList :: Ord k => [(k, [a])] -> MultiMap k a
fromList = foldr (uncurry insert) empty . listToAssocs


-- | Converts a multimap into an association list, with the keys in
-- ascending order.
--
-- >>> toList $ fromList [(4, "dca"), (1, "aba"), (2, "b"), (1, "ac"), (3, "")]
-- [(1,"abaac"),(2,"b"),(4,"dca")]
--
toList :: MultiMap k a -> [(k, [a])]
toList (MultiMap m) = Map.assocs m


instance (Show k, Show a) => Show (MultiMap k a) where
    -- Adds parentheses if the precedence of the enclosing context
    -- is greater than that of the top-level constructor (function
    -- application).
    showsPrec d m = showParen (d > 10) $
                    showString "fromList " . shows (toList m)



-- | The empty multimap.
--
empty :: MultiMap k a
empty = MultiMap (Map.empty)


-- | A multimap with a single entry.
--
singleton :: k -> a -> MultiMap k a
singleton k a = MultiMap (Map.singleton k [a])


-- | Inserts a new key-value pair. If other entries already exist
-- with the same key, then the new entry is inserted just before them.
--
-- >>> insert 2 'a' $ fromList [(1, "efg"), (2, "jzw"), (3, "abc")]
-- fromList [(1,"efg"),(2,"ajzw"),(3,"abc")]
--
insert :: Ord k => k -> a -> MultiMap k a -> MultiMap k a
insert k a (MultiMap m) = MultiMap (Map.alter f k m)
    where f Nothing = Just [a]
          f (Just as) = Just (a : as)


-- | Passes down the list from left to right, inserting each entry into
-- the multimap.
--
-- >>> insertMany [(1, 'a'), (5, 'a'), (1, 'a'), (1, 'b')] empty
-- fromList [(1,"baa"),(5,"a")]
--
insertMany :: Ord k => [(k, a)] -> MultiMap k a -> MultiMap k a
insertMany xs mm = foldl (flip . uncurry $ insert) mm xs


-- | Prepends a list of values onto the entry with the given key.
--
-- >>> insertList 7 "hello" $ fromList [(5, "ab"), (7, "efg")]
-- fromList [(5,"ab"),(7,"helloefg")]
--
insertList :: Ord k => k -> [a] -> MultiMap k a -> MultiMap k a
insertList k as (MultiMap m) = MultiMap (Map.alter f k m)
    where f Nothing = Just as
          f (Just as2) = Just (as ++ as2)


-- | Passes down the given list from left to right invoking 'insertList'.
--
insertManyLists :: Ord k => [(k, [a])] -> MultiMap k a -> MultiMap k a
insertManyLists xs mm = foldl (flip . uncurry $ insertList) mm xs



-- | Tests if the multimap is empty.
--
null :: MultiMap k a -> Bool
null (MultiMap m) = Map.null m


-- | Returns the list of values associated with the given key.
--
-- >>> lookup 5 $ fromList [(1, "abc"), (5, "aagf"), (6, "c")]
-- "aagf"
--
lookup :: Ord k => k -> MultiMap k a -> [a]
lookup k (MultiMap m) = Map.findWithDefault [] k m


-- | Deletes all the values associated with the given key.
--
deleteList :: Ord k => k -> MultiMap k a -> MultiMap k a
deleteList k (MultiMap m) = MultiMap (Map.delete k m)


-- | Simultaneous lookup and deleteList.
--
extractList :: Ord k => k -> MultiMap k a -> ([a], MultiMap k a)
extractList k (MultiMap m)
    = pairApply (fromMaybe []) MultiMap (Map.extract k m)


-- | For each key that maps to a non-empty list of values, returns
-- that key and its corresponding values as well as the multimap with
-- those values removed. The keys are enumerated in ascending order.
--
extractEachListWithKey :: Ord k =>
                          MultiMap k a -> [((k, [a]), MultiMap k a)]
extractEachListWithKey m = map f (toList m)
    where f x@(k, _) = (x, deleteList k m)


-- | Modifies the list of values associated with a given key.
--
alter :: Ord k => ([a] -> [a]) -> k -> MultiMap k a -> MultiMap k a
alter f k (MultiMap m) = MultiMap (Map.alter g k m)
    where g = toMaybe (not . Prelude.null) . f . fromMaybe []



-- | Returns 'Nothing' if the multimap is empty, otherwise returns the
-- first value associated with the maximal key of the multimap, and
-- the multimap stripped of that value.
--
-- >>> maxView $ fromList [(1, "ab"), (2, "efg")]
-- Just ('e',fromList [(1,"ab"),(2,"fg")])
--
maxView :: MultiMap k a -> Maybe (a, MultiMap k a)
maxView (MultiMap m) = fmap f (Map.maxView m)
    where f ([], _) = error "maxView: multimap is invalid"
          f ([x], rest) = (x, MultiMap rest)
          f (x : _, _) = (x, MultiMap (Map.updateMax (Just . tail) m))


-- | Returns all of the values in the multimap in ascending order of
-- their keys.
--
-- >>> elems $ fromList [(1, "aba"), (2, "adf"), (3, "z")]
-- "abaadfz"
--
elems :: MultiMap k a -> [a]
elems (MultiMap m) = concat (Map.elems m)


-- | Returns all of the values in the multimap in descending order of
-- their keys. The values are enumerated in the same order as with 'maxView'.
--
-- >>> descElems $ fromList [(1, "aba"), (2, "adf"), (3, "z")]
-- "zadfaba"
--
descElems :: MultiMap k a -> [a]
descElems (MultiMap m) = concat . reverse $ Map.elems m


-- | Returns all of the key-value pairs in the multimap in ascending order
-- of keys.
--
-- >>> assocs $ fromList [(1, "ab"), (4, "cda")]
-- [(1,'a'),(1,'b'),(4,'c'),(4,'d'),(4,'a')]
--
assocs :: MultiMap k a -> [(k, a)]
assocs = listToAssocs . toList


-- | Returns all of the key-value pairs in the multimap in descending order
-- of keys. The values are enumerated in the same order as with 'maxView'.
--
-- >>> descAssocs (fromList [(1, "ab"), (4, "cda")])
-- [(4,'c'),(4,'d'),(4,'a'),(1,'a'),(1,'b')]
--
descAssocs :: MultiMap k a -> [(k, a)]
descAssocs = listToAssocs . reverse . toList



-----------------------------------------------------------
{- UNIT TESTS

*Util.MultiMap1e> fromList []
fromList []
*Util.MultiMap1e> fromList [(3, ""), (1, "")]
fromList []
*Util.MultiMap1e> fromList [(4, "dca"), (1, "aba"), (2, "b"), (1, "ac"), (3, "")]
fromList [(1,"abaac"),(2,"b"),(4,"dca")]

*Util.MultiMap1e> toList (fromList [(1,"abaac"),(2,"b"),(4,"dca")])
[(1,"abaac"),(2,"b"),(4,"dca")]
*Util.MultiMap1e> toList (fromList [(4, "dca"), (1, "aba"), (2, "b"), (1, "ac"), (3, "")])
[(1,"abaac"),(2,"b"),(4,"dca")]

*Util.MultiMap1e> empty
fromList []
*Util.MultiMap1e> singleton 1 'a'
fromList [(1,"a")]

*Util.MultiMap1e> insert 1 'a' empty
fromList [(1,"a")]
*Util.MultiMap1e> insert 2 'a' $ fromList [(1, "efg"), (3, "abc")]
fromList [(1,"efg"),(2,"a"),(3,"abc")]
*Util.MultiMap1e> insert 2 'a' $ fromList [(1, "efg"), (2, "jzw"), (3, "abc")]
fromList [(1,"efg"),(2,"ajzw"),(3,"abc")]
*Util.MultiMap1e> insert 1 'a' $ insert 5 'a' $ insert 1 'a' $ insert 1 'b' $ empty
fromList [(1,"aab"),(5,"a")]

*Util.MultiMap1e> insertMany [(1, 'a'), (5, 'a'), (1, 'a'), (1, 'b')] empty
fromList [(1,"baa"),(5,"a")]
*Util.MultiMap1e> insertMany (zip (cycle "abc") [1..8]) empty
fromList [('a',[7,4,1]),('b',[8,5,2]),('c',[6,3])]
*Util.MultiMap1e> insertMany (zip (cycle "abc") [1..8]) $ fromList [('b', [2, 1])]
fromList [('a',[7,4,1]),('b',[8,5,2,2,1]),('c',[6,3])]

*Rosso.MultiMap1e> insertList 3 "hello" $ fromList [(5, "ab"), (7, "efg")]
fromList [(3,"hello"),(5,"ab"),(7,"efg")]
*Rosso.MultiMap1e> insertList 7 "hello" $ fromList [(5, "ab"), (7, "efg")]
fromList [(5,"ab"),(7,"helloefg")]

*Rosso.MultiMap1e> insertManyLists [(5, "abcd"), (3, "xxy"), (5, "fa")] $ fromList [(1, "ab"), (5, "z")]
fromList [(1,"ab"),(3,"xxy"),(5,"faabcdz")]

*Rosso.MultiMap1e> null empty
True
*Rosso.MultiMap1e> null $ fromList [(3, "")]
True
*Rosso.MultiMap1e> null $ fromList [(3, "abc")]
False

*Rosso.MultiMap1e> lookup 3 empty
[]
*Rosso.MultiMap1e> lookup 5 $ fromList [(1, "abc"), (5, "aagf"), (6, "c")]
"aagf"
*Rosso.MultiMap1e> lookup 2 $ fromList [(1, "abc"), (5, "aagf"), (6, "c")]
""

*Rosso.MultiMap1e> deleteList 3 $ fromList [(2, "abc"), (5, "ef")]
fromList [(2,"abc"),(5,"ef")]
*Rosso.MultiMap1e> deleteList 2 $ fromList [(2, "abc"), (5, "ef")]
fromList [(5,"ef")]

*Rosso.MultiMap1e> extractList 3 $ fromList [(2, "abc"), (5, "ef")]
("",fromList [(2,"abc"),(5,"ef")])
*Rosso.MultiMap1e> extractList 2 $ fromList [(2, "abc"), (5, "ef")]
("abc",fromList [(5,"ef")])

*Rosso.MultiMap1e> extractEachListWithKey empty
[]
*Rosso.MultiMap1e> extractEachListWithKey $ singleton 3 'a'
[((3,"a"),fromList [])]
*Rosso.MultiMap1e> extractEachListWithKey $ fromList [(2, "abc"), (5, "e"), (6, "zzzz")]
[((2,"abc"),fromList [(5,"e"),(6,"zzzz")]),((5,"e"),fromList [(2,"abc"),(6,"zzzz")]),((6,"zzzz"),fromList [(2,"abc"),(5,"e")])]

*Rosso.MultiMap1e> alter reverse 2 $ fromList [(2, "abc"), (5, "ef")]
fromList [(2,"cba"),(5,"ef")]
*Rosso.MultiMap1e> alter tail 5 $ fromList [(2, "abc"), (5, "e")]
fromList [(2,"abc")]
*Rosso.MultiMap1e> alter reverse 6 $ fromList [(2, "abc"), (5, "e")]
fromList [(2,"abc"),(5,"e")]
*Rosso.MultiMap1e> alter (++ "hello") 6 $ fromList [(2, "abc"), (5, "e")]
fromList [(2,"abc"),(5,"e"),(6,"hello")]

*Util.MultiMap1e> maxView empty
Nothing
*Util.MultiMap1e> maxView $ singleton 1 'a'
Just ('a',fromList [])
*Util.MultiMap1e> maxView $ fromList [(1, "ab"), (2, "g")]
Just ('g',fromList [(1,"ab")])
*Util.MultiMap1e> maxView $ fromList [(1, "ab"), (2, "efg")]
Just ('e',fromList [(1,"ab"),(2,"fg")])

*Util.MultiMap1e> elems empty
[]
*Util.MultiMap1e> elems $ fromList [(1, "aba"), (2, "adf"), (3, "z")]
"abaadfz"

*Util.MultiMap1e> descElems empty
[]
*Util.MultiMap1e> descElems $ fromList [(1, "aba"), (2, "adf"), (3, "z")]
"zadfaba"

*Util.MultiMap1e> assocs empty
[]
*Util.MultiMap1e> assocs (fromList [(1, "ab"), (4, "cda")])
[(1,'a'),(1,'b'),(4,'c'),(4,'d'),(4,'a')]

*Util.MultiMap1e> descAssocs empty
[]
*Util.MultiMap1e> descAssocs (fromList [(1, "ab"), (4, "cda")])
[(4,'c'),(4,'d'),(4,'a'),(1,'a'),(1,'b')]

-}