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

module Data.TrieMap.Regular.IdMap where

import Data.TrieMap.TrieKey
import Data.TrieMap.Regular.Base
import Data.TrieMap.Regular.Class

import Control.Applicative
import Control.Arrow
import Control.Monad

newtype IMap k a ix = IMap (TrieMap k a ix)
type instance TrieMapT I0 = IMap
type instance TrieMap (I0 k) = IMap k

instance TrieKeyT I0 IMap where
	emptyT = IMap emptyM
	nullT (IMap m) = nullM m
	sizeT s (IMap m) = sizeM s m
	lookupT (I0 k) (IMap m) = lookupM k m
	lookupIxT s (I0 k) (IMap m) = lookupIxM s k m
	assocAtT s i (IMap m) = case assocAtM s i m of
		(i', k, a) -> (i', I0 k, a)
	updateAtT s f i (IMap m) = IMap (updateAtM s (\ i -> f i . I0) i m)
	alterT s f (I0 k) (IMap m) = IMap (alterM s f k m)
	traverseWithKeyT s f (IMap m) = IMap <$> traverseWithKeyM s (f . I0) m
	foldWithKeyT f (IMap m) = foldWithKeyM (f . I0) m
	foldlWithKeyT f (IMap m) = foldlWithKeyM (f . I0) m
	mapEitherT s1 s2 f (IMap m) = (IMap *** IMap) (mapEitherM s1 s2 (f . I0) m)
	splitLookupT s f (I0 k) (IMap m) = IMap `sides` splitLookupM s f k m
	unionT s f (IMap m1) (IMap m2) = IMap (unionM s (f . I0) m1 m2)
	isectT s f (IMap m1) (IMap m2) = IMap (isectM s (f . I0) m1 m2)
	diffT s f (IMap m1) (IMap m2) = IMap (diffM s (f . I0) m1 m2)
	extractMinT s (IMap m) = (first I0 *** IMap) `liftM` extractMinM s m
	extractMaxT s (IMap m) = (first I0 *** IMap) `liftM` extractMaxM s m
	alterMinT s f (IMap m) = IMap (alterMinM s (f . I0) m)
	alterMaxT s f (IMap m) = IMap (alterMaxM s (f . I0) m)
	isSubmapT (<=) (IMap m1) (IMap m2) = isSubmapM (<=) m1 m2
	fromListT s f xs = IMap (fromListM s (f . I0) [(k, a) | (I0 k, a) <- xs])
	fromAscListT s f xs = IMap (fromAscListM s (f . I0) [(k, a) | (I0 k, a) <- xs])
	fromDistAscListT s xs = IMap (fromDistAscListM s [(k, a) | (I0 k, a) <- xs])

instance TrieKey k (TrieMap k) => TrieKey (I0 k) (IMap k) where
	emptyM = emptyT
	nullM = nullT
	sizeM = sizeT
	lookupM = lookupT
	lookupIxM = lookupIxT
	assocAtM = assocAtT
	updateAtM = updateAtT
	alterM = alterT
	traverseWithKeyM = traverseWithKeyT
	foldWithKeyM = foldWithKeyT
	foldlWithKeyM = foldlWithKeyT
	mapEitherM = mapEitherT
	splitLookupM = splitLookupT
	unionM = unionT
	isectM = isectT
	diffM = diffT
	extractMinM = extractMinT
	extractMaxM = extractMaxT
	alterMinM = alterMinT
	alterMaxM = alterMaxT
	isSubmapM = isSubmapT
	fromListM = fromListT
	fromAscListM = fromAscListT
	fromDistAscListM = fromDistAscListT