{-# 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 a ix
	nullT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> m r a ix -> Bool
	sizeT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => HSized phi a -> m r a ix -> Int
	lookupT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> f r ix -> m r a ix -> Maybe (a ix)
	lookupIxT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> HSized phi a -> f r ix -> m r a ix -> Maybe (Int, a ix)
	assocAtT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => phi ix -> HSized phi a -> Int -> m r a ix -> (Int, f r ix, a ix)
	updateAtT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> (Int -> f r ix -> a ix -> Maybe (a ix)) -> Int -> m r a ix -> m r a ix
	alterT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> (Maybe (a ix) -> Maybe (a ix)) -> f r ix ->
			m r a ix -> m r a ix
	{-# SPECIALIZE traverseWithKeyT :: HTrieKey phi r =>
		phi ix -> HSized phi b -> (f r ix -> a ix -> Id (b ix)) -> m r a ix -> Id (m r b ix) #-}
	traverseWithKeyT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r), Applicative t) =>
		phi ix -> HSized phi b -> (f r ix -> a ix -> t (b ix)) -> m r a ix -> t (m r b ix)
	foldWithKeyT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> (f r ix -> a ix -> b -> b) -> m r a ix -> b -> b
	foldlWithKeyT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) =>
		phi ix -> (f r ix -> b -> a ix -> b) -> m r a ix -> 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 ix) (b ix) (c ix) -> m r a ix -> (m r b ix, m r c ix)
	splitLookupT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> SplitMap (a ix) x -> f r ix ->
			m r a ix -> (m r a ix, Maybe x, m r a ix)
	unionT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> UnionFunc (f r ix) (a ix) ->
			m r a ix -> m r a ix -> m r a ix
	isectT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi c -> IsectFunc (f r ix) (a ix) (b ix) (c ix) -> m r a ix -> m r b ix -> m r c ix
	diffT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) =>
		phi ix -> HSized phi a -> DiffFunc (f r ix) (a ix) (b ix) -> m r a ix -> m r b ix -> m r a ix
	extractMinT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> ExtractFunc (f r ix) First (a ix) (m r a ix)
	extractMaxT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> ExtractFunc (f r ix) Last (a ix) (m r a ix)
	alterMinT, alterMaxT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> (f r ix -> a ix -> Maybe (a ix)) -> m r a ix -> m r a ix
	isSubmapT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> LEq (a ix) (b ix) -> LEq (m r a ix) (m r b ix)
	fromListT, fromAscListT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> (f r ix -> a ix -> a ix -> a ix) -> [(f r ix, a ix)] -> m r a ix
	fromDistAscListT :: (m ~ HTrieMapT phi f, HTrieKey phi r (HTrieMap phi r)) => 
		phi ix -> HSized phi a -> [(f r ix, a ix)] -> m r a ix
-- 	sizeT pf s m = foldWithKeyT pf (\ _ x n -> s pf 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)
	updateAtT pf s f i m = case assocAtT pf s i m of
		(i', k, a) -> alterT pf s (const (f i' k a)) k m

class HOrd0 phi r => HTrieKey (phi :: * -> *) (r :: * -> *) m | m -> phi r where
	emptyH :: m ~ HTrieMap phi r => phi ix -> m a ix
	nullH :: m ~ HTrieMap phi r => phi ix -> m a ix -> Bool
	sizeH :: (m ~ HTrieMap phi r) => HSized phi a -> m a ix -> Int
	lookupH :: m ~ HTrieMap phi r => phi ix -> r ix -> m a ix -> Maybe (a ix)
	alterH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> (Maybe (a ix) -> Maybe (a ix)) -> r ix -> m a ix -> m a ix
	lookupIxH :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> r ix -> m a ix -> Maybe (Int, a ix)
	assocAtH :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> Int -> m a ix -> (Int, r ix, a ix)
	updateAtH :: m ~ HTrieMap phi r => phi ix -> HSized phi a -> (Int -> r ix -> a ix -> Maybe (a ix)) -> Int -> m a ix -> m a ix
	{-# SPECIALIZE traverseWithKeyH :: phi ix -> (r ix -> a ix -> Id (b ix)) ->
		m a ix -> Id (m b ix) #-}
	traverseWithKeyH :: (m ~ HTrieMap phi r, Applicative f) => 
		phi ix -> HSized phi b -> (r ix -> a ix -> f (b ix)) -> m a ix -> f (m b ix)
	foldWithKeyH :: m ~ HTrieMap phi r => phi ix -> (r ix -> a ix -> b -> b) -> m a ix -> b -> b
	foldlWithKeyH :: m ~ HTrieMap phi r => phi ix -> (r ix -> b -> a ix -> b) -> m a ix -> b -> b
	mapEitherH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi b -> HSized phi c ->
		EitherMap (r ix) (a ix) (b ix) (c ix) -> m a ix -> (m b ix, m c ix)
	splitLookupH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> SplitMap (a ix) x -> r ix -> m a ix ->
				(m a ix, Maybe x, m a ix)
	unionH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> UnionFunc (r ix) (a ix) -> m a ix -> m a ix
			-> m a ix
	isectH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi c -> IsectFunc (r ix) (a ix) (b ix) (c ix) ->
			m a ix -> m b ix -> m c ix
	diffH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> DiffFunc (r ix) (a ix) (b ix) ->
			m a ix -> m b ix -> m a ix
	extractMinH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> ExtractFunc (r ix) First (a ix) (m a ix)
	extractMaxH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> ExtractFunc (r ix) Last (a ix) (m a ix)
	alterMinH, alterMaxH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> (r ix -> a ix -> Maybe (a ix)) ->
		m a ix -> m a ix
	isSubmapH :: m ~ HTrieMap phi r => 
		phi ix -> LEq (a ix) (b ix) -> LEq (m a ix) (m b ix)
	fromListH, fromAscListH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> (r ix -> a ix -> a ix -> a ix) ->
		[(r ix, a ix)] -> m a ix
	fromDistAscListH :: (m ~ HTrieMap phi r) => phi ix -> HSized phi a -> [(r ix, a ix)] -> m a ix
-- 	sizeH pf s m = foldWithKeyH pf (\ _ x n -> s pf 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)
	updateAtH pf s f i m = case assocAtH pf s i m of
		(i', k, a) -> alterH pf s (const (f i' k a)) k m

mapWithKeyT :: (HTrieKeyT phi f (HTrieMapT phi f), HTrieKey phi r (HTrieMap phi r)) =>
	phi ix -> HSized phi b -> (f r ix -> a ix -> b ix) -> HTrieMapT phi f r a ix -> HTrieMapT phi f r b ix
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 ix -> b ix) -> HTrieMap phi r a ix -> HTrieMap phi r b ix
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 a ix -> Maybe (m r a ix)
guardNullT pf m
	| nullT pf m	= Nothing
	| otherwise	= Just m