{-# LANGUAGE FlexibleInstances, UndecidableInstances, KindSignatures, StandaloneDeriving, GeneralizedNewtypeDeriving, IncoherentInstances, TypeOperators, FlexibleContexts, StandaloneDeriving, ExistentialQuantification #-}

module TrieMap.MapTypes where

import Data.Foldable
import Data.Traversable
import Control.Applicative hiding (Const)
import Prelude hiding (foldl, foldr)
import qualified Data.IntMap as IMap

data (f :*: g) a = f a :*: g a deriving (Eq, Ord, Show)
data (f :+: g) a = A (f a) | B (g a) deriving (Eq, Ord, Show)
newtype Const a b = Const {unConst :: a} deriving (Eq, Ord, Show)
newtype Id a = Id {unId :: a} deriving (Eq, Ord, Show)
newtype Fix f = Fix (f (Fix f))
newtype FixMap f t a = FixMap (t (Fix f) (FixMap f t) a)

newtype O f g a = O (f (App g a))
newtype App f a = App {unApp :: f a}

o :: Functor f => f (g a) -> (f `O` g) a
o = O . fmap App

unO :: Functor f => (f `O` g) a -> f (g a)
unO (O x) = fmap unApp x

-- | 'ProdMap' is used to hold a map on the product of two key types.
newtype ProdMap t1 t2 k (m :: * -> *) a = PMap {unPMap :: t1 k m (t2 k m a)}
data UnionMap t1 t2 k (m :: * -> *) a = UMap (t1 k m a) (t2 k m a)
newtype CProdMap m1 k2 m2 a = CPMap {unCPMap :: m1 (m2 a)}
data CUnionMap m1 k2 m2 a = CUMap (m1 a) (m2 a)

data Edge k m v = Edge {-# UNPACK #-} !Int [k] (Maybe v) (m (Edge k m v))
type MEdge k m v = Maybe (Edge k m v)

-- | 'RadixTrie' is used to hold a map on a list of keys.
newtype RadixTrie k m v = Radix {unRad :: MEdge k m v} 

newtype IdMap k m a = IdMap {unIdMap :: m a} 

newtype ConstMap (m :: * -> *) k (x :: * -> *) a = ConstMap {unConstMap :: m a}

newtype CompMap t1 f2 (t2 :: * -> (* -> *) -> * -> *) k (m :: * -> *) a = CompMap (t1 (App f2 k) (App (t2 k m)) a)

-- newtype FixMap (m :: (* -> *) -> * -> *) a = FixMap (m (FixMap m) a)

newtype Elem a = Elem {getElem :: a} deriving (Eq, Ord)

instance Functor Elem where
	fmap f (Elem x) = Elem (f x)

instance Foldable Elem where
	foldr f z (Elem a) = a `f` z
	foldl f z (Elem a) = z `f` a

instance Traversable Elem where
	traverse f (Elem x) = Elem <$> f x

infixr 5 `ProdMap`
infixr 5 :+:
infixr 8 :*:
infixr 9 `O`

class Sized a where
	getSize :: a -> Int

instance Sized (Elem a) where
	getSize _ = 1

instance Functor m => Functor (Edge k m) where
	fmap f (Edge n ks v ts) = Edge n ks (fmap f v) (fmap (fmap f) ts)
raverse f (Radix e) = Radix <$> traverse (traverse f) e

instance (Functor f, Functor g) => Functor (f :*: g) where
	fmap f (a :*: b) = fmap f a :*: fmap f b

instance (Foldable f, Foldable g) => Foldable (f :*: g) where
	foldr f z (a :*: b) = foldr f (foldr f z b) a
	foldl f z (a :*: b) = foldl f (foldl f z a) b

instance (Traversable f, Traversable g) => Traversable (f :*: g) where
	traverse f (a :*: b) = liftA2 (:*:) (traverse f a) (traverse f b)

instance (Functor f, Functor g) => Functor (f :+: g) where
	fmap f (A a) = A (fmap f a)
	fmap f (B b) = B (fmap f b)

instance (Foldable f, Foldable g) => Foldable (f :+: g) where
	foldr f z (A a) = foldr f z a
	foldr f z (B b) = foldr f z b
	foldl f z (A a) = foldl f z a
	foldl f z (B b) = foldl f z b

instance (Traversable f, Traversable g) => Traversable (f :+: g) where
	traverse f (A a) = A <$> traverse f a
	traverse f (B b) = B <$> traverse f b

instance Functor (Const a) where
	fmap f (Const x) = Const x

instance Foldable (Const a) where
	foldr f z _ = z
	foldl f z _ = z

instance Traversable (Const a) where
	traverse f (Const x) = pure (Const x)

instance Functor Id where
	fmap f (Id a) = Id (f a)

instance Foldable Id where
	foldr f z (Id a) = a `f` z
	foldl f z (Id a) = z `f` a

instance Traversable Id where
	traverse f (Id a) = Id <$> f a

class EqT f where
	eq :: Eq a => f a -> f a -> Bool

instance EqT f => Eq (Fix f) where
	Fix x == Fix y = x `eq` y

instance (EqT f, EqT g) => EqT (f :*: g) where
	(a :*: x) `eq` (b :*: y) = a `eq` b && x `eq` y

instance (EqT f, EqT g) => EqT (f :+: g) where
	A a `eq` A b = a `eq` b
	B x `eq` B y = x `eq` y
	_ `eq` _ = False

instance Eq a => EqT (Const a) where
	Const a `eq` Const b = a == b

instance EqT Id where
	Id a `eq` Id b = a == b

instance EqT [] where
	eq = (==)

instance EqT Maybe where
	eq = (==)

instance Eq a => EqT ((,) a) where
	eq = (==)

instance Eq a => EqT (Either a) where
	eq = (==)

instance EqT f => EqT (App f) where
	App a `eq` App b = a `eq` b

instance (EqT f, Eq a) => Eq (App f a) where
	(==) = eq

instance (EqT f, EqT g) => EqT (f `O` g) where
	O a `eq` O b = a `eq` b

instance (EqT f, EqT g, Eq a) => Eq ((f `O` g) a) where
	(==) = eq

instance (Functor f, Functor g) => Functor (f `O` g) where
	fmap f (O x) = O (fmap (\ (App x) -> App (fmap f x)) x)

instance Traversable IMap.IntMap where
	traverse f m = IMap.fromDistinctAscList <$> traverse (\ (k, v) -> ((,) k) <$> f v) (IMap.assocs m)