{-# LANGUAGE TemplateHaskell, TypeFamilies, KindSignatures, FlexibleContexts, FlexibleInstances, UndecidableInstances, PatternGuards, MultiParamTypeClasses, TypeOperators #-}

module Data.TrieMap.MultiRec.UnionMap () where

import Data.TrieMap.MultiRec.Class
import Data.TrieMap.MultiRec.Eq
import Data.TrieMap.MultiRec.Base
import Data.TrieMap.Applicative
import Data.TrieMap.TrieKey
-- import Data.TrieMap.Rep
-- import Data.TrieMap.Rep.TH
import Data.TrieMap.MultiRec.TH
import qualified Data.TrieMap.Regular.Base as Reg

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

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

import Prelude hiding (foldr)

data UnionMap (phi :: * -> *) f g (r :: * -> *) ix a = HTrieMapT phi f r ix a :&: HTrieMapT phi g r ix a
type instance HTrieMapT phi (f :+: g) = UnionMap phi f g--HTrieMap phi (f r) :*: HTrieMap phi (g r)
type instance HTrieMap phi ((f :+: g) r) = HTrieMapT phi (f :+: g) r

-- type instance RepT (UnionMap phi f g r ix) = (Reg.:*:) (RepT (HTrieMapT phi f r ix)) (RepT (HTrieMapT phi g r ix))
-- type instance Rep (UnionMap phi f g r ix a) = RepT (UnionMap phi f g r ix) (Rep a)

-- -- $(genRepT [d|
--     instance (ReprT (HTrieMapT phi f r ix), ReprT (HTrieMapT phi g r ix)) => ReprT (UnionMap phi f g r ix) where
-- 	toRepT (m1 :&: m2) = (Reg.:*:) (toRepT m1) (toRepT m2)
-- 	fromRepT ((Reg.:*:) m1 m2) = fromRepT m1 :&: fromRepT m2
-- 	|])

$(inferH [d|
    instance (HTrieKeyT phi f (HTrieMapT phi f), HTrieKeyT phi g (HTrieMapT phi g)) => HTrieKeyT phi (f :+: g) (UnionMap phi f g) where
	emptyT = liftM2 (:&:) emptyT emptyT
	nullT pf (m1 :&: m2) = nullT pf m1 && nullT pf m2
	sizeT pf s (m1 :&: m2) = sizeT pf s m1 + sizeT pf s m2
	lookupT pf k (m1 :&: m2)
		| L k <- k	= lookupT pf k m1
		| R k <- k	= lookupT pf k m2
	lookupIxT pf s k (m1 :&: m2)
		| L k <- k	= case onKey L (lookupIxT pf s k m1) of
			(lb, x, ub) -> (lb, x, ub <|> ((onKeyA R . onIndexA (+ sizeT pf s m1)) <$> getMin pf s m2))
		| R k <- k	= case onIndex (sizeT pf s m1 +) (onKey R (lookupIxT pf s k m2)) of
			(lb, x, ub) -> ((onKeyA L <$> getMax pf s m1) <|> lb, x, ub)
			where	getMin pf s m = aboutT pf (\ k a -> return $ Asc 0 k a) m
				getMax pf s m = aboutT pf (\ k a -> return $ Asc (sizeT pf s m - s a) k a) m
	assocAtT pf s i (m1 :&: m2)
		| i < s1	= case onKey L (assocAtT pf s i m1) of
			(lb, x, ub) -> (lb, x, ub <|> ((onKeyA R . onIndexA (+ s1)) <$> getMin pf s m2))
		| otherwise	= case onKey R (onIndex (s1 +) (assocAtT pf s (i - s1) m2)) of
			(lb, x, ub) -> ((onKeyA L <$> getMax pf s m1) <|> lb, x, ub)
		where	getMin pf s m = aboutT pf (\ k a -> return $ Asc 0 k a) m
			getMax pf s m = aboutT pf (\ k a -> return $ Asc (sizeT pf s m - s a) k a) m
			s1 = sizeT pf s m1
{-	updateAtT pf s r f i (m1 :&: m2)
		| not r && i >= lastIx m1
			= m1 :&: updateAtT pf s r (\ i' -> f (i' + s1) . R) (i - s1) m2
		| i < s1
			= updateAtT pf s r (\ i' -> f i' . L) i m1 :&: m2
		| otherwise
			= m1 :&: updateAtT pf s r (\ i' -> f (i' + s1) . R) (i - s1) m2
		where	s1 = sizeT pf s m1
			lastIx m = case extractMaxT pf s (\ _ v -> (v, Just v)) m of
				Last (Just (v, _)) -> sizeT pf s m - s v
				_			-> sizeT pf s m-}
	alterT pf s f k (m1 :&: m2)
		| L k <- k	= alterT pf s f k m1 :&: m2
		| R k <- k	= m1 :&: alterT pf s f k m2
	traverseWithKeyT pf s f (m1 :&: m2)
		= (:&:) <$> traverseWithKeyT pf s (f . L) m1 <*> traverseWithKeyT pf s (f . R) m2
	foldWithKeyT pf f (m1 :&: m2) 
		= foldWithKeyT pf (f . L) m1 . foldWithKeyT pf (f . R) m2
	foldlWithKeyT pf f (m1 :&: m2)
		= foldlWithKeyT pf (f . R) m2 . foldlWithKeyT pf (f . L) m1
	mapEitherT pf s1 s2 f (m1 :&: m2) = case (mapEitherT pf s1 s2 (f . L) m1, mapEitherT pf s1 s2 (f . R) m2) of
		((m1L, m1R), (m2L, m2R)) -> (m1L :&: m2L, m1R :&: m2R)
	splitLookupT pf s f k0 (m1 :&: m2)
		| L k <- k0, (m1L, x, m1R) <- splitLookupT pf s f k m1
			= (m1L :&: emptyT pf, x, m1R :&: m2)
		| R k <- k0, (m2L, x, m2R) <- splitLookupT pf s f k m2
			= (m1 :&: m2L, x, emptyT pf :&: m2R)
	unionT pf s f (m11 :&: m12) (m21 :&: m22)
		= unionT pf s (f . L) m11 m21 :&: unionT pf s (f . R) m12 m22
	isectT pf s f (m11 :&: m12) (m21 :&: m22)
		= isectT pf s (f . L) m11 m21 :&: isectT pf s (f . R) m12 m22
	diffT pf s f (m11 :&: m12) (m21 :&: m22)
		= diffT pf s (f . L) m11 m21 :&: diffT pf s (f . R) m12 m22
	extractT pf s f (m1 :&: m2) = second (:&: m2) <$> extractT pf s (f . L) m1 <|>
		second (m1 :&:) <$> extractT pf s (f . R) m2
-- 	extractMinT pf s f (m1 :&: m2) = second (:&: m2) <$> extractMinT pf s (f . L) m1 <|>
-- 		second (m1 :&:) <$> extractMinT pf s (f . R) m2
-- 	extractMaxT pf s f (m1 :&: m2) = second (:&: m2) <$> extractMaxT pf s (f . L) m1 <|>
-- 		second (m1 :&:) <$> extractMaxT pf s (f . R) m2
-- 	alterMinT pf s f (m1 :&: m2)
-- 		| nullT pf m1	= m1 :&: alterMinT pf s (f . R) m2
-- 		| otherwise	= alterMinT pf s (f . L) m1 :&: m2
-- 	alterMaxT pf s f (m1 :&: m2)
-- 		| nullT pf m2	= alterMaxT pf s (f . L) m1 :&: m2
-- 		| otherwise	= m1 :&: alterMaxT pf s (f . R) m2
	isSubmapT pf (<=) (m11 :&: m12) (m21 :&: m22)
		= isSubmapT pf (<=) m11 m21 && isSubmapT pf (<=) m12 m22
	fromListT pf s f xs = case breakEither xs of
		(ys, zs) -> fromListT pf s (f . L) ys :&: fromListT pf s (f . R) zs
	fromAscListT pf s f xs = case breakEither xs of
		(ys, zs) -> fromAscListT pf s (f . L) ys :&: fromAscListT pf s (f . R) zs
	fromDistAscListT pf s xs = case breakEither xs of
		(ys, zs) -> fromDistAscListT pf s ys :&: fromDistAscListT pf s zs |])