{-# LANGUAGE MultiParamTypeClasses, ViewPatterns, ScopedTypeVariables #-}
module Data.Containers(
  -- * The basic data class
  DataMap(..),Indexed(..),OrderedMap(..),
  
  member,delete,touch,insert,singleton,fromList,
  _set,_map,cached,

  -- * Map instances
  -- ** Sets and maps
  Set,Map,
  
  -- ** Bimaps
  Bimap(..),toMap,keysSet,

  -- ** Relations
  Relation(..),domains,ranges,related,link
  )
  where

import Algebra
import qualified Data.Set as S
import qualified Data.Map as M
import Data.Map (Map)
import Data.Set (Set)
import Control.Concurrent.MVar

class Monoid m => DataMap m k a | m -> k a where
  at :: k -> Lens' m (Maybe a)
class Indexed f i | f -> i where
  keyed :: Iso (f (i,a)) (f (i,b)) (f a) (f b) 
class OrderedMap m k a m' k' a' | m -> k a, m' -> k' a' where
  ascList :: Iso [(k,a)] [(k',a')] m m'

_set :: Set a -> Set a
_set = id
_map :: Map a b -> Map a b
_map = id

member :: DataMap m k Void => k -> Lens' m Bool
member k = at k.from _maybe
delete :: DataMap m k a => k -> m -> m
delete k = at k %- Nothing
insert :: DataMap m k a => k -> a -> m -> m
insert k a = at k %- Just a
touch :: (Monoid a, DataMap m k a) => k -> m -> m
touch k = insert k zero
singleton :: DataMap m k a => k -> a -> m
singleton = map2 ($zero) insert
fromList :: DataMap m k a => [(k,a)] -> m
fromList l = compose (uncurry insert<$>l) zero

instance Ord a => DataMap (Set a) a Void where
  at k = lens (S.member k) (flip (bool (S.insert k) (S.delete k)))._maybe
instance Eq b => OrderedMap (Set a) a Void (Set b) b Void where
  ascList = iso S.toAscList S.fromAscList . mapping (_iso.commuted)
instance Ord k => DataMap (Map k a) k a where
  at k = lens (M.lookup k) (\m a -> M.alter (const a) k m)
instance Eq k' => OrderedMap (Map k a) k a (Map k' a') k' a' where 
  ascList = iso M.toAscList M.fromAscList
  
instance Ord a => Semigroup (Set a) where (+) = S.union
instance Ord a => Monoid (Set a) where zero = S.empty
instance Ord a => Disjonctive (Set a) where (-) = S.difference
instance Ord a => Semiring (Set a) where (*) = S.intersection
instance Functor Set where map = S.mapMonotonic
instance Foldable Set where fold = S.foldr (+) zero

instance Ord k => Semigroup (Map k a) where (+) = M.union
instance Ord k => Monoid (Map k a) where zero = M.empty
instance Ord k => Disjonctive (Map k a) where (-) = M.difference
instance (Ord k,Semigroup a) => Semiring (Map k a) where (*) = M.intersectionWith (+)
instance Functor (Map k) where map = M.map
instance Foldable (Map k) where fold = M.foldr (+) zero
instance Eq k => Traversable (Map k) where sequence = (ascList._Compose) sequence
instance Indexed (Map k) k where keyed = iso (M.mapWithKey (,)) (map snd)

-- |An invertible map
newtype Bimap a b = Bimap (Map a b,Map b a)
                  deriving (Show,Semigroup,Monoid,Disjonctive,Semiring)
instance Commutative Bimap where
  commute (Bimap (b,a)) = Bimap (a,b)

instance (Ord a,Ord b) => DataMap (Bimap a b) a b where
  at a = lens lookup setAt
    where lookup ma = toMap ma^.at a
          setAt (Bimap (ma,mb)) b' = Bimap (
            maybe id delete (b' >>= \b'' -> mb^.at b'') ma & at a %- b',
            mb & maybe id delete b >>> maybe id (flip insert a) b')
            where b = ma^.at a 
instance (Ord b,Ord a) => DataMap (Flip Bimap b a) b a where
  at k = from (commuted._Flip).at k
instance (Ord a,Ord b,Ord c,Ord d) => OrderedMap (Bimap a b) a b (Bimap c d) c d where
  ascList = iso (toMap >>> \m -> m^.ascList) (\l -> Bimap (l^..ascList,l^..ascList.mapping commuted))
toMap :: Bimap a b -> Map a b
toMap (Bimap (a,_)) = a

keysSet :: (Eq a,Eq b) => Iso (Set a) (Set b) (Map a Void) (Map b Void)
keysSet = ascList.from ascList

--- |The Relation type
newtype Relation a b = Relation (Map a (Set b),Map b (Set a))
                     deriving (Show,Semigroup,Monoid,Eq,Ord)
_Relation :: Iso (Relation a b) (Relation c d) (Map a (Set b),Map b (Set a)) (Map c (Set d),Map d (Set c))
_Relation = iso Relation (\(Relation r) -> r)
instance Commutative Relation where
  commute (Relation (a,b)) = Relation (b,a)

-- |Define a Relation from its ranges. O(1) <-> O(1,n*ln(n)) 
ranges :: (Ord c,Ord d) => Iso (Map a (Set b)) (Map c (Set d)) (Relation a b) (Relation c d)
ranges = iso (\(Relation (rs,_)) -> rs) fromRanges
  where fromRanges rs = Relation (rs,compose (rs^.keyed <&> \ (a,bs) -> compose $ bs <&> \b ->
                                              at b%~Just . touch a . fold) zero)
-- |Define a Relation from its domain (uses the Commutative instance)
domains :: (Ord c,Ord d) => Iso (Map b (Set a)) (Map d (Set c)) (Relation a b) (Relation c d)
domains = commuted.ranges

instance (Ord k,Ord a) => DataMap (Relation k a) k (Set a) where
  at a = lens (\(Relation (rs,_)) -> rs^.at a) setRan
    where setRan (Relation (rs,ds)) (fold -> ran) = Relation (
            rs & at a %- if empty ran then Nothing else Just ran,
            adjust ds)
            where oldRan = fold $ rs^.at a
                  adjust = compose ((oldRan-ran) <&> \b -> at b.traverse.member a %- False)
                           >>> compose ((ran-oldRan) <&> \b -> at b %~ Just . touch a . fold)

may :: (Monoid (f b),Foldable f) => Iso (Maybe (f a)) (Maybe (f b)) (f a) (f b)
may = iso (\f -> if empty f then Nothing else Just f) (maybe zero id)

related :: (Ord a,Ord b) => a -> Lens' (Relation a b) (Set b)
related a = at a.from may

link :: (Ord a,Ord b) => a -> b -> Lens' (Relation a b) Bool
link a b = related a.member b

cached :: forall a b. Ord a => (a -> b) -> a -> b
cached f = \a -> g a^.thunk
  where g a = do
          m <- vm `seq` takeMVar vm
          case m^.at a of
            Just b -> putMVar vm m >> return b
            Nothing -> let b = f a in putMVar vm (insert a b m) >> return b
        vm = newMVar (zero :: Map a b)^.thunk