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

module Data.TrieMap.MultiRec.IMap where

import Data.TrieMap.MultiRec.Class
import Data.TrieMap.MultiRec.Sized
import Data.TrieMap.TrieKey

import Control.Applicative
import Control.Arrow

import Generics.MultiRec

newtype IMap phi xi r a ix = IMap (HTrieMap phi r (I ix a) xi)
type instance HTrieMapT phi (I xi) = IMap phi xi
type instance HTrieMap phi (I xi r) = HTrieMapT phi (I xi) r

combineI :: (I xi r ix -> a ix -> b ix -> Maybe (c ix)) -> r xi -> I ix a xi -> I ix b xi -> Maybe (I ix c xi)
combineI f k (I a) (I b) = I <$> f (I k) a b

mapI :: Functor f => (I xi r ix -> a ix -> f (b ix)) -> r xi -> I ix a xi -> f (I ix b xi)
mapI f k (I a) = I <$> f (I k) a

sizeI :: HSized phi r -> HSized phi (I xi r)
sizeI s (I x) = s x

instance El phi xi => HTrieKeyT phi (I xi) (IMap phi xi) where
	emptyT _ = IMap (emptyH proof)
	nullT _ (IMap m) = nullH proof m
	sizeT s (IMap m) = sizeH (sizeI s) m
	lookupT _ (I k) (IMap m) = unI <$> lookupH proof k m
	lookupIxT _ s (I k) (IMap m) = fmap unI <$> lookupIxH proof (sizeI s) k m
	assocAtT _ s i (IMap m) = case assocAtH proof (sizeI s) i m of
		(i, k, I a) -> (i, I k, a)
	updateAtT _ s f i (IMap m) = IMap (updateAtH proof (sizeI s) (\ i' k (I a) -> I <$> f i' (I k) a) i m)
	alterT _ s f (I k) (IMap m) = IMap (alterH proof (sizeI s) f' k m) where
		f' = fmap I . f . fmap unI
	traverseWithKeyT _ s f (IMap m) = IMap <$> traverseWithKeyH proof (sizeI s) (mapI f) m
	foldWithKeyT _ f (IMap m) = foldWithKeyH proof (\ k (I a) -> f (I k) a) m
	foldlWithKeyT _ f (IMap m) = foldlWithKeyH proof (\ k z (I a) -> f (I k) z a) m
	mapEitherT _ s1 s2 f (IMap m) = (IMap *** IMap) (mapEitherH proof (sizeI s1) (sizeI s2) f' m) where
		f' k (I a) = (fmap I *** fmap I) (f (I k) a)
	splitLookupT pf s f (I k) (IMap m) = IMap `sides` splitLookupH proof (sizeI s) f' k m
		where	f' = sides (I <$>) . f . unI
	unionT pf s f (IMap m1) (IMap m2) = IMap (unionH proof (sizeI s) (combineI f) m1 m2)
	isectT pf s f (IMap m1) (IMap m2) = IMap (isectH proof (sizeI s) (combineI f) m1 m2)
	diffT pf s f (IMap m1) (IMap m2) = IMap (diffH proof (sizeI s) (combineI f) m1 m2)
	extractMinT pf s (IMap m) = do
		((k, I a), m') <- extractMinH proof (sizeI s) m
		return ((I k, a), IMap m')
	extractMaxT pf s (IMap m) = do
		((k, I a), m') <- extractMaxH proof (sizeI s) m
		return ((I k, a), IMap m')
	alterMinT pf s f (IMap m) = IMap (alterMinH proof (sizeI s) (mapI f) m)
	alterMaxT pf s f (IMap m) = IMap (alterMaxH proof (sizeI s) (mapI f) m)
	isSubmapT pf (<=) (IMap m1) (IMap m2) = isSubmapH proof (<<=) m1 m2 where
		I a <<= I b = a <= b
	fromListT _ s f xs = IMap (fromListH proof (sizeI s) (\ k (I a) (I b) -> I (f (I k) a b)) [(k, I a) | (I k, a) <- xs])
	fromAscListT _ s f xs = IMap (fromAscListH proof (sizeI s) (\ k (I a) (I b) -> I (f (I k) a b)) [(k, I a) | (I k, a) <- xs])
	fromDistAscListT _ s xs = IMap (fromDistAscListH proof (sizeI s) [(k, I a) | (I k, a) <- xs])

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