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

module Data.TrieMap.MultiRec.FamMap () 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.Base
-- import Data.TrieMap.Sized
-- import Data.TrieMap.Applicative
import Data.TrieMap.TrieKey
-- import qualified Data.TrieMap.Regular.Class as Reg

import Control.Applicative
import Control.Arrow

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

import Generics.MultiRec

newtype FamMap (phi :: * -> *) m ix a = FamMap (m (Family phi) ix a)
type instance HTrieMap phi (Family phi) = FamMap phi (HTrieMapT phi (PF phi))
type instance TrieMap (Family phi ix) = FamMap phi (HTrieMapT phi (PF phi)) ix

to' :: (Fam phi, HFunctor phi (PF phi)) => phi ix -> PF phi (Family phi) ix -> Family phi ix
to' pf = F . to pf . hmap (\ _ (F x) -> I0 x) pf

push :: (Fam phi, HFunctor phi (PF phi)) => phi ix -> (Family phi ix -> a) -> PF phi (Family phi) ix -> a
push pf f = f . to' pf

from' :: (Fam phi, HFunctor phi (PF phi)) => phi ix -> Family phi ix -> PF phi (Family phi) ix
from' pf (F x) = hmap (const (F . unI0)) pf (from pf x)

from'' :: (Fam phi, HFunctor phi (PF phi), El phi ix) => Family phi ix -> PF phi (Family phi) ix
from'' = from' proof

instance (Fam phi, HFunctor phi (PF phi), HEq phi (PF phi)) => HEq0 phi (Family phi) where
	heq0 pf a b = heqT pf (from' pf a) (from' pf b)

instance (Fam phi, HFunctor phi (PF phi), HOrd phi (PF phi)) => HOrd0 phi (Family phi) where
	compare0 pf a b = hcompare pf (from' pf a) (from' pf b)

instance (Fam phi, HFunctor phi (PF phi), HEq phi (PF phi), El phi ix) => Eq (Family phi ix) where
	a == b = heq0 (prove a) a b where
		prove :: El phi ix => Family phi ix -> phi ix
		prove _ = proof

instance (Fam phi, HFunctor phi (PF phi), HOrd phi (PF phi), El phi ix) => Ord (Family phi ix) where
	compare a b = compare0 (prove a) a b where
		prove :: El phi ix => Family phi ix -> phi ix
		prove _ = proof

instance (Fam phi, HFunctor phi (PF phi), HTrieKeyT phi (PF phi) m) => HTrieKey phi (Family phi) (FamMap phi m) where
	empty0 pf = FamMap (emptyH pf)
	null0 pf (FamMap m) = nullH pf m 
	size0 pf s (FamMap m) = sizeH pf s m
	lookup0 pf k (FamMap m) = lookupH pf (from' pf k) m
	lookupIx0 pf s k (FamMap m) = onKey (to' pf) (lookupIxH pf s (from' pf k) m)
	assocAt0 pf s i (FamMap m) = onKey (to' pf) (assocAtH pf s i m)
	alter0 pf s f k (FamMap m) = FamMap (alterH pf s f (from' pf k) m)
	extract0 pf s f (FamMap m) = fmap FamMap <$> extractH pf s (push pf f) m
	alterLookup0 pf s f k (FamMap m) = FamMap <$> alterLookupH pf s f (from' pf k) m
	traverseWithKey0 pf s f (FamMap m) = FamMap <$> traverseWithKeyH pf s (push pf f) m
	foldWithKey0 pf f (FamMap m) = foldWithKeyH pf (push pf f) m
	foldlWithKey0 pf f (FamMap m) = foldlWithKeyH pf (push pf f) m
	mapEither0 pf s1 s2 f (FamMap m) = (FamMap *** FamMap) (mapEitherH pf s1 s2 (push pf f) m)
	splitLookup0 pf s f k (FamMap m) = FamMap `sides` splitLookupH pf s f (from' pf k) m
	union0 pf s f (FamMap m1) (FamMap m2) = FamMap (unionH pf s (push pf f) m1 m2)
	isect0 pf s f (FamMap m1) (FamMap m2) = FamMap (isectH pf s (push pf f) m1 m2)
	diff0 pf s f (FamMap m1) (FamMap m2) = FamMap (diffH pf s (push pf f) m1 m2)
	isSubmap0 pf (<=) (FamMap m1) (FamMap m2) = isSubmapH pf (<=) m1 m2
	fromList0 pf s f xs = FamMap (fromListH pf s (push pf f) [(from' pf k, a) | (k, a) <- xs])
	fromAscList0 pf s f xs = FamMap (fromAscListH pf s (push pf f) [(from' pf k, a) | (k, a) <- xs])
	fromDistAscList0 pf s xs = FamMap (fromDistAscListH pf s [(from' pf k, a) | (k, a) <- xs])

instance (Fam phi, HFunctor phi (PF phi), El phi ix, HTrieKeyT phi (PF phi) m) => TrieKey (Family phi ix) (FamMap phi m ix) where
	emptyM = empty0 proof
	nullM = null0 proof
	sizeM = size0 proof
	lookupM = lookup0 proof
	lookupIxM = lookupIx0 proof
	assocAtM = assocAt0 proof
	alterM = alter0 proof
	alterLookupM = alterLookup0 proof
	extractM = extract0 proof
	traverseWithKeyM = traverseWithKey0 proof
	foldWithKeyM = foldWithKey0 proof
	foldlWithKeyM = foldlWithKey0 proof
	mapEitherM = mapEither0 proof
	splitLookupM = splitLookup0 proof
	unionM = union0 proof
	isectM = isect0 proof
	diffM = diff0 proof
	isSubmapM = isSubmap0 proof
	fromListM = fromList0 proof
	fromAscListM = fromAscList0 proof
	fromDistAscListM = fromDistAscList0 proof