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

module Data.TrieMap.Regular.UnitMap() where

import Data.TrieMap.Regular.Class
import Data.TrieMap.Regular.Base
import Data.TrieMap.TrieKey
-- import Data.TrieMap.Rep
-- import Data.TrieMap.Rep.Instances
-- import Data.TrieMap.Rep.TH
import Data.TrieMap.Applicative

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

import Data.Foldable
import Data.Maybe
import Data.Monoid
import Data.Traversable

import Prelude hiding (foldr, foldl)

newtype M k a = M (Maybe a)
type instance TrieMapT U0 = M
type instance TrieMap (U0 r) = M r

instance TrieKey (U0 r) (M r) where
	emptyM = M Nothing
	nullM (M a) = isNothing a
	sizeM s (M a) = maybe 0 s a
	lookupM _ (M a) = a
	lookupIxM s _ (M a) = (mzero, Asc 0 U0 <$> a, mzero)
	assocAtM s i (M a)
		| i < 0	= (mzero, mzero, Asc 0 U0 <$> First a)
		| i > maybe 0 s a
			= (Asc 0 U0 <$> Last a, mzero, mzero)
		| otherwise
			= (mzero, Asc 0 U0 <$> a, mzero)
-- 	updateAtM s r f i (M v) = case v of
-- 		Just a	| not r && i <= 0	-> M (v >>= f 0 U0)
-- 			| r && i >= 0		-> M (v >>= f 0 U0)
-- 		_ -> M v
	alterM _ f _ (M a) = M (f a)
	alterLookupM _ f _ (M a) = M <$> f a
	traverseWithKeyM _ f (M a) = M <$> traverse (f U0) a
	foldWithKeyM f (M a) z = foldr (f U0) z a
	foldlWithKeyM f (M a) z = foldl (f U0) z a
	mapEitherM _ _ f (M Nothing) = (M Nothing, M Nothing)
	mapEitherM _ _ f (M (Just a)) = (M *** M) (f U0 a)
	splitLookupM _ f _ (M a) = M `sides` maybe (Nothing, Nothing, Nothing) f a
	unionM _ f (M a) (M b) = M (unionMaybe (f U0) a b)
	isectM _ f (M a) (M b) = M (isectMaybe (f U0) a b)
	diffM _ f (M a) (M b) = M (diffMaybe (f U0) a b)
	extractM _ f (M a) = maybe empty (fmap M <.> f U0) a
-- 	extractMinM _ f (M a) = fmap (second M . f U0) (First a)
-- 	extractMaxM _ f (M a) = fmap (second M . f U0) (Last a)
-- 	alterMinM _ f (M a) = M (a >>= f U0)
-- 	alterMaxM = alterMinM
	isSubmapM (<=) (M a) (M b) = subMaybe (<=) a b
	fromListM _ f = M . foldr (\ (_, a) -> Just . maybe a (f U0 a)) Nothing
	fromDistAscListM _ = M . fmap snd . listToMaybe

instance TrieKeyT U0 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