{-# LANGUAGE Rank2Types, FlexibleContexts, TypeFamilies, MultiParamTypeClasses, FunctionalDependencies #-}

module Data.TrieMap.Regular.Class where

import Data.TrieMap.Sized
import Data.TrieMap.Applicative
import Data.TrieMap.TrieKey
import Data.TrieMap.Regular.Eq
import Data.TrieMap.Regular.Ord

import Data.Monoid

import Control.Applicative

type family TrieMapT (f :: * -> *) :: * -> (* -> *) -> * -> *

class OrdT f => TrieKeyT (f :: * -> *) (m :: * -> (* -> *) -> * -> *) | m -> f, f -> m where
	emptyT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => m k a ix
	nullT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => m k a ix -> Bool
	sizeT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> m k a ix -> Int
	lookupT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => f k -> m k a ix -> Maybe (a ix)
	lookupIxT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> f k -> m k a ix -> Maybe (Int, a ix)
	assocAtT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> Int -> m k a ix -> (Int, f k, a ix)
	updateAtT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> (Int -> f k -> a ix -> Maybe (a ix)) -> Int -> m k a ix -> m k a ix
	alterT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> (Maybe (a ix) -> Maybe (a ix)) -> f k -> m k a ix -> m k a ix
	traverseWithKeyT :: (TrieMapT f ~ m, TrieKey k (TrieMap k), Applicative t) => 
		Sized b -> (f k -> a ix -> t (b ix)) -> m k a ix -> t (m k b ix)
	foldWithKeyT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => 
		(f k -> a ix -> b -> b) -> m k a ix -> b -> b
	foldlWithKeyT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) =>
		(f k -> b -> a ix -> b) -> m k a ix -> b -> b
	mapEitherT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => 
		Sized b -> Sized c -> EitherMap (f k) (a ix) (b ix) (c ix) -> m k a ix -> (m k b ix, m k c ix)
	splitLookupT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> SplitMap (a ix) x -> f k ->
		m k a ix -> (m k a ix, Maybe x, m k a ix)
	unionT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> UnionFunc (f k) (a ix) ->
		m k a ix -> m k a ix -> m k a ix
	isectT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized c -> IsectFunc (f k) (a ix) (b ix) (c ix) ->
		m k a ix -> m k b ix -> m k c ix
	diffT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> DiffFunc (f k) (a ix) (b ix) ->
		m k a ix -> m k b ix -> m k a ix
	extractMinT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> m k a ix -> First ((f k, a ix), m k a ix)
	extractMaxT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> m k a ix -> Last ((f k, a ix), m k a ix)
	alterMinT, alterMaxT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> (f k -> a ix -> Maybe (a ix)) ->
		m k a ix -> m k a ix
	isSubmapT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => LEq (a ix) (b ix) -> LEq (m k a ix) (m k b ix)
	fromListT, fromAscListT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> (f k -> a ix -> a ix -> a ix) ->
		[(f k, a ix)] -> m k a ix
	fromDistAscListT :: (TrieMapT f ~ m, TrieKey k (TrieMap k)) => Sized a -> [(f k, a ix)] -> m k a ix
	fromListT s f = foldr (\ (k, a) -> alterT s (Just . maybe a (f k a)) k) emptyT
	fromAscListT = fromListT
	fromDistAscListT s = fromAscListT s (const const)
	updateAtT s f i m = case assocAtT s i m of
		(i, k, a) -> alterT s (const (f i k a)) k m

guardNullT :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => TrieMapT f k a ix -> Maybe (TrieMapT f k a ix)
guardNullT m
	| nullT m	= Nothing
	| otherwise	= Just m

assocsT :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => TrieMapT f k a ix -> [(f k, a ix)]
assocsT m = foldWithKeyT (\ k a -> ((k, a):)) m []

singletonT :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => Sized a -> f k -> a ix -> TrieMapT f k a ix
singletonT s k a = alterT s (const (Just a)) k emptyT

mapWithKeyT :: (TrieKeyT f (TrieMapT f), TrieKey k (TrieMap k)) => 
	Sized b -> (f k -> a ix -> b ix) -> TrieMapT f k a ix -> TrieMapT f k b ix
mapWithKeyT s f m = unId (traverseWithKeyT s (Id .: f) m)