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

module Data.TrieMap.TrieKey where

import Data.TrieMap.Applicative
import Data.TrieMap.Sized

import Control.Applicative

import Data.Monoid
import Data.List

type family TrieMap k :: (* -> *) -> * -> *

type family MapPF (m :: (* -> *) -> * -> *) ix :: (* -> *) -> *
-- data Fixer f

type EitherMap k a b c = k -> a -> (Maybe b, Maybe c)
type SplitMap a x = a -> (Maybe a, Maybe x, Maybe a)
type UnionFunc k a = k -> a -> a -> Maybe a
type IsectFunc k a b c = k -> a -> b -> Maybe c
type DiffFunc k a b = k -> a -> b -> Maybe a
type ExtractFunc k f a m = m -> f ((k, a), m)
type LEq a b = a -> b -> Bool
-- type Sized f = forall ix . f ix -> Int

-- toFixer :: a -> Fixer a
-- toFixer _ = undefined

class Ord k => TrieKey k m | m -> k where
	emptyM :: TrieMap k ~ m => m a ix
	nullM :: TrieMap k ~ m => m a ix -> Bool
	sizeM :: (TrieMap k ~ m) => Sized a -> m a ix -> Int
	lookupM :: TrieMap k ~ m => k -> m a ix -> Maybe (a ix)
	lookupIxM :: TrieMap k ~ m => Sized a -> k -> m a ix -> Maybe (Int, a ix)
	assocAtM :: TrieMap k ~ m => Sized a -> Int -> m a ix -> (Int, k, a ix)
	updateAtM :: TrieMap k ~ m => Sized a -> (Int -> k -> a ix -> Maybe (a ix)) -> Int -> m a ix -> m a ix
	alterM :: (TrieMap k ~ m) => Sized a -> (Maybe (a ix) -> Maybe (a ix)) -> k -> m a ix -> m a ix
	{-# SPECIALIZE traverseWithKeyM :: (k -> a ix -> Id (b ix)) -> m a ix -> Id (m b ix) #-}
	traverseWithKeyM :: (TrieMap k ~ m, Applicative f) => (forall ix . b ix -> Int) -> 
		(k -> a ix -> f (b ix)) -> m a ix -> f (m b ix)
	foldWithKeyM :: TrieMap k ~ m => (k -> a ix -> b -> b) -> m a ix -> b -> b
	foldlWithKeyM :: TrieMap k ~ m => (k -> b -> a ix -> b) -> m a ix -> b -> b
	mapEitherM :: (TrieMap k ~ m) => Sized b -> Sized c -> EitherMap k (a ix) (b ix) (c ix) -> m a ix -> (m b ix, m c ix)
	splitLookupM :: (TrieMap k ~ m) => Sized a -> SplitMap (a ix) x -> k -> m a ix -> (m a ix, Maybe x, m a ix)
	unionM :: (TrieMap k ~ m) => Sized a -> UnionFunc k (a ix) -> m a ix -> m a ix -> m a ix
	isectM :: (TrieMap k ~ m) => Sized c -> IsectFunc k (a ix) (b ix) (c ix) -> m a ix -> m b ix -> m c ix
	diffM :: (TrieMap k ~ m) => Sized a -> DiffFunc k (a ix) (b ix) -> m a ix -> m b ix -> m a ix
	extractMinM :: (TrieMap k ~ m) => Sized a -> ExtractFunc k First (a ix) (m a ix)
	extractMaxM :: (TrieMap k ~ m) => Sized a -> ExtractFunc k Last (a ix) (m a ix)
	alterMinM, alterMaxM :: (TrieMap k ~ m) => Sized a -> (k -> a ix -> Maybe (a ix)) -> m a ix -> m a ix
	isSubmapM :: TrieMap k ~ m => LEq (a ix) (b ix) -> LEq (m a ix) (m b ix)
	fromListM, fromAscListM :: (TrieMap k ~ m) => Sized a -> (k -> a ix -> a ix -> a ix) -> [(k, a ix)] -> m a ix
	fromDistAscListM :: (TrieMap k ~ m) => Sized a -> [(k, a ix)] -> m a ix
	
	sizeM s m = foldWithKeyM (\ _ a n -> s a + n) m 0
	fromListM s f = foldl' (flip (uncurry (insertWithKeyM s f))) emptyM
	fromAscListM = fromListM
	fromDistAscListM s = fromAscListM s (const const)
	updateAtM s f i m = case assocAtM s i m of
		(i', k, a)	-> alterM s (const (f i' k a)) k m

guardNullM :: (TrieKey k m, m ~ TrieMap k) => m a ix -> Maybe (m a ix)
guardNullM m
	| nullM m	= Nothing
	| otherwise	= Just m

sides :: (a -> c) -> (a, b, a) -> (c, b, c)
sides f (l, x, r) = (f l, x, f r)

mapMaybeM :: (TrieKey k m, m ~ TrieMap k) => Sized b -> (k -> a ix -> Maybe (b ix)) -> m a ix -> m b ix
mapMaybeM s f = snd . mapEitherM elemSize s (((,) (Nothing :: Maybe (Elem ix))) .: f)

mapWithKeyM :: (TrieKey k m, m ~ TrieMap k) => Sized b -> (k -> a ix -> b ix) -> m a ix -> m b ix
mapWithKeyM s f  = unId . traverseWithKeyM s (Id .: f)

mapM :: (TrieKey k m, m ~ TrieMap k) => Sized b -> (a ix -> b ix) -> m a ix -> m b ix
mapM s = mapWithKeyM s . const

assocsM :: (TrieKey k m, m ~ TrieMap k) => m a ix -> [(k, a ix)]
assocsM m = foldWithKeyM (\ k a xs -> (k, a):xs) m []

insertM :: (TrieKey k m, m ~ TrieMap k) => Sized a -> k -> a ix -> m a ix -> m a ix
insertM s = insertWithKeyM s (const const)

insertWithKeyM :: (TrieKey k m, m ~ TrieMap k) => Sized a -> (k -> a ix -> a ix -> a ix) -> k -> a ix -> m a ix -> m a ix
insertWithKeyM s f k a = alterM s f' k where
	f' = Just . maybe a (f k a)

singletonM :: (TrieKey k m, m ~ TrieMap k) => Sized a -> k -> a ix -> m a ix
singletonM s k a = insertM s k a emptyM

fromListM' :: (TrieKey k m, m ~ TrieMap k) => Sized a -> [(k, a ix)] -> m a ix
fromListM' s = fromListM s (const const) --xs = foldr (uncurry insertM) emptyM xs

unionMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a
unionMaybe _ Nothing y = y
unionMaybe _ x Nothing = x
unionMaybe f (Just x) (Just y) = f x y

isectMaybe :: (a -> b -> Maybe c) -> Maybe a -> Maybe b -> Maybe c
isectMaybe f (Just x) (Just y) = f x y
isectMaybe _ _ _ = Nothing

diffMaybe :: (a -> b -> Maybe a) -> Maybe a -> Maybe b -> Maybe a
diffMaybe f Nothing = const Nothing
diffMaybe f (Just x) = maybe (Just x) (f x)

subMaybe :: (a -> b -> Bool) -> Maybe a -> Maybe b -> Bool
subMaybe _ Nothing _ = True
subMaybe (<=) (Just a) (Just b) = a <= b
subMaybe _ _ _ = False