fix subtree typing in lambda expr

fix subtree typing in lambda expr

module TypeChecker where

import Prelude (class Monad, class Eq, class Show, (&&), (==), (>>=), map, (++), ($), pure, (<*>), (<$>), return, bind, const, otherwise, show, (+), div, mod, flip)

import Control.Monad.Except.Trans (ExceptT, runExceptT, throwError)
import Control.Monad.State (State, evalState, put, get)

import Data.Either (Either(Left, Right))
import Data.List (List(..), length, delete, concat, unzip, foldM, toList, (:), zip)
import Data.Map as Map
import Data.Tuple (Tuple(Tuple), snd, fst)
import Data.Set as Set
import Data.Foldable (foldl, foldr)
import Data.Maybe (Maybe(..), fromMaybe)

import Data.Tuple (Tuple(..),fst,snd)
import Data.Map (Map(..), insert, lookup, empty)
import Data.List.WordsLines (fromCharList)
import Data.List (List(Nil, Cons), singleton, fromList, toList, length, (..), zipWithA,concat,zip, (!!))
import Data.String (joinWith, toCharArray)

import Control.Apply ((*>))
import Control.Bind ((=<<), (>=>))
import Control.Monad.State

import AST (AD(..), Atom(..), Binding(..), Definition(..), Expr(..), Op(..), TVar(..), Type(..), TypeBinding(..), TypeTree(..),TypeError(..),Path(..))

data Scheme = Forall (List TVar) Type

instance showScheme :: Show Scheme where
  show (Forall a b) = "Forall (" ++ show a ++ ") " ++ show b

data TypeEnv = TypeEnv (Map.Map Atom Scheme)

instance showTypeEnv :: Show TypeEnv where
  show (TypeEnv a) = show a

data Unique = Unique { count :: Int }

type Infer a = ExceptT TypeError (State Unique) a

data InferState = InferState { count  :: Int }

data Env = TypEnv {types:: Map.Map Atom Scheme}

type Constraint = Tuple Type Type
type Subst = Map.Map TVar Type

class Substitutable a where
  apply :: Subst -> a -> a
  ftv   :: a -> Set.Set TVar

instance subScheme :: Substitutable Scheme where
    apply s (Forall as t)   = Forall as $ apply s' t
                              where s' = foldr Map.delete s as
    ftv (Forall as t) = (ftv t) `Set.difference` Set.fromList as

instance subType :: Substitutable Type where
   apply _ (TypCon a) = TypCon a
   apply s t@(TypVar a) = fromMaybe t $ Map.lookup a s
   apply s (TypArr t1 t2) =  TypArr (apply s t1) (apply s t2)
   apply s (AD a) = AD (apply s a)

   ftv (TypCon  _)         = Set.empty
   ftv (TypVar a)       = Set.singleton a
   ftv (TypArr t1 t2) =  Set.union (ftv t1) (ftv t2)
   ftv (AD a) = ftv a

instance listSub :: (Substitutable a) => Substitutable (List a) where
  apply s xs = map (apply s) xs
  ftv   = foldr (\a b -> Set.union (ftv a) b) Set.empty

instance subTypeEnv :: Substitutable TypeEnv where
  apply s (TypeEnv env) =  TypeEnv $ map (apply s) env
  ftv (TypeEnv env) = ftv $ snd $ unzip $ Map.toList env

instance subAD :: Substitutable AD where
  apply s (TList t) = TList (apply s t)
  apply s(TTuple t) = TTuple (apply s t)

  ftv (TList t) = ftv t
  ftv (TTuple t) = ftv t

instance subTupAtomScheme :: Substitutable (Tuple Atom Scheme) where
  apply s (Tuple a b) = Tuple a (apply s b)

  ftv (Tuple _ b) = ftv b

instance subTypeTree :: Substitutable TypeTree where
  apply s (TAtom t) = TAtom $ apply s t
  apply s (TListTree l t) = TListTree (apply s l) (apply s t)
  apply s (TNTuple l t) = TNTuple (apply s l) (apply s t)
  apply s (TBinary op tt1 tt2 t) =  TBinary (apply s op) (apply s tt1) (apply s tt2) (apply s t)
  apply s (TUnary op tt t) = TUnary (apply s op) (apply s tt) (apply s t)
  apply s (TSectL tt op t) = TSectL (apply s tt) (apply s op) (apply s t)
  apply s (TSectR op tt t) = TSectR (apply s op) (apply s tt) (apply s t)
  apply s (TPrefixOp t) = TPrefixOp $ apply s t
  apply s (TIfExpr tt1 tt2 tt3 t) = TIfExpr (apply s tt1) (apply s tt2) (apply s tt3) (apply s t)
  apply s (TLetExpr b tt1 tt2 t) = TLetExpr (apply s b) (apply s tt1) (apply s tt2) (apply s t)
  apply s (TLambda lb tt t) = TLambda (apply s lb) (apply s tt) (apply s t)
  apply s (TApp tt1 l t) = TApp (apply s tt1) (apply s l) (apply s t)

  ftv (TAtom t)  = ftv t
  ftv (TListTree _ t)  = ftv t
  ftv (TNTuple _ t)  = ftv t
  ftv (TBinary _ _ _ t)  = ftv t
  ftv (TUnary _ _ t)  = ftv t
  ftv (TSectL _ _ t)  = ftv t
  ftv (TSectR _ _ t)  = ftv t
  ftv (TPrefixOp t)  = ftv t
  ftv (TIfExpr _ _ _ t)  = ftv t
  ftv (TLetExpr _ _ _ t)  = ftv t
  ftv (TLambda _ _ t)  = ftv t
  ftv (TApp _ _ t)  = ftv t

instance subTypeBinding :: Substitutable TypeBinding where
  apply s (TLit t) = TLit $ apply s t
  apply s (TConsLit b1 b2 t) = TConsLit (apply s b1) (apply s b2) (apply s t)
  apply s (TListLit lb t) = TListLit (apply s lb) (apply s t)
  apply s (TNTupleLit lb t) = TListLit (apply s lb) (apply s t)

  ftv (TLit t) = ftv t
  ftv (TConsLit _ _ t) = ftv t
  ftv (TListLit _ t) = ftv t
  ftv (TNTupleLit _ t) = ftv t

initInfer :: InferState
initInfer = InferState { count : 0}

initUnique :: Unique
initUnique = Unique { count : 0 }

extend :: TypeEnv -> Tuple Atom Scheme -> TypeEnv
extend (TypeEnv env) (Tuple x s) = TypeEnv $ Map.insert x s env

nullSubst:: Subst
nullSubst = Map.empty

fresh :: Infer Type
fresh = do
  (Unique s) <- get
  put (Unique {count:(s.count+1)})
  return $ TypVar $ TVar $ "t_" ++ show s.count

unify ::  Type -> Type -> Infer Subst
unify (TypArr l r) (TypArr l' r')  = do
    s1 <- unify l l'
    s2 <- unify (apply s1 r) (apply s1 r')
    return (s2 `compose` s1)

unify (TypVar a) t = bindVar a t
unify t (TypVar a) = bindVar a t
unify (TypCon a) (TypCon b) | a == b = return nullSubst
unify (AD a) (AD b) | a == b = return nullSubst
unify (AD a) (AD b) = unifyAD a b
unify t1 t2 = throwError $ UnificationFail t1 t2

unifyAD :: AD -> AD -> Infer Subst
unifyAD (TList a) (TList b) = unify a b
unifyAD (TTuple (Cons a as)) (TTuple (Cons b bs)) = do
  s1 <- unifyAD (TTuple as) (TTuple bs)
  s2 <- unify a b
  return (s1 `compose` s2)
unifyAD (TTuple Nil) (TTuple Nil) = return nullSubst

unifyAD a b = throwError $ UnificationFail (AD a) (AD b)

bindVar ::  TVar -> Type -> Infer Subst
bindVar a t | t == (TypVar a)     = return nullSubst
         | occursCheck a t = throwError $ InfiniteType a t
         | otherwise       = return $ Map.singleton a t

occursCheck :: forall a.  (Substitutable a) => TVar -> a -> Boolean
occursCheck a t = a `Set.member` ftv t

compose :: Subst -> Subst -> Subst
compose s1 s2 = map (apply s1) s2 `Map.union` s1

instantiate ::  Scheme -> Infer Type
instantiate (Forall as t) = do
  as' <- mapM (const fresh) as
  let s = Map.fromList $ zip as as'
  return $ apply s t

generalize :: TypeEnv -> Type -> Scheme
generalize env t  = Forall as t
    where as = Set.toList $ ftv t `Set.difference` ftv env

runInfer :: Infer (Tuple Subst Type) -> Either TypeError Scheme
runInfer m = case evalState (runExceptT m) initUnique of
  Left err -> Left err
  Right res -> Right $ closeOver res

closeOver :: (Tuple (Map.Map TVar Type) Type) -> Scheme
closeOver (Tuple sub ty) = sc --TODO add normalize back in
  where sc = generalize emptyTyenv (apply sub ty)

emptyTyenv :: TypeEnv
emptyTyenv = TypeEnv Map.empty

-- TODO find Purescript equivalent
mapM :: forall a b m. (Monad m) => (a -> m b) -> List a -> m (List b)
mapM f as = foldr k (return Nil) as
            where
              k a r = do
                x <- f a
                xs <- r
                return (x:xs)

lookupEnv :: TypeEnv -> Atom -> Infer (Tuple Subst Type)
lookupEnv (TypeEnv env) x = do
  case Map.lookup x env of
    Nothing -> throwError $ UnboundVariable (f x)
    Just s  -> do t <- instantiate s
                  return (Tuple nullSubst t)
  where
    f (Name x) = x
    f x = show x

extractType:: TypeTree -> Type
extractType (TAtom t)  =  t
extractType (TListTree _ t)  =  t
extractType (TNTuple _ t)  =  t
extractType (TBinary _ _ _ t)  =  t
extractType (TUnary _ _ t)  =  t
extractType (TSectL _ _ t)  =  t
extractType (TSectR _ _ t)  =  t
extractType (TPrefixOp t)  =  t
extractType (TIfExpr _ _ _ t)  =  t
extractType (TLetExpr _ _ _ t)  =  t
extractType (TLambda _ _ t)  =  t
extractType (TApp tt1 _ t)  =  t

extractBindingType:: TypeBinding -> Type
extractBindingType (TLit t) = t
extractBindingType (TConsLit _ _ t) = t
extractBindingType (TListLit _ t) = t
extractBindingType (TNTupleLit _ t) = t

inferType :: TypeEnv -> Expr -> Infer (Tuple Subst Type)
inferType env exp = do
  Tuple s t <- infer env exp
  let t' = extractType t
  return $ Tuple s t'

infer :: TypeEnv -> Expr -> Infer (Tuple Subst TypeTree)
infer env ex = case ex of

  (Atom x@(Name _)) -> do
    Tuple s t <- lookupEnv env x
    return (Tuple s $ TAtom t)
  (Atom (Bool _)) -> return (Tuple nullSubst $ TAtom (TypCon "Bool"))
  (Atom (Char _)) -> return (Tuple nullSubst $ TAtom (TypCon "Char"))
  (Atom (AInt _)) -> return (Tuple nullSubst $ TAtom (TypCon "Int"))

  Lambda (Cons bin Nil) e -> do
    tv <- fresh
    Tuple envL tvB <- extractBinding bin
    let env' = foldr (\a b -> extend b a) env envL
    s0 <- unify tv (extractBindingType tvB)
    (Tuple s1 t1) <- infer env' e
    return (Tuple s1 $ apply (s0 `compose` s1) (TLambda (Cons tvB Nil) t1 $ tv `TypArr` (extractType t1)))

  Lambda (Cons bin xs) e -> do
    tv <- fresh
    Tuple envL tvB <- extractBinding bin
    let env' = foldr (\a b -> extend b a) env envL
    s0 <- unify tv (extractBindingType tvB)
    (Tuple s1 (TLambda tb tt t1)) <- infer env' (Lambda xs e)
    return (Tuple s1 $ apply (s0 `compose` s1) (TLambda (Cons tvB tb) tt (tv `TypArr` t1)))

  Lambda Nil e -> infer env e

    -- one element list
  App e1 (Cons e2 Nil) -> do
    tv <- fresh
    (Tuple s1 t1) <- infer env e1
    (Tuple s2 t2) <- infer (apply s1 env) e2
    s3       <- unify (apply (s1  `compose` s2) (extractType t1)) (TypArr (extractType t2) tv)
    return (Tuple (s3 `compose` s2 `compose` s1) (apply s3 $ TApp t1 (Cons t2 Nil) tv))

  App e1 (Cons e2 xs) -> do
    Tuple s (TApp (TApp tt lt _) lt' t') <- infer env  (App (App e1 (Cons e2 Nil)) xs)
    return $ Tuple s (TApp tt (lt++lt') t')

  App _ Nil -> throwError $ UnknownError "congrats you found a bug TypeChecker.infer (App Nil)"

  LetExpr bin e1 e2 -> do
    (Tuple s1 t1) <- infer env e1
    (Tuple list typ) <- extractBinding bin
    s2 <- unify (extractBindingType typ) (extractType t1)
    let env' = apply (s1 `compose` s2) env
        t'   = generalize env' (apply (s1 `compose` s2) (extractType t1))
        env'' = apply s2 (foldr (\a b -> extend b a) env' list)
    (Tuple s3 t2) <- infer env'' e2
    let sC = (s1 `compose` s2 `compose` s3)
    return (Tuple sC $ apply sC (TLetExpr typ t1 t2 (extractType t2)))

  IfExpr cond tr fl -> do
    tv <- fresh
    t@(TypVar t') <- fresh
    let name = Name $ "if"
    let env' = env `extend` (Tuple name (Forall (Cons t' Nil) (TypArr (TypCon "Bool") (TypArr t (TypArr t  t)))))
    Tuple s (TApp tt (Cons tcond (Cons ttr (Cons tfl Nil))) ift) <- infer env' (App (Atom  name) (toList [cond, tr, fl]))
    return (Tuple s $ apply s (TIfExpr tcond ttr tfl ift))

  PrefixOp op -> do
    Tuple s t <- inferOp op
    return (Tuple s $ TPrefixOp t)

  SectL e op -> do
    tv <- fresh
    (Tuple s1 t1) <- inferOp op
    (Tuple s2 t2) <- infer (apply s1 env) e
    s3       <- unify (apply s2 t1) (TypArr (extractType t2) tv)
    return (Tuple (s3 `compose` s2 `compose` s1) (apply s3 (TSectL t2 t1 tv)))

  SectR op e -> do
    (Tuple s1 t1@(TypArr a (TypArr b c))) <- inferOp op
    (Tuple s2 t2) <- infer env e
    s3       <- unify (apply s2 b) (extractType t2)
    let s4 = (s3 `compose` s2 `compose` s1)
    return (Tuple s4 (apply s4 (TSectR t1 t2 (TypArr a c))))

  Unary op e -> do
    tv <- fresh
    (Tuple s1 t1) <- inferOp op
    (Tuple s2 t2) <- infer (apply s1 env) e
    s3       <- unify (apply s2 t1) (TypArr (extractType t2) tv)
    return (Tuple (s3 `compose` s2 `compose` s1) (apply s3 (TUnary t1 t2 tv)))

  Binary op e1 e2 -> do
    (Tuple s (TApp tt (Cons tt1 (Cons tt2 Nil)) t)) <- infer env (App (PrefixOp op) (Cons e1 (Cons e2 Nil)))
    return $ Tuple s (TBinary (extractType tt) tt1 tt2 t)

  List (Cons e1 xs) -> do
    (Tuple s1 (TListTree ltt (AD (TList t1)))) <- infer env (List xs)
    (Tuple s2 t2) <- infer (apply s1 env) e1
    s3 <- unify (apply s2 t1) (extractType t2)
    return (Tuple (s3 `compose` s2 `compose` s1) (apply s3 (TListTree (Cons t2 ltt) (AD $ TList (extractType t2)))))

  List Nil -> do
    tv <- fresh
    return (Tuple nullSubst $ TListTree Nil (AD $ TList tv))

  NTuple (Cons e Nil) -> do
    (Tuple s t) <- infer env e
    return $ Tuple s $ TNTuple (Cons t Nil) (AD $ TTuple $ Cons (extractType t) Nil)

  NTuple (Cons e1 xs) -> do
    (Tuple s1 (TNTuple lt (AD (TTuple t1)))) <- infer env (NTuple xs)
    (Tuple s2 t2) <- infer (apply s1 env) e1
    return (Tuple (s2 `compose` s1) $ TNTuple (Cons t2 lt) (AD $ TTuple (Cons (extractType t2) t1)))

  NTuple Nil -> throwError $ UnknownError "congrats you found a bug in TypeChecker.infer (NTuple Nil)"

inferOp :: Op -> Infer (Tuple Subst Type)
inferOp op = do
  a <- fresh
  b <- fresh
  c <- fresh
  case op of
    Composition ->
      f ((b `TypArr` c) `TypArr` ((a `TypArr` b) `TypArr` (a `TypArr` c)))
    Power -> int3
    Mul ->  int3
    Div -> int3
    Mod -> int3
    Add -> int3
    Sub -> int3
    Colon -> f (a `TypArr` ((AD $ TList a) `TypArr` (AD $ TList a)))
    Append -> f ((AD $ TList a) `TypArr` ((AD $ TList a) `TypArr` (AD $ TList a)))
    Equ -> aBool a
    Neq -> aBool a
    Lt -> aBool a
    Leq -> aBool a
    Gt -> aBool a
    Geq -> aBool a
    And -> f (TypCon "Bool" `TypArr` (TypCon "Bool" `TypArr` TypCon "Bool"))
    Or -> f (TypCon "Bool" `TypArr` (TypCon "Bool" `TypArr` TypCon "Bool"))
    Dollar -> f ((a `TypArr` b) `TypArr` (a `TypArr` b))
 where
  f typ = return (Tuple nullSubst typ)
  int3 = f (TypCon "Int" `TypArr` (TypCon "Int" `TypArr` TypCon "Int"))
  aBool a = f (a `TypArr` (a `TypArr` TypCon "Bool"))

inferDef :: TypeEnv -> Definition -> Infer (Tuple Subst Type)
inferDef env (Def str bin exp) = do
    tv <- fresh
    let env' = env `extend` (Tuple (Name str) (Forall Nil tv))
    let exp' = Lambda bin exp
    Tuple s1 t1' <- infer env' exp'
    let t1 = extractType t1'
    let env'' = env `extend` (Tuple (Name str) (Forall Nil (apply s1 t1)))
    Tuple s2 t2 <- infer env'' exp'
    s3 <- unify (apply s1 t1) (apply s2 (extractType t2))
    return $ Tuple (s3 `compose` s1) (apply s3 (apply s1 t1))

extractConsLit:: Type -> Binding -> Infer (Tuple (List (Tuple Atom Scheme)) TypeBinding)
extractConsLit tv (ConsLit a b@(ConsLit _ _)) = do
  Tuple list1 typ1 <- extractBinding a
  s1 <- unify (extractBindingType typ1) tv
  let list1' = apply s1 list1
  Tuple listB b@(TConsLit b1 b2 (AD ( TList typB))) <- extractConsLit tv b
  s2 <- unify typB (apply s1 tv)
  return $ Tuple (apply s2 (list1' ++ listB)) $ apply s2 (TConsLit typ1 b (AD $ TList (apply s2 typB)))

extractConsLit tv (ConsLit a (Lit b)) = do
  Tuple list typ <- extractBinding a
  s1 <- unify tv (extractBindingType typ)
  let list' = apply s1 list
  let ltyp = AD $ TList (apply s1 tv)
  return $ Tuple (Cons (Tuple b $ Forall Nil ltyp) list') $ (apply s1 (TConsLit typ (TLit ltyp) ltyp))

extractConsLit tv (ConsLit a b) = do
  Tuple list typ <- extractBinding a
  Tuple list2 typ2 <- extractBinding b
  s1 <- unify (extractBindingType typ2) (extractBindingType typ)
  s2 <- unify (apply s1 (extractBindingType typ)) tv
  let sC =  s2 `compose` s1
  let list' = map (\(Tuple a b) -> Tuple a $ apply sC b) list
  let list2' = map (\ (Tuple a b) -> Tuple a $  apply sC b) list2
  return $ Tuple (list' ++ list2') $ apply sC $ TConsLit typ typ2 $ AD $ TList tv

extractConsLit _ _ = throwError $ UnknownError "congrats you found a bug in TypeChecker.extractConsLit"

extractListLit :: List Binding -> Infer (List (Tuple (List (Tuple Atom Scheme)) TypeBinding))
extractListLit (Cons a Nil) = do
  b1 <- extractBinding a
  return (Cons b1 Nil)

extractListLit (Cons a b) = do
  b1 <- extractBinding a
  bs <- extractListLit b
  return (Cons b1 bs)

extractListLit Nil = return Nil

extractBinding:: Binding -> Infer (Tuple (List (Tuple Atom Scheme)) TypeBinding)
extractBinding (Lit a@(Name _)) = do
  tv <- fresh
  return $ Tuple (toList [(Tuple a (Forall Nil tv))]) (TLit tv)
extractBinding (Lit (Bool _)) = return $ Tuple Nil $ TLit (TypCon "Bool")
extractBinding (Lit (Char _)) = return $ Tuple Nil $ TLit (TypCon "Char")
extractBinding (Lit (AInt _)) = return $ Tuple Nil $ TLit (TypCon "Int")

extractBinding a@(ConsLit _ _) = do
  tv <- fresh
  extractConsLit tv a

extractBinding (ListLit a) = do -- Tuple TypEnv  (TListLit (List TypeBinding) Type)
  bs <- extractListLit a
  tv <- fresh
  let ini = Tuple Nil (TListLit Nil tv)
  Tuple list (TListLit lb typ) <- foldM f ini bs
  return $ Tuple list (TListLit lb (AD $ TList typ))
  where
-- :: Tuple (List (Tuple Atom Scheme)) BindingType -> Tuple (List (Tuple Atom Scheme)) BindingType
--           -> Infer (Tuple (List (Tuple Atom Scheme)) BindingType)
    f (Tuple l1 (TListLit lb1 t1)) (Tuple l2 b) = do
      let ls = l1 ++ l2
      s1 <- unify t1 (extractBindingType b)
      return $ Tuple (apply s1 ls) (apply s1 (TListLit (Cons b lb1) t1))
    f _ _ = throwError $ UnknownError "congrats you found a bug in TypeChecker.extractBinding"

extractBinding (NTupleLit a) = do
  bs <- extractListLit a
  let tup = unzip bs
  return $ Tuple (concat $ fst tup) (TNTupleLit (snd tup) (AD $ TTuple $ (map extractBindingType (snd tup))))

getTypEnv:: Binding -> TypeEnv -> Maybe TypeEnv
getTypEnv b  env= case evalState (runExceptT (extractBinding b)) initUnique of
    Left _ -> Nothing
    Right (Tuple bs _) -> Just $ foldr (\a b -> extend b a) env bs

getTypEnvFromList::  List Binding -> TypeEnv-> Maybe TypeEnv
getTypEnvFromList bs env = do
                  mTypList <-  mapM (flip getTypEnv emptyTyenv) bs
                  return $ foldr (\a b -> unionTypeEnv a b) env mTypList

unionTypeEnv :: TypeEnv -> TypeEnv -> TypeEnv
unionTypeEnv (TypeEnv a) (TypeEnv b) = TypeEnv (Map.union a b)

inferGroup:: TypeEnv -> List Definition -> Infer (Tuple Subst Type)
inferGroup _ Nil = throwError $ UnknownError "Cant type empty Group"
inferGroup env (Cons def1 Nil) = inferDef env def1
inferGroup env (Cons def1 defs) = do
  Tuple s1 t1 <- inferDef env def1
  Tuple s2 t2 <- inferGroup env defs
  s3 <- unify t1 t2
  return $ Tuple s3 (apply s3 t1)

buildTypeEnv:: List Definition -> Either TypeError TypeEnv
buildTypeEnv Nil = Right emptyTyenv
buildTypeEnv defs = buildTypeEnvFromGroups emptyTyenv groupMap keys
  where
    groupMap = buildGroups defs
    keys = Map.keys groupMap

buildGroups:: List Definition -> Map.Map String (List Definition)
buildGroups Nil = Map.empty
buildGroups (Cons def@(Def str bin exp) Nil) =
  Map.singleton str (Cons def Nil)

buildGroups (Cons def@(Def str bin exp) defs) =
  case binList of
    Just list -> Map.insert str (Cons def list) defMap
    Nothing -> Map.insert str (Cons def Nil) defMap
  where
    defMap = buildGroups defs
    binList = Map.lookup str defMap
--                         old -> grouped Definitions -> keys -> new
buildTypeEnvFromGroups:: TypeEnv -> Map.Map String (List Definition)
  -> List String -> Either TypeError TypeEnv
buildTypeEnvFromGroups env _ Nil = Right env
buildTypeEnvFromGroups env groupMap k@(Cons key keys) =
  case mayDefs of
    Nothing -> Left $ UnboundVariable key
    Just defs -> case runInfer $ inferGroup env defs of
      Right scheme -> buildTypeEnvFromGroups
        (env `extend` (Tuple (Name key) scheme)) (Map.delete key groupMap) keys
      Left (UnboundVariable var) -> buildTypeEnvFromGroups
        env groupMap (Cons var (delete var k))
      Left err -> Left err
  where
    mayDefs = Map.lookup key groupMap

typeProgramn:: List Definition -> Expr -> Either TypeError Scheme
typeProgramn defs exp = case runInfer <$>
    (inferType <$> (buildTypeEnv defs) <*> pure exp) of
  Left err -> Left err
  Right a -> a

typeTreeProgramn:: List Definition -> Expr -> Either TypeError TypeTree
typeTreeProgramn defs exp = case m of
  Left e -> Left e
  Right m' -> case evalState (runExceptT m') initUnique of
    Left err -> Left err
    Right res -> Right $ closeOver' res
  where
    m = (infer <$> (buildTypeEnv defs) <*> pure exp)

typeTreeProgramnEnv:: TypeEnv -> Expr -> Either TypeError TypeTree
typeTreeProgramnEnv env expr = case evalState (runExceptT (infer env expr)) initUnique of
      Left err -> Left err
      Right res -> Right $ closeOver' res

closeOver' :: (Tuple (Map.Map TVar Type) TypeTree) -> TypeTree
closeOver' (Tuple s tt) = apply s tt

prettyPrintType:: Type -> String
prettyPrintType (TypVar (TVar a)) = a
prettyPrintType (TypCon a) =  a
prettyPrintType (TypArr t1@(TypArr _ _) t2) = "(" ++ prettyPrintType t1 ++ ")" ++ " -> " ++ prettyPrintType t2
prettyPrintType (TypArr t1 t2) = prettyPrintType t1 ++ " -> " ++ prettyPrintType t2
prettyPrintType (AD a) = prettyPrintAD a
prettyPrintType (TypeError a) = prettyPrintTypeError a

prettyPrintAD:: AD -> String
prettyPrintAD (TList a) = "[" ++ prettyPrintType a ++ "]"
prettyPrintAD (TTuple (Cons t1 a)) = foldl (\s t -> s ++ "," ++ prettyPrintType t) ("(" ++ prettyPrintType t1 ) a  ++ ")"
prettyPrintAD (TTuple Nil) = "()"

emptyType:: Type
emptyType = TypCon ""

-- typeTreeProgramnEnv env expr
-- types subtree if typ correct
buildEmptyTypeTree:: TypeEnv -> Expr  -> TypeTree
buildEmptyTypeTree env e = case typeTreeProgramnEnv env e of
  Right tt -> tt
  Left err -> f err e
  where
  f err (Atom _) = TAtom (TypeError err)
  f err (List ls) = TListTree (map (buildEmptyTypeTree env) ls) (TypeError err)
  f err (NTuple ls) = TNTuple (map (buildEmptyTypeTree env) ls) (TypeError err)
  f err (Binary op e1 e2) = TBinary (typeOP op) (buildEmptyTypeTree env e1) (buildEmptyTypeTree env e2) (TypeError err)
  f err (SectL e op) = TSectL (buildEmptyTypeTree env e) (typeOP op) (TypeError err)
  f err (SectR op e) = TSectR (typeOP op) (buildEmptyTypeTree env e) (TypeError err)
  f err (PrefixOp _) = TPrefixOp (TypeError err)
  f err (IfExpr e1 e2 e3) = TIfExpr (buildEmptyTypeTree env e1) (buildEmptyTypeTree env e2) (buildEmptyTypeTree env e3) (TypeError err)
  f err (LetExpr b1 e1 e2) =let f env' =  TLetExpr (g b1) (buildEmptyTypeTree env' e1) (buildEmptyTypeTree env' e2) (TypeError err) in
                        case getTypEnv b1 env of
                          Nothing -> f env
                          Just env' -> f env'
  f err (Lambda bs e) = let f env' = TLambda (map g bs) (buildEmptyTypeTree env' e) (TypeError err) in
                    case getTypEnvFromList bs env of
                          Nothing -> f env
                          Just env' -> f env'
  f err (App e es) = TApp (buildEmptyTypeTree env e) (map (buildEmptyTypeTree env) es) (TypeError err)

-- Binding to BindingType
  g (Lit _) = TLit emptyType
  g (ConsLit b1 b2) = TConsLit (g b1) (g b2) emptyType
  g (ListLit bs) = TListLit (map g bs) emptyType
  g (NTupleLit bs) = TNTupleLit (map g bs) emptyType

  typeOP op = case typeTreeProgramnEnv env (PrefixOp op) of
    Left err -> TypeError err
    Right (TPrefixOp typ) -> typ

eqScheme :: Scheme -> Scheme -> Boolean
eqScheme (Forall l1 t1) (Forall l2 t2)
  = ((length l1) == (length l2)) && (fst (eqType Map.empty t1 t2))

eqType:: Map.Map TVar TVar -> Type -> Type ->Tuple Boolean (Map.Map TVar TVar)
eqType map (TypVar a) (TypVar b) = case  Map.lookup a map of
  (Just b') -> Tuple (b' == b) (map)
  Nothing -> Tuple true (Map.insert a b map)
eqType map (TypCon a) (TypCon b) = Tuple (a == b) map
eqType map (TypArr a b) (TypArr a' b') = Tuple (fst tup1 && fst tup2) (snd tup2)
  where
  tup1 = eqType map a a'
  tup2 = eqType (snd tup1) b b'
eqType map (AD (TList a)) (AD (TList b)) = eqType map a b
eqType map (AD (TTuple a)) (AD (TTuple b)) = eqTypeList map a b
eqType map _ _ = Tuple false map

eqTypeList:: Map.Map TVar TVar -> List Type -> List Type -> Tuple Boolean (Map.Map TVar TVar)
eqTypeList map (Cons a as) (Cons b bs) = let tup1 = eqType map a b in if (fst tup1)
  then eqTypeList (snd tup1) as bs
  else Tuple false (snd tup1)
eqTypeList map Nil Nil = Tuple true map
eqTypeList map _ _ = Tuple false map

txToABC:: TypeTree -> TypeTree
txToABC tt = fst $ runState (helptxToABC tt) {count : 0, env : empty}

helptxToABC:: TypeTree -> State {count:: Int, env:: Map String String} TypeTree
helptxToABC tt = go tt
  where
    go (TAtom t) = helpTypeToABC t >>= \t -> return $ TAtom t
    go (TListTree tts t) = do
      tts' <- mapM helptxToABC tts
      t' <- helpTypeToABC t
      return $ TListTree tts' t'
    go (TNTuple tts t) = do
      tts' <- mapM helptxToABC tts
      t' <- helpTypeToABC t
      return $ TNTuple tts' t'
    go (TBinary t1 tt1 tt2 t) = do
      t1' <- helpTypeToABC t1
      tt1' <- helptxToABC tt1
      tt2' <- helptxToABC tt2
      t' <- helpTypeToABC t
      return $ TBinary t1' tt1' tt2' t'
    go (TUnary t1 tt t) = do
      t1' <- helpTypeToABC t1
      tt' <- helptxToABC tt
      t' <- helpTypeToABC t
      return $ (TUnary t1' tt' t')
    go (TSectL tt t1 t) = do
      t1' <- helpTypeToABC t1
      tt' <- helptxToABC tt
      t' <- helpTypeToABC t
      return $ TSectL tt' t1' t'
    go (TSectR t1 tt t) = do
      t1' <- helpTypeToABC t1
      tt' <- helptxToABC tt
      t' <- helpTypeToABC t
      return $ TSectR t1' tt' t'
    go (TPrefixOp t) = helpTypeToABC t >>= \t -> return $ TPrefixOp t
    go (TIfExpr tt1 tt2 tt3 t) = do
      tt1' <- helptxToABC tt1
      tt2' <- helptxToABC tt2
      tt3' <- helptxToABC tt3
      t' <- helpTypeToABC t
      return $ TIfExpr tt1' tt2' tt3' t'
    go (TLetExpr tb tt1 tt2 t) = do
      tb' <- helpBindingToABC tb
      tt1' <- helptxToABC tt1
      tt2' <- helptxToABC tt2
      t' <- helpTypeToABC t
      return $ TLetExpr tb' tt1' tt2' t'
    go (TLambda tbs tt t) = do
      tbs' <- mapM helpBindingToABC tbs
      tt' <- helptxToABC tt
      t' <- helpTypeToABC t
      return $ TLambda tbs' tt' t'
    go (TApp tt tts t) = do
      tt' <- helptxToABC tt
      tts' <- mapM helptxToABC tts
      t' <- helpTypeToABC t
      return $ TApp tt' tts' t'

typeToABC:: Type -> Type
typeToABC t = fst $ runState (helpTypeToABC t) {count: 0, env: empty}

helpTypeToABC:: Type  -> State {count :: Int, env:: (Map String String)} Type
helpTypeToABC t = go t
  where
   go (TypVar (TVar var)) = do
      {env = env :: Map String String} <- get
      case lookup var env of
        Just var' -> return $ TypVar (TVar var')
        Nothing -> do
          var' <- freshLetter
          let env' = (insert var var' env)
          {count=count:: Int} <- get
          put {count:count, env:env'}
          return $ TypVar (TVar var')
   go (TypArr t1 t2) = do
        t1' <- helpTypeToABC t1
        t2' <- helpTypeToABC t2
        return $ TypArr t1' t2'
   go (AD a) = helpADTypeToABC a >>= \a -> return $ AD a
   go a = return a

helpADTypeToABC:: AD -> State {count :: Int, env:: (Map String String)} AD
helpADTypeToABC (TList t) = helpTypeToABC t >>= \t -> return $ TList t
helpADTypeToABC (TTuple ts) = mapM helpTypeToABC ts >>= \ts -> return $ TTuple ts

helpBindingToABC:: TypeBinding -> State {count :: Int, env:: (Map String String)} TypeBinding
helpBindingToABC bin = go bin
  where
    go (TLit t) = helpTypeToABC t >>= \t ->return $ TLit t
    go (TConsLit tb1 tb2 t) = do
      tb1' <- helpBindingToABC tb1
      tb2' <- helpBindingToABC tb2
      t' <- helpTypeToABC t
      return $ TConsLit tb1' tb2' t'
    go (TListLit tbs t) = do
      tbs' <- mapM helpBindingToABC tbs
      t' <- helpTypeToABC t
      return $ TListLit tbs' t'
    go (TNTupleLit tbs t) = do
      tbs' <- mapM helpBindingToABC tbs
      t' <- helpTypeToABC t
      return $ TNTupleLit tbs' t'

freshLetter:: State {count:: Int, env:: Map String String} String
freshLetter = do
    {count = count, env = env :: Map String String} <- get
    put {count:count+1, env:env}
    return $ fromCharList $ newTypVar count

letters:: List Char
letters = (toList $ toCharArray "abcdefghijklmnopqrstuvwxyz")

letters1:: List Char
letters1 = (toList $ toCharArray " abcdefghijklmnopqrstuvwxyz")

newTypVar:: Int -> List Char
newTypVar i = case (letters !! i) of
  Just c ->  Cons c Nil
  Nothing -> let i1 = (i `div` 26) in let i2 = (i `mod` 26) in (newTypVar1 i1) ++ (newTypVar i2)

-- workaround
-- if i  subtract one from i => stack overflow at ~50
newTypVar1:: Int -> List Char
newTypVar1 i = case (letters1 !! (i)) of
  Just c ->  Cons c Nil
  Nothing -> let i1 = (i `div` 26) in let i2 = (i `mod` 26) in (newTypVar1 i1) ++ (newTypVar i2)

prettyPrintTypeError:: TypeError -> String
prettyPrintTypeError (UnificationFail t1 t2) = let t1t2 = twoTypesToABC t1 t2  in
                                                "UnificationFail: Can't unify "
                                                ++ prettyPrintType (fst t1t2) ++ " with "
                                                ++ prettyPrintType (snd t1t2)
prettyPrintTypeError (InfiniteType a b) = let ab = twoTypesToABC (TypVar a) b in
                                              "InfiniteType: cannot construct the infinite type: "
                                              ++ prettyPrintType (fst ab)++ " ~ "
                                              ++ prettyPrintType (snd ab)
prettyPrintTypeError (UnboundVariable var) = "UnboundVariable: Not in scope "
                                              ++ var
prettyPrintTypeError (UnificationMismatch ts1 ts2) = let ts1ts2 = twoTypeListsToABC ts1 ts2 in
    "UnificationMismatch: " ++ toStr (fst ts1ts2) ++ " " ++ toStr (snd ts1ts2)
  where
    toStr ts = "[" ++ (foldr (\t s -> t ++ "," ++ s) "]" (map (\a -> prettyPrintType (typeToABC a)) ts))
prettyPrintTypeError (UnknownError str) = "UnknownError: " ++ str

twoTypesToABC t1 t2 = (\(TypArr t1' t2') -> Tuple t1' t2') (typeToABC (TypArr t1 t2))
twoTypeListsToABC t1 t2 = (\(TypArr (AD (TTuple t1')) (AD (TTuple t2'))) -> Tuple t1' t2') (typeToABC (TypArr (AD (TTuple t1)) (AD (TTuple t2)) ))

checkForError:: Path -> TypeTree -> Boolean
checkForError p' tt = case p' of
  (End) -> isTypeError $ extractType tt
  (Fst p) -> case tt of
      TBinary op e1 e2 _-> checkForError p e1
      TUnary op e      _-> checkForError p e
      TSectL e op      _-> checkForError p e
      TIfExpr ce te ee _-> checkForError p ce
      TLambda bs body  _-> checkForError p body
      TApp e es        _-> checkForError p e
      _               -> true
  (Snd p) -> case tt of
      TBinary op e1 e2 _-> checkForError p e2
      TSectR op e      _-> checkForError p e
      TIfExpr ce te ee _-> checkForError p te
      _               -> true
  (Thrd p) -> case tt of
      TIfExpr ce te ee _-> checkForError p ee
      _ -> true
  (Nth n p) -> case tt of
      TListTree es  _-> nth n es p
      TNTuple es _-> nth n es p
      TApp e' es _-> nth n es p
      _        -> true
  where
    nth n es p = case (!!) es n of
                  Nothing -> true
                  Just e -> checkForError p e

isTypeError:: Type -> Boolean
isTypeError t = case t of
  (TypeError _) -> true
  _ -> false

module TypeChecker where

import Prelude (class Monad, class Eq, class Show, (&&), (==), (>>=), map, (++), ($), pure, (<*>), (<$>), return, bind, const, otherwise, show, (+), div, mod)

import Control.Monad.Except.Trans (ExceptT, runExceptT, throwError)
import Control.Monad.State (State, evalState, put, get)

import Data.Either (Either(Left, Right))
import Data.List (List(..), length, delete, concat, unzip, foldM, toList, (:), zip)
import Data.Map as Map
import Data.Tuple (Tuple(Tuple), snd, fst)
import Data.Set as Set
import Data.Foldable (foldl, foldr)
import Data.Maybe (Maybe(..), fromMaybe)

import Data.Tuple (Tuple(..),fst,snd)
import Data.Map (Map(..), insert, lookup, empty)
import Data.List.WordsLines (fromCharList)
import Data.List (List(Nil, Cons), singleton, fromList, toList, length, (..), zipWithA,concat,zip, (!!))
import Data.String (joinWith, toCharArray)

import Control.Apply ((*>))
import Control.Bind ((=<<), (>=>))
import Control.Monad.State

import AST (AD(..), Atom(..), Binding(..), Definition(..), Expr(..), Op(..), TVar(..), Type(..), TypeBinding(..), TypeTree(..),TypeError(..),Path(..))

data Scheme = Forall (List TVar) Type

instance showScheme :: Show Scheme where
  show (Forall a b) = "Forall (" ++ show a ++ ") " ++ show b

data TypeEnv = TypeEnv (Map.Map Atom Scheme)

instance showTypeEnv :: Show TypeEnv where
  show (TypeEnv a) = show a

data Unique = Unique { count :: Int }

type Infer a = ExceptT TypeError (State Unique) a

data InferState = InferState { count  :: Int }

data Env = TypEnv {types:: Map.Map Atom Scheme}

type Constraint = Tuple Type Type
type Subst = Map.Map TVar Type

class Substitutable a where
  apply :: Subst -> a -> a
  ftv   :: a -> Set.Set TVar

instance subScheme :: Substitutable Scheme where
    apply s (Forall as t)   = Forall as $ apply s' t
                              where s' = foldr Map.delete s as
    ftv (Forall as t) = (ftv t) `Set.difference` Set.fromList as

instance subType :: Substitutable Type where
   apply _ (TypCon a) = TypCon a
   apply s t@(TypVar a) = fromMaybe t $ Map.lookup a s
   apply s (TypArr t1 t2) =  TypArr (apply s t1) (apply s t2)
   apply s (AD a) = AD (apply s a)

   ftv (TypCon  _)         = Set.empty
   ftv (TypVar a)       = Set.singleton a
   ftv (TypArr t1 t2) =  Set.union (ftv t1) (ftv t2)
   ftv (AD a) = ftv a

instance listSub :: (Substitutable a) => Substitutable (List a) where
  apply s xs = map (apply s) xs
  ftv   = foldr (\a b -> Set.union (ftv a) b) Set.empty

instance subTypeEnv :: Substitutable TypeEnv where
  apply s (TypeEnv env) =  TypeEnv $ map (apply s) env
  ftv (TypeEnv env) = ftv $ snd $ unzip $ Map.toList env

instance subAD :: Substitutable AD where
  apply s (TList t) = TList (apply s t)
  apply s(TTuple t) = TTuple (apply s t)

  ftv (TList t) = ftv t
  ftv (TTuple t) = ftv t

instance subTupAtomScheme :: Substitutable (Tuple Atom Scheme) where
  apply s (Tuple a b) = Tuple a (apply s b)

  ftv (Tuple _ b) = ftv b

instance subTypeTree :: Substitutable TypeTree where
  apply s (TAtom t) = TAtom $ apply s t
  apply s (TListTree l t) = TListTree (apply s l) (apply s t)
  apply s (TNTuple l t) = TNTuple (apply s l) (apply s t)
  apply s (TBinary op tt1 tt2 t) =  TBinary (apply s op) (apply s tt1) (apply s tt2) (apply s t)
  apply s (TUnary op tt t) = TUnary (apply s op) (apply s tt) (apply s t)
  apply s (TSectL tt op t) = TSectL (apply s tt) (apply s op) (apply s t)
  apply s (TSectR op tt t) = TSectR (apply s op) (apply s tt) (apply s t)
  apply s (TPrefixOp t) = TPrefixOp $ apply s t
  apply s (TIfExpr tt1 tt2 tt3 t) = TIfExpr (apply s tt1) (apply s tt2) (apply s tt3) (apply s t)
  apply s (TLetExpr b tt1 tt2 t) = TLetExpr (apply s b) (apply s tt1) (apply s tt2) (apply s t)
  apply s (TLambda lb tt t) = TLambda (apply s lb) (apply s tt) (apply s t)
  apply s (TApp tt1 l t) = TApp (apply s tt1) (apply s l) (apply s t)

  ftv (TAtom t)  = ftv t
  ftv (TListTree _ t)  = ftv t
  ftv (TNTuple _ t)  = ftv t
  ftv (TBinary _ _ _ t)  = ftv t
  ftv (TUnary _ _ t)  = ftv t
  ftv (TSectL _ _ t)  = ftv t
  ftv (TSectR _ _ t)  = ftv t
  ftv (TPrefixOp t)  = ftv t
  ftv (TIfExpr _ _ _ t)  = ftv t
  ftv (TLetExpr _ _ _ t)  = ftv t
  ftv (TLambda _ _ t)  = ftv t
  ftv (TApp _ _ t)  = ftv t

instance subTypeBinding :: Substitutable TypeBinding where
  apply s (TLit t) = TLit $ apply s t
  apply s (TConsLit b1 b2 t) = TConsLit (apply s b1) (apply s b2) (apply s t)
  apply s (TListLit lb t) = TListLit (apply s lb) (apply s t)
  apply s (TNTupleLit lb t) = TListLit (apply s lb) (apply s t)

  ftv (TLit t) = ftv t
  ftv (TConsLit _ _ t) = ftv t
  ftv (TListLit _ t) = ftv t
  ftv (TNTupleLit _ t) = ftv t

initInfer :: InferState
initInfer = InferState { count : 0}

initUnique :: Unique
initUnique = Unique { count : 0 }

extend :: TypeEnv -> Tuple Atom Scheme -> TypeEnv
extend (TypeEnv env) (Tuple x s) = TypeEnv $ Map.insert x s env

nullSubst:: Subst
nullSubst = Map.empty

fresh :: Infer Type
fresh = do
  (Unique s) <- get
  put (Unique {count:(s.count+1)})
  return $ TypVar $ TVar $ "t_" ++ show s.count

unify ::  Type -> Type -> Infer Subst
unify (TypArr l r) (TypArr l' r')  = do
    s1 <- unify l l'
    s2 <- unify (apply s1 r) (apply s1 r')
    return (s2 `compose` s1)

unify (TypVar a) t = bindVar a t
unify t (TypVar a) = bindVar a t
unify (TypCon a) (TypCon b) | a == b = return nullSubst
unify (AD a) (AD b) | a == b = return nullSubst
unify (AD a) (AD b) = unifyAD a b
unify t1 t2 = throwError $ UnificationFail t1 t2

unifyAD :: AD -> AD -> Infer Subst
unifyAD (TList a) (TList b) = unify a b
unifyAD (TTuple (Cons a as)) (TTuple (Cons b bs)) = do
  s1 <- unifyAD (TTuple as) (TTuple bs)
  s2 <- unify a b
  return (s1 `compose` s2)
unifyAD (TTuple Nil) (TTuple Nil) = return nullSubst

unifyAD a b = throwError $ UnificationFail (AD a) (AD b)

bindVar ::  TVar -> Type -> Infer Subst
bindVar a t | t == (TypVar a)     = return nullSubst
         | occursCheck a t = throwError $ InfiniteType a t
         | otherwise       = return $ Map.singleton a t

occursCheck :: forall a.  (Substitutable a) => TVar -> a -> Boolean
occursCheck a t = a `Set.member` ftv t

compose :: Subst -> Subst -> Subst
compose s1 s2 = map (apply s1) s2 `Map.union` s1

instantiate ::  Scheme -> Infer Type
instantiate (Forall as t) = do
  as' <- mapM (const fresh) as
  let s = Map.fromList $ zip as as'
  return $ apply s t

generalize :: TypeEnv -> Type -> Scheme
generalize env t  = Forall as t
    where as = Set.toList $ ftv t `Set.difference` ftv env

runInfer :: Infer (Tuple Subst Type) -> Either TypeError Scheme
runInfer m = case evalState (runExceptT m) initUnique of
  Left err -> Left err
  Right res -> Right $ closeOver res

closeOver :: (Tuple (Map.Map TVar Type) Type) -> Scheme
closeOver (Tuple sub ty) = sc --TODO add normalize back in
  where sc = generalize emptyTyenv (apply sub ty)

emptyTyenv :: TypeEnv
emptyTyenv = TypeEnv Map.empty

-- TODO find Purescript equivalent
mapM :: forall a b m. (Monad m) => (a -> m b) -> List a -> m (List b)
mapM f as = foldr k (return Nil) as
            where
              k a r = do
                x <- f a
                xs <- r
                return (x:xs)

lookupEnv :: TypeEnv -> Atom -> Infer (Tuple Subst Type)
lookupEnv (TypeEnv env) x = do
  case Map.lookup x env of
    Nothing -> throwError $ UnboundVariable (f x)
    Just s  -> do t <- instantiate s
                  return (Tuple nullSubst t)
  where
    f (Name x) = x
    f x = show x

extractType:: TypeTree -> Type
extractType (TAtom t)  =  t
extractType (TListTree _ t)  =  t
extractType (TNTuple _ t)  =  t
extractType (TBinary _ _ _ t)  =  t
extractType (TUnary _ _ t)  =  t
extractType (TSectL _ _ t)  =  t
extractType (TSectR _ _ t)  =  t
extractType (TPrefixOp t)  =  t
extractType (TIfExpr _ _ _ t)  =  t
extractType (TLetExpr _ _ _ t)  =  t
extractType (TLambda _ _ t)  =  t
extractType (TApp tt1 _ t)  =  t

extractBindingType:: TypeBinding -> Type
extractBindingType (TLit t) = t
extractBindingType (TConsLit _ _ t) = t
extractBindingType (TListLit _ t) = t
extractBindingType (TNTupleLit _ t) = t

inferType :: TypeEnv -> Expr -> Infer (Tuple Subst Type)
inferType env exp = do
  Tuple s t <- infer env exp
  let t' = extractType t
  return $ Tuple s t'

infer :: TypeEnv -> Expr -> Infer (Tuple Subst TypeTree)
infer env ex = case ex of

  (Atom x@(Name _)) -> do
    Tuple s t <- lookupEnv env x
    return (Tuple s $ TAtom t)
  (Atom (Bool _)) -> return (Tuple nullSubst $ TAtom (TypCon "Bool"))
  (Atom (Char _)) -> return (Tuple nullSubst $ TAtom (TypCon "Char"))
  (Atom (AInt _)) -> return (Tuple nullSubst $ TAtom (TypCon "Int"))

  Lambda (Cons bin Nil) e -> do
    tv <- fresh
    Tuple envL tvB <- extractBinding bin
    let env' = foldr (\a b -> extend b a) env envL
    s0 <- unify tv (extractBindingType tvB)
    (Tuple s1 t1) <- infer env' e
    return (Tuple s1 $ apply (s0 `compose` s1) (TLambda (Cons tvB Nil) t1 $ tv `TypArr` (extractType t1)))

  Lambda (Cons bin xs) e -> do
    tv <- fresh
    Tuple envL tvB <- extractBinding bin
    let env' = foldr (\a b -> extend b a) env envL
    s0 <- unify tv (extractBindingType tvB)
    (Tuple s1 (TLambda tb tt t1)) <- infer env' (Lambda xs e)
    return (Tuple s1 $ apply (s0 `compose` s1) (TLambda (Cons tvB tb) tt (tv `TypArr` t1)))

  Lambda Nil e -> infer env e

    -- one element list
  App e1 (Cons e2 Nil) -> do
    tv <- fresh
    (Tuple s1 t1) <- infer env e1
    (Tuple s2 t2) <- infer (apply s1 env) e2
    s3       <- unify (apply (s1  `compose` s2) (extractType t1)) (TypArr (extractType t2) tv)
    return (Tuple (s3 `compose` s2 `compose` s1) (apply s3 $ TApp t1 (Cons t2 Nil) tv))

  App e1 (Cons e2 xs) -> do
    Tuple s (TApp (TApp tt lt _) lt' t') <- infer env  (App (App e1 (Cons e2 Nil)) xs)
    return $ Tuple s (TApp tt (lt++lt') t')

  App _ Nil -> throwError $ UnknownError "congrats you found a bug TypeChecker.infer (App Nil)"

  LetExpr bin e1 e2 -> do
    (Tuple s1 t1) <- infer env e1
    (Tuple list typ) <- extractBinding bin
    s2 <- unify (extractBindingType typ) (extractType t1)
    let env' = apply (s1 `compose` s2) env
        t'   = generalize env' (apply (s1 `compose` s2) (extractType t1))
        env'' = apply s2 (foldr (\a b -> extend b a) env' list)
    (Tuple s3 t2) <- infer env'' e2
    let sC = (s1 `compose` s2 `compose` s3)
    return (Tuple sC $ apply sC (TLetExpr typ t1 t2 (extractType t2)))

  IfExpr cond tr fl -> do
    tv <- fresh
    t@(TypVar t') <- fresh
    let name = Name $ "if"
    let env' = env `extend` (Tuple name (Forall (Cons t' Nil) (TypArr (TypCon "Bool") (TypArr t (TypArr t  t)))))
    Tuple s (TApp tt (Cons tcond (Cons ttr (Cons tfl Nil))) ift) <- infer env' (App (Atom  name) (toList [cond, tr, fl]))
    return (Tuple s $ apply s (TIfExpr tcond ttr tfl ift))

  PrefixOp op -> do
    Tuple s t <- inferOp op
    return (Tuple s $ TPrefixOp t)

  SectL e op -> do
    tv <- fresh
    (Tuple s1 t1) <- inferOp op
    (Tuple s2 t2) <- infer (apply s1 env) e
    s3       <- unify (apply s2 t1) (TypArr (extractType t2) tv)
    return (Tuple (s3 `compose` s2 `compose` s1) (apply s3 (TSectL t2 t1 tv)))

  SectR op e -> do
    (Tuple s1 t1@(TypArr a (TypArr b c))) <- inferOp op
    (Tuple s2 t2) <- infer env e
    s3       <- unify (apply s2 b) (extractType t2)
    let s4 = (s3 `compose` s2 `compose` s1)
    return (Tuple s4 (apply s4 (TSectR t1 t2 (TypArr a c))))

  Unary op e -> do
    tv <- fresh
    (Tuple s1 t1) <- inferOp op
    (Tuple s2 t2) <- infer (apply s1 env) e
    s3       <- unify (apply s2 t1) (TypArr (extractType t2) tv)
    return (Tuple (s3 `compose` s2 `compose` s1) (apply s3 (TUnary t1 t2 tv)))

  Binary op e1 e2 -> do
    (Tuple s (TApp tt (Cons tt1 (Cons tt2 Nil)) t)) <- infer env (App (PrefixOp op) (Cons e1 (Cons e2 Nil)))
    return $ Tuple s (TBinary (extractType tt) tt1 tt2 t)

  List (Cons e1 xs) -> do
    (Tuple s1 (TListTree ltt (AD (TList t1)))) <- infer env (List xs)
    (Tuple s2 t2) <- infer (apply s1 env) e1
    s3 <- unify (apply s2 t1) (extractType t2)
    return (Tuple (s3 `compose` s2 `compose` s1) (apply s3 (TListTree (Cons t2 ltt) (AD $ TList (extractType t2)))))

  List Nil -> do
    tv <- fresh
    return (Tuple nullSubst $ TListTree Nil (AD $ TList tv))

  NTuple (Cons e Nil) -> do
    (Tuple s t) <- infer env e
    return $ Tuple s $ TNTuple (Cons t Nil) (AD $ TTuple $ Cons (extractType t) Nil)

  NTuple (Cons e1 xs) -> do
    (Tuple s1 (TNTuple lt (AD (TTuple t1)))) <- infer env (NTuple xs)
    (Tuple s2 t2) <- infer (apply s1 env) e1
    return (Tuple (s2 `compose` s1) $ TNTuple (Cons t2 lt) (AD $ TTuple (Cons (extractType t2) t1)))

  NTuple Nil -> throwError $ UnknownError "congrats you found a bug in TypeChecker.infer (NTuple Nil)"

inferOp :: Op -> Infer (Tuple Subst Type)
inferOp op = do
  a <- fresh
  b <- fresh
  c <- fresh
  case op of
    Composition ->
      f ((b `TypArr` c) `TypArr` ((a `TypArr` b) `TypArr` (a `TypArr` c)))
    Power -> int3
    Mul ->  int3
    Div -> int3
    Mod -> int3
    Add -> int3
    Sub -> int3
    Colon -> f (a `TypArr` ((AD $ TList a) `TypArr` (AD $ TList a)))
    Append -> f ((AD $ TList a) `TypArr` ((AD $ TList a) `TypArr` (AD $ TList a)))
    Equ -> aBool a
    Neq -> aBool a
    Lt -> aBool a
    Leq -> aBool a
    Gt -> aBool a
    Geq -> aBool a
    And -> f (TypCon "Bool" `TypArr` (TypCon "Bool" `TypArr` TypCon "Bool"))
    Or -> f (TypCon "Bool" `TypArr` (TypCon "Bool" `TypArr` TypCon "Bool"))
    Dollar -> f ((a `TypArr` b) `TypArr` (a `TypArr` b))
 where
  f typ = return (Tuple nullSubst typ)
  int3 = f (TypCon "Int" `TypArr` (TypCon "Int" `TypArr` TypCon "Int"))
  aBool a = f (a `TypArr` (a `TypArr` TypCon "Bool"))

inferDef :: TypeEnv -> Definition -> Infer (Tuple Subst Type)
inferDef env (Def str bin exp) = do
    tv <- fresh
    let env' = env `extend` (Tuple (Name str) (Forall Nil tv))
    let exp' = Lambda bin exp
    Tuple s1 t1' <- infer env' exp'
    let t1 = extractType t1'
    let env'' = env `extend` (Tuple (Name str) (Forall Nil (apply s1 t1)))
    Tuple s2 t2 <- infer env'' exp'
    s3 <- unify (apply s1 t1) (apply s2 (extractType t2))
    return $ Tuple (s3 `compose` s1) (apply s3 (apply s1 t1))

extractConsLit:: Type -> Binding -> Infer (Tuple (List (Tuple Atom Scheme)) TypeBinding)
extractConsLit tv (ConsLit a b@(ConsLit _ _)) = do
  Tuple list1 typ1 <- extractBinding a
  s1 <- unify (extractBindingType typ1) tv
  let list1' = apply s1 list1
  Tuple listB b@(TConsLit b1 b2 (AD ( TList typB))) <- extractConsLit tv b
  s2 <- unify typB (apply s1 tv)
  return $ Tuple (apply s2 (list1' ++ listB)) $ apply s2 (TConsLit typ1 b (AD $ TList (apply s2 typB)))

extractConsLit tv (ConsLit a (Lit b)) = do
  Tuple list typ <- extractBinding a
  s1 <- unify tv (extractBindingType typ)
  let list' = apply s1 list
  let ltyp = AD $ TList (apply s1 tv)
  return $ Tuple (Cons (Tuple b $ Forall Nil ltyp) list') $ (apply s1 (TConsLit typ (TLit ltyp) ltyp))

extractConsLit tv (ConsLit a b) = do
  Tuple list typ <- extractBinding a
  Tuple list2 typ2 <- extractBinding b
  s1 <- unify (extractBindingType typ2) (extractBindingType typ)
  s2 <- unify (apply s1 (extractBindingType typ)) tv
  let sC =  s2 `compose` s1
  let list' = map (\(Tuple a b) -> Tuple a $ apply sC b) list
  let list2' = map (\ (Tuple a b) -> Tuple a $  apply sC b) list2
  return $ Tuple (list' ++ list2') $ apply sC $ TConsLit typ typ2 $ AD $ TList tv

extractConsLit _ _ = throwError $ UnknownError "congrats you found a bug in TypeChecker.extractConsLit"

extractListLit :: List Binding -> Infer (List (Tuple (List (Tuple Atom Scheme)) TypeBinding))
extractListLit (Cons a Nil) = do
  b1 <- extractBinding a
  return (Cons b1 Nil)

extractListLit (Cons a b) = do
  b1 <- extractBinding a
  bs <- extractListLit b
  return (Cons b1 bs)

extractListLit Nil = return Nil

extractBinding:: Binding -> Infer (Tuple (List (Tuple Atom Scheme)) TypeBinding)
extractBinding (Lit a@(Name _)) = do
  tv <- fresh
  return $ Tuple (toList [(Tuple a (Forall Nil tv))]) (TLit tv)
extractBinding (Lit (Bool _)) = return $ Tuple Nil $ TLit (TypCon "Bool")
extractBinding (Lit (Char _)) = return $ Tuple Nil $ TLit (TypCon "Char")
extractBinding (Lit (AInt _)) = return $ Tuple Nil $ TLit (TypCon "Int")

extractBinding a@(ConsLit _ _) = do
  tv <- fresh
  extractConsLit tv a

extractBinding (ListLit a) = do -- Tuple TypEnv  (TListLit (List TypeBinding) Type)
  bs <- extractListLit a
  tv <- fresh
  let ini = Tuple Nil (TListLit Nil tv)
  Tuple list (TListLit lb typ) <- foldM f ini bs
  return $ Tuple list (TListLit lb (AD $ TList typ))
  where
-- :: Tuple (List (Tuple Atom Scheme)) BindingType -> Tuple (List (Tuple Atom Scheme)) BindingType
--           -> Infer (Tuple (List (Tuple Atom Scheme)) BindingType)
    f (Tuple l1 (TListLit lb1 t1)) (Tuple l2 b) = do
      let ls = l1 ++ l2
      s1 <- unify t1 (extractBindingType b)
      return $ Tuple (apply s1 ls) (apply s1 (TListLit (Cons b lb1) t1))
    f _ _ = throwError $ UnknownError "congrats you found a bug in TypeChecker.extractBinding"

extractBinding (NTupleLit a) = do
  bs <- extractListLit a
  let tup = unzip bs
  return $ Tuple (concat $ fst tup) (TNTupleLit (snd tup) (AD $ TTuple $ (map extractBindingType (snd tup))))

inferGroup:: TypeEnv -> List Definition -> Infer (Tuple Subst Type)
inferGroup _ Nil = throwError $ UnknownError "Cant type empty Group"
inferGroup env (Cons def1 Nil) = inferDef env def1
inferGroup env (Cons def1 defs) = do
  Tuple s1 t1 <- inferDef env def1
  Tuple s2 t2 <- inferGroup env defs
  s3 <- unify t1 t2
  return $ Tuple s3 (apply s3 t1)

buildTypeEnv:: List Definition -> Either TypeError TypeEnv
buildTypeEnv Nil = Right emptyTyenv
buildTypeEnv defs = buildTypeEnvFromGroups emptyTyenv groupMap keys
  where
    groupMap = buildGroups defs
    keys = Map.keys groupMap

buildGroups:: List Definition -> Map.Map String (List Definition)
buildGroups Nil = Map.empty
buildGroups (Cons def@(Def str bin exp) Nil) =
  Map.singleton str (Cons def Nil)

buildGroups (Cons def@(Def str bin exp) defs) =
  case binList of
    Just list -> Map.insert str (Cons def list) defMap
    Nothing -> Map.insert str (Cons def Nil) defMap
  where
    defMap = buildGroups defs
    binList = Map.lookup str defMap
--                         old -> grouped Definitions -> keys -> new
buildTypeEnvFromGroups:: TypeEnv -> Map.Map String (List Definition)
  -> List String -> Either TypeError TypeEnv
buildTypeEnvFromGroups env _ Nil = Right env
buildTypeEnvFromGroups env groupMap k@(Cons key keys) =
  case mayDefs of
    Nothing -> Left $ UnboundVariable key
    Just defs -> case runInfer $ inferGroup env defs of
      Right scheme -> buildTypeEnvFromGroups
        (env `extend` (Tuple (Name key) scheme)) (Map.delete key groupMap) keys
      Left (UnboundVariable var) -> buildTypeEnvFromGroups
        env groupMap (Cons var (delete var k))
      Left err -> Left err
  where
    mayDefs = Map.lookup key groupMap

typeProgramn:: List Definition -> Expr -> Either TypeError Scheme
typeProgramn defs exp = case runInfer <$>
    (inferType <$> (buildTypeEnv defs) <*> pure exp) of
  Left err -> Left err
  Right a -> a

typeTreeProgramn:: List Definition -> Expr -> Either TypeError TypeTree
typeTreeProgramn defs exp = case m of
  Left e -> Left e
  Right m' -> case evalState (runExceptT m') initUnique of
    Left err -> Left err
    Right res -> Right $ closeOver' res
  where
    m = (infer <$> (buildTypeEnv defs) <*> pure exp)

typeTreeProgramnEnv:: TypeEnv -> Expr -> Either TypeError TypeTree
typeTreeProgramnEnv env expr = case evalState (runExceptT (infer env expr)) initUnique of
      Left err -> Left err
      Right res -> Right $ closeOver' res

closeOver' :: (Tuple (Map.Map TVar Type) TypeTree) -> TypeTree
closeOver' (Tuple s tt) = apply s tt

prettyPrintType:: Type -> String
prettyPrintType (TypVar (TVar a)) = a
prettyPrintType (TypCon a) =  a
prettyPrintType (TypArr t1@(TypArr _ _) t2) = "(" ++ prettyPrintType t1 ++ ")" ++ " -> " ++ prettyPrintType t2
prettyPrintType (TypArr t1 t2) = prettyPrintType t1 ++ " -> " ++ prettyPrintType t2
prettyPrintType (AD a) = prettyPrintAD a
prettyPrintType (TypeError a) = prettyPrintTypeError a

prettyPrintAD:: AD -> String
prettyPrintAD (TList a) = "[" ++ prettyPrintType a ++ "]"
prettyPrintAD (TTuple (Cons t1 a)) = foldl (\s t -> s ++ "," ++ prettyPrintType t) ("(" ++ prettyPrintType t1 ) a  ++ ")"
prettyPrintAD (TTuple Nil) = "()"

emptyType:: Type
emptyType = TypCon ""

-- typeTreeProgramnEnv env expr
-- types subtree if typ correct
buildEmptyTypeTree:: TypeEnv -> Expr  -> TypeTree
buildEmptyTypeTree env e = case typeTreeProgramnEnv env e of
  Right tt -> tt
  Left err -> f err e
  where
  f err (Atom _) = TAtom (TypeError err)
  f err (List ls) = TListTree (map (buildEmptyTypeTree env) ls) (TypeError err)
  f err (NTuple ls) = TNTuple (map (buildEmptyTypeTree env) ls) (TypeError err)
  f err (Binary op e1 e2) = TBinary (typeOP op) (buildEmptyTypeTree env e1) (buildEmptyTypeTree env e2) (TypeError err)
  f err (SectL e op) = TSectL (buildEmptyTypeTree env e) (typeOP op) (TypeError err)
  f err (SectR op e) = TSectR (typeOP op) (buildEmptyTypeTree env e) (TypeError err)
  f err (PrefixOp _) = TPrefixOp (TypeError err)
  f err (IfExpr e1 e2 e3) = TIfExpr (buildEmptyTypeTree env e1) (buildEmptyTypeTree env e2) (buildEmptyTypeTree env e3) (TypeError err)
  f err (LetExpr b1 e1 e2) = TLetExpr (g b1) (buildEmptyTypeTree env e1) (buildEmptyTypeTree env e2) (TypeError err)
  f err (Lambda bs e) = TLambda (map g bs) (buildEmptyTypeTree env e) (TypeError err)
  f err (App e es) = TApp (buildEmptyTypeTree env e) (map (buildEmptyTypeTree env) es) (TypeError err)

-- Binding to BindingType
  g (Lit _) = TLit emptyType
  g (ConsLit b1 b2) = TConsLit (g b1) (g b2) emptyType
  g (ListLit bs) = TListLit (map g bs) emptyType
  g (NTupleLit bs) = TNTupleLit (map g bs) emptyType

  typeOP op = case typeTreeProgramnEnv env (PrefixOp op) of
    Left err -> TypeError err
    Right (TPrefixOp typ) -> typ

eqScheme :: Scheme -> Scheme -> Boolean
eqScheme (Forall l1 t1) (Forall l2 t2)
  = ((length l1) == (length l2)) && (fst (eqType Map.empty t1 t2))

eqType:: Map.Map TVar TVar -> Type -> Type ->Tuple Boolean (Map.Map TVar TVar)
eqType map (TypVar a) (TypVar b) = case  Map.lookup a map of
  (Just b') -> Tuple (b' == b) (map)
  Nothing -> Tuple true (Map.insert a b map)
eqType map (TypCon a) (TypCon b) = Tuple (a == b) map
eqType map (TypArr a b) (TypArr a' b') = Tuple (fst tup1 && fst tup2) (snd tup2)
  where
  tup1 = eqType map a a'
  tup2 = eqType (snd tup1) b b'
eqType map (AD (TList a)) (AD (TList b)) = eqType map a b
eqType map (AD (TTuple a)) (AD (TTuple b)) = eqTypeList map a b
eqType map _ _ = Tuple false map

eqTypeList:: Map.Map TVar TVar -> List Type -> List Type -> Tuple Boolean (Map.Map TVar TVar)
eqTypeList map (Cons a as) (Cons b bs) = let tup1 = eqType map a b in if (fst tup1)
  then eqTypeList (snd tup1) as bs
  else Tuple false (snd tup1)
eqTypeList map Nil Nil = Tuple true map
eqTypeList map _ _ = Tuple false map

txToABC:: TypeTree -> TypeTree
txToABC tt = fst $ runState (helptxToABC tt) {count : 0, env : empty}

helptxToABC:: TypeTree -> State {count:: Int, env:: Map String String} TypeTree
helptxToABC tt = go tt
  where
    go (TAtom t) = helpTypeToABC t >>= \t -> return $ TAtom t
    go (TListTree tts t) = do
      tts' <- mapM helptxToABC tts
      t' <- helpTypeToABC t
      return $ TListTree tts' t'
    go (TNTuple tts t) = do
      tts' <- mapM helptxToABC tts
      t' <- helpTypeToABC t
      return $ TNTuple tts' t'
    go (TBinary t1 tt1 tt2 t) = do
      t1' <- helpTypeToABC t1
      tt1' <- helptxToABC tt1
      tt2' <- helptxToABC tt2
      t' <- helpTypeToABC t
      return $ TBinary t1' tt1' tt2' t'
    go (TUnary t1 tt t) = do
      t1' <- helpTypeToABC t1
      tt' <- helptxToABC tt
      t' <- helpTypeToABC t
      return $ (TUnary t1' tt' t')
    go (TSectL tt t1 t) = do
      t1' <- helpTypeToABC t1
      tt' <- helptxToABC tt
      t' <- helpTypeToABC t
      return $ TSectL tt' t1' t'
    go (TSectR t1 tt t) = do
      t1' <- helpTypeToABC t1
      tt' <- helptxToABC tt
      t' <- helpTypeToABC t
      return $ TSectR t1' tt' t'
    go (TPrefixOp t) = helpTypeToABC t >>= \t -> return $ TPrefixOp t
    go (TIfExpr tt1 tt2 tt3 t) = do
      tt1' <- helptxToABC tt1
      tt2' <- helptxToABC tt2
      tt3' <- helptxToABC tt3
      t' <- helpTypeToABC t
      return $ TIfExpr tt1' tt2' tt3' t'
    go (TLetExpr tb tt1 tt2 t) = do
      tb' <- helpBindingToABC tb
      tt1' <- helptxToABC tt1
      tt2' <- helptxToABC tt2
      t' <- helpTypeToABC t
      return $ TLetExpr tb' tt1' tt2' t'
    go (TLambda tbs tt t) = do
      tbs' <- mapM helpBindingToABC tbs
      tt' <- helptxToABC tt
      t' <- helpTypeToABC t
      return $ TLambda tbs' tt' t'
    go (TApp tt tts t) = do
      tt' <- helptxToABC tt
      tts' <- mapM helptxToABC tts
      t' <- helpTypeToABC t
      return $ TApp tt' tts' t'

typeToABC:: Type -> Type
typeToABC t = fst $ runState (helpTypeToABC t) {count: 0, env: empty}

helpTypeToABC:: Type  -> State {count :: Int, env:: (Map String String)} Type
helpTypeToABC t = go t
  where
   go (TypVar (TVar var)) = do
      {env = env :: Map String String} <- get
      case lookup var env of
        Just var' -> return $ TypVar (TVar var')
        Nothing -> do
          var' <- freshLetter
          let env' = (insert var var' env)
          {count=count:: Int} <- get
          put {count:count, env:env'}
          return $ TypVar (TVar var')
   go (TypArr t1 t2) = do
        t1' <- helpTypeToABC t1
        t2' <- helpTypeToABC t2
        return $ TypArr t1' t2'
   go (AD a) = helpADTypeToABC a >>= \a -> return $ AD a
   go a = return a

helpADTypeToABC:: AD -> State {count :: Int, env:: (Map String String)} AD
helpADTypeToABC (TList t) = helpTypeToABC t >>= \t -> return $ TList t
helpADTypeToABC (TTuple ts) = mapM helpTypeToABC ts >>= \ts -> return $ TTuple ts

helpBindingToABC:: TypeBinding -> State {count :: Int, env:: (Map String String)} TypeBinding
helpBindingToABC bin = go bin
  where
    go (TLit t) = helpTypeToABC t >>= \t ->return $ TLit t
    go (TConsLit tb1 tb2 t) = do
      tb1' <- helpBindingToABC tb1
      tb2' <- helpBindingToABC tb2
      t' <- helpTypeToABC t
      return $ TConsLit tb1' tb2' t'
    go (TListLit tbs t) = do
      tbs' <- mapM helpBindingToABC tbs
      t' <- helpTypeToABC t
      return $ TListLit tbs' t'
    go (TNTupleLit tbs t) = do
      tbs' <- mapM helpBindingToABC tbs
      t' <- helpTypeToABC t
      return $ TNTupleLit tbs' t'

freshLetter:: State {count:: Int, env:: Map String String} String
freshLetter = do
    {count = count, env = env :: Map String String} <- get
    put {count:count+1, env:env}
    return $ fromCharList $ newTypVar count

letters:: List Char
letters = (toList $ toCharArray "abcdefghijklmnopqrstuvwxyz")

letters1:: List Char
letters1 = (toList $ toCharArray " abcdefghijklmnopqrstuvwxyz")

newTypVar:: Int -> List Char
newTypVar i = case (letters !! i) of
  Just c ->  Cons c Nil
  Nothing -> let i1 = (i `div` 26) in let i2 = (i `mod` 26) in (newTypVar1 i1) ++ (newTypVar i2)

-- workaround
-- if i  subtract one from i => stack overflow at ~50
newTypVar1:: Int -> List Char
newTypVar1 i = case (letters1 !! (i)) of
  Just c ->  Cons c Nil
  Nothing -> let i1 = (i `div` 26) in let i2 = (i `mod` 26) in (newTypVar1 i1) ++ (newTypVar i2)

prettyPrintTypeError:: TypeError -> String
prettyPrintTypeError (UnificationFail t1 t2) = let t1t2 = twoTypesToABC t1 t2  in
                                                "UnificationFail: Can't unify "
                                                ++ prettyPrintType (fst t1t2) ++ " with "
                                                ++ prettyPrintType (snd t1t2)
prettyPrintTypeError (InfiniteType a b) = let ab = twoTypesToABC (TypVar a) b in
                                              "InfiniteType: cannot construct the infinite type: "
                                              ++ prettyPrintType (fst ab)++ " ~ "
                                              ++ prettyPrintType (snd ab)
prettyPrintTypeError (UnboundVariable var) = "UnboundVariable: Not in scope "
                                              ++ var
prettyPrintTypeError (UnificationMismatch ts1 ts2) = let ts1ts2 = twoTypeListsToABC ts1 ts2 in
    "UnificationMismatch: " ++ toStr (fst ts1ts2) ++ " " ++ toStr (snd ts1ts2)
  where
    toStr ts = "[" ++ (foldr (\t s -> t ++ "," ++ s) "]" (map (\a -> prettyPrintType (typeToABC a)) ts))
prettyPrintTypeError (UnknownError str) = "UnknownError: " ++ str

twoTypesToABC t1 t2 = (\(TypArr t1' t2') -> Tuple t1' t2') (typeToABC (TypArr t1 t2))
twoTypeListsToABC t1 t2 = (\(TypArr (AD (TTuple t1')) (AD (TTuple t2'))) -> Tuple t1' t2') (typeToABC (TypArr (AD (TTuple t1)) (AD (TTuple t2)) ))

checkForError:: Path -> TypeTree -> Boolean
checkForError p' tt = case p' of
  (End) -> isTypeError $ extractType tt
  (Fst p) -> case tt of
      TBinary op e1 e2 _-> checkForError p e1
      TUnary op e      _-> checkForError p e
      TSectL e op      _-> checkForError p e
      TIfExpr ce te ee _-> checkForError p ce
      TLambda bs body  _-> checkForError p body
      TApp e es        _-> checkForError p e
      _               -> true
  (Snd p) -> case tt of
      TBinary op e1 e2 _-> checkForError p e2
      TSectR op e      _-> checkForError p e
      TIfExpr ce te ee _-> checkForError p te
      _               -> true
  (Thrd p) -> case tt of
      TIfExpr ce te ee _-> checkForError p ee
      _ -> true
  (Nth n p) -> case tt of
      TListTree es  _-> nth n es p
      TNTuple es _-> nth n es p
      TApp e' es _-> nth n es p
      _        -> true
  where
    nth n es p = case (!!) es n of
                  Nothing -> true
                  Just e -> checkForError p e

isTypeError:: Type -> Boolean
isTypeError t = case t of
  (TypeError _) -> true
  _ -> false


copy paste error

copy paste error

module Main where

import           Control.Monad.Eff.Class
import           Control.Monad.Eff.Console
import           Control.Monad.Eff.Random (RANDOM())
import           Data.Foreign.EasyFFI
import           Data.Maybe.Unsafe (fromJust)
import           Node.Express.App
import           Node.Express.Handler
import           Node.Express.Response
import           Prelude hiding (apply)

import ScienceQuestions.Question
import ScienceQuestions.Transitions

-- | The handler for the /question endpoint. Returns a random question.
questionHandler :: forall e. Transitions ->
                             Transitions ->
                             Handler (random :: RANDOM | e)
questionHandler qTransitions aTransitions = do
  -- create a random question (in the Eff context) and lift it into the
  -- Handler context
  question <- liftEff $ randomQuestion 4 (randomNextToken qTransitions)
                                         (randomNextToken aTransitions)
  sendJson question

appSetup :: forall e. Transitions ->
                      Transitions ->
                      App (random :: RANDOM, console :: CONSOLE | e)
appSetup qTransitions aTransitions = do
  liftEff $ log "Setting up"
  -- add our handler at the "/question" endpoint
  get "/question"   (questionHandler qTransitions aTransitions)

main = do
  qTransitions <- fromJust <$> loadTransitions "questions.json"
  aTransitions <- fromJust <$> loadTransitions "answers.json"
  port <- unsafeForeignFunction [""] "process.env.PORT || 8080"
  listenHttp (appSetup qTransitions aTransitions) port \_ ->
    log $ "Listening on " ++ show port

module Main where

import           Control.Monad.Eff.Class
import           Control.Monad.Eff.Console
import           Control.Monad.Eff.Random (RANDOM())
import           Data.Foreign.EasyFFI
import           Data.Maybe.Unsafe (fromJust)
import           Node.Express.App
import           Node.Express.Handler
import           Node.Express.Response
import           Prelude hiding (apply)

import ScienceQuestions.Question
import ScienceQuestions.Transitions

-- | The handler for the /question endpoint. Returns a random question.
type Handler e = HandlerM (express :: EXPRESS | e) Unit
-- |    sendJson :: forall e a. a -> Handler e
questionHandler :: forall e. Transitions ->
                             Transitions ->
                             Handler (random :: RANDOM | e)
questionHandler qTransitions aTransitions = do
  -- create a random question (in the Eff context) and lift it into the
  -- Handler context
  question <- liftEff $ randomQuestion 4 (randomNextToken qTransitions)
                                         (randomNextToken aTransitions)
  sendJson question

appSetup :: forall e. Transitions ->
                      Transitions ->
                      App (random :: RANDOM, console :: CONSOLE | e)
appSetup qTransitions aTransitions = do
  liftEff $ log "Setting up"
  -- add our handler at the "/question" endpoint
  get "/question"   (questionHandler qTransitions aTransitions)

main = do
  qTransitions <- fromJust <$> loadTransitions "questions.json"
  aTransitions <- fromJust <$> loadTransitions "answers.json"
  port <- unsafeForeignFunction [""] "process.env.PORT || 8080"
  listenHttp (appSetup qTransitions aTransitions) port \_ ->
    log $ "Listening on " ++ show port


Make Swift nice again

Make Swift nice again

Original commit 1fa5c854d3f8452745a9f81f0be6414c26f7f510

module Language.Swift
    ( renderer
    ) where

import Doc
import IRGraph
import Prelude

import Data.Array as A
import Data.Char.Unicode (isAlphaNum, isDigit)
import Data.Foldable (for_, null)
import Data.Map (Map)
import Data.Map as M
import Data.Maybe (Maybe(..))
import Data.Set (Set)
import Data.Set as S
import Data.String as String
import Data.String.Util (camelCase, capitalize, decapitalize, genericStringEscape, intToHex, legalizeCharacters, startWithLetter)
import Data.Tuple (Tuple(..))
import Utils (removeElement)

keywords :: Array String
keywords =
    [ "associatedtype", "class", "deinit", "enum", "extension", "fileprivate", "func", "import", "init", "inout", "internal", "let", "open", "operator", "private", "protocol", "public", "static", "struct", "subscript", "typealias", "var"
    , "break", "case", "continue", "default", "defer", "do", "else", "fallthrough", "for", "guard", "if", "in", "repeat", "return", "switch", "where", "while"
    , "as", "Any", "catch", "false", "is", "nil", "rethrows", "super", "self", "Self", "throw", "throws", "true", "try"
    , "_"
    , "associativity", "convenience", "dynamic", "didSet", "final", "get", "infix", "indirect", "lazy", "left", "mutating", "nonmutating", "optional", "override", "postfix", "precedence", "prefix", "Protocol", "required", "right", "set", "Type", "unowned", "weak", "willSet"
    , "String", "Int", "Double", "Bool", "Data", "CommandLine", "FileHandle", "JSONSerialization"
    , "checkNull", "removeNSNull", "nilToNSNull", "convertArray", "convertOptional", "convertDict", "convertDouble"
    ]

renderer :: Renderer
renderer =
    { name: "Swift"
    , aceMode: "swift"
    , extension: "swift"
    , doc: swiftDoc
    , transforms:
        { nameForClass: simpleNamer nameForClass
        , nextName: \s -> "Other" <> s
        , forbiddenNames: keywords
        , topLevelName: forbidNamer (swiftNameStyle true) (\n -> [swiftNameStyle true n])
        , unions: Just
            { predicate: unionIsNotSimpleNullable
            , properName: simpleNamer (swiftNameStyle true <<< combineNames)
            , nameFromTypes: simpleNamer (unionNameIntercalated (swiftNameStyle true) "Or")
            }
        }
    }

legalize :: String -> String
legalize = legalizeCharacters isPartCharacter
    where
        isPartCharacter :: Char -> Boolean
        isPartCharacter c = c == '_' || isAlphaNum c

swiftNameStyle :: Boolean -> String -> String
swiftNameStyle isUpper =
    legalize >>> camelCase >>> startWithLetter isStartCharacter isUpper
    where
        isStartCharacter :: Char -> Boolean
        isStartCharacter c = c == '_' || (isAlphaNum c && not (isDigit c))

nameForClass :: IRClassData -> String
nameForClass (IRClassData { names }) = swiftNameStyle true $ combineNames names

stringEscape :: String -> String
stringEscape =
    genericStringEscape unicodeEscape
    where
        unicodeEscape i =
            "\\u{" <> (String.fromCharArray $ intToHex 0 i) <> "}"

swiftDoc :: Doc Unit
swiftDoc = do
    line "// To parse the JSON, add this file to your project and do:"
    line "//"
    forEachTopLevel_ \topLevelName topLevelType -> do
        typ <- renderType topLevelType
        line $ "//   let " <> decapitalize topLevelName <> " = " <> topLevelName <> "(fromString: jsonString)!"
    blank
    line "import Foundation"
    blank

    renderRenderItems blank (Just renderTopLevelAlias) renderClassDefinition (Just renderUnionDefinition)

    blank
    line $ "// Serialization extensions"

    forEachTopLevel_ renderTopLevelExtensions

    forEachClass_ \className properties -> do
        blank
        renderClassExtension className properties

    forEachUnion_ \unionName unionTypes -> do
        blank
        renderUnionExtension unionName unionTypes

    blank
    supportFunctions

supportFunctions :: Doc Unit
supportFunctions = do
    line """// Helpers

fileprivate func convertArray<T>(_ converter: (Any) -> T?, _ json: Any) -> [T]? {
    guard let jsonArr = json as? [Any] else { return nil }
    var arr = [T]()
    for v in jsonArr {
        if let converted = converter(v) {
            arr.append(converted)
        } else {
            return nil
        }
    }
    return arr
}

fileprivate func convertOptional<T>(_ converter: (Any) -> T?, _ json: Any?) -> T?? {
    guard let v = json else { return .some(nil) }
    return converter(v)
}

fileprivate func convertDict<T>(_ converter: (Any) -> T?, _ json: Any?) -> [String: T]? {
    guard let jsonDict = json as? [String: Any] else { return nil }
    var dict = [String: T]()
    for (k, v) in jsonDict {
        if let converted = converter(v) {
            dict[k] = converted
        } else {
            return nil
        }
    }
    return dict
}

fileprivate func convertToAny<T>(_ dictionary: [String: T], _ converter: (T) -> Any) -> Any {
    var result = [String: Any]()
    for (k, v) in dictionary {
        result[k] = converter(v)
    }
    return result
}

fileprivate func convertDouble(_ v: Any) -> Double? {
    if let w = v as? Double { return w }
    if let w = v as? Int { return Double(w) }
    return nil
}

fileprivate let falseType = NSNumber(value: false).objCType
fileprivate let trueNumber = NSNumber(value: true)
fileprivate let trueType = trueNumber.objCType

fileprivate func convertBool(_ v: Any?) -> Bool? {
    guard let number = v as? NSNumber
    else {
        if let b = v as? Bool {
            return b
        }
        return nil
    }

    if number.objCType != falseType && number.objCType != trueType {
        return nil
    }
    return number.isEqual(trueNumber)
}

fileprivate func removeNSNull(_ v: Any?) -> Any? {
    if let w = v {
        if w is NSNull {
            return nil
        }
        return w
    }
    return nil
}

fileprivate func checkNull(_ v: Any?) -> Any?? {
    if v != nil { return .none }
    return .some(nil)
}"""

renderUnion :: IRUnionRep -> Doc String
renderUnion ur =
    case nullableFromSet $ unionToSet ur of
    Just r -> do
        rendered <- renderType r
        pure $ rendered <> "?"
    Nothing -> lookupUnionName ur

renderType :: IRType -> Doc String
renderType = case _ of
    IRNothing -> pure "Any?"
    IRNull -> pure "Any?"
    IRInteger -> pure "Int"
    IRDouble -> pure "Double"
    IRBool -> pure "Bool"
    IRString -> pure "String"
    IRArray a -> do
        rendered <- renderType a
        pure $ "[" <> rendered <> "]"
    IRClass i -> lookupClassName i
    IRMap t -> do
        rendered <- renderType t
        pure $ "[String: " <> rendered <> "]"
    IRUnion ur -> renderUnion ur

convertAny :: IRType -> String -> Doc String
convertAny (IRArray a) var = do
    converter <- convertAnyFunc a
    pure $ "convertArray(" <> converter <> ", " <> var <> ")"
convertAny (IRMap m) var = do
    converter <- convertAnyFunc m
    pure $ "convertDict(" <> converter <> ", " <> var <> ")"
convertAny (IRUnion ur) var =
    case nullableFromSet $ unionToSet ur of
    Just t -> do
        converter <- convertAnyFunc t
        pure $ "convertOptional(" <> converter <> ", " <> var <> ")"
    Nothing -> do
        name <- lookupUnionName ur
        pure $ name <> ".fromJson(" <> var <> ")"
convertAny IRNothing var =
    pure var
convertAny IRBool var =
    pure $ "convertBool(" <> var <> ")"
convertAny IRInteger var =
    pure $ var <> " as? Int"
convertAny IRString var =
    pure $ var <> " as? String"
convertAny t var = do
    converter <- convertAnyFunc t
    pure $ converter <> "(" <> var <> ")"

convertAnyFunc :: IRType -> Doc String
convertAnyFunc = case _ of
    IRClass i -> do
        name <- lookupClassName i
        pure $ name <> ".init(fromAny:)"
    IRUnion ur ->
        case nullableFromSet $ unionToSet ur of
        Just t -> do
            converter <- convertAnyFunc t
            pure $ "{ (json: Any) in convertOptional(" <> converter <> ", json) }"
        Nothing -> do
            name <- lookupUnionName ur
            pure $ name <> ".fromJson"
    IRDouble -> pure "convertDouble"
    IRNull -> pure "checkNull"
    t -> do
        converted <- convertAny t "$0"
        pure $ "{ " <> converted <> " }"

convertToAny :: IRType -> String -> Doc String
convertToAny (IRArray a) var = do
    convertCode <- convertToAny a "$0"
    pure $ var <> ".map({ " <> convertCode <> " }) as Any"
convertToAny (IRMap m) var = do
    convertCode <- convertToAny m "$0"
    pure $ "convertToAny(" <> var <> ", { "<> convertCode <> " })"
convertToAny (IRClass i) var =
    pure $ var <> ".any"
convertToAny (IRUnion ur) var =
    case nullableFromSet $ unionToSet ur of
    Just t -> do
        convertCode <- convertToAny t "$0"
        pure $ var <> ".map({ " <> convertCode  <> " }) ?? NSNull()"
    Nothing ->
        pure $ var <> ".any"
convertToAny IRNothing var =
    pure $ var <> " ?? NSNull()"
convertToAny IRNull var =
    pure $ "NSNull() as Any"
convertToAny _ var =
    pure $ var <> " as Any"

renderTopLevelAlias :: String -> IRType -> Doc Unit
renderTopLevelAlias topLevelName topLevelType = do
    top <- renderType topLevelType
    line $ "typealias "<> topLevelName <> " = " <> top

renderClassDefinition :: String -> Map String IRType -> Doc Unit
renderClassDefinition className properties = do
    let forbidden = keywords <> ["json", "any"]
    let propertyNames = makePropertyNames properties "" forbidden
    line $ "struct " <> className <> " {"
    indent do
        forEachProperty_ properties propertyNames \_ ptype fieldName _ -> do
            rendered <- renderType ptype
            line $ "let " <> fieldName <> ": " <> rendered
    line "}"

renderTopLevelExtensions :: String -> IRType -> Doc Unit
renderTopLevelExtensions topLevelName topLevelType = do
    blank

    topLevelRendered <- renderType topLevelType
    extensionType <- case topLevelType of
        IRArray t -> ("Array where Element == " <> _) <$> renderType t
        IRMap t -> ("Dictionary where Key == String, Value == " <> _) <$> renderType t
        _ -> pure topLevelRendered

    line $ "extension " <> extensionType <> " {"
    indent do
        line $ "init?(fromString json: String, using encoding: String.Encoding = .utf8) {"
        indent do
            line "guard let data = json.data(using: encoding) else { return nil }"
            line "self.init(fromData: data)"
        line "}"
        blank
        line $ "init?(fromData data: Data) {"
        indent do
            line "guard let json = try? JSONSerialization.jsonObject(with: data, options: []) else { return nil }"
            line "self.init(fromAny: json)"
        line "}"

        case topLevelType of
            IRArray _ ->  do
                blank
                line $ "init?(fromAny untyped: Any) {"
                indent do
                    convertCode <- convertAny topLevelType "untyped"
                    line $ "guard let elements = " <> convertCode <> " else { return nil }"
                    line $ "self.init(elements)"
                line "}"
            IRMap _ -> do
                blank
                line $ "init?(fromAny untyped: Any) {"
                indent do
                    convertCode <- convertAny topLevelType "untyped"
                    line $ "guard let elements = " <> convertCode <> " else { return nil }"
                    line $ "self.init()"
                    line $ "elements.forEach { self[$0.key] = $0.value }"
                line "}"
            _ -> pure unit

        blank
        line $ "var jsonData: Data? {"
        indent do
            convertCode <- convertToAny topLevelType "self"
            line $ "let json = " <> convertCode
            line "return try? JSONSerialization.data(withJSONObject: json, options: [])"
        line "}"

        blank
        line $ "var jsonString: String? {"
        indent do
            line $ "guard let data = self.jsonData else { return nil }"
            line $ "return String(data: data, encoding: .utf8)"
        line "}"

    line "}"

renderClassExtension :: String -> Map String IRType -> Doc Unit
renderClassExtension className properties = do
    let forbidden = keywords <> ["jsonUntyped", "json"]
    let propertyNames = makePropertyNames properties "" forbidden
    line $ "extension " <> className <> " {"
    indent do
        line $ "fileprivate init?(fromAny any: Any) {"
        unless (M.isEmpty properties) $ indent do
            line "guard let json = any as? [String: Any] else { return nil }"
            let forbiddenForUntyped = forbidden <> (A.fromFoldable $ M.keys propertyNames)
            let untypedNames = makePropertyNames properties "Any" forbiddenForUntyped
            let forbiddenForConverted = forbiddenForUntyped <> (A.fromFoldable $ M.keys untypedNames)
            forEachProperty_ properties untypedNames \pname ptype untypedName _ -> do
                when (canBeNull ptype) do
                    line $ "let " <> untypedName <> " = removeNSNull(json[\"" <> stringEscape pname <> "\"])"
            line "guard"
            indent do
                forEachProperty_ properties untypedNames \pname ptype untypedName isLast -> do
                    let convertedName = lookupName pname propertyNames
                    unless (canBeNull ptype) do
                        line $ "let " <> untypedName <> " = removeNSNull(json[\"" <> stringEscape pname <> "\"]),"
                    convertCode <- convertAny ptype untypedName
                    line $ "let " <> convertedName <> " = " <> convertCode <> (if isLast then "" else ",")
                line "else { return nil }"
            forEachProperty_ properties propertyNames \pname _ fieldName _ -> do
                let convertedName = lookupName pname propertyNames
                line $ "self." <> fieldName <> " = " <> convertedName
        line "}"
        blank
        line "fileprivate var any: Any {"
        indent do
            if null properties
                then line "return [String: Any]()"
                else do
                    line "return ["
                    indent do
                        forEachProperty_ properties propertyNames \pname ptype fieldName _ -> do
                            convertCode <- convertToAny ptype ("self." <> fieldName)
                            line $ "\"" <> stringEscape pname <> "\": " <> convertCode <> ","
                    line "]"
        line "}"
    line "}"

makePropertyNames :: Map String IRType -> String -> Array String -> Map String String
makePropertyNames properties suffix forbidden =
    transformPropertyNames (fieldNamer suffix) otherField forbidden properties
    where
        fieldNamer :: String -> Namer String
        fieldNamer suffix = simpleNamer \name -> swiftNameStyle false name <> suffix

        otherField :: String -> String
        otherField name = "other" <> capitalize name

renderUnionDefinition :: String -> Set IRType -> Doc Unit
renderUnionDefinition unionName unionTypes = do
    let { element: emptyOrNull, rest: nonNullTypes } = removeElement (_ == IRNull) unionTypes
    line $ "enum " <> unionName <> " {"
    indent do
        for_ nonNullTypes \typ -> do
            name <- caseName typ
            rendered <- renderType typ
            line $ "case " <> name <> "(" <> rendered <> ")"
        case emptyOrNull of
            Just t -> do
                name <- caseName t
                line $ "case " <> name
            Nothing -> pure unit
    line "}"

renderUnionExtension :: String -> Set IRType -> Doc Unit
renderUnionExtension unionName unionTypes = do
    let { element: emptyOrNull, rest: nonNullTypes } = removeElement (_ == IRNull) unionTypes
    line $ "extension " <> unionName <> " {"
    indent do
        line $ "fileprivate static func fromJson(_ v: Any) -> " <> unionName <> "? {"
        indent do
            case emptyOrNull of
                Just t -> do
                    name <- caseName t
                    line "guard let v = removeNSNull(v)"
                    line "else {"
                    indent do
                        line $ "return ." <> name
                    line "}"
                Nothing -> pure unit
            when (S.member IRBool unionTypes) do
                renderCase IRBool
            when (S.member IRInteger unionTypes) do
                renderCase IRInteger
            for_ (S.difference nonNullTypes $ S.fromFoldable [IRBool, IRInteger]) \typ -> do
                renderCase typ
            line "return nil"
        line "}"
        blank
        line $ "fileprivate var any: Any {"
        indent do
            line $ "switch self {"
            for_ unionTypes \typ -> do
                name <- caseName typ
                let letString = if typ == IRNull then "" else "(let x)"
                convertCode <- convertToAny typ "x"
                line $ "case ." <> name <> letString <> ": return " <> convertCode
            line "}"
        line "}"
    line "}"
    where
    renderCase :: IRType -> Doc Unit
    renderCase t = do
        convertCode <- convertAny t "v"
        name <- caseName t
        line $ "if let x = " <> convertCode <> " { return ." <> name <> "(x) }"

caseName :: IRType -> Doc String
caseName t = swiftNameStyle false <$> getTypeNameForUnion t

module Language.Swift
    ( renderer
    ) where

import Doc
import IRGraph
import Prelude

import Data.Array as A
import Data.Char.Unicode (isAlphaNum, isDigit)
import Data.Foldable (for_, null)
import Data.Map (Map)
import Data.Map as M
import Data.Maybe (Maybe(..))
import Data.Set (Set)
import Data.Set as S
import Data.String as String
import Data.String.Util (camelCase, capitalize, decapitalize, genericStringEscape, intToHex, legalizeCharacters, startWithLetter)
import Data.Tuple (Tuple(..))
import Utils (removeElement)

keywords :: Array String
keywords =
    [ "associatedtype", "class", "deinit", "enum", "extension", "fileprivate", "func", "import", "init", "inout", "internal", "let", "open", "operator", "private", "protocol", "public", "static", "struct", "subscript", "typealias", "var"
    , "break", "case", "continue", "default", "defer", "do", "else", "fallthrough", "for", "guard", "if", "in", "repeat", "return", "switch", "where", "while"
    , "as", "Any", "catch", "false", "is", "nil", "rethrows", "super", "self", "Self", "throw", "throws", "true", "try"
    , "_"
    , "associativity", "convenience", "dynamic", "didSet", "final", "get", "infix", "indirect", "lazy", "left", "mutating", "nonmutating", "optional", "override", "postfix", "precedence", "prefix", "Protocol", "required", "right", "set", "Type", "unowned", "weak", "willSet"
    , "String", "Int", "Double", "Bool", "Data", "CommandLine", "FileHandle", "JSONSerialization"
    , "checkNull", "removeNSNull", "nilToNSNull", "convertArray", "convertOptional", "convertDict", "convertDouble"
    ]

renderer :: Renderer
renderer =
    { name: "Swift"
    , aceMode: "swift"
    , extension: "swift"
    , doc: swiftDoc
    , transforms:
        { nameForClass: simpleNamer nameForClass
        , nextName: \s -> "Other" <> s
        , forbiddenNames: keywords
        , topLevelName: forbidNamer (swiftNameStyle true) (\n -> [swiftNameStyle true n])
        , unions: Just
            { predicate: unionIsNotSimpleNullable
            , properName: simpleNamer (swiftNameStyle true <<< combineNames)
            , nameFromTypes: simpleNamer (unionNameIntercalated (swiftNameStyle true) "Or")
            }
        }
    }

legalize :: String -> String
legalize = legalizeCharacters isPartCharacter
    where
        isPartCharacter :: Char -> Boolean
        isPartCharacter c = c == '_' || isAlphaNum c

swiftNameStyle :: Boolean -> String -> String
swiftNameStyle isUpper =
    legalize >>> camelCase >>> startWithLetter isStartCharacter isUpper
    where
        isStartCharacter :: Char -> Boolean
        isStartCharacter c = c == '_' || (isAlphaNum c && not (isDigit c))

nameForClass :: IRClassData -> String
nameForClass (IRClassData { names }) = swiftNameStyle true $ combineNames names

stringEscape :: String -> String
stringEscape =
    genericStringEscape unicodeEscape
    where
        unicodeEscape i =
            "\\u{" <> (String.fromCharArray $ intToHex 0 i) <> "}"

swiftDoc :: Doc Unit
swiftDoc = do
    line "// To parse the JSON, add this file to your project and do:"
    line "//"
    forEachTopLevel_ \topLevelName topLevelType -> do
        typ <- renderType topLevelType
        line $ "//   let " <> decapitalize topLevelName <> " = " <> topLevelName <> "(fromString: jsonString)!"
    blank
    line "import Foundation"
    blank

    renderRenderItems blank (Just renderTopLevelAlias) renderClassDefinition (Just renderUnionDefinition)

    blank
    line $ "// Serialization extensions"

    forEachTopLevel_ renderTopLevelExtensions

    forEachClass_ \className properties -> do
        blank
        renderClassExtension className properties

    forEachUnion_ \unionName unionTypes -> do
        blank
        renderUnionExtension unionName unionTypes

    blank
    supportFunctions

supportFunctions :: Doc Unit
supportFunctions = do
    line """// Helpers

fileprivate func convertArray<T>(_ converter: (Any) -> T?, _ json: Any) -> [T]? {
    guard let jsonArr = json as? [Any] else { return nil }
    var arr = [T]()
    for v in jsonArr {
        if let converted = converter(v) {
            arr.append(converted)
        } else {
            return nil
        }
    }
    return arr
}

fileprivate func convertOptional<T>(_ converter: (Any) -> T?, _ json: Any?) -> T?? {
    guard let v = json else { return .some(nil) }
    return converter(v)
}

fileprivate func convertDict<T>(_ converter: (Any) -> T?, _ json: Any?) -> [String: T]? {
    guard let jsonDict = json as? [String: Any] else { return nil }
    var dict = [String: T]()
    for (k, v) in jsonDict {
        if let converted = converter(v) {
            dict[k] = converted
        } else {
            return nil
        }
    }
    return dict
}

fileprivate func convertToAny<T>(_ dictionary: [String: T], _ converter: (T) -> Any) -> Any {
    var result = [String: Any]()
    for (k, v) in dictionary {
        result[k] = converter(v)
    }
    return result
}

fileprivate func convertDouble(_ v: Any) -> Double? {
    if let w = v as? Double { return w }
    if let w = v as? Int { return Double(w) }
    return nil
}

fileprivate let falseType = NSNumber(value: false).objCType
fileprivate let trueNumber = NSNumber(value: true)
fileprivate let trueType = trueNumber.objCType

fileprivate func convertBool(_ v: Any?) -> Bool? {
    guard let number = v as? NSNumber
    else {
        if let b = v as? Bool {
            return b
        }
        return nil
    }

    if number.objCType != falseType && number.objCType != trueType {
        return nil
    }
    return number.isEqual(trueNumber)
}

fileprivate func removeNSNull(_ v: Any?) -> Any? {
    if let w = v {
        if w is NSNull {
            return nil
        }
        return w
    }
    return nil
}

fileprivate func checkNull(_ v: Any?) -> Any?? {
    if v != nil { return .none }
    return .some(nil)
}"""

renderUnion :: IRUnionRep -> Doc String
renderUnion ur =
    case nullableFromSet $ unionToSet ur of
    Just r -> do
        rendered <- renderType r
        pure $ rendered <> "?"
    Nothing -> lookupUnionName ur

renderType :: IRType -> Doc String
renderType = case _ of
    IRNothing -> pure "Any?"
    IRNull -> pure "Any?"
    IRInteger -> pure "Int"
    IRDouble -> pure "Double"
    IRBool -> pure "Bool"
    IRString -> pure "String"
    IRArray a -> do
        rendered <- renderType a
        pure $ "[" <> rendered <> "]"
    IRClass i -> lookupClassName i
    IRMap t -> do
        rendered <- renderType t
        pure $ "[String: " <> rendered <> "]"
    IRUnion ur -> renderUnion ur

convertAny :: IRType -> String -> Doc String
convertAny (IRArray a) var = do
    converter <- convertAnyFunc a
    pure $ "convertArray(" <> converter <> ", " <> var <> ")"
convertAny (IRMap m) var = do
    converter <- convertAnyFunc m
    pure $ "convertDict(" <> converter <> ", " <> var <> ")"
convertAny (IRUnion ur) var =
    case nullableFromSet $ unionToSet ur of
    Just t -> do
        converter <- convertAnyFunc t
        pure $ "convertOptional(" <> converter <> ", " <> var <> ")"
    Nothing -> do
        name <- lookupUnionName ur
        pure $ name <> ".fromJson(" <> var <> ")"
convertAny IRNothing var =
    pure var
convertAny IRBool var =
    pure $ "convertBool(" <> var <> ")"
convertAny IRInteger var =
    pure $ var <> " as? Int"
convertAny IRString var =
    pure $ var <> " as? String"
convertAny t var = do
    converter <- convertAnyFunc t
    pure $ converter <> "(" <> var <> ")"

convertAnyFunc :: IRType -> Doc String
convertAnyFunc = case _ of
    IRClass i -> do
        name <- lookupClassName i
        -- TODO make this look less alien
        pure $ "{ " <> name <> "(fromAny: $0) }"
    IRUnion ur ->
        case nullableFromSet $ unionToSet ur of
        Just t -> do
            converter <- convertAnyFunc t
            pure $ "{ (json: Any) in convertOptional(" <> converter <> ", json) }"
        Nothing -> do
            name <- lookupUnionName ur
            pure $ name <> ".fromJson"
    IRDouble -> pure "convertDouble"
    IRNull -> pure "checkNull"
    t -> do
        converted <- convertAny t "$0"
        pure $ "{ " <> converted <> " }"

convertToAny :: IRType -> String -> Doc String
convertToAny (IRArray a) var = do
    convertCode <- convertToAny a "$0"
    pure $ var <> ".map({ " <> convertCode <> " }) as Any"
convertToAny (IRMap m) var = do
    convertCode <- convertToAny m "$0"
    pure $ "convertToAny(" <> var <> ", { "<> convertCode <> " })"
convertToAny (IRClass i) var =
    pure $ var <> ".any"
convertToAny (IRUnion ur) var =
    case nullableFromSet $ unionToSet ur of
    Just t -> do
        convertCode <- convertToAny t "$0"
        pure $ var <> ".map({ " <> convertCode  <> " }) ?? NSNull()"
    Nothing ->
        pure $ var <> ".any"
convertToAny IRNothing var =
    pure $ var <> " ?? NSNull()"
convertToAny IRNull var =
    pure $ "NSNull() as Any"
convertToAny _ var =
    pure $ var <> " as Any"

renderTopLevelAlias :: String -> IRType -> Doc Unit
renderTopLevelAlias topLevelName topLevelType = do
    top <- renderType topLevelType
    line $ "typealias "<> topLevelName <> " = " <> top

renderClassDefinition :: String -> Map String IRType -> Doc Unit
renderClassDefinition className properties = do
    let forbidden = keywords <> ["json", "any"]
    let propertyNames = makePropertyNames properties "" forbidden
    line $ "struct " <> className <> " {"
    indent do
        forEachProperty_ properties propertyNames \_ ptype fieldName _ -> do
            rendered <- renderType ptype
            line $ "let " <> fieldName <> ": " <> rendered
    line "}"

renderTopLevelExtensions :: String -> IRType -> Doc Unit
renderTopLevelExtensions topLevelName topLevelType = do
    blank

    topLevelRendered <- renderType topLevelType
    extensionType <- case topLevelType of
        IRArray t -> ("Array where Element == " <> _) <$> renderType t
        IRMap t -> ("Dictionary where Key == String, Value == " <> _) <$> renderType t
        _ -> pure topLevelRendered

    line $ "extension " <> extensionType <> " {"
    indent do
        line $ "init?(fromString json: String, using encoding: String.Encoding = .utf8) {"
        indent do
            line "guard let data = json.data(using: encoding) else { return nil }"
            line "self.init(fromData: data)"
        line "}"
        blank
        line $ "init?(fromData data: Data) {"
        indent do
            line "guard let json = try? JSONSerialization.jsonObject(with: data, options: []) else { return nil }"
            line "self.init(fromAny: json)"
        line "}"

        case topLevelType of
            IRArray _ ->  do
                blank
                line $ "init?(fromAny untyped: Any) {"
                indent do
                    convertCode <- convertAny topLevelType "untyped"
                    line $ "guard let elements = " <> convertCode <> " else { return nil }"
                    line $ "self.init(elements)"
                line "}"
            IRMap _ -> do
                blank
                line $ "init?(fromAny untyped: Any) {"
                indent do
                    convertCode <- convertAny topLevelType "untyped"
                    line $ "guard let elements = " <> convertCode <> " else { return nil }"
                    line $ "self.init()"
                    line $ "elements.forEach { self[$0.key] = $0.value }"
                line "}"
            _ -> pure unit

        blank
        line $ "var jsonData: Data? {"
        indent do
            convertCode <- convertToAny topLevelType "self"
            line $ "let json = " <> convertCode
            line "return try? JSONSerialization.data(withJSONObject: json, options: [])"
        line "}"

        blank
        line $ "var jsonString: String? {"
        indent do
            line $ "guard let data = self.jsonData else { return nil }"
            line $ "return String(data: data, encoding: .utf8)"
        line "}"

    line "}"

renderClassExtension :: String -> Map String IRType -> Doc Unit
renderClassExtension className properties = do
    let forbidden = keywords <> ["jsonUntyped", "json"]
    let propertyNames = makePropertyNames properties "" forbidden
    line $ "extension " <> className <> " {"
    indent do
        line $ "fileprivate init?(fromAny any: Any) {"
        unless (M.isEmpty properties) $ indent do
            line "guard let json = any as? [String: Any] else { return nil }"
            let forbiddenForUntyped = forbidden <> (A.fromFoldable $ M.keys propertyNames)
            let untypedNames = makePropertyNames properties "Any" forbiddenForUntyped
            let forbiddenForConverted = forbiddenForUntyped <> (A.fromFoldable $ M.keys untypedNames)
            forEachProperty_ properties untypedNames \pname ptype untypedName _ -> do
                when (canBeNull ptype) do
                    line $ "let " <> untypedName <> " = removeNSNull(json[\"" <> stringEscape pname <> "\"])"
            line "guard"
            indent do
                forEachProperty_ properties untypedNames \pname ptype untypedName isLast -> do
                    let convertedName = lookupName pname propertyNames
                    unless (canBeNull ptype) do
                        line $ "let " <> untypedName <> " = removeNSNull(json[\"" <> stringEscape pname <> "\"]),"
                    convertCode <- convertAny ptype untypedName
                    line $ "let " <> convertedName <> " = " <> convertCode <> (if isLast then "" else ",")
                line "else { return nil }"
            forEachProperty_ properties propertyNames \pname _ fieldName _ -> do
                let convertedName = lookupName pname propertyNames
                line $ "self." <> fieldName <> " = " <> convertedName
        line "}"
        blank
        line "fileprivate var any: Any {"
        indent do
            if null properties
                then line "return [String: Any]()"
                else do
                    line "return ["
                    indent do
                        forEachProperty_ properties propertyNames \pname ptype fieldName _ -> do
                            convertCode <- convertToAny ptype ("self." <> fieldName)
                            line $ "\"" <> stringEscape pname <> "\": " <> convertCode <> ","
                    line "]"
        line "}"
    line "}"

makePropertyNames :: Map String IRType -> String -> Array String -> Map String String
makePropertyNames properties suffix forbidden =
    transformPropertyNames (fieldNamer suffix) otherField forbidden properties
    where
        fieldNamer :: String -> Namer String
        fieldNamer suffix = simpleNamer \name -> swiftNameStyle false name <> suffix

        otherField :: String -> String
        otherField name = "other" <> capitalize name

renderUnionDefinition :: String -> Set IRType -> Doc Unit
renderUnionDefinition unionName unionTypes = do
    let { element: emptyOrNull, rest: nonNullTypes } = removeElement (_ == IRNull) unionTypes
    line $ "enum " <> unionName <> " {"
    indent do
        for_ nonNullTypes \typ -> do
            name <- caseName typ
            rendered <- renderType typ
            line $ "case " <> name <> "(" <> rendered <> ")"
        case emptyOrNull of
            Just t -> do
                name <- caseName t
                line $ "case " <> name
            Nothing -> pure unit
    line "}"

renderUnionExtension :: String -> Set IRType -> Doc Unit
renderUnionExtension unionName unionTypes = do
    let { element: emptyOrNull, rest: nonNullTypes } = removeElement (_ == IRNull) unionTypes
    line $ "extension " <> unionName <> " {"
    indent do
        line $ "fileprivate static func fromJson(_ v: Any) -> " <> unionName <> "? {"
        indent do
            case emptyOrNull of
                Just t -> do
                    name <- caseName t
                    line "guard let v = removeNSNull(v)"
                    line "else {"
                    indent do
                        line $ "return ." <> name
                    line "}"
                Nothing -> pure unit
            when (S.member IRBool unionTypes) do
                renderCase IRBool
            when (S.member IRInteger unionTypes) do
                renderCase IRInteger
            for_ (S.difference nonNullTypes $ S.fromFoldable [IRBool, IRInteger]) \typ -> do
                renderCase typ
            line "return nil"
        line "}"
        blank
        line $ "fileprivate var any: Any {"
        indent do
            line $ "switch self {"
            for_ unionTypes \typ -> do
                name <- caseName typ
                let letString = if typ == IRNull then "" else "(let x)"
                convertCode <- convertToAny typ "x"
                line $ "case ." <> name <> letString <> ": return " <> convertCode
            line "}"
        line "}"
    line "}"
    where
    renderCase :: IRType -> Doc Unit
    renderCase t = do
        convertCode <- convertAny t "v"
        name <- caseName t
        line $ "if let x = " <> convertCode <> " { return ." <> name <> "(x) }"

caseName :: IRType -> Doc String
caseName t = swiftNameStyle false <$> getTypeNameForUnion t