-- | Partly invertible finite maps.
--
-- Time complexities are given under the assumption that all relevant
-- instance functions, as well as arguments of function type, take
-- constant time, and "n" is the number of keys involved in the
-- operation.

module Agda.Utils.BiMap where

import Prelude hiding (null, lookup)

import Control.Monad.Identity
import Control.Monad.State

import Data.Function
import qualified Data.List as List
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe
import Data.Ord
import Data.Tuple

import GHC.Generics (Generic)

import Agda.Utils.List
import Agda.Utils.Null

-- | Partial injections from a type to some tag type.
--
-- The idea is that 'tag' should be injective on its domain: if
-- @'tag' x = 'tag' y = 'Just' i@, then @x = y@. However, this
-- property does not need to hold globally. The preconditions of the
-- 'BiMap' operations below specify for which sets of values 'tag'
-- must be injective.

class HasTag a where
  type Tag a
  tag :: a -> Maybe (Tag a)

-- | Checks if the function 'tag' is injective for the values in the
-- given list for which the function is defined.

tagInjectiveFor ::
  (Eq v, Eq (Tag v), HasTag v) =>
  [v] -> Bool
tagInjectiveFor vs = and
  [ v1 == v2
  | v1 <- vs
  , v2 <- vs
  , isJust (tag v1)
  , tag v1 == tag v2
  ]

-- | Finite maps from @k@ to @v@, with a way to quickly get from @v@
-- to @k@ for certain values of type @v@ (those for which 'tag' is
-- defined).
--
-- Every value of this type must satisfy 'biMapInvariant'.

data BiMap k v = BiMap
  { biMapThere :: Map k v
  , biMapBack  :: Map (Tag v) k
  }
  deriving Generic

-- | The invariant for 'BiMap'.

biMapInvariant ::
  (Eq k, Eq v, Ord (Tag v), HasTag v) =>
  BiMap k v -> Bool
biMapInvariant m@(BiMap t u) =
  u ==
  Map.fromList
    [ (k', k)
    | (k, v) <- Map.toList t
    , Just k' <- [tag v]
    ]
    &&
  tagInjectiveFor (map snd $ toList m)

instance Null (BiMap k v) where
  empty = BiMap Map.empty Map.empty
  null  = null . biMapThere

-- | Lookup. O(log n).

lookup :: Ord k => k -> BiMap k v -> Maybe v
lookup a = Map.lookup a . biMapThere

-- | Inverse lookup. O(log n).

invLookup :: Ord (Tag v) => Tag v -> BiMap k v -> Maybe k
invLookup k = Map.lookup k . biMapBack

-- | Singleton map. O(1).

singleton :: HasTag v => k -> v -> BiMap k v
singleton k v =
  BiMap
    (Map.singleton k v)
    (case tag v of
       Nothing -> Map.empty
       Just k' -> Map.singleton k' k)

-- | Insertion. Overwrites existing values. O(log n).
--
-- Precondition: See 'insertPrecondition'.

insert ::
  (Ord k, HasTag v, Ord (Tag v)) =>
  k -> v -> BiMap k v -> BiMap k v
insert k v (BiMap t b) =
  BiMap
    (Map.insert k v t)
    (case tag v of
       Nothing -> b'
       Just k' -> Map.insert k' k b')
  where
  b' = case tag =<< Map.lookup k t of
    Nothing -> b
    Just k' -> Map.delete k' b

-- | The precondition for @'insert' k v m@: If @v@ has a 'tag' (@'tag'
-- v ≠ 'Nothing'@), then @m@ must not contain any mapping @k' ↦ v'@
-- for which @k ≠ k'@ and @'tag' v = 'tag' v'@.

insertPrecondition ::
  (Eq k, Eq v, Eq (Tag v), HasTag v) =>
  k -> v -> BiMap k v -> Bool
insertPrecondition k v m =
  case tag v of
    Nothing -> True
    Just _  ->
      not $ any (\(k', v') -> k' /= k && tag v == tag v') $ toList m

-- | Modifies the value at the given position, if any. If the function
-- returns 'Nothing', then the value is removed. O(log n).
--
-- The precondition for @'alterM' f k m@ is that, if the value @v@ is
-- inserted into @m@, and @'tag' v@ is defined, then no key other than
-- @k@ may map to a value @v'@ for which @'tag' v' = 'tag' v@.

alterM ::
  forall k v m. (Ord k, Ord (Tag v), HasTag v, Monad m) =>
  (Maybe v -> m (Maybe v)) -> k -> BiMap k v -> m (BiMap k v)
alterM f k m@(BiMap t b) = do
  (t', r) <- runStateT (Map.alterF f' k t) Nothing
  return $ case r of
    Nothing -> m
    Just r  -> BiMap t' (updateBack r b)
  where
  f' ::
    Maybe v ->
    StateT (Maybe (Maybe (Tag v), Maybe (Tag v))) m (Maybe v)
  f' v = do
    r <- lift (f v)
    put $ Just (tag =<< v, tag =<< r)
    return r

  updateBack (k'1, k'2) =
    if k'1 == k'2
    then id
    else maybe id (flip Map.insert k) k'2 .
         maybe id Map.delete k'1

-- | Modifies the value at the given position, if any. If the function
-- returns 'Nothing', then the value is removed. O(log n).
--
-- Precondition: See 'alterPrecondition'.

alter ::
  forall k v. (Ord k, Ord (Tag v), HasTag v) =>
  (Maybe v -> Maybe v) -> k -> BiMap k v -> BiMap k v
alter f k m = runIdentity $ alterM (Identity . f) k m

-- | The precondition for @'alter' f k m@ is that, if the value @v@ is
-- inserted into @m@, and @'tag' v@ is defined, then no key other than
-- @k@ may map to a value @v'@ for which @'tag' v' = 'tag' v@.

alterPrecondition ::
  (Ord k, Eq v, Eq (Tag v), HasTag v) =>
  (Maybe v -> Maybe v) -> k -> BiMap k v -> Bool
alterPrecondition f k m =
  case tag =<< f (lookup k m) of
    Nothing -> True
    Just k' -> and
      [ Just k' /= tag v
      | (k'', v) <- toList m
      , k'' /= k
      ]

-- | Modifies the value at the given position, if any. If the function
-- returns 'Nothing', then the value is removed. O(log n).
--
-- Precondition: See 'updatePrecondition'.

update ::
  (Ord k, Ord (Tag v), HasTag v) =>
  (v -> Maybe v) -> k -> BiMap k v -> BiMap k v
update f = alter (f =<<)

-- | The precondition for @'update' f k m@ is that, if the value @v@
-- is inserted into @m@, and @'tag' v@ is defined, then no key other
-- than @k@ may map to a value @v'@ for which @'tag' v' = 'tag' v@.

updatePrecondition ::
  (Ord k, Eq v, Eq (Tag v), HasTag v) =>
  (v -> Maybe v) -> k -> BiMap k v -> Bool
updatePrecondition f = alterPrecondition (f =<<)

-- | Modifies the value at the given position, if any. O(log n).
--
-- Precondition: See 'adjustPrecondition'.

adjust ::
  (Ord k, Ord (Tag v), HasTag v) =>
  (v -> v) -> k -> BiMap k v -> BiMap k v
adjust f = update (Just . f)

-- | The precondition for @'adjust' f k m@ is that, if the value @v@
-- is inserted into @m@, and @'tag' v@ is defined, then no key other
-- than @k@ may map to a value @v'@ for which @'tag' v' = 'tag' v@.

adjustPrecondition ::
  (Ord k, Eq v, Eq (Tag v), HasTag v) =>
  (v -> v) -> k -> BiMap k v -> Bool
adjustPrecondition f = updatePrecondition (Just . f)

-- | Inserts a binding into the map. If a binding for the key already
-- exists, then the value obtained by applying the function to the
-- key, the new value and the old value is inserted, and the old value
-- is returned.
--
-- Precondition: See 'insertLookupWithKeyPrecondition'.

insertLookupWithKey ::
  forall k v. (Ord k, Ord (Tag v), HasTag v) =>
  (k -> v -> v -> v) -> k -> v -> BiMap k v -> (Maybe v, BiMap k v)
insertLookupWithKey f k v m = swap $ runState (alterM f' k m) Nothing
  where
  f' :: Maybe v -> State (Maybe v) (Maybe v)
  f' Nothing     = return $ Just v
  f' r@(Just v') = do
    put r
    return $ Just (f k v v')

-- | The precondition for @'insertLookupWithKey' f k v m@ is that, if
-- the value @v'@ is inserted into @m@, and @'tag' v'@ is defined,
-- then no key other than @k@ may map to a value @v''@ for which
-- @'tag' v'' = 'tag' v'@.

insertLookupWithKeyPrecondition ::
  (Ord k, Eq v, Eq (Tag v), HasTag v) =>
  (k -> v -> v -> v) -> k -> v -> BiMap k v -> Bool
insertLookupWithKeyPrecondition f k v =
  alterPrecondition (Just . maybe v (f k v)) k

-- | Changes all the values using the given function, which is also
-- given access to keys. O(n log n).
--
-- Precondition: See 'mapWithKeyPrecondition'.

mapWithKey ::
  (Ord k, Ord (Tag v), HasTag v) =>
  (k -> v -> v) -> BiMap k v -> BiMap k v
mapWithKey f = fromList . map (\(k, v) -> (k, f k v)) . toList

-- | The precondition for @'mapWithKey' f m@: For any two distinct
-- mappings @k₁ ↦ v₁@, @k₂ ↦ v₂@ in @m@ for which the tags of
-- @f k₁ v₁@ and @f k₂ v₂@ are defined the values of @f@ must be
-- distinct (@f k₁ v₁ ≠ f k₂ v₂@). Furthermore 'tag' must be injective
-- for @{ f k v | (k, v) ∈ m }@.

mapWithKeyPrecondition ::
  (Eq k, Eq v, Eq (Tag v), HasTag v) =>
  (k -> v -> v) -> BiMap k v -> Bool
mapWithKeyPrecondition f =
  fromListPrecondition . map (\(k, v) -> (k, f k v)) . toList

-- | Changes all the values using the given function, which is also
-- given access to keys. O(n).
--
-- Precondition: See 'mapWithKeyFixedTagsPrecondition'. Note that tags
-- must not change.

mapWithKeyFixedTags :: (k -> v -> v) -> BiMap k v -> BiMap k v
mapWithKeyFixedTags f (BiMap t b) = BiMap (Map.mapWithKey f t) b

-- | The precondition for @'mapWithKeyFixedTags' f m@ is that, if @m@
-- maps @k@ to @v@, then @'tag' (f k v) == 'tag' v@.

mapWithKeyFixedTagsPrecondition ::
  (Eq v, Eq (Tag v), HasTag v) =>
  (k -> v -> v) -> BiMap k v -> Bool
mapWithKeyFixedTagsPrecondition f m = and
  [ tag (f k v) == tag v
  | (k, v) <- toList m
  ]

-- | Left-biased union. For the time complexity, see 'Map.union'.
--
-- Precondition: See 'unionPrecondition'.

union :: (Ord k, Ord (Tag v)) => BiMap k v -> BiMap k v -> BiMap k v
union (BiMap t1 b1) (BiMap t2 b2) =
  BiMap (Map.union t1 t2) (Map.union b1 b2)

-- The precondition for @'union' m₁ m₂@: If @k@ is mapped to @v₁@ in
-- @m₁@ and @v₂@ in @m₂@, then @'tag' v₂ = 'Nothing'@ or @'tag' v₁ =
-- 'tag' v₂@. Furthermore, if @k₁@ is mapped to @v₁@ in @m₁@ and @k₂@
-- is mapped to @v₂@ in @m₂@, where @'tag' v₁ = 'tag' v₂ = 'Just' k@,
-- then @k₁ = k₂@. Finally 'tag' must be injective for
-- @{v₁ | (k₁, v₁) ∈ m₁} ∪ {v₂ | (k₂, v₂) ∈ m₂, k₂ ∉ m₁}@.

unionPrecondition ::
  (Ord k, Eq v, Eq (Tag v), HasTag v) =>
  BiMap k v -> BiMap k v -> Bool
unionPrecondition m1@(BiMap t1 _) m2@(BiMap t2 _) =
  and
    [ tag v2 == Nothing || tag v1 == tag v2
    | (v1, v2) <- Map.elems $ Map.intersectionWith (,) t1 t2
    ] &&
  and
    [ k1 == k2
    | (k1, v1) <- toList m1
    , (k2, v2) <- toList m2
    , tag v1 == tag v2
    , isJust (tag v1)
    ]
    &&
  tagInjectiveFor
    ([ v1
     | (_, v1) <- toList m1
     ] ++
     [ v2
     | (k2, v2) <- toList m2
     , not (k2 `elem` ks1)
     ])
  where
  ks1 = map fst (toList m1)

-- | Conversion from lists of pairs. Later entries take precedence
-- over earlier ones. O(n log n).
--
-- Precondition: See 'fromListPrecondition'.

fromList ::
  (Ord k, Ord (Tag v), HasTag v) =>
  [(k, v)] -> BiMap k v
fromList = List.foldr (uncurry insert) empty

-- The precondition for @'fromList' kvs@: For all pairs @(k₁, v₁)@,
-- @(k₂, v₂)@ in @kvs@ for which the tags of @v₁@ and @v₂@ are
-- defined, if @v₁ = v₂@ then @k₁ = k₂@. Furthermore 'tag' must be
-- injective for the values in the list.

fromListPrecondition ::
  (Eq k, Eq v, Eq (Tag v), HasTag v) =>
  [(k, v)] -> Bool
fromListPrecondition kvs =
  and
    [ k1 == k2
    | (k1, v1) <- kvs
    , (k2, v2) <- kvs
    , isJust (tag v1)
    , isJust (tag v2)
    , v1 == v2
    ]
    &&
  tagInjectiveFor (map snd kvs)

-- | Conversion to lists of pairs, with the keys in ascending order.
-- O(n).

toList :: BiMap k v -> [(k, v)]
toList = Map.toAscList . biMapThere

-- | The keys, in ascending order. O(n).

keys :: BiMap k v -> [k]
keys = Map.keys . biMapThere

-- | The values, ordered according to the corresponding keys. O(n).

elems :: BiMap k v -> [v]
elems = Map.elems . biMapThere

-- | Conversion from two lists that contain distinct keys/tags, with
-- the keys/tags in ascending order. O(n).
--
-- Precondition: See 'fromDistinctAscendingListsPrecondition'.

fromDistinctAscendingLists ::
  ([(k, v)], [(Tag v, k)]) -> BiMap k v
fromDistinctAscendingLists (t, b) =
  BiMap (Map.fromDistinctAscList t) (Map.fromDistinctAscList b)

-- The precondition for @'fromDistinctAscendingLists' (kvs, kks)@: The
-- lists must contain distinct keys/tags, and must be sorted according
-- to the keys/tags. Furthermore, for every pair @(k, v)@ in the first
-- list for which @'tag' v = 'Just' k'@ there must be a pair @(k', k)@
-- in the second list, and there must not be any other pairs in that
-- list. Finally 'tag' must be injective for @{v | (_, v) ∈ kvs }@.

fromDistinctAscendingListsPrecondition ::
  (Ord k, Eq v, Ord (Tag v), HasTag v) =>
  ([(k, v)], [(Tag v, k)]) -> Bool
fromDistinctAscendingListsPrecondition (kvs, kks) =
  fastDistinct (map fst kvs) && sorted (map fst kvs)
    &&
  fastDistinct (map fst kks) && sorted (map fst kks)
    &&
  kks ==
  List.sortBy (comparing fst)
    [ (k', k)
    | (k, v) <- kvs
    , Just k' <- [tag v]
    ]
    &&
  tagInjectiveFor [ v | (_, v) <- kvs ]

-- | Generates input suitable for 'fromDistinctAscendingLists'. O(n).

toDistinctAscendingLists :: BiMap k v -> ([(k, v)], [(Tag v, k)])
toDistinctAscendingLists (BiMap t b) =
  (Map.toAscList t, Map.toAscList b)

------------------------------------------------------------------------
-- Instances
------------------------------------------------------------------------

instance (Eq k, Eq v) => Eq (BiMap k v) where
  (==) = (==) `on` biMapThere

instance (Ord k, Ord v) => Ord (BiMap k v) where
  compare = compare `on` biMapThere

instance (Show k, Show v) => Show (BiMap k v) where
  show bimap = "Agda.Utils.BiMap.fromList " ++ show (toList bimap)