-----------------------------------------------------------------------------
--- A finite map is an efficient purely functional data structure
--- to store a mapping from keys to values.
---
--- This version was ported from a corresponding Haskell library.
---
--- @author Frank Huch, Bernd Brassel
--- @version July 2024
-----------------------------------------------------------------------------

module Data.Map (
        -- Although this type is abstract, we export the constructors
        -- to enable the definition of specific `Map` instances,
        -- like a `RW.Base.ReadWrite` instance.
        Map(..),

        empty,
        singleton,
        fromList,

        insert,
        insertWith,
        insertList,
        insertListWith,
        delete,
        deleteAll,
        adjust,
        splitLookup,

        union,
        unionWith,
        difference,
        intersection,
        intersectionWith,

        foldrWithKey,
        mapWithKey,
        filterWithKey,

        size,
        null,
        member,
        lookup,
        findWithDefault,

        toList,
        keys,
        elems,
        sortWithMap,

        lookupMin,
        lookupMax,
        toPreOrderList

    ) where

import Data.Maybe
import Prelude hiding (empty)

-----------------------------------------------
--        BUILDING finite maps
-----------------------------------------------

--- The empty map.
--- @result an empty map
empty :: Map _ _
empty = Tip

--- Construct a map with only a single element.
--- @param k key of
--- @param a the single element to form
--- @result a finite map with only a single element
singleton :: k -> a -> Map k a
singleton k a = Bin k a 1 Tip Tip


--- Builts a map from given list of tuples (key,element).
--- For multiple occurences of key, the last corresponding
--- element of the list is taken.
fromList :: Ord k => [(k,a)] -> Map k a
fromList xs = insertList xs empty

-----------------------------------------------
--        ADDING AND DELETING
-----------------------------------------------

--- Throws away any previous binding and stores the new one given.
insert :: Ord k => k -> a -> Map k a -> Map k a
insert k a m = insertWith (\ _ new -> new) k a m

--- Instead of throwing away the old binding,
--- insertWith combines the new element with the old one.
--- @param combiner a function combining to elements
--- @param k the key of the elements to be combined
--- @param a the new element
--- @param m a map
--- @result a modified map
insertWith :: Ord k => (a -> a -> a)
           -> k -> a -> Map k a -> Map k a
insertWith _ k a Tip = singleton k a
insertWith combiner new_k new_a (Bin k a sizeM m_l m_r)
  = if new_k < k
    then mkBalBranch k a (insertWith combiner new_k new_a m_l) m_r
    else
      if new_k == k
      then Bin new_k (combiner a new_a) sizeM m_l m_r
      else mkBalBranch k a m_l (insertWith combiner new_k new_a m_r)


--- Throws away any previous bindings and stores the new ones given.
--- The items are added starting with the first one in the list
insertList :: Ord k => [(k, a)] -> Map k a -> Map k a
insertList k_a_pairs m = insertListWith (\ _ new -> new) k_a_pairs m

--- Combine with a list of tuples (key, element), cf. insertWith
insertListWith :: Ord k => (a -> a -> a)
               -> [(k, a)] -> Map k a -> Map k a
insertListWith combiner k_a_pairs m
  = foldl add m k_a_pairs        -- foldl adds from the left
  where
    add m' (k, a) = insertWith combiner k a m'

--- Deletes key from map.
--- Deletion doesn't complain if you try to delete something
--- which isn't there
delete :: Ord k => k -> Map k a -> Map k a
delete _ Tip = Tip
delete del_k (Bin k a _ m_l m_r)
  = if del_k < k
    then mkBalBranch k a (delete del_k m_l) m_r
    else
      if del_k == k
        then glueBal m_l m_r
        else mkBalBranch k a m_l (delete del_k m_r)

--- Deletes a list of keys from map.
--- Deletion doesn't complain if you try to delete something
--- which isn't there
deleteAll :: Ord k => [k] -> Map k a -> Map k a
deleteAll ks m = foldl (flip delete) m ks

--- Applies a function to element bound to given key.
adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a
adjust _ _ Tip          = Tip
adjust f i (Bin k x h l r)
            | i == k    =  Bin k (f x) h l r
            | i <  k    =  Bin k x h (adjust f i l) r
            | otherwise =  Bin k x h l (adjust f i r)

--- Combines delFrom and lookup.
splitLookup :: Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a)
splitLookup _ Tip = (Tip, Nothing, Tip)
splitLookup i (Bin k a _ m_l m_r)
            | i == k    = (m_l, Just a, m_r)
            | i <  k    = let (m_l', v, m_r') = splitLookup i m_l
                          in (m_l', v, glueBal m_r' m_r)
            | otherwise = let (m_l', v, m_r') = splitLookup i m_r
                          in (glueBal m_l m_l', v, m_r')

-------------------------------------------------
-- COMBINING finite maps
-------------------------------------------------


--- Efficiently add key/element mappings of two maps into a single one.
--- CHANGED: Bindings in left argument shadow those in the right
union :: Ord k => Map k a -> Map k a -> Map k a
union  Tip m2 = m2
union  m1@(Bin _ _ _ _ _) Tip = m1
union  (Bin split_key1 a1 _ left1 right1) m2@(Bin _ _ _ _ _)
  = mkVBalBranch split_key1 a1 (union left1 lts) (union right1 gts)
  where
    lts     = splitLT m2 split_key1
    gts     = splitGT m2 split_key1

--- Efficiently combine key/element mappings of two maps into a single one,
--- cf. insertWith
unionWith :: Ord k => (a -> a -> a)
          -> Map k a -> Map k a -> Map k a
unionWith _        Tip m2 = m2
unionWith _        m1@(Bin _ _ _ _ _) Tip = m1
unionWith combiner (Bin split_key1 a1 _ left1 right1) m2@(Bin _ _ _ _ _)
  = mkVBalBranch split_key1 new_a
                 (unionWith combiner left1 lts)
                 (unionWith combiner right1 gts)
  where
    lts   = splitLT m2 split_key1
    gts   = splitGT m2 split_key1
    new_a = case lookup split_key1 m2 of
              Nothing -> a1
              Just a2 -> combiner a2 a1

--- (difference a1 a2) deletes from a1 any bindings which are bound in a2
difference :: Ord k => Map k a -> Map k b -> Map k a
difference Tip _ = Tip
difference m1@(Bin _ _ _ _ _) Tip = m1
difference m1@(Bin _ _ _ _ _) (Bin split_key2 _ _ left2 right2)
  = glueVBal (difference lts left2) (difference gts right2)
       -- The two can be way different, so we need glueVBal
  where
    lts = splitLT m1 split_key2  -- NB gt and lt, so the equal ones
    gts = splitGT m1 split_key2  -- are not in either.

--- Filters only those keys that are bound in both of the given maps.
--- CHANGED: The elements will be taken from the first map.
intersection :: Ord k => Map k a -> Map k a -> Map k a
intersection m1 m2 = intersectionWith (\ left _ -> left) m1 m2

--- Filters only those keys that are bound in both of the given maps
--- and combines the elements as in insertWith.
intersectionWith :: Ord k => (a -> b -> c)
                 -> Map k a -> Map k b -> Map k c
intersectionWith _        _   Tip             = Tip
intersectionWith _        Tip (Bin _ _ _ _ _) = Tip
intersectionWith combiner m1@(Bin _ _ _ _ _) (Bin split_key2 a2 _  left2 right2)

  | isJust maybe_a1   -- split_a *is* in intersection
  = mkVBalBranch split_key2 (combiner a1' a2)
                 (intersectionWith combiner lts left2)
                 (intersectionWith combiner gts right2)

  | otherwise           -- split_a is *not* in intersection
  = glueVBal (intersectionWith combiner lts left2)
             (intersectionWith combiner gts right2)

  where
    lts = splitLT m1 split_key2      -- NB gt and lt, so the equal ones
    gts = splitGT m1 split_key2      -- are not in either.

    maybe_a1 = lookup split_key2 m1
    Just a1'  = maybe_a1

-------------------------------------------------------------
--  MAPPING, FOLDING, FILTERING on maps
-------------------------------------------------------------

--- Folds map by given function.
foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
foldrWithKey _ z Tip = z
foldrWithKey k z (Bin key elt _ fm_l fm_r)
  = foldrWithKey k (k key elt (foldrWithKey k z fm_r)) fm_l

--- Applies a given function on every element in the map.
mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
mapWithKey _ Tip = Tip
mapWithKey f (Bin key a size' m_l m_r)
  = Bin key (f key a) size' (mapWithKey f m_l) (mapWithKey f m_r)

--- Yields a new map with only those key/element pairs matching the
--- given predicate.
filterWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> Map k a
filterWithKey _ Tip = Tip
filterWithKey p (Bin key a _ m_l m_r)
  | p key a          -- Keep the item
  = mkVBalBranch key a (filterWithKey p m_l) (filterWithKey p m_r)

  | otherwise          -- Drop the item
  = glueVBal (filterWithKey p m_l) (filterWithKey p m_r)

-----------------------------------------------------
-- INTERROGATING maps
-----------------------------------------------------

--- How many elements does given map contain?
size :: Map _ _ -> Int
size Tip                 = 0
size (Bin _ _ size' _ _) = size'

--- Is the given finite map empty?
null :: Map _ _ -> Bool
null m = size m == 0

--- Does given map contain given key?
member :: Ord k => k -> Map k _ -> Bool
member k m = isJust (lookup k m)

--- Retrieves element bound to given key
lookup :: Ord k => k -> Map k a -> Maybe a
lookup _ Tip = Nothing
lookup key_to_find (Bin k a _ m_l m_r)
  = if key_to_find < k
    then lookup key_to_find m_l
    else if key_to_find == k
         then Just a
         else lookup key_to_find m_r

--- Retrieves element bound to given key.
--- If the element is not contained in map, return
--- default value.
findWithDefault :: Ord k => a -> k -> Map k a -> a
findWithDefault deflt k m
  = case lookup k m of
      Nothing -> deflt
      Just a  -> a

--- Retrieves the smallest key/element pair in the finite map
--- according to the basic key ordering.
lookupMin :: Map k a -> Maybe (k, a)
lookupMin  Tip                         = Nothing
lookupMin (Bin k x _ l _) | isBranch l = lookupMin l
                          | otherwise  = Just (k, x)

--- Retrieves the greatest key/element pair in the finite map
--- according to the basic key ordering.
lookupMax :: Map k a -> Maybe (k, a)
lookupMax Tip                          = Nothing
lookupMax (Bin k x _ _ r) | isBranch r = lookupMax r
                          | otherwise  = Just (k, x)



----------------------------------------------------
-- LISTIFYING: transform finite maps to lists
----------------------------------------------------

--- Builds a list of key/element pairs. The list is ordered
--- by the "Ord" context on keys.
toList :: Map k a -> [(k, a)]
toList m = foldrWithKey (\ k a rest -> (k, a) : rest) [] m

--- Retrieves a list of keys contained in map.
--- The list is ordered
--- by the "Ord" context on keys.
keys :: Map k _ -> [k]
keys m = foldrWithKey (\ k _ rest -> k : rest) [] m

--- Retrieves a list of elements contained in map.
--- The list is ordered
--- by the "Ord" context on keys.
elems :: Map _ a -> [a]
elems m = foldrWithKey (\ _ a rest -> a : rest) [] m

--- Retrieves list of key/element pairs in preorder of the internal tree.
--- Useful for lists that will be retransformed into a tree or to match
--- any elements regardless of basic order.
toPreOrderList :: Map k a -> [(k, a)]
toPreOrderList m = pre m []
   where
     pre Tip xs = xs
     pre (Bin k x _ l r) xs = (k, x) : pre l (pre r xs)

--- Sorts a given list by inserting and retrieving from map.
--- Duplicates are deleted.
sortWithMap :: Ord k => [k] -> [k]
sortWithMap l = keys (fromList (zip l (repeat ())))

-----------------------------------------------------
-- internal Implementation
-----------------------------------------------------
data Map k a
  = Tip
  | Bin k a -- Key and element stored here
    Int{-STRICT-}    -- Size >= 1
    (Map k a)        -- Children
    (Map k a)
  deriving (Show, Read)

instance (Eq k, Eq a) => Eq (Map k a) where
  m_1 == m_2 =
    (size   m_1 == size   m_2) &&   -- quick test
    (toList m_1 == toList m_2)

isBranch :: Map _ _ -> Bool
isBranch (Bin _ _ _ _ _) = True
isBranch Tip             = False

-------------------------------------------------------------------------
--                                                                      -
--  The implementation of balancing                                     -
--                                                                      -
-------------------------------------------------------------------------
-------------------------------------------------------------------------
--                                                                      -
--  Basic construction of a FiniteMap                                   -
--                                                                      -
-------------------------------------------------------------------------
sIZE_RATIO :: Int
sIZE_RATIO = 5

mkBranch :: Int
         -> key -> elt
         -> Map key elt -> Map key elt
         -> Map key elt

mkBranch _{-which-} key elt fm_l fm_r =
    let result = Bin key elt (unbox (1 + left_size + right_size)) fm_l fm_r
    in
      result
      --    if size result <= 8 then
      --     result
      --    else
      --      pprTrace ("mkBranch:"++(show which)) (ppr result) (
      --      result
      --      )
  where
    {-left_ok  = case fm_l of
                 Tip            -> True
                 Bin _ _ _ _ _  -> cmpWithBiggest_left_key key

    cmpWithBiggest_left_key key' = (fst (findMax fm_l)) < key'

    right_ok = case fm_r of
                 Tip           -> True
                 Bin _ _ _ _ _ -> cmpWithSmallest_right_key key

    cmpWithSmallest_right_key key' = key' < (fst (findMin fm_r))

    balance_ok = True -- sigh-}












    left_size  = size fm_l
    right_size = size fm_r


    unbox :: Int -> Int
    unbox x = x


-------------------------------------------------------------------------
--                                                                        -
-- Balanced construction of a FiniteMap                                 -
--                                                                        -
-------------------------------------------------------------------------
mkBalBranch :: key -> elt
            -> Map key elt -> Map key elt
            -> Map key elt

mkBalBranch key elt fm_L fm_R

  | size_l + size_r < 2
  = mkBranch 1{-which-} key elt fm_L fm_R

  | size_r > sIZE_RATIO * size_l        -- Right tree too big
  = case fm_R of
        Bin _ _ _ fm_rl fm_rr ->
              if size fm_rl < 2 * size fm_rr
                then single_L fm_L fm_R
                else double_L fm_L fm_R
        -- Other case impossible
        Tip -> error "Data.Map.mkBalBranch"

  | size_l > sIZE_RATIO * size_r        -- Left tree too big
  = case fm_L of
        Bin _ _ _ fm_ll fm_lr ->
              if size fm_lr < 2 * size fm_ll
                then single_R fm_L fm_R
                else double_R fm_L fm_R
        -- Other case impossible
        Tip -> error "Data.Map.mkBalBranch"

  | otherwise                                -- No imbalance
  = mkBranch 2{-which-} key elt fm_L fm_R

  where
    size_l   = size fm_L
    size_r   = size fm_R

    single_L fm_l (Bin key_r elt_r _ fm_rl fm_rr)
        = mkBranch 3{-which-} key_r elt_r (mkBranch 4{-which-} key elt fm_l fm_rl) fm_rr
    single_L _ Tip = error "Data.Map.single_L"

    double_L fm_l (Bin key_r elt_r _ (Bin key_rl elt_rl _ fm_rll fm_rlr) fm_rr)
        = mkBranch 5{-which-} key_rl elt_rl (mkBranch 6{-which-} key   elt   fm_l   fm_rll)
                                 (mkBranch 7{-which-} key_r elt_r fm_rlr fm_rr)
    double_L _ Tip = error "Data.Map.double_L"
    double_L _ (Bin _ _ _ Tip _) = error "Data.Map.double_L"

    single_R (Bin key_l elt_l _ fm_ll fm_lr) fm_r
        = mkBranch 8{-which-} key_l elt_l fm_ll (mkBranch 9{-which-} key elt fm_lr fm_r)
    single_R Tip _ = error "Data.Map.single_R"

    double_R (Bin key_l elt_l _ fm_ll (Bin key_lr elt_lr _ fm_lrl fm_lrr)) fm_r
        = mkBranch 10{-which-} key_lr elt_lr (mkBranch 11{-which-} key_l elt_l fm_ll  fm_lrl)
                                 (mkBranch 12{-which-} key   elt   fm_lrr fm_r)
    double_R Tip _ = error "Data.Map.double_R"
    double_R (Bin _ _ _ _ Tip) _ = error "Data.Map.double_R"


mkVBalBranch :: Ord key => key -> elt -> Map key elt
             -> Map key elt -> Map key elt

-- Assert: in any call to (mkVBalBranch_C comb key elt l r),
--           (a) all keys in l are < all keys in r
--           (b) all keys in l are < key
--           (c) all keys in r are > key

mkVBalBranch key elt Tip fm_r = insert key elt fm_r
mkVBalBranch key elt (Bin key_l elt_l s_l fm_ll fm_lr) Tip =
   insert key elt (Bin key_l elt_l s_l fm_ll fm_lr)

mkVBalBranch key elt (Bin key_l elt_l s_l fm_ll fm_lr)
                     (Bin key_r elt_r s_r fm_rl fm_rr)
  | sIZE_RATIO * size_l < size_r
  = mkBalBranch key_r elt_r (mkVBalBranch key elt fm_l fm_rl) fm_rr

  | sIZE_RATIO * size_r < size_l
  = mkBalBranch key_l elt_l fm_ll (mkVBalBranch key elt fm_lr fm_r)

  | otherwise
  = mkBranch 13{-which-} key elt fm_l fm_r

  where
    fm_l = Bin key_l elt_l s_l fm_ll fm_lr
    fm_r = Bin key_r elt_r s_r fm_rl fm_rr
    size_l = size fm_l
    size_r = size fm_r

-------------------------------------------------------------------------
--                                                                        -
-- Gluing two trees together                                            -
--                                                                        -
-------------------------------------------------------------------------
glueBal :: Map key elt -> Map key elt -> Map key elt
glueBal fm1 fm2 =
  if null fm1
    then fm2
    else if null fm2
           then fm1
           else
        -- The case analysis here (absent in Adams' program) is really to deal
        -- with the case where fm2 is a singleton. Then deleting the minimum means
        -- we pass an empty tree to mkBalBranch, which breaks its invariant.
             let (mid_key1, mid_elt1) = findMax fm1
                 (mid_key2, mid_elt2) = findMin fm2
             in
             if size fm2 > size fm1
               then mkBalBranch mid_key2 mid_elt2 fm1 (deleteMin fm2)
               else mkBalBranch mid_key1 mid_elt1 (deleteMax fm1) fm2

glueVBal :: Map key elt -> Map key elt -> Map key elt
glueVBal fm_l fm_r =
  if null fm_l
    then fm_r
    else if null fm_r
           then fm_l
           else
             let Bin key_l elt_l _ fm_ll fm_lr = fm_l
                 Bin key_r elt_r _ fm_rl fm_rr = fm_r
                 --(mid_key_l,mid_elt_l) = findMax fm_l
                 --(mid_key_r,mid_elt_r) = findMin fm_r
                 size_l = size fm_l
                 size_r = size fm_r
             in
               if sIZE_RATIO * size_l < size_r
               then
                 mkBalBranch key_r elt_r (glueVBal fm_l fm_rl) fm_rr
                else if sIZE_RATIO * size_r < size_l
                    then
                      mkBalBranch key_l elt_l fm_ll (glueVBal fm_lr fm_r)

                      -- We now need the same two cases as in glueBal above.
                    else glueBal fm_l fm_r

-------------------------------------------------------------------------
--                                                                        -
-- Local utilities                                                      -
--                                                                        -
-------------------------------------------------------------------------

splitLT, splitGT :: Ord key => Map key elt -> key -> Map key elt

-- splitLT fm split_key  =  fm restricted to keys <  split_key
-- splitGT fm split_key  =  fm restricted to keys >  split_key

splitLT  Tip _ = Tip
splitLT (Bin key elt _ fm_l fm_r) split_key
  = if split_key < key
    then splitLT fm_l split_key
    else if split_key == key
         then fm_l
         else mkVBalBranch key elt fm_l (splitLT fm_r split_key)

splitGT Tip _ = Tip
splitGT (Bin key elt _ fm_l fm_r) split_key
  = if split_key < key
    then mkVBalBranch key elt (splitGT fm_l split_key) fm_r
    else if split_key == key
         then fm_r
         else splitGT fm_r split_key

findMin :: Map key elt -> (key, elt)
findMin Tip = error "Data.Map.findMin: empty map"
findMin (Bin key elt _ Tip _) = (key, elt)
findMin (Bin _   _   _ (Bin key_l elt_l s_l fm_ll fm_lr)_) =
      findMin (Bin key_l elt_l s_l fm_ll fm_lr)

deleteMin :: Map key elt -> Map key elt
deleteMin Tip                      = error "Data.Map.deleteMin: empty map"
deleteMin (Bin _   _   _ Tip fm_r) = fm_r
deleteMin (Bin key elt _ (Bin key_l elt_l s_l fm_ll fm_lr) fm_r) =
  mkBalBranch key elt (deleteMin (Bin key_l elt_l s_l fm_ll fm_lr))
                         fm_r

findMax :: Map key elt -> (key,elt)
findMax Tip = error "Data.Map.findMax: empty map"
findMax (Bin key elt _ _ Tip) = (key, elt)
findMax (Bin _   _   _ _  (Bin key_r elt_r s_r fm_rl fm_rr)) =
  findMax (Bin key_r elt_r s_r fm_rl fm_rr)

deleteMax :: Map key elt -> Map key elt
deleteMax Tip                      = error "FiniteMap.deleteMax: empty map"
deleteMax (Bin _   _   _ fm_l Tip) = fm_l
deleteMax (Bin key elt _ fm_l (Bin key_r elt_r s_r fm_rl fm_rr)) =
  mkBalBranch key elt fm_l
              (deleteMax (Bin key_r elt_r s_r fm_rl fm_rr))

------------------------------------------------------------------------------