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

module Data.TrieMap.Rep.Instances()  where

import Data.TrieMap.Rep
import Data.TrieMap.Rep.TH
import Data.TrieMap.Regular.Base
import Data.TrieMap.OrdMap
import Data.TrieMap.Modifiers
-- import Language.Haskell.TH

import Control.Arrow

import Data.Char
import Data.Int
import Data.Word
import Data.Foldable (toList)
import Data.Bits
import Data.Array.IArray
import Data.ByteString hiding (map)
import qualified Data.ByteString as BS

import Data.Sequence (Seq, (|>))
import qualified Data.Sequence as Seq
import qualified Data.Foldable as Fold

import qualified Data.Map as Map
import qualified Data.Set as Set

import Prelude hiding (concat, take, length)

type Pair a = (,) a
type Sum a = Either a

type instance RepT Maybe = U0 :+: I0
type instance Rep (Maybe a) = RepT Maybe (Rep a)

$(genRepT [d|
   instance ReprT Maybe where
	toRepTMap f = maybe (L U0) (R . I0 . f)
	fromRepTMap _ L{} = Nothing
	fromRepTMap f (R (I0 a)) = Just (f a) |])

type instance RepT [] = L I0
type instance Rep [a] = L I0 (Rep a)

$(genRepT [d| 
   instance ReprT [] where
	toRepTMap f = List . map (I0 . f)
	fromRepTMap f (List xs) = map (f . unI0) xs |])

type instance RepT ((,) a) = Pair (Rep a)
type instance Rep (a, b) = RepT ((,) a) (Rep b)

$(genRepT [d|
    instance Repr a => ReprT ((,) a) where
	toRepTMap f = toRep *** f
	fromRepTMap f = fromRep *** f |])

-- instance (ReprT ((,) a), Repr b) => Repr ((,) a b) where

-- instance (Repr a, Repr b) => Repr (a, b) where
-- 	toRep = fmap toRep . toRepT
-- 	fromRep = fromRepT . fmap fromRep

type instance RepT ((,,) a b) = K0 (Rep a) :*: K0 (Rep b) :*: I0
type instance Rep (a, b, c) = RepT ((,,) a b) (Rep c)

$(genRepT [d|
   instance (Repr a, Repr b) => ReprT ((,,) a b) where
	toRepTMap f (a, b, c) = K0 (toRep a) :*: K0 (toRep b) :*: I0 (f c)
	fromRepTMap f (K0 a :*: K0 b :*: I0 c) = (fromRep a, fromRep b, f c) |])

type instance RepT ((,,,) a b c) = K0 (Rep a) :*: K0 (Rep b) :*: K0 (Rep c) :*: I0
type instance Rep (a, b, c, d) = RepT ((,,,) a b c) (Rep d)

$(genRepT [d|
   instance (Repr a, Repr b, Repr c) => ReprT ((,,,) a b c) where
	toRepTMap f (a, b, c, d) = K0 (toRep a) :*: K0 (toRep b) :*: K0 (toRep c) :*: I0 (f d)
	fromRepTMap f (K0 a :*: K0 b :*: K0 c :*: I0 d) = (fromRep a, fromRep b, fromRep c, f d) |])

type instance RepT (Either a) = Sum (Rep a)
type instance Rep (Either a b) = RepT (Either a) (Rep b)

$(genRepT [d|
  instance Repr a => ReprT (Either a) where
	toRepTMap f = either (Left . toRep) (Right . f)
	fromRepTMap f = either (Left . fromRep) (Right . f) |])

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

type instance Rep Char = Word32
instance Repr Char where
	toRep = fromIntegral . ord
	fromRep = chr . fromIntegral

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

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

type instance Rep Int = Rep Int32
instance Repr Int where
	toRep = toSigned
	fromRep = fromSigned

type instance Rep Word8 = Word32
instance Repr Word8 where
	toRep = fromIntegral
	fromRep = fromIntegral

type instance Rep Word16 = Word32
instance Repr Word16 where
	toRep = fromIntegral
	fromRep = fromIntegral

type instance Rep Word = Word32
instance Repr Word where
	toRep = fromIntegral
	fromRep = fromIntegral

type instance Rep Int8 = Rep Int32

instance Repr Int8 where
	toRep = toSigned
	fromRep = fromSigned

type instance Rep Int16 = Rep Int32
instance Repr Int16 where
	toRep = toSigned
	fromRep = fromSigned

type instance Rep Int32 = Sum (Rev Word32) Word32
instance Repr Int32 where
	toRep = toSigned
	fromRep = fromSigned

type instance Rep Word64 = Pair Word32 Word32
instance Repr Word64 where
	toRep x = (fromIntegral (x `shiftR` 32), fromIntegral x)
	fromRep (x, y) = fromIntegral x `shiftL` 32 .|. fromIntegral y

type instance Rep Int64 = Sum (Rev (Rep Word64)) (Rep Word64)
instance Repr Int64 where
	toRep x | x < 0	= Left (Rev (toRep' (fromIntegral (-x))))
		| otherwise = Right (toRep' (fromIntegral x))
		where toRep' = toRep :: Word64 -> Rep Word64
	fromRep (Left (Rev x)) = - fromIntegral ((fromRep :: Rep Word64 -> Word64) x)
	fromRep (Right x) = fromIntegral ((fromRep :: Rep Word64 -> Word64) x)

{-# INLINE toSigned #-}
toSigned :: Integral a => a -> Sum (Rev Word32) Word32
toSigned x
	| x < 0		= Left (Rev (fromIntegral (-x)))
	| otherwise	= Right (fromIntegral x)

{-# INLINE fromSigned #-}
fromSigned :: Integral a => Sum (Rev Word32) Word32 -> a
fromSigned = either (\ (Rev x) -> - fromIntegral x) fromIntegral

type instance Rep Word32 = Word32
instance Repr Word32 where
	toRep = id
	fromRep = id

type instance Rep ByteString = (L I0 :*: I0) Word32
instance Repr ByteString where
	toRep xs = List (toList64 xs) :*: I0 (fromIntegral (length xs))
	fromRep (List xs :*: I0 n) = case xs of
		[]	-> BS.empty
		(I0 x:xs) -> fst (unfoldrN (fromIntegral n) toBlock (W (Words 3 x) xs))

data Words = Words {ix :: {-# UNPACK #-} !Int, word32 :: {-# UNPACK #-} !Word32}
data Words' = W {-# UNPACK #-} !Words [I0 Word32]

toList64 :: ByteString -> [I0 Word32]
toList64 xs = case BS.foldl c (Words 4 0, Seq.empty) xs of
	(Words i w32, ys) -> toList ys ++ [I0 w32]
	where	fS :: Word8 -> Int -> Word32
		fS w x = fromIntegral w `shiftL` x
		(Words 0 w, xs) `c` w8
			= (Words 3 (w .|. sL w8 24), xs |> I0 w)
		(Words (i+1) w, xs) `c` w8
			= (Words i (w .|. sL w8 (8 * i)), xs)
		sL :: Word8 -> Int -> Word32
		w `sL` x = fromIntegral w `shiftL` x

toBlock :: Words' -> Maybe (Word8, Words')
toBlock (W (Words i0@(i+1) w) xs) = Just (extract w (8 * i0), (W (Words i w) xs))
	where	extract :: Word32 -> Int -> Word8
		extract w x = fromIntegral (w `shiftR` x)
toBlock (W (Words 0 w) (I0 x:xs)) = Just (fromIntegral w, (W (Words 3 x) xs))
toBlock _ = Nothing

type instance RepT (Array i) = L (Pair (Rep i)) :*: K0 (Pair (Rep i) (Rep i))
type instance Rep (Array i e) = RepT (Array i) (Rep e)

$(genRepT [d| 
   instance (Repr i, Ix i) => ReprT (Array i) where
	toRepTMap f arr = List [(toRep i, f a) | (i, a) <- assocs arr] :*: K0 (toRep l, toRep u)
		where (l, u) = bounds arr
	fromRepTMap f (List xs :*: K0 (l, r))
		= array (fromRep l, fromRep r) [(fromRep k, f a) | (k, a) <- xs] |])

type instance RepT Set.Set = L I0
type instance RepT (Map.Map k) = L (Pair (Rep k))
type instance Rep (Set.Set a) = L I0 (Rep a)
type instance Rep (Map.Map k a) = RepT (Map.Map k) (Rep a)

$(genRepT [d|
   instance ReprT Set.Set where
	toRepTMap f s = List (Fold.foldr (\ a xs -> I0 (f a):xs) [] s)
	fromRepTMap f (List xs) = Set.fromDistinctAscList [f x | I0 x <- xs] |])

$(genRepT [d|
   instance Repr k => ReprT (Map.Map k) where
	toRepTMap f m = List (Map.foldWithKey (\ k a xs -> (toRep k, f a):xs) [] m)
	fromRepTMap f (List xs) = Map.fromDistinctAscList [(fromRep k, f x) | (k, x) <- xs] |])

type instance RepT Rev = Rev
type instance Rep (Rev a) = Rev (Rep a)

$(genRepT [d|
   instance ReprT Rev where
   	toRepTMap f (Rev m) = Rev (f m)
	fromRepTMap f (Rev m) = Rev (f m) |])