{-# LANGUAGE CPP, Rank2Types, TypeOperators, TypeFamilies, FlexibleContexts, UndecidableInstances #-}

module Data.TrieMap.Class.Instances where

import Data.TrieMap.Class
import Data.TrieMap.TrieKey
-- import Data.TrieMap.RadixTrie()
import Data.TrieMap.MultiRec.Instances
import Data.TrieMap.IntMap()
import Data.TrieMap.OrdMap(Ordered(..))
import Data.TrieMap.Class
import Data.TrieMap.Regular.Base
import Data.TrieMap.Regular.Class
import Data.TrieMap.Regular.Instances
-- import Data.TrieMap.UnionMap()
-- import Data.TrieMap.UnitMap()

import Data.Bits
import Data.Char
import Data.Complex
import Data.Either
import Data.Foldable
import Data.Int 
import Data.List hiding (foldr)
import Data.Word

import Prelude hiding (foldr)

#if __GLASGOW_HASKELL__
import GHC.Exts (build)
#else

build :: (forall b . (a -> b -> b) -> b -> b) -> [a]
build f = f (:) []
#endif
{-
instance TKey k => TKey [k] where
	type Rep [k] = L I0 (Rep k)
	toRep = map toRep
	fromRep = map fromRep-}

type instance Rep Int = Int
instance TKey Int where
	toRep = id
	fromRep = id

type instance Rep Double = Ordered Double
instance TKey Double where
	toRep = Ord
	fromRep = unOrd

type instance Rep Char = Int
instance TKey Char where
	toRep = ord
	fromRep = chr

type instance Rep Word = Int
instance TKey Word where
	toRep = fromEnum
	fromRep = toEnum

type instance Rep Word8 = Int
instance TKey Word8 where
	toRep = fromEnum
	fromRep = toEnum

type instance Rep Word16 = Int
instance TKey Word16 where
	toRep = fromEnum
	fromRep = toEnum

type instance Rep Word32 = Int
instance TKey Word32 where
	toRep = fromEnum
	fromRep = toEnum

type instance Rep Int8 = Int
instance TKey Int8 where
	toRep = fromIntegral
	fromRep = fromIntegral

type instance Rep Int16 = Int
instance TKey Int16 where
	toRep = fromIntegral
	fromRep = fromIntegral

type instance Rep Int32 = Int
instance TKey Int32 where
	toRep = fromIntegral
	fromRep = fromIntegral
-- 
-- type instance Rep (Complex a) = Rep (a, a)
-- instance (RealFloat a, TKey a) => TKey (Complex a) where
-- 	toRep (a :+ b) = toRep (a, b)
-- 	fromRep = uncurry (:+) . fromRep

type instance Rep Integer = (I0 :+: I0) (Rep [Int32])
instance TKey Integer where
	toRep x = (if x >= 0 then R . I0 else L . I0) (toRep (unroll x))
	fromRep (L (I0 xs)) = - roll (map negate (fromRep xs))
	fromRep (R (I0 xs)) = roll (fromRep xs)

unroll :: Integer -> [Int32]
unroll x = if x >= 0 then unfoldr step x else map negate (unfoldr step (negate x)) where 
	step 0 = Nothing
	step i = Just (fromIntegral i,  i `shiftR` 32)

roll :: [Int32] -> Integer
roll = foldr unstep 0 where
	unstep b a = a `shiftL` 32 .|. fromIntegral b

type instance Rep () = U0 ()
instance TKey () where
	toRep _ = U0
	fromRep _ = ()

type instance Rep (Either a b) = (K0 (Rep a) :+: I0) (Rep b)
instance (TKey a, TKey b) => TKey (Either a b) where
	toRep = either (L . K0 . toRep) (R . I0 . toRep)
	fromRep = either' (Left . unK0 . fromRep) (Right . unI0 . fromRep)

either' :: (f r -> a) -> (g r -> a) -> (f :+: g) r -> a
either' f g x = case x of
	L x	-> f x
	R x	-> g x

type instance Rep (a, b) = (K0 (Rep a) :*: I0) (Rep b)
instance (TKey a, TKey b) => TKey (a, b) where
	toRep (a, b) = K0 (toRep a) :*: I0 (toRep b)
	fromRep (K0 a :*: I0 b) = (fromRep a, fromRep b)

type instance Rep (a, b, c) = (K0 (Rep a) :*: K0 (Rep b) :*: I0) (Rep c)
instance (TKey a, TKey b, TKey c) => TKey (a, b, c) where
	toRep (a, b, c) = K0 (toRep a) :*: K0 (toRep b) :*: I0 (toRep c)
	fromRep (K0 a :*: K0 b :*: I0 c) = (fromRep a, fromRep b, fromRep c)

type instance Rep (a, b, c, d) = (K0 (Rep a) :*: K0 (Rep b) :*: K0 (Rep c) :*: I0) (Rep d)
instance (TKey a, TKey b, TKey c, TKey d) => TKey (a, b, c, d) where
	toRep (a, b, c, d) = K0 (toRep a) :*: K0 (toRep b) :*: K0 (toRep c) :*: I0 (toRep d)
	fromRep (K0 a :*: K0 b :*: K0 c :*: I0 d) = (fromRep a, fromRep b, fromRep c, fromRep d)

type instance Rep (a, b, c, d, e) = (K0 (Rep a) :*: K0 (Rep b) :*: K0 (Rep c) :*: K0 (Rep d) :*: I0) (Rep e)
instance (TKey a, TKey b, TKey c, TKey d, TKey e) => TKey (a, b, c, d, e) where
	toRep (a, b, c, d, e) = K0 (toRep a) :*: K0 (toRep b) :*: K0 (toRep c) :*: K0 (toRep d) :*: I0 (toRep e)
	fromRep (K0 a :*: K0 b :*: K0 c :*: K0 d :*: I0 e) = (fromRep a, fromRep b, fromRep c, fromRep d, fromRep e)

type instance Rep Bool = (U0 :+: U0) (U0 ())
instance TKey Bool where
	toRep False = L U0
	toRep True = R U0
	fromRep L{} = False
	fromRep R{} = True

type instance Rep (Maybe a) = (U0 :+: I0) (Rep a)
instance TKey a => TKey (Maybe a) where
	toRep = maybe (L U0) (R . I0 . toRep)
	fromRep = either' (const Nothing) (Just . fromRep . unI0)

type instance Rep [a] = L I0 (Rep a)
instance TKey a => TKey [a] where
	toRep xs = List [I0 (toRep x) | x <- xs]
	fromRep (List xs) = [fromRep x | I0 x <- xs]

type instance Rep ((f :*: g) r) = (f :*: g) (Rep r)
instance (TKey a, TrieKeyT f (TrieMapT f), TrieKeyT g (TrieMapT g), Functor f, Functor g) => TKey ((f :*: g) a) where
	toRep = fmap toRep
	fromRep = fmap fromRep

type instance Rep ((f :+: g) r) = (f :+: g) (Rep r)
instance (TKey a, TrieKeyT f (TrieMapT f), TrieKeyT g (TrieMapT g), Functor f, Functor g) => TKey ((f :+: g) a) where
	toRep = fmap toRep
	fromRep = fmap fromRep
{-
type instance Rep [r] = L I0 (Rep r)
instance TKey r => TKey [r] where
	toRep = List . map (I0 . toRep)
	fromRep (List xs) = [fromRep x | I0 x <- xs]-}

type instance Rep (L f r) = L f (Rep r)
instance (TKey a, TrieKeyT f (TrieMapT f), Functor f) => TKey (L f a) where
	toRep = fmap toRep
	fromRep = fmap fromRep

type instance Rep (U0 r) = U0 r
instance TKey (U0 r) where
	toRep _ = U0
	fromRep _ = U0

type instance Rep (K0 k r) = K0 (Rep k) r
instance TKey k => TKey (K0 k r) where
	toRep (K0 a) = K0 (toRep a)
	fromRep (K0 a) = K0 (fromRep a)

type instance Rep (I0 r) = I0 (Rep r)
instance TKey r => TKey (I0 r) where
	toRep = fmap toRep
	fromRep = fmap fromRep

type instance Rep (TMap k a) = L (K0 (Rep k) :*: I0) (Rep a)
instance (TKey k, TKey a) => TKey (TMap k a) where
	toRep (TMap m) = List [K0 k :*: I0 (toRep a) | (k, K0 a) <- foldWithKeyM (curry (:)) m []]
	fromRep (List xs) = TMap (fromDistAscListM (const 1) [(k, K0 (fromRep a)) | K0 k :*: I0 a <- xs])