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

module Data.TrieMap.MultiRec.Class where

import Data.TrieMap.MultiRec.Sized
import Data.TrieMap.MultiRec.Eq
import Data.TrieMap.MultiRec.Ord
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
	emptyT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> m r ix a
	nullT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> m r ix a -> Bool
	sizeT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> HSized phi a -> m r ix a -> Int
	lookupT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> f r ix -> m r ix a -> Maybe a
	lookupIxT :: (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
	assocAtT :: (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
-- 	updateAtT :: (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
	alterT :: (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
	{-# SPECIALIZE traverseWithKeyT :: HTrieKey phi r =>
		phi ix -> HSized phi b -> (f r ix -> ix a -> Id b) -> m r ix a -> Id (m r ix b) #-}
	traverseWithKeyT :: (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)
	foldWithKeyT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> (f r ix -> a -> b -> b) -> m r ix a -> b -> b
	foldlWithKeyT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) =>
		phi ix -> (f r ix -> b -> a -> b) -> m r ix a -> b -> b
	mapEitherT :: (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)
	splitLookupT :: (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)
	unionT :: (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
	isectT :: (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
	diffT :: (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
	extractT :: (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
-- 	extractMinT :: (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
-- 	extractMaxT :: (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)
-- 	alterMaxT :: (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)
	isSubmapT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> LEq a b -> LEq (m r ix a) (m r ix b)
	fromListT, fromAscListT :: (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
	fromDistAscListT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> [(f r ix, a )] -> m r ix a
	sizeT pf s m = foldWithKeyT pf (\ _ x n -> s x + n) m 0
	fromListT pf s f = foldr (\ (k, a) -> alterT pf s (Just . maybe a (f k a)) k) (emptyT pf)
	fromAscListT = fromListT
	fromDistAscListT pf s = fromAscListT pf s (const const)

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)

mapWithKeyT :: (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)) =>
	phi ix -> HSized phi b -> (r ix -> a -> b) -> HTrieMap phi r ix a -> HTrieMap phi r ix b
mapWithKeyH pf s f m = unId (traverseWithKeyH pf s (Id .: f) m)

guardNullT :: (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)
guardNullT pf m
	| nullT pf m	= Nothing
	| otherwise	= Just m

-- alterMaxT, alterMinT :: (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)

aboutT :: (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
aboutT pf f m = fst <$> extractT pf (const 0) (\ k a -> fmap (flip (,) Nothing) (f k a)) m