{-# 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))

instance (Fam phi, HEq phi (PF phi), HFunctor phi (PF phi)) => HEq0 phi (Family phi) where
	heqH pf (F x) (F y) = heqT pf (from' pf x) (from' pf y)

instance (Fam phi, HOrd phi (PF phi), HFunctor phi (PF phi)) => HOrd0 phi (Family phi) where
	compareH0 pf (F x) (F y) = hcompare pf (from' pf x) (from' pf y)

instance (El phi ix, Fam phi, HEq phi (PF phi), HFunctor phi (PF phi)) => Eq (Family phi ix) where
	x == y = heqH (prove x) x y

instance (El phi ix, Fam phi, HOrd phi (PF phi), HFunctor phi (PF phi)) => Ord (Family phi ix) where
	x `compare` y = compareH0 (prove x) x y

instance HEq0 phi r => HEq0 phi (FamT phi r) where
	heqH pf (FamT x) (FamT y) = heqH pf x y

instance HOrd0 phi r => HOrd0 phi (FamT phi r) where
	compareH0 pf (FamT x) (FamT y) = compareH0 pf x y

instance (El phi ix, HEq0 phi r) => Eq (FamT phi r ix) where
	x == y = heqH (prove' x) x y

instance (El phi ix, HOrd0 phi r) => Ord (FamT phi r ix) where
	x `compare` y = compareH0 (prove' x) x y

prove' :: El phi ix => FamT phi r ix -> phi ix
prove' _ = proof

prove :: El phi ix => Family phi ix -> phi ix
prove _ = proof

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

to' :: (Fam phi, HFunctor phi (PF phi)) => phi ix -> PF phi (Family phi) ix -> ix
to' pf = 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 . F . to' pf

instance (Fam phi, HFunctor phi (PF phi), HTrieKeyT phi (PF phi) m) => HTrieKey phi (Family phi) (FamMap phi m) where
	emptyH pf = FamMap (emptyT pf)
	nullH pf (FamMap m) = nullT pf m
	sizeH pf s (FamMap m) = sizeT pf s m
	lookupH pf (F k) (FamMap m) = lookupT pf (from' pf k) m
	lookupIxH pf s (F k) (FamMap m) = onKey (F . to' pf) (lookupIxT pf s (from' pf k) m)
	assocAtH pf s i (FamMap m) = onKey (F . to' pf) (assocAtT pf s i m)
-- 	updateAtH pf s r f i (FamMap m) = FamMap (updateAtT pf s r (\ i -> f i . F . to' pf) i m)
	alterH pf s f (F k) (FamMap m) = FamMap (alterT pf s f (from' pf k) m)
	traverseWithKeyH pf s f (FamMap m) =
		FamMap <$> traverseWithKeyT pf s (push pf f) m
	foldWithKeyH pf f (FamMap m) = foldWithKeyT pf (push pf f) m
	foldlWithKeyH pf f (FamMap m) = foldlWithKeyT pf (push pf f) m
	mapEitherH pf s1 s2 f (FamMap m) = (FamMap *** FamMap) (mapEitherT pf s1 s2 (push pf f) m)
	splitLookupH pf s f (F k) (FamMap m) = FamMap `sides` splitLookupT pf s f (from' pf k) m
	unionH pf s f (FamMap m1) (FamMap m2) = FamMap (unionT pf s (push pf f) m1 m2)
	isectH pf s f (FamMap m1) (FamMap m2) = FamMap (isectT pf s (push pf f) m1 m2)
	diffH pf s f (FamMap m1) (FamMap m2) = FamMap (diffT pf s (push pf f) m1 m2)
	extractH pf s f (FamMap m) = second FamMap <$> extractT pf s (f . F . to' pf) m
-- 	extractMinH pf s f (FamMap m) = second FamMap <$> extractMinT pf s (f . F . to' pf) m
-- 	extractMaxH pf s f (FamMap m) = second FamMap <$> extractMaxT pf s (f . F . to' pf) m
-- 	alterMinH pf s f (FamMap m) = FamMap (alterMinT pf s (push pf f) m)
-- 	alterMaxH pf s f (FamMap m) = FamMap (alterMaxT pf s (push pf f) m)
	isSubmapH pf (<=) (FamMap m1) (FamMap m2) = isSubmapT pf (<=) m1 m2
	fromListH pf s f xs = FamMap (fromListT pf s (push pf f) [(from' pf k, a) | (F k, a) <- xs])
	fromAscListH pf s f xs = FamMap (fromAscListT pf s (push pf f) [(from' pf k, a) | (F k, a) <- xs])
	fromDistAscListH pf s xs = FamMap (fromDistAscListT pf s [(from' pf k, a) | (F k, a) <- xs])

-- type family UniqueFam ix :: * -> *
newtype FMap (phi :: * -> *) m ix a = FMap (m ix a)
type instance TrieMap (Family phi ix) = FMap phi (HTrieMap phi (Family phi)) ix

instance (El phi ix, Fam phi, HFunctor phi (PF phi), HTrieKey phi (Family phi) m, m ~ HTrieMap phi (Family phi),
		HOrd phi (PF phi)) => TrieKey (Family phi ix) (FMap phi m ix) where
	emptyM = FMap (emptyH proof)
	nullM (FMap m) = nullH proof m
	sizeM s (FMap m) = sizeH proof s m
	lookupM k (FMap m) = lookupH proof k m
	lookupIxM s k (FMap m) = lookupIxH proof s k m
	assocAtM s i (FMap m) = assocAtH proof s i m
-- 	updateAtM s r f i (FMap m) = FMap (updateAtH proof s r f i m)
	alterM s f k (FMap m) = FMap (alterH proof s f k m)
	traverseWithKeyM s f (FMap m) = FMap <$> traverseWithKeyH proof s f m
	foldWithKeyM f (FMap m) = foldWithKeyH proof f m
	foldlWithKeyM f (FMap m) = foldlWithKeyH proof f m
	mapEitherM s1 s2 f (FMap m) = 
		(FMap *** FMap) (mapEitherH proof s1 s2 f m)
	splitLookupM s f k (FMap m) = FMap `sides` splitLookupH proof s (f) k m
	unionM s f (FMap m1) (FMap m2) = FMap (unionH proof s f m1 m2) 
	isectM s f (FMap m1) (FMap m2) = FMap (isectH proof s f m1 m2) 
	diffM s f (FMap m1) (FMap m2) = FMap (diffH proof s f m1 m2) 
	extractM s f (FMap m) = second FMap <$> extractH proof s f m
-- 	extractMinM s f (FMap m) = second FMap <$> extractMinH proof s f m
-- 	extractMaxM s f (FMap m) = second FMap <$> extractMaxH proof s f m
-- 	alterMinM s f (FMap m) = FMap (alterMinH proof s f m)
-- 	alterMaxM s f (FMap m) = FMap (alterMaxH proof s f m)
	isSubmapM (<=) (FMap m1) (FMap m2) = isSubmapH proof (<=) m1 m2 
	fromListM s f xs = FMap (fromListH proof s f xs)
	fromAscListM s f xs = FMap (fromAscListH proof s f xs)
	fromDistAscListM s xs = FMap (fromDistAscListH proof s xs)

newtype FTMap (phi :: * -> *) (r :: * -> *) ix a = FTMap (HTrieMap phi r ix a)
type instance TrieMap (FamT phi r ix) = FTMap phi r ix

-- instance (HTrieKey KeyFam r (HTrieMap KeyFam r)) => Reg.TrieKeyT (FamT KeyFam r) (FTMap KeyFam r) where
-- 	emptyT = FTMap (emptyH KF)

instance (El phi ix, HTrieKey phi r (HTrieMap phi r)) => TrieKey (FamT phi r ix) (FTMap phi r ix) where
	emptyM = FTMap (emptyH proof)
	nullM (FTMap m) = nullH proof m
	sizeM s (FTMap m) = sizeH proof s m
	lookupM (FamT k) (FTMap m) = lookupH proof k m
	lookupIxM s (FamT k) (FTMap m) = onKey FamT (lookupIxH proof s k m)
	assocAtM s i (FTMap m) = onKey FamT (assocAtH proof s i m)
-- 	updateAtM s r f i (FTMap m) = FTMap (updateAtH proof s r (\ i' -> f i' . FamT) i m)
	alterM s f (FamT k) (FTMap m) = FTMap (alterH proof s f k m)
	foldWithKeyM f (FTMap m) = foldWithKeyH proof (f . FamT) m
	foldlWithKeyM f (FTMap m) = foldlWithKeyH proof (f . FamT) m
	traverseWithKeyM s f (FTMap m) = FTMap <$> traverseWithKeyH proof s (f . FamT) m
	mapEitherM s1 s2 f (FTMap m) = (FTMap *** FTMap) (mapEitherH proof s1 s2 (f . FamT) m)
	splitLookupM s f (FamT k) (FTMap m) = FTMap `sides` splitLookupH proof s f k m
	unionM s f (FTMap m1) (FTMap m2) = FTMap (unionH proof s (f . FamT) m1 m2)
	isectM s f (FTMap m1) (FTMap m2) = FTMap (isectH proof s (f . FamT) m1 m2)
	diffM s f (FTMap m1) (FTMap m2) = FTMap (diffH proof s (f . FamT) m1 m2)
	isSubmapM (<=) (FTMap m1) (FTMap m2) = isSubmapH proof (<=) m1 m2
	extractM s f (FTMap m) = second FTMap <$> extractH proof s (f . FamT) m
-- 	extractMinM s f (FTMap m){--} = second FTMap <$> extractMinH proof s (f . FamT) m
-- 	extractMaxM s f (FTMap m) = second FTMap <$> extractMaxH proof s (f . FamT) m
-- 	alterMinM s f (FTMap m) = FTMap (alterMinH proof s (f . FamT) m)
-- 	alterMaxM s f (FTMap m) = FTMap (alterMaxH proof s (f . FamT) m)