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

-- 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 a -> R ((Either () a), [(Char, Trie a)])
trieRep a = Prod (Either Unit a) (List (Prod Char (Trie a)))

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 _ = 1
size Unit _ = 0   -- takes no space to store
size Bool _ = 1
size (Either a b) (Left x) = 1 + size a x
size (Either a b) (Right x) = 1 + size b x
size (Prod a b) (x,y) = size a x + size b y  -- don't need to store tag
size (List a) l = size (listRep a) (fromList l)
size (Trie a) t = size (trieRep a) (fromTrie t)

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 [t1] v           = FMapList t1 v
data FMapList t1 v = FMapList (FMap (Either () (t1, [t1])) v) 

instance Show (FMapList Char Bool) where
   show (FMapList m) = show m


fmlookup :: R t -> t -> FMap t v -> Maybe v
fmlookup Unit  () m = m
fmlookup Char  c m = lookup c 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 = (fmlookup t1 k1 <> fmlookup t2 k2) m
fmlookup t@(List t1) k (FMapList m) = fmlookup (listRep t1) (fromList k) m

-- for some reason, I need to add this 'v' argument to 
-- make sure that two different recursive calls produce 
-- maps with the same codomain. 
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 :: R a) (t2 :: R b)) = fmempty' (undefined :: FMap b v) t1
fmempty' v t@(List t1) = FMapList (fmempty' v (listRep t1))

fmempty = fmempty' undefined

fminsert :: R t -> t -> v -> FMap t v -> FMap t v
fminsert Unit () v m = Just v
fminsert Char c  v m = insertAL c 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' v t2) (fmlookup t1 k1 m) in
  fminsert t1 k1 (fminsert t2 k2 v m2) m
fminsert t@(List t1) k v (FMapList l) = 
    FMapList (fminsert (listRep t1) (fromList k) v l)

-----------------------------------------------------------------------
-- 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?
-}