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

module Data.TrieMap.Regular.ConstMap() where

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

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

-- import Data.Monoid

newtype KMap m k a = KMap (m a)
type instance TrieMapT (K0 a) = KMap (TrieMap a)
type instance TrieMap (K0 a r) = TrieMapT (K0 a) r

instance (TrieKey k m, m ~ TrieMap k) => TrieKey (K0 k r) (KMap m r) where
	emptyM = KMap emptyM
	nullM (KMap m) = nullM m
	sizeM s (KMap m) = sizeM s m
	lookupM (K0 k) (KMap m) = lookupM k m
	lookupIxM s (K0 k) (KMap m) = onKey K0 (lookupIxM s k m)
	assocAtM s i (KMap m) = onKey K0 (assocAtM s i m)
-- 	updateAtM s r f i (KMap m) = KMap (updateAtM s r (\ i -> f i . K0) i m)
	alterM s f (K0 k) (KMap m) = KMap (alterM s f k m)
	alterLookupM s f (K0 k) (KMap m) = KMap <$> alterLookupM s f k m
	traverseWithKeyM s f (KMap m) = KMap <$> traverseWithKeyM s (f . K0) m
	foldWithKeyM f (KMap m) = foldWithKeyM (f . K0) m
	foldlWithKeyM f (KMap m) = foldlWithKeyM (f . K0) m
	mapEitherM s1 s2 f (KMap m) = (KMap *** KMap) (mapEitherM s1 s2 (f . K0) m)
	splitLookupM s f (K0 k) (KMap m) = KMap `sides` splitLookupM s f k m
	unionM s f (KMap m1) (KMap m2) = KMap (unionM s (f . K0) m1 m2)
	isectM s f (KMap m1) (KMap m2) = KMap (isectM s (f . K0) m1 m2)
	diffM s f (KMap m1) (KMap m2) = KMap (diffM s (f . K0) m1 m2)
	extractM s f (KMap m) = fmap KMap <$> extractM s (f . K0) m
-- 	extractMinM s f (KMap m) = second KMap <$> extractMinM s (f . K0) m
-- 	extractMaxM s f (KMap m) = second KMap <$> extractMaxM s (f . K0) m
-- 	alterMinM s f (KMap m) = KMap (alterMinM s (f . K0) m) 
-- 	alterMaxM s f (KMap m) = KMap (alterMaxM s (f . K0) m)
	isSubmapM (<=) (KMap m1) (KMap m2) = isSubmapM (<=) m1 m2
	fromListM s f xs = KMap (fromListM s (f . K0) [(k, a) | (K0 k, a) <- xs])
	fromAscListM s f xs = KMap (fromAscListM s (f . K0) [(k, a) | (K0 k, a) <- xs])
	fromDistAscListM s xs = KMap (fromDistAscListM s [(k, a) | (K0 k, a) <- xs])

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