(*type system*)
type expType = TyVar of int | TyConst of (string * expType list)
type substitution = int -> expType


(*unification algorithm*)
(* Problem 1 *)
let rec contains n ty =
  match ty with
    TyVar m -> n=m
  | TyConst(st, typelst) -> List.fold_left (fun xl x -> if xl then xl else contains n x) false typelst
;;

(* Problem 2 *)
let rec substitute ie ty = 
  let n,sub = ie 
  in match ty with
       TyVar m -> if n=m then sub else ty
     | TyConst(st, typelist) -> TyConst(st, List.map (fun t -> substitute ie t) typelist)
;;

(* Problem 3 *)
let rec lift_subst s ty =
  match ty with
    TyVar m -> s m
  | TyConst(st, typelst) ->  TyConst(st, List.map (fun t -> lift_subst s t) typelst)
;;

(* Problem 4 *)
let rec unify eqlst : substitution option =
  let rec addNewEqs lst1 lst2 acc =
    match lst1,lst2 with
      [],[] -> Some acc
    | t::tl, t'::tl' -> addNewEqs tl tl' ((t,t')::acc)
    | _ -> None
  in
  match eqlst with
    [] -> Some((fun n -> TyVar n))
    (* Delete *)
  | (s,t)::eqs when s=t -> unify eqs
    (* Eliminate *)
  | (TyVar(n),t)::eqs when not(contains n t)-> 
      let eqs' = List.map (fun (t1,t2) -> (substitute (n,t) t1 , substitute (n,t) t2)) eqs
      in (match unify eqs' with
           None -> None
         | Some(phi) -> Some(fun m -> if n=m then lift_subst phi t
                                      else phi m))
    (* Orient *)
  | (TyConst(str, tl), TyVar(m))::eqs -> unify ((TyVar(m), TyConst(str, tl))::eqs)
    (* Decompose *)
  | (TyConst(str, tl), TyConst(str', tl'))::eqs when str=str' -> 
      (match (addNewEqs tl tl' eqs) with
        None -> None
      | Some l -> unify l)
    (* Other *)
  | _ -> None
;;


(*-----------------------------------------------*)

(* type printing *)
let rec print_type t =
  let rec print_tlst = function
     []     -> ()
   | t'::[] -> print_type t'
   | t'::ts -> print_type t'; print_string ","; print_tlst ts
  in
    match t with
       TyVar n         -> print_int n
     | TyConst("prod",[t1;t2])->print_string "("; print_type t1;
                                print_string " * "; print_type t2;
                                print_string ")"
     | TyConst("->",[t1;t2])->print_type t1; print_string " -> ";
                              print_type t2
     | TyConst(s,[])   -> print_string s
     | TyConst(s,tlst) -> print_string s; print_string "("; print_tlst tlst;
                          print_string ")";

(*abstract syntax*)
type dec = Valbind of (string * exp) | ValbindRec of (string * exp)
         | Local of (dec * dec) | Seq of (dec * dec)
and exp =
 (*special constants*)
 Bool of bool | Char of char | Int of int | Real of float
              | String of string
 (*pair*)
 | Pair of exp * exp
 (*identifier*)
 | Id of string
 (*let in, application, functions*)
 | LetIn of dec * exp | App of exp * exp | Fn of string * exp
 | IfThenElse of exp * exp * exp

let rec print_exp = function 
   Bool b     -> print_string(if b then "true" else "false")
 | Int i      -> print_int i
 | Real f     -> print_float f
 | String s   -> print_string "\""; print_string s; print_string "\"";
 | Char c     -> print_string "'"; print_char c; print_string "'";
 | Pair(e1,e2)-> print_string "("; print_exp e1; print_string ",";
                 print_exp e2; print_string ")";
 | Id x       -> print_string x
 | LetIn(d,e) -> print_string "let ("; print_dec d; print_string ") in (";
                 print_exp e; print_string ")"
 | App(e1,e2) -> print_string "("; print_exp e1; print_string ")(";
                 print_exp e2; print_string ")"
 | Fn(s,e)    -> print_string "(fn "; print_string s; print_string " => ";
                 print_exp e; print_string ")"
 | IfThenElse(e1,e2,e3)->print_string "if ("; print_exp e1;
                 print_string ") then ("; print_exp e2; 
                 print_string ") else ("; print_exp e3;
                 print_string ")"
and print_dec = function
   Valbind(s,e)-> print_string "val "; print_string s; print_string " = ";
                  print_exp e
 | ValbindRec(s,e)->print_string "val rec "; print_string s;
                  print_string " = "; print_exp e
 | Local(d1,d2)-> print_string "local "; print_dec d1; print_string " in ";
                  print_dec d2
 | Seq(d1,d2)  -> print_dec d1; print_string " "; print_dec d2

(*judgment*) 
type env = (string * expType) list
type judgment_exp = { gamma:env; exp:exp; expType:expType }
type judgment_dec = { gamma1:env; dec:dec; delta2:env }

(*printing*)
let print_env gamma = 
  let rec print_env_aux gamma =
    match gamma with
       []        -> ()
     | (x,y)::xs -> print_string x; print_string "->"; print_type y;
                    match xs with [] -> () | _  -> print_string ",";
                                                   print_env_aux xs
  in
    print_string "{"; print_env_aux gamma; print_string "}"
let print_jexp {gamma=gamma;exp=exp;expType=expType} =
  print_env gamma; print_string " |= "; print_exp exp;
  print_string " : "; print_type expType
let print_jdec {gamma1=gamma1;dec=dec;delta2=delta2} =
  print_env gamma1; print_string " |= "; print_dec dec;
  print_string " -> "; print_env delta2

(*proof tree*)
type proof = ExpProof of judgment_exp * proof list
           | DecProof of judgment_dec * proof list

(*constraint list*)
type consList = (expType * expType) list

(*fresh type variable*)
let nxt = ref 0
let fresh () = nxt := !nxt + 1; TyVar(!nxt)

(*fresh type*)
let fresh_type ty =
  let rec get_subst f = function
     TyVar n -> if f n = TyVar n then
                  let tyFresh = fresh() 
                  in 
                    (fun x -> if x=n then tyFresh else f x)
                else f
   | TyConst(s,lst) -> List.fold_left get_subst f lst
  in
    lift_subst (get_subst (fun x -> TyVar x) ty) ty

(*environment operations*)
let make_env x y = ([(x,y)]:env)
let rec lookup_env (gamma:env) x =
  match gamma with
     []        -> None
   | (y,z)::ys -> if x = y then Some z else lookup_env ys x
let sum_env (delta:env) (gamma:env) = ((delta@gamma):env)
let ins_env (gamma:env) x y = sum_env (make_env x y) gamma

(*top level environment*)
let pervasives = 
  let basic_type s = TyConst(s, [])
  in let unaryfn_type s = TyConst("->",[basic_type s;basic_type s])
  in let binaryfn_type s = 
    TyConst("->",[TyConst("prod",[basic_type s;basic_type s]);
    basic_type s])
  in
  [
    ("true", basic_type "bool");
    ("false", basic_type "bool");
    ("and", binaryfn_type "bool");
    ("or", binaryfn_type "bool");
    ("not", unaryfn_type "bool");
    ("+", binaryfn_type "int");
    ("-", binaryfn_type "int");
    ("*", binaryfn_type "int");
    ("/", binaryfn_type "int");
    ("+.", binaryfn_type "real");
    ("-.", binaryfn_type "real");
    ("*.", binaryfn_type "real");
    ("**", binaryfn_type "real");
    ("nil", TyConst("list",[fresh()]));
    ("::", let tyList = TyConst("list",[TyVar 0]) in
        TyConst("->",[TyConst("prod",[TyVar 0;tyList]);tyList]));
    ("head", TyConst("->",[TyConst("list",[TyVar 0]);TyVar 0]));
    ("tail", TyConst("->",[TyConst("list",[TyVar 0]);TyVar 0]));
    ("=", TyConst("->",[TyConst("prod",[TyVar 0;TyVar 0]);basic_type "bool"]))
  ]


(*applying a substitution to a proof*)
let rec subst_proof f = function
   ExpProof({gamma=gamma;exp=exp;expType=expType},assum) -> 
   ExpProof({gamma=subst_env f gamma;exp=exp;expType=lift_subst f expType},
           List.map (subst_proof f) assum)
 | DecProof({gamma1=gamma1;dec=dec;delta2=delta2},assum) ->
   DecProof({gamma1=subst_env f gamma1;dec=dec;delta2=subst_env f delta2},
           List.map (subst_proof f) assum)
     
(*applying a substitution to an environment*)
and subst_env f gamma =
  match gamma with
     []        -> []
   | (x,y)::xs -> (x,lift_subst f y)::subst_env f xs

(*proof printing*)
let print_proof p =
  let depth_max = 10 in
  let rec print_struts = function
     []    -> ()
   | x::[] -> print_string (if x then "|-" else "|-")
   | x::xs -> print_string (if x then "  " else "| "); print_struts xs
  in let rec print_proof_aux p depth lst =
    print_newline(); print_string "  "; print_struts lst;
    if (depth > 0) then print_char('-'); 
    let assum = (match p with
        ExpProof(jdg,assum) -> print_jexp jdg; assum
      | DecProof(jdg,assum) -> print_jdec jdg; assum)
    in if depth <= depth_max then
      print_assum depth lst assum
  and print_assum depth lst assum =
    match assum with 
       []     -> ()
     | p'::ps -> print_proof_aux p' (depth + 1) (lst@[ps=[]]);
                 print_assum depth lst ps
  in
    print_proof_aux p 0 []; print_newline(); print_newline()


let get_ty = function
   None       -> raise(Failure "None")
 | Some(ty,p) -> ty

let get_proof = function
   None       -> raise(Failure "None")
 | Some(ty,p) -> p