{-# 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 = IMap (TrieMap k a)
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) = onKey I0 (lookupIxM s k m)
	assocAtT s i (IMap m) = onKey I0 (assocAtM s i m)
-- 	updateAtT s r f i (IMap m) = IMap (updateAtM s r (\ i -> f i . I0) i m)
	alterT s f (I0 k) (IMap m) = IMap (alterM s f k m)
	alterLookupT s f (I0 k) (IMap m) = IMap <$> alterLookupM 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)
	extractT s f (IMap m) = fmap IMap <$> extractM s (f . I0) m
-- 	extractMinT s f (IMap m) = second IMap <$> extractMinM s (f . I0) m
-- 	extractMaxT s f (IMap m) = second IMap <$> extractMaxM s (f . I0) 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
	alterLookupM = alterLookupT
	traverseWithKeyM = traverseWithKeyT
	foldWithKeyM = foldWithKeyT
	foldlWithKeyM = foldlWithKeyT
	mapEitherM = mapEitherT
	splitLookupM = splitLookupT
	unionM = unionT
	isectM = isectT
	diffM = diffT
	extractM = extractT
-- 	extractMinM = extractMinT
-- 	extractMaxM = extractMaxT
-- 	alterMinM = alterMinT 
-- 	alterMaxM = alterMaxT
	isSubmapM = isSubmapT
	fromListM = fromListT
	fromAscListM = fromAscListT
	fromDistAscListM = fromDistAscListT