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

module Data.TrieMap.Regular.RegMap where

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

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

newtype RegMap k m (a :: * -> *) ix = RegMap (m (Reg k) a ix)

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