{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveGeneric #-}
module Data.BimapMany
(
-- * BimapMany type
BimapMany
-- * Construction
, empty
, singleton
, fromMap
, fromSet
-- ** From unordered lists
, fromList
-- * Insertion
, insert
-- * Deletion/Update
, delete
, deleteL
, deleteR
-- * Query
-- ** Lookup
, lookup
, lookupL
, lookupR
, lookupL'
, lookupR'
-- ** Size
, null
, size
, sizeL
, sizeR
-- * Combine
, union
-- * Conversion
-- ** Maps
, toMap
-- ** Lists
, toList
-- * Debugging
, valid
) where

import Prelude hiding (abs, lookup, null)

import Data.Map.Signature (Map)
import qualified Data.Map.Signature as M
import qualified Data.Map.Strict as MS
import Data.Set (Set)
import qualified Data.Set as S

import Control.DeepSeq (NFData)
import GHC.Generics (Generic)

import Data.Function (on)
import Data.List (foldl', groupBy, sort)
import Data.Maybe (fromMaybe)

-- * BimapMany type
-------------------

-- NOTE/TODO: considering different internal representations
--
-- The current implementation uses a Map (a, b) c, but it's far from being the
-- only solution.
--
-- (Map a (Map b c)) (Map b (Set a)
-- It would take less space, having less set-like structures
-- It's asymmetric, which would make a few operations more expensive:
-- * lookupL (but lookupL' could be used instead)
-- Other operations are more expensive due to the lack of Map (a, b) c
-- * toMap
-- * size (but it could be stored separately)
--
-- (Map a (Map b c)) (Map b (Map a c)
-- It would take less space, having less set-like structures
-- Some operations are more expensive due to the lack of Map (a, b) c
-- * toMap
-- * size (but it could be stored separately)
-- It's symetric but it "duplicates" c. On one hand it gets shared anyway, but
-- on the other hand the less used map could accumulate thunks if care is not
-- taken.
--
-- Without analyzing this more in depth, the current option looks safest.

-- NOTE: operations on l and r are always strict
-- When operating on l and r, the Map value is another set structure, not a
-- BimapMany value (c), so operations should be strict.

data BimapMany a b c = BimapMany
  !(MS.Map a (Set b)) -- l
  !(MS.Map b (Set a)) -- r
  !(Map (a, b) c) -- m
  -- Invariants (checked by 'valid'):
  --   (a, bs) ∈ l, b ∈ bs,   => (a, b) ∈ m, (b, as) ∈ r, a ∈ as
  --   (b, as) ∈ r, a ∈ as,   => (a, b) ∈ m, (a, bs) ∈ l, b ∈ bs
  --   ((a, b), _) ∈ m        => (a, bs) ∈ l, (b, as) ∈ r, a ∈ as, b ∈ bs
  -- TODO use proper lazy/strict functions for these instances
  -- MAYBE move the data definition to a separate module and share it, like Map
  -- Or better, provide coercion functions
  deriving (Functor, Foldable, Traversable, Generic)

instance (Show a, Show b, Show c) => Show (BimapMany a b c) where
  show x = "fromList " ++ show (toList x)

instance (Eq a, Eq b, Eq c) => Eq (BimapMany a b c) where
  {-# INLINABLE (==) #-}
  (==) = (==) `on` toMap

instance (Ord a, Ord b, Ord c) => Ord (BimapMany a b c) where
  {-# INLINABLE compare #-}
  compare = compare `on` toMap

instance (Ord a, Ord b) => Semigroup (BimapMany a b c) where
  {-# INLINABLE (<>) #-}
  (<>) = union

instance (Ord a, Ord b) => Monoid (BimapMany a b c) where
  {-# INLINABLE mempty #-}
  mempty = empty

instance (NFData a, NFData b, NFData c) => NFData (BimapMany a b c)

-- * Construction
-----------------

{-# INLINABLE empty #-}
empty :: BimapMany a b c
empty = BimapMany MS.empty MS.empty M.empty

{-# INLINABLE singleton #-}
singleton :: a -> b -> c -> BimapMany a b c
singleton a b c = BimapMany l r m
  where
    l = MS.singleton a $ S.singleton b
    r = MS.singleton b $ S.singleton a
    m = M.singleton (a, b) c

-- TODO write a warning about strict/lazy map that must correspond to BimapMany
{-# INLINABLE fromMap #-}
fromMap :: (Ord a, Ord b) => Map (a, b) c -> BimapMany a b c
fromMap m = BimapMany l r m
  where
    abs = M.keys m
    l = ascListToMapSet abs
    bas = sort $ (\(a, b) -> (b, a)) <$> abs
    r = ascListToMapSet bas

{-# INLINABLE fromSet #-}
fromSet :: (Ord a, Ord b) => (a -> b -> c) -> Set (a, b) -> BimapMany a b c
fromSet f s = BimapMany l r m
  where
    abs = S.toAscList s
    l = ascListToMapSet abs
    bas = sort $ (\(a, b) -> (b, a)) <$> abs
    r = ascListToMapSet bas
    m = M.fromSet (uncurry f) s

-- ** From unordered lists
--------------------------

{-# INLINABLE fromList #-}
fromList :: (Ord a, Ord b) => [(a, b, c)] -> BimapMany a b c
fromList = foldl' (\m (a, b, c) -> insert a b c m) empty

-- TODO fromListWith(Key)
-- TODO asc and desc unsafe construction

-- * Insertion
--------------

{-# INLINABLE insert #-}
insert :: (Ord a, Ord b) => a -> b -> c -> BimapMany a b c -> BimapMany a b c
insert a b c (BimapMany l r m) = BimapMany l' r' m'
  where
    l' = MS.alter (Just . maybe (S.singleton b) (S.insert b)) a l
    r' = MS.alter (Just . maybe (S.singleton a) (S.insert a)) b r
    m' = M.insert (a, b) c m

-- TODO insertWith etc

-- * Deletion/Update
--------------------

{-# INLINABLE delete #-}
delete :: (Ord a, Ord b) => a -> b -> BimapMany a b c -> BimapMany a b c
delete a b (BimapMany l r m) = BimapMany l' r' m'
  where
    l' = MS.update (setDelete' b) a l
    r' = MS.update (setDelete' a) b r
    m' = M.delete (a, b) m

{-# INLINABLE deleteL #-}
deleteL :: (Ord a, Ord b) => a -> BimapMany a b c -> BimapMany a b c
deleteL a (BimapMany l r m) = BimapMany l' r' m'
  where
    bs = fromMaybe S.empty $ MS.lookup a l
    l' = MS.delete a l
    r' = S.foldr' (MS.update $ setDelete' a) r bs
    m' = S.foldr' (\b -> M.delete (a, b)) m bs

{-# INLINABLE deleteR #-}
deleteR :: (Ord a, Ord b) => b -> BimapMany a b c -> BimapMany a b c
deleteR b (BimapMany l r m) = BimapMany l' r' m'
  where
    as = fromMaybe S.empty $ MS.lookup b r
    r' = MS.delete b r
    l' = S.foldr' (MS.update $ setDelete' b) l as
    m' = S.foldr' (\a -> M.delete (a, b)) m as

-- TODO updating functions

-- * Query
----------

-- ** Lookup
------------

{-# INLINABLE lookup #-}
lookup :: (Ord a, Ord b) => a -> b -> BimapMany a b c -> Maybe c
lookup a b (BimapMany _ _ m) = M.lookup (a, b) m

{-# INLINABLE lookupL #-}
lookupL :: Ord a => a -> BimapMany a b c -> Set b
lookupL a (BimapMany l _ _) = mSetToSet $ MS.lookup a l

{-# INLINABLE lookupR #-}
lookupR :: Ord b => b -> BimapMany a b c -> Set a
lookupR b (BimapMany _ r _) = mSetToSet $ MS.lookup b r

{-# INLINABLE lookupL' #-}
lookupL' :: (Ord a, Ord b) => a -> BimapMany a b c -> Map b c
lookupL' a (BimapMany l _ m) = M.fromSet (\b -> m M.! (a, b)) bs
  where bs = mSetToSet $ MS.lookup a l

{-# INLINABLE lookupR' #-}
lookupR' :: (Ord a, Ord b) => b -> BimapMany a b c -> Map a c
lookupR' b (BimapMany _ r m) = M.fromSet (\a -> m M.! (a, b)) as
  where as = mSetToSet $ MS.lookup b r

-- ** Size
----------

{-# INLINABLE null #-}
null :: BimapMany a b c -> Bool
null (BimapMany _ _ m) = M.null m

{-# INLINABLE size #-}
size :: BimapMany a b c -> Int
size (BimapMany _ _ m) = M.size m

{-# INLINABLE sizeL #-}
sizeL :: BimapMany a b c -> Int
sizeL (BimapMany l _ _) = MS.size l

{-# INLINABLE sizeR #-}
sizeR :: BimapMany a b c -> Int
sizeR (BimapMany _ r _) = MS.size r

-- * Combine
------------

{-# INLINABLE union #-}
union :: (Ord a, Ord b) => BimapMany a b c -> BimapMany a b c -> BimapMany a b c
union (BimapMany l1 r1 m1) (BimapMany l2 r2 m2) = BimapMany l r m
  where
    l = MS.unionWith S.union l1 l2
    r = MS.unionWith S.union r1 r2
    m = M.union m1 m2

-- * Conversion
---------------

-- ** Maps
----------

{-# INLINABLE toMap #-}
toMap :: BimapMany a b c -> Map (a, b) c
toMap (BimapMany _ _ m) = m

-- MAYBE BimapMany a b c -> Map a (Set b) (and mirrored)
-- MAYBE toNestedMapL :: BimapMany a b c -> Map a (Map b c) (and mirrored)

-- ** Lists
-----------

{-# INLINABLE toList #-}
toList :: BimapMany a b c -> [(a, b, c)]
toList (BimapMany _ _ m) = (\((a, b), c) -> (a, b, c)) <$> M.toList m

-- * Debugging
--------------

valid :: (Ord a, Ord b) => BimapMany a b c -> Bool
valid (BimapMany l r m) = prop1 && prop2 && prop3
  where
    prop1 = all (\(a, b) -> M.member (a, b) m &&
                            maybe False (S.member a) (r MS.!? b)) $
                foldMap (\(a, bs) -> (,) a <$> S.toList bs) $ MS.toList l
    prop2 = all (\(b, a) -> M.member (a, b) m &&
                            maybe False (S.member b) (l MS.!? a)) $
                foldMap (\(b, as) -> (,) b <$> S.toList as) $ MS.toList r
    prop3 = all (\(a, b) -> maybe False (S.member b) (l MS.!? a) &&
                            maybe False (S.member a) (r MS.!? b)) $
                M.keys m

-- Internal utilities
---------------------

{-# INLINE mSetToSet #-}
mSetToSet :: Maybe (Set a) -> Set a
mSetToSet Nothing = S.empty
mSetToSet (Just set) = set

{-# INLINE ascListToMapSet #-}
ascListToMapSet :: (Ord a, Ord b) => [(a, b)] -> MS.Map a (Set b)
ascListToMapSet abs = MS.fromDistinctAscList sets
  where
    grouped = groupBy ((==) `on` fst) abs
    sets = (\xs -> (fst $ head xs, S.fromAscList (snd <$> xs))) <$> grouped

{-# INLINE setDelete' #-}
setDelete' :: Ord a => a -> Set a -> Maybe (Set a)
setDelete' x s = if S.null s' then Nothing else Just s'
  where s' = S.delete x s