{-| Module  : WeakSets
Description : A `WeakMap` is a Data.Map which does not require the keys to implement the Ord typeclass.
Copyright   : Guillaume Sabbagh 2022
License     : LGPL-3.0-or-later
Maintainer  : guillaumesabbagh@protonmail.com
Stability   : experimental
Portability : portable

A `WeakMap` is a Data.Map which does not require the keys to implement the Ord typeclass. It is a weak set of pairs (key,value).

The datatype only assumes its keys are equatable, it is therefore slower to access data than the Data.Map datatype.

We use this datatype because most of the datatypes we care about are not orderable.

Almost all Data.WeakMap functions are implemented so that you can replace a Data.Map import such as 

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

by a Data.WeakMap import such as

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

without breaking anything in your code.

The only functions for which this would fail are the functions converting maps back into list (they require the Eq typeclass unlike in Data.Map). `size` is one of them.

If a function really requires the Ord typeclass to even make sense, then it is not defined in this package, you should use Data.Map.

Note that, just like in Data.Map, the implementation is generally left-biased. Functions that take two maps as arguments and combine them, such as union and intersection, prefer the entries in the first argument to those in the second. 

Functions with non colliding names are defined in Data.WeakMap.Safe. Inline functions are written between pipes @|@.

This module is intended to be imported qualified, to avoid name clashes with Prelude functions, except for functions in Data.WeakSet.Map, e.g.

> import Data.WeakMap (Map)
> import qualified Data.WeakMap as Map
> import Data.WeakMap.Safe

Unlike Data.Map, we defer the removing of duplicate keys to the conversion back to a list. 

Beware if the map is supposed to contain a lot of duplicate keys, you should purge them yourself by transforming the map into a list and back into a map. The time complexity is always given in function of the number of pairs in the map including the duplicate pairs.
-}

module Data.WeakMap
(
    -- * Map type
    AssociationList(..)
    , Map

    -- * Construction
    , weakMap
    , weakMapFromSet
    , empty
    , singleton
    , fromSet

    -- ** From Unordered Lists
    , fromList
    , fromListWith
    , fromListWithKey

    -- ** From Ascending Lists
    , fromAscList
    , fromAscListWith
    , fromAscListWithKey
    , fromDistinctAscList

    -- ** From Descending Lists
    , fromDescList
    , fromDescListWith
    , fromDescListWithKey
    , fromDistinctDescList

    -- * Insertion
    , insert
    , insertWith
    , insertWithKey
    , insertLookupWithKey
    , insertMaybe

    -- * Deletion\/Update
    , delete
    , adjust
    , adjustWithKey
    , update
    , updateWithKey
    , updateLookupWithKey
    , alter
    , alterF

    -- * Query
    -- ** Lookup
    , lookup
    , (!?)
    , (!)
    , (|?|), (|!|)
    , findWithDefault
    , member
    , notMember

    -- ** Size
    , size
    , null

    -- * Combine

    -- ** Union
    , union
    , unionWith
    , unionWithKey
    , unions
    , unionsWith

    -- ** Difference
    , difference
    , (\\)
    , differenceWith
    , differenceWithKey

    -- ** Intersection
    , intersection
    , intersectionWith
    , intersectionWithKey

    -- ** Disjoint
    , disjoint

    -- -- ** Compose
    , (|.|)
    , compose

    , mergeWithKey

    -- * Traversal
    -- ** Map
    , map
    , mapWithKey
    , traverseWithKey
    , traverseMaybeWithKey
    , mapAccum
    , mapAccumWithKey
    , mapAccumRWithKey
    , mapKeys
    , mapKeysWith
    , mapKeysMonotonic

    -- * Folds
    , foldr
    , foldl
    , foldrWithKey
    , foldlWithKey
    , foldMapWithKey

    -- ** Strict folds
    , foldr'
    , foldl'
    , foldrWithKey'
    , foldlWithKey'

    -- * Conversion
    , mapToList
    , mapToSet
    , elems
    , elems'
    , values
    , image
    , keys
    , keys'
    , domain
    , assocs
    , keysSet
    , elemsSet

    -- ** Lists
    , toList

    -- ** Ordered lists
    , toAscList
    , toDescList

    -- * Filter
    , filter
    , filterWithKey
    , restrictKeys
    , withoutKeys
    , partition
    , partitionWithKey

    , mapMaybe
    , mapMaybeWithKey
    , mapEither
    , mapEitherWithKey

    -- * Submap
    , isSubmapOf, isSubmapOfBy
    , isProperSubmapOf, isProperSubmapOfBy

    -- * Indexed
    , lookupIndex
    , findIndex
    , elemAt
    , updateAt
    , deleteAt
    , take
    , drop
    , splitAt

    -- * Others
    , idFromSet
    , memorizeFunction
    , inverse
    , pseudoInverse
    , enumerateMaps
) where
import           Prelude  hiding      (lookup, map, filter, drop, take, splitAt, foldr, foldl, null)
import           Data.WeakSet         (Set)
import qualified Data.WeakSet       as Set
import           Data.WeakSet.Safe
import qualified Data.List          as L
import           Data.Maybe           (fromJust)
import           Control.Applicative  (liftA3)
import qualified Data.Foldable      as Foldable

-- | An association list is a list of pairs (key,value).
type AssociationList k v = [(k,v)]

-- | A weak map is a weak set of pairs (key,value) such that their should only be one pair with a given key.
--
-- It is an abstract type, the smart constructor is `weakMap`.
data Map k v = Map {-# UNPACK #-} !(Set (k,v))

instance (Eq k, Eq v) => Eq (Map k v) where
    m1 == m2 = (mapToSet m1) == (mapToSet m2)

instance (Show k, Show v) => Show (Map k v) where
    show (Map al) = "(weakMap "++ init rest ++")"
        where
            ('(':'s':'e':'t':' ':rest)= show al

instance Semigroup (Map k v) where
    (Map al1) <> (Map al2) = Map $ al1 <> al2

instance Monoid (Map k v) where
    mempty = Map (set [])

instance Functor (Map k) where
    fmap f (Map al) = Map $ (\(k,v) -> (k,f v)) <$> al


-- Construction

-- | /O(1)/. The smart constructor of weak maps. This is the only way of instantiating a `Map`.
--
-- Takes an association list and returns a function which maps to each key the value associated.
--
-- If several pairs have the same keys, they are kept until the user wants an association list back.
weakMap :: AssociationList k v -> Map k v
weakMap al = Map $ set $ al

-- | /O(1)/. Construct a 'Map' from a 'Set' of pairs (key,value).
weakMapFromSet :: Set (k,v) -> Map k v
weakMapFromSet s = Map s

-- | /O(1)/. Alias of `weakMap` for backward compatibility with Data.Map.
fromList :: AssociationList k v -> Map k v
fromList = weakMap

-- | Alias of mempty for backward compatibility with Data.Map.
empty :: Map k a
empty = mempty

-- | /O(1)/. A map with a single pair (key,value).
singleton :: k -> a -> Map k a
singleton k v = Map $ set [(k,v)]

-- | /O(n)/.  Build a map from a set of keys and a function which for each key computes its value.
fromSet :: (k -> a) -> Set k -> Map k a
fromSet f s = Map $ (\x -> (x, f x)) <$> s

-- | /O(n)/. Build a map from a list of key/value pairs with a combining function. 
fromListWith :: (Eq k) => (a -> a -> a) -> [(k, a)] -> Map k a 
fromListWith f xs = fromListWithKey (\_ x y -> f x y) xs

-- | /O(n)/. Build a map from a list of key/value pairs with a combining function. 
fromListWithKey :: (Eq k) => (k -> a -> a -> a) -> [(k, a)] -> Map k a
fromListWithKey c [] = empty
fromListWithKey c ((k,v):xs) = insertWithKey c k v (fromListWithKey c xs)

-- | Alias for backward compatibility.
fromAscList :: Eq k => [(k, a)] -> Map k a 
fromAscList = fromList

-- | Alias for backward compatibility.
fromAscListWith :: Eq k => (a -> a -> a) -> [(k, a)] -> Map k a 
fromAscListWith = fromListWith

-- | Alias for backward compatibility.
fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k, a)] -> Map k a 
fromAscListWithKey = fromListWithKey 

-- | Alias for backward compatibility.
fromDistinctAscList :: [(k, a)] -> Map k a 
fromDistinctAscList = fromList

-- | Alias for backward compatibility.
fromDescList :: Eq k => [(k, a)] -> Map k a
fromDescList = fromList

-- | Alias for backward compatibility.
fromDescListWith :: Eq k => (a -> a -> a) -> [(k, a)] -> Map k a 
fromDescListWith = fromListWith

-- | Alias for backward compatibility.
fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k, a)] -> Map k a 
fromDescListWithKey = fromListWithKey

-- | Alias for backward compatibility.
fromDistinctDescList :: [(k, a)] -> Map k a 
fromDistinctDescList = fromList

-- Operators


-- | /O(n)/.  Lookup the value at a key in the map. If the map is not defined on the given value returns `Nothing`, otherwise returns `Just` the image.
--
-- This function is like `lookup` in Data.Map for function (beware: the order of the argument are reversed).
(|?|) :: (Eq k) => Map k v -> k -> Maybe v
(|?|) (Map al) key = (Set.setToMaybe).(Set.catMaybes) $ (\(k,v) -> if k == key then Just v else Nothing) <$> al

-- | /O(n)/. Unsafe version of `(|?|)`.
--
-- This function is like `(!)` in Data.Map, it is renamed to avoid name collisions.
(|!|) :: (Eq k) => Map k v -> k -> v
(|!|) f key
    | Foldable.null safeResult = error "WeakMap.|!|: element not in the map"
    | otherwise = result
    where
        safeResult = f |?| key
        Just result = safeResult
        
-- | O(n). Find the value at a key. Calls `error` when the element can not be found.
-- 
-- Alias of `(|!|)` for backward compatibility purposes. 
(!) :: (Eq k) => Map k a -> k -> a
(!) = (|!|)

-- | Alias for backward compatibility.
(!?) :: Eq k => Map k a -> k -> Maybe a
(!?) = (|?|)


-- | See `difference`.
(\\) :: (Eq k) => Map k a -> Map k b -> Map k a
(\\) = difference

-- | /O(n*m)/. Difference of two maps. Return elements of the first map not existing in the second map.
difference :: (Eq k) => Map k a -> Map k b -> Map k a
difference (Map al) m2 = Map $ Set.filter (\(k,v) -> not $ k `isIn` (keys' m2)) al


-- Query


-- | /O(n^2)/. The number of elements in the map. 
size :: (Eq k) => Map k a -> Int
size m = length $ mapToList m

-- | /O(1)/. Return wether the map is empty.
null :: Map k a -> Bool
null (Map al) = Set.null al

-- | /O(n)/. Is the key a member of the map? See also `notMember`. 
member :: (Eq k) => k -> Map k a -> Bool
member k m = not.(Foldable.null) $ m |?| k

-- | /O(n)/. Negation of `member`. 
notMember :: (Eq k) => k -> Map k a -> Bool
notMember k m = Foldable.null $ m |?| k

-- | /O(n)/. Just like `(|?|)` but the order of argument is reversed. For backward compatibility with Data.Map.
lookup :: (Eq k) => k -> Map k a -> Maybe a
lookup x y = (|?|) y x

-- | /O(n)/. Apply a map to a given value, if the key is in the domain returns the image, otherwise return a default value.
--
-- This function is like `findWithDefault` (the order of the argument are reversed though).
findWithDefault' :: (Eq k) => Map k v -> v -> k -> v
findWithDefault' f d key
    | Foldable.null safeResult = d
    | otherwise = result
    where
        safeResult = f |?| key
        Just result = safeResult
        
-- | /O(n)/. The expression @(findWithDefault def k map)@ returns the value at key k or returns default value def when the key is not in the map. 
findWithDefault :: (Eq k) => a -> k -> Map k a -> a
findWithDefault d key f = findWithDefault' f d key


-- Insertion


-- | /O(1)/. Insert a new key and value in the map. If the key is already present in the map, the associated value is replaced with the supplied value. `insert` is equivalent to @`insertWith` `const`@. 
insert :: k -> a -> Map k a -> Map k a
insert k v (Map al) = Map $ Set.insert (k,v) al

-- | /O(n)/. Insert with a function, combining new value and old value. @insertWith f key value mp@ will insert the pair (key, value) into mp if key does not exist in the function. If the key does exist, the function will insert the pair (key, f new_value old_value). 
insertWith :: (Eq k) => (v -> v -> v) -> k -> v -> Map k v -> Map k v
insertWith comb k v f
    | Foldable.null prev = insert k v f
    | otherwise = insert k (comb v prev_value) f
    where
        prev = f |?| k
        Just prev_value = prev

-- | /O(n)/. Alias for `insertWith` for backward compatibility purposes.
insertWith' :: (Eq k) => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWith' = insertWith

-- | /O(n)/. Insert with a function, combining key, new value and old value. @insertWithKey f key value mp@ will insert the pair (key, value) into mp if key does not exist in the function. If the key does exist, the function will insert the pair (key,f key new_value old_value). Note that the key passed to f is the same key passed to `insertWithKey`. 
insertWithKey :: (Eq k) => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWithKey comb k v f
    | Foldable.null prev = insert k v f
    | otherwise = insert k (comb k v prev_value) f
    where
        prev = f |?| k
        Just prev_value = prev

-- | /O(n)/. Alias for `insertWithKey` for backward compatibility purposes.
insertWithKey' :: (Eq k) => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWithKey' = insertWithKey

-- | /O(n)/.  Combines insert operation with old value retrieval. The expression @(insertLookupWithKey f k x map)@ is a pair where the first element is equal to @(lookup k map)@ and the second element equal to @(insertWithKey f k x map)@. 
insertLookupWithKey :: (Eq k) => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a)
insertLookupWithKey f k x map = (lookup k map, insertWithKey f k x map)

-- | /O(n)/. Alias for `insertLookupWithKey` for backward compatibility purposes.
insertLookupWithKey' :: (Eq k) => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a)
insertLookupWithKey' = insertLookupWithKey

-- | /O(1)/. Insert a new key and value if it is Just in the map. If the key is already present in the map, the associated value is replaced with the supplied value.
insertMaybe :: (Eq k) => k -> Maybe a -> Map k a -> Map k a
insertMaybe _ Nothing m = m
insertMaybe k (Just v) (Map al) = Map $ Set.insert (k,v) al


-- Delete/update


-- | /O(n)/. Delete a key and its value from the map. When the key is not a member of the map, the original map is returned. 
delete :: (Eq k) => k -> Map k a -> Map k a
delete key (Map al) = Map $ Set.filter (\(k,v) -> key /= k) al

-- | /O(n)/. Update a value at a specific key with the result of the provided function. When the key is not a member of the map, the original map is returned. 
adjust :: (Eq k) => (a -> a) -> k -> Map k a -> Map k a
adjust func key (Map al) = Map $ (\(k,v) -> if key == k then (k, func v) else (k,v)) <$> al

-- | /O(n)/. Adjust a value at a specific key. When the key is not a member of the map, the original map is returned. 
adjustWithKey :: (Eq k) => (k -> a -> a) -> k -> Map k a -> Map k a
adjustWithKey func key (Map al) = Map $ (\(k,v) -> if key == k then (k, func k v) else (k,v)) <$> al

-- | /O(n)/. The expression (`alter` f k map) alters the value x at k, or absence thereof. alter can be used to insert, delete, or update a value in a `Map`. In short : `lookup` k (`alter` f k m) = f (`lookup` k m). 
alter :: (Eq k) => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
alter func key f
    | Foldable.null lookupKey = insert key unpackedImageNothing f
    | Foldable.null result = delete key f
    | otherwise = insert key unpackedResult f
    where
        lookupKey = f |?| key
        result = func lookupKey
        Just unpackedResult = result
        Just unpackedImageNothing = func Nothing
      
-- | /O(n)/. The expression (`alterF` f k map) alters the value x at k, or absence thereof. `alterF` can be used to inspect, insert, delete, or update a value in a Map.      
alterF :: (Functor f, Eq k) => (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
alterF func key f
    | Foldable.null lookupKey = add <$> imageNothing
    | otherwise = treat <$> result
    where
        lookupKey = f |?| key
        result = func lookupKey
        treat x = if Foldable.null x then delete key f else insert key (fromJust x) f
        imageNothing = func Nothing
        add x = if Foldable.null x then f else insert key (fromJust x) f


-- | /O(n)/. The expression @(update f k map)@ updates the value x at k (if it is in the map). If (f x) is Nothing, the element is deleted. If it is (Just y), the key k is bound to the new value y. 
update :: (Eq k) => (a -> Maybe a) -> k -> Map k a -> Map k a
update f = updateWithKey (\_ x -> f x)

-- | /O(n)/. The expression @(updateWithKey f k map)@ updates the value x at k (if it is in the map). If (f k x) is Nothing, the element is deleted. If it is (Just y), the key k is bound to the new value y. 
updateWithKey :: (Eq k) => (k -> a -> Maybe a) -> k -> Map k a -> Map k a
updateWithKey f key m
    | Foldable.null look = m
    | Foldable.null res = delete key m
    | otherwise = insert key r m
    where
        look = m |?| key
        Just v = look
        res = f key v
        Just r = res

-- | /O(n)/. Lookup and update. See also `updateWithKey`. The function returns changed value, if it is updated. Returns the original key value if the map entry is deleted. 
updateLookupWithKey :: (Eq k) => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a, Map k a)
updateLookupWithKey f key m
    | Foldable.null look = (Nothing, m)
    | Foldable.null res = (Just v, delete key m)
    | otherwise = (Just v, insert key r m)
    where
        look = m |?| key
        Just v = look
        res = f key v
        Just r = res


-- Combine


-- | /O(n)/. The expression @(`union` t1 t2)@ takes the left-biased union of t1 and t2. It prefers t1 when duplicate keys are encountered. 
union :: (Eq k) => Map k a -> Map k a -> Map k a
union (Map al1) (Map al2) = Map $ al1 ||| al2

-- | /O(n)/. Union with a combining function.
unionWith :: (Eq k) => (a -> a -> a) -> Map k a -> Map k a -> Map k a
unionWith f m1 m2 = unionWithKey (\_ x y -> f x y) m1 m2

-- | /O(n)/.
-- Union with a combining function. 
--
-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value
-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
unionWithKey :: (Eq k) => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a
unionWithKey f (Map al) m2 = new_m1 `union` m2
    where
        new_m1 = Map $ (\(k,v) -> if member k m2 then (k,f k v (m2 |!| k)) else (k,v)) <$> al

-- | The union of a list of maps: @(unions == foldl union empty)@. 
unions :: (Eq k) => [Map k a] -> Map k a
unions = L.foldl union empty

-- | The union of a list of maps, with a combining operation: @(unionsWith f == foldl (unionWith f) empty)@. 
unionsWith :: (Eq k) => (a -> a -> a) -> [Map k a] -> Map k a
unionsWith f = L.foldl (unionWith f) empty


-- Difference


-- | /O(n*m)/. Difference with a combining function. When two equal keys are encountered, the combining function is applied to the values of these keys. If it returns Nothing, the element is discarded (proper set difference). If it returns (Just y), the element is updated with a new value y
differenceWith :: (Eq k) => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
differenceWith f m1 m2 = differenceWithKey (\_ x y -> f x y) m1 m2


-- | /O(n*m)/. Difference with a combining function. When two equal keys are encountered, the combining function is applied to the key and both values. If it returns Nothing, the element is discarded (proper set difference). If it returns (Just y), the element is updated with a new value y. 
differenceWithKey :: (Eq k) => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
differenceWithKey f (Map al) m2 = Map $ Set.mapMaybe (\(k,v) -> if member k m2 then (sequence (k, f k v (m2 |!| k))) else Just (k, v)) al


-- Intersection


-- | /O(n*m)/. Intersection of two maps. Return data in the first map for the keys existing in both maps. 
intersection :: (Eq) k => Map k a -> Map k b -> Map k a
intersection = intersectionWith const      
    
-- | /O(n*m)/. Intersection with a combining function. 
intersectionWith :: (Eq k) => (a -> b -> c) -> Map k a -> Map k b -> Map k c
intersectionWith f m1 m2 = intersectionWithKey (\_ x y -> f x y) m1 m2

-- | /O(n*m)/. Intersection with a combining function.
intersectionWithKey :: (Eq k) => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c
intersectionWithKey f (Map al) m2 = Map $ Set.mapMaybe (\(k,v) -> if member k m2 then Just (k, f k v (m2 |!| k)) else Nothing) al

-- | Check whether the key sets of two maps are disjoint.
disjoint :: Eq k => Map k a -> Map k b -> Bool
disjoint m1 m2 = null $ intersection m1 m2

-- | Relate the keys of one map to the values of the other, by using the values of the former as keys for lookups in the latter.
compose :: Eq b => Map b c -> Map a b -> Map a c 
compose = (|.|)

-- | Compose two functions. If the two functions are not composable, strips the functions until they can compose.
(|.|) :: (Eq b) => Map b c -> Map a b -> Map a c
(|.|) f2 (Map al) = Map $ Set.mapMaybe (\(k,v) -> sequence (k,(f2 |?| v))) al

-- Traversal

-- | /O(n)/. Map a function over all values in the map. 
map :: (a -> b) -> Map k a -> Map k b
map = fmap

-- | /O(n)/. Map a function over all values in the map. 
mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
mapWithKey f (Map al) = Map $ (\(k,a) -> (k, f k a)) <$> al

-- | /O(n)/. The function `mapAccum` threads an accumulating argument through the map. 
mapAccum :: (Eq k) => (a -> b -> (a, c)) -> a -> Map k b -> (a, Map k c)
mapAccum f d m = mapAccumWithKey (\a k b -> f a b) d m

-- | /O(n^2)/. The function `mapAccumWithKey` threads an accumulating argument through the map. 
mapAccumWithKey :: (Eq k) => (a -> k -> b -> (a, c)) -> a -> Map k b -> (a, Map k c)
mapAccumWithKey f d m = Map <$> set <$> helper d al []
    where 
        al = mapToList m
        helper d [] l = (d,l)
        helper d ((k,v):xs) l = helper new xs ((k,val):l) 
            where
                (new,val) = f d k v


-- | /O(n)/. Alias of `mapAccumWithKey` for backward compatibility purposes. We don't implement it because order of pairs should not matter.
mapAccumRWithKey :: (Eq k) => (a -> k -> b -> (a, c)) -> a -> Map k b -> (a, Map k c)
mapAccumRWithKey = mapAccumWithKey


-- | /O(n)/. @mapKeys f s@ is the map obtained by applying f to each key of s. 
mapKeys :: (k1 -> k2) -> Map k1 a -> Map k2 a
mapKeys f (Map al) = Map $ (\(k,v) -> (f k,v)) <$> al

-- | /O(n^2)/. @mapKeysWith c f s@ is the map obtained by applying f to each key of s. 
--
-- The size of the result may be smaller if f maps two or more distinct keys to the same new key. In this case the associated values will be combined using c. 
mapKeysWith :: (Eq k1, Eq k2) => (a -> a -> a) -> (k1 -> k2) -> Map k1 a -> Map k2 a
mapKeysWith f g m = helper al empty
    where
        al = mapToList m
        helper [] r = r
        helper ((k,v):xs) r
            | g k `member` r = helper xs (insert (g k) (f v (r |!| (g k))) r)
            | otherwise = helper xs (insert (g k) v r)

-- | Alias of `mapKeys` defined for backward compatibility.
mapKeysMonotonic ::  (k1 -> k2) -> Map k1 a -> Map k2 a
mapKeysMonotonic = mapKeys

-- | Eq typeclass must be added.
--
-- It behaves much like a regular traverse except that the traversing function also has access to the key associated with a value and the values are forced before they are installed in the result map.
traverseWithKey :: (Applicative t, Eq k, Eq a) => (k -> a -> t b) -> Map k a -> t (Map k b) 
traverseWithKey f m = foldrWithKey (\k v ys -> liftA3 insert (pure k) (f k v) ys) (pure mempty) (fromList.mapToList $ m)

-- | Eq typeclass must be added. Traverse keys/values and collect the Just results.
traverseMaybeWithKey :: (Applicative f, Eq k, Eq a) => (k -> a -> f (Maybe b)) -> Map k a -> f (Map k b) 
traverseMaybeWithKey f m = foldrWithKey (\k v ys -> liftA3 insertMaybe (pure k) (f k v) ys) (pure mempty) (fromList.mapToList $ m)

-- Folds

-- | /O(n^2)/. Fold the values in the map using the given right-associative binary operator.
--
-- Note that an Eq constraint must be added.
foldr :: (Eq k) => (a -> b -> b) -> b -> Map k a -> b 
foldr f d m = L.foldr (\(k,v) a -> f v a) d (mapToList m)

-- | /O(n^2)/. Fold the values in the map using the given left-associative binary operator.
--
-- Note that an Eq constraint must be added.
foldl :: (Eq k) => (a -> b -> a) -> a -> Map k b -> a 
foldl f d m = L.foldl (\a (k,v) -> f a v) d (mapToList m)

-- | `foldrWithKey` from the left.
foldlWithKey :: (Eq k, Eq a) => (b -> k -> a -> b) -> b -> Map k a -> b
foldlWithKey f d (Map al) = Set.foldl (\b (k,v) -> f b k v) d al

-- | Fold with key.
foldrWithKey :: (Eq a, Eq k) =>(k -> a -> b -> b) -> b -> Map k a -> b
foldrWithKey f d (Map al) = Set.foldr (\(k,v) -> f k v) d al

-- |  Fold the keys and values in the map using the given monoid.
foldMapWithKey :: (Eq a, Eq k, Monoid m) => (k -> a -> m) -> Map k a -> m
foldMapWithKey f (Map al) = Set.foldr (\(k,v) m -> (f k v) <> m) mempty al


-- Strict folds



-- | Strict foldr.
foldr' :: (Eq k) => (a -> b -> b) -> b -> Map k a -> b
foldr' f d m = L.foldr (\(k,v) a -> f v a) d (mapToList m)

-- | Strict foldl.
foldl' :: (Eq k) => (b -> a -> b) -> b -> Map k a -> b
foldl' f d m = L.foldl' (\a (k,v) -> f a v) d (mapToList m)

-- | Strict `foldrWithKey`.
foldlWithKey' :: (b -> k -> a -> b) -> b -> Map k a -> b
foldlWithKey' f d (Map al) = Set.foldl' (\b (k,v) -> f b k v) d al

-- | Strict `foldrWithKey`.
foldrWithKey' :: (k -> a -> b -> b) -> b -> Map k a -> b
foldrWithKey' f d (Map al) = Set.foldr' (\(k,v) -> f k v) d al



-- | A universal combining function.
mergeWithKey :: Eq k => (k -> a -> b -> Maybe c) -> (Map k a -> Map k c) -> (Map k b -> Map k c) -> Map k a -> Map k b -> Map k c 
mergeWithKey c f g f1 f2 = Map $ Set.mapMaybe helper ((keys' f1) ||| (keys' f2))
    where
        newF1 = f f1
        newF2 = g f2
        helper k
            | k `member` f1 && k `member` f2 = sequence (k, c k (f1 |!| k) (f2 |!| k))
            | k `member` f1 = sequence (k, newF1 |?| k)
            | k `member` f2 = sequence (k, newF2 |?| k)


-- Conversion


-- | /O(n^2)/. Transform a function back into its underlying association list.
mapToList :: (Eq k) => Map k v -> AssociationList k v
mapToList (Map al) = nubSetBy (\x y -> (fst x) == (fst y)) al

-- | /O(n^2)/. Transform a function back into its underlying set of pairs.
mapToSet :: (Eq k) => Map k v -> Set (k,v)
mapToSet m = set $ mapToList m


-- | /O(n^2)/. Return all values of the map. Beware that an Eq typeclass must be added.
elems :: (Eq k) => Map k a -> [a]
elems m = snd <$> mapToList m

-- | /O(n^2)/. Same as `elems` but returns a `Set`. Beware that an Eq typeclass must be added.
elems' :: (Eq k) => Map k a -> Set a
elems' = set.elems

-- | /O(n^2)/. Alias of `elems'`.
values :: (Eq k) => Map k a -> Set a
values = elems'

-- | /O(n^2)/. Alias of `elems'`.
image :: (Eq k) => Map k a -> Set a
image = elems'

-- | /O(n^2)/. Return the keys of a map. Beware that an Eq typeclass must be added.
keys :: (Eq k) => Map k v -> [k]
keys m = fst <$> mapToList m

-- | /O(n)/. Same as `keys` but returns a `Set`. No Eq typeclass required.
keys' :: Map k v -> Set k
keys' (Map al) = fst <$> al

-- | /O(n^2)/. Alias of `keys'`.
domain :: Map k a -> Set k
domain = keys'

-- | Alias of `mapToList` for backward compatibility. Beware that an Eq typeclass must be added.
assocs :: (Eq k) => Map k a -> [(k, a)]
assocs = mapToList

-- | Alias of `keys'` for backward compatibility.
keysSet :: Map k a -> Set k 
keysSet = keys'

-- | /O(n^2)/. Return the set of values of a map.
elemsSet :: (Eq k) => Map k a -> Set a
elemsSet = values

-- | /O(n^2)/. Alias of `mapToList` for backward compatibility. Beware that an Eq typeclass must be added.
toList :: (Eq k) => Map k a -> [(k, a)]
toList = mapToList

-- | Alias of `toList` for backward compatibility.
toAscList :: (Eq k) => Map k a -> [(k, a)]
toAscList = toList

-- | Alias of `toList` for backward compatibility.
toDescList :: (Eq k) => Map k a -> [(k, a)]
toDescList = toList



-- Filter


-- | /O(n)/. Filter all values that satisfy the predicate.
filter :: (a -> Bool) -> Map k a -> Map k a
filter p (Map al) = Map $ Set.filter (\(k,v) -> p v) al
   
-- | /O(n)/. Filter all keys/values that satisfy the predicate.
filterWithKey :: (k -> a -> Bool) -> Map k a -> Map k a 
filterWithKey p (Map al) = Map $ Set.filter (uncurry p) al

-- | /O(n*m)/. Restrict a `Map` to only those keys found in a `Set`.
restrictKeys :: Eq k => Map k a -> Set k -> Map k a
restrictKeys m s = filterWithKey (\k v -> k `isIn` s) m

-- | /O(n*m)/. Remove all keys in a `Set` from a `Map`.
withoutKeys :: Eq k => Map k a -> Set k -> Map k a 
withoutKeys m s = filterWithKey (\k v -> not $ k `isIn` s) m

-- | /O(n)/. Partition the map according to a predicate. The first map contains all elements that satisfy the predicate, the second all elements that fail the predicate.
partition :: (a -> Bool) -> Map k a -> (Map k a, Map k a)
partition p m = (filter p m, filter (not.p) m)

-- | /O(n)/. Partition the map according to a predicate. The first map contains all elements that satisfy the predicate, the second all elements that fail the predicate. 
partitionWithKey :: (k -> a -> Bool) -> Map k a -> (Map k a, Map k a)
partitionWithKey p m = (filterWithKey p m, filterWithKey ((fmap not).p) m)

-- | /O(n)/. Map values and collect the Just results.
mapMaybe :: (a -> Maybe b) -> Map k a -> Map k b
mapMaybe f m = mapMaybeWithKey (const f) m

-- | /O(n)/. Map keys/values and collect the Just results.
mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b 
mapMaybeWithKey f (Map al) = Map $ Set.mapMaybe customF al
    where
        customF (k,v) = sequence (k, f k v)

-- | /O(n)/. Map values and separate the Left and Right results.
mapEither :: (a -> Either b c) -> Map k a -> (Map k b, Map k c) 
mapEither f m = mapEitherWithKey (const f) m

-- | /O(n)/. Map keys/values and separate the Left and Right results.
mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) 
mapEitherWithKey f (Map al) = (Map ls, Map rs)
    where
        (ls,rs) = Set.mapEither customF al
        customF (k,v)
            | Foldable.null result = Left (k,l)
            | otherwise = Right (k,r)
            where
                result = f k v
                Left l = result
                Right r = result



-- Submap


-- | /O(max(m^2,n^2))/. This function is defined as @(isSubmapOf = isSubmapOfBy (==))@.
isSubmapOf :: (Eq k, Eq a) => Map k a -> Map k a -> Bool
isSubmapOf = isSubmapOfBy (==)

-- | /O(max(m^2,n^2))/. Returns True if the keys of the first map is included in the keys of the second and the predicate evaluation at their value is True.
isSubmapOfBy :: Eq k => (a -> b -> Bool) -> Map k a -> Map k b -> Bool
isSubmapOfBy p m1 m2 = (keys' m1) `isIncludedIn` (keys' m2) && (Set.and $ test <$> keys' m1)
    where
        test k = p (m1 |!| k) (m2 |!| k)

-- | /O(max(m^2,n^2))/. This function is defined as @(isProperSubmapOf = isProperSubmapOfBy (==))@.
isProperSubmapOf :: (Eq k, Eq a) => Map k a -> Map k a -> Bool
isProperSubmapOf = isProperSubmapOfBy (==)

-- | /O(max(m^2,n^2))/. Returns True if the keys of the first map is strictly included in the keys of the second and the predicate evaluation at their value is True.
isProperSubmapOfBy :: Eq k => (a -> b -> Bool) -> Map k a -> Map k b -> Bool
isProperSubmapOfBy p m1 m2 = (keys' m1) `Set.isProperSubsetOf` (keys' m2) && (Set.and $ test <$> keys' m1)
    where
        test k = p (m1 |!| k) (m2 |!| k)



-- Indexed


-- | /O(n^2)/. Lookup the index of a key, which is its zero-based index in the sequence. The index is a number from 0 up to, but not including, the size of the map.
lookupIndex :: Eq k => k -> Map k a -> Maybe Int
lookupIndex k m = L.elemIndex k (keys m)

-- | /O(n^2)/. Return the index of a key, which is its zero-based index in the sequence. The index is a number from 0 up to, but not including, the size of the map. Calls error when the key is not a member of the map.
findIndex :: Eq k => k -> Map k a -> Int
findIndex k m
    | Foldable.null index = error "WeakSet.findIndex: element is not in the set"
    | otherwise = i
    where
        index = lookupIndex k m
        Just i = index
        
-- | /O(n^2)/. Retrieve an element by its index, i.e. by its zero-based index in the sequence. If the index is out of range (less than zero, greater or equal to size of the map), error is called.
elemAt :: Eq k => Int -> Map k a -> (k, a)
elemAt i m
    | i < 0 || i >= (length xs) = error "WeakSet.elemAt: index out of range"
    | otherwise = (L.!!) xs i
    where
        xs = mapToList m

-- | Helper function for updateAt.
updateListAt :: Int -> (a -> Maybe a) -> [a] -> [a]
updateListAt _ _ [] = []
updateListAt 0 f (x : xs)
    | Foldable.null result = xs
    | otherwise = r : xs
    where
        result = f x
        Just r = result
updateListAt n f (x : xs) = x : updateListAt (n-1) f xs

-- | /O(n^2)/. Update the element at index. Calls error when an invalid index is used.
updateAt :: Eq k => (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
updateAt f i m = Map $ set $ updateListAt i (\(k,v) -> sequence (k, f k v)) (mapToList m)

-- | /O(n^2)/. Delete the element at index, i.e. by its zero-based index in the sequence. If the index is out of range (less than zero, greater or equal to size of the map), error is called.
deleteAt :: (Eq k) => Int -> Map k a -> Map k a
deleteAt i m = Map $ set $ helper i al
    where
        al = mapToList m
        helper _ [] = []
        helper 0 (x:xs) = xs
        helper n (x:xs) = helper (n-1) xs

-- | /O(n^2)/. Take a given number of pairs to create a new map.
take :: (Eq k) => Int -> Map k a -> Map k a
take n m = fst $ splitAt n m

-- | /O(n^2)/. Drop a given number of pairs to create a new map.
drop :: (Eq k) => Int -> Map k a -> Map k a
drop n m = snd $ splitAt n m

-- | /O(n^2)/. Split a map at a particular index.
splitAt :: (Eq k) => Int -> Map k a -> (Map k a, Map k a)
splitAt n m = (Map $ set l, Map $ set r)
    where
        al = mapToList m
        (l,r) = L.splitAt n al


-- Others



-- | /O(n)/. Return the identity function associated to a `Set`.
idFromSet :: Set a -> Map a a
idFromSet set = Map $ (\x -> (x,x)) <$> set


-- | /O(n)/. Memorize a Haskell function on a given finite domain. Alias of `fromSet`.
memorizeFunction :: (k -> v) -> Set k -> Map k v
memorizeFunction f xs = Map $ (\k -> (k, f k)) <$> xs

-- | /O(n^2)/. Try to construct an inverse map.
inverse :: (Eq k, Eq v) => Map k v -> Maybe (Map v k)
inverse m@(Map al)
    | (cardinal.image $ m) == (cardinal.domain $ m) = Just $ Map $ (\(k,v) -> (v,k)) <$> al
    | otherwise = Nothing
    
-- | /O(n)/. Return a pseudo inverse /g/ of a 'Map' /f/ such that @f |.| g |.| f == f@.
pseudoInverse :: Map k v -> Map v k
pseudoInverse (Map al) = Map $ (\(k,v) -> (v,k)) <$> al

-- | Return all Functions from a domain to a codomain.
enumerateMaps :: (Eq a, Eq b) =>
                    Set a -- ^ Domain.
                 -> Set b -- ^ Codomain.
                 -> Set (Map a b) -- ^ All maps from domain to codomain.
enumerateMaps dom codom = set $ weakMap <$> zip (setToList dom) <$> setToList (codom |^| (cardinal dom))