{-| Module  : WeakSets
Description : Homogeneous functions are functions between `HomogeneousSet`s. They are more flexible than Data.Map because they do not require the keys to be orderable.
Copyright   : Guillaume Sabbagh 2022
License     : LGPL-3.0-or-later
Maintainer  : guillaumesabbagh@protonmail.com
Stability   : experimental
Portability : portable

Homogeneous functions are functions between `HomogeneousSet`s.

They are more flexible than Data.Map because they do not require the keys to be orderable.

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

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

Inline functions related to homogeneous functions are written between pipes @|@.

Function names should not collide with Prelude but may collide with Data.Map.
-}

module Data.WeakSets.HomogeneousFunction
(
    -- * Function datatype and smart constructor
    AssociationList(..),
    Function, -- abstract type, the smart constructor is `function`
    function, -- the smart constructor for `Function`
    -- * Function related functions
    domain,
    image,
    idFromSet,
    (|.|),
    -- * Query
    size,
    member,
    notMember,
    (|?|),
    (|!|),
    findWithDefault,
    -- * Construction
    -- ** Insertion
    insert,
    insertWith,
    insertWithKey,
    -- ** Delete/Update 
    delete,
    adjust,
    adjustWithKey,
    alter,
    -- * Combine
    -- ** Union
    union,
    -- * Traversal
    mapKeys,
    -- * Conversion 
    keys,
    elems,
    functionToSet,
    memorizeFunction,
)
where
    import              Data.WeakSets.HomogeneousSet
    
    -- | A function of homogeneous sets. It is a 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 `function`.
    data Function k v = Function (Set (k,v)) deriving (Eq)
    
    instance (Show k, Show v) => Show (Function k v) where
        show (Function al) = "(function "++show al++")"
    
    instance Semigroup (Function k v) where
        (Function al1) <> (Function al2) = Function $ al1 <> al2
    
    instance Monoid (Function k v) where
        mempty = Function (set [])
        
    instance Foldable (Function k) where
        foldr f d (Function al) = foldr (\(k,v) -> f v) d al

    instance Functor (Function k) where
        fmap f (Function al) = Function $ (\(k,v) -> (k,f v)) <$> al
    
    -- | An association list is a list of pairs (key,value).
    type AssociationList k v = [(k,v)]
    
    -- | /O(1)/. The smart constructor of functions. This is the only way of instantiating a `Function`.
    --
    -- 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.
    function :: AssociationList k v -> Function k v
    function al = Function $ set $ al
    
    -- | /O(n)/. Return the domain of a function.
    domain :: Function k v -> Set k
    domain (Function al) = fst <$> al
    
    -- | /O(n)/. Return the image of a function. The image of a function is the set of values which are reachable by applying the function.
    image :: Function k v -> Set v
    image (Function al) = snd <$> al
    
    -- | /O(n)/. Return the identity function associated to a `Set`.
    idFromSet :: Set a -> Function a a
    idFromSet set = Function $ (\x -> (x,x)) <$> set
    
    -- | Compose two functions. If the two functions are not composable, strips the functions until they can compose.
    (|.|) :: (Eq a, Eq b) => Function b c -> Function a b -> Function a c
    (|.|) f2 f1 = Function $ set [(k,(f2 |!| (f1 |!| k))) | k <- (setToList.domain $ f1), f1 |!| k `isIn` (domain f2)]
        
{--------------------------------------------------------------------
  Query
--------------------------------------------------------------------}
    
    -- | /O(n)/. The number of entries in the function.
    size :: (Eq k) => Function k v -> Int
    size f = length.setToList.keys $ f
    
    -- | /O(n)/. Return wether a key is in the function domain or not.
    member :: (Eq k) => Function k v -> k -> Bool
    member f k = k `isIn` (domain f)
    
    -- | /O(n)/. Negation of member.
    notMember :: (Eq k) => Function k v -> k -> Bool
    notMember f k = not $ member f k
    
    -- | /O(n)/. Apply a function to a given value. If the function is not defined on the given value returns `Nothing`, otherwise returns `Just` the image.
    --
    -- This function is like `lookup` in Data.Map for function (the order of the argument are reversed though).
    (|?|) :: (Eq k) => Function k v -> k -> Maybe v
    (|?|) (Function al) key = setToMaybe.catMaybesToSet $ (\(k,v) -> if k == key then Just v else Nothing) <$> al
    
    -- | /O(n)/. Unsafe version of `(|?|)`.
    --
    -- This function is like `(!)` in Data.Map for function.
    (|!|) :: (Eq k) => Function k v -> k -> v
    (|!|) f key
        | null safeResult = error "Function applied on a value not in the domain."
        | otherwise = result
        where
            safeResult = f |?| key
            Just result = safeResult
    
    -- | /O(n)/. Apply a function to a given value, if the value is in the domain returns the image, otherwise return a default value.
    --
    -- This function is like `findWithDefault` in Data.Map for function (the order of the argument are reversed though).
    findWithDefault :: (Eq k) => Function k v -> v -> k -> v
    findWithDefault f d key
        | null safeResult = d
        | otherwise = result
        where
            safeResult = f |?| key
            Just result = safeResult
            
{--------------------------------------------------------------------
  Insertion
--------------------------------------------------------------------}

    -- | O(1). Insert a new key and value in the function. If the key is already present in the function, the associated value is replaced with the supplied value. insert is equivalent to insertWith const. 
    insert :: k -> v -> Function k v -> Function k v
    insert k v (Function al) = Function $ (set [(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 -> Function k v -> Function k v
    insertWith comb k v f
        | null prev = insert k v f
        | otherwise = insert k (comb v prev_value) f
        where
            prev = f |?| k
            Just prev_value = prev

    -- | 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 -> Function k a -> Function k a
    insertWithKey comb k v f
        | null prev = insert k v f
        | otherwise = insert k (comb k v prev_value) f
        where
            prev = f |?| k
            Just prev_value = prev

{--------------------------------------------------------------------
  Conversion
--------------------------------------------------------------------}

    -- | /O(n)/. Transform a function back into its underlying association list.
    functionToSet :: (Eq k) => Function k v -> Set (k,v)
    functionToSet (Function al) = nubSetBy (\x y -> (fst x) == (fst y)) al

    -- | /O(n)/. Alias of domain.
    keys :: Function k v -> Set k
    keys = domain
    
    -- | /O(n)/. Alias of image.
    elems :: Function k v -> Set v
    elems = image

    -- | /O(n)/. Memorize a Haskell function on a given finite domain.
    memorizeFunction :: (k -> v) -> Set k -> Function k v
    memorizeFunction f xs = Function $ (\k -> (k, f k)) <$> xs
    
{--------------------------------------------------------------------
  Delete/Update 
--------------------------------------------------------------------}
    -- | O(n). Delete a key and its value from the function. When the key is not a member of the function, the original function is returned. 
    delete :: Eq k => k -> Function k a -> Function k a
    delete key (Function al) = Function $ filterSet (\(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 function, the original function is returned. 
    adjust :: Eq k => (a -> a) -> k -> Function k a -> Function k a
    adjust func key (Function al) = Function $ (\(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 function, the original function is returned. 
    adjustWithKey :: Eq k => (k -> a -> a) -> k -> Function k a -> Function k a
    adjustWithKey func key (Function al) = Function $ (\(k,v) -> if key == k then (k, func k v) else (k,v)) <$> al
    
    -- | O(n). The expression (`alter` f k function) alters the value x at k, or absence thereof. alter can be used to insert, delete, or update a value in a `Function`. In short : `lookup` k (`alter` f k m) = f (`lookup` k m). 
    alter :: Eq k => (Maybe a -> Maybe a) -> k -> Function k a -> Function k a
    alter func key f
        | null lookupKey = insert key unpackedImageNothing f
        | null result = delete key f
        | otherwise = adjust (const unpackedResult) key f
        where
            lookupKey = f |?| key
            result = func lookupKey
            Just unpackedResult = result
            Just unpackedImageNothing = func Nothing
            
    -- | /O(n)/. Map a function over the keys of a function.
    mapKeys :: (k1 -> k2) -> Function k1 v -> Function k2 v
    mapKeys f (Function al) = Function $ (\(k,v) -> (f k,v)) <$> al
    
    
    -- | /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 => Function k a -> Function k a -> Function k a
    union (Function al1) (Function al2) = Function $ al1 ||| al2