{- | Representation of a structure mapping unique keys to values. The internal
structure is an AVL tree.
-}
module Mini.Data.Map (
  -- * Type
  Map,

  -- * Combination
  difference,
  intersection,
  union,

  -- * Construction
  empty,
  fromList,
  singleton,

  -- * Conversion
  toAscList,
  toDescList,

  -- * Fold
  foldlWithKey,
  foldrWithKey,

  -- * Modification
  adjust,
  delete,
  filter,
  filterWithKey,
  insert,
  update,

  -- * Query
  isSubmapOf,
  lookup,
  lookupMax,
  lookupMin,
  member,
  null,
  size,

  -- * Traversal
  traverseWithKey,

  -- * Validation
  valid,
) where

import Control.Monad (
  liftM2,
 )
import Data.Bool (
  bool,
 )
import Prelude hiding (
  filter,
  lookup,
  map,
  null,
 )

{-
 - Type
 -}

{- | A map from keys of type /k/ to values of type /a/.

The internal structure is an AVL tree; a tree that is always height-balanced
(the absolute value of the level difference between the left and right
subtrees is at most 1).
-}
data Map k a
  = -- | Empty bin
    E
  | -- | Left-heavy bin
    L (Map k a) k a (Map k a)
  | -- | Balanced bin
    B (Map k a) k a (Map k a)
  | -- | Right-heavy bin
    R (Map k a) k a (Map k a)
  deriving (Eq, Ord)

instance (Show k, Show a) => Show (Map k a) where
  show = curl . map [] go go go
   where
    go _ k a _ recl recr = recl <> show (k, a) <> "," <> recr
    curl = wrap "{" "}" . removeTrailingComma
    wrap open close s = open <> s <> close
    removeTrailingComma s = case s of
      [] -> []
      [_] -> []
      (c : cs) -> c : removeTrailingComma cs

instance Functor (Map k) where
  fmap f =
    map
      E
      (\_ k a _ recl recr -> L recl k (f a) recr)
      (\_ k a _ recl recr -> B recl k (f a) recr)
      (\_ k a _ recl recr -> R recl k (f a) recr)

instance Foldable (Map k) where
  foldr f b = map b go go go where go _ _ a r recl _ = foldr f (f a recl) r

instance Traversable (Map k) where
  traverse = traverseWithKey . const

instance (Ord k) => Semigroup (Map k a) where
  (<>) = union

instance (Ord k) => Monoid (Map k a) where
  mempty = empty

{-
 - Primitive recursion
 -}

-- | Primitive recursion on maps
map
  :: b
  -- ^ Empty bin
  -> (Map k a -> k -> a -> Map k a -> b -> b -> b)
  -- ^ Left-heavy bin: left child, key, value, right child, left recursion,
  -- right recursion
  -> (Map k a -> k -> a -> Map k a -> b -> b -> b)
  -- ^ Balanced bin: left child, key, value, right child, left recursion, right
  -- recursion
  -> (Map k a -> k -> a -> Map k a -> b -> b -> b)
  -- ^ Right-heavy bin: left child, key, value, right child, left recursion,
  -- right recursion
  -> Map k a
  -- ^ Map
  -> b
map e _ _ _ E = e
map e f g h (L l k a r) = f l k a r (map e f g h l) (map e f g h r)
map e f g h (B l k a r) = g l k a r (map e f g h l) (map e f g h r)
map e f g h (R l k a r) = h l k a r (map e f g h l) (map e f g h r)

{-
 - Combination
 -}

-- | \(O(n \log n)\) Map difference (matching only on keys)
difference :: (Ord k) => Map k a -> Map k b -> Map k a
difference = foldrWithKey (\k _ b -> delete k b)

-- | \(O(n \log n)\) Left-biased map intersection (matching only on keys)
intersection :: (Ord k) => Map k a -> Map k b -> Map k a
intersection t1 t2 =
  foldrWithKey
    (\k a b -> bool b (insert k a b) $ k `member` t2)
    empty
    t1

-- | \(O(n \log n)\) Left-biased map union (matching only on keys)
union :: (Ord k) => Map k a -> Map k a -> Map k a
union t = foldrWithKey (\k a b -> bool b (insert k a b) . not $ k `member` t) t

{-
 - Construction
 -}

-- | \(O(1)\) The empty map
empty :: Map k a
empty = E

{- | \(O(n \log n)\) From a tail-biased list of @(key, value)@ pairs to a map
with bins containing the keys and values
-}
fromList :: (Ord k) => [(k, a)] -> Map k a
fromList = foldl (flip $ uncurry insert) empty

-- | \(O(1)\) From a key and a value to a map with a single bin
singleton :: k -> a -> Map k a
singleton k a = B E k a E

{-
 - Conversion
 -}

{- | \(O(n)\) From a map to a list of @(key, value)@ pairs in key-ascending
order
-}
toAscList :: Map k a -> [(k, a)]
toAscList = foldlWithKey (\b k a -> (k, a) : b) []

{- | \(O(n)\) From a map to a list of @(key, value)@ pairs in key-descending
order
-}
toDescList :: Map k a -> [(k, a)]
toDescList = foldrWithKey (\k a b -> (k, a) : b) []

{-
 - Fold
 -}

{- | \(O(n)\) From a left-associative operation on keys and values, a starting
accumulator and a map to a thing
-}
foldlWithKey :: (b -> k -> a -> b) -> b -> Map k a -> b
foldlWithKey f b = map b go go go
 where
  go l k a _ _ recr = foldlWithKey f (f recr k a) l

{- | \(O(n)\) From a right-associative operation on keys and values, a starting
accumulator and a map to a thing
-}
foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
foldrWithKey f b = map b go go go
 where
  go _ k a r recl _ = foldrWithKey f (f k a recl) r

{-
 - Modification
 -}

{- | \(O(\log n)\) From an operation, a key and a map to the map adjusted by
applying the operation to the value associated with the key
-}
adjust :: (Ord k) => (a -> a) -> k -> Map k a -> Map k a
adjust f k0 =
  map
    E
    ( \l k a r recl recr ->
        case compare k0 k of
          LT -> L recl k a r
          EQ -> L l k (f a) r
          GT -> L l k a recr
    )
    ( \l k a r recl recr ->
        case compare k0 k of
          LT -> B recl k a r
          EQ -> B l k (f a) r
          GT -> B l k a recr
    )
    ( \l k a r recl recr ->
        case compare k0 k of
          LT -> R recl k a r
          EQ -> R l k (f a) r
          GT -> R l k a recr
    )

-- | \(O(\log n)\) From a key and a map to the map without the key
delete :: (Ord k) => k -> Map k a -> Map k a
delete k0 t = bool t (go t) (k0 `member` t)
 where
  go =
    map
      (error "Map.delete: L0")
      ( \l k a r _ _ ->
          case compare k0 k of
            LT -> deleteLl l k a r
            EQ -> substituteL l r
            GT -> deleteLr l k a r
      )
      ( \l k a r _ _ ->
          case compare k0 k of
            LT -> deleteBl l k a r
            EQ -> substituteBr l r
            GT -> deleteBr l k a r
      )
      ( \l k a r _ _ ->
          case compare k0 k of
            LT -> deleteRl l k a r
            EQ -> substituteR l r
            GT -> deleteRr l k a r
      )
  deleteRl l k a r =
    map
      (error "Map.delete: L1")
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> checkLeftR (deleteLl ll lk la lr) k a r
            EQ -> checkLeftR (substituteL ll lr) k a r
            GT -> checkLeftR (deleteLr ll lk la lr) k a r
      )
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> R (deleteBl ll lk la lr) k a r
            EQ -> checkLeftR' (substituteBr ll lr) k a r
            GT -> R (deleteBr ll lk la lr) k a r
      )
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> checkLeftR (deleteRl ll lk la lr) k a r
            EQ -> checkLeftR (substituteR ll lr) k a r
            GT -> checkLeftR (deleteRr ll lk la lr) k a r
      )
      l
  deleteRr l k a =
    map
      (error "Map.delete: L2")
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> checkRightR l k a (deleteLl rl rk ra rr)
            EQ -> checkRightR l k a (substituteL rl rr)
            GT -> checkRightR l k a (deleteLr rl rk ra rr)
      )
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> R l k a (deleteBl rl rk ra rr)
            EQ -> checkRightR' l k a (substituteBl rl rr)
            GT -> R l k a (deleteBr rl rk ra rr)
      )
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> checkRightR l k a (deleteRl rl rk ra rr)
            EQ -> checkRightR l k a (substituteR rl rr)
            GT -> checkRightR l k a (deleteRr rl rk ra rr)
      )
  deleteBl l k a r =
    map
      (error "Map.delete: L3")
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> checkLeftB (deleteLl ll lk la lr) k a r
            EQ -> checkLeftB (substituteL ll lr) k a r
            GT -> checkLeftB (deleteLr ll lk la lr) k a r
      )
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> B (deleteBl ll lk la lr) k a r
            EQ -> checkLeftB' (substituteBr ll lr) k a r
            GT -> B (deleteBr ll lk la lr) k a r
      )
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> checkLeftB (deleteRl ll lk la lr) k a r
            EQ -> checkLeftB (substituteR ll lr) k a r
            GT -> checkLeftB (deleteRr ll lk la lr) k a r
      )
      l
  deleteBr l k a =
    map
      (error "Map.delete: L4")
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> checkRightB l k a (deleteLl rl rk ra rr)
            EQ -> checkRightB l k a (substituteL rl rr)
            GT -> checkRightB l k a (deleteLr rl rk ra rr)
      )
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> B l k a (deleteBl rl rk ra rr)
            EQ -> checkRightB' l k a (substituteBl rl rr)
            GT -> B l k a (deleteBr rl rk ra rr)
      )
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> checkRightB l k a (deleteRl rl rk ra rr)
            EQ -> checkRightB l k a (substituteR rl rr)
            GT -> checkRightB l k a (deleteRr rl rk ra rr)
      )
  deleteLl l k a r =
    map
      (error "Map.delete: L5")
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> checkLeftL (deleteLl ll lk la lr) k a r
            EQ -> checkLeftL (substituteL ll lr) k a r
            GT -> checkLeftL (deleteLr ll lk la lr) k a r
      )
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> L (deleteBl ll lk la lr) k a r
            EQ -> checkLeftL' (substituteBr ll lr) k a r
            GT -> L (deleteBr ll lk la lr) k a r
      )
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> checkLeftL (deleteRl ll lk la lr) k a r
            EQ -> checkLeftL (substituteR ll lr) k a r
            GT -> checkLeftL (deleteRr ll lk la lr) k a r
      )
      l
  deleteLr l k a =
    map
      (error "Map.delete: L6")
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> checkRightL l k a (deleteLl rl rk ra rr)
            EQ -> checkRightL l k a (substituteL rl rr)
            GT -> checkRightL l k a (deleteLr rl rk ra rr)
      )
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> L l k a (deleteBl rl rk ra rr)
            EQ -> checkRightL' l k a (substituteBl rl rr)
            GT -> L l k a (deleteBr rl rk ra rr)
      )
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> checkRightL l k a (deleteRl rl rk ra rr)
            EQ -> checkRightL l k a (substituteR rl rr)
            GT -> checkRightL l k a (deleteRr rl rk ra rr)
      )
  rebalanceR l k a =
    map
      (error "Map.delete: L7")
      ( \rl rk ra rr _ _ ->
          map
            (error "Map.delete: L8")
            (\rll rlk rla rlr _ _ -> B (B l k a rll) rlk rla (R rlr rk ra rr))
            (\rll rlk rla rlr _ _ -> B (B l k a rll) rlk rla (B rlr rk ra rr))
            (\rll rlk rla rlr _ _ -> B (L l k a rll) rlk rla (B rlr rk ra rr))
            rl
      )
      (\rl rk ra rr _ _ -> L (R l k a rl) rk ra rr)
      (\rl rk ra rr _ _ -> B (B l k a rl) rk ra rr)
  rebalanceL l k a r =
    map
      (error "Map.delete: L9")
      (\ll lk la lr _ _ -> B ll lk la (B lr k a r))
      (\ll lk la lr _ _ -> R ll lk la (L lr k a r))
      ( \ll lk la lr _ _ ->
          map
            (error "Map.delete: L10")
            (\lrl lrk lra lrr _ _ -> B (B ll lk la lrl) lrk lra (R lrr k a r))
            (\lrl lrk lra lrr _ _ -> B (B ll lk la lrl) lrk lra (B lrr k a r))
            (\lrl lrk lra lrr _ _ -> B (L ll lk la lrl) lrk lra (B lrr k a r))
            lr
      )
      l
  checkLeftR l k a r =
    map
      (error "Map.delete: L11")
      (\_ _ _ _ _ _ -> R l k a r)
      (\_ _ _ _ _ _ -> rebalanceR l k a r)
      (\_ _ _ _ _ _ -> R l k a r)
      l
  checkLeftB l k a r =
    map
      (error "Map.delete: L12")
      (\_ _ _ _ _ _ -> B l k a r)
      (\_ _ _ _ _ _ -> R l k a r)
      (\_ _ _ _ _ _ -> B l k a r)
      l
  checkLeftL l k a r =
    map
      (error "Map.delete: L13")
      (\_ _ _ _ _ _ -> L l k a r)
      (\_ _ _ _ _ _ -> B l k a r)
      (\_ _ _ _ _ _ -> L l k a r)
      l
  checkRightR l k a r =
    map
      (error "Map.delete: L14")
      (\_ _ _ _ _ _ -> R l k a r)
      (\_ _ _ _ _ _ -> B l k a r)
      (\_ _ _ _ _ _ -> R l k a r)
      r
  checkRightB l k a r =
    map
      (error "Map.delete: L15")
      (\_ _ _ _ _ _ -> B l k a r)
      (\_ _ _ _ _ _ -> L l k a r)
      (\_ _ _ _ _ _ -> B l k a r)
      r
  checkRightL l k a r =
    map
      (error "Map.delete: L16")
      (\_ _ _ _ _ _ -> L l k a r)
      (\_ _ _ _ _ _ -> rebalanceL l k a r)
      (\_ _ _ _ _ _ -> L l k a r)
      r
  substituteR l =
    map
      (error "Map.delete: L17")
      ( \rl rk ra rr _ _ ->
          (\(k, a, r) -> checkRightR l k a r) $
            popLeftL rl rk ra rr
      )
      ( \rl rk ra rr _ _ ->
          (\(k, a, r) -> checkRightR' l k a r) $
            popLeftB rl rk ra rr
      )
      ( \rl rk ra rr _ _ ->
          (\(k, a, r) -> checkRightR l k a r) $
            popLeftR rl rk ra rr
      )
  substituteBr l =
    map
      E
      ( \rl rk ra rr _ _ ->
          (\(k, a, r) -> checkRightB l k a r) $
            popLeftL rl rk ra rr
      )
      ( \rl rk ra rr _ _ ->
          (\(k, a, r) -> checkRightB' l k a r) $
            popLeftB rl rk ra rr
      )
      ( \rl rk ra rr _ _ ->
          (\(k, a, r) -> checkRightB l k a r) $
            popLeftR rl rk ra rr
      )
  substituteBl l r =
    map
      E
      ( \ll lk la lr _ _ ->
          (\(l', k, a) -> checkLeftB l' k a r) $
            popRightL ll lk la lr
      )
      ( \ll lk la lr _ _ ->
          (\(l', k, a) -> checkLeftB' l' k a r) $
            popRightB ll lk la lr
      )
      ( \ll lk la lr _ _ ->
          (\(l', k, a) -> checkLeftB l' k a r) $
            popRightR ll lk la lr
      )
      l
  substituteL l r =
    map
      (error "Map.delete: L18")
      ( \ll lk la lr _ _ ->
          (\(l', k, a) -> checkLeftL l' k a r) $
            popRightL ll lk la lr
      )
      ( \ll lk la lr _ _ ->
          (\(l', k, a) -> checkLeftL' l' k a r) $
            popRightB ll lk la lr
      )
      ( \ll lk la lr _ _ ->
          (\(l', k, a) -> checkLeftL l' k a r) $
            popRightR ll lk la lr
      )
      l
  checkLeftR' l k a r =
    map
      (rebalanceR l k a r)
      (\_ _ _ _ _ _ -> R l k a r)
      (\_ _ _ _ _ _ -> R l k a r)
      (\_ _ _ _ _ _ -> R l k a r)
      l
  checkLeftB' l k a r =
    map
      (R l k a r)
      (\_ _ _ _ _ _ -> B l k a r)
      (\_ _ _ _ _ _ -> B l k a r)
      (\_ _ _ _ _ _ -> B l k a r)
      l
  checkLeftL' l k a r =
    map
      (B l k a r)
      (\_ _ _ _ _ _ -> L l k a r)
      (\_ _ _ _ _ _ -> L l k a r)
      (\_ _ _ _ _ _ -> L l k a r)
      l
  checkRightR' l k a r =
    map
      (B l k a r)
      (\_ _ _ _ _ _ -> R l k a r)
      (\_ _ _ _ _ _ -> R l k a r)
      (\_ _ _ _ _ _ -> R l k a r)
      r
  checkRightB' l k a r =
    map
      (L l k a r)
      (\_ _ _ _ _ _ -> B l k a r)
      (\_ _ _ _ _ _ -> B l k a r)
      (\_ _ _ _ _ _ -> B l k a r)
      r
  checkRightL' l k a r =
    map
      (rebalanceL l k a r)
      (\_ _ _ _ _ _ -> L l k a r)
      (\_ _ _ _ _ _ -> L l k a r)
      (\_ _ _ _ _ _ -> L l k a r)
      r
  popLeftR l k a r =
    map
      (k, a, r)
      ( \ll lk la lr _ _ ->
          (\(k', a', l') -> (k', a', checkLeftR l' k a r)) $
            popLeftL ll lk la lr
      )
      (\ll lk la lr _ _ -> popLeftRB ll lk la lr k a r)
      ( \ll lk la lr _ _ ->
          (\(k', a', l') -> (k', a', checkLeftR l' k a r)) $
            popLeftR ll lk la lr
      )
      l
  popLeftB l k a r =
    map
      (k, a, E)
      (\ll lk la lr _ _ -> popLeftBL ll lk la lr k a r)
      (\ll lk la lr _ _ -> popLeftBB ll lk la lr k a r)
      (\ll lk la lr _ _ -> popLeftBR ll lk la lr k a r)
      l
  popLeftL l k a r =
    map
      (error "Map.delete: L19")
      ( \ll lk la lr _ _ ->
          (\(k', a', l') -> (k', a', checkLeftL l' k a r)) $
            popLeftL ll lk la lr
      )
      (\ll lk la lr _ _ -> popLeftLB ll lk la lr k a r)
      ( \ll lk la lr _ _ ->
          (\(k', a', l') -> (k', a', checkLeftL l' k a r)) $
            popLeftR ll lk la lr
      )
      l
  popLeftRB ll lk la lr k a r =
    map
      (lk, la, rebalanceR E k a r)
      ( \lll llk lla llr _ _ ->
          (\(k', a', l) -> (k', a', R l k a r)) $
            popLeftBL lll llk lla llr lk la lr
      )
      ( \lll llk lla llr _ _ ->
          (\(k', a', l) -> (k', a', R l k a r)) $
            popLeftBB lll llk lla llr lk la lr
      )
      ( \lll llk lla llr _ _ ->
          (\(k', a', l) -> (k', a', R l k a r)) $
            popLeftBR lll llk lla llr lk la lr
      )
      ll
  popLeftBB ll lk la lr k a r =
    map
      (lk, la, R E k a r)
      ( \lll llk lla llr _ _ ->
          (\(k', a', l) -> (k', a', B l k a r)) $
            popLeftBL lll llk lla llr lk la lr
      )
      ( \lll llk lla llr _ _ ->
          (\(k', a', l) -> (k', a', B l k a r)) $
            popLeftBB lll llk lla llr lk la lr
      )
      ( \lll llk lla llr _ _ ->
          (\(k', a', l) -> (k', a', B l k a r)) $
            popLeftBR lll llk lla llr lk la lr
      )
      ll
  popLeftLB ll lk la lr k a r =
    map
      (lk, la, B E k a E)
      ( \lll llk lla llr _ _ ->
          (\(k', a', l) -> (k', a', L l k a r)) $
            popLeftBL lll llk lla llr lk la lr
      )
      ( \lll llk lla llr _ _ ->
          (\(k', a', l) -> (k', a', L l k a r)) $
            popLeftBB lll llk lla llr lk la lr
      )
      ( \lll llk lla llr _ _ ->
          (\(k', a', l) -> (k', a', L l k a r)) $
            popLeftBR lll llk lla llr lk la lr
      )
      ll
  popLeftBR ll lk la lr k a r =
    (\(k', a', l) -> (k', a', checkLeftB l k a r)) $
      popLeftR ll lk la lr
  popLeftBL ll lk la lr k a r =
    (\(k', a', l) -> (k', a', checkLeftB l k a r)) $
      popLeftL ll lk la lr
  popRightR l k a =
    map
      (error "Map.delete: L20")
      ( \rl rk ra rr _ _ ->
          (\(r, k', a') -> (checkRightR l k a r, k', a')) $
            popRightL rl rk ra rr
      )
      (\rl rk ra rr _ _ -> popRightRB l k a rl rk ra rr)
      ( \rl rk ra rr _ _ ->
          (\(r, k', a') -> (checkRightR l k a r, k', a')) $
            popRightR rl rk ra rr
      )
  popRightB l k a =
    map
      (E, k, a)
      (\rl rk ra rr _ _ -> popRightBL l k a rl rk ra rr)
      (\rl rk ra rr _ _ -> popRightBB l k a rl rk ra rr)
      (\rl rk ra rr _ _ -> popRightBR l k a rl rk ra rr)
  popRightL l k a =
    map
      (l, k, a)
      ( \rl rk ra rr _ _ ->
          (\(r, k', a') -> (checkRightL l k a r, k', a')) $
            popRightL rl rk ra rr
      )
      (\rl rk ra rr _ _ -> popRightLB l k a rl rk ra rr)
      ( \rl rk ra rr _ _ ->
          (\(r, k', a') -> (checkRightL l k a r, k', a')) $
            popRightR rl rk ra rr
      )
  popRightRB l k a rl rk ra =
    map
      (B E k a E, rk, ra)
      ( \rrl rrk rra rrr _ _ ->
          (\(r, k', a') -> (R l k a r, k', a')) $
            popRightBL rl rk ra rrl rrk rra rrr
      )
      ( \rrl rrk rra rrr _ _ ->
          (\(r, k', a') -> (R l k a r, k', a')) $
            popRightBB rl rk ra rrl rrk rra rrr
      )
      ( \rrl rrk rra rrr _ _ ->
          (\(r, k', a') -> (R l k a r, k', a')) $
            popRightBR rl rk ra rrl rrk rra rrr
      )
  popRightBB l k a rl rk ra =
    map
      (L l k a E, rk, ra)
      ( \rrl rrk rra rrr _ _ ->
          (\(r, k', a') -> (B l k a r, k', a')) $
            popRightBL rl rk ra rrl rrk rra rrr
      )
      ( \rrl rrk rra rrr _ _ ->
          (\(r, k', a') -> (B l k a r, k', a')) $
            popRightBB rl rk ra rrl rrk rra rrr
      )
      ( \rrl rrk rra rrr _ _ ->
          (\(r, k', a') -> (B l k a r, k', a')) $
            popRightBR rl rk ra rrl rrk rra rrr
      )
  popRightLB l k a rl rk ra =
    map
      (rebalanceL l k a E, rk, ra)
      ( \rrl rrk rra rrr _ _ ->
          (\(r, k', a') -> (L l k a r, k', a')) $
            popRightBL rl rk ra rrl rrk rra rrr
      )
      ( \rrl rrk rra rrr _ _ ->
          (\(r, k', a') -> (L l k a r, k', a')) $
            popRightBB rl rk ra rrl rrk rra rrr
      )
      ( \rrl rrk rra rrr _ _ ->
          (\(r, k', a') -> (L l k a r, k', a')) $
            popRightBR rl rk ra rrl rrk rra rrr
      )
  popRightBR l k a rl rk ra rr =
    (\(r, k', a') -> (checkRightB l k a r, k', a')) $
      popRightR rl rk ra rr
  popRightBL l k a rl rk ra rr =
    (\(r, k', a') -> (checkRightB l k a r, k', a')) $
      popRightL rl rk ra rr

{- | \(O(n)\) From a predicate and a map to the map with values satisfying the
predicate
-}
filter :: (Ord k) => (a -> Bool) -> Map k a -> Map k a
filter p = foldrWithKey (\k a b -> bool b (insert k a b) $ p a) empty

{- | \(O(n)\) From a predicate and a map to the map with keys and values
satisfying the predicate
-}
filterWithKey :: (Ord k) => (k -> a -> Bool) -> Map k a -> Map k a
filterWithKey p = foldrWithKey (\k a b -> bool b (insert k a b) $ p k a) empty

{- | \(O(\log n)\) From a key, a value and a map to the map including a bin
containing the key and the value
-}
insert :: (Ord k) => k -> a -> Map k a -> Map k a
insert k0 a0 =
  map
    (B E k0 a0 E)
    (\l k a r _ _ -> insertL l k a r)
    (\l k a r _ _ -> insertB l k a r)
    (\l k a r _ _ -> insertR l k a r)
 where
  insertR l k a r =
    case compare k0 k of
      LT -> insertRl l k a r
      EQ -> R l k0 a0 r
      GT -> insertRr l k a r
  insertB l k a r =
    case compare k0 k of
      LT -> insertBl l k a r
      EQ -> B l k0 a0 r
      GT -> insertBr l k a r
  insertL l k a r =
    case compare k0 k of
      LT -> insertLl l k a r
      EQ -> L l k0 a0 r
      GT -> insertLr l k a r
  insertRl l k a r =
    map
      (B (B E k0 a0 E) k a r)
      (\ll lk la lr _ _ -> R (insertL ll lk la lr) k a r)
      ( \ll lk la lr _ _ ->
          let l' = insertB ll lk la lr
           in map
                (error "Map.insert: L0")
                (\_ _ _ _ _ _ -> B l' k a r)
                (\_ _ _ _ _ _ -> R l' k a r)
                (\_ _ _ _ _ _ -> B l' k a r)
                l'
      )
      (\ll lk la lr _ _ -> R (insertR ll lk la lr) k a r)
      l
  insertBl l k a r =
    map
      (L (B E k0 a0 E) k a r)
      (\ll lk la lr _ _ -> B (insertL ll lk la lr) k a r)
      ( \ll lk la lr _ _ ->
          let l' = insertB ll lk la lr
           in map
                (error "Map.insert: L1")
                (\_ _ _ _ _ _ -> L l' k a r)
                (\_ _ _ _ _ _ -> B l' k a r)
                (\_ _ _ _ _ _ -> L l' k a r)
                l'
      )
      (\ll lk la lr _ _ -> B (insertR ll lk la lr) k a r)
      l
  insertBr l k a =
    map
      (R l k a (B E k0 a0 E))
      (\rl rk ra rr _ _ -> B l k a (insertL rl rk ra rr))
      ( \rl rk ra rr _ _ ->
          let r = insertB rl rk ra rr
           in map
                (error "Map.insert: L2")
                (\_ _ _ _ _ _ -> R l k a r)
                (\_ _ _ _ _ _ -> B l k a r)
                (\_ _ _ _ _ _ -> R l k a r)
                r
      )
      (\rl rk ra rr _ _ -> B l k a (insertR rl rk ra rr))
  insertLr l k a =
    map
      (B l k a (B E k0 a0 E))
      (\rl rk ra rr _ _ -> L l k a (insertL rl rk ra rr))
      ( \rl rk ra rr _ _ ->
          let r = insertB rl rk ra rr
           in map
                (error "Map.insert: L3")
                (\_ _ _ _ _ _ -> B l k a r)
                (\_ _ _ _ _ _ -> L l k a r)
                (\_ _ _ _ _ _ -> B l k a r)
                r
      )
      (\rl rk ra rr _ _ -> L l k a (insertR rl rk ra rr))
  insertRr l k a =
    map
      (error "Map.insert: L4")
      (\rl rk ra rr _ _ -> R l k a (insertL rl rk ra rr))
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> insertRrl l k a rl rk ra rr
            EQ -> R l k a (B rl k0 a0 rr)
            GT -> insertRrr l k a rl rk ra rr
      )
      (\rl rk ra rr _ _ -> R l k a (insertR rl rk ra rr))
  insertLl l k a r =
    map
      (error "Map.insert: L5")
      (\ll lk la lr _ _ -> L (insertL ll lk la lr) k a r)
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> insertLll ll lk la lr k a r
            EQ -> L (B ll k0 a0 lr) k a r
            GT -> insertLlr ll lk la lr k a r
      )
      (\ll lk la lr _ _ -> L (insertR ll lk la lr) k a r)
      l
  insertRrr l k a rl rk ra =
    map
      (B (B l k a rl) rk ra (B E k0 a0 E))
      (\rrl rrk rra rrr _ _ -> R l k a (B rl rk ra (insertL rrl rrk rra rrr)))
      ( \rrl rrk rra rrr _ _ ->
          let rr = insertB rrl rrk rra rrr
           in map
                (error "Map.insert: L6")
                (\_ _ _ _ _ _ -> B (B l k a rl) rk ra rr)
                (\_ _ _ _ _ _ -> R l k a (B rl rk ra rr))
                (\_ _ _ _ _ _ -> B (B l k a rl) rk ra rr)
                rr
      )
      (\rrl rrk rra rrr _ _ -> R l k a (B rl rk ra (insertR rrl rrk rra rrr)))
  insertLll ll lk la lr k a r =
    map
      (B (B E k0 a0 E) lk la (B lr k a r))
      (\lll llk lla llr _ _ -> L (B (insertL lll llk lla llr) lk la lr) k a r)
      ( \lll llk lla llr _ _ ->
          let ll' = insertB lll llk lla llr
           in map
                (error "Map.insert: L7")
                (\_ _ _ _ _ _ -> B ll' lk la (B lr k a r))
                (\_ _ _ _ _ _ -> L (B ll' lk la lr) k a r)
                (\_ _ _ _ _ _ -> B ll' lk la (B lr k a r))
                ll'
      )
      (\lll llk lla llr _ _ -> L (B (insertR lll llk lla llr) lk la lr) k a r)
      ll
  insertRrl l k a rl rk ra rr =
    map
      (B (B l k a E) k0 a0 (B E rk ra rr))
      (\rll rlk rla rlr _ _ -> R l k a (B (insertL rll rlk rla rlr) rk ra rr))
      ( \rll rlk rla rlr _ _ ->
          let rl' = insertB rll rlk rla rlr
           in map
                (error "Map.insert: L8")
                ( \rll' rlk' rla' rlr' _ _ ->
                    B
                      (B l k a rll')
                      rlk'
                      rla'
                      (R rlr' rk ra rr)
                )
                (\_ _ _ _ _ _ -> R l k a (B rl' rk ra rr))
                ( \rll' rlk' rla' rlr' _ _ ->
                    B
                      (L l k a rll')
                      rlk'
                      rla'
                      (B rlr' rk ra rr)
                )
                rl'
      )
      (\rll rlk rla rlr _ _ -> R l k a (B (insertR rll rlk rla rlr) rk ra rr))
      rl
  insertLlr ll lk la lr k a r =
    map
      (B (B ll lk la E) k0 a0 (B E k a r))
      (\lrl lrk lra lrr _ _ -> L (B ll lk la (insertL lrl lrk lra lrr)) k a r)
      ( \lrl lrk lra lrr _ _ ->
          let lr' = insertB lrl lrk lra lrr
           in map
                (error "Map.insert: L9")
                ( \lrl' lrk' lra' lrr' _ _ ->
                    B
                      (B ll lk la lrl')
                      lrk'
                      lra'
                      (R lrr' k a r)
                )
                (\_ _ _ _ _ _ -> L (B ll lk la lr') k a r)
                ( \lrl' lrk' lra' lrr' _ _ ->
                    B
                      (L ll lk la lrl')
                      lrk'
                      lra'
                      (B lrr' k a r)
                )
                lr'
      )
      (\lrl lrk lra lrr _ _ -> L (B ll lk la (insertR lrl lrk lra lrr)) k a r)
      lr

{- | \(O(\log n)\) From an operation, a key and a map to the map updated by
applying the operation to the value associated with the key (setting if
'Just', deleting if 'Nothing')
-}
update :: (Ord k) => (a -> Maybe a) -> k -> Map k a -> Map k a
update f k t =
  maybe
    t
    ( maybe
        (delete k t)
        (\a -> insert k a t)
        . f
    )
    $ lookup k t

{-
 - Query
 -}

{- | \(O(n \log n)\) From a map and another map to whether the former is a
submap of the latter (matching on keys and values)
-}
isSubmapOf :: (Ord k, Eq a) => Map k a -> Map k a -> Bool
isSubmapOf p q =
  foldrWithKey
    (\k a b -> maybe False ((&& b) . (== a)) $ lookup k q)
    True
    p

{- | \(O(\log n)\) From a key and a map to the value associated with the key in
the map
-}
lookup :: (Ord k) => k -> Map k a -> Maybe a
lookup k = map Nothing go go go
 where
  go _ k' a _ recl recr = case compare k k' of
    LT -> recl
    EQ -> Just a
    GT -> recr

{- | \(O(\log n)\) From a map to the value associated with the maximum key in
the map
-}
lookupMax :: Map k a -> Maybe a
lookupMax = map Nothing go go go
 where
  go _ _ a r _ recr = map (Just a) go' go' go' r where go' _ _ _ _ _ _ = recr

{- | \(O(\log n)\) From a map to the value associated with the minimum key in
the map
-}
lookupMin :: Map k a -> Maybe a
lookupMin = map Nothing go go go
 where
  go l _ a _ recl _ = map (Just a) go' go' go' l where go' _ _ _ _ _ _ = recl

-- | \(O(\log n)\) From a key and a map to whether the key is in the map
member :: (Ord k) => k -> Map k a -> Bool
member k = map False go go go
 where
  go _ k' _ _ recl recr = case compare k k' of
    LT -> recl
    EQ -> True
    GT -> recr

-- | \(O(1)\) From a map to whether the map is empty
null :: Map k a -> Bool
null = map True go go go where go _ _ _ _ _ _ = False

-- | \(O(n)\) From a map to the size of the map
size :: Map k a -> Int
size =
  map
    0
    (\_ _ _ _ recl recr -> 1 + recl + recr)
    (\_ _ _ _ recl recr -> 1 + recl + recr)
    (\_ _ _ _ recl recr -> 1 + recl + recr)

{-
 - Traversal
 -}

{- | \(O(n)\) From a lifting operation on keys and values and a map to a lifted
map
-}
traverseWithKey :: (Applicative f) => (k -> a -> f b) -> Map k a -> f (Map k b)
traverseWithKey f =
  map
    (pure E)
    (\_ k a _ recl recr -> L <$> recl <*> pure k <*> f k a <*> recr)
    (\_ k a _ recl recr -> B <$> recl <*> pure k <*> f k a <*> recr)
    (\_ k a _ recl recr -> R <$> recl <*> pure k <*> f k a <*> recr)

{-
 - Validation
 -}

{- | \(O(n)\) From a map to whether its internal structure is valid, i.e.
height-balanced and ordered
-}
valid :: (Ord k) => Map k a -> Bool
valid = liftM2 (&&) balanced ordered
 where
  balanced =
    map
      True
      (\l _ _ r recl recr -> levels l - levels r == 1 && recl && recr)
      (\l _ _ r recl recr -> levels l - levels r == 0 && recl && recr)
      (\l _ _ r recl recr -> levels r - levels l == 1 && recl && recr)
  levels = map 0 go go go where go _ _ _ _ recl recr = 1 + max recl recr :: Int
  ordered = map True go go go
   where
    go l k _ r recl recr =
      map
        True
        (\_ lk _ _ _ _ -> lk < k && recl && recr)
        (\_ lk _ _ _ _ -> lk < k && recl && recr)
        (\_ lk _ _ _ _ -> lk < k && recl && recr)
        l
        && map
          True
          (\_ rk _ _ _ _ -> rk > k && recl && recr)
          (\_ rk _ _ _ _ -> rk > k && recl && recr)
          (\_ rk _ _ _ _ -> rk > k && recl && recr)
          r