type lambda_exp = Var of string
  | Lambda of string * lambda_exp
  | Appl of lambda_exp * lambda_exp
  | IntConstant of int
  | BoolConstant of bool
  | Empty
  | Cons of lambda_exp * lambda_exp
  | Plus of lambda_exp * lambda_exp
  | ITE of lambda_exp * lambda_exp * lambda_exp;;

let rec union l1 l2 =
  match (l1,l2) with
  | ([],_) -> l2
  | (_,[]) -> l1
  | (h1::t1,h2::t2) -> if (h1 < h2) then h1::(union t1 l2)
                       else if (h1 > h2) then h2::(union l1 t2)
                       else h1::(union t1 t2);;

let rec remove x l =
  match l with
    [] -> []
  | h::t -> if (x=h) then t else h::(remove x t);;

let rec contains x l =
  match l with
    [] -> false
  | h::t -> if (x=h) then true else contains x t;;

let rec fv exp =
  match exp with
    Var v -> [v]
  | Lambda(v,exp1) -> remove v (fv exp1)
  | Appl(exp1,exp2) -> union (fv exp1) (fv exp2)
  | IntConstant(i) -> []
  | BoolConstant(b) -> []
  | Empty -> []
  | Cons(exp1,exp2) -> union (fv exp1) (fv exp2)
  | Plus(exp1,exp2) -> union (fv exp1) (fv exp2)
  | ITE(exp1,exp2,exp3) -> union (union (fv exp1) (fv exp2)) (fv exp3);;

let longer s1 s2 =
  if ((String.length s1) > (String.length s2)) then s1 else s2;;

let rec longest exp =
  match exp with
    Var s -> s
  | Lambda(s,exp1) -> longer s (longest exp1)
  | Appl(exp1,exp2) -> longer (longest exp1) (longest exp2)
  | IntConstant(i) -> ""
  | BoolConstant(b) -> ""
  | Empty -> ""
  | Cons(exp1,exp2) -> longer (longest exp1) (longest exp2)
  | Plus(exp1,exp2) -> longer (longest exp1) (longest exp2)
  | ITE(exp1,exp2,exp3) -> longer (longer (longest exp1) (longest exp2))
                              (longest exp3);;

let fresh exp = (longest exp)^"f";;

let rec subst m x n =
  match m with
    Var(y) -> if (y = x) then n else m
  | Appl(m1,m2) -> Appl(subst m1 x n, subst m2 x n)
  | Lambda(y,m1) -> if (y = x) then m
                    else if (contains y (fv n)) then
                        let fr = fresh (Appl(m,n)) in
                            Lambda(fr,subst (subst m1 y (Var fr)) x n)
                    else Lambda(y,subst m1 x n)
  | IntConstant(i) -> m
  | BoolConstant(b) -> m
  | Empty -> m
  | Cons(m1,m2) -> Cons(subst m1 x n,subst m2 x n)
  | Plus(m1,m2) -> Plus(subst m1 x n,subst m2 x n)
  | ITE(m1,m2,m3) -> ITE(subst m1 x n,subst m2 x n,subst m3 x n);;

let rec append l1 l2 =
  match l1 with
    [] -> l2
  | h::t -> h::(append t l2);;

type ml_type =
   Int
 | Bool
 | List of ml_type
 | TypeVar of int
 | Arrow of ml_type * ml_type;;

(* restituisce il tipo dell'espressione, un insieme di
   vincoli generati per il tipaggio e il prossimo numero di
   variabile di tipo disponibile *)

let rec inferAux exp env next =
  match exp with
    IntConstant(i) -> (Int,[],next)
  | BoolConstant(b) -> (Bool,[],next)
  | Var(s) -> (env(s),[],next)
  | Plus(m1,m2) ->
     let (t1,c1,n1) = inferAux m1 env next
     in let (t2,c2,n2) = inferAux m2 env n1
     in (Int,(t1,Int) :: (t2,Int) :: (append c1 c2),n2)
  | ITE(m1,m2,m3) ->
     let (t1,c1,n1) = inferAux m1 env next
     in let (t2,c2,n2) = inferAux m2 env n1
     in let (t3,c3,n3) = inferAux m3 env n2
     in (t2,(t1,Bool) :: (t2,t3) :: (append (append c1 c2) c3),n3)
  | Empty -> (List(TypeVar(next)),[],next + 1)
  | Cons(h,t) ->
     let (t1,c1,n1) = inferAux h env next
     in let (t2,c2,n2) = inferAux t env n1
     in (List(t1),(t2,List(t1)) :: (append c1 c2),n2)
  | Lambda(s,m) ->
     let newEnv = function
                     s1 -> if (s1 = s) then TypeVar(next)
                           else env s1
     in let (t,c,n) = inferAux m newEnv (next + 1)
     in (Arrow(TypeVar(next),t),c,n)
  | Appl(m1,m2) ->
     let (t1,c1,n1) = inferAux m1 env (next + 1)
     in let (t2,c2,n2) = inferAux m2 env n1
     in let t = TypeVar(next)
     in (t,(t1,Arrow(t2,t)) :: (append c1 c2),n2);;

exception UnificationException;;

let rec occurs_in_type n t =
  match t with
    Int -> false
  | Bool -> false
  | List(elements) -> occurs_in_type n elements
  | TypeVar(n1) -> (n1 = n)
  | Arrow(left,right) -> (occurs_in_type n left) or (occurs_in_type n right);;

let rec occurs n l =
  match l with
    [] -> false
  | (left,right) :: t -> (occurs_in_type n left) or (occurs_in_type n right)
                                or (occurs n t);;

let rec substitute_in_type n t1 t2 =
  match t2 with
    Int -> Int
  | Bool -> Bool
  | List(elements) -> List(substitute_in_type n t1 elements)
  | TypeVar(n1) -> if (n1 = n) then t1 else t2
  | Arrow(t3,t4) ->
       Arrow(substitute_in_type n t1 t3,substitute_in_type n t1 t4);;

let rec substitute n t l =
  match l with
    [] -> []
  | (left,right) :: t1 ->
      (substitute_in_type n t left,substitute_in_type n t right)
                              :: (substitute n t t1);;

let rec unify c =
  match c with
    [] -> []
  | h::t -> unifyAux [] h t
and unifyAux past current future =
  match current with
    (left,right) ->
      if (left = right) then match future with
          [] -> past
        | h :: t -> unifyAux past h t
      else match (left,right) with
        (Arrow(t1,t2),Arrow(t3,t4)) ->
           unifyAux past (t1,t3) ((t2,t4) :: future)
      | (List(t1),List(t2)) ->
           unifyAux past (t1,t2) future
      | (TypeVar(n),_) ->
          if (occurs_in_type n right) then raise UnificationException
          else if ((occurs n past) or (occurs n future)) then
                  unify (current :: (substitute n right (append past future)))
               else (match future with
                   [] -> current :: past
                 | h1 :: t1 -> unifyAux (current :: past) h1 t1)
      | (_,TypeVar(n1)) ->
          unifyAux past (TypeVar(n1),left) future
      | _ -> raise UnificationException;;

exception UndefinedVariable;;

let rec apply t c =
  match c with
    [] -> t
  | (TypeVar(n),right) :: t1 -> apply (substitute_in_type n right t) t1;;
 
let infer exp =
  let (t,c,_) = inferAux exp (function s -> raise UndefinedVariable) 0
  in (t,c);; (* apply t (unify c);; *)

let exp1 = Lambda("z",Plus(IntConstant(3),Var "z"));;
let exp2 = Lambda("z",Lambda("y",Lambda("v",ITE(Var "z",Var "v",Var "z"))));;
let exp3 = Lambda("x",Plus(Var("x"),IntConstant(3)));;
let exp4 = Appl(Lambda("x",Plus(Var("x"),IntConstant(3))),IntConstant(4));;
let exp5 = Appl(Lambda("x",Plus(Var("x"),IntConstant(3))),BoolConstant(true));;
let omega = Appl(Lambda("x",Appl(Var "x",Var "x")),
                 Lambda("x",Appl(Var "x",Var "x")));;
let exp6 = Lambda("x",Cons(IntConstant(3),Var "x"));;
let exp7 = Empty;;
let exp8 = Lambda("x",Lambda("y",ITE(Var "x",Empty,Cons(Var "y",Cons(Var "x",Empty)))));;
let exp9 = Lambda("x",Lambda("y",Lambda("z",Appl(Var "x",ITE(Var "y",Appl(Var "z",Var "x"),IntConstant(0))))));;