module Language.Sifflet.TypeCheck
    (valueType, typeCheckValues
    , typeVarsIn
    , Subst
    , subst, substComp, idSubst, deltaSubst
    , extend, unify
    , TypeScheme(..), schematicsTS, unknownsTS, substTS
    , TypeEnv, emptyTypeEnv, install, unknownsTE, substTE
    , NameSupply(..), newNameSupply, nameSupplyNext, nameSupplyTake
    , tcExpr, tcExprs
    , tcVar, newInstance
    , arrow
    , envToTypeEnv, baseTypeEnv
    , fromLambdaType, decideTypes
    )

where

-- drop after debugging:
import Debug.Trace

import Control.Monad (foldM)

import qualified Data.List as AList (lookup)

import Data.Map as Map hiding (filter, foldr, lookup, map, null)
import qualified Data.Map as Map (map, lookup)

import Text.Printf

import Language.Sifflet.Expr
import Text.Sifflet.Repr
import Language.Sifflet.Util


-- | Determine the type of a value.
-- May result in a type variable.

valueType :: Value -> SuccFail Type
valueType v =
    case v of
      VBool _ -> Succ typeBool
      VChar _ -> Succ typeChar
      VNumber _ -> Succ typeNum
      VString _ -> Succ typeString
      VFun (Function _ atypes rtype _) -> Succ $ typeFunction atypes rtype

      VList []  -> Succ $ typeList $ TypeVar "list_element"
      VList (x:xs) -> 
          do
            xtype <- valueType x
            xstypes <- mapM valueType xs
            if filter (/= xtype) xstypes == []
               then Succ $ typeList xtype
               else Fail "list with diverse element types"

-- | Check whether the values agree with the types (which may be abstract)
--
-- This is *probably* too lenient in the case of type variables:
-- it can pass a mixed-type list.

typeCheckValues :: [String] -> [Type] -> [Value] -> SuccFail [Value]
typeCheckValues names types values =
    let check :: (Map String Type) -> [String] -> [Type] -> [Value] 
              -> SuccFail [Value]
        check _ [] [] [] = Succ []
        check env (n:ns) (t:ts) (v:vs) = 
            case typeMatch t v env of
              Succ env' -> check env' ns ts vs >>= Succ . (v:)
              Fail msg -> Fail $ "For variable " ++ n ++ ":\n" ++ msg
        check _ _ _ _ = error "typeCheckValues: mismatched list lengths"
    in check empty names types values
       


-- | Try to match a single type and value,
-- may result in binding a type variable in a new environment
-- or just the old environment
-- (Note that this tries to produce a map of type variable names to Types,
-- instead of TypeSchemes as with TypeEnv below)
typeMatch :: Type -> Value -> (Map String Type) -> SuccFail (Map String Type)
typeMatch atype value env = 
    let sorry x etype =
            Fail $ repr x ++ ": " ++ etype ++ " expected"
    in case (atype, value) of
      -- easy cases
      (TypeCons "Bool" [], VBool _) -> Succ env
      (TypeCons "Bool" [], x) -> sorry x "True or False"
      (TypeCons "Char" [], VChar _) -> Succ env
      (TypeCons "Char" [], x) -> sorry x "character"
      (TypeCons "Num" [], VNumber _) -> Succ env
      (TypeCons "Num" [], x) -> sorry x "number"
      (TypeCons "String" [], VString _) -> Succ env
      (TypeCons "String" [], x) -> sorry x "string"
      -- VV Harder
      -- VV Are the avalues below supposed to be equal to the value above?
      (TypeVar name, avalue) -> 
          case Map.lookup name env of
            Nothing -> 
                -- bind type variable
                valueType avalue >>= \ vtype -> 
                    Succ $ Map.insert name vtype env
            Just concreteType -> typeMatch concreteType avalue env
      (TypeCons "List" [etype], VList lvalues) ->
          case lvalues of
            [] -> Succ env
            v:vs -> 
                typeMatch etype v env >>= 
                typeMatch (typeList etype) (VList vs)
      (TypeCons "Function" [_, _], _) ->
          -- this will require matching type variables with type variables!
          error "typeMatch: unimplemented case for TypeCons Function"
      _ -> Fail $ "type mismatch: " ++ show (atype, value)


-- The rest of this is based (loosely and conceptually) on the type checker
-- chapters 8-9 of Simon L. Peyton Jones, The Implementation of Functional
-- Programming Languages, Prentice Hall, 1987.
-- Out of print but available on the web at
-- http://research.microsoft.com/en-us/um/people/simonpj/papers/slpj-book-1987/
--
-- Chapter 8: Polymorphic Type Checking, by Peter Hancock.
-- Chapter 9: A Type-Checker, by Peter Hancock.

-- 9.2 Representation of Type Expressions
-- ('a -> b' means the textbook notation 'a' corresponds (at least roughly)
-- to the Sifflet code 'b'.)
-- tvname -> Language.Sifflet.Expr.TypeVarName
-- type_exp -> Language.Sifflet.Expr.Type
-- tvars_in -> typeVarsIn

typeVarsIn :: Type -> [TypeVarName]
typeVarsIn atype =
    let tvarsIn atype' result =
            case atype' of
              TypeVar vname -> vname : result
              TypeCons _ctor ts -> foldr tvarsIn result ts
              -- TypeList t -> foldr tvarsIn result [t]
              -- TypeFunction ats rt -> foldr tvarsIn result (rt:ats)
              -- _ -> []
    in tvarsIn atype []

-- 9.3 Success and Failure
-- reply -> Language.Sifflet.Util.SuccFail

-- 9.4 Solving Equations
-- 9.4.1 Substitutions
-- subst -> Subst
-- sub_type -> subst
-- scomp -> substComp
-- id_subst -> idSubst
-- delta -> deltaSubst

-- | The type of a substitution function, which maps type variables to types

type Subst = TypeVarName -> Type

-- | Apply a substitution function to a type (possibly a type with variables)

subst :: Subst -> Type -> Type
subst phi typeExpr =
    case typeExpr of
      TypeVar tvn -> phi tvn
      -- TypeList etype -> TypeList (subst phi etype)
      -- TypeFunction atypes rtype -> 
      --    TypeFunction (map (subst phi) atypes) (subst phi rtype)
      TypeCons ctor atypes ->
         TypeCons ctor (map (subst phi) atypes)

-- | Composing two substitutions

substComp :: Subst -> Subst -> Subst
substComp sub2 sub1 tvn = subst sub2 (sub1 tvn)

-- | Identity substitution

idSubst :: Subst
idSubst tvn = TypeVar tvn

-- | A delta substitution is one that affects just a single variable.

deltaSubst :: TypeVarName -> Type -> Subst
deltaSubst tvn typeExpr tvn' =
    if tvn == tvn'
    then typeExpr
    else TypeVar tvn'

-- 9.4.2.  Unification
-- extend -> extend
-- unify -> unify
-- unifyl -> unifyList

-- | Try to extend a substitution by adding a single variable-value pair

extend :: Subst -> TypeVarName -> Type -> SuccFail Subst
extend phi tvn t =
    if t == TypeVar tvn
    then Succ phi
    else if tvn `elem` typeVarsIn t
         then Fail ("occurs check: " ++ tvn ++ " is in " ++ (show t))
         else Succ (deltaSubst tvn t `substComp` phi)

-- extend is untested; we can wait until unify to test

unify :: Subst -> (Type, Type) -> SuccFail Subst
unify phi (TypeVar tvn, t) =
    if phi tvn == TypeVar tvn
    then extend phi tvn (subst phi t)
    else unify phi (phi tvn, subst phi t)
unify phi (TypeCons tcn ts, TypeVar tvn) =
    unify phi (TypeVar tvn, TypeCons tcn ts)
unify phi types@(TypeCons tcn ts, TypeCons tcn' ts') =
    if tcn == tcn'
    then unifyList phi (zip ts ts')
    else Fail ("cannot unify " ++ show (fst types) ++ 
               " with " ++ show (snd types))

unifyList :: Subst -> [(Type, Type)] -> SuccFail Subst
unifyList = foldM unify

-- 9.5 Keeping track of types
-- 9.5.2 Look to the variables
--   type_scheme -> TypeScheme
--   SCHEME -> TypeScheme
--   unknowns_scheme -> unknownsTS
--   sub_scheme -> substTS

-- | A TypeScheme (TypeScheme schematicVariables typeExpression)
-- is a sort of type template, in which
-- schematic variables (i.e., type parameters or "generics")
-- are made explicit; any other type variables
-- in the type expression are unknowns

data TypeScheme = TypeScheme [TypeVarName] Type
    deriving (Eq, Show)

schematicsTS :: TypeScheme -> [TypeVarName]
schematicsTS (TypeScheme schVars _) = schVars

unknownsTS :: TypeScheme -> [TypeVarName]
unknownsTS (TypeScheme schVars t) =
    [tv | tv <- typeVarsIn t, not (tv `elem` schVars)]

-- | Apply a substitution to a type scheme, taking care to affect
-- only the unknowns, not the schematic variables

substTS :: Subst -> TypeScheme -> TypeScheme
substTS phi (TypeScheme schVars t) =
    let phi' :: Subst
        phi' tvn =
            -- phi' is like phi but excludes the schematic variables
            if tvn `elem` schVars
            then TypeVar tvn
            else phi tvn
    in TypeScheme schVars (subst phi' t)

-- 9.5.3 Type environments
-- I'll use maps instead of association lists
--   type_env -> TypeEnv
--   unknowns_te -> unknownsTE
--   sub_te -> substTE

-- | A type environment maps type variable names to type schemes.
-- Oh, is it really type variable names, or just names?
-- It seems to me that in envToTypeEnv, it's function names,
-- instead of type variable names.

type TypeEnv = Map TypeVarName TypeScheme

emptyTypeEnv :: TypeEnv
emptyTypeEnv = Map.empty

-- -- | The domain of a type environment

-- dom :: TypeEnv -> [TypeVarName]
-- dom = keys

-- | The range of a type environment

rng :: TypeEnv -> [TypeScheme]
rng = elems

-- -- | The type scheme value of a type variable name
-- -- This will fail, of course, if the variable name is not
-- -- bound in the type environment.
-- val :: TypeEnv -> TypeVarName -> TypeScheme
-- val = (!)

-- | Insert a new type variable and its type scheme value
install :: TypeEnv -> TypeVarName -> TypeScheme -> TypeEnv
install te tvn ts = insert tvn ts te

-- | The unknowns of a type environment

unknownsTE :: TypeEnv -> [TypeVarName]
unknownsTE gamma = concatMap unknownsTS (rng gamma)

substTE :: Subst -> TypeEnv -> TypeEnv
substTE phi gamma = Map.map (substTS phi) gamma

-- 9.6 New variables
-- I diverge from the "easy" way given in the book

-- | A source of variable names

data NameSupply = NameSupply TypeVarName Int -- prefix, count
                deriving (Eq, Show)

-- | Creates a name supply

newNameSupply :: TypeVarName -> NameSupply
newNameSupply prefix = NameSupply prefix 1

-- | Produces next variable from a name supply

nameSupplyNext :: NameSupply -> (TypeVarName, NameSupply)
nameSupplyNext (NameSupply prefix count) =
    (prefix ++ show count, NameSupply prefix (count + 1))

-- | Produces next several variables from a name supply

nameSupplyTake :: NameSupply -> Int -> ([TypeVarName], NameSupply)
nameSupplyTake (NameSupply prefix count) n =
    ([prefix ++ show i | i <- [count .. count+n-1]],
     NameSupply prefix (count + n))

-- 9.7 The type checker
--   tc -> tcExpr
--   tcl -> tcExprs (type checks a list of Expr)
--   
--   Additional functions (?) tcExprTree ??

-- | Type check an expression

tcExpr :: TypeEnv -> NameSupply -> Expr -> SuccFail (Subst, Type, NameSupply)
tcExpr te ns expr =
    let prim ptype = Succ (idSubst, ptype, ns)
    in case expr of
         EUndefined -> 
             let (tvn, ns') = nameSupplyNext ns
             in Succ (idSubst, TypeVar tvn, ns')

         EBool _ -> prim typeBool
         EChar _ -> prim typeChar
         ENumber _ -> prim typeNum
         EString _ -> prim typeString

         ESymbol (Symbol name) -> tcVar te ns name
         EIf c a b -> tcIf te ns c a b
         EList exprs -> tcList te ns exprs
         ELambda (Symbol x) body -> tcLambda te ns x body
         EApp f arg -> tcApp te ns f arg
         ECall _ _ -> tcExpr te ns (callToApp expr)

         -- EOp and EGroup are not needed in Sifflet itself,
         -- only for export
         EOp _op _ _ -> Fail "tcExpr: not implemented for EOp"
         EGroup e -> tcExpr te ns e

-- 9.7.1 Type checking lists of expressions
--   tcl -> tcExprs
--   tcl1 -> tcExprs1
--   tcl2 -> tcExprs2
-- This occurs in "two stages":
-- 1.  Type check each expression in the list, building up the 
--     type environment as we scan through it.
-- 2.  Cumulatively compose all the substitutions found,
--     and make sure each type found is a fixed point of the
--     final substitution.
-- NOTE: This is to check a list of Exprs, not an Expr of the type EList.
-- See tcList.

tcExprs :: TypeEnv -> NameSupply -> [Expr] 
        -> SuccFail (Subst, [Type], NameSupply)
tcExprs _te ns [] = Succ (idSubst, [], ns)
tcExprs te ns (e:es) = do
  {
    (phi, t, ns') <- tcExpr te ns e
  ; tcExprs1 te ns' es phi t
  }

tcExprs1 :: TypeEnv -> NameSupply -> [Expr] -> Subst -> Type
        -> SuccFail (Subst, [Type], NameSupply)
tcExprs1 te ns es phi t = do
  {
    (psi, ts, ns') <- tcExprs (substTE phi te) ns es
  ; tcExprs2 phi t psi ts ns'
  }

tcExprs2 :: Subst -> Type -> Subst -> [Type] -> NameSupply
         -> SuccFail (Subst, [Type], NameSupply)
tcExprs2 phi t psi ts ns =
    Succ (substComp psi phi, subst psi t : ts, ns)

-- 9.7.2 Type checking variables
--   tcvar -> tcVar
--   newinstance -> newInstance
--   al_to_subst -> alistToSubst

-- | Type check a variable (actually the variable name).
-- Find the type scheme associated with the variable in the type environment.
-- Return a "new instance" of the type scheme, in which its schematic variables
-- are replaced by new variables, and the unknowns are left as they are.

tcVar :: TypeEnv -> NameSupply -> String -> SuccFail (Subst, Type, NameSupply)
tcVar te ns name =
    case Map.lookup name te of
      Nothing -> 
          Fail ("BUG in type checker: tcVar: unable to find " ++ name ++ 
                " in type environment")
      Just ntype ->
          let (t, ns') = newInstance ns ntype
          in Succ (idSubst, t, ns')

newInstance :: NameSupply -> TypeScheme -> (Type, NameSupply)
newInstance ns (TypeScheme schVars t) =
    let (names, ns') = nameSupplyTake ns (length schVars)
        alist = zip schVars names
        phi = alistToSubst alist
    in (subst phi t, ns')

-- | Convert an association list of type variable names into a substitution

alistToSubst :: [(TypeVarName, TypeVarName)] -> Subst
alistToSubst alist tvn =
    case AList.lookup tvn alist of
      Just tvn' -> TypeVar tvn'
      Nothing -> TypeVar tvn

-- 9.7.3 Type checking application
--   tcap -> tcApp
--   a $arrow b -> a `arrow` b = TypeCons "Function" [a, b]

tcApp :: TypeEnv -> NameSupply -> Expr -> Expr
      -> SuccFail (Subst, Type, NameSupply)
tcApp te ns e1 e2 = 
    let (tvn, ns') = nameSupplyNext ns
    in do
      (phi, [t1, t2], ns'') <- tcExprs te ns' [e1, e2]
      phi' <- unify phi (t1, t2 `arrow` (TypeVar tvn))
      Succ (phi', phi' tvn, ns'')

arrow :: Type -> Type -> Type
arrow a b = TypeCons "Function" [a, b]
    

-- 9.7.4 Type checking lambda abstractions
-- 
--   LAMBDA -> ELambda

tcLambda :: TypeEnv -> NameSupply -> String -> Expr
         -> SuccFail (Subst, Type, NameSupply)
tcLambda te ns x e =
    let (tvn, ns') = nameSupplyNext ns
        te' = install te x (TypeScheme [] (TypeVar tvn))
    in do
      {
        (phi, t, ns'') <- tcExpr te' ns' e
      ; Succ (phi, phi tvn `arrow` t, ns'')
      } 

-- 9.7.5 Type checking let expressions
-- Not needed, at least for the present 

-- 9.7.6 Type checking letrec expressions
-- Not needed, at least for the present 



-- Extras (not in the textbook)

-- | Decide the type of a function -- called by the function editor
-- when the Apply button is clicked.
-- decideTypes tries to find the argument types and return type
-- of an expression considered as the body of a function,
-- at the same time checking for consistency of inputs and
-- outputs between the parts of the expression.
-- It returns Succ (argtypes, returntype) if successful;
-- Fail errormessage otherwise.

decideTypes :: String -> Expr -> [String] -> Env -> SuccFail([Type], Type)
decideTypes name body args env =
    let te = envToTypeEnv env
        -- extend te to bind name to a type variable that will not be used
        te' = install te name (TypeScheme [] (TypeVar "u"))
        result :: SuccFail (Type, ([Type], Type))
        result = do
          expr <- traceValue "decideTypes.expr" $ toLambdaExpr args body
          (_, t, _) <- tcExpr te' (newNameSupply "t") expr
          _ <- traceValue "decideTypes.t" $ return t
          res <- traceValue "decideTypes.res" $ fromLambdaType t
          Succ (t, res)
        finalResult = 
          case result of
            Succ (t, res@(argTypes, _retType)) ->
              if length args == length argTypes
              then Succ res
              else Fail ("Inferred function type " ++ show t ++
                         " does not agree with number of arguments")
            Fail msg -> Fail msg
    in traceValue (printf "decideTypes: %s ::\n" (show (args, body))) 
                  finalResult


traceValue :: (Show a) => String -> a -> a
traceValue label value =
  trace (printf "\n%s: %s\n" label (show value)) value

-- | Convert a function type, such as a -> (b -> c), to
-- a pair consisting of the list of argument types and the
-- result type, such as ([a, b], c).

fromLambdaType :: Type -> SuccFail ([Type], Type)
fromLambdaType t = 
    case t of
      TypeCons "Function" [t1, t2] -> 
          case t2 of
            TypeCons "Function" _ -> do
              (atypes, rtype) <- fromLambdaType t2
              return (t1:atypes, rtype)
            _ -> Succ ([t1], t2)
      TypeCons "Function" _ -> 
          error ("fromLambdaType: invalid function type " ++ show t)
      _ -> error ("fromLambdaType: non-function type " ++ show t)

-- | Type check an IF expression
-- *** Failure messages need improvement ***

tcIf :: TypeEnv -> NameSupply -> Expr -> Expr -> Expr 
     -> SuccFail (Subst, Type, NameSupply)
tcIf te ns expr1 expr2 expr3 = do
  {
    (phi, ts, ns') <- tcExprs te ns [expr1, expr2, expr3]
  ; -- assert ts has length 3
    phi' <- unifyList phi [(head ts, typeBool), (ts !! 1, ts !! 2)]
  ; return (phi', subst phi' (ts !! 1), ns')
  }

-- | Type check a list expression.
-- NOTE: This is to check a single Expr with the EList type, 
-- in which all elements must be of the same type,
-- not a general list of Exprs of possibly mixed type.  
-- See tcExprs.
--
-- *** Failure messages need improvement ***
--
-- IDEA: can this be reduced to [] and application of (:)?

tcList :: TypeEnv -> NameSupply -> [Expr]
       -> SuccFail (Subst, Type, NameSupply)
tcList te ns exprs =
    case exprs of
      [] -> 
          let (tvn, ns') = nameSupplyNext ns
          in Succ (idSubst, typeList (TypeVar tvn), ns')
      e:es ->
          do
            {
              (phi, ts, ns') <- tcExprs te ns [e, EList es]
            ; -- assert ts has length 2
              phi' <- unify phi (typeList (head ts), ts !! 1)
            ; return (phi', subst phi' (ts !! 1), ns')
            }



-- | Create a TypeEnv from an Env
-- Bindings on the left replace bindings of the same name on the right,
-- if any.

envToTypeEnv :: Env -> TypeEnv
envToTypeEnv env =
    let frameTE frame = Map.map vftype frame
    in Map.unions (Prelude.map frameTE env)

-- | The base type environment, derived from the base environment
-- of the Sifflet language (built-in functions)


baseTypeEnv :: TypeEnv
baseTypeEnv = Map.map vftype (head baseEnv)

vftype :: Value -> TypeScheme
vftype (VFun f) = ftype f
vftype _ = error "vftype: non-function value type"

ftype :: Function -> TypeScheme
ftype f = 
    let ft = uncurry typeFunction (functionArgResultTypes f)
        -- treat /all/ type variables in ft as schematic
        scvars = typeVarsIn ft
    in TypeScheme scvars ft