{-# LANGUAGE KindSignatures, TypeFamilies, MultiParamTypeClasses, FlexibleInstances #-}

module Data.TrieMap.MultiRec.UnitMap where

import Data.TrieMap.MultiRec.Class
import Data.TrieMap.MultiRec.Eq
import Data.TrieMap.Applicative
import Data.TrieMap.TrieKey

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

import Data.Maybe
import Data.Monoid
import Data.Foldable
import Data.Traversable
import Generics.MultiRec

import Prelude hiding (foldr, foldl)

newtype UMap (phi :: * -> *) (r :: * -> *) a ix = UMap (Maybe (a ix))
type instance HTrieMapT phi U = UMap phi
type instance HTrieMap phi (U r) = UMap phi r

instance HTrieKeyT phi U (UMap phi) where
	emptyT = emptyH
	nullT = nullH
	sizeT = sizeH
	lookupT = lookupH
	lookupIxT = lookupIxH
	assocAtT = assocAtH
	updateAtT = updateAtH
	alterT = alterH
	traverseWithKeyT = traverseWithKeyH
	foldWithKeyT = foldWithKeyH
	foldlWithKeyT = foldlWithKeyH
	mapEitherT = mapEitherH
	splitLookupT = splitLookupH
	unionT = unionH
	isectT = isectH
	diffT = diffH
	extractMinT = extractMinH
	extractMaxT = extractMaxH
	alterMinT = alterMinH
	alterMaxT = alterMaxH
	isSubmapT = isSubmapH
	fromListT = fromListH
	fromAscListT = fromAscListH
	fromDistAscListT = fromDistAscListH

instance HTrieKey phi (U r) (UMap phi r) where
	emptyH _ = UMap Nothing
	nullH _ (UMap m) = isNothing m
	sizeH s (UMap m) = maybe 0 s m
	lookupH _ _ (UMap m) = m
	lookupIxH _ _ _ (UMap m) = fmap ((,) 0) m
	assocAtH _ _ _ (UMap (Just a)) = (0, U, a)
	updateAtH _ _ f _ (UMap m) = UMap (m >>= f 0 U)
	alterH _ _ f _ (UMap m) = UMap (f m)
	traverseWithKeyH _ _ f (UMap m) = UMap <$> traverse (f U) m
	foldWithKeyH _ f (UMap m) z = foldr (f U) z m
	foldlWithKeyH _ f (UMap m) z = foldl (f U) z m
	mapEitherH _ _ _ f (UMap m) = (UMap *** UMap) (maybe (Nothing, Nothing) (f U) m)
	splitLookupH _ _ f _ (UMap m) = UMap `sides` maybe (Nothing, Nothing, Nothing) f m
	unionH _ _ f (UMap m1) (UMap m2) = UMap (unionMaybe (f U) m1 m2)
	isectH _ _ f (UMap m1) (UMap m2) = UMap (isectMaybe (f U) m1 m2)
	diffH _ _ f (UMap m1) (UMap m2) = UMap (diffMaybe (f U) m1 m2)
	extractMinH _ _ (UMap m) = do	v <- First m
					return ((U, v), UMap Nothing)
	extractMaxH _ _ (UMap m) = do	v <- Last m
					return ((U, v), UMap Nothing)
	alterMinH _ _ f (UMap m) = UMap (m >>= f U)
	alterMaxH = alterMinH
	isSubmapH _ _ (UMap Nothing) _ = True
	isSubmapH _ (<=) (UMap m1) (UMap m2) = subMaybe (<=) m1 m2
	fromListH _ _ f xs = UMap (foldr (\ (_, a) -> Just . maybe a (f U a)) Nothing xs)
	fromAscListH = fromListH
	fromDistAscListH _ _ xs = UMap (fmap snd (listToMaybe xs))