type expr = Int of int 
          | Add of expr * expr 
          | Mul of expr * expr
          | Lam of (expr -> expr) 
          | App of expr * expr
          | Var of string
          | Let of string * expr * expr

type typ = Integer | TVar of string | Arrow of typ * typ

type stat = TypDecl of string * typ
          | FuncDefn of string * expr

type prog = stat list

let lookup x env = env x
let add x y env = fun z -> if z = x then y else env z

exception LookingFor of string
let emptyVE z = raise (LookingFor z)
let emptyTE z = raise (LookingFor z)

let getint = function
    | Int x -> x
    | Add _ | Mul _ | App _ | Lam _ | Var _ | Let _ -> failwith "hey, I didn't get an integer?!?!"

let rec getfun env = function
    | Lam f -> f
    | Var _ -> failwith "trying to apply an unreduced variable"
    | Add _ | Mul _ | Int _ -> failwith "applying a non-function"
    | Let _ -> failwith "trying to apply an unreduced Let"
    | App _ -> failwith "trying to apply an unreduced application"

let rec eval2 env = function
    | Int x -> Int x
    | Add (x,y) -> 
        let xx = getint (eval2 env x) and 
            yy = getint (eval2 env y) in Int (xx + yy)
    | Mul (x,y) ->
        let xx = getint (eval2 env x) and 
            yy = getint (eval2 env y) in Int (xx * yy)
    | App (f, x) ->
        let g = getfun env (eval2 env f) and 
            y = eval2 env x in eval2 env (g y)
    | Lam f -> Lam f
    | Var x -> lookup x env
    | Let (x,y,z) -> eval2 (add x y env) z

let decls (ve, te) = function
    | TypDecl (n,t) -> (ve, add n t te)
    | FuncDefn (n,f) -> (add n f ve, te)

let defns prog = List.fold_left decls (emptyVE, emptyTE) prog

(* does not deal with recursive definitions (will loop) *)
let rec check env prog = List.iter (checkeach env) prog
and checkeach env = function
    | TypDecl (n,t) -> let _ = check_defined env n in check_typ env t
    | FuncDefn (n,f) -> let _ = check_typed env n in check_func env f
and check_defined (ve,_) n = 
    try lookup n ve
    with LookingFor _ -> failwith (n^" is not defined")
and check_typed (_,te) n = 
    try lookup n te
    with LookingFor _ -> failwith (n^" has no type")
and check_typ env = function
    | Integer -> ()
    | TVar t  -> begin
       try (let def = lookup t (snd env) in check_typ env def)
       with LookingFor _ -> failwith (t^" is an undeclared type")
    end
    | Arrow (a,b) -> check_typ env a; check_typ env b
and check_func env = function
    | Int _ -> ()
    | Add (x,y) -> check_func env x; check_func env y;
    | Mul (x,y) -> check_func env x; check_func env y;
    | App (f,y) -> check_func env f; check_func env y;
    | Lam _ -> ()
    | Var x -> begin
        try ignore (lookup x (snd env))
        with LookingFor _ -> failwith (x^" is a variable without a type")
    end
    | Let (x,y,z) -> 
      check_func (fst env, add x (TVar x) (snd env)) z

let rec typecheck (ve,te) prog = List.iter (checkall (ve,te)) prog
and checkall (ve,te) = function
    | TypDecl (n,t) -> ()
    | FuncDefn (n,f) -> 
        let typ = lookup n te in
        check_type (ve, te) f typ 
and check_type (ve,te) f typ = match f with
    | Int _ -> if typ = Integer then () else failwith "foo"
    | Add (x,y) ->
        check_type (ve,te) x Integer; 
        check_type (ve,te) y Integer; 
        if typ = Integer then () else failwith "foo"
    | Mul (x,y) ->
        check_type (ve,te) x Integer; 
        check_type (ve,te) y Integer; 
        if typ = Integer then () else failwith "foo"
    | App (f,y) -> ()
    | Lam _ -> ()
    | Var x -> ()
    | Let (x,y,z) -> ()

let eval_stat prog =
    let (ve,te) = defns prog in
    begin
        check (ve,te) prog;
        typecheck (ve,te) prog;
        eval2 ve (lookup "main" ve) 
    end
;;

let test1 = Add( Int 1, Mul(Int 2, Int 3)) (* 1 + (2*3) *)
let test2 = App ((Lam (fun x -> x)),test1)
let test3 = Lam (fun x -> Add(x, Int 1))
let test4 = App (test3, test1)
let test5 = App (Int 1, Int 2)
let test6 = Let ("x",Int 5, Add(Var "x", Int 3))
let test7 = Let ("y",test3, App(Var "y", test1))

let res1 = eval2 emptyVE test1
let res4 = eval2 emptyVE test4
let res6 = eval2 emptyVE test6
let res7 = eval2 emptyVE test7
;;
let addt = TypDecl("add", Arrow (Integer, Arrow(Integer, Integer)))
let addf = FuncDefn("add", Lam (fun x -> Lam (fun y -> Add(x,y))))
let maint = TypDecl("main", Integer)
let mainf = FuncDefn("main", App(App(Var "add", Int 2), Int 3))
;;

let ans = eval_stat [addf; addt; maint; mainf];;