{-# LANGUAGE TemplateHaskell, UndecidableInstances, TypeFamilies, MultiParamTypeClasses, FlexibleContexts #-}

module Data.TrieMap.ReverseMap (reverse, unreverse) where

import Data.TrieMap.TrieKey
import Data.TrieMap.Modifiers
import Data.TrieMap.Applicative
import Data.TrieMap.Regular.Class
import Data.TrieMap.Regular.TH

import Control.Applicative
import Control.Arrow

import Data.Monoid hiding (Dual)

import Prelude hiding (reverse)
import qualified Data.List as L

newtype ReverseMap k a = RMap (TrieMap k a)

type instance TrieMapT Rev = ReverseMap
type instance TrieMap (Rev k) = ReverseMap k

instance TrieKey k (TrieMap k) => TrieKey (Rev k) (ReverseMap k) where
	emptyM = emptyT
	nullM = nullT
	lookupM = lookupT
	lookupIxM = lookupIxT
	assocAtM = assocAtT
	alterM = alterT
	alterLookupM = alterLookupT
	traverseWithKeyM = traverseWithKeyT
	foldWithKeyM = foldWithKeyT
	foldlWithKeyM = foldlWithKeyT
	mapEitherM = mapEitherT
	splitLookupM = splitLookupT
	unionM = unionT
	isectM = isectT
	diffM = diffT
	extractM = extractT
	isSubmapM = isSubmapT
	fromListM = fromListT
	fromAscListM = fromAscListT
	fromDistAscListM = fromDistAscListT


instance TrieKeyT Rev ReverseMap where
	emptyT = RMap emptyM
	nullT (RMap m) = nullM m
	sizeT s (RMap m) = sizeM s m
	lookupT (Rev k) (RMap m) = lookupM k m
	lookupIxT s (Rev k) (RMap m) = case lookupIxM s k m of
		(Last lb, x, First ub) -> onKey Rev (onIndex (sizeM s m - 1 -) (Last ub, x, First lb))
	assocAtT s i (RMap m) = case assocAtM s (sz - 1 - i) m of
		(Last lb, x, First ub) -> onKey Rev (onIndex (sz -) (Last ub, x, First lb))
		where 	sz = sizeM s m
-- 	updateAtM s r f i (RMap m) = RMap (updateAtM s r' f' (sz - i) m) where
-- 		r' = not r
-- 		f' i = f (sz - 1 - i) . Rev
-- 		sz = sizeM s m
	traverseWithKeyT s f (RMap m) = RMap <$> runDual (traverseWithKeyM s (\ k a -> Dual (f (Rev k) a)) m)
	alterT s f (Rev k) (RMap m) = RMap (alterM s f k m)
	alterLookupT s f (Rev k) (RMap m) = RMap <$> alterLookupM s f k m
	splitLookupT s f (Rev k) (RMap m) = case splitLookupM s f' k m of
		(mL, x, mR) -> (RMap mR, x, RMap mL)
		where f' x = case f x of
			(xL, ans, xR) -> (xR, ans, xL)
	mapEitherT s1 s2 f (RMap m) = (RMap *** RMap) (mapEitherM s1 s2 (f . Rev) m)
	foldWithKeyT f (RMap m) = foldlWithKeyM (flip . f . Rev) m
	foldlWithKeyT f (RMap m) = foldWithKeyM (flip . f . Rev) m
	unionT s f (RMap m1) (RMap m2) = RMap (unionM s (f . Rev) m1 m2)
	isectT s f (RMap m1) (RMap m2) = RMap (isectM s (f . Rev) m1 m2)
	diffT s f (RMap m1) (RMap m2) = RMap (diffM s (f . Rev) m1 m2)
	extractT s f (RMap m) = fmap RMap <$> runDual (extractM s (\ k a -> Dual (f (Rev k) a)) m)
-- 	extractMinM s f (RMap m) = second RMap <$> First (getLast (extractMaxM s (f . Rev) m))
-- 	extractMaxM s f (RMap m) = second RMap <$> Last (getFirst (extractMinM s (f . Rev) m))
-- 	alterMinM s f (RMap m) = RMap (alterMaxM s (f . Rev) m)
-- 	alterMaxM s f (RMap m) = RMap (alterMinM s (f . Rev) m)
	isSubmapT (<=) (RMap m1) (RMap m2) = isSubmapM (<=) m1 m2
	fromListT s f xs = RMap (fromListM s (f . Rev) [(k, a) | (Rev k, a) <- xs])
	fromAscListT s f xs = RMap (fromAscListM s (\ k -> flip (f (Rev k))) [(k, a) | (Rev k, a) <- L.reverse xs])
	fromDistAscListT s xs = RMap (fromDistAscListM s [(k, a) | (Rev k, a) <- L.reverse xs])

reverse :: TrieMap k a -> TrieMap (Rev k) a
reverse = RMap

unreverse :: TrieMap (Rev k) a -> TrieMap k a
unreverse (RMap m) = m