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

module TrieMap.Algebraic (Algebraic(..), AlgebraicT(..), SAlgebraicT(..), Ordered(..), AlgWrap (..)) where

import Control.Arrow
import Data.Bits
import Data.ByteString (ByteString, pack, unpack)
import Data.Char
import Data.Maybe
import Data.List (unfoldr)
import Data.Word
import Data.IntSet (IntSet)
import Data.Set(Set)
import qualified Data.IntSet as ISet
import qualified Data.Set as Set
import Data.IntMap (IntMap)
import Data.Map (Map)
import qualified Data.IntMap as IMap
import qualified Data.Map as Map
import qualified Data.Foldable as Fold
import GHC.Exts (build)

import TrieMap.TrieAlgebraic
import TrieMap.MapTypes

newtype AlgWrap t a = AlgWrap {unAlgWrap :: t a}

-- | 'Algebraic' refers to a type with an algebraic representation, armed with methods to convert in each direction.
-- 'toAlg' and 'fromAlg' should preserve equality and ordering.
class Algebraic k where
	-- | @'AlgRep' k@ is a fully decomposed representation of k into algebraic pieces.
	type AlgRep k
	toAlg :: k -> AlgRep k
	fromAlg :: AlgRep k -> k

class Functor (AlgRepT t) => AlgebraicT t where
	type AlgRepT t :: * -> *
	toAlgT :: t a -> AlgRepT t a
	fromAlgT :: AlgRepT t a -> t a

class Functor (SAlgRepT t) => SAlgebraicT t where
	type SAlgRepT t :: * -> *
	toSAlgT :: Sized a => t a -> SAlgRepT t a
	fromSAlgT :: Sized a => SAlgRepT t a -> t a

instance AlgebraicT Id where
	type AlgRepT Id = Id
	toAlgT = id
	fromAlgT = id

instance (AlgebraicT t, Algebraic a) => Algebraic (AlgWrap t a) where
	type AlgRep (AlgWrap t a) = AlgRepT t (AlgRep a)
	toAlg = fmap toAlg . toAlgT . unAlgWrap
	fromAlg = AlgWrap . fromAlgT . fmap fromAlg

instance (AlgebraicT f, AlgebraicT g) => AlgebraicT (f `O` g) where
	type AlgRepT (f `O` g) = AlgRepT f `O` AlgRepT g
	toAlgT (O x) = O (fmap (\ (App y) -> App (toAlgT y)) (toAlgT x))
	fromAlgT (O x) = O (fromAlgT (fmap (\ (App y) -> App (fromAlgT y)) x))

instance (Algebraic (f (g a)), Functor f) => Algebraic ((f `O` g) a) where
	type AlgRep ((f `O` g) a) = AlgRep (f (g a))
	toAlg = toAlg . unO
	fromAlg = o . fromAlg

instance (AlgebraicT f, AlgebraicT g) => AlgebraicT (f :*: g) where
	type AlgRepT (f :*: g) = AlgRepT f :*: AlgRepT g
	toAlgT (a :*: b) = toAlgT a :*: toAlgT b
	fromAlgT (a :*: b) = fromAlgT a :*: fromAlgT b

instance (AlgebraicT f, AlgebraicT g, Algebraic a) => Algebraic ((f :*: g) a) where
	type AlgRep ((f :*: g) a) = (AlgRepT f :*: AlgRepT g) (AlgRep a)
	toAlg (a :*: b) = fmap toAlg (toAlgT a :*: toAlgT b)
	fromAlg (a :*: b) = fromAlgT (fmap fromAlg a) :*: fromAlgT (fmap fromAlg b)

instance (AlgebraicT f, AlgebraicT g) => AlgebraicT (f :+: g) where
	type AlgRepT (f :+: g) = AlgRepT f :+: AlgRepT g
	toAlgT (A a) = A (toAlgT a)
	toAlgT (B b) = B (toAlgT b)
	fromAlgT (A a) = A (fromAlgT a)
	fromAlgT (B b) = B (fromAlgT b)

instance (AlgebraicT f, AlgebraicT g, Algebraic a) => Algebraic ((f :+: g) a) where
	type AlgRep ((f :+: g) a) = AlgRep (AlgWrap (f :+: g) a)
	toAlg = toAlg . AlgWrap
	fromAlg = unAlgWrap . fromAlg

instance AlgebraicT f => Algebraic (Fix f) where
	type AlgRep (Fix f) = Fix (AlgRepT f)
	toAlg (Fix x) = Fix (fmap toAlg (toAlgT x))
	fromAlg (Fix x) = Fix (fromAlgT (fmap fromAlg x))

instance Algebraic a => AlgebraicT (Const a) where
	type AlgRepT (Const a) = Const (AlgRep a)
	toAlgT (Const a) = Const (toAlg a)
	fromAlgT (Const a) = Const (fromAlg a)

instance Algebraic a => Algebraic (Const a b) where
	type AlgRep (Const a b) = Const (AlgRep a) b
	toAlg (Const a) = Const (toAlg a)
	fromAlg (Const a) = fromAlg (Const a)

instance Algebraic a => AlgebraicT ((,) a) where
	type AlgRepT ((,) a) = (,) (AlgRep a)
	toAlgT = first toAlg
	fromAlgT = first fromAlg

instance (Algebraic a, Algebraic b) => Algebraic (a, b) where
	type AlgRep (a, b) = AlgRep (AlgWrap ((,) a) b)
	toAlg = toAlg . AlgWrap
	fromAlg = unAlgWrap . fromAlg

instance (Algebraic a, Algebraic b) => AlgebraicT ((,,) a b) where
	type AlgRepT ((,,) a b) = (,) (AlgRep (a, b))
	toAlgT (a, b, c) = (toAlg (a, b), c)
	fromAlgT (ab, c) = case fromAlg ab of
		(a, b)	-> (a, b, c)

instance (Algebraic a, Algebraic b, Algebraic c) => Algebraic (a, b, c) where
	type AlgRep (a, b, c) = AlgRep (AlgWrap ((,,) a b) c)
	toAlg = toAlg . AlgWrap
	fromAlg = unAlgWrap . fromAlg

instance (Algebraic a, Algebraic b, Algebraic c) => AlgebraicT ((,,,) a b c) where
	type AlgRepT ((,,,) a b c) = (,) (AlgRep (a, b, c))
	toAlgT (a, b, c, d) = (toAlg (a, b, c), d)
	fromAlgT (abc, d) = case fromAlg abc of
		(a, b, c) -> (a, b, c, d)

instance (Algebraic a, Algebraic b, Algebraic c, Algebraic d) => Algebraic (a, b, c, d) where
	type AlgRep (a, b, c, d) = AlgRep (AlgWrap ((,,,) a b c) d)
	toAlg = toAlg . AlgWrap
	fromAlg = unAlgWrap . fromAlg

instance Algebraic a => AlgebraicT (Either a) where
	type AlgRepT (Either a) = Either (AlgRep a)
	toAlgT = either (Left . toAlg) Right
	fromAlgT = either (Left . fromAlg) Right

instance (Algebraic a, Algebraic b) => Algebraic (Either a b) where
	type AlgRep (Either a b) = AlgRep (AlgWrap (Either a) b)
	toAlg = toAlg . AlgWrap
	fromAlg = unAlgWrap . fromAlg

instance AlgebraicT [] where
	type AlgRepT [] = []
	toAlgT = id
	fromAlgT = id

instance Algebraic k => Algebraic [k] where
	type AlgRep [k] = [AlgRep k]
	toAlg = map toAlg
	fromAlg = map fromAlg

instance Algebraic () where
	type AlgRep () = ()
	toAlg = id
	fromAlg = id

instance AlgebraicT Maybe where
	type AlgRepT Maybe = Either ()
	toAlgT = maybe (Left ()) Right
	fromAlgT = either (const Nothing) Just

instance SAlgebraicT Maybe where
	type SAlgRepT Maybe = AlgRepT Maybe
	toSAlgT = toAlgT 
	fromSAlgT = fromAlgT

instance Algebraic a => Algebraic (Maybe a) where
	type AlgRep (Maybe a) = AlgRep (AlgWrap Maybe a)
	toAlg = toAlg . AlgWrap
	fromAlg = unAlgWrap . fromAlg

instance Algebraic Bool where
	type AlgRep Bool = AlgRep (Maybe ())
	toAlg b = toAlg $ if b then Just () else Nothing
	fromAlg = maybe False (const True) . fromAlg'
		where	fromAlg' = fromAlg :: AlgRep (Maybe ()) -> Maybe ()

instance Algebraic Int where
	type AlgRep Int = Int
	toAlg = id
	fromAlg = id

instance Algebraic Char where
	type AlgRep Char = Int
	toAlg = ord
	fromAlg = chr

instance Algebraic Float where
	type AlgRep Float = Ordered Float
	toAlg = Ord
	fromAlg = unOrd

instance Algebraic Double where
	type AlgRep Double = Ordered Double
	toAlg = Ord
	fromAlg = unOrd

instance Algebraic Rational where
	type AlgRep Rational = Ordered Rational
	toAlg = Ord
	fromAlg = unOrd

instance Algebraic a => Algebraic (Ordered a) where
	type AlgRep (Ordered a) = AlgRep a
	toAlg = toAlg . unOrd
	fromAlg = Ord . fromAlg

instance (Algebraic k, Algebraic v) => Algebraic (Map k v) where
	type AlgRep (Map k v) = AlgRep (AlgWrap (Map k) v) 
	toAlg = toAlg . AlgWrap
	fromAlg = unAlgWrap . fromAlg

instance Algebraic k => AlgebraicT (Map k) where
	type AlgRepT (Map k) = [] `O` ((,) k)
	toAlgT = o . Map.assocs
	fromAlgT = Map.fromDistinctAscList . unO

instance Algebraic k => SAlgebraicT (Map k) where
	type SAlgRepT (Map k) = [] `O` ((,) k)
	toSAlgT = o . Map.assocs
	fromSAlgT = Map.fromDistinctAscList . unO

instance Algebraic v => Algebraic (IntMap v) where
	type AlgRep (IntMap v) = AlgRep (AlgWrap IntMap v)
	toAlg = toAlg . AlgWrap
	fromAlg = unAlgWrap . fromAlg

instance AlgebraicT IntMap where
	type AlgRepT IntMap = AlgRepT ([] `O` ((,) Int))
	toAlgT = toAlgT . o . IMap.assocs
	fromAlgT = IMap.fromDistinctAscList . unO . fromAlgT

instance SAlgebraicT IntMap where
	type SAlgRepT IntMap = AlgRepT ([] `O` ((,) Int))
	toSAlgT = toAlgT . o . IMap.assocs
	fromSAlgT = IMap.fromDistinctAscList . unO . fromAlgT

instance Algebraic a => Algebraic (Set a) where
	type AlgRep (Set a) = AlgRep (AlgWrap Set a)
	toAlg = toAlg . AlgWrap
	fromAlg = unAlgWrap . fromAlg 

instance AlgebraicT Set where
	type AlgRepT Set = AlgRepT []
	toAlgT = toAlgT . Fold.toList
	fromAlgT = Set.fromDistinctAscList . fromAlgT

instance Algebraic IntSet where
	type AlgRep IntSet = AlgRep [Int]
	toAlg = toAlg . ISet.toList
	fromAlg = ISet.fromDistinctAscList . fromAlg

{-# RULES
	"map/id" forall xs . map id xs = xs;
	#-}

instance SAlgebraicT m => SAlgebraicT (ConstMap m k m') where
	type SAlgRepT (ConstMap m k m') = SAlgRepT m
	toSAlgT (ConstMap m) = toSAlgT m
	fromSAlgT = ConstMap . fromSAlgT

instance Algebraic (m a) => Algebraic (ConstMap m k m' a) where
	type AlgRep (ConstMap m k m' a) = AlgRep (m a)
	toAlg (ConstMap m) = toAlg m
	fromAlg = ConstMap . fromAlg

instance SAlgebraicT m => SAlgebraicT (IdMap k m) where
	type SAlgRepT (IdMap k m) = SAlgRepT m
	toSAlgT (IdMap m) = toSAlgT m
	fromSAlgT = IdMap . fromSAlgT

instance Algebraic (m a) => Algebraic (IdMap k m a) where
	type AlgRep (IdMap k m a) = AlgRep (m a)
	toAlg (IdMap m) = toAlg m
	fromAlg = IdMap . fromAlg

instance (SAlgebraicT (t1 k m), SAlgebraicT (t2 k m), TrieKey k m, TrieKeyT f2 t2) => SAlgebraicT (ProdMap t1 t2 k m) where
	type SAlgRepT (ProdMap t1 t2 k m) = (SAlgRepT (t1 k m) `O` SAlgRepT (t2 k m))
	toSAlgT (PMap m) = O (fmap (App . toSAlgT) (toSAlgT m))
	fromSAlgT (O m) = PMap (fromSAlgT (fmap (\ (App x) -> fromSAlgT x) m))

instance Algebraic (t1 k m (t2 k m a)) => Algebraic (ProdMap t1 t2 k m a) where
	type AlgRep (ProdMap t1 t2 k m a) = AlgRep (t1 k m (t2 k m a))
	toAlg (PMap m) = toAlg m
	fromAlg = PMap . fromAlg

instance (SAlgebraicT m1, SAlgebraicT m2, TrieKey k2 m2) => SAlgebraicT (CProdMap m1 k2 m2) where
	type SAlgRepT (CProdMap m1 k2 m2) = SAlgRepT m1 `O` SAlgRepT m2
	toSAlgT (CPMap m) = O (fmap (App . toSAlgT) (toSAlgT m))
	fromSAlgT (O m) = CPMap (fromSAlgT (fmap (fromSAlgT . unApp) m))

instance (Algebraic (m1 (m2 a))) => Algebraic (CProdMap m1 k2 m2 a) where
	type AlgRep (CProdMap m1 k2 m2 a) = AlgRep (m1 (m2 a))
	toAlg (CPMap m) = toAlg m
	fromAlg = CPMap . fromAlg

{-
instance Algebraic (t1 k m (t2 k m a)) => Algebraic (ProdMap t1 t2 k m a) where
	type AlgRep (ProdMap t1 t2 k m a) = AlgRep (t1 k m (t2 k m a))
	toAlg = toAlg . unPMap
	fromAlg = PMap . fromAlg-}

instance (SAlgebraicT (t1 k m), SAlgebraicT (t2 k m)) => SAlgebraicT (UnionMap t1 t2 k m) where
	type SAlgRepT (UnionMap t1 t2 k m) = SAlgRepT (t1 k m) :*: SAlgRepT (t2 k m)
	toSAlgT (UMap m1 m2) = toSAlgT m1 :*: toSAlgT m2
	fromSAlgT (m1 :*: m2) = UMap (fromSAlgT m1) (fromSAlgT m2)

instance (Algebraic (t1 k m a), Algebraic (t2 k m a)) => Algebraic (UnionMap t1 t2 k m a) where	
	type AlgRep (UnionMap t1 t2 k m a) = AlgRep (t1 k m a, t2 k m a)
	toAlg (UMap m1 m2) = toAlg (m1, m2)
	fromAlg = uncurry UMap . fromAlg

instance (SAlgebraicT m1, SAlgebraicT m2) => SAlgebraicT (CUnionMap m1 k2 m2) where
	type SAlgRepT (CUnionMap m1 k2 m2) = SAlgRepT m1 :*: SAlgRepT m2
	toSAlgT (CUMap m1 m2) = toSAlgT m1 :*: toSAlgT m2
	fromSAlgT (m1 :*: m2) = CUMap (fromSAlgT m1) (fromSAlgT m2)

instance (Algebraic (m1 a), Algebraic (m2 a)) => Algebraic (CUnionMap m1 k2 m2 a) where
	type AlgRep (CUnionMap m1 k2 m2 a) = AlgRep (m1 a, m2 a)
	toAlg (CUMap m1 m2) = toAlg (m1, m2)
	fromAlg = uncurry CUMap . fromAlg

-- instance (Algebraic (t1 k m a), Algebraic (t2 k m a)) => Algebraic (UnionMap t1 t2 k m a) where
-- 	type AlgRep (UnionMap t1 t2 k m a) = AlgRep (t1 k m a, t2 k m a)
-- 	toAlg (UMap m1 m2) = toAlg (m1, m2)
-- 	fromAlg = uncurry UMap . fromAlg

instance SAlgebraicT f => SAlgebraicT (App f) where
	type SAlgRepT (App f) = SAlgRepT f
	toSAlgT = toSAlgT . unApp
	fromSAlgT = App . fromSAlgT

instance AlgebraicT f => AlgebraicT (App f) where
	type AlgRepT (App f) = AlgRepT f
	toAlgT = toAlgT . unApp
	fromAlgT = App . fromAlgT

instance Algebraic (f a) => Algebraic (App f a) where
	type AlgRep (App f a) = AlgRep (f a)
	toAlg = toAlg . unApp
	fromAlg = App . fromAlg

instance SAlgebraicT (t1 (App f2 k) (App (t2 k m))) => SAlgebraicT (CompMap t1 f2 t2 k m) where
	type SAlgRepT (CompMap t1 f2 t2 k m) = SAlgRepT (t1 (App f2 k) (App (t2 k m)))
	toSAlgT (CompMap m) = toSAlgT m
	fromSAlgT = CompMap . fromSAlgT

instance Algebraic (t1 (App f2 k) (App (t2 k m)) a) => Algebraic (CompMap t1 f2 t2 k m a) where
	type AlgRep (CompMap t1 f2 t2 k m a) = AlgRep (t1 (App f2 k) (App (t2 k m)) a)
	toAlg (CompMap m) = toAlg m
	fromAlg = CompMap . fromAlg

-- instance (AlgebraicT (t1 (App f2 k) (App (t2 k m))), Algebraic a) => Algebraic (CompMap t1 f2 t2 k m a) where
-- 	type AlgRep (CompMap t1 f2 t2 k m a) = AlgRep (AlgWrap (CompMap t1 f2 t2 k m) a) 
-- 	toAlg = toAlg . AlgWrap
-- 	fromAlg = unAlgWrap . fromAlg

-- newtype   f t a = FixMap (t (Fix f) (FixMap f t) a)

instance (TrieKeyT f t) => SAlgebraicT (FixMap f t) where
	type SAlgRepT (FixMap f t) = [] `O` ((,) (Fix f))
	toSAlgT m = o (assocsAlg m)
	fromSAlgT = fromDistAscListAlg . unO

instance (TrieKeyT f t, AlgebraicT f, Sized a, Algebraic a) => Algebraic (FixMap f t a) where
	type AlgRep (FixMap f t a) = AlgRep [(Fix f, a)]
	toAlg = toAlg . assocsAlg
	fromAlg = fromDistAscListAlg . fromAlg

-- instance (AlgebraicT f, TrieKeyT f t, Sized a, Algebraic a) => Algebraic (FixMap f t a) where
-- 	type AlgRep (FixMap f t a) = AlgRep [(Fix f, a)]
-- 	toAlg = toAlg . assocsAlg
-- 	fromAlg = fromDistAscListAlg . fromAlg

instance Algebraic Word8 where
	type AlgRep Word8 = Int
	toAlg = fromIntegral
	fromAlg = fromIntegral

instance Algebraic Word16 where
	type AlgRep Word16 = Int
	toAlg = fromIntegral
	fromAlg = fromIntegral

instance Algebraic Word32 where
	type AlgRep Word32 = Int
	toAlg = fromIntegral
	fromAlg = fromIntegral

instance Algebraic Integer where
	type AlgRep Integer = AlgRep [Word8]
	toAlg = toAlg . unroll
	fromAlg = roll . fromAlg

instance Algebraic ByteString where
	type AlgRep ByteString = AlgRep [Word8]
	toAlg = toAlg . unpack
	fromAlg = pack . fromAlg

unroll :: Integer -> [Word8]
unroll = unfoldr step
  where
    step 0 = Nothing
    step i = Just (fromIntegral i, i `shiftR` 8)

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


{-# RULES
	"toAlg/fromAlg" forall x . toAlg (fromAlg x) = x;
 #-}