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

module Data.TrieMap.MultiRec.ProdMap () where

import Data.TrieMap.MultiRec.Class
import Data.TrieMap.MultiRec.Eq
import Data.TrieMap.MultiRec.Ord
import Data.TrieMap.MultiRec.Sized
import Data.TrieMap.MultiRec.TH
import Data.TrieMap.Regular.Base (O(..))
import Data.TrieMap.Applicative
import Data.TrieMap.TrieKey
-- import Data.TrieMap.Rep
-- import Data.TrieMap.Rep.TH

import Control.Applicative
import Control.Arrow

import Data.Maybe
import Data.Monoid
import Data.Foldable
import Data.Sequence ((|>))
import qualified Data.Sequence as Seq

import Generics.MultiRec

newtype ProdMap (phi :: * -> *) f g (r :: * -> *) ix a = PMap (HTrieMapT phi f r ix (HTrieMapT phi g r ix a))
type instance HTrieMapT phi (f :*: g) = ProdMap phi f g--(HTrieMapT phi f) (HTrieMapT phi g)
type instance HTrieMap phi ((f :*: g) r) = HTrieMapT phi (f :*: g) r

-- type instance RepT (ProdMap phi f g r ix) = RepT (HTrieMapT phi f r ix) `O` RepT (HTrieMapT phi g r ix)
-- type instance Rep (ProdMap phi f g r ix a) = RepT (ProdMap phi f g r ix) (Rep a)

-- -- $(genRepT [d|
-- 	instance (ReprT (HTrieMapT phi f r ix), ReprT (HTrieMapT phi g r ix)) =>
-- 			ReprT (ProdMap phi f g r ix) where
-- 		toRepT (PMap m) = O (fmap toRepT (toRepT m))
-- 		fromRepT (O m) = PMap (fromRepT (fmap fromRepT m)) |] )

maxIx :: (HTrieKeyT phi f (HTrieMapT phi f), HTrieKey phi r (HTrieMap phi r)) => phi ix -> HSized phi a -> 
		HTrieMapT phi f r ix a -> Int
maxIx pf s m = fromMaybe (sizeT pf s m) (getFirst (aboutT pf (\ _ a -> return (sizeT pf s m - s a)) m))

$(inferH [d|
	instance (HTrieKeyT phi f (HTrieMapT phi f), HTrieKeyT phi g (HTrieMapT phi g)) => 
			HTrieKeyT phi (f :*: g) (ProdMap phi f g) where
		emptyT = PMap . emptyT
		nullT pf (PMap m) = nullT pf m
		sizeT pf s (PMap m) = sizeT pf (sizeT pf s) m
		lookupT pf (a :*: b) (PMap m) = lookupT pf a m >>= lookupT pf b
		lookupIxT pf s (a :*: b) (PMap m) = case lookupIxT pf (sizeT pf s) a m of
			(lb, x, rb) -> let lookupX = do	Asc i a' m' <- x
							let (lb', x', rb') = lookupIxT pf s b m'
							let f = onIndexA (i +) . onKeyA (a' :*:)
							return (f <$> lb', f <$> x', f <$> rb')
			   in 	((do	Asc iA aL mL <- lb
					fmap (onIndexA (iA +) . onKeyA (aL :*:)) (getLast pf s mL)) <|>
				 (do	(lb', _, _) <- Last lookupX
					lb'),
				 (do	(_, x', _) <- lookupX
					x'),
				 (do	(_, _, rb') <- First lookupX
					rb') <|>
				 (do	Asc iA aR mR <- rb
				  	fmap (onIndexA (iA +) . onKeyA (aR :*:)) (getFirst pf s mR)))
			where 	getLast pf s m = aboutT pf (\ k a -> return (Asc (sizeT pf s m - s a) k a)) m
				getFirst pf s m = aboutT pf (\ k a -> return (Asc 0 k a)) m
		assocAtT pf s i (PMap m) = case assocAtT pf (sizeT pf s) i m of
			(lb, x, rb) -> let lookupX = do	Asc i' a' m' <- x
							let (lb', x', rb') = assocAtT pf s (i - i') m'
							let f = onIndexA (i' +) . onKeyA (a' :*:)
							return (f <$> lb', f <$> x', f <$> rb')
				in ((do	Asc iA aL mL <- lb
					fmap (onIndexA (iA +) . onKeyA (aL :*:)) (getLast pf s mL)) <|>
				    (do	(lb', _, _) <- Last lookupX
				    	lb'),
				    (do	(_, x', _) <- lookupX
				    	x'),
				    (do	(_, _, rb') <- First lookupX
				    	rb') <|>
				    (do	Asc iA aR mR <- rb
				    	fmap (onIndexA (iA +) . onKeyA (aR :*:)) (getFirst pf s mR)))
			where 	getLast pf s m = aboutT pf (\ k a -> return (Asc (sizeT pf s m - s a) k a)) m
				getFirst pf s m = aboutT pf (\ k a -> return (Asc 0 k a)) m
-- 		updateAtT pf s r f i (PMap m) = PMap (updateAtT pf (sizeT pf s) r g i m) where
-- 			g iA a m 
-- 				| i >= iA && i <= iA + maxIx pf s m
-- 					= (guardNullT pf . updateAtT pf s r (\ iB b -> f (iA + iB) (a :*: b)) (i - iA)) m
-- 				| i < iA
-- 					= guardNullT pf $
-- 						alterMaxT pf s (\ b v -> f (iA + sizeT pf s m - s v) (a :*: b) v) m
-- 				| otherwise
-- 					= guardNullT pf $ alterMinT pf s (f iA . (a :*:)) m
		alterT pf s f (a :*: b) (PMap m) = PMap (alterT pf (sizeT pf s) (guardNullT pf . g) a m) where
			g = alterT pf s f b . fromMaybe (emptyT pf)
		traverseWithKeyT pf s f (PMap m) = 
			PMap <$> traverseWithKeyT pf (sizeT pf s) (\ a -> traverseWithKeyT pf s (\ b -> f (a :*: b))) m
		foldWithKeyT pf f (PMap m) =
			foldWithKeyT pf (\ a -> foldWithKeyT pf (\ b -> f (a :*: b))) m
		foldlWithKeyT pf f (PMap m) =
			foldlWithKeyT pf (\ a -> flip (foldlWithKeyT pf (\ b -> f (a :*: b)))) m
		mapEitherT pf s1 s2 f (PMap m) = (PMap *** PMap) (mapEitherT pf (sizeT pf s1) (sizeT pf s2) g m) where
			g a = (guardNullT pf *** guardNullT pf) . mapEitherT pf s1 s2 (\ b -> f (a :*: b))
		splitLookupT pf s f (a :*: b) (PMap m) = PMap `sides` splitLookupT pf (sizeT pf s) g a m where
			g = sides (guardNullT pf) . splitLookupT pf s f b
		unionT pf s f (PMap m1) (PMap m2) = PMap (unionT pf (sizeT pf s) g m1 m2) where
			g a = guardNullT pf .: unionT pf s (\ b -> f (a :*: b))
		isectT pf s f (PMap m1) (PMap m2) = PMap (isectT pf (sizeT pf s) g m1 m2) where
			g a = guardNullT pf .: isectT pf s (\ b -> f (a :*: b))
		diffT pf s f (PMap m1) (PMap m2) = PMap (diffT pf (sizeT pf s) g m1 m2) where
			g a = guardNullT pf .: diffT pf s (\ b -> f (a :*: b))
		extractT pf s f (PMap m) = second PMap <$> extractT pf (sizeT pf s) g m where
			g a = second (guardNullT pf) <.> extractT pf s (\ b -> f (a :*: b))
-- 		extractMinT pf s f (PMap m) = second PMap <$> extractMinT pf (sizeT pf s) g m where 
-- 			g a m1 = fromJust $ getFirst $ second (guardNullT pf) <$> extractMinT pf s (f . (a :*:)) m1
-- 		extractMaxT pf s f (PMap m) = second PMap <$> extractMaxT pf (sizeT pf s) g m where 
-- 			g a m1 = fromJust $ getLast $ second (guardNullT pf) <$> extractMaxT pf s (f . (a :*:)) m1
-- 		alterMinT pf s f (PMap m) = PMap (alterMinT pf (sizeT pf s) g m) where
-- 			g a = guardNullT pf . alterMinT pf s (\ b -> f (a :*: b))
-- 		alterMaxT pf s f (PMap m) = PMap (alterMaxT pf (sizeT pf s) g m) where
-- 			g a = guardNullT pf . alterMaxT pf s (\ b -> f (a :*: b))
		isSubmapT pf (<=) (PMap m1) (PMap m2) = isSubmapT pf (isSubmapT pf (<=)) m1 m2
		fromListT pf s f xs = PMap (mapWithKeyT pf (sizeT pf s) (\ a -> fromListT pf s (\ b -> f (a :*: b)))
					(fromListT pf (const 1) (\ _ (xs) (ys) -> (xs ++ ys))
						[(a, ts) | (a, ts) <- breakFst pf xs]))
		fromAscListT pf s f xs = PMap (fromDistAscListT pf (sizeT pf s)
			[(a, fromAscListT pf s (\ b -> f (a :*: b)) ts) | (a, ts) <- breakFst pf xs])
		fromDistAscListT pf s xs = PMap (fromDistAscListT pf (sizeT pf s)
			[(a, fromDistAscListT pf s ts) | (a, ts) <- breakFst pf xs])

	breakFst :: (HEq phi f, HEq0 phi r) => phi ix -> [((f :*: g) r ix, a)] -> [(f r ix, [(g r ix, a)])]
	breakFst pf [] = []
	breakFst pf ((a :*: b, x):xs) = breakFst' a (Seq.singleton (b, x)) xs where
		breakFst' a0 ts ((a :*: b, x):xs)
			| heqT pf a0 a	= breakFst' a0 (ts |> (b, x)) xs
			| otherwise	= (a0, toList ts):breakFst' a (Seq.singleton (b,x)) xs
		breakFst' a ts [] = [(a, toList ts)]
	|])