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

module Data.TrieMap.MultiRec.TagMap () where

import Data.TrieMap.MultiRec.Class
import Data.TrieMap.MultiRec.Eq
import Data.TrieMap.MultiRec.Sized
import Data.TrieMap.CPair
-- import Data.TrieMap.MultiRec.TH
-- import Data.TrieMap.Applicative
import Data.TrieMap.TrieKey
-- import Data.TrieMap.Rep

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

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

data TagF a ix xi where
	TagF :: a -> TagF a ix ix

unTagF :: TagF a ix xi -> a
unTagF (TagF x) = x

newtype TagMap (phi :: * -> *) f ix (r :: * -> *) xi a = TagMap (HTrieMapT phi f r xi (TagF a ix xi))
type instance HTrieMapT phi (f :>: ix) = TagMap phi f ix
-- type instance HTrieMap phi ((f :>: ix) r) = HTrieMapT phi (f :>: ix) r

-- type instance RepT (TagMap phi f ix r xi) = RepT (HTrieMapT phi f r xi)
-- type instance Rep (TagMap phi f ix r xi a) = RepT (HTrieMapT phi f r xi) (Rep a)

-- instance (ReprT (HTrieMapT phi f r xi), ix ~ xi) => ReprT (TagMap phi f ix r xi) where
-- 	toRepT (TagMap m) = fmap unTagF (toRepT m)
-- 	fromRepT = TagMap . fromRepT . fmap TagF
-- 
-- instance (ReprT (HTrieMapT phi f r xi), ix ~ xi, Repr a) => Repr (TagMap phi f ix r xi a) where
-- 	toRep (TagMap m) = fmap (toRep . unTagF) (toRepT m)
-- 	fromRep = TagMap . fromRepT . fmap (TagF . fromRep)

combineTag :: IsectFunc ((f :>: ix) r xi) (a) (b) (c) ->
	IsectFunc (f r xi) (TagF a ix xi) (TagF b ix xi) (TagF c ix xi)
combineTag f k (TagF a) (TagF b) = TagF <$> f (Tag k) a b

mapTag :: Functor t => ((f :>: ix) r xi -> a -> t (b)) -> f r xi -> TagF a ix xi -> t (TagF b ix xi)
mapTag f k (TagF a) = TagF <$> f (Tag k) a

sizeTag :: HSized phi a -> HSized phi (TagF a ix xi)
sizeTag s (TagF x) = s x

restructure :: HTrieKeyT phi f (HTrieMapT phi f) =>
	((f r ix, TagF a xi ix), HTrieMapT phi f r ix (TagF a xi ix)) -> (((f :>: xi) r ix, a), TagMap phi f xi r ix a)
restructure ((k, TagF a), m) = ((Tag k, a), TagMap m)

restructure' :: Applicative t => ((f :>: xi) r ix -> a -> t (CPair x (Maybe a))) -> f r ix -> TagF a xi ix -> t (CPair x (Maybe (TagF a xi ix)))
restructure' f k (TagF a) = fmap (fmap TagF) <$> f (Tag k) a

retag :: (f r ix, TagF a xi ix) -> ((f :>: xi) r ix, a)
retag (k, TagF a) = (Tag k, a)

instance (HTrieKeyT phi f (HTrieMapT phi f)) => HTrieKeyT phi (f :>: ix) (TagMap phi m ix) where
	emptyH = TagMap . emptyH
	nullH pf (TagMap m) = nullH pf m
	sizeH pf s (TagMap m) = sizeH pf (sizeTag s) m
	lookupH pf (Tag k) (TagMap m) = unTagF <$> lookupH pf k m
	lookupIxH pf s (Tag k) (TagMap m) = onValue retag (lookupIxH pf (sizeTag s) k m)
	assocAtH pf s i (TagMap m) = onValue retag (assocAtH pf (sizeTag s) i m) 
-- 	updateAtT pf s r f i (TagMap m) = TagMap (updateAtT pf (sizeTag s) r (f' f) i m) where
-- 		f' :: (Int -> (f :>: xi) r ix -> a -> Maybe (a)) -> Int -> f r ix -> TagF a xi ix -> Maybe (TagF a xi ix)
-- 		f' f i k (TagF a) = TagF <$> f i (Tag k) a
	alterH pf s f (Tag k) (TagMap m) = TagMap (alterH pf (sizeTag s) (fmap TagF . f . fmap unTagF) k m)
	alterLookupH pf s f (Tag k) (TagMap m) = TagMap <$> alterLookupH pf (sizeTag s) (fmap (fmap TagF) . f . fmap unTagF) k m
	traverseWithKeyH pf s f (TagMap m) = TagMap <$> traverseWithKeyH pf (sizeTag s) (mapTag f) m where
		f' :: Applicative t => ((f :>: ix) r xi -> a -> t (b )) -> f r xi -> TagF a ix xi -> t (TagF b ix xi)
		f' f k (TagF a) = TagF <$> f (Tag k) a
	foldWithKeyH pf f (TagMap m) = foldWithKeyH pf (f' f) m where
		f' :: ((f :>: ix) r xi -> a -> b -> b) -> f r xi -> TagF a ix xi -> b -> b
		f' f k (TagF a) = f (Tag k) a
	foldlWithKeyH pf f (TagMap m) = foldlWithKeyH pf (f' f) m where
		f' :: ((f :>: ix) r xi -> b -> a -> b) -> f r xi -> b -> TagF a ix xi -> b
		f' f k z (TagF a) = f (Tag k) z a
	mapEitherH pf s1 s2 f (TagMap m) = (TagMap *** TagMap) (mapEitherH pf (sizeTag s1) (sizeTag s2) (f' f) m) where
		f' :: EitherMap ((f :>: ix) r xi) (a ) (b) (c) -> EitherMap (f r xi) (TagF a ix xi) (TagF b ix xi) (TagF c ix xi)
		f' f k (TagF a) = (fmap TagF *** fmap TagF) (f (Tag k) a)
	splitLookupH pf s f (Tag k) (TagMap m) = TagMap `sides` splitLookupH pf (sizeTag s) (f' f) k m where
		f' :: SplitMap (a) x -> SplitMap (TagF a xi ix) x
		f' f (TagF a) = fmap TagF `sides` f a
	unionH pf s f (TagMap m1) (TagMap m2) = TagMap (unionH pf (sizeTag s) (combineTag f) m1 m2) 
	isectH pf s f (TagMap m1) (TagMap m2) = TagMap (isectH pf (sizeTag s) (combineTag f) m1 m2)
	diffH pf s f (TagMap m1) (TagMap m2) = TagMap (diffH pf (sizeTag s) (combineTag f) m1 m2)
-- 	extractMinT pf s f (TagMap m) = second TagMap <$> extractMinT pf (sizeTag s) (restructure' f) m
-- 	extractMaxT pf s f (TagMap m) = second TagMap <$> extractMaxT pf (sizeTag s) (restructure' f) m
	extractH pf s f (TagMap m) = fmap TagMap <$> extractH pf (sizeTag s) (restructure' f) m
-- 	alterMinT pf s f (TagMap m) = TagMap <$> alterMinT pf (sizeTag s) (mapTag f) m
-- 	alterMaxT pf s f (TagMap m) = TagMap <$> alterMaxT pf (sizeTag s) (mapTag f) m
	isSubmapH pf (<=) (TagMap m1) (TagMap m2) = isSubmapH pf (le (<=)) m1 m2 where
		le :: LEq a b -> LEq (TagF a xi ix) (TagF b xi ix)
		le (<=) (TagF a) (TagF b) = a <= b
	fromListH pf s f xs = TagMap (fromListH pf (sizeTag s) (f' f) [(k, TagF a) | (Tag k, a) <- xs]) where
		f' :: ((f :>: ix) r xi -> a -> a -> a) -> f r xi -> TagF a ix xi -> TagF a ix xi -> TagF a ix xi
		f' f k (TagF a) (TagF b) = TagF (f (Tag k) a b)
	fromAscListH pf s f xs = TagMap (fromAscListH pf (sizeTag s) (f' f) [(k, TagF a) | (Tag k, a) <- xs]) where
		f' :: ((f :>: ix) r xi -> a -> a -> a ) -> f r xi -> TagF a ix xi -> TagF a ix xi -> TagF a ix xi
		f' f k (TagF a) (TagF b) = TagF (f (Tag k) a b)
	fromDistAscListH pf s xs = TagMap (fromDistAscListH pf (sizeTag s) (map f xs)) where
		f :: ((f :>: ix) r xi, a) -> (f r xi, TagF a ix xi)
		f (Tag k, a) = (k, TagF a)
{-
instance (HTrieKeyT phi f m, m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		HTrieKey phi ((f :>: ix) r) (TagMap phi f ix 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-}