{-# 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

type instance RepT (M k) = RepT Maybe
type instance Rep (M k a) = RepT Maybe (Rep a)

$(genRepT [d|
   instance ReprT (M k) where
	toRepT (M a) = toRepT a
	fromRepT = M . fromRepT |])

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)
	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 (second 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
	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