{-# 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 :: * -> *) ix = KMap (m a ix)
type instance TrieMapT (K0 a) = KMap (TrieMap a)
type instance TrieMap (K0 a r) = TrieMapT (K0 a) r
type instance PF (KMap m k a ix) = PF (m a ix)

instance (Regular (m a ix), Functor (PF (m a ix))) => Regular (KMap m k a ix) where
	from (KMap m) = fmap KMap (from m)
	to = KMap . to . fmap (\ (KMap m) -> m)

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) = lookupIxM s k m
	assocAtM s i (KMap m) = case assocAtM s i m of
		(i', k, a) -> (i', K0 k, a)
	updateAtM s f i (KMap m) = KMap (updateAtM s (\ i -> f i . K0) i m)
	alterM s f (K0 k) (KMap m) = KMap (alterM 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)
	extractMinM s (KMap m) = (first K0 *** KMap) `liftM` extractMinM s m
	extractMaxM s (KMap m) = (first K0 *** KMap) `liftM` extractMaxM s 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
	traverseWithKeyT = traverseWithKeyM
	foldWithKeyT = foldWithKeyM
	foldlWithKeyT = foldlWithKeyM
	mapEitherT = mapEitherM
	splitLookupT = splitLookupM
	unionT = unionM
	isectT = isectM
	diffT = diffM
	extractMinT = extractMinM
	extractMaxT = extractMaxM
	alterMinT = alterMinM
	alterMaxT = alterMaxM
	isSubmapT = isSubmapM
	fromListT = fromListM
	fromAscListT = fromAscListM
	fromDistAscListT = fromDistAscListM