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

module TrieMap.Algebraic (Algebraic(..), Ordered(..)) where

import Data.Char
import Data.Maybe
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

-- | '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
	-- | @'Alg' k@ is a fully decomposed representation of k into algebraic pieces.
	type Alg k
	toAlg :: k -> Alg k
	fromAlg :: Alg k -> k

instance (Algebraic k1, Algebraic k2) => Algebraic (k1, k2) where
	type Alg (k1, k2) = (Alg k1, Alg k2)
	toAlg (k1, k2) = (toAlg k1, toAlg k2)
	fromAlg (k1, k2) = (fromAlg k1, fromAlg k2)

instance (Algebraic a, Algebraic b, Algebraic c) => Algebraic (a, b, c) where
	type Alg (a, b, c) = (Alg a, (Alg b, Alg c))
	toAlg (a, b, c) = toAlg (a, (b, c))
	fromAlg x = case fromAlg x of
		(a, (b, c)) -> (a, b, c) 

instance (Algebraic a, Algebraic b, Algebraic c, Algebraic d) => Algebraic (a, b, c, d) where
	type Alg (a, b, c, d) = (Alg a, (Alg b, (Alg c, Alg d)))
	toAlg (a, b, c, d) = toAlg (a, (b, (c, d)))
	fromAlg x = case fromAlg x of
		(a, (b, (c, d))) -> (a, b, c, d)

instance (Algebraic a, Algebraic b, Algebraic c, Algebraic d, Algebraic e) => Algebraic (a, b, c, d, e) where
	type Alg (a, b, c, d, e) = (Alg a, (Alg b, (Alg c, (Alg d, Alg e))))
	toAlg (a, b, c, d, e) = toAlg (a, (b, (c, (d, e))))
	fromAlg x = case fromAlg x of
		(a, (b, (c, (d, e)))) -> (a, b, c, d, e)

instance (Algebraic k1, Algebraic k2) => Algebraic (Either k1 k2) where
	type Alg (Either k1 k2) = Either (Alg k1) (Alg k2)
	toAlg = either (Left . toAlg) (Right . toAlg)
	fromAlg = either (Left . fromAlg) (Right . fromAlg)

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

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

instance Algebraic a => Algebraic (Maybe a) where
	type Alg (Maybe a) = Either () (Alg a)
	toAlg Nothing = Left ()
	toAlg (Just a) = Right (toAlg a)
	fromAlg (Left _) = Nothing
	fromAlg (Right a) = Just (fromAlg a)

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

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

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

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

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

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

instance (Algebraic k, Algebraic v) => Algebraic (Map k v) where
	type Alg (Map k v) = [(Alg k, Alg v)]
	toAlg m = build (\ c n -> Map.foldWithKey (\ k v -> c (toAlg k, toAlg v)) n m)
	fromAlg xs = Map.fromDistinctAscList [(fromAlg k, fromAlg v) | (k, v) <- xs]

instance Algebraic v => Algebraic (IntMap v) where
	type Alg (IntMap v) = [(Int, Alg v)]
	toAlg m = build (\ c n -> IMap.foldWithKey (\ k v -> c (k, toAlg v)) n m)
	fromAlg xs = IMap.fromDistinctAscList [(k, fromAlg v) | (k, v) <- xs]

instance Algebraic a => Algebraic (Set a) where
	type Alg (Set a) = [Alg a]
	toAlg s = build (\ c n -> Fold.foldr (c . toAlg) n s)
	fromAlg = Set.fromDistinctAscList . map fromAlg

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