{-# LANGUAGE GADTs, TypeFamilies, UndecidableInstances, FlexibleInstances, ScopedTypeVariables, 
    TypeSynonymInstances #-}

-- Generic programming / type families example, based on Generalized Tries
-- Adapted from Hinze, Loeh "Type-indexed Datatypes"

module FMap where

import Maybe


-- PART 1: Standard Tries, an optimized dictionary from strings to values.

-- A tree is a tree structure, where each node contains a character and 
-- perhaps a value (if the string from the root to this node is mapped).

data Trie v = T (Maybe v, [(Char, Trie v)]) deriving Show

-- basic operations

emptyTrie :: Trie v
emptyTrie = T (Nothing, [])

insertTrie :: [Char] -> v -> Trie v -> Trie v
insertTrie [] v (T (_, t)) = T (Just v, t)
insertTrie (a:as) v (T (m, ls)) = 
   let inner = fromMaybe emptyTrie (lookup a ls) in
   T (m, insertAL a (insertTrie as v inner) ls)

lookupTrie :: [Char] -> Trie v -> Maybe v
lookupTrie [] (T (m, _)) = m
lookupTrie (a:as) (T (_, t)) = (lookup a <> lookupTrie as) t

-- helper functions

-- update an association list with a new mapping, removing
-- the old one.
insertAL :: (Eq a) => a -> v -> [(a, v)] -> [(a, v)]
insertAL c v [] = [(c,v)]
insertAL c v ((c',_):m) | c == c' = (c, v):m
insertAL c v ((c',v'):m) = (c',v'):insertAL c v m

-- reverse bind
(<>) :: (a -> Maybe b) -> (b -> Maybe c) -> a -> Maybe c
(f <> g) a = case f a of { Nothing -> Nothing; Just b -> g b }



-- Example trie

alist = [("c", True), 
         ("ca", False), 
         ("cab", True),
         ("car", True),
         ("care", True), 
         ("cax", False), 
         ("carrot", True),
         ("cd", False),
         ("cdr", True)]

trie = foldr (uncurry insertTrie) emptyTrie alist










-------------------------------------------------------------
-- Part 2: How do we know that a trie is more efficient?

-- Lets compare the sizes, using generic programming 
--
-- This is a simple example of generic programming to 
-- get us warmed up. It also demonstrates (one way) to 
-- extend generic programs to arbitrary datatypes. 

-- A GADT of type representations (also called Ty a, or Type a)
data R t where
  Char   :: R Char
  Unit   :: R ()
  Bool   :: R Bool
  Either :: R a -> R b -> R (Either a b)
  Prod   :: R a -> R b -> R (a, b)
  List   :: R a -> R [a]
  Trie   :: R a -> R (Trie a)

-- We'll define the cases for List and Trie in terms of the 
-- other cases (Unit, Either, Prod, etc.) So we need some 
-- helper functions that let us convert between Lists/Tries 
-- and their structures.

--- isomorphisms for lists 

-- convert representation
listRep :: R a -> R (Either () (a, [a]))
listRep ra = Either Unit (Prod ra (List ra))

-- convert values
toList :: Either () (a, [a]) -> [a]
toList (Left _) = []
toList (Right (a,as)) = (a:as)

fromList :: [a] -> Either () (a,[a])
fromList [] = Left ()
fromList (a:as) = Right (a,as)

--- isomorphisms for tries
-- recall Trie v = T (Maybe v, [(Char, Trie v)])
trieRep :: R v -> R  ((Either () v) , ([(Char, Trie v)]))
trieRep rv = Prod (Either Unit rv) (List (Prod Char (Trie rv)))

toTrie  :: ((Either () v) , ([(Char, Trie v)])) -> Trie v
toTrie (Left _,alist) = T (Nothing, alist)
toTrie (Right v, alist) = T (Just v, alist)

fromTrie :: Trie v -> ((Either () v) , ([(Char, Trie v)]))
fromTrie (T (Nothing, alist)) = (Left (), alist)
fromTrie (T (Just v,  alist)) = (Right v, alist)
------------------------------------------------
-- Generic function for calculating the size of a data-structure
size :: R a -> a -> Int 
size Char   x = 1 
size Unit   x = 1
size Bool   x = 1
size (Either t1 t2) x = 1 + case x of 
                              Left y -> size t1 y
                              Right y -> size t2 y
size (Prod t1 t2) (x,y) = 1 + size t1 x + size t2 y  
size (List t1) x = size (listRep t1) (fromList x)
size (Trie t1) x = size (trieRep t1) (fromTrie x)
 
s1 = size (List (Prod (List Char) Bool)) alist
s2 = size (Trie Bool) trie

-------------------------------------------------------------
-- Part 3: Generlized tries

-- How to use generic programming to derive the definitions of 
-- tries above?

-- A generic trie, indexed by types

type family FMap t v
type instance FMap ()             v = Maybe v
type instance FMap Char           v = [(Char,v)]
type instance FMap (Either t1 t2) v = (FMap t1 v, FMap t2 v)
type instance FMap (t1, t2)       v = FMap t1 (FMap t2 v)
type instance FMap [a]            v = FMapList a v

data FMapList a v = FMapList (FMap (Either () (a, [a])) v)

fmlookup :: R t -> t -> FMap t v -> Maybe v
fmlookup Unit k m = m 
fmlookup Char k m = lookup k m
fmlookup (Either t1 t2) (Left k1) (m1, m2) = fmlookup t1 k1 m1 
fmlookup (Either t1 t2) (Right k2) (m1, m2) = fmlookup t2 k2 m2
fmlookup (Prod t1 t2) (k1,k2) m = do m2 <- (fmlookup t1 k1 m) 
                                     fmlookup t2 k2 m2

-- v is "type argument"
fmempty' :: forall t v. v -> R t -> FMap t v
fmempty' v Unit = Nothing
fmempty' v Char = []
fmempty' v (Either t1 t2) = (fmempty' v t1, fmempty' v t2) 
fmempty' v (Prod t1 t2) = fmempty' (fmempty' v t2)  t1

fminsert :: forall t v. R t -> t -> v -> FMap t v -> FMap t v
fminsert Unit k v _ = Just v
fminsert Char k v m = insertAL k v m
fminsert (Either t1 t2) (Left k1) v (m1, m2) = (fminsert t1 k1 v m1, m2)
fminsert (Either t1 t2) (Right k2) v (m1, m2) = (m1, fminsert t2 k2 v m2)
fminsert (Prod t1 t2) (k1,k2) v m =
    let m2 = fromMaybe (fmempty' (undefined::v) t2)  (fmlookup t1 k1 m) in
    fminsert t1 k1 (fminsert t2 k2 v m2) m    
    
-----------------------------------------------------------------------
-- Examples

fmap1 :: FMap String Bool
fmap1 =  foldr (uncurry (fminsert (List Char))) (fmempty' True (List Char)) alist

---------------------------------------------------------------
{-
Final questions (just for thinking, not HW)

1. What if we left the v out of the definition of FMap? How does that
   change the interaction with type inference?
2. There is a way to get a uniform representation of 
   datatypes that we can use a single constructor of R for both List/Trie.
   (It embeds the toList/fromList/toListRep info in that constructor)
   I can see how to make it work for generic size, but I don't know how 
   to make it work for FMap b/c of the type transformations. Are we 
   out of luck?
-}