{-# LANGUAGE Haskell2010 #-}

-- | A very simple Multimap, based on 'Data.Map.Map' from the containers package.
module Data.Multimap (

    -- * Multimap type
    Multimap,

    -- * Query
    null,
    size,
    numKeys,
    numValues,

    member,
    notMember,
    lookup,

    -- * Operators
    (!),

    -- * Construction
    empty,
    
    -- ** Insertion
    insert,

    -- ** Delete
    delete,

    -- * Traversal
    map,
    mapKeys,
    mapWithKey,
    
    -- * Folds
    foldr,
    foldl,
    foldrWithKey,
    foldlWithKey,

    -- * Conversion
    elems,
    keys,
    keysSet,
    assocs,

    toMap,
    toMapOfSets,
    toList,
    fromList,
    
    -- * Min/Max
    findMin,
    findMax,
    findMinWithValues,
    findMaxWithValues

  ) where

import Prelude hiding (lookup, map, null, foldr, foldl)
import qualified Prelude as P

import qualified Data.Set as Set
import Data.Set (Set)

import qualified Data.Map as Map
import Data.Map (Map)

import Data.Word

-- | A Multimap with keys @k@ and values @v@.
--
-- A key can have multiple values (but not zero).
-- The same value can be added multiple times (thus no
-- constraints are ever imposed on @v@).
--
-- Internally this is simply a @Map k [v]@.
-- See 'toMap' for accessing the underlying 'Map'.
newtype Multimap k v = Multimap (Word32, Word32, Map k [v])


null :: Multimap k a -> Bool
-- ^ /O(1)./ Check whether the multimap is empty or not.
null (Multimap (_, _, m)) = Map.null m

size :: Multimap k a -> Int
-- ^ /O(1)./ The number of elements in the multimap.
size (Multimap (_, size, _)) = fromIntegral size

numKeys :: Multimap k a -> Word32
-- ^ /O(1)./ The number of keys in the multimap.
-- 
-- As this is a multimap, the number of keys is not
-- necessarily equal to the number of values.
numKeys (Multimap (nk, _, _)) = nk

numValues :: Multimap k a -> Word32
-- ^ /O(1)./ The number of values in the multimap.
--
-- As this is a multimap, the number of keys is not
-- necessarily equal to the number of values.
numValues (Multimap (_, nv, _)) = nv

notMember, member :: Ord k => Multimap k a -> k -> Bool
-- | /O(log n)./ Is the key a member of the multimap?
member (Multimap (_, _, map)) key = Map.member key map
-- | /O(log n)./ Is the key not a member of the multimap?
notMember key = not . member key

(!) :: Ord k => k -> Multimap k a -> [a]
(!) = lookup

lookup :: Ord k => k -> Multimap k a -> [a]
-- ^ /O(log n)./ Lookup the value at a key in the map.
--
-- The function will return the corrsponding values as a List,
-- or the empty list if no values are associated witht the
-- given key.
lookup key (Multimap (_, _, map)) = maybe [] id (Map.lookup key map)


empty :: Multimap k a
-- ^ /O(1)./ The empty multimap.
empty = Multimap (0, 0, Map.empty)


insert :: Ord k => k -> a -> Multimap k a -> Multimap k a
-- ^ /O(log n)./ Insert a new key and value in the map.
insert k v (Multimap (nk, nv, map))
    | Map.member k map = Multimap (nk, succ nv, Map.insert k (v : map Map.! k) map)
    | otherwise = Multimap (succ nk, succ nv, Map.insert k [v] map)

delete :: Ord k => k -> a -> Multimap k a -> Multimap k a
-- ^ /O(log n)./ Delete a key and all its values from the map.
delete k v m@(Multimap (nk, nv, map)) = case Map.lookup k map of
    Just v -> Multimap (pred nk, nv - fromIntegral (length v), Map.delete k map)
    _      -> m

map :: (a -> b) -> Multimap k a -> Multimap k b
-- ^ Map a function over all values in the map.
map f (Multimap (nk, nv, map)) = Multimap (nk, nv, Map.map (P.map f) map)

mapKeys :: Ord k2 => (k1 -> k2) -> Multimap k1 a -> Multimap k2 a
-- ^ mapKeys f s is the multimap obtained by applying f to each key of s.
mapKeys f (Multimap (nk, nv, map)) = Multimap (nk, nv, Map.mapKeys f map)

mapWithKey :: (k -> a -> b) -> Multimap k a -> Multimap k b
-- ^ Map a function over all key/value pairs in the map.
mapWithKey f (Multimap (nk, nv, map))
  = Multimap (nk, nv, Map.mapWithKey (\k -> P.map (f k)) map)


foldr :: (a -> b -> b) -> b -> Multimap k a -> b
foldr f e = P.foldr f e . concat . elems

foldl :: (a -> b -> a) -> a -> Multimap k b -> a
foldl f e = P.foldl f e . concat . elems

foldrWithKey :: (k -> a -> b -> b) -> b -> Multimap k a -> b
foldrWithKey f e = P.foldr (uncurry f) e . toList

foldlWithKey :: (a -> k -> b -> a) -> a -> Multimap k b -> a
foldlWithKey f e = P.foldl (\a (k,v) -> f a k v) e . toList


elems :: Multimap k a -> [[a]]
-- ^ /O(n)./ Return all elements of the multimap in the
-- ascending order of their keys.
--
-- A list of lists is returned since a key can have
-- multiple values. Use 'concat' to flatten.
elems (Multimap (_, _, map)) = Map.elems map

keys :: Multimap k a -> [k]
-- ^ /O(n)./ Return all keys of the multimap in ascending order.
keys (Multimap (_, _, map)) = Map.keys map

keysSet :: Multimap k a -> Set k
-- ^ /O(n)./ The set of all keys of the multimap.
keysSet (Multimap (_, _, map)) = Map.keysSet map

assocs :: Multimap k a -> [(k, [a])]
-- ^ /O(n)./ Return all key/value pairs in the multimap
-- in ascending key order.
assocs (Multimap (_, _, map)) = Map.assocs map


toMap :: Multimap k a -> Map k [a]
-- ^ /O(1)./ Return the map of lists.
toMap (Multimap (_, _, theUnderlyingMap)) = theUnderlyingMap

toMapOfSets :: Ord a => Multimap k a -> Map k (Set a)
-- ^ /O(k*m*log m) where k is the number of keys and m the
-- maximum number of elements associated with a single key/
toMapOfSets (Multimap (_, _, map)) = Map.map Set.fromList map

toList :: Multimap k a -> [(k, a)]
-- ^ Return a flattened list of key/value pairs.
toList (Multimap (_, _, map))
  = concat $ Map.elems $ Map.mapWithKey (\k -> zip (repeat k)) map

fromList :: Ord k => [(k, a)] -> Multimap k a
-- ^ /O(n*log n)/ Create a multimap from a list of key/value pairs.
--
-- > fromList xs == foldr (uncurry insert) empty
fromList = P.foldr (uncurry insert) empty


findMin :: Multimap k a -> Maybe k
-- ^ /O(log n)/ Find the minimal key of the multimap.
findMin (Multimap (_, _, map))
    | Map.null map = Nothing
    | otherwise    = Just $ fst $ Map.findMin map

findMax :: Multimap k a -> Maybe k
-- ^ /O(log n)/ Find the maximal key of the multimap.
findMax (Multimap (_, _, map))
    | Map.null map = Nothing
    | otherwise    = Just $ fst $ Map.findMax map

findMinWithValues :: Multimap k a -> Maybe (k, [a])
-- ^ /O(log n)/ Find the minimal key and the values associated with it.
findMinWithValues (Multimap (_, _, map))
    | Map.null map = Nothing
    | otherwise    = Just $ Map.findMin map

findMaxWithValues :: Multimap k a -> Maybe (k, [a])
-- ^ /O(log n)/ Find the maximal key and the values associated with it.
findMaxWithValues (Multimap (_, _, map))
    | Map.null map = Nothing
    | otherwise    = Just $ Map.findMax map