{-# LANGUAGE TypeOperators, Rank2Types, FunctionalDependencies, FlexibleContexts, KindSignatures, TypeFamilies, MultiParamTypeClasses #-}

module Data.TrieMap.MultiRec.Class where

-- import Data.TrieMap.Regular.Class
import Data.TrieMap.CPair
import Data.TrieMap.MultiRec.Sized
-- import Data.TrieMap.MultiRec.Eq
import Data.TrieMap.MultiRec.Ord
-- import Data.TrieMap.Regular.Ord
import Data.TrieMap.MultiRec.Base
-- import Data.TrieMap.MultiRec.KeyFam
import Data.TrieMap.TrieKey
import Data.TrieMap.Applicative

import Control.Applicative
-- import Data.Monoid
-- import Generics.MultiRec.Eq

type family HTrieMapT (phi :: * -> *) (f :: (* -> *) -> * -> *) :: (* -> *) -> * -> * -> *
type family HTrieMap (phi :: * -> *) (r :: * -> *) :: * -> * -> *

class HOrd phi f => HTrieKeyT (phi :: * -> *) (f :: (* -> *) -> * -> *) m | m -> phi f where
	emptyH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> m r ix a
	nullH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> m r ix a -> Bool
	sizeH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> HSized phi a -> m r ix a -> Int
	lookupH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> f r ix -> m r ix a -> Maybe a
	lookupIxH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> HSized phi a -> f r ix -> m r ix a -> IndexPos (f r ix) a
	assocAtH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> HSized phi a -> Int -> m r ix a -> IndexPos (f r ix) a
-- 	updateAtH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
-- 		phi ix -> HSized phi a -> Round -> (Int -> f r ix -> a -> Maybe a) -> Int -> m r ix a -> m r ix a
	alterH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> (Maybe a -> Maybe a) -> f r ix -> m r ix a -> m r ix a
	alterLookupH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) =>
		phi ix -> HSized phi a -> (Maybe a -> CPair x (Maybe a)) -> f r ix ->
			m r ix a -> CPair x (m r ix a)
-- 	{-# SPECIALIZE traverseWithKeyH :: HTrieKey phi r (HTrieMap phi r) =>
-- 		phi ix -> HSized phi b -> (f r ix -> ix a -> Id b) -> m r ix a -> Id (m r ix b) #-}
	traverseWithKeyH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r), Applicative t) =>
		phi ix -> HSized phi b -> (f r ix -> a -> t b) -> m r ix a -> t (m r ix b)
	foldWithKeyH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> (f r ix -> a -> b -> b) -> m r ix a -> b -> b
	foldlWithKeyH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) =>
		phi ix -> (f r ix -> b -> a -> b) -> m r ix a -> b -> b
	mapEitherH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> 
		HSized phi b -> HSized phi c -> EitherMap (f r ix) a b c -> m r ix a -> (m r ix b, m r ix c)
	splitLookupH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> SplitMap a x -> f r ix ->
			m r ix a -> (m r ix a, Maybe x, m r ix a)
	unionH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> UnionFunc (f r ix) a ->
			m r ix a -> m r ix a -> m r ix a
	isectH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi c -> IsectFunc (f r ix) a b c -> m r ix a -> m r ix b -> m r ix c
	diffH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) =>
		phi ix -> HSized phi a -> DiffFunc (f r ix) a b -> m r ix a -> m r ix b -> m r ix a
	extractH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r), Alternative t) =>
		phi ix -> HSized phi a -> ExtractFunc t (m r ix a) (f r ix) a x
-- 	extractMinH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
-- 		phi ix -> HSized phi a -> ExtractFunc (f r ix) First a (m r ix a) x
-- 	extractMaxH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
-- 		phi ix -> HSized phi a -> ExtractFunc (f r ix) Last a (m r ix a) x
-- 	alterMinT:: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
-- 		phi ix -> HSized phi a -> (f r ix -> a -> Maybe a) -> m r ix a -> First (m r ix a)
-- 	alterMaxH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
-- 		phi ix -> HSized phi a -> (f r ix -> a -> Maybe a) -> m r ix a -> Last (m r ix a)
	isSubmapH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> LEq a b -> LEq (m r ix a) (m r ix b)
	fromListH, fromAscListH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> (f r ix -> a -> a -> a ) -> [(f r ix, a )] -> m r ix a
	fromDistAscListH :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> [(f r ix, a )] -> m r ix a
	sizeH pf s m = foldWithKeyH pf (\ _ x n -> s x + n) m 0
	fromListH pf s f = foldr (\ (k, a) -> alterH pf s (Just . maybe a (f k a)) k) (emptyH pf)
	fromAscListH = fromListH
	fromDistAscListH pf s = fromAscListH pf s (const const)

class HOrd0 phi r => HTrieKey (phi :: * -> *) (r :: * -> *) m | m -> phi r where
	empty0 :: m ~ HTrieMap phi r => phi ix -> m ix a
	null0 :: m ~ HTrieMap phi r => phi ix -> m ix a -> Bool
	size0 :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> m ix a -> Int
	lookup0 :: m ~ HTrieMap phi r => phi ix -> r ix -> m ix a -> Maybe a
	lookupIx0 :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> r ix -> m ix a -> IndexPos (r ix) a
	assocAt0 :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> Int -> m ix a -> IndexPos (r ix) a
	alter0 :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> (Maybe a -> Maybe a) -> r ix -> m ix a -> m ix a
	alterLookup0 :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> (Maybe a -> CPair z (Maybe a))
				-> r ix -> m ix a -> CPair z (m ix a)
	extract0 :: (m ~ HTrieMap phi r, Alternative t) => phi ix -> HSized phi a ->
		ExtractFunc t (m ix a) (r ix) a x
	traverseWithKey0 :: (m ~ HTrieMap phi r, Applicative t) => phi ix -> HSized phi b ->
		(r ix -> a -> t b) -> m ix a -> t (m ix b)
	foldWithKey0 :: m ~ HTrieMap phi r => phi ix -> (r ix -> a -> b -> b) -> m ix a -> b -> b
	foldlWithKey0 :: m ~ HTrieMap phi r => phi ix -> (r ix -> b -> a -> b) -> m ix a -> b -> b
	mapEither0 :: m ~ HTrieMap phi r => phi ix -> HSized phi b -> HSized phi c -> EitherMap (r ix) a b c -> m ix a -> (m ix b, m ix c)
	splitLookup0 :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> SplitMap a x ->
		r ix -> m ix a -> (m ix a, Maybe x, m ix a)
	union0 :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> UnionFunc (r ix) a ->
		m ix a -> m ix a -> m ix a
	isect0 :: m ~ HTrieMap phi r => phi ix -> HSized phi c -> IsectFunc (r ix) a b c->
		m ix a -> m ix b -> m ix c
	diff0 :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> DiffFunc (r ix) a b ->
		m ix a -> m ix b -> m ix a
	isSubmap0 :: m ~ HTrieMap phi r => phi ix -> LEq a b -> LEq (m ix a) (m ix b)
	fromList0, fromAscList0 :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> (r ix -> a -> a -> a) -> [(r ix, a)] -> m ix a
	fromDistAscList0 :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> [(r ix, a)] -> m ix a
	

-- class HOrd0 phi r => HTrieKey (phi :: * -> *) (r :: * -> *) m | m -> phi r where
-- 	emptyH :: m ~ HTrieMap phi r => phi ix -> m ix a
-- 	nullH :: m ~ HTrieMap phi r => phi ix -> m ix a -> Bool
-- 	sizeH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> m ix a -> Int
-- 	lookupH :: m ~ HTrieMap phi r => phi ix -> r ix -> m ix a -> Maybe a
-- 	alterH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> (Maybe a -> Maybe a) -> r ix -> m ix a -> m ix a
-- 	lookupIxH :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> r ix -> m ix a -> IndexPos (r ix) a
-- 	assocAtH :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> Int -> m ix a -> IndexPos (r ix) a
-- -- 	updateAtH :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> Round -> (Int -> r ix -> a -> Maybe a) -> Int -> m ix a -> m ix a
-- 	{-# SPECIALIZE traverseWithKeyH :: phi ix -> (r ix -> ix a -> Id b) ->
-- 		m ix a -> Id (m ix b) #-}
-- 	traverseWithKeyH :: (m ~ HTrieMap phi r, Applicative f) => 
-- 		phi ix -> HSized phi b -> (r ix -> a -> f b) -> m ix a -> f (m ix b)
-- 	foldWithKeyH :: m ~ HTrieMap phi r => phi ix -> (r ix -> a -> b -> b) -> m ix a -> b -> b
-- 	foldlWithKeyH :: m ~ HTrieMap phi r => phi ix -> (r ix -> b -> a -> b) -> m ix a -> b -> b
-- 	mapEitherH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi b -> HSized phi c ->
-- 		EitherMap (r ix) a b c -> m ix a -> (m ix b, m ix c)
-- 	splitLookupH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> SplitMap a x -> r ix -> m ix a ->
-- 				(m ix a, Maybe x, m ix a)
-- 	unionH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> UnionFunc (r ix) a -> m ix a -> m ix a
-- 			-> m ix a
-- 	isectH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi c -> IsectFunc (r ix) a b c ->
-- 			m ix a -> m ix b -> m ix c
-- 	diffH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> DiffFunc (r ix) a b ->
-- 			m ix a -> m ix b -> m ix a
-- 	extractH :: (m ~ HTrieMap phi r, Alternative t) =>
-- 		phi ix -> HSized phi a -> ExtractFunc t (m ix a) (r ix) a x
-- -- 	extractMinH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> ExtractFunc (r ix) First a (m ix a) x
-- -- 	extractMaxH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> ExtractFunc (r ix) Last a (m ix a) x
-- -- 	alterMinH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> (r ix -> a -> Maybe a) ->
-- -- 		m ix a -> First (m ix a)
-- -- 	alterMaxH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> (r ix -> a -> Maybe a) ->
-- -- 		m ix a -> Last (m ix a)
-- 	isSubmapH :: m ~ HTrieMap phi r => 
-- 		phi ix -> LEq a b -> LEq (m ix a) (m ix b)
-- 	fromListH, fromAscListH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> (r ix -> a -> a -> a) ->
-- 		[(r ix, a)] -> m ix a
-- 	fromDistAscListH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> [(r ix, a)] -> m ix a
-- 	sizeH pf s m = foldWithKeyH pf (\ _ x n -> s x + n) m 0
-- 	fromListH pf s f = foldr (\ (k, a) -> alterH pf s (Just . maybe a (f k a)) k) (emptyH pf)
-- 	fromAscListH = fromListH
-- 	fromDistAscListH pf s = fromAscListH pf s (const const)
-- 
-- mapWithKeyH :: (HTrieKeyT phi f (HTrieMapT phi f), HTrieKey phi r (HTrieMap phi r)) =>
-- 	phi ix -> HSized phi b -> (f r ix -> a -> b ) -> HTrieMapT phi f r ix a -> HTrieMapT phi f r ix b
-- mapWithKeyT pf s f m = unId (traverseWithKeyT pf s (Id .: f) m)
-- 
-- mapWithKeyH :: (HTrieKey phi r (HTrieMap phi r), HTrieKeyT phi f (HTrieMapT phi f)) =>
-- 	phi ix -> HSized phi b -> (r ix -> a -> b) -> HTrieMapT phi f r ix a -> HTrieMapT phi f r ix b
mapWithKeyH pf s f m = unId (traverseWithKeyH pf s (Id .: f) m)
-- 
guardNullH :: (m ~ HTrieMapT phi f, HTrieKeyT phi f m, HTrieKey phi r (HTrieMap phi r)) => 
	phi ix -> m r ix a -> Maybe (m r ix a)
guardNullH pf m
	| nullH pf m	= Nothing
	| otherwise	= Just m
-- 
-- -- alterMaxT, alterMinH :: (m ~ HTrieMapT phi f, HTrieKeyT phi f m, HTrieKey phi r (HTrieMap phi r)) =>
-- -- 	phi ix -> HSized phi a -> (f r ix -> a -> Maybe a) -> m r ix a -> m r ix a
-- -- alterMaxT pf s f m = maybe m snd $ getLast (extractMaxT pf s (\ k a -> ((), f k a)) m)
-- -- alterMinT pf s f m = maybe m snd $ getFirst (extractMinT pf s (\ k a -> ((), f k a)) m)
-- 
aboutH :: (m ~ HTrieMapT phi f, HTrieKeyT phi f m, HTrieKey phi r (HTrieMap phi r), Alternative t) =>
	phi ix -> (f r ix -> a -> t z) -> m r ix a -> t z
aboutH pf f m = cpFst <$> extractH pf (const 0) (\ k a -> fmap (flip cP Nothing) (f k a)) m

breakEither :: [((f :+: g) r ix, a)] -> ([(f r ix, a)], [(g r ix, a)])
breakEither [] = ([], [])
breakEither ((L k, a):xs) = case breakEither xs of
	(ys, zs) -> ((k, a):ys, zs)
breakEither ((R k, a):xs) = case breakEither xs of
	(ys, zs) -> (ys, (k, a):zs)